fix(scope): enforce AGENT_ID isolation on mem::search (closes #817)#821
fix(scope): enforce AGENT_ID isolation on mem::search (closes #817)#821wyh0626 wants to merge 1 commit into
Conversation
|
@wyh0626 is attempting to deploy a commit to the rohitg00's projects Team on Vercel. A member of the Team first needs to authorize it. |
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (1)
🚧 Files skipped from review as they are similar to previous changes (1)
📝 WalkthroughWalkthroughAdds an ChangesAgent observation enrichment
Estimated code review effort🎯 2 (Simple) | ⏱️ ~10 minutes Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (1)
test/search-agent-scope.test.ts (1)
156-162: 💤 Low valueConsider consolidating duplicate test coverage.
This test verifies the same isolation behavior as the previous test (lines 148-154), just with a negative assertion style. The previous test already confirms only
obs_bis returned, which inherently provesobs_ais not leaked.While explicit negative assertions for security properties have value, the redundancy adds minimal benefit.
💡 Optional consolidation
Consider either removing this test or combining it with the previous one:
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"]); + // Verify no cross-agent leakage + expect(results.some((r) => r.obsId === "obs_a")).toBe(false); }); - 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); - });🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@test/search-agent-scope.test.ts` around lines 156 - 162, The test "isolated mode does not leak another agent's memory" duplicates the previous isolation test's coverage; remove this redundant test or merge its negative assertion into the preceding test so isolation is asserted once. Specifically, either delete the "isolated mode does not leak another agent's memory" it(...) block in test/search-agent-scope.test.ts, or update the prior test (the other it(...) that checks returned results for only obs_b) to also assert results.some(r => r.obsId === "obs_a") is false, keeping a single test that verifies both positive and negative expectations for isolation.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@src/functions/search.ts`:
- Around line 380-381: Split the single "filtering" boolean into two flags: one
for session-based filtering (e.g., sessionFiltering = !!(projectFilter ||
cwdFilter)) and one for agent-only filtering (e.g., agentFiltering =
!!filterAgentId); keep a combined boolean (e.g., anyFiltering = sessionFiltering
|| agentFiltering) for calculating fetchLimit (replace the current filtering use
in fetchLimit with anyFiltering) but use sessionFiltering when deciding to load
sessions from KV.sessions (the block around the code previously at lines 414-416
that reads KV.sessions should only run if sessionFiltering is true), and leave
the agent check where it currently is (around line 464) using agentFiltering;
update all references to the old filtering variable accordingly.
In `@test/search-agent-scope.test.ts`:
- Around line 69-133: The test is bypassing TypeScript checks by passing mocks
with "as never" to registerSearchFunction; update the mock factories so
mockSdk() returns ISdk and mockKV() returns StateKV (or use a narrower assertion
like mockSdk() as unknown as ISdk) and then call registerSearchFunction(sdk, kv)
without "as never"; ensure mockSdk and mockKV signatures and returned objects
implement the methods used by registerSearchFunction so the compiler enforces
interface compatibility.
---
Nitpick comments:
In `@test/search-agent-scope.test.ts`:
- Around line 156-162: The test "isolated mode does not leak another agent's
memory" duplicates the previous isolation test's coverage; remove this redundant
test or merge its negative assertion into the preceding test so isolation is
asserted once. Specifically, either delete the "isolated mode does not leak
another agent's memory" it(...) block in test/search-agent-scope.test.ts, or
update the prior test (the other it(...) that checks returned results for only
obs_b) to also assert results.some(r => r.obsId === "obs_a") is false, keeping a
single test that verifies both positive and negative expectations for isolation.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 250263de-cb8f-4fa6-8f17-5669c7af0601
📒 Files selected for processing (6)
src/functions/search.tssrc/mcp/server.tssrc/state/memory-utils.tssrc/triggers/api.tstest/cross-project-isolation.test.tstest/search-agent-scope.test.ts
| 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(); | ||
| }); |
There was a problem hiding this comment.
Consider improving type safety of mock injection.
The test setup correctly seeds agent-scoped sessions and observations, and properly clears the SearchIndex singleton to prevent cross-test pollution. However, line 75 uses as never to bypass TypeScript's type checking when passing mocks to registerSearchFunction.
While this works, it could hide type mismatches if the mock implementations drift from the actual interfaces.
🔧 Suggested improvement
Consider typing the mocks to properly match ISdk and StateKV interfaces, or use a type assertion that preserves some checking:
- registerSearchFunction(sdk as never, kv as never);
+ registerSearchFunction(sdk as any, kv as any);Or better yet, properly type the mock factories to return the correct interfaces (though this may require additional method stubs).
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@test/search-agent-scope.test.ts` around lines 69 - 133, The test is bypassing
TypeScript checks by passing mocks with "as never" to registerSearchFunction;
update the mock factories so mockSdk() returns ISdk and mockKV() returns StateKV
(or use a narrower assertion like mockSdk() as unknown as ISdk) and then call
registerSearchFunction(sdk, kv) without "as never"; ensure mockSdk and mockKV
signatures and returned objects implement the methods used by
registerSearchFunction so the compiler enforces interface compatibility.
wyh0626
force-pushed
the
fix/817-mem-search-isolation
branch
2 times, most recently
from
674bf46 to
cc338e9
Compare
… follow-up) The rohitg00#817 fix enforces AGENT_ID isolation in mem::search by post-filtering loaded rows on `obs.agentId !== filterAgentId`. But memoryToObservation() — the path that wraps /remember-saved Memory rows for the index — never copied agentId, so memories load with agentId=undefined and get dropped by the isolated-mode filter. Effect: in AGENTMEMORY_AGENT_SCOPE=isolated an agent cannot recall its own remembered memories (false negative, not a leak). Copy agentId in memoryToObservation so isolation applies to memories too.
wyh0626
force-pushed
the
fix/817-mem-search-isolation
branch
from
cc338e9 to
55443a6
Compare
|
Hi @rohitg00 — rebased onto latest main and trimmed this down to a one-liner, since the main #817 leak is already fixed in #849. That fix filters on obs.agentId !== filterAgentId, but memoryToObservation() never copies agentId. So /remember-saved memories load with agentId === undefined and get dropped in isolated mode — meaning an agent can't recall its own memories. |
1 participant
Summary
Closes #817.
AGENTMEMORY_AGENT_SCOPE=isolatedfiltered recall onmem::smart-search,/memories,/observations, and/sessions(PR #654), butmem::searchapplied no agent-scope filtering at all. It backsPOST /agentmemory/searchplus the MCP toolsmemory_recallandrecall_context, so an isolated worker for agent B could read agent A's memories.memory_recallis a primary recall tool, so this leaked in normal agent usage, not just via a raw REST call.Root cause: isolation in #654 (commit
1aec56a) was added per-function —smart-search.tsand the list endpoints gotgetAgentId()/isAgentScopeIsolated(), butsrc/functions/search.tswas never touched. The #654 CHANGELOG enumerates the filtered paths and omitsmem::search, matching the code.Fix
mem::searchnow mirrorsmem::smart-search. The BM25 index doesn't carryagentId, so it over-fetches and post-filters on the loaded observation:api::searchforwardsagentId— body first, then?agentId=— consistent with/memories,/observations,/sessions. When omitted,mem::searchfalls back to the envAGENT_ID.memory_recall/recall_contextforwardagentId.recall_contextalso scopes its directkv.list(KV.memories)pass, which previously filtered only byisLatestand leaked across agents.memoryToObservation()now copiesagentId. Memories saved via/rememberare indexed through this helper; without it, isolated recall dropped the owner's own remembered memories — exactly the/remember→/searchpath in the AGENT_ID isolated scope not enforced on mem::search (REST /search, memory_recall, recall_context) → cross-agent memory leak #817 repro.Verify
Reproducer from #817:
Unit:
I also reverted each source change and confirmed the new tests fail on the pre-fix code (3 leak cases + the
/rememberMemory case), so they actually catch the regression.Tests
6 cases in
test/search-agent-scope.test.ts:AGENT_ID's hitsagentIdoverride wins over the env defaultagentId: "*"opts out of isolationKV.memories(the/rememberpath throughmemoryToObservation)test/cross-project-isolation.test.ts: its partialconfig.jsmock gained the now-importedisAgentScopeIsolatedexport.Known limitations / follow-ups
mem::smart-search, agent-filtered recall over-fetches BM25 hits then post-filters. With a very smalllimitand a large out-of-scope head, an in-scope match beyond the fetch window can be missed — a false negative, not a leak. A durable fix indexesagentIdin BM25 or makes the over-fetch adaptive.agentIdoverride discoverability. The override is wired at the function and REST-body level (matchingsmart-search), but isn't exposed in the MCP tool schema or forwarded by the standalone MCP proxy. This mirrorssmart-search's existing behavior; default env isolation is unaffected.Summary by CodeRabbit