Merge 9ddd057b8c into cfa6c7cb17

packages/ai-sdk/src/index.ts

@@ -1 +1,2 @@
 export * from "./tools"
+export * from "./working-memory"

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<string>({ 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)
+	})
+})

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<T = unknown> = {
+	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<T = unknown> {
+	private readonly maxEntries: number
+	private readonly ttlMs: number
+	private readonly cache = new Map<string, WorkingMemoryEntry<T>>()
+	private readonly inflight = new Map<string, Promise<T[]>>()
+	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<T[]>,
+		opts?: { limit?: number; ttlMs?: number },
+	): Promise<T[]> {
+		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<T> = {
+			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<T>): 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<any> } },
+>(
+	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
+}