fix(scope): enforce AGENT_ID isolation on mem::search (closes #817)#821
Open
wyh0626 wants to merge 1 commit into
Open
fix(scope): enforce AGENT_ID isolation on mem::search (closes #817)#821wyh0626 wants to merge 1 commit into
wyh0626 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. |
Contributor
|
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 (6)
✅ Files skipped from review due to trivial changes (1)
🚧 Files skipped from review as they are similar to previous changes (4)
📝 WalkthroughWalkthroughThe PR adds optional agent-scoping to the ChangesAgent-scoped memory search and isolation
Sequence DiagramsequenceDiagram
participant Caller
participant Handler
participant MemSearch as "mem::search"
participant SearchIndex
participant KV
participant Enricher
Caller->>Handler: request (optional agentId)
Handler->>MemSearch: invoke mem::search (forward agentId only if explicit)
MemSearch->>SearchIndex: query index (over-fetch candidates if agent filter)
SearchIndex->>KV: load observation/memory blobs
KV->>MemSearch: return raw observations
MemSearch->>Enricher: enrich observations -> SearchResult (includes obs.agentId)
Enricher->>MemSearch: enriched results
MemSearch->>MemSearch: filter results where obs.agentId == filterAgentId
MemSearch->>Handler: return filtered results
Handler->>Caller: respond with scoped results
Estimated code review effort🎯 4 (Complex) | ⏱️ ~40 minutes Possibly related issues
Possibly related PRs
Suggested reviewers
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}
Contributor
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
Comment on lines
+69
to
+133
| 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(); | ||
| }); |
Contributor
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.
…0#817) mem::search applied no agent-scope filtering, so the REST /search endpoint and the memory_recall / recall_context MCP tools that route through it leaked memories across agents under AGENTMEMORY_AGENT_SCOPE=isolated. Mirror the mem::smart-search isolation logic: filter to the env AGENT_ID, support a per-call agentId override, and agentId="*" opt-out, over-fetching BM25 hits and post-filtering on the loaded observation. - api::search forwards agentId (body, then ?agentId=) - memory_recall / recall_context forward agentId; recall_context also scopes its direct kv.list(KV.memories) pass (previously isLatest-only) - memoryToObservation copies agentId so /remember memories aren't dropped in isolated recall (the rohitg00#817 repro path) Adds test/search-agent-scope.test.ts (6 cases) and gives the cross-project config.js mock the now-imported isAgentScopeIsolated export.
wyh0626
force-pushed
the
fix/817-mem-search-isolation
branch
from
9ffe441 to
674bf46
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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
New Features
/agentmemory/searchand recall prompts accept an optional agent identifier parameter.Tests