diff --git a/src/functions/context.ts b/src/functions/context.ts index c258a45b..8a25f9b8 100644 --- a/src/functions/context.ts +++ b/src/functions/context.ts @@ -5,11 +5,17 @@ import type { SessionSummary, ContextBlock, ProjectProfile, + MemorySlot, } from "../types.js"; import { KV } from "../state/schema.js"; import { StateKV } from "../state/kv.js"; import { recordAccessBatch } from "./access-tracker.js"; import { logger } from "../logger.js"; +import { + isSlotsEnabled, + listPinnedSlots, + renderPinnedContext, +} from "./slots.js"; function estimateTokens(text: string): number { return Math.ceil(text.length / 3); @@ -33,9 +39,24 @@ export function registerContextFunction( const budget = data.budget || tokenBudget; const blocks: ContextBlock[] = []; - const profile = await kv - .get(KV.profiles, data.project) - .catch(() => null); + const [pinnedSlots, profile] = await Promise.all([ + isSlotsEnabled() + ? listPinnedSlots(kv).catch(() => [] as MemorySlot[]) + : Promise.resolve([] as MemorySlot[]), + kv + .get(KV.profiles, data.project) + .catch(() => null), + ]); + + const slotContent = renderPinnedContext(pinnedSlots); + if (slotContent) { + blocks.push({ + type: "memory", + content: slotContent, + tokens: estimateTokens(slotContent), + recency: Date.now(), + }); + } if (profile) { const profileParts = []; if (profile.topConcepts.length > 0) { diff --git a/test/context-slots.test.ts b/test/context-slots.test.ts new file mode 100644 index 00000000..808597c4 --- /dev/null +++ b/test/context-slots.test.ts @@ -0,0 +1,177 @@ +import { describe, it, expect, beforeEach, afterEach, vi } from "vitest"; +import { registerContextFunction } from "../src/functions/context.js"; +import { KV } from "../src/state/schema.js"; + +function mockKV() { + const store = new Map>(); + return { + get: async (scope: string, key: string): Promise => { + return (store.get(scope)?.get(key) as T) ?? null; + }, + set: async (scope: string, key: string, data: T): Promise => { + if (!store.has(scope)) store.set(scope, new Map()); + store.get(scope)!.set(key, data); + return data; + }, + delete: async (scope: string, key: string): Promise => { + store.get(scope)?.delete(key); + }, + list: async (scope: string): Promise => { + if (!store.has(scope)) return []; + return Array.from(store.get(scope)!.values()) as T[]; + }, + }; +} + +type ContextHandler = (data: { + sessionId: string; + project: string; + budget?: number; +}) => Promise<{ context: string; blocks: number; tokens: number }>; + +function wireContext(kv: ReturnType) { + let handler: ContextHandler | undefined; + const sdk = { + registerFunction: vi.fn((id: string, cb: ContextHandler) => { + if (id === "mem::context") handler = cb; + }), + } as unknown as import("iii-sdk").ISdk; + registerContextFunction(sdk, kv as never, 2000); + if (!handler) throw new Error("mem::context not registered"); + return handler; +} + +async function seedPinnedSlot( + kv: ReturnType, + label: string, + content: string, + scope: "project" | "global" = "global", +) { + const target = scope === "global" ? KV.globalSlots : KV.slots; + await kv.set(target, label, { + label, + content, + description: "", + sizeLimit: 2000, + pinned: true, + readOnly: false, + scope, + createdAt: new Date().toISOString(), + updatedAt: new Date().toISOString(), + }); +} + +describe("mem::context — pinned slot injection", () => { + const ORIGINAL_SLOTS_ENV = process.env["AGENTMEMORY_SLOTS"]; + + afterEach(() => { + if (ORIGINAL_SLOTS_ENV === undefined) { + delete process.env["AGENTMEMORY_SLOTS"]; + } else { + process.env["AGENTMEMORY_SLOTS"] = ORIGINAL_SLOTS_ENV; + } + }); + + describe("when AGENTMEMORY_SLOTS=true", () => { + let kv: ReturnType; + let handler: ContextHandler; + + beforeEach(() => { + process.env["AGENTMEMORY_SLOTS"] = "true"; + kv = mockKV(); + handler = wireContext(kv); + }); + + it("includes pinned global slot content in returned context", async () => { + await seedPinnedSlot(kv, "tool_guidelines", "rule-alpha", "global"); + + const result = await handler({ + sessionId: "ses_a", + project: "/tmp/proj", + }); + + expect(result.context).toContain("tool_guidelines"); + expect(result.context).toContain("rule-alpha"); + expect(result.blocks).toBeGreaterThan(0); + }); + + it("renders multiple pinned slots, sorted by label", async () => { + await seedPinnedSlot(kv, "user_preferences", "pref-alpha", "global"); + await seedPinnedSlot(kv, "tool_guidelines", "rule-alpha", "global"); + + const result = await handler({ + sessionId: "ses_b", + project: "/tmp/proj", + }); + + const guidelinesIdx = result.context.indexOf("tool_guidelines"); + const prefsIdx = result.context.indexOf("user_preferences"); + expect(guidelinesIdx).toBeGreaterThan(-1); + expect(prefsIdx).toBeGreaterThan(-1); + expect(guidelinesIdx).toBeLessThan(prefsIdx); + }); + + it("skips unpinned slots even when they have content", async () => { + await kv.set(KV.globalSlots, "self_notes", { + label: "self_notes", + content: "unpinned-content-alpha", + description: "", + sizeLimit: 1500, + pinned: false, + readOnly: false, + scope: "global", + createdAt: new Date().toISOString(), + updatedAt: new Date().toISOString(), + }); + + const result = await handler({ + sessionId: "ses_c", + project: "/tmp/proj", + }); + + expect(result.context).not.toContain("unpinned-content-alpha"); + }); + + it("skips empty pinned slots (the seeded defaults)", async () => { + await seedPinnedSlot(kv, "persona", "", "global"); + + const result = await handler({ + sessionId: "ses_d", + project: "/tmp/proj", + }); + + expect(result.context).not.toContain("persona"); + }); + + it("project-scoped slot shadows global slot with the same label", async () => { + await seedPinnedSlot(kv, "tool_guidelines", "global-value", "global"); + await seedPinnedSlot(kv, "tool_guidelines", "project-value", "project"); + + const result = await handler({ + sessionId: "ses_e", + project: "/tmp/proj", + }); + + expect(result.context).toContain("project-value"); + expect(result.context).not.toContain("global-value"); + }); + }); + + describe("when AGENTMEMORY_SLOTS is off", () => { + it("does not include any slot content", async () => { + delete process.env["AGENTMEMORY_SLOTS"]; + const kv = mockKV(); + const handler = wireContext(kv); + + await seedPinnedSlot(kv, "tool_guidelines", "rule-alpha", "global"); + + const result = await handler({ + sessionId: "ses_f", + project: "/tmp/proj", + }); + + expect(result.context).not.toContain("tool_guidelines"); + expect(result.context).not.toContain("rule-alpha"); + }); + }); +});