Skip to content

[codex] Add memory query result cache #59

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
saurabhhhcodes wants to merge 1 commit into
base: main
Choose a base branch
from
Open

[codex] Add memory query result cache #59

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
130 changes: 130 additions & 0 deletions src/mcp-server-cache.test.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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<string, unknown>) {
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);
});
});
65 changes: 64 additions & 1 deletion src/mcp-server.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<string, { result: unknown; ts: number; sessionId: string }>();
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<string, unknown>,
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)}`;
}
Expand All @@ -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;
}

Expand All @@ -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 };
}

Expand All @@ -1273,12 +1330,14 @@ export async function dispatch(
const metadata = typeof params.metadata === "object" && params.metadata !== null
? (params.metadata as Record<string, unknown>) : 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 };
}

Expand Down Expand Up @@ -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": {
Expand Down