diff --git a/src/functions/search.ts b/src/functions/search.ts index 655ea93d..1f7c219b 100644 --- a/src/functions/search.ts +++ b/src/functions/search.ts @@ -7,6 +7,7 @@ import { VectorIndex } from '../state/vector-index.js' import type { EmbeddingProvider } from '../types.js' import { memoryToObservation } from '../state/memory-utils.js' import { recordAccessBatch } from './access-tracker.js' +import { getAgentId, isAgentScopeIsolated } from '../config.js' import { logger } from "../logger.js"; let index: SearchIndex | null = null @@ -328,6 +329,7 @@ export function registerSearchFunction(sdk: ISdk, kv: StateKV): void { cwd?: string format?: string token_budget?: number + agentId?: string }) => { const idx = getSearchIndex() @@ -363,10 +365,25 @@ export function registerSearchFunction(sdk: ISdk, kv: StateKV): void { logger.info('Search index rebuilt', { entries: count }) } - // When filtering by project/cwd, over-fetch from the index so the + const isolated = isAgentScopeIsolated() + const explicitAgentId = + typeof data.agentId === 'string' && data.agentId.trim().length > 0 + ? data.agentId.trim() + : undefined + const wildcardAgent = explicitAgentId === '*' + const filterAgentId = wildcardAgent + ? undefined + : explicitAgentId ?? (isolated ? getAgentId() : undefined) + + // When filtering by project/cwd/agent, over-fetch from the index so the // post-filter still has a chance of returning `effectiveLimit` results. - const filtering = !!(projectFilter || cwdFilter) - const fetchLimit = filtering ? Math.max(effectiveLimit * 10, 100) : effectiveLimit + // Session reads only matter for project/cwd; agentId is filtered later + // off the loaded observation, so don't load sessions for agent-only scope. + const needsSessionFiltering = !!(projectFilter || cwdFilter) + const fetchLimit = + needsSessionFiltering || filterAgentId + ? Math.max(effectiveLimit * 10, 100) + : effectiveLimit const results = idx.search(query, fetchLimit) // Resolve session -> project/cwd once per sessionId we touch. @@ -395,10 +412,11 @@ export function registerSearchFunction(sdk: ISdk, kv: StateKV): void { // First pass: filter by session (sequential — benefits from session cache). // Memory entries with a synthetic sessionId take a secondary KV.memories // path so project filtering works correctly for them too. + const candidateCap = filterAgentId ? fetchLimit : effectiveLimit const candidates: typeof results = [] for (const r of results) { - if (candidates.length >= effectiveLimit) break - if (filtering) { + if (candidates.length >= candidateCap) break + if (needsSessionFiltering) { const s = await loadSession(r.sessionId) if (s) { if (projectFilter && s.project !== projectFilter) continue @@ -448,6 +466,7 @@ export function registerSearchFunction(sdk: ISdk, kv: StateKV): void { for (let i = 0; i < candidates.length; i++) { const obs = obsResults[i] if (obs) { + if (filterAgentId && obs.agentId !== filterAgentId) continue enriched.push({ observation: obs, score: candidates[i].score, @@ -455,6 +474,7 @@ export function registerSearchFunction(sdk: ISdk, kv: StateKV): void { }) } } + if (enriched.length > effectiveLimit) enriched.length = effectiveLimit void recordAccessBatch( kv, diff --git a/src/mcp/server.ts b/src/mcp/server.ts index e2144567..ff2e1f34 100644 --- a/src/mcp/server.ts +++ b/src/mcp/server.ts @@ -10,6 +10,7 @@ import type { } from "../types.js"; import { getVisibleTools } from "./tools-registry.js"; import { timingSafeCompare } from "../auth.js"; +import { getAgentId, isAgentScopeIsolated } from "../config.js"; type McpResponse = { status_code: number; @@ -117,6 +118,9 @@ export function registerMcpEndpoints( limit: typeof args.limit === "number" ? args.limit : 10, format, token_budget: tokenBudget, + ...(typeof args.agentId === "string" && args.agentId.trim() + ? { agentId: args.agentId.trim() } + : {}), } }); const text = format === "narrative" && @@ -1635,14 +1639,29 @@ export function registerMcpEndpoints( }, }; } + const recallAgentId = + typeof promptArgs.agentId === "string" && promptArgs.agentId.trim() + ? promptArgs.agentId.trim() + : undefined; + const wildcardAgent = recallAgentId === "*"; + const filterAgentId = wildcardAgent + ? undefined + : recallAgentId ?? (isAgentScopeIsolated() ? getAgentId() : undefined); const searchResult = await sdk .trigger({ function_id: "mem::search", - payload: { query: taskDesc, limit: 10 }, + payload: { + query: taskDesc, + limit: 10, + ...(recallAgentId ? { agentId: recallAgentId } : {}), + }, }) .catch(() => ({ results: [] })); const memories = await kv.list(KV.memories); - const relevant = memories.filter((m) => m.isLatest).slice(0, 5); + const relevant = memories + .filter((m) => m.isLatest) + .filter((m) => !filterAgentId || m.agentId === filterAgentId) + .slice(0, 5); return { status_code: 200, body: { diff --git a/src/state/memory-utils.ts b/src/state/memory-utils.ts index aa0bcc5b..1483060a 100644 --- a/src/state/memory-utils.ts +++ b/src/state/memory-utils.ts @@ -20,5 +20,6 @@ export function memoryToObservation(memory: Memory): CompressedObservation { concepts: memory.concepts, files: memory.files, importance: memory.strength, + agentId: memory.agentId, }; } diff --git a/src/triggers/api.ts b/src/triggers/api.ts index 6bcf12a1..9f45c11d 100644 --- a/src/triggers/api.ts +++ b/src/triggers/api.ts @@ -370,6 +370,7 @@ export function registerApiTriggers( cwd?: string; format?: string; token_budget?: number; + agentId?: string; }>, ): Promise => { const body = (req.body ?? {}) as Record; @@ -407,6 +408,13 @@ export function registerApiTriggers( body: { error: "token_budget must be a positive integer" }, }; } + const requestAgentId = + typeof body.agentId === "string" && body.agentId.trim().length > 0 + ? body.agentId.trim().slice(0, 128) + : typeof req.query_params?.["agentId"] === "string" && + req.query_params["agentId"].trim().length > 0 + ? req.query_params["agentId"].trim().slice(0, 128) + : undefined; const payload = { query: body.query.trim(), limit: body.limit as number | undefined, @@ -417,6 +425,7 @@ export function registerApiTriggers( ? body.format.trim().toLowerCase() : undefined, token_budget: body.token_budget as number | undefined, + ...(requestAgentId ? { agentId: requestAgentId } : {}), }; const result = await sdk.trigger({ function_id: "mem::search", payload: payload }); return { status_code: 200, body: result }; diff --git a/test/cross-project-isolation.test.ts b/test/cross-project-isolation.test.ts index 8fc37498..6d76da3e 100644 --- a/test/cross-project-isolation.test.ts +++ b/test/cross-project-isolation.test.ts @@ -19,6 +19,7 @@ vi.mock("../src/functions/access-tracker.js", () => ({ vi.mock("../src/config.js", () => ({ getAgentId: () => undefined, + isAgentScopeIsolated: () => false, })); import { registerRememberFunction } from "../src/functions/remember.js"; diff --git a/test/search-agent-scope.test.ts b/test/search-agent-scope.test.ts new file mode 100644 index 00000000..5d394b5c --- /dev/null +++ b/test/search-agent-scope.test.ts @@ -0,0 +1,222 @@ +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; + +vi.mock("../src/logger.js", () => ({ + logger: { info: vi.fn(), warn: vi.fn(), error: vi.fn() }, +})); + +import { registerSearchFunction, getSearchIndex } from "../src/functions/search.js"; +import { KV } from "../src/state/schema.js"; +import type { CompressedObservation, Memory, Session } from "../src/types.js"; + +function mockKV() { + const store = new Map>(); + return { + get: async (scope: string, key: string): Promise => + (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 => { + const entries = store.get(scope); + return entries ? (Array.from(entries.values()) as T[]) : []; + }, + }; +} + +function mockSdk() { + const functions = new Map(); + return { + registerFunction: (idOrOpts: string | { id: string }, handler: Function) => { + const id = typeof idOrOpts === "string" ? idOrOpts : idOrOpts.id; + functions.set(id, handler); + }, + registerTrigger: () => {}, + trigger: async ( + idOrInput: string | { function_id: string; payload: unknown }, + data?: unknown, + ) => { + const id = typeof idOrInput === "string" ? idOrInput : idOrInput.function_id; + const payload = typeof idOrInput === "string" ? data : idOrInput.payload; + const fn = functions.get(id); + if (!fn) throw new Error(`No function: ${id}`); + return fn(payload); + }, + }; +} + +// #817: mem::search must honor AGENTMEMORY_AGENT_SCOPE=isolated so the +// REST /search endpoint and the memory_recall / recall_context MCP tools +// (which all route through mem::search) don't leak memories across agents. +describe("mem::search agent-scope isolation (#817)", () => { + const ORIG_ID = process.env["AGENT_ID"]; + const ORIG_MODE = process.env["AGENTMEMORY_AGENT_SCOPE"]; + + let sdk: ReturnType; + let kv: ReturnType; + + async function search(payload: Record) { + return (await sdk.trigger("mem::search", { + format: "compact", + ...payload, + })) as { results: Array<{ obsId: string }> }; + } + + beforeEach(async () => { + delete process.env["AGENT_ID"]; + delete process.env["AGENTMEMORY_AGENT_SCOPE"]; + + sdk = mockSdk(); + kv = mockKV(); + registerSearchFunction(sdk as never, kv as never); + + // Two sessions owned by two different agents, each with one matching + // observation. The query hits both, so agentId is the only thing that + // can separate them. + const sessionA: Session = { + id: "ses_a", + project: "demo", + cwd: "/tmp/demo", + startedAt: "2026-01-01T00:00:00Z", + status: "completed", + observationCount: 1, + agentId: "agent_a", + }; + const sessionB: Session = { + id: "ses_b", + project: "demo", + cwd: "/tmp/demo", + startedAt: "2026-01-02T00:00:00Z", + status: "completed", + observationCount: 1, + agentId: "agent_b", + }; + await kv.set(KV.sessions, sessionA.id, sessionA); + await kv.set(KV.sessions, sessionB.id, sessionB); + + const obsA: CompressedObservation = { + id: "obs_a", + sessionId: "ses_a", + timestamp: "2026-01-01T00:00:00Z", + type: "decision", + title: "Secret marker alpha", + facts: ["alpha private fact"], + narrative: "secret marker private to agent_a", + concepts: ["secret"], + files: [], + importance: 8, + agentId: "agent_a", + }; + const obsB: CompressedObservation = { + id: "obs_b", + sessionId: "ses_b", + timestamp: "2026-01-02T00:00:00Z", + type: "decision", + title: "Secret marker beta", + facts: ["beta private fact"], + narrative: "secret marker private to agent_b", + concepts: ["secret"], + files: [], + importance: 8, + agentId: "agent_b", + }; + await kv.set(KV.observations("ses_a"), obsA.id, obsA); + await kv.set(KV.observations("ses_b"), obsB.id, obsB); + + // Module-level SearchIndex singleton leaks across tests; reset so each + // case triggers a fresh rebuild from the mock KV. + getSearchIndex().clear(); + }); + + afterEach(() => { + if (ORIG_ID === undefined) delete process.env["AGENT_ID"]; + else process.env["AGENT_ID"] = ORIG_ID; + if (ORIG_MODE === undefined) delete process.env["AGENTMEMORY_AGENT_SCOPE"]; + else process.env["AGENTMEMORY_AGENT_SCOPE"] = ORIG_MODE; + }); + + it("shared mode (default) returns hits from every agent", async () => { + const { results } = await search({ query: "secret marker" }); + const ids = results.map((r) => r.obsId).sort(); + expect(ids).toEqual(["obs_a", "obs_b"]); + }); + + it("isolated mode returns only the env AGENT_ID's hits", async () => { + process.env["AGENT_ID"] = "agent_b"; + process.env["AGENTMEMORY_AGENT_SCOPE"] = "isolated"; + + const { results } = await search({ query: "secret marker" }); + expect(results.map((r) => r.obsId)).toEqual(["obs_b"]); + }); + + it("isolated mode does not leak another agent's memory", async () => { + process.env["AGENT_ID"] = "agent_b"; + process.env["AGENTMEMORY_AGENT_SCOPE"] = "isolated"; + + const { results } = await search({ query: "secret marker" }); + expect(results.some((r) => r.obsId === "obs_a")).toBe(false); + }); + + it("explicit agentId override wins over the env default", async () => { + process.env["AGENT_ID"] = "agent_b"; + process.env["AGENTMEMORY_AGENT_SCOPE"] = "isolated"; + + const { results } = await search({ query: "secret marker", agentId: "agent_a" }); + expect(results.map((r) => r.obsId)).toEqual(["obs_a"]); + }); + + it("wildcard agentId '*' opts out of isolation and returns all agents", async () => { + process.env["AGENT_ID"] = "agent_b"; + process.env["AGENTMEMORY_AGENT_SCOPE"] = "isolated"; + + const { results } = await search({ query: "secret marker", agentId: "*" }); + const ids = results.map((r) => r.obsId).sort(); + expect(ids).toEqual(["obs_a", "obs_b"]); + }); + + // The #817 repro uses /remember, not raw observations. Remembered + // Memory rows are indexed via memoryToObservation(); agentId must + // survive that conversion or isolated recall drops the owner's own + // memories (and the cross-agent filter can't distinguish them). + it("isolated mode scopes memories saved via KV.memories by agentId", async () => { + const memA: Memory = { + id: "mem_a", + title: "Remembered marker alpha", + content: "secret marker remembered by agent_a", + type: "fact", + concepts: ["secret"], + files: [], + strength: 1, + createdAt: "2026-01-03T00:00:00Z", + isLatest: true, + agentId: "agent_a", + } as Memory; + const memB: Memory = { + id: "mem_b", + title: "Remembered marker beta", + content: "secret marker remembered by agent_b", + type: "fact", + concepts: ["secret"], + files: [], + strength: 1, + createdAt: "2026-01-04T00:00:00Z", + isLatest: true, + agentId: "agent_b", + } as Memory; + await kv.set(KV.memories, memA.id, memA); + await kv.set(KV.memories, memB.id, memB); + getSearchIndex().clear(); + + process.env["AGENT_ID"] = "agent_b"; + process.env["AGENTMEMORY_AGENT_SCOPE"] = "isolated"; + + const { results } = await search({ query: "secret marker remembered" }); + const ids = results.map((r) => r.obsId); + expect(ids).toContain("mem_b"); // own remembered memory still recalled + expect(ids).not.toContain("mem_a"); // other agent's is not leaked + }); +});