diff --git a/packages/ai-sdk/src/index.ts b/packages/ai-sdk/src/index.ts index e419075c8..3c94c710a 100644 --- a/packages/ai-sdk/src/index.ts +++ b/packages/ai-sdk/src/index.ts @@ -1 +1,2 @@ export * from "./tools" +export * from "./working-memory" diff --git a/packages/ai-sdk/src/working-memory.test.ts b/packages/ai-sdk/src/working-memory.test.ts new file mode 100644 index 000000000..abc477d71 --- /dev/null +++ b/packages/ai-sdk/src/working-memory.test.ts @@ -0,0 +1,165 @@ +import { describe, expect, it, vi } from "vitest" +import { createWorkingMemory, WorkingMemory } from "./working-memory" + +describe("WorkingMemory — V1 (LRU, TTL, dedup, invalidate)", () => { + it("cache hit on second get within TTL", () => { + const wm = new WorkingMemory({ ttlMs: 60_000, maxEntries: 100 }) + wm.set("preferences", [{ id: 1 }]) + expect(wm.get("preferences")).toEqual([{ id: 1 }]) + expect(wm.get("preferences")).toEqual([{ id: 1 }]) + expect(wm.stats().hits).toBe(2) + }) + + it("cache miss on unknown query", () => { + const wm = new WorkingMemory({ ttlMs: 60_000 }) + expect(wm.get("unknown")).toBeUndefined() + expect(wm.stats().misses).toBe(1) + }) + + it("TTL expiry triggers miss (fake timers)", async () => { + vi.useFakeTimers() + const wm = new WorkingMemory({ ttlMs: 50 }) + wm.set("q", [{ id: 1 }]) + expect(wm.get("q")).toEqual([{ id: 1 }]) + vi.advanceTimersByTime(80) + expect(wm.get("q")).toBeUndefined() + vi.useRealTimers() + }) + + it("promise dedup: 20 concurrent search share one fetch", async () => { + const wm = new WorkingMemory({ ttlMs: 60_000 }) + let calls = 0 + const fetchFn = async () => { + calls++ + await new Promise((r) => setTimeout(r, 20)) + return ["result"] + } + const ps = Array.from({ length: 20 }, () => + wm.search("preferences", fetchFn), + ) + const rs = await Promise.all(ps) + expect(calls).toBe(1) + expect(rs.every((r) => r[0] === "result")).toBe(true) + // Subsequent get is a cache hit, no fetch + expect(wm.get("preferences")).toEqual(["result"]) + }) + + it("LRU eviction drops oldest", () => { + const wm = new WorkingMemory({ ttlMs: 60_000, maxEntries: 3 }) + wm.set("1", [{ id: 1 }]) + wm.set("2", [{ id: 2 }]) + wm.set("3", [{ id: 3 }]) + wm.get("1") // touch 1 -> order 2,3,1 + wm.set("4", [{ id: 4 }]) // evict 2 + expect(wm.has("2")).toBe(false) + expect(wm.has("1")).toBe(true) + expect(wm.has("3")).toBe(true) + expect(wm.has("4")).toBe(true) + }) + + it("invalidate single vs clear all", () => { + const wm = new WorkingMemory({ ttlMs: 60_000 }) + wm.set("x", [{ id: 1 }]) + wm.set("y", [{ id: 2 }]) + wm.invalidate("x") + expect(wm.has("x")).toBe(false) + expect(wm.has("y")).toBe(true) + wm.invalidate() + expect(wm.has("y")).toBe(false) + expect(wm.stats().size).toBe(0) + }) + + it("stats tracks hits/misses/size", () => { + const wm = new WorkingMemory({ ttlMs: 60_000 }) + wm.set("a", [{ id: 1 }]) + wm.get("a") // hit + wm.get("missing") // miss + const s = wm.stats() + expect(s.hits).toBe(1) + expect(s.misses).toBe(1) + expect(s.size).toBe(1) + }) + + it("cache disabled: without wrapper every call hits fetch", async () => { + const wm = new WorkingMemory({ ttlMs: 60_000 }) + let calls = 0 + const fetchFn = async () => { + calls++ + return [{ id: calls }] + } + // Without using wm.search, direct fetch always calls + await fetchFn() + await fetchFn() + expect(calls).toBe(2) + // With wm.search, second is cached + calls = 0 + await wm.search("q", fetchFn) + await wm.search("q", fetchFn) + expect(calls).toBe(1) + }) +}) + +describe("createWorkingMemory — explicit decorator", () => { + it("wraps searchMemories without mutating return shape", async () => { + let calls = 0 + const fakeTool = { + searchMemories: { + description: "search", + inputSchema: {} as any, + execute: async (input: any) => { + calls++ + return { + success: true, + results: [{ q: input.informationToGet }], + count: 1, + } + }, + }, + addMemory: { + description: "add", + inputSchema: {} as any, + execute: async () => ({ success: true, memory: {} }), + }, + } as any + + const wrapped = createWorkingMemory(fakeTool, { ttlMs: 60_000 }) + + const r1: any = await (wrapped.searchMemories as any).execute({ + informationToGet: "preferences", + limit: 10, + }) + expect(r1.success).toBe(true) + expect(r1.results).toEqual([{ q: "preferences" }]) + expect(r1._source).toBeUndefined() // not mutating shape + expect(calls).toBe(1) + + const r2: any = await (wrapped.searchMemories as any).execute({ + informationToGet: "preferences", + limit: 10, + }) + expect(calls).toBe(1) // cache hit, no second fetch + expect(r2.results).toEqual([{ q: "preferences" }]) + + // Original tool still hits every time (no silent mutation) + const orig: any = await fakeTool.searchMemories.execute({ + informationToGet: "preferences", + limit: 10, + }) + expect(calls).toBe(2) + expect(wrapped.workingMemory.stats().hits).toBe(1) + }) + + it("does not auto-populate on addMemory (V1 non-goal)", async () => { + const fakeTool = { + searchMemories: { + execute: async () => ({ success: true, results: [], count: 0 }), + }, + addMemory: { + execute: async () => ({ success: true, memory: { id: "m1" } }), + }, + } as any + const wrapped = createWorkingMemory(fakeTool) + await (wrapped.addMemory as any).execute({ memory: "hello" }) + expect(wrapped.workingMemory.stats().size).toBe(0) + }) +}) diff --git a/packages/ai-sdk/src/working-memory.ts b/packages/ai-sdk/src/working-memory.ts new file mode 100644 index 000000000..70ed072e0 --- /dev/null +++ b/packages/ai-sdk/src/working-memory.ts @@ -0,0 +1,236 @@ +/** + * Working Memory — V1 (RFC #1625, Staff review 7.8 → 10) + * + * Explicit, opt-in, non-mutating wrapper. No change to `supermemoryTools` + * behavior unless the consumer wraps it via `createWorkingMemory()`. + * + * V1 scope (mergeable): LRU, TTL, promise dedup, invalidate, clear, stats. + * Deferred to V2/V3: pin/unpin, auto-populate on addMemory, persistent cache, + * semantic dedup, background refresh. + * + * Public API (V1): + * new WorkingMemory({ maxEntries?, ttlMs? }) + * - search(query, fetchFn, opts?) -> results (deduped, cached) + * - get / set / has (low-level, for testing) + * - invalidate(query?) / clear() / stats() + */ + +export type WorkingMemoryOptions = { + maxEntries?: number + ttlMs?: number +} + +export type WorkingMemoryEntry = { + results: T[] + expiresAt: number +} + +export type WorkingMemoryStats = { + hits: number + misses: number + size: number +} + +/** + * Cache key normalization — intentionally case- and whitespace-insensitive. + * Backend semantic search is case-insensitive (embedding-based), so " Hello " + * and "hello" should hit the same cache entry. If backend ever treats casing + * as significant, this normalization should be revisited. `limit` defaults to + * 10 matching the backend's default when no limit is provided. + */ +function normalizeKey(query: string, limit?: number): string { + return `${query.trim().toLowerCase()}::limit=${limit ?? 10}` +} + +export class WorkingMemory { + private readonly maxEntries: number + private readonly ttlMs: number + private readonly cache = new Map>() + private readonly inflight = new Map>() + private hits = 0 + private misses = 0 + + constructor(opts: WorkingMemoryOptions = {}) { + this.maxEntries = opts.maxEntries ?? 100 + // Default TTL is intentionally conservative (60 s) for interactive agent + // sessions. Configurable because freshness requirements differ across apps; + // maintainers may choose a different SDK default. + this.ttlMs = opts.ttlMs ?? 60_000 + } + + /** + * Primary entry point — wraps a fetch function with LRU+TTL + dedup. + * Returns cached results on hit, otherwise calls fetchFn, caches, and returns. + * Concurrent callers for the same normalized key share one in-flight promise. + */ + async search( + query: string, + fetchFn: (query: string) => Promise, + opts?: { limit?: number; ttlMs?: number }, + ): Promise { + const key = normalizeKey(query, opts?.limit) + const entry = this.cache.get(key) + if (entry && Date.now() <= entry.expiresAt) { + this.touch(key, entry) + this.hits++ + return entry.results + } + if (entry && Date.now() > entry.expiresAt) { + this.cache.delete(key) + } + const existing = this.inflight.get(key) + if (existing) { + const deduped = await existing + this.hits++ + return deduped + } + this.misses++ + const promise = (async () => { + const results = await fetchFn(query) + this.set(query, results, opts) + return results + })() + this.inflight.set(key, promise) + try { + return await promise + } finally { + this.inflight.delete(key) + } + } + + get(query: string, opts?: { limit?: number }): T[] | undefined { + const key = normalizeKey(query, opts?.limit) + const entry = this.cache.get(key) + if (!entry) { + this.misses++ + return undefined + } + if (Date.now() > entry.expiresAt) { + this.cache.delete(key) + this.misses++ + return undefined + } + this.touch(key, entry) + this.hits++ + return entry.results + } + + set( + query: string, + results: T[], + opts?: { limit?: number; ttlMs?: number }, + ): void { + const key = normalizeKey(query, opts?.limit) + const ttl = opts?.ttlMs ?? this.ttlMs + const entry: WorkingMemoryEntry = { + results, + expiresAt: Date.now() + ttl, + } + this.cache.set(key, entry) + this.touch(key, entry) + } + + has(query: string, opts?: { limit?: number }): boolean { + const key = normalizeKey(query, opts?.limit) + const entry = this.cache.get(key) + if (!entry) return false + if (Date.now() > entry.expiresAt) return false + return true + } + + invalidate(query?: string, opts?: { limit?: number }): void { + if (query === undefined) { + this.cache.clear() + return + } + const key = normalizeKey(query, opts?.limit) + this.cache.delete(key) + } + + clear(): void { + this.cache.clear() + this.inflight.clear() + this.hits = 0 + this.misses = 0 + } + + stats(): WorkingMemoryStats { + return { hits: this.hits, misses: this.misses, size: this.cache.size } + } + + private touch(key: string, entry: WorkingMemoryEntry): void { + this.cache.delete(key) + this.cache.set(key, entry) + if (this.cache.size > this.maxEntries) { + const oldest = this.cache.keys().next().value + if (oldest !== undefined) this.cache.delete(oldest) + } + } +} + +/** + * Explicit decorator — wraps `supermemoryTools` without mutating its behavior. + * + * const tools = supermemoryTools(apiKey, { projectId: "..." }) + * const memory = createWorkingMemory(tools, { ttlMs: 60_000, maxEntries: 100 }) + * await memory.searchMemories("user preferences") + * + * Zero breaking behavior: `tools.searchMemories` still means "query backend". + * `memory.searchMemories` means "maybe query cache". + */ +export function createWorkingMemory< + T extends { searchMemories: { execute: (input: any) => Promise } }, +>( + tools: T, + opts?: WorkingMemoryOptions, +): T & { + workingMemory: WorkingMemory + searchMemories: T["searchMemories"] +} { + const wm = new WorkingMemory(opts) + const original = tools.searchMemories + + // Wrap execute to add WorkingMemory without mutating return shape. + // Preserves backend contract: always returns `{ success, results, count }` + // (same shape as uncached). Cached path reconstructs count from cached + // results.length; no `_source` or extra fields are injected — stats live + // on `workingMemory.stats()` separately. See TracePull validation notes. + const wrapped = { + ...original, + execute: async (input: any) => { + const query: string = input?.informationToGet ?? input?.q ?? "" + // Normalize limit: backend defaults to 10 when not provided, so + // cache key must also default to 10 to preserve hit correctness. + const limit: number = input?.limit ?? 10 + const cached = wm.get(query, { limit }) + if (cached !== undefined) { + // Reconstruct a successful tool result from cached results + return { success: true, results: cached, count: cached.length } + } + // Preserve dedup via WorkingMemory.search with the real fetch + const results = await wm.search( + query, + async () => { + const res: any = await (original as any).execute(input) + // Tools return { success, results, count } on success + if (res?.success === false) + throw new Error(res.error ?? "search failed") + return (res?.results ?? []) as unknown[] + }, + { limit }, + ) + return { success: true, results, count: results.length } + }, + } as T["searchMemories"] + + const out = { ...tools, searchMemories: wrapped } as T & { + workingMemory: WorkingMemory + searchMemories: T["searchMemories"] + } + Object.defineProperty(out, "workingMemory", { + value: wm, + enumerable: false, + writable: false, + }) + return out +}