fix(mcp): proxy standalone tools to running server when reachable#161
fix(mcp): proxy standalone tools to running server when reachable#161
Conversation
Previously @agentmemory/mcp's handleToolCall only ever used a local InMemoryKV, so in the supported Claude Code setup (plugin + @agentmemory/agentmemory server) the MCP tools never saw the data hooks had captured — memory_sessions returned [] even when the REST endpoint had 69 observations, memory_smart_search searched a standalone.json that only contained explicit memory_save calls, and the server's BM25 + vector + graph indexing was inaccessible through MCP. Now handleToolCall resolves a Handle once per process: if a server responds to GET /agentmemory/livez within 500ms, all subsequent tool calls route through the REST endpoints and see exactly what hooks and the viewer see. If the probe fails, the caller falls back to the existing local InMemoryKV path so users who run MCP alone still work. - New src/mcp/rest-proxy.ts: probe + cached Handle with 15s call timeout. Honours AGENTMEMORY_URL (defaults to http://localhost:3111) and attaches Bearer AGENTMEMORY_SECRET when set. - src/mcp/standalone.ts: split handleProxy and handleLocal branches; announce active mode once to stderr so users can see which store their tools are reading. On per-request proxy failure, log and fall back to local KV for that single call. - test/mcp-standalone-proxy.test.ts: five cases covering proxy happy path, auth header, unreachable server fallback, and post-probe proxy failure fallback.
|
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 (4)
🚧 Files skipped from review as they are similar to previous changes (3)
📝 WalkthroughWalkthroughAdded a runtime resolver that probes a configured AgentMemory REST server and returns either a proxy handle (forwarding MCP tool calls to REST endpoints with auth and timeouts) or a local handle (InMemoryKV). Standalone MCP tool handler now uses this resolver, forwarding supported tools to the server and falling back to local KV per-request. Changes
Sequence Diagram(s)sequenceDiagram
participant Client as MCP Client
participant Handler as handleToolCall
participant Resolver as resolveHandle
participant Probe as Health Probe
participant ProxyAPI as REST API
participant LocalKV as InMemoryKV
Client->>Handler: invoke toolName, args
Handler->>Resolver: resolveHandle()
Resolver->>Probe: GET /agentmemory/livez (500ms)
alt Server reachable
Probe-->>Resolver: 200 OK
Resolver-->>Handler: ProxyHandle {mode: "proxy", call}
Handler->>ProxyAPI: call specific endpoint (e.g., POST /agentmemory/smart-search)
alt Proxy call succeeds
ProxyAPI-->>Handler: 200 + JSON
Handler-->>Client: MCP text response
else Proxy call fails
ProxyAPI--xHandler: error/status
Handler->>LocalKV: fallback local memory_* operation
LocalKV-->>Handler: local data
Handler-->>Client: MCP text response (from local)
end
else Server unreachable
Probe-xResolver: timeout/error
Resolver-->>Handler: LocalHandle {mode: "local"}
Handler->>LocalKV: local memory_* operation
LocalKV-->>Handler: local data
Handler-->>Client: MCP text response
end
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 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 docstrings
🧪 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 |
There was a problem hiding this comment.
Actionable comments posted: 4
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@src/mcp/rest-proxy.ts`:
- Around line 17-18: The cached Handle (cached) is never invalidated causing
permanent local-mode after an initial probe; update resolveHandle()/probe logic
to allow re-probing by adding an invalidateHandle() that clears cached and
probeInFlight (or set a short TTL on cached to auto-expire local-mode), and call
invalidateHandle() from the proxy-failure catch in handleToolCall (in
standalone.ts) so any proxy error clears the cache and triggers a fresh probe on
the next resolveHandle() attempt; ensure probeInFlight handling prevents
duplicate probes while allowing recovery after failures.
In `@src/mcp/standalone.ts`:
- Around line 79-156: handleProxy returns raw REST bodies via textResponse which
mismatches the shaped objects returned by handleLocal (causing inconsistent
contracts for memory_recall/memory_smart_search, memory_export, memory_audit,
memory_save, memory_governance_delete). Update handleProxy to normalize each
proxied response to the same shapes handleLocal produces: for
memory_recall/memory_smart_search return an array of Memory (extract hits or map
result.hits → Memory[]), for memory_export wrap into { version, memories,
sessions }, for memory_audit return an array of Entry, for memory_save return {
saved: id }, and for memory_governance_delete return { deleted, requested,
reason }; adjust usages of textResponse/handle.call inside handleProxy to parse
the proxied JSON and rebuild these shaped objects before returning so both proxy
and local branches present a stable contract.
- Around line 271-279: The catch in the call site that invokes handleProxy (in
the block that may fall back to handleLocal) is too broad and treats synchronous
validation errors from handleProxy as transport failures; update handling so
only actual proxy/transport failures trigger the fallback: introduce a dedicated
error class (e.g., ProxyTransportError) and ensure handleProxy throws that only
around the async transport/fetch/handle.call(...) sections while rethrowing
validation errors as-is; then change the try/catch here to catch only
ProxyTransportError (so validation errors propagate and are not retried locally)
and keep the existing fallback to handleLocal for true transport failures.
In `@test/mcp-standalone-proxy.test.ts`:
- Around line 74-86: The test currently allows a false positive because
seen.every(...) can pass when seen is empty or only contains the probe; update
the assertions in the "attaches Bearer token when AGENTMEMORY_SECRET is set"
test to explicitly verify that the proxied sessions request was recorded and was
authorized: after invoking handleToolCall("memory_sessions", {}), assert that
seen includes an entry for the sessions endpoint (e.g. contains
"/agentmemory/sessions" or not ending with "/agentmemory/livez"), assert
seen.length > 0 (or at least one non-probe entry exists), and assert that the
specific sessions entry's authorization portion ends with "Bearer s3cret" (or
that all non-probe entries end with that token); reference the seen array,
installFetch stub and handleToolCall invocation to locate and update the test.
🪄 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: 48b1ec45-f82f-4a53-b142-7ba010b45960
📒 Files selected for processing (3)
src/mcp/rest-proxy.tssrc/mcp/standalone.tstest/mcp-standalone-proxy.test.ts
| let cached: Handle | null = null; | ||
| let probeInFlight: Promise<Handle> | null = null; |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
fd -t f 'package.json' -d 2 --exec sh -c 'echo "== $0 =="; jq ".engines // {}" "$0"'Repository: rohitg00/agentmemory
Length of output: 210
🏁 Script executed:
#!/bin/bash
# Find all package.json files
fd -t f 'package.json' -x cat {} \; | head -50Repository: rohitg00/agentmemory
Length of output: 1547
🏁 Script executed:
#!/bin/bash
# Get the full package.json to see engines field
find . -name "package.json" -type f -exec cat {} \; | jq '.engines // "no engines specified"'Repository: rohitg00/agentmemory
Length of output: 113
🏁 Script executed:
#!/bin/bash
# Read rest-proxy.ts file
wc -l src/mcp/rest-proxy.tsRepository: rohitg00/agentmemory
Length of output: 88
🏁 Script executed:
#!/bin/bash
# Read the entire rest-proxy.ts file
cat -n src/mcp/rest-proxy.tsRepository: rohitg00/agentmemory
Length of output: 2970
🏁 Script executed:
#!/bin/bash
# Search for issue `#159` reference in the codebase
rg "#159|issue.*159|159.*issue" --type ts --type js --type md -iRepository: rohitg00/agentmemory
Length of output: 173
🏁 Script executed:
#!/bin/bash
# Read the test file to understand issue `#159` context
cat -n test/mcp-standalone-proxy.test.ts | head -100Repository: rohitg00/agentmemory
Length of output: 4788
🏁 Script executed:
#!/bin/bash
# Continue reading the test file to see more tests
cat -n test/mcp-standalone-proxy.test.ts | sed -n '100,150p'Repository: rohitg00/agentmemory
Length of output: 915
Resolved handle is cached permanently — no recovery from startup-order or server bouncing.
Once the first probe resolves, the cache is never invalidated. If the server was down at MCP startup but brought up later, MCP stays on local KV forever (writes never reach the server). If the server goes down after being probed as up, every call still hits proxy first, fails, and silently falls back to local; once it's back up, there's no re-sync and any memories written during the outage live only in local KV.
This is the split-data scenario issue #159 addresses. Consider invalidating cached when a proxy call throws (so the next resolveHandle() re-probes), and/or adding a short TTL on the local-mode cache so a later-started server is eventually discovered.
♻️ Sketch: invalidate on proxy failure + short local-mode TTL
-let cached: Handle | null = null;
+let cached: Handle | null = null;
+let cachedAt = 0;
+const LOCAL_MODE_TTL_MS = 30_000;
+
+export function invalidateHandle(): void {
+ cached = null;
+}
...
export async function resolveHandle(): Promise<Handle> {
- if (cached) return cached;
+ const now = new Date().toISOString();
+ if (cached) {
+ if (cached.mode === "local" && /* elapsed check */ ) {
+ cached = null;
+ } else {
+ return cached;
+ }
+ }
if (probeInFlight) return probeInFlight;
...
- cached = handle;
+ cached = handle;
+ cachedAt = now;
return handle;
}
const local: LocalHandle = { mode: "local" };
- cached = local;
+ cached = local;
+ cachedAt = now;
return local;Then call invalidateHandle() from handleToolCall's proxy-failure catch in standalone.ts.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@src/mcp/rest-proxy.ts` around lines 17 - 18, The cached Handle (cached) is
never invalidated causing permanent local-mode after an initial probe; update
resolveHandle()/probe logic to allow re-probing by adding an invalidateHandle()
that clears cached and probeInFlight (or set a short TTL on cached to
auto-expire local-mode), and call invalidateHandle() from the proxy-failure
catch in handleToolCall (in standalone.ts) so any proxy error clears the cache
and triggers a fresh probe on the next resolveHandle() attempt; ensure
probeInFlight handling prevents duplicate probes while allowing recovery after
failures.
| async function handleProxy( | ||
| toolName: string, | ||
| args: Record<string, unknown>, | ||
| kvInstance: InMemoryKV = kv, | ||
| handle: ProxyHandle, | ||
| ): Promise<{ content: Array<{ type: string; text: string }> }> { | ||
| switch (toolName) { | ||
| case "memory_save": { | ||
| const content = args["content"]; | ||
| if (typeof content !== "string" || !content.trim()) { | ||
| throw new Error("content is required"); | ||
| } | ||
| const payload = { | ||
| content, | ||
| type: (args["type"] as string) || "fact", | ||
| concepts: normalizeList(args["concepts"]), | ||
| files: normalizeList(args["files"]), | ||
| }; | ||
| const result = await handle.call("/agentmemory/remember", { | ||
| method: "POST", | ||
| body: JSON.stringify(payload), | ||
| }); | ||
| return textResponse(result); | ||
| } | ||
|
|
||
| case "memory_recall": | ||
| case "memory_smart_search": { | ||
| const query = args["query"]; | ||
| if (typeof query !== "string" || !query.trim()) { | ||
| throw new Error("query is required"); | ||
| } | ||
| const limit = parseLimit(args["limit"]); | ||
| const result = await handle.call("/agentmemory/smart-search", { | ||
| method: "POST", | ||
| body: JSON.stringify({ query: query.trim(), limit }), | ||
| }); | ||
| return textResponse(result, true); | ||
| } | ||
|
|
||
| case "memory_sessions": { | ||
| const limit = parseLimit(args["limit"], 20); | ||
| const result = await handle.call( | ||
| `/agentmemory/sessions?limit=${limit}`, | ||
| { method: "GET" }, | ||
| ); | ||
| return textResponse(result, true); | ||
| } | ||
|
|
||
| case "memory_governance_delete": { | ||
| const ids = normalizeList(args["memoryIds"]); | ||
| if (ids.length === 0) throw new Error("memoryIds is required"); | ||
| const result = await handle.call("/agentmemory/governance/memories", { | ||
| method: "POST", | ||
| body: JSON.stringify({ | ||
| memoryIds: ids, | ||
| reason: (args["reason"] as string) || "plugin skill request", | ||
| }), | ||
| }); | ||
| return textResponse(result); | ||
| } | ||
|
|
||
| case "memory_export": { | ||
| const result = await handle.call("/agentmemory/export", { method: "GET" }); | ||
| return textResponse(result, true); | ||
| } | ||
|
|
||
| case "memory_audit": { | ||
| const limit = parseLimit(args["limit"], 50); | ||
| const result = await handle.call( | ||
| `/agentmemory/audit?limit=${limit}`, | ||
| { method: "GET" }, | ||
| ); | ||
| return textResponse(result, true); | ||
| } | ||
|
|
||
| default: | ||
| throw new Error(`Unknown tool: ${toolName}`); | ||
| } | ||
| } |
There was a problem hiding this comment.
Response shapes diverge between proxy and local paths — per-request fallback will return different JSON structures for the same tool call.
Because handleProxy wraps the raw REST body in textResponse(...) while handleLocal builds a locally-shaped payload, callers see different structures depending on which branch served the request. Examples:
| tool | proxy returns (body of REST endpoint) | local returns |
|---|---|---|
memory_recall / memory_smart_search |
{ query, hits: [...] } (per your own test at lines 62, 70-71) |
Memory[] (array) — see handleLocal line 215 |
memory_export |
whatever /agentmemory/export returns |
{ version, memories, sessions } |
memory_audit |
whatever /agentmemory/audit returns |
Entry[] |
memory_save |
raw /remember body |
{ saved: id } |
memory_governance_delete |
raw body | { deleted, requested, reason } |
This matters because the PR explicitly introduces a per-request fallback (lines 271-279): if the server is up at probe time but a single call fails, the next call is served locally and will return a different shape than the previous successful proxied call. Agents/clients that parsed .hits[0] suddenly get an array; code that expected { saved } gets something else. This is a silent contract break.
Normalize the proxy responses to match the local shapes (or formally document a single shape and adapt both sides). Suggested minimum for recall/search since it's the hottest path and is already covered by divergent tests:
♻️ Normalize memory_recall / memory_smart_search response
case "memory_recall":
case "memory_smart_search": {
const query = args["query"];
if (typeof query !== "string" || !query.trim()) {
throw new Error("query is required");
}
const limit = parseLimit(args["limit"]);
const result = await handle.call("/agentmemory/smart-search", {
method: "POST",
body: JSON.stringify({ query: query.trim(), limit }),
});
- return textResponse(result, true);
+ const hits =
+ result && typeof result === "object" && Array.isArray((result as { hits?: unknown }).hits)
+ ? (result as { hits: unknown[] }).hits
+ : Array.isArray(result)
+ ? result
+ : [];
+ return textResponse(hits, true);
}Apply the same normalization pattern to memory_export (wrap into { version, memories, sessions }) and memory_audit (unwrap to array) so consumers see one stable contract regardless of mode.
Also applies to: 263-281
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@src/mcp/standalone.ts` around lines 79 - 156, handleProxy returns raw REST
bodies via textResponse which mismatches the shaped objects returned by
handleLocal (causing inconsistent contracts for
memory_recall/memory_smart_search, memory_export, memory_audit, memory_save,
memory_governance_delete). Update handleProxy to normalize each proxied response
to the same shapes handleLocal produces: for memory_recall/memory_smart_search
return an array of Memory (extract hits or map result.hits → Memory[]), for
memory_export wrap into { version, memories, sessions }, for memory_audit return
an array of Entry, for memory_save return { saved: id }, and for
memory_governance_delete return { deleted, requested, reason }; adjust usages of
textResponse/handle.call inside handleProxy to parse the proxied JSON and
rebuild these shaped objects before returning so both proxy and local branches
present a stable contract.
Four issues flagged by review.
1. Handle cache was never invalidated. Once resolveHandle() set
cached = local, a later server start would never be picked up
until the process restarted. Add invalidateHandle() + a 30s TTL
on local-mode caching so the next call after proxy failure or
timeout re-probes.
2. Response shapes diverged between proxy and local. memory_recall
and memory_smart_search returned { mode, query, results } from
the server but Memory[] from local; memory_audit returned an
array from local and whatever the server sent from proxy.
Normalize local responses to server-shaped envelopes
({ mode: compact, results: [...] } and { entries: [...] }) so
a per-request fallback no longer flips structure under the caller.
3. Validation errors from handleProxy were logged as 'proxy call
failed' and then retried against local, which re-threw the same
validation error. Extract validate(toolName, args) and run it
once before mode resolution, so empty content/query/memoryIds
throw cleanly with no misleading stderr line.
4. Bearer-token test was vacuous — seen.every() passes trivially
on an empty list and only the probe was guaranteed to record.
Replace with a Map<path, auth> assertion that pins both the
probe at /livez and the proxied tool call at /sessions to
Bearer s3cret.
Updates test/mcp-standalone.test.ts to match the new local
response shape (previously asserted raw Memory[] length, now
asserts results.length).
Bump version + ship CHANGELOG covering everything that merged since v0.8.13: - #118 security advisory drafts for v0.8.2 CVEs - #132 semantic eviction routing + batched retention audit - #157 iii console docs + vendored screenshots in README - #160 (#158) health gated on RSS floor - #161 (#159) standalone MCP proxies to the running server - #162 (#125) mem::forget audit coverage + policy doc - #163 (#62) @agentmemory/fs-watcher filesystem connector - #164 Next.js website (website/ root, ship to Vercel) Version bumps (8 files): - package.json / package-lock.json (top + packages['']) - plugin/.claude-plugin/plugin.json - packages/mcp/package.json (self + ~0.9.0 dep pin) - src/version.ts (union extended, assigned 0.9.0) - src/types.ts (ExportData.version union) - src/functions/export-import.ts (supportedVersions set) - test/export-import.test.ts (export assertion) Tests: 777 passing. Build clean.
Fixes #159.
Problem
In v0.8.13,
@agentmemory/mcpand@agentmemory/agentmemoryran two independent KV stores. When the supported Claude Code setup was active (plugin + server), MCP tools only saw the standalone JSON file — they never consulted the running server. So:memory_sessionsreturned[]even whenGET /agentmemory/sessionshad four sessions and 69 observations.memory_smart_searchfell back to a substring match on standalone data only, missing every hooks-captured observation.Reporter reproduced on Windows 11, Node v24.13.0, plugin v0.8.13.
Fix
handleToolCallnow resolves a Handle once per process:GET /agentmemory/livezatAGENTMEMORY_URL || http://localhost:3111with a 500ms timeout.InMemoryKVpath so pure-standalone users stay supported.Per-request proxy failures after a successful probe fall back to local KV for that call and log a stderr warning — no silent drop. Active mode is announced once to stderr so users can tell which store their tools are reading.
Bearer
AGENTMEMORY_SECRETis attached automatically when the server requires auth.Endpoints routed
memory_savePOST /agentmemory/remembermemory_recall/memory_smart_searchPOST /agentmemory/smart-searchmemory_sessionsGET /agentmemory/sessions?limit=Nmemory_governance_deletePOST /agentmemory/governance/memoriesmemory_exportGET /agentmemory/exportmemory_auditGET /agentmemory/audit?limit=NTests
test/mcp-standalone-proxy.test.ts— five cases:memory_sessionsproxies to GET/agentmemory/sessionswhen server is up.memory_smart_searchproxies to POST/agentmemory/smart-search.AGENTMEMORY_SECRETis set.InMemoryKV.All 759 tests still pass (the one pre-existing failure needs a live server on 3111 and is unrelated).
Summary by CodeRabbit
New Features
Changes
Tests