UNPKG

@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
#!/usr/bin/env tsx 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