feat: implement concept graph search with depth-2 BFS expansion#349
feat: implement concept graph search with depth-2 BFS expansion#349Tanmay-008 wants to merge 3 commits into
Conversation
This feature improves search recall by utilizing concept co-occurrence relationships: - Added 'mem:concept-edges' KV scope to store graph data. - Automatically extract and reinforce concept pairs on 'mem::remember'. - Added 'mem::concept-graph-search' with Depth-2 BFS, lazy decay, and Top-10 node cap. - Integrated graph mode into 'smart-search' with a 0.8 relevance penalty. - Registered new 'memory_graph_search' MCP tool. - Added first-run auto-backfill migration for existing memories.
|
@Tanmay-008 is attempting to deploy a commit to the rohitg00's projects Team on Vercel. A member of the Team first needs to authorize it. |
📝 WalkthroughWalkthroughThis PR adds concept-edge storage and reinforcement, a bounded BFS concept-graph search that scores latest memories, smart-search graph-mode integration, a concept-edge backfill, an MCP tool and server handler, viewer UI for graph visualization, tests, and related wiring and timeout tweaks. ChangesConcept Graph Search Implementation
Sequence Diagram(s)(omitted — interactions are captured in the hidden review stack artifact) Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 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)
Warning There were issues while running some tools. Please review the errors and either fix the tool's configuration or disable the tool if it's a critical failure. 🔧 ESLint
ESLint skipped: no ESLint configuration detected in root package.json. To enable, add 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: 9
🤖 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 `@AGENTS.md`:
- Line 114: AGENTS.md currently lists 45 MCP tools but README.md and plugin.json
list 52; update AGENTS.md to say 52 and ensure the MCP tool count and any
related lists are synchronized across all locations (tools-registry.ts,
server.ts, triggers/api.ts, index.ts, test/mcp-standalone.test.ts, README.md,
plugin/.claude-plugin/plugin.json and AGENTS.md) so every reference and test
reflects 52 tools.
In `@src/functions/concept-backfill.ts`:
- Around line 14-17: The current backfill marks CONFIG_KEY done non-atomically
and counts failed upserts as processed; fix by making CONFIG_KEY updates atomic
and only setting done when all upserts succeed: implement a distributed
lock/lease or CAS on CONFIG_KEY (use kv.compare-and-set / put-if-absent
semantics) so the initial kv.get/flag check and the transition to
in_progress/done are atomic and prevent concurrent runs, change the loop that
calls the upsert function (concept upsert / whatever method currently swallowing
errors) to propagate or record failures instead of incrementing the processed
count on exception, and only write { done: true } to CONFIG_KEY when no upsert
errors remain (or persist a retryable cursor/state rather than done=true on
partial failure). Ensure references to CONFIG_KEY, kv.get, kv.put/CAS and the
upsert function are updated accordingly.
In `@src/functions/concept-graph-search.ts`:
- Around line 9-14: decayedStrength can produce NaN when edge.lastSeenAt is
malformed; update the function to validate the parsed date and fallback to a
safe value (e.g., treat daysSinceLastSeen as 0 or use Date.now()) when new
Date(edge.lastSeenAt) is invalid, and also guard edge.strength (ensure it's a
finite number) before computing decay so the final return Math.max(0.05, ...)
never receives NaN. Change decayedStrength to check Number.isFinite on the
parsed time and computed daysSinceLastSeen/decay and use sensible defaults so
malformed lastSeenAt won't pollute scores.
- Around line 24-31: The current depth validation in concept-graph-search.ts
allows zero, negative, and non-integer values which can break BFS; tighten
validation in the handler before running BFS by verifying data.depth (and the
local const depth) is an integer within 1..MAX_BFS_DEPTH: reject non-numeric,
non-integer, <=0, or >MAX_BFS_DEPTH with an error (e.g., "depth_out_of_range" or
"invalid_depth") and clear message; update the code paths that use depth (the
const depth, the BFS loop) to rely on this validated integer and ensure you do
not call sdk.trigger or run traversal with unvalidated raw request values.
In `@src/functions/smart-search.ts`:
- Around line 89-137: The graph-mode try/catch currently swallows errors but the
function still reports mode: "graph"; modify the block around sdk.trigger (and
the catch) to detect failure: log the caught error via logger.error including
context (e.g., data.query), set a local flag like graphSucceeded = false
(default true before the try) or set data.modeFallback = "compact" on failure,
and ensure the final return uses that flag (return mode: graphSucceeded ?
"graph" : "compact"); keep existing graph expansion logic intact but do not
suppress errors silently.
In `@src/index.ts`:
- Line 373: The sdk.trigger call for the concept backfill (sdk.trigger({
function_id: "mem::concept-backfill", payload: {} })) is currently swallowing
errors with .catch(() => {}) — change this to handle failures explicitly by
catching the error and logging/reporting it (e.g., use your logger or
processLogger to emit a descriptive error that includes the function_id and the
caught error), and optionally surface or retry the failure (simple retry/backoff
or rethrow to crash on startup) so first-run migration failures are visible and
diagnosable; update the catch handler on the sdk.trigger invocation to implement
this logging/reporting and any desired retry/backoff behavior.
In `@src/mcp/server.ts`:
- Around line 470-482: The handler currently allows
non-integer/zero/negative/fractional depth and can call sdk.trigger with an
empty concepts array; update validation so depth = asNumber(args.depth, 2) must
be an integer between 1 and 2 (reject if not Number.isInteger(depth) or depth <
1 or depth > 2 with status_code 400 and a clear error code like
"depth_out_of_range"), and after tokenizing args.query into concepts (the
existing map/filter), ensure concepts.length > 0 and reject with status_code 400
and an error like "empty_concepts" if no valid tokens remain; only call
sdk.trigger({ function_id: "mem::concept-graph-search", payload: { concepts,
depth, limit } }) after these checks.
In `@src/triggers/api.ts`:
- Around line 852-858: The payload for sdk.trigger("mem::smart-search")
currently forwards body.expandIds without type checks, which can break
mem::smart-search when it expects an array; update the request handling in
src/triggers/api.ts (the block that builds payload and calls sdk.trigger) to
validate and whitelist types: ensure body.expandIds is an array (use
Array.isArray) before assigning payload.expandIds (otherwise omit it or coerce
to an array), validate body.query is a string before assigning payload.query,
and validate body.limit is a number before assigning payload.limit; keep the
existing check for body.mode === "graph" and do not pass any raw body fields to
sdk.trigger.
In `@test/mcp-standalone.test.ts`:
- Line 34: The test uses a weak assertion
expect(tools.length).toBeGreaterThanOrEqual(42); which allows regressions —
change this to assert the exact current expected tool count (e.g.,
expect(tools.length).toBe(<CURRENT_EXPECTED_COUNT>)) or a much tighter baseline,
replacing the 42 literal; update the test variable reference (tools) accordingly
and, when changing the canonical tool count, also update the corresponding
locations mentioned in the comment (tools-registry.ts, server.ts,
triggers/api.ts, index.ts, README.md, plugin/.claude-plugin/plugin.json) so the
codebase stays consistent.
🪄 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: 56560752-f5ec-4911-a174-47c5a75ac35f
📒 Files selected for processing (17)
AGENTS.mdREADME.mdplugin/.claude-plugin/plugin.jsonsrc/functions/concept-backfill.tssrc/functions/concept-edges.tssrc/functions/concept-graph-search.tssrc/functions/remember.tssrc/functions/smart-search.tssrc/index.tssrc/mcp/server.tssrc/mcp/tools-registry.tssrc/state/schema.tssrc/triggers/api.tssrc/types.tstest/concept-edges.test.tstest/concept-graph-search.test.tstest/mcp-standalone.test.ts
| ## Current Stats (v0.8.9) | ||
|
|
||
| - 44 MCP tools (8 visible by default, `AGENTMEMORY_TOOLS=all` for all) | ||
| - 45 MCP tools (8 visible by default, `AGENTMEMORY_TOOLS=all` for all) |
There was a problem hiding this comment.
MCP tool count inconsistent with README and plugin.json.
AGENTS.md shows 45 MCP tools, but README.md and plugin.json both show 52 MCP tools. Per the review stack context, Layer 1 should update documentation to reflect 52 tools.
📝 Suggested fix
-- 45 MCP tools (8 visible by default, `AGENTMEMORY_TOOLS=all` for all)
+- 52 MCP tools (8 visible by default, `AGENTMEMORY_TOOLS=all` for all)Based on learnings: When adding or removing MCP tools, update all six locations: tools-registry.ts, server.ts, triggers/api.ts, index.ts, test/mcp-standalone.test.ts, README.md, and plugin/.claude-plugin/plugin.json.
📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| - 45 MCP tools (8 visible by default, `AGENTMEMORY_TOOLS=all` for all) | |
| - 52 MCP tools (8 visible by default, `AGENTMEMORY_TOOLS=all` for all) |
🤖 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 `@AGENTS.md` at line 114, AGENTS.md currently lists 45 MCP tools but README.md
and plugin.json list 52; update AGENTS.md to say 52 and ensure the MCP tool
count and any related lists are synchronized across all locations
(tools-registry.ts, server.ts, triggers/api.ts, index.ts,
test/mcp-standalone.test.ts, README.md, plugin/.claude-plugin/plugin.json and
AGENTS.md) so every reference and test reflects 52 tools.
| const flag = await kv.get<{ done: boolean }>(KV.config, CONFIG_KEY); | ||
| if (flag?.done) { | ||
| return { success: true, skipped: true, reason: "already completed" }; | ||
| } |
There was a problem hiding this comment.
Backfill can be permanently marked “done” after partial failure and can double-run concurrently.
Line 36 swallows upsert failures, Line 39 still counts them as processed, and Line 42 sets done: true anyway. Also, the check at Line 14 and final set at Line 42 are non-atomic, so concurrent invocations can both process and reinforce the same edges.
Suggested direction
- const flag = await kv.get<{ done: boolean }>(KV.config, CONFIG_KEY);
+ const flag = await kv.get<{ done: boolean; inProgress?: boolean }>(KV.config, CONFIG_KEY);
if (flag?.done) {
return { success: true, skipped: true, reason: "already completed" };
}
+ if (flag?.inProgress) {
+ return { success: true, skipped: true, reason: "already running" };
+ }
+ await kv.set(KV.config, CONFIG_KEY, { done: false, inProgress: true });
const memories = await kv.list<Memory>(KV.memories);
@@
- let processed = 0;
+ let processed = 0;
+ let failed = 0;
@@
- await Promise.all(
- batch.map((m) =>
- sdk
- .trigger({
- function_id: "mem::concept-edge-upsert",
- payload: { concepts: m.concepts },
- })
- .catch(() => {}),
- ),
- );
- processed += batch.length;
+ const settled = await Promise.allSettled(
+ batch.map((m) =>
+ sdk.trigger({
+ function_id: "mem::concept-edge-upsert",
+ payload: { concepts: m.concepts },
+ }),
+ ),
+ );
+ for (const r of settled) {
+ if (r.status === "fulfilled") processed += 1;
+ else failed += 1;
+ }
}
- await kv.set(KV.config, CONFIG_KEY, { done: true, completedAt: new Date().toISOString() });
+ const completedAt = new Date().toISOString();
+ await kv.set(KV.config, CONFIG_KEY, {
+ done: failed === 0,
+ inProgress: false,
+ completedAt,
+ failed,
+ });Also applies to: 29-43
🤖 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 `@src/functions/concept-backfill.ts` around lines 14 - 17, The current backfill
marks CONFIG_KEY done non-atomically and counts failed upserts as processed; fix
by making CONFIG_KEY updates atomic and only setting done when all upserts
succeed: implement a distributed lock/lease or CAS on CONFIG_KEY (use
kv.compare-and-set / put-if-absent semantics) so the initial kv.get/flag check
and the transition to in_progress/done are atomic and prevent concurrent runs,
change the loop that calls the upsert function (concept upsert / whatever method
currently swallowing errors) to propagate or record failures instead of
incrementing the processed count on exception, and only write { done: true } to
CONFIG_KEY when no upsert errors remain (or persist a retryable cursor/state
rather than done=true on partial failure). Ensure references to CONFIG_KEY,
kv.get, kv.put/CAS and the upsert function are updated accordingly.
There was a problem hiding this comment.
Actionable comments posted: 10
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
plugin/scripts/session-end.mjs (1)
27-61:⚠️ Potential issue | 🟠 Major | 🏗️ Heavy liftSequential execution of fire-and-forget operations blocks session-end for up to 4 minutes.
The four fetch calls are awaited sequentially (30s + 60s + 120s + 30s = 240s worst-case). Since all catch blocks are empty, these independent operations should run in parallel using
Promise.allSettled()instead, allowing session-end to complete in roughly the longest individual timeout (~120s) rather than the sum of all timeouts.Proposed refactor to parallel execution
- try { - await fetch(`${REST_URL}/agentmemory/session/end`, { - method: "POST", - headers: authHeaders(), - body: JSON.stringify({ sessionId }), - signal: AbortSignal.timeout(3e4) - }); - } catch {} - if (process.env["CONSOLIDATION_ENABLED"] === "true") { - try { - await fetch(`${REST_URL}/agentmemory/crystals/auto`, { - method: "POST", - headers: authHeaders(), - body: JSON.stringify({ olderThanDays: 0 }), - signal: AbortSignal.timeout(6e4) - }); - } catch {} - try { - await fetch(`${REST_URL}/agentmemory/consolidate-pipeline`, { - method: "POST", - headers: authHeaders(), - body: JSON.stringify({ - tier: "all", - force: true - }), - signal: AbortSignal.timeout(12e4) - }); - } catch {} - } - if (process.env["CLAUDE_MEMORY_BRIDGE"] === "true") try { - await fetch(`${REST_URL}/agentmemory/claude-bridge/sync`, { - method: "POST", - headers: authHeaders(), - signal: AbortSignal.timeout(3e4) - }); - } catch {} + const requests = [ + fetch(`${REST_URL}/agentmemory/session/end`, { + method: "POST", + headers: authHeaders(), + body: JSON.stringify({ sessionId }), + signal: AbortSignal.timeout(3e4) + }) + ]; + + if (process.env["CONSOLIDATION_ENABLED"] === "true") { + requests.push( + fetch(`${REST_URL}/agentmemory/crystals/auto`, { + method: "POST", + headers: authHeaders(), + body: JSON.stringify({ olderThanDays: 0 }), + signal: AbortSignal.timeout(6e4) + }), + fetch(`${REST_URL}/agentmemory/consolidate-pipeline`, { + method: "POST", + headers: authHeaders(), + body: JSON.stringify({ + tier: "all", + force: true + }), + signal: AbortSignal.timeout(12e4) + }) + ); + } + + if (process.env["CLAUDE_MEMORY_BRIDGE"] === "true") { + requests.push( + fetch(`${REST_URL}/agentmemory/claude-bridge/sync`, { + method: "POST", + headers: authHeaders(), + signal: AbortSignal.timeout(3e4) + }) + ); + } + + await Promise.allSettled(requests);🤖 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 `@plugin/scripts/session-end.mjs` around lines 27 - 61, The current session-end logic awaits four independent fetch calls sequentially (the calls to `${REST_URL}/agentmemory/session/end`, `/agentmemory/crystals/auto`, `/agentmemory/consolidate-pipeline`, and `/agentmemory/claude-bridge/sync`), which can block for the sum of their timeouts; change this to fire them in parallel and ignore individual failures by collecting the fetch Promises into an array and awaiting Promise.allSettled(...) once. Keep the same conditional checks for process.env["CONSOLIDATION_ENABLED"] and process.env["CLAUDE_MEMORY_BRIDGE"], create each fetch call (with authHeaders, body and AbortSignal) without awaiting immediately, push them into the array, then after building the list call await Promise.allSettled(requests) inside a single try/catch (or no-op catch) so session-end completes in roughly the longest timeout instead of the sum.src/viewer/index.html (1)
1523-1532:⚠️ Potential issue | 🔴 Critical | ⚡ Quick winRestore the
resize()helper name here.This block defines
checkRestOnline(), but the code below still callsresize()and registersresizeas the window handler. Opening the Graph tab now throws immediately.Suggested fix
- function checkRestOnline() { + function resize() { var r = canvas.parentElement.getBoundingClientRect(); canvas.width = r.width * window.devicePixelRatio; canvas.height = r.height * window.devicePixelRatio; canvas.style.width = r.width + 'px'; canvas.style.height = r.height + 'px'; graphSim.ctx.setTransform(window.devicePixelRatio, 0, 0, window.devicePixelRatio, 0, 0); }🤖 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 `@src/viewer/index.html` around lines 1523 - 1532, The function was accidentally named checkRestOnline() while the code below calls resize() and registers resize as the window handler; rename the function back to resize (or change the call/handler to use checkRestOnline consistently) so the function referenced by window.addEventListener('resize', ...) and the resize() invocation exists; update the function declaration that manipulates canvas (the one using canvas.parentElement.getBoundingClientRect(), canvas.width/height, canvas.style.width/height, and graphSim.ctx.setTransform) to match the handler name.
♻️ Duplicate comments (2)
src/functions/concept-backfill.ts (1)
14-48:⚠️ Potential issue | 🟠 Major | ⚖️ Poor tradeoffConcurrency TOCTOU between flag read and flag write is still unresolved.
The partial-failure path is now correctly handled by the
throwat Line 44–46 before thedoneflag is written. However, two concurrent invocations ofmem::concept-backfill(e.g., a startup trigger racing with a manually issued call) can both readflag?.done === falseat Line 14, both pass the guard at Line 15, and both fan outmem::concept-edge-upserttriggers across the same eligible set — double-reinforcing every edge strength. AninProgresslease set immediately after thedonecheck (or a CAS/put-if-absent) would prevent this.🤖 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 `@src/functions/concept-backfill.ts` around lines 14 - 48, Add an atomic in-progress lease/check before processing to avoid TOCTOU: after reading CONFIG_KEY from KV.config in mem::concept-backfill, attempt an atomic "put-if-absent" for an inProgress marker (e.g., set CONFIG_KEY.inProgress = { owner: <id>, expiresAt: <ISO> }) using your KV client's compare-and-set/put-if-absent API; if the atomic put fails because inProgress is already present and not expired, return skipped; if it succeeds proceed with the batch work, and ensure you remove/clear the inProgress marker on both success (before writing done=true) and on error (so retries can run), and keep the existing error throw behavior so partial failures still do not write done=true.src/triggers/api.ts (1)
852-860:⚠️ Potential issue | 🟠 Major | ⚡ Quick winPayload validation is incomplete and doesn't fully address past review feedback.
The past review comment requested validation to ensure
queryis a non-empty string,expandIdsis an array of strings, andlimitis a positive integer. The current implementation:
- Line 854: Accepts any string
query, including empty strings that will fail downstream (line 68 ofsmart-search.tsrejects empty queries).- Lines 855-857: Filters
expandIdsto strings but doesn't validate array element structure (could be objects that coerce to strings).- Line 858: Accepts any number
limit, including negatives, zero, or non-integers.🔒 Suggested fix matching the past review
const body = (req.body ?? {}) as Record<string, unknown>; const payload: Record<string, unknown> = {}; - if (typeof body.query === "string") payload.query = body.query; + if (typeof body.query === "string") { + const trimmed = body.query.trim(); + if (!trimmed) { + return { status_code: 400, body: { error: "query must be a non-empty string" } }; + } + payload.query = trimmed; + } if (Array.isArray(body.expandIds)) { - payload.expandIds = body.expandIds.filter((id) => typeof id === "string"); + if (!body.expandIds.every((id) => typeof id === "string" || (typeof id === "object" && id !== null && typeof id.obsId === "string"))) { + return { status_code: 400, body: { error: "expandIds must be an array of strings or {obsId, sessionId} objects" } }; + } + payload.expandIds = body.expandIds; } - if (typeof body.limit === "number") payload.limit = body.limit; + if (body.limit !== undefined) { + if (!Number.isInteger(body.limit) || (body.limit as number) < 1) { + return { status_code: 400, body: { error: "limit must be a positive integer" } }; + } + payload.limit = body.limit; + } if (body.mode === "graph") payload.mode = "graph";As per coding guidelines: "Validate inputs at system boundaries (MCP handlers, REST endpoints) — never pass raw request body to sdk.trigger()."
🤖 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 `@src/triggers/api.ts` around lines 852 - 860, Validate and sanitize the request before calling sdk.trigger for the "mem::smart-search" payload: ensure body.query is a non-empty string (reject or return 400 if missing/empty), ensure body.expandIds is an actual array and that every element is a primitive string (filtering alone isn't enough—validate each item type), and ensure body.limit is a positive integer (coerce/validate and reject non-integers, <=0, or NaN); build payload only from these validated values and return an error response when validation fails instead of passing invalid data to sdk.trigger({ function_id: "mem::smart-search", payload }).
🧹 Nitpick comments (1)
src/functions/smart-search.ts (1)
108-108: ⚖️ Poor tradeoffLoading all memories into memory could be expensive.
kv.list<Memory>(KV.memories)loads the entire memory collection. For repos with thousands of memories, this creates memory pressure and latency.Consider batching or streaming if
KV.memoriesgrows large, or fetch only the memories referenced bygraphResult.results[].memoryId.🤖 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 `@src/functions/smart-search.ts` at line 108, The current call using kv.list<Memory>(KV.memories) loads the entire memories collection into the memories variable; instead, change the logic in smart-search.ts to only fetch referenced memories (from graphResult.results[].memoryId) and avoid kv.list for large datasets: collect the unique memoryId values from graphResult.results, then load them in batches (or use a bulk/GET-many API or paginated/streamed fetch) to populate the memories map used downstream (replace direct use of kv.list with batched kv.get/kv.getMany calls keyed by those IDs and handle missing items).
🤖 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/concept-backfill.ts`:
- Around line 57-60: The summary log mixes counts from different sets:
`processed` is based on the filtered `eligible` array while `total:
memories.length` is the full unfiltered list; update the `logger.info` call (the
line that currently logs "Concept backfill completed") to either log `total:
eligible.length` instead of `memories.length` or include both values (e.g.,
`eligible: eligible.length, totalMemories: memories.length`) so the output
unambiguously reports both the eligible count and the overall memory count.
In `@src/functions/smart-search.ts`:
- Around line 134-136: The code leaves finalMode set to "graph" when
graphSuccess is false but compact is empty, causing a misleading mode in the
response; update the fallback logic in the smart-search handler so that when
graphSuccess === false you set finalMode = "compact" if compact.length > 0,
otherwise set finalMode to a distinct no-result value (e.g., "none" or "empty")
instead of leaving "graph"; change the conditional that currently references
graphSuccess and compact to explicitly handle the empty-compact branch so
responses reflect the true mode (use the finalMode, graphSuccess, and compact
symbols to locate and modify the logic).
- Line 99: The current naive tokenization in smart-search.ts (where concepts is
derived via data.query.split(/\s+/).filter((t) => t.length > 1)) should be
replaced with a normalized tokenizer: lowercase the query, remove or normalize
punctuation/possessives, remove stop words, and apply stemming or lemmatization
(or use a small tokenizer library) before filtering short tokens; update the
code that builds the concepts array (the concepts variable) to perform these
steps so tokens like "User's" become "user" and plural/singular forms match.
- Line 121: The code is using an unsafe cast "type: mem.type as any" which
bypasses TypeScript checks and risks runtime errors if Memory.type diverges from
CompactSearchResult['type']; fix by aligning types or validating at runtime:
replace the cast by either narrowing/mapping Memory.type to the exact
CompactSearchResult['type'] (e.g., use an explicit mapping or switch between
known enum/union cases) or implement a type guard function (e.g.,
isValidCompactType(value): value is CompactSearchResult['type']) and set type
only after the guard passes (handle invalid cases explicitly). Target the
assignment where mem.type is used in smart-search.ts and update Memory.type /
CompactSearchResult['type'] handling accordingly.
In `@src/triggers/api.ts`:
- Line 919: Replace the raw parseInt call with the existing
parseOptionalPositiveInt helper: call
parseOptionalPositiveInt(req.query_params["limit"]) and set limit = result ??
100; then validate that limit is a positive integer within acceptable bounds
(e.g. enforce an upper cap like MAX_LIMIT or 1000) and return a 400 error if the
parsed value is invalid or out-of-range; update references to the variable name
limit and use parseOptionalPositiveInt (defined earlier in the file) to locate
the helper.
- Around line 915-928: Update the published endpoint counts to reflect the newly
added api::concept-graph-viewer endpoint: increment the "107 REST" count to "108
REST" in the index file and change the README badge string "104-endpoints" to
"105-endpoints"; verify the two documented totals (the REST count and the badge)
now match the actual number of endpoints and adjust if the discrepancy persists.
In `@src/viewer/index.html`:
- Line 1103: The tab-switch case references state.concepts.loaded and
loadConcepts() which don't exist; add a concepts sub-state and loader hookup:
ensure the global state object contains a concepts = { loaded: false, data: null
} (or equivalent), implement or import an async loadConcepts() function that
fetches/initializes concept data and sets state.concepts.loaded = true and
state.concepts.data accordingly, and after awaiting loadConcepts() call the
existing renderConcepts() (or trigger the renderer) so the tab has valid data
before switching; update the 'concepts' case to guard for missing state.concepts
(e.g., initialize lazily) and only await loadConcepts() when
!state.concepts.loaded.
- Around line 1519-1521: The Concepts tab is only shown conditionally using the
literal '__AGENTMEMORY_CONCEPT_GRAPH_VIEWER__' check at bootstrap, but the tab
remains hidden until initGraph() runs later when the Graph tab is opened; move
or duplicate the visibility logic to run during initial UI bootstrap so the
element with id 'tab-concepts' is set to display 'inline-block' as soon as the
page loads (i.e., evaluate the feature flag at startup, not inside initGraph()),
referencing the feature flag '__AGENTMEMORY_CONCEPT_GRAPH_VIEWER__', the DOM id
'tab-concepts', and the initGraph() flow to ensure the tab is visible
immediately when the flag is enabled.
- Around line 2062-2068: The code marks the Concepts tab loaded too early by
setting el.dataset.loaded = 'true' before the API call; change this so the
loaded flag is only set after a successful response: keep the initial guard (if
(el.dataset.loaded) return) but remove the early assignment, call
apiGet('concept-graph-viewer?limit=50') as before, and inside the .then handler
set el.dataset.loaded = 'true' only when data && data.success (after updating
the `#concepts-sidebar`); do not set the flag on failure so retries remain
possible.
- Around line 1133-1137: The code references an undefined variable tabId causing
a ReferenceError; remove the conditional branches that use tabId (the lines
calling renderDashboard/renderGraph/renderConcepts/renderMemories guarded by
tabId) or replace them with a real source-of-truth (e.g., a defined
currentTab/getActiveTab function or reading from location/hash) and then call
the corresponding render function; ensure renderDashboard() remains as the
fallback call (the lone renderDashboard() line) and update any event wiring to
use the chosen source-of-truth instead of tabId.
---
Outside diff comments:
In `@plugin/scripts/session-end.mjs`:
- Around line 27-61: The current session-end logic awaits four independent fetch
calls sequentially (the calls to `${REST_URL}/agentmemory/session/end`,
`/agentmemory/crystals/auto`, `/agentmemory/consolidate-pipeline`, and
`/agentmemory/claude-bridge/sync`), which can block for the sum of their
timeouts; change this to fire them in parallel and ignore individual failures by
collecting the fetch Promises into an array and awaiting Promise.allSettled(...)
once. Keep the same conditional checks for process.env["CONSOLIDATION_ENABLED"]
and process.env["CLAUDE_MEMORY_BRIDGE"], create each fetch call (with
authHeaders, body and AbortSignal) without awaiting immediately, push them into
the array, then after building the list call await Promise.allSettled(requests)
inside a single try/catch (or no-op catch) so session-end completes in roughly
the longest timeout instead of the sum.
In `@src/viewer/index.html`:
- Around line 1523-1532: The function was accidentally named checkRestOnline()
while the code below calls resize() and registers resize as the window handler;
rename the function back to resize (or change the call/handler to use
checkRestOnline consistently) so the function referenced by
window.addEventListener('resize', ...) and the resize() invocation exists;
update the function declaration that manipulates canvas (the one using
canvas.parentElement.getBoundingClientRect(), canvas.width/height,
canvas.style.width/height, and graphSim.ctx.setTransform) to match the handler
name.
---
Duplicate comments:
In `@src/functions/concept-backfill.ts`:
- Around line 14-48: Add an atomic in-progress lease/check before processing to
avoid TOCTOU: after reading CONFIG_KEY from KV.config in mem::concept-backfill,
attempt an atomic "put-if-absent" for an inProgress marker (e.g., set
CONFIG_KEY.inProgress = { owner: <id>, expiresAt: <ISO> }) using your KV
client's compare-and-set/put-if-absent API; if the atomic put fails because
inProgress is already present and not expired, return skipped; if it succeeds
proceed with the batch work, and ensure you remove/clear the inProgress marker
on both success (before writing done=true) and on error (so retries can run),
and keep the existing error throw behavior so partial failures still do not
write done=true.
In `@src/triggers/api.ts`:
- Around line 852-860: Validate and sanitize the request before calling
sdk.trigger for the "mem::smart-search" payload: ensure body.query is a
non-empty string (reject or return 400 if missing/empty), ensure body.expandIds
is an actual array and that every element is a primitive string (filtering alone
isn't enough—validate each item type), and ensure body.limit is a positive
integer (coerce/validate and reject non-integers, <=0, or NaN); build payload
only from these validated values and return an error response when validation
fails instead of passing invalid data to sdk.trigger({ function_id:
"mem::smart-search", payload }).
---
Nitpick comments:
In `@src/functions/smart-search.ts`:
- Line 108: The current call using kv.list<Memory>(KV.memories) loads the entire
memories collection into the memories variable; instead, change the logic in
smart-search.ts to only fetch referenced memories (from
graphResult.results[].memoryId) and avoid kv.list for large datasets: collect
the unique memoryId values from graphResult.results, then load them in batches
(or use a bulk/GET-many API or paginated/streamed fetch) to populate the
memories map used downstream (replace direct use of kv.list with batched
kv.get/kv.getMany calls keyed by those IDs and handle missing items).
🪄 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: 0737e0e6-5821-4102-a9fb-b6a0ab172daa
📒 Files selected for processing (11)
plugin/scripts/session-end.mjsplugin/scripts/stop.mjssrc/functions/concept-backfill.tssrc/functions/concept-graph-search.tssrc/functions/smart-search.tssrc/index.tssrc/mcp/server.tssrc/triggers/api.tssrc/viewer/document.tssrc/viewer/index.htmltest/mcp-standalone.test.ts
✅ Files skipped from review due to trivial changes (1)
- plugin/scripts/stop.mjs
🚧 Files skipped from review as they are similar to previous changes (3)
- src/mcp/server.ts
- src/index.ts
- src/functions/concept-graph-search.ts
There was a problem hiding this comment.
Actionable comments posted: 3
🤖 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/smart-search.ts`:
- Around line 89-90: The code sets let finalMode = data.mode but data.mode is
optional; change initialization to default to "compact" when undefined (e.g.,
finalMode = data.mode ?? "compact") so the normal query path preserves the
existing compact-search response contract; update both occurrences where
finalMode is assigned (the first assignment around variable finalMode and the
second occurrence near the later branch at line 146) and use the same defaulting
pattern so callers that omit mode continue to receive "compact".
- Around line 109-111: The code currently calls kv.list(KV.memories) to build
memoryMap which scans all memories; instead, retrieve only the memory IDs
returned by the graph search (the compact array / mem::concept-graph-search
results) and fetch those Memory entries in parallel with Promise.all using
kv.get for each id, then build memoryMap from those fetched results; keep
existingObsIds and compact usage but replace memories/memoryMap population with
the targeted parallel gets to avoid full-store scanning.
In `@src/viewer/index.html`:
- Line 1104: Add an "in-flight" guard to prevent starting multiple fetch/render
pipelines for Concepts: in the code path where you check state.concepts.loaded
and call renderConcepts(), introduce a boolean like state.concepts.loading (or
conceptsInFlight) that is checked before calling renderConcepts() and set to
true immediately when starting the pipeline (e.g., at the top of renderConcepts
or before initiating the fetch). Ensure this flag is cleared in a finally block
after the fetch/render completes or errors so subsequent tab switches can retry;
also ensure any RAF loops or window listeners started by renderConcepts are only
registered if loading was not already in-flight and are cleaned up on
completion.
🪄 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: 70cedc9f-0cfc-4634-9222-b9da0816c54a
📒 Files selected for processing (6)
README.mdsrc/functions/concept-backfill.tssrc/functions/smart-search.tssrc/index.tssrc/triggers/api.tssrc/viewer/index.html
✅ Files skipped from review due to trivial changes (1)
- README.md
🚧 Files skipped from review as they are similar to previous changes (3)
- src/functions/concept-backfill.ts
- src/index.ts
- src/triggers/api.ts
| let finalMode = data.mode; | ||
| if (data.mode === "graph") { |
There was a problem hiding this comment.
Default finalMode to "compact".
data.mode is optional, so the normal query path now returns { mode: undefined } whenever the caller omits mode. That changes the existing compact-search response contract and can break callers that branch on mode.
🛠 Suggested fix
- let finalMode = data.mode;
+ let finalMode: "compact" | "graph" = data.mode === "graph" ? "graph" : "compact";Also applies to: 146-146
🤖 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 `@src/functions/smart-search.ts` around lines 89 - 90, The code sets let
finalMode = data.mode but data.mode is optional; change initialization to
default to "compact" when undefined (e.g., finalMode = data.mode ?? "compact")
so the normal query path preserves the existing compact-search response
contract; update both occurrences where finalMode is assigned (the first
assignment around variable finalMode and the second occurrence near the later
branch at line 146) and use the same defaulting pattern so callers that omit
mode continue to receive "compact".
| const existingObsIds = new Set(compact.map((r) => r.obsId)); | ||
| const memories = await kv.list<Memory>(KV.memories); | ||
| const memoryMap = new Map(memories.map((m) => [m.id, m])); |
There was a problem hiding this comment.
Avoid a full memory scan on every graph search.
kv.list(KV.memories) makes graph-mode latency grow with total memory count even though this branch only needs the small set of memoryIds returned by mem::concept-graph-search. Fetch those IDs directly and in parallel instead of materializing the whole store on the request path.
♻️ Suggested change
- const memories = await kv.list<Memory>(KV.memories);
- const memoryMap = new Map(memories.map((m) => [m.id, m]));
+ const memories = await Promise.all(
+ graphResult.results.map(({ memoryId }) =>
+ kv.get<Memory>(KV.memories, memoryId).catch(() => null),
+ ),
+ );
+ const memoryMap = new Map(
+ memories
+ .filter((m): m is Memory => m !== null)
+ .map((m) => [m.id, m]),
+ );As per coding guidelines, "Use parallel operations where possible with Promise.all for independent kv writes/reads".
📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| const existingObsIds = new Set(compact.map((r) => r.obsId)); | |
| const memories = await kv.list<Memory>(KV.memories); | |
| const memoryMap = new Map(memories.map((m) => [m.id, m])); | |
| const existingObsIds = new Set(compact.map((r) => r.obsId)); | |
| const memories = await Promise.all( | |
| graphResult.results.map(({ memoryId }) => | |
| kv.get<Memory>(KV.memories, memoryId).catch(() => null), | |
| ), | |
| ); | |
| const memoryMap = new Map( | |
| memories | |
| .filter((m): m is Memory => m !== null) | |
| .map((m) => [m.id, m]), | |
| ); |
🤖 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 `@src/functions/smart-search.ts` around lines 109 - 111, The code currently
calls kv.list(KV.memories) to build memoryMap which scans all memories; instead,
retrieve only the memory IDs returned by the graph search (the compact array /
mem::concept-graph-search results) and fetch those Memory entries in parallel
with Promise.all using kv.get for each id, then build memoryMap from those
fetched results; keep existingObsIds and compact usage but replace
memories/memoryMap population with the targeted parallel gets to avoid
full-store scanning.
| case 'activity': if (!state.activity.loaded) await loadActivity(); break; | ||
| case 'profile': if (!state.profile.loaded) await loadProfile(); break; | ||
| case 'replay': if (!state.replay.loaded) await loadReplay(); break; | ||
| case 'concepts': if (!state.concepts.loaded) renderConcepts(); break; |
There was a problem hiding this comment.
Prevent duplicate Concepts graph initialization during in-flight loads.
Switching to the Concepts tab multiple times before the first request resolves can start multiple fetch/render pipelines (including multiple RAF loops and window listeners). Add an in-flight guard and clear it in finally.
Proposed fix
- concepts: { loaded: false },
+ concepts: { loaded: false, loading: false },
@@
function renderConcepts() {
var el = document.getElementById('view-concepts');
- if (state.concepts.loaded) return;
+ if (state.concepts.loaded || state.concepts.loading) return;
+ state.concepts.loading = true;
apiGet('concept-graph-viewer?limit=50').then(function(data) {
if (!data || !data.success) {
document.getElementById('concepts-sidebar').innerHTML = '<div style="color:var(--accent);font-size:12px;">Failed to load.</div>';
return;
}
state.concepts.loaded = true;
@@
- }).catch(function() {
+ }).catch(function() {
document.getElementById('concepts-sidebar').innerHTML = '<div style="color:var(--accent);font-size:12px;">Failed to load graph.</div>';
+ }).finally(function() {
+ state.concepts.loading = false;
});
}Also applies to: 2053-2157
🤖 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 `@src/viewer/index.html` at line 1104, Add an "in-flight" guard to prevent
starting multiple fetch/render pipelines for Concepts: in the code path where
you check state.concepts.loaded and call renderConcepts(), introduce a boolean
like state.concepts.loading (or conceptsInFlight) that is checked before calling
renderConcepts() and set to true immediately when starting the pipeline (e.g.,
at the top of renderConcepts or before initiating the fetch). Ensure this flag
is cleared in a finally block after the fetch/render completes or errors so
subsequent tab switches can retry; also ensure any RAF loops or window listeners
started by renderConcepts are only registered if loading was not already
in-flight and are cleaned up on completion.
This feature improves search recall by utilizing concept co-occurrence relationships:
Added 'mem:concept-edges' KV scope to store graph data.
Automatically extract and reinforce concept pairs on 'mem::remember'.
Added 'mem::concept-graph-search' with Depth-2 BFS, lazy decay, and Top-10 node cap.
Integrated graph mode into 'smart-search' with a 0.8 relevance penalty.
Registered new 'memory_graph_search' MCP tool.
Added first-run auto-backfill migration for existing memories.
close concept_edges table + BFS-depth-2 recall mode #345
Summary by CodeRabbit
New Features
Tests
Chores / Documentation