From 3fe066b53fbafee491952b64d6e3a0364be6c95b Mon Sep 17 00:00:00 2001 From: Saurabh Kumar Bajpai Date: Thu, 28 May 2026 18:11:18 +0530 Subject: [PATCH] Add memory query result cache --- src/mcp-server-cache.test.ts | 130 +++++++++++++++++++++++++++++++++++ src/mcp-server.ts | 65 +++++++++++++++++- 2 files changed, 194 insertions(+), 1 deletion(-) create mode 100644 src/mcp-server-cache.test.ts diff --git a/src/mcp-server-cache.test.ts b/src/mcp-server-cache.test.ts new file mode 100644 index 0000000..498256d --- /dev/null +++ b/src/mcp-server-cache.test.ts @@ -0,0 +1,130 @@ +import { afterEach, describe, expect, test } from "bun:test"; +import { dispatch, type MemoryStore } from "./mcp-server"; + +const originalTtl = process.env.MEMORY_QUERY_CACHE_TTL_MS; + +class FakeMemoryStore { + sessionId = ""; + queryCalls = 0; + + setSessionId(sessionId: string) { + this.sessionId = sessionId; + } + + queryMemories() { + this.queryCalls++; + return [{ id: `result-${this.queryCalls}`, sessionId: this.sessionId }]; + } + + addMemory() { + return "mem_1"; + } + + updateMemory() { + return true; + } + + deleteMemory() { + return true; + } + + getAccessAnalytics() { + return { totalQueries: 0 }; + } +} + +function callDispatch(store: FakeMemoryStore, method: string, params: Record) { + return dispatch(store as unknown as MemoryStore, method, params); +} + +afterEach(() => { + if (originalTtl === undefined) { + delete process.env.MEMORY_QUERY_CACHE_TTL_MS; + } else { + process.env.MEMORY_QUERY_CACHE_TTL_MS = originalTtl; + } +}); + +describe("memory_query result cache", () => { + test("returns cached results for identical queries within the TTL", async () => { + process.env.MEMORY_QUERY_CACHE_TTL_MS = "30000"; + const store = new FakeMemoryStore(); + const params = { session_id: "cache-hit-session", query: "typescript skill", limit: 3 }; + + const first = await callDispatch(store, "memory_query", params); + const second = await callDispatch(store, "memory_query", params); + + expect(second).toBe(first); + expect(store.queryCalls).toBe(1); + }); + + test("invalidates cached results after memory_add for the same session", async () => { + process.env.MEMORY_QUERY_CACHE_TTL_MS = "30000"; + const store = new FakeMemoryStore(); + const params = { session_id: "cache-invalidate-session", query: "testing skill", limit: 2 }; + + const first = await callDispatch(store, "memory_query", params); + await callDispatch(store, "memory_add", { + session_id: "cache-invalidate-session", + type: "fact", + content: "new testing skill memory", + }); + const second = await callDispatch(store, "memory_query", params); + + expect(second).not.toBe(first); + expect(store.queryCalls).toBe(2); + }); + + test("invalidates cached results after memory_update and memory_delete for the same session", async () => { + process.env.MEMORY_QUERY_CACHE_TTL_MS = "30000"; + const store = new FakeMemoryStore(); + const params = { session_id: "cache-update-delete-session", query: "mutable skill", limit: 2 }; + + await callDispatch(store, "memory_query", params); + await callDispatch(store, "memory_query", params); + expect(store.queryCalls).toBe(1); + + await callDispatch(store, "memory_update", { + session_id: "cache-update-delete-session", + id: "mem_1", + content: "updated skill memory", + }); + await callDispatch(store, "memory_query", params); + expect(store.queryCalls).toBe(2); + + await callDispatch(store, "memory_delete", { + session_id: "cache-update-delete-session", + id: "mem_1", + }); + await callDispatch(store, "memory_query", params); + expect(store.queryCalls).toBe(3); + }); + + test("honors MEMORY_QUERY_CACHE_TTL_MS", async () => { + process.env.MEMORY_QUERY_CACHE_TTL_MS = "1"; + const store = new FakeMemoryStore(); + const params = { session_id: "cache-ttl-session", query: "short ttl", limit: 1 }; + + const first = await callDispatch(store, "memory_query", params); + await Bun.sleep(5); + const second = await callDispatch(store, "memory_query", params); + + expect(second).not.toBe(first); + expect(store.queryCalls).toBe(2); + }); + + test("surfaces query cache stats in memory_access_analytics", async () => { + process.env.MEMORY_QUERY_CACHE_TTL_MS = "30000"; + const store = new FakeMemoryStore(); + const params = { session_id: "cache-stats-session", query: "stats", limit: 1 }; + + await callDispatch(store, "memory_query", params); + await callDispatch(store, "memory_query", params); + const analytics = await callDispatch(store, "memory_access_analytics", { + session_id: "cache-stats-session", + }) as { queryCacheHits?: number; queryCacheMisses?: number }; + + expect(analytics.queryCacheHits).toBeGreaterThanOrEqual(1); + expect(analytics.queryCacheMisses).toBeGreaterThanOrEqual(1); + }); +}); diff --git a/src/mcp-server.ts b/src/mcp-server.ts index 4df7e36..188e4ee 100644 --- a/src/mcp-server.ts +++ b/src/mcp-server.ts @@ -1224,6 +1224,50 @@ const TOOL_NAMES = [ "skill_search", "skill_apply", "skill_check_compatibility", ]; +const DEFAULT_QUERY_CACHE_TTL_MS = 30_000; +const MAX_QUERY_CACHE_ENTRIES = 200; + +const _queryCache = new Map(); +let queryCacheHits = 0; +let queryCacheMisses = 0; + +function queryCacheTtlMs(): number { + const raw = process.env.MEMORY_QUERY_CACHE_TTL_MS; + if (raw === undefined) return DEFAULT_QUERY_CACHE_TTL_MS; + + const parsed = Number(raw); + return Number.isFinite(parsed) && parsed >= 0 ? parsed : DEFAULT_QUERY_CACHE_TTL_MS; +} + +function memoryQueryCacheKey( + params: Record, + sessionId: string, + layers: MemoryLayer[] +): string { + return JSON.stringify({ + q: String(params.query ?? ""), + t: params.type ?? null, + l: params.limit ?? 10, + layers, + s: sessionId, + }); +} + +function pruneQueryCache(now = Date.now()): void { + if (_queryCache.size <= MAX_QUERY_CACHE_ENTRIES) return; + + const ttl = queryCacheTtlMs(); + for (const [key, value] of _queryCache) { + if (now - value.ts > ttl) _queryCache.delete(key); + } +} + +function invalidateQueryCache(sessionId: string): void { + for (const [key, value] of _queryCache) { + if (value.sessionId === sessionId) _queryCache.delete(key); + } +} + export function defaultSessionId(): string { return `project_${Buffer.from(process.cwd()).toString("base64").slice(0, 16)}`; } @@ -1245,7 +1289,19 @@ export async function dispatch( const layers = Array.isArray(params.layers) ? (params.layers as MemoryLayer[]) : (["project", "global"] as MemoryLayer[]); + const cacheKey = memoryQueryCacheKey(params, sessionId, layers); + const cached = _queryCache.get(cacheKey); + const now = Date.now(); + const ttl = queryCacheTtlMs(); + if (cached && now - cached.ts < ttl) { + queryCacheHits++; + return cached.result; + } + + queryCacheMisses++; const results = store.queryMemories({ search: query, type, limit, layers }); + _queryCache.set(cacheKey, { result: results, ts: now, sessionId }); + pruneQueryCache(now); return results; } @@ -1256,6 +1312,7 @@ export async function dispatch( const layer = (params.layer as MemoryLayer) ?? "project"; const id = store.addMemory({ type, content, metadata: { addedBy: "mcp" }, importance }, layer); if (!id) throw new Error("Failed to store memory"); + invalidateQueryCache(sessionId); return { id, layer }; } @@ -1273,12 +1330,14 @@ export async function dispatch( const metadata = typeof params.metadata === "object" && params.metadata !== null ? (params.metadata as Record) : undefined; const ok = store.updateMemory(id, { content, importance, metadata }); + if (ok) invalidateQueryCache(sessionId); return { success: ok, id }; } case "memory_delete": { const id = String(params.id ?? ""); const ok = store.deleteMemory(id); + if (ok) invalidateQueryCache(sessionId); return { success: ok, id }; } @@ -1364,7 +1423,11 @@ export async function dispatch( case "memory_access_analytics": { const layer = (params.layer as "project" | "global") ?? "project"; - return store.getAccessAnalytics(layer); + return { + ...store.getAccessAnalytics(layer), + queryCacheHits, + queryCacheMisses, + }; } case "memory_validation_report": {