@stackmemoryai/stackmemory
Version:
Project-scoped memory for AI coding tools. Durable context across sessions with MCP integration, frames, smart retrieval, Claude Code skills, and automatic hooks.
390 lines (389 loc) • 13.4 kB
JavaScript
import { fileURLToPath as __fileURLToPath } from 'url';
import { dirname as __pathDirname } from 'path';
const __filename = __fileURLToPath(import.meta.url);
const __dirname = __pathDirname(__filename);
import { createClient } from "redis";
import Database from "better-sqlite3";
import { v4 as uuidv4 } from "uuid";
import chalk from "chalk";
import ora from "ora";
import { join } from "path";
import { existsSync, mkdirSync } from "fs";
import { RailwayOptimizedStorage } from "../src/core/storage/railway-optimized-storage.js";
import { ConfigManager } from "../src/core/config/config-manager.js";
import { TraceType } from "../src/core/trace/types.js";
import dotenv from "dotenv";
function getEnv(key, defaultValue) {
const value = process.env[key];
if (value === void 0) {
if (defaultValue !== void 0) return defaultValue;
throw new Error(`Environment variable ${key} is required`);
}
return value;
}
function getOptionalEnv(key) {
return process.env[key];
}
dotenv.config();
async function testRedisConnection() {
const spinner = ora("Testing Redis connection...").start();
try {
const redisUrl = process.env["REDIS_URL"] || "redis://localhost:6379";
console.log(
chalk.gray(
` Using Redis URL: ${redisUrl.replace(/:[^:@]+@/, ":****@")})`
)
);
const client = createClient({ url: redisUrl });
await client.connect();
const testKey = "test:connection";
await client.set(testKey, "connected");
const result = await client.get(testKey);
await client.del(testKey);
if (result === "connected") {
spinner.succeed(`Redis connected successfully at ${redisUrl}`);
const info = await client.info("memory");
const memoryUsed = info.match(/used_memory_human:(\S+)/)?.[1];
console.log(chalk.gray(` Memory used: ${memoryUsed || "unknown"}`));
} else {
spinner.fail("Redis connection test failed");
return false;
}
await client.quit();
return true;
} catch (error) {
spinner.fail(`Redis connection failed: ${error}`);
return false;
}
}
function createMockTrace(index) {
const now = Date.now() - index * 60 * 60 * 1e3;
const tools = [
{
id: uuidv4(),
tool: "search",
timestamp: now,
arguments: { query: `test query ${index}` },
filesAffected: ["src/test.ts", "src/index.ts"]
},
{
id: uuidv4(),
tool: "read",
timestamp: now + 1e3,
arguments: { file: "src/test.ts" },
result: "file contents"
},
{
id: uuidv4(),
tool: "edit",
timestamp: now + 2e3,
arguments: { file: "src/test.ts", changes: "some changes" },
filesAffected: ["src/test.ts"]
},
{
id: uuidv4(),
tool: "test",
timestamp: now + 3e3,
arguments: { command: "npm test" },
result: "tests passed"
}
];
const trace = {
id: uuidv4(),
type: TraceType.SEARCH_DRIVEN,
tools,
score: 0.5 + Math.random() * 0.5,
// Random score 0.5-1.0
summary: `Test trace #${index}: Search-driven modification`,
metadata: {
startTime: now,
endTime: now + 4e3,
filesModified: ["src/test.ts"],
errorsEncountered: index % 3 === 0 ? ["Test error"] : [],
decisionsRecorded: index % 2 === 0 ? ["Use async pattern"] : [],
causalChain: index % 3 === 0
}
};
return trace;
}
async function testStorageOperations() {
console.log(chalk.blue("\n\u{1F4E6} Testing Storage Operations"));
console.log(chalk.gray("\u2501".repeat(50)));
const dbDir = join(process.cwd(), ".stackmemory");
if (!existsSync(dbDir)) {
mkdirSync(dbDir, { recursive: true });
}
const dbPath = join(dbDir, "test-context.db");
const db = new Database(dbPath);
db.exec(`
CREATE TABLE IF NOT EXISTS traces (
id TEXT PRIMARY KEY,
type TEXT NOT NULL,
score REAL NOT NULL,
summary TEXT NOT NULL,
start_time INTEGER NOT NULL,
end_time INTEGER NOT NULL,
frame_id TEXT,
user_id TEXT,
files_modified TEXT,
errors_encountered TEXT,
decisions_recorded TEXT,
causal_chain INTEGER,
compressed_data TEXT,
created_at INTEGER DEFAULT (unixepoch())
)
`);
db.exec(`
CREATE TABLE IF NOT EXISTS tool_calls (
id TEXT PRIMARY KEY,
trace_id TEXT NOT NULL,
tool TEXT NOT NULL,
arguments TEXT,
timestamp INTEGER NOT NULL,
result TEXT,
error TEXT,
files_affected TEXT,
duration INTEGER,
sequence_number INTEGER NOT NULL,
FOREIGN KEY (trace_id) REFERENCES traces(id) ON DELETE CASCADE
)
`);
db.exec(`
CREATE TABLE IF NOT EXISTS storage_tiers (
trace_id TEXT PRIMARY KEY,
tier TEXT NOT NULL,
location TEXT NOT NULL,
original_size INTEGER,
compressed_size INTEGER,
compression_ratio REAL,
access_count INTEGER DEFAULT 0,
last_accessed INTEGER DEFAULT (unixepoch()),
created_at INTEGER DEFAULT (unixepoch()),
migrated_at INTEGER,
score REAL,
migration_score REAL,
metadata TEXT,
FOREIGN KEY (trace_id) REFERENCES traces(id) ON DELETE CASCADE
)
`);
db.exec(`
CREATE INDEX IF NOT EXISTS idx_storage_tier ON storage_tiers(tier);
CREATE INDEX IF NOT EXISTS idx_storage_created ON storage_tiers(created_at);
CREATE INDEX IF NOT EXISTS idx_storage_accessed ON storage_tiers(last_accessed);
`);
const configManager = new ConfigManager();
const storage = new RailwayOptimizedStorage(db, configManager, {
redis: {
url: process.env["REDIS_URL"],
ttlSeconds: 24 * 60 * 60,
// 24 hours
maxMemory: "100mb"
}
});
const traces = [];
const results = [];
console.log(chalk.yellow("\n\u270F\uFE0F Creating and storing test traces..."));
for (let i = 0; i < 10; i++) {
const trace = createMockTrace(i);
traces.push(trace);
const spinner = ora(`Storing trace #${i + 1}...`).start();
try {
const insertTrace = db.prepare(`
INSERT INTO traces (
id, type, score, summary, start_time, end_time,
files_modified, errors_encountered, decisions_recorded, causal_chain,
created_at
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
`);
insertTrace.run(
trace.id,
trace.type,
trace.score,
trace.summary,
trace.metadata.startTime,
trace.metadata.endTime,
JSON.stringify(trace.metadata.filesModified || []),
JSON.stringify(trace.metadata.errorsEncountered || []),
JSON.stringify(trace.metadata.decisionsRecorded || []),
trace.metadata.causalChain ? 1 : 0,
Date.now()
);
const insertToolCall = db.prepare(`
INSERT INTO tool_calls (
id, trace_id, tool, arguments, timestamp, result, files_affected, sequence_number
) VALUES (?, ?, ?, ?, ?, ?, ?, ?)
`);
trace.tools.forEach((tool, index) => {
insertToolCall.run(
tool.id,
trace.id,
tool.tool,
JSON.stringify(tool.arguments),
tool.timestamp,
tool.result || null,
JSON.stringify(tool.filesAffected || []),
index
);
});
const tier = await storage.storeTrace(trace);
results.push({ id: trace.id, tier, score: trace.score });
const tierIcon = tier === "hot" ? "\u{1F525}" : tier === "warm" ? "\u2601\uFE0F" : "\u2744\uFE0F";
spinner.succeed(
`Trace #${i + 1} stored in ${tierIcon} ${tier} tier (score: ${trace.score.toFixed(2)})`
);
} catch (error) {
spinner.fail(`Failed to store trace #${i + 1}: ${error}`);
}
await new Promise((resolve) => setTimeout(resolve, 100));
}
console.log(chalk.yellow("\n\u{1F50D} Testing trace retrieval..."));
for (let i = 0; i < 3; i++) {
const result = results[i];
const spinner = ora(
`Retrieving trace ${result.id.substring(0, 8)}...`
).start();
try {
const retrieved = await storage.retrieveTrace(result.id);
if (retrieved) {
spinner.succeed(
`Retrieved from ${result.tier} tier: ${retrieved.summary}`
);
} else {
spinner.fail("Trace not found");
}
} catch (error) {
spinner.fail(`Retrieval failed: ${error}`);
}
}
console.log(chalk.yellow("\n\u{1F4CA} Storage Statistics:"));
const stats = storage.getStorageStats();
console.log(chalk.gray("\u2501".repeat(50)));
for (const tier of stats.byTier) {
const icon = tier.tier === "hot" ? "\u{1F525}" : tier.tier === "warm" ? "\u2601\uFE0F" : "\u2744\uFE0F";
console.log(`${icon} ${chalk.bold(tier.tier.toUpperCase())} Tier:`);
console.log(` Traces: ${tier.count}`);
console.log(` Original Size: ${formatBytes(tier.total_original || 0)}`);
console.log(` Compressed: ${formatBytes(tier.total_compressed || 0)}`);
if (tier.avg_compression) {
console.log(
` Compression: ${(tier.avg_compression * 100).toFixed(1)}%`
);
}
}
console.log(chalk.yellow("\n\u{1F525} Testing Redis Hot Tier..."));
const redisClient = createClient({ url: process.env["REDIS_URL"] });
await redisClient.connect();
const keys = await redisClient.keys("trace:*");
console.log(` Traces in Redis: ${chalk.green(keys.length)}`);
const byScore = await redisClient.zCard("traces:by_score");
const byTime = await redisClient.zCard("traces:by_time");
console.log(` Score index: ${chalk.green(byScore)} entries`);
console.log(` Time index: ${chalk.green(byTime)} entries`);
const topTraces = await redisClient.zRangeWithScores(
"traces:by_score",
-3,
-1
);
if (topTraces.length > 0) {
console.log(chalk.yellow("\n\u{1F3C6} Top Traces by Score:"));
for (const trace of topTraces.reverse()) {
console.log(
` ${trace.value.substring(0, 8)}... - Score: ${trace.score.toFixed(3)}`
);
}
}
const memInfo = await redisClient.memoryUsage("trace:" + results[0]?.id);
if (memInfo) {
console.log(chalk.yellow("\n\u{1F4BE} Memory Usage:"));
console.log(` Sample trace memory: ${formatBytes(memInfo)}`);
console.log(` Estimated total: ${formatBytes(memInfo * keys.length)}`);
}
await redisClient.quit();
console.log(chalk.yellow("\n\u{1F504} Testing tier migration..."));
const migrationResults = await storage.migrateTiers();
console.log(
` Hot \u2192 Warm: ${chalk.yellow(migrationResults.hotToWarm)} traces`
);
console.log(
` Warm \u2192 Cold: ${chalk.cyan(migrationResults.warmToCold)} traces`
);
if (migrationResults.errors.length > 0) {
console.log(chalk.red(` Errors: ${migrationResults.errors.length}`));
}
db.close();
console.log(chalk.green("\n\u2705 Storage tests completed successfully!"));
}
function formatBytes(bytes) {
if (bytes === 0) return "0 B";
const k = 1024;
const sizes = ["B", "KB", "MB", "GB"];
const i = Math.floor(Math.log(bytes) / Math.log(k));
return `${(bytes / Math.pow(k, i)).toFixed(2)} ${sizes[i]}`;
}
async function main() {
console.log(chalk.blue.bold("\n\u{1F9EA} StackMemory Redis Storage Test\n"));
const redisConnected = await testRedisConnection();
if (!redisConnected) {
console.log(chalk.red("\n\u274C Cannot proceed without Redis connection"));
console.log(chalk.yellow("\nTo fix:"));
console.log("1. Ensure Redis is running");
console.log("2. Check REDIS_URL in .env file");
console.log("3. For Railway: Ensure Redis addon is provisioned");
process.exit(1);
}
await testStorageOperations();
console.log(chalk.blue("\n\u{1F3AE} Interactive Test"));
console.log(chalk.gray("\u2501".repeat(50)));
console.log(
chalk.cyan("You can now use the CLI to interact with the stored traces:")
);
console.log();
console.log(
" " + chalk.white("stackmemory storage status") + " - View storage statistics"
);
console.log(
" " + chalk.white("stackmemory storage migrate") + " - Migrate traces between tiers"
);
console.log(
" " + chalk.white("stackmemory storage retrieve <id>") + " - Retrieve a specific trace"
);
console.log();
console.log(chalk.gray("Trace IDs from this test:"));
const dbPath = join(process.cwd(), ".stackmemory", "test-context.db");
const db = new Database(dbPath);
db.exec(`
CREATE TABLE IF NOT EXISTS storage_tiers (
trace_id TEXT PRIMARY KEY,
tier TEXT NOT NULL,
location TEXT NOT NULL,
original_size INTEGER,
compressed_size INTEGER,
compression_ratio REAL,
access_count INTEGER DEFAULT 0,
last_accessed INTEGER DEFAULT (unixepoch()),
created_at INTEGER DEFAULT (unixepoch()),
migrated_at INTEGER,
score REAL,
migration_score REAL,
metadata TEXT
)
`);
const recentTraces = db.prepare(
`
SELECT trace_id, tier FROM storage_tiers
ORDER BY created_at DESC LIMIT 3
`
).all();
for (const trace of recentTraces) {
const tierIcon = trace.tier === "hot" ? "\u{1F525}" : trace.tier === "warm" ? "\u2601\uFE0F" : "\u2744\uFE0F";
console.log(` ${tierIcon} ${trace.trace_id}`);
}
db.close();
console.log(chalk.green("\n\u2728 Test complete!"));
}
main().catch((error) => {
console.error(chalk.red("Test failed:"), error);
process.exit(1);
});
//# sourceMappingURL=test-redis-storage.js.map