diff --git a/agent-memory-mesh/src/config.ts b/agent-memory-mesh/src/config.ts index 9c507dc..8a6b99b 100644 --- a/agent-memory-mesh/src/config.ts +++ b/agent-memory-mesh/src/config.ts @@ -32,6 +32,8 @@ export interface MemoryConfig { * Leave empty for rule-based synthesis (no LLM required). */ consolidationModel: string; + /** Directory for memory state snapshots. */ + snapshotsDir: string; } export function loadConfig(): MemoryConfig { @@ -50,5 +52,6 @@ export function loadConfig(): MemoryConfig { port: Number(env("MEMORY_PORT", "8377")), apiKey: env("MEMORY_API_KEY", ""), consolidationModel: env("CONSOLIDATION_MODEL", ""), + snapshotsDir: env("SNAPSHOTS_DIR") || join(base, "snapshots"), }; } diff --git a/agent-memory-mesh/src/memory/snapshots.ts b/agent-memory-mesh/src/memory/snapshots.ts new file mode 100644 index 0000000..b576bf6 --- /dev/null +++ b/agent-memory-mesh/src/memory/snapshots.ts @@ -0,0 +1,110 @@ +import { + readFileSync, + writeFileSync, + existsSync, + mkdirSync, + readdirSync, + copyFileSync, + rmSync, +} from "node:fs"; +import { join } from "node:path"; +import { randomUUID } from "node:crypto"; + +export interface SnapshotManifest { + id: string; + createdAt: string; + label?: string; + files: { + workMemoryPath: string; + graphPath: string; + feedbackPath: string; + workMemoryCopy: string; + graphCopy: string; + feedbackCopy: string; + }; +} + +export class SnapshotStore { + constructor(private snapshotsDir: string) { + mkdirSync(snapshotsDir, { recursive: true }); + } + + async create(opts: { + workMemoryPath: string; + graphPath: string; + feedbackPath: string; + label?: string; + }): Promise { + const id = randomUUID(); + const dir = join(this.snapshotsDir, id); + mkdirSync(dir, { recursive: true }); + + const safeCopy = (src: string, dest: string) => { + if (existsSync(src)) copyFileSync(src, dest); + else writeFileSync(dest, "[]"); + }; + + const wCopy = join(dir, "work-memory.json"); + const gCopy = join(dir, "graph.json"); + const fCopy = join(dir, "feedback.json"); + safeCopy(opts.workMemoryPath, wCopy); + safeCopy(opts.graphPath, gCopy); + safeCopy(opts.feedbackPath, fCopy); + + const manifest: SnapshotManifest = { + id, + createdAt: new Date().toISOString(), + label: opts.label, + files: { + workMemoryPath: opts.workMemoryPath, + graphPath: opts.graphPath, + feedbackPath: opts.feedbackPath, + workMemoryCopy: wCopy, + graphCopy: gCopy, + feedbackCopy: fCopy, + }, + }; + writeFileSync(join(dir, "manifest.json"), JSON.stringify(manifest, null, 2)); + return manifest; + } + + list(): SnapshotManifest[] { + const subdirs = readdirSync(this.snapshotsDir, { withFileTypes: true }) + .filter((d) => d.isDirectory()) + .map((d) => join(this.snapshotsDir, d.name, "manifest.json")) + .filter((p) => existsSync(p)) + .map((p) => JSON.parse(readFileSync(p, "utf8")) as SnapshotManifest); + return subdirs.sort((a, b) => b.createdAt.localeCompare(a.createdAt)); + } + + get(id: string): SnapshotManifest | undefined { + const manifestPath = join(this.snapshotsDir, id, "manifest.json"); + if (!existsSync(manifestPath)) return undefined; + return JSON.parse(readFileSync(manifestPath, "utf8")); + } + + async restore(id: string): Promise { + const manifest = this.get(id); + if (!manifest) throw new Error(`Snapshot ${id} not found`); + // Auto-snapshot current state before restoring + await this.create({ + workMemoryPath: manifest.files.workMemoryPath, + graphPath: manifest.files.graphPath, + feedbackPath: manifest.files.feedbackPath, + label: `pre-restore-${id.slice(0, 8)}`, + }); + if (existsSync(manifest.files.workMemoryCopy)) + copyFileSync(manifest.files.workMemoryCopy, manifest.files.workMemoryPath); + if (existsSync(manifest.files.graphCopy)) + copyFileSync(manifest.files.graphCopy, manifest.files.graphPath); + if (existsSync(manifest.files.feedbackCopy)) + copyFileSync(manifest.files.feedbackCopy, manifest.files.feedbackPath); + } + + delete(id: string): boolean { + const dir = join(this.snapshotsDir, id); + if (!existsSync(dir)) return false; + rmSync(dir, { recursive: true, force: true }); + return true; + } +} diff --git a/agent-memory-mesh/src/service/engine.ts b/agent-memory-mesh/src/service/engine.ts index 3e39b20..9280161 100644 --- a/agent-memory-mesh/src/service/engine.ts +++ b/agent-memory-mesh/src/service/engine.ts @@ -17,8 +17,10 @@ import { Consolidator, type ConsolidationResult } from "../memory/consolidator.j import { ContextGraph, type ContextGraphEntity, type ContextGraphEdge, type EntityType, type NeighborResult } from "../memory/context-graph.js"; import { RetrievalPolicyStore, applyRecencyBoost, type RetrievalPolicy } from "../memory/retrieval-policy.js"; import { FeedbackStore, applyFeedbackScoring, type FeedbackSignal, type NoteFeedback } from "../memory/feedback.js"; +import { SnapshotStore, type SnapshotManifest } from "../memory/snapshots.js"; +import { metrics, type MetricsSnapshot } from "./metrics.js"; -export { type WorkMemoryEntry, type WorkMemoryQuery, type ConsolidationResult, type ContextGraphEntity, type ContextGraphEdge, type EntityType, type NeighborResult, type RetrievalPolicy, type FeedbackSignal, type NoteFeedback }; +export { type WorkMemoryEntry, type WorkMemoryQuery, type ConsolidationResult, type ContextGraphEntity, type ContextGraphEdge, type EntityType, type NeighborResult, type RetrievalPolicy, type FeedbackSignal, type NoteFeedback, type SnapshotManifest, type MetricsSnapshot }; export interface WikiSummary { entityId: string; @@ -41,6 +43,7 @@ export class MemoryEngine { private graph: ContextGraph; private policies: RetrievalPolicyStore; private feedback: FeedbackStore; + private snapshots: SnapshotStore; private opened = false; constructor(private cfg: MemoryConfig) { @@ -58,6 +61,7 @@ export class MemoryEngine { }); this.policies = new RetrievalPolicyStore(cfg.policiesPath); this.feedback = new FeedbackStore(cfg.feedbackPath); + this.snapshots = new SnapshotStore(cfg.snapshotsDir); } /** Open the store lazily, sizing the table from a probe embedding the first time. */ @@ -70,23 +74,40 @@ export class MemoryEngine { /** Hybrid (vector + keyword) search. Optional metadata filter, e.g. type/tags. */ async search(query: string, k = 8, filter?: string): Promise { - await this.ensureOpen(); - const qvec = await this.embedder.embed(query); - return this.store.retrieve(query, qvec, k, filter); + const start = Date.now(); + try { + await this.ensureOpen(); + const qvec = await this.embedder.embed(query); + const hits = await this.store.retrieve(query, qvec, k, filter); + metrics.recordSearch(Date.now() - start); + return hits; + } catch (err) { + metrics.recordSearch(Date.now() - start, true); + throw err; + } } /** Rebuild the index from the vault. */ async reindex(onProgress?: (m: string) => void): Promise { - const indexer = new Indexer(this.cfg.vaultPath, this.store, this.embedder); - const res = await indexer.indexAll(onProgress); - this.opened = true; // indexer opens the store as part of its run - return res; + const start = Date.now(); + try { + const indexer = new Indexer(this.cfg.vaultPath, this.store, this.embedder); + const res = await indexer.indexAll(onProgress); + this.opened = true; + metrics.recordIndex(Date.now() - start); + return res; + } catch (err) { + metrics.recordIndex(Date.now() - start, true); + throw err; + } } // Work memory — episodic record of agent actions, outputs, and corrections. recordWork(entry: Omit): WorkMemoryEntry { - return this.workMemory.record(entry); + const result = this.workMemory.record(entry); + metrics.recordWorkMemory(); + return result; } queryWork(q: WorkMemoryQuery = {}): WorkMemoryEntry[] { @@ -98,17 +119,23 @@ export class MemoryEngine { } recordCorrection(sessionId: string, note: string, sourceEntryId?: string): WorkMemoryEntry { - return this.workMemory.recordCorrection(sessionId, note, sourceEntryId); + const result = this.workMemory.recordCorrection(sessionId, note, sourceEntryId); + metrics.recordWorkMemory(); + return result; } // Consolidation — synthesise sessions into vault lesson notes. async consolidateAll(): Promise { - return this.consolidator.consolidateAll(); + const results = await this.consolidator.consolidateAll(); + results.forEach(() => metrics.recordConsolidation()); + return results; } async consolidateSession(sessionId: string): Promise { - return this.consolidator.consolidateSession(sessionId); + const result = await this.consolidator.consolidateSession(sessionId); + metrics.recordConsolidation(); + return result; } // Context graph — entities, edges, and LLM-wiki summaries. @@ -226,9 +253,11 @@ export class MemoryEngine { vote: "up" | "down", opts: { sessionId?: string; query?: string; note?: string; workMemoryEntryId?: string } = {} ): FeedbackSignal { - return vote === "up" + const signal = vote === "up" ? this.feedback.upvote(notePath, opts) : this.feedback.downvote(notePath, opts); + metrics.recordFeedback(); + return signal; } listFeedbackSignals(notePath?: string): FeedbackSignal[] { @@ -267,6 +296,43 @@ export class MemoryEngine { return { processed: processedIds.length, signals }; } + // Metrics + + getMetrics(): MetricsSnapshot { + return metrics.snapshot(); + } + + resetMetrics(): void { + metrics.reset(); + } + + // Snapshots + + async createSnapshot(label?: string): Promise { + return this.snapshots.create({ + workMemoryPath: this.cfg.workMemoryPath, + graphPath: this.cfg.graphPath, + feedbackPath: this.cfg.feedbackPath, + label, + }); + } + + listSnapshots(): SnapshotManifest[] { + return this.snapshots.list(); + } + + getSnapshot(id: string): SnapshotManifest | undefined { + return this.snapshots.get(id); + } + + async restoreSnapshot(id: string): Promise { + return this.snapshots.restore(id); + } + + deleteSnapshot(id: string): boolean { + return this.snapshots.delete(id); + } + private collectWikis(query: string, types?: EntityType[]): WikiSummary[] { const matches = this.graph.findByName(query); const candidates = types ? matches.filter((e) => types.includes(e.type)) : matches; diff --git a/agent-memory-mesh/src/service/http.ts b/agent-memory-mesh/src/service/http.ts index f1750c1..0d4963f 100644 --- a/agent-memory-mesh/src/service/http.ts +++ b/agent-memory-mesh/src/service/http.ts @@ -279,6 +279,46 @@ export function startHttp( return send(res, 200, { entity }); } + // Metrics + if (req.method === "GET" && url.pathname === "/metrics") { + return send(res, 200, { metrics: engine.getMetrics() }); + } + if (req.method === "POST" && url.pathname === "/metrics/reset") { + engine.resetMetrics(); + return send(res, 200, { ok: true }); + } + + // Snapshots + if (req.method === "POST" && url.pathname === "/snapshots") { + const body = await readJson(req); + const manifest = await engine.createSnapshot(body?.label); + return send(res, 201, { manifest }); + } + if (req.method === "GET" && url.pathname === "/snapshots") { + return send(res, 200, { snapshots: engine.listSnapshots() }); + } + const snapIdMatch = url.pathname.match(/^\/snapshots\/([^/]+)$/); + if (snapIdMatch) { + const id = decodeURIComponent(snapIdMatch[1]); + if (req.method === "GET") { + const manifest = engine.getSnapshot(id); + if (!manifest) return send(res, 404, { error: "snapshot not found" }); + return send(res, 200, { manifest }); + } + if (req.method === "DELETE") { + const ok = engine.deleteSnapshot(id); + return send(res, 200, { ok }); + } + } + const snapRestoreMatch = url.pathname.match(/^\/snapshots\/([^/]+)\/restore$/); + if (req.method === "POST" && snapRestoreMatch) { + const id = decodeURIComponent(snapRestoreMatch[1]); + const manifest = engine.getSnapshot(id); + if (!manifest) return send(res, 404, { error: "snapshot not found" }); + await engine.restoreSnapshot(id); + return send(res, 200, { ok: true }); + } + return send(res, 404, { error: "not found" }); } catch (err) { return send(res, 500, { error: String(err) }); diff --git a/agent-memory-mesh/src/service/mcp.ts b/agent-memory-mesh/src/service/mcp.ts index 6560f08..170dc48 100644 --- a/agent-memory-mesh/src/service/mcp.ts +++ b/agent-memory-mesh/src/service/mcp.ts @@ -359,6 +359,58 @@ export async function startMcpStdio(cfg: MemoryConfig, engine: MemoryEngine): Pr } ); + server.registerTool( + "create_snapshot", + { + title: "Create memory snapshot", + description: + "Snapshot the current work memory, context graph, and feedback state. " + + "Use before risky operations or experiments to enable rollback.", + inputSchema: { + label: z.string().optional().describe("Human-readable label for this snapshot."), + }, + }, + async ({ label }) => { + const manifest = await engine.createSnapshot(label); + return { content: [{ type: "text", text: JSON.stringify(manifest, null, 2) }] }; + } + ); + + server.registerTool( + "list_snapshots", + { + title: "List memory snapshots", + description: "List all saved memory snapshots, newest first.", + inputSchema: {}, + }, + async () => { + const snapshots = engine.listSnapshots(); + const text = snapshots.length + ? snapshots.map((s) => `[${s.createdAt}] ${s.id}${s.label ? ` — ${s.label}` : ""}`).join("\n") + : "No snapshots found."; + return { content: [{ type: "text", text }] }; + } + ); + + server.registerTool( + "restore_snapshot", + { + title: "Restore memory snapshot", + description: + "Restore work memory, context graph, and feedback to a previously saved snapshot. " + + "A safety backup is automatically created before the restore.", + inputSchema: { + id: z.string().describe("Snapshot ID to restore."), + }, + }, + async ({ id }) => { + const manifest = engine.getSnapshot(id); + if (!manifest) return { content: [{ type: "text", text: `Snapshot ${id} not found.` }] }; + await engine.restoreSnapshot(id); + return { content: [{ type: "text", text: `Restored snapshot ${id}${manifest.label ? ` (${manifest.label})` : ""}.` }] }; + } + ); + const transport = new StdioServerTransport(); await server.connect(transport); console.error("[mcp] agent-memory-mesh stdio server ready"); diff --git a/agent-memory-mesh/src/service/metrics.ts b/agent-memory-mesh/src/service/metrics.ts new file mode 100644 index 0000000..740309d --- /dev/null +++ b/agent-memory-mesh/src/service/metrics.ts @@ -0,0 +1,84 @@ +export interface MetricsSnapshot { + searchCount: number; + searchErrors: number; + totalSearchLatencyMs: number; + avgSearchLatencyMs: number; + indexCount: number; + indexErrors: number; + totalIndexLatencyMs: number; + workMemoryCount: number; + consolidationCount: number; + feedbackCount: number; + uptimeMs: number; + startedAt: string; +} + +export class MetricsCollector { + private startedAt = new Date(); + private searchCount = 0; + private searchErrors = 0; + private totalSearchLatencyMs = 0; + private indexCount = 0; + private indexErrors = 0; + private totalIndexLatencyMs = 0; + private workMemoryCount = 0; + private consolidationCount = 0; + private feedbackCount = 0; + + recordSearch(latencyMs: number, error = false): void { + this.searchCount++; + this.totalSearchLatencyMs += latencyMs; + if (error) this.searchErrors++; + } + + recordIndex(latencyMs: number, error = false): void { + this.indexCount++; + this.totalIndexLatencyMs += latencyMs; + if (error) this.indexErrors++; + } + + recordWorkMemory(): void { + this.workMemoryCount++; + } + + recordConsolidation(): void { + this.consolidationCount++; + } + + recordFeedback(): void { + this.feedbackCount++; + } + + snapshot(): MetricsSnapshot { + return { + searchCount: this.searchCount, + searchErrors: this.searchErrors, + totalSearchLatencyMs: this.totalSearchLatencyMs, + avgSearchLatencyMs: this.searchCount > 0 ? this.totalSearchLatencyMs / this.searchCount : 0, + indexCount: this.indexCount, + indexErrors: this.indexErrors, + totalIndexLatencyMs: this.totalIndexLatencyMs, + workMemoryCount: this.workMemoryCount, + consolidationCount: this.consolidationCount, + feedbackCount: this.feedbackCount, + uptimeMs: Date.now() - this.startedAt.getTime(), + startedAt: this.startedAt.toISOString(), + }; + } + + reset(): void { + this.startedAt = new Date(); + this.searchCount = 0; + this.searchErrors = 0; + this.totalSearchLatencyMs = 0; + this.indexCount = 0; + this.indexErrors = 0; + this.totalIndexLatencyMs = 0; + this.workMemoryCount = 0; + this.consolidationCount = 0; + this.feedbackCount = 0; + } +} + +/** Process-level singleton — one collector for the entire service lifetime. */ +export const metrics = new MetricsCollector(); diff --git a/agent-memory-mesh/test/metrics.test.ts b/agent-memory-mesh/test/metrics.test.ts new file mode 100644 index 0000000..6b4d116 --- /dev/null +++ b/agent-memory-mesh/test/metrics.test.ts @@ -0,0 +1,62 @@ +import assert from "node:assert/strict"; +import { MetricsCollector } from "../src/service/metrics.js"; + +export async function runMetricsTests(): Promise { + console.log(" [metrics] running..."); + + const m = new MetricsCollector(); + + // Fresh collector has zero counts + let snap = m.snapshot(); + assert.equal(snap.searchCount, 0); + assert.equal(snap.avgSearchLatencyMs, 0); + + // recordSearch increments count and latency + m.recordSearch(50); + m.recordSearch(100); + snap = m.snapshot(); + assert.equal(snap.searchCount, 2); + assert.equal(snap.totalSearchLatencyMs, 150); + assert.equal(snap.avgSearchLatencyMs, 75); + assert.equal(snap.searchErrors, 0); + + // error flag increments error count + m.recordSearch(20, true); + snap = m.snapshot(); + assert.equal(snap.searchErrors, 1); + + // recordIndex + m.recordIndex(200); + snap = m.snapshot(); + assert.equal(snap.indexCount, 1); + assert.equal(snap.totalIndexLatencyMs, 200); + assert.equal(snap.indexErrors, 0); + + m.recordIndex(10, true); + snap = m.snapshot(); + assert.equal(snap.indexErrors, 1); + + // other counters + m.recordWorkMemory(); + m.recordWorkMemory(); + m.recordConsolidation(); + m.recordFeedback(); + snap = m.snapshot(); + assert.equal(snap.workMemoryCount, 2); + assert.equal(snap.consolidationCount, 1); + assert.equal(snap.feedbackCount, 1); + + // uptimeMs is non-negative + assert.ok(snap.uptimeMs >= 0); + assert.ok(snap.startedAt.length > 0); + + // reset clears everything + m.reset(); + snap = m.snapshot(); + assert.equal(snap.searchCount, 0); + assert.equal(snap.indexCount, 0); + assert.equal(snap.workMemoryCount, 0); + assert.equal(snap.avgSearchLatencyMs, 0); + + console.log(" [metrics] all tests passed"); +} diff --git a/agent-memory-mesh/test/run-tests.ts b/agent-memory-mesh/test/run-tests.ts index d1ca5c2..92b8166 100644 --- a/agent-memory-mesh/test/run-tests.ts +++ b/agent-memory-mesh/test/run-tests.ts @@ -5,6 +5,8 @@ import { runConsolidatorTests } from "./consolidator.test.js"; import { runContextGraphTests } from "./context-graph.test.js"; import { runRetrievalPolicyTests } from "./retrieval-policy.test.js"; import { runFeedbackTests } from "./feedback.test.js"; +import { runMetricsTests } from "./metrics.test.js"; +import { runSnapshotsTests } from "./snapshots.test.js"; async function main() { let exitCode = 0; @@ -14,6 +16,8 @@ async function main() { { name: "ContextGraph", fn: runContextGraphTests }, { name: "RetrievalPolicy", fn: runRetrievalPolicyTests }, { name: "Feedback", fn: runFeedbackTests }, + { name: "Metrics", fn: runMetricsTests }, + { name: "Snapshots", fn: runSnapshotsTests }, ]; for (const suite of suites) { diff --git a/agent-memory-mesh/test/snapshots.test.ts b/agent-memory-mesh/test/snapshots.test.ts new file mode 100644 index 0000000..f958cc9 --- /dev/null +++ b/agent-memory-mesh/test/snapshots.test.ts @@ -0,0 +1,92 @@ +import assert from "node:assert/strict"; +import { mkdirSync, writeFileSync, readFileSync, existsSync } from "node:fs"; +import { join } from "node:path"; +import { tmpdir } from "node:os"; +import { randomUUID } from "node:crypto"; +import { SnapshotStore } from "../src/memory/snapshots.js"; + +function makeTmpDir(): string { + const dir = join(tmpdir(), "snap-test-" + randomUUID()); + mkdirSync(dir, { recursive: true }); + return dir; +} + +export async function runSnapshotsTests(): Promise { + console.log(" [snapshots] running..."); + + const snapshotsDir = makeTmpDir(); + const sourceDir = makeTmpDir(); + const store = new SnapshotStore(snapshotsDir); + + const wmPath = join(sourceDir, "work-memory.json"); + const gPath = join(sourceDir, "graph.json"); + const fbPath = join(sourceDir, "feedback.json"); + writeFileSync(wmPath, JSON.stringify([{ id: "e1" }])); + writeFileSync(gPath, JSON.stringify({ entities: [] })); + writeFileSync(fbPath, JSON.stringify({ signals: [] })); + + // create snapshot + const manifest = await store.create({ + workMemoryPath: wmPath, + graphPath: gPath, + feedbackPath: fbPath, + label: "test-snap", + }); + assert.ok(manifest.id); + assert.equal(manifest.label, "test-snap"); + assert.ok(existsSync(manifest.files.workMemoryCopy)); + assert.ok(existsSync(manifest.files.graphCopy)); + assert.ok(existsSync(manifest.files.feedbackCopy)); + + // list returns the snapshot + const list = store.list(); + assert.equal(list.length, 1); + assert.equal(list[0].id, manifest.id); + + // get by id + const fetched = store.get(manifest.id); + assert.ok(fetched); + assert.equal(fetched.label, "test-snap"); + + // get non-existent + assert.equal(store.get("non-existent-id"), undefined); + + // create second snapshot for ordering + await store.create({ workMemoryPath: wmPath, graphPath: gPath, feedbackPath: fbPath, label: "snap2" }); + const list2 = store.list(); + assert.equal(list2.length, 2); + // newest first + assert.equal(list2[0].label, "snap2"); + + // restore — mutate the source files then verify restore brings them back + writeFileSync(wmPath, JSON.stringify([{ id: "e1" }, { id: "e2" }])); + await store.restore(manifest.id); + const restored = JSON.parse(readFileSync(wmPath, "utf8")); + assert.equal(restored.length, 1); // back to original single entry + + // restore auto-creates a safety snapshot (so count is now 4: original + snap2 + safety) + const list3 = store.list(); + assert.ok(list3.length >= 3); + const safetySnap = list3.find((s) => s.label?.startsWith("pre-restore-")); + assert.ok(safetySnap); + + // delete + const ok = store.delete(manifest.id); + assert.equal(ok, true); + assert.equal(store.get(manifest.id), undefined); + + // delete non-existent returns false + assert.equal(store.delete("does-not-exist"), false); + + // missing source file handled gracefully (safeCopy writes empty array) + const missingPath = join(sourceDir, "missing.json"); + const manifest2 = await store.create({ + workMemoryPath: missingPath, + graphPath: gPath, + feedbackPath: fbPath, + }); + const copyContent = readFileSync(manifest2.files.workMemoryCopy, "utf8"); + assert.equal(copyContent, "[]"); + + console.log(" [snapshots] all tests passed"); +}