From 20b46bbb3b0f9eaf2d7d6d3aa7e36516377f4da5 Mon Sep 17 00:00:00 2001 From: Peter Steinberger Date: Tue, 18 Aug 2026 02:00:38 -0700 Subject: [PATCH] perf(memory): avoid duplicate embeddings during filtered recall (#125735) * perf(memory): search filtered sessions in one pass * refactor(memory): model transcript capture outcome --- config/max-lines-baseline.txt | 1 - .../src/memory-search-tool-query.ts | 188 +++++++++++++ .../src/memory/manager-index.test-support.ts | 12 +- .../src/tools.real-manager.test.ts | 108 ++++++++ extensions/memory-core/src/tools.test.ts | 37 ++- extensions/memory-core/src/tools.ts | 250 ++++-------------- src/hooks/bundled/session-memory/handler.ts | 56 ++-- 7 files changed, 428 insertions(+), 224 deletions(-) create mode 100644 extensions/memory-core/src/memory-search-tool-query.ts create mode 100644 extensions/memory-core/src/tools.real-manager.test.ts diff --git a/config/max-lines-baseline.txt b/config/max-lines-baseline.txt index 34fdb6c20358..3ccd0f7e6728 100644 --- a/config/max-lines-baseline.txt +++ b/config/max-lines-baseline.txt @@ -150,7 +150,6 @@ extensions/memory-core/src/memory/manager-search.ts extensions/memory-core/src/rem-evidence.ts extensions/memory-core/src/short-term-promotion.test.ts extensions/memory-core/src/tools.test.ts -extensions/memory-core/src/tools.ts extensions/memory-lancedb/index.test.ts extensions/memory-wiki/src/chatgpt-import.ts extensions/memory-wiki/src/cli.ts diff --git a/extensions/memory-core/src/memory-search-tool-query.ts b/extensions/memory-core/src/memory-search-tool-query.ts new file mode 100644 index 000000000000..91bb4b9a4e9c --- /dev/null +++ b/extensions/memory-core/src/memory-search-tool-query.ts @@ -0,0 +1,188 @@ +// Memory Core plugin module owns ranked search-window filtering and diagnostics. +import { formatErrorMessage } from "openclaw/plugin-sdk/error-runtime"; +import type { + MemorySearchManager, + MemorySearchRuntimeDebug, + MemorySource, +} from "openclaw/plugin-sdk/memory-core-host-engine-storage"; +import type { OpenClawConfig } from "openclaw/plugin-sdk/memory-core-host-runtime-core"; +import type { OpenClawPluginToolContext } from "openclaw/plugin-sdk/plugin-entry"; +import { asNullableRecord } from "openclaw/plugin-sdk/string-coerce-runtime"; +import { filterMemorySearchHitsBySessionVisibility } from "./session-search-visibility.js"; +import { buildMemorySearchUnavailableResult } from "./tools.shared.js"; + +const MEMORY_SEARCH_POST_FILTER_MAX_CANDIDATES = 200; +const PAUSED_MEMORY_INDEX_WARNING = + "Tell the user: memory search is paused because the memory index was built with a different embedding provider/model/settings."; +const PAUSED_MEMORY_INDEX_ACTION = + "Tell the user to run: openclaw memory status --index or openclaw memory index --force."; + +export function buildPausedMemoryIndexUnavailableResult(reason: string) { + return buildMemorySearchUnavailableResult(reason, { + warning: PAUSED_MEMORY_INDEX_WARNING, + action: PAUSED_MEMORY_INDEX_ACTION, + }); +} + +type ManagerState = { manager: MemorySearchManager; managerMs?: number }; + +type MemorySearchToolQuery = { + text: string; + resultLimit: number; + minScore?: number; + explicitSources?: MemorySource[]; + defaultSources?: MemorySource[]; + indexedSources?: MemorySource[]; + requestedCorpus?: "memory" | "wiki" | "all" | "sessions"; + sessionKey?: string; + activeProjectKeys?: readonly string[]; + conversationRecall?: OpenClawPluginToolContext["conversationRecall"]; +}; + +type MemorySearchToolVisibility = { + cfg: OpenClawConfig; + agentId: string; + sandboxed: boolean; +}; + +function isClosedMemoryStoreError(error: unknown): boolean { + const message = formatErrorMessage(error).toLowerCase(); + return ( + message.includes("database is not open") || + message.includes("database connection is not open") || + message.includes("database handle is closed") || + message.includes("memory search manager is closed") + ); +} + +export async function executeMemorySearchToolQuery(params: { + initialManager: ManagerState; + refreshManager: () => Promise; + query: MemorySearchToolQuery; + visibility: MemorySearchToolVisibility; + runWithDeadline: (task: (signal: AbortSignal) => Promise) => Promise; +}) { + const startedAt = Date.now(); + const runtimeDebug: MemorySearchRuntimeDebug[] = []; + let active = params.initialManager; + const { query, runWithDeadline, visibility } = params; + // Product recall may index transcripts without adding them to ordinary model search. + // Explicit corpus selection is authorized by the tool owner before this point. + const searchSources = + query.explicitSources ?? + (query.requestedCorpus === "sessions" + ? query.defaultSources + : query.requestedCorpus == null || query.requestedCorpus === "all" + ? query.conversationRecall?.corpus === "configured" + ? query.indexedSources + : query.defaultSources + : undefined); + + const searchOnce = async () => { + const allowedSources = searchSources ? new Set(searchSources) : undefined; + const searchesSessions = searchSources?.includes("sessions") === true; + const indexedCandidateCount = searchesSessions + ? (active.manager.status().sourceCounts ?? []) + .filter((entry) => allowedSources?.has(entry.source)) + .reduce((total, entry) => total + entry.chunks, 0) + : query.resultLimit; + // A zero-count index can populate during first-search bootstrap. Reserve the + // full bounded window so that bootstrap cannot recreate post-filter starvation. + const availableCandidates = + indexedCandidateCount > 0 ? indexedCandidateCount : MEMORY_SEARCH_POST_FILTER_MAX_CANDIDATES; + const searchWindow = searchesSessions + ? Math.min(MEMORY_SEARCH_POST_FILTER_MAX_CANDIDATES, availableCandidates) + : query.resultLimit; + const candidates = await runWithDeadline( + async (signal) => + await active.manager.search(query.text, { + maxResults: searchWindow, + minScore: query.minScore, + sessionKey: query.sessionKey, + activeProjectKeys: query.activeProjectKeys ? [...query.activeProjectKeys] : undefined, + signal, + onDebug: (debug) => runtimeDebug.push(debug), + ...(searchSources ? { sources: searchSources } : {}), + }), + ); + return { candidates, searchWindow }; + }; + + let searched: Awaited>; + try { + searched = await searchOnce(); + } catch (error) { + if (!isClosedMemoryStoreError(error)) { + throw error; + } + const refreshed = await params.refreshManager(); + if (!refreshed) { + throw error; + } + active = refreshed; + searched = await searchOnce(); + } + + const status = active.manager.status(); + const indexIdentity = asNullableRecord(asNullableRecord(status.custom)?.indexIdentity); + const pausedIndexIdentityReason = + indexIdentity?.status === "mismatched" || indexIdentity?.status === "missing" + ? typeof indexIdentity.reason === "string" && indexIdentity.reason.trim() + ? indexIdentity.reason.trim() + : "memory index identity is missing or mismatched" + : undefined; + if (pausedIndexIdentityReason) { + return { + status, + rawResults: [], + pausedIndexIdentityReason, + searchMode: undefined, + debug: undefined, + }; + } + + let filtered = await runWithDeadline( + async () => + await filterMemorySearchHitsBySessionVisibility({ + cfg: visibility.cfg, + agentId: visibility.agentId, + requesterSessionKey: query.sessionKey, + sandboxed: visibility.sandboxed, + hits: searched.candidates, + conversationRecall: query.conversationRecall, + }), + ); + if (searchSources) { + const allowedSources = new Set(searchSources); + filtered = filtered.filter((hit) => allowedSources.has(hit.source)); + } + if (query.requestedCorpus === "sessions") { + filtered = filtered.filter((hit) => hit.source === "sessions"); + } else if (query.requestedCorpus === "memory") { + filtered = filtered.filter((hit) => hit.source === "memory"); + } + + const postFilterHits = filtered.length; + const rawResults = filtered.slice(0, query.resultLimit); + const latestDebug = runtimeDebug.at(-1); + return { + status, + rawResults, + pausedIndexIdentityReason: undefined, + searchMode: latestDebug?.effectiveMode, + debug: { + backend: status.backend, + configuredMode: latestDebug?.configuredMode, + effectiveMode: "n/a", + fallback: latestDebug?.fallback, + managerMs: active.managerMs, + searchMs: Math.max(0, Date.now() - startedAt), + embeddingBootstrap: runtimeDebug.findLast((entry) => entry.embeddingBootstrap) + ?.embeddingBootstrap, + hits: rawResults.length, + candidateHits: searched.candidates.length, + withheldHits: Math.max(0, searched.candidates.length - postFilterHits), + searchWindow: searched.searchWindow, + }, + }; +} diff --git a/extensions/memory-core/src/memory/manager-index.test-support.ts b/extensions/memory-core/src/memory/manager-index.test-support.ts index 603aced6e84d..c410b6474a39 100644 --- a/extensions/memory-core/src/memory/manager-index.test-support.ts +++ b/extensions/memory-core/src/memory/manager-index.test-support.ts @@ -59,6 +59,8 @@ type ProviderCall = { }; type ProviderControls = { + embedQueryCalls: number; + embeddedQueryTexts: string[]; embedBatchCalls: number; embeddedBatchTexts: string[]; embedBatchInputCalls: number; @@ -117,6 +119,8 @@ export type ManagerIndexFixture = { }; const providerState = vi.hoisted(() => ({ + embedQueryCalls: 0, + embeddedQueryTexts: [] as string[], embedBatchCalls: 0, embeddedBatchTexts: [] as string[], embedBatchInputCalls: 0, @@ -257,7 +261,11 @@ vi.mock("./embeddings.js", async (importOriginal) => { throw providerState.providerCloseFailure; } }, - embedQuery: async (text: string) => embedText(text), + embedQuery: async (text: string) => { + providerState.embedQueryCalls += 1; + providerState.embeddedQueryTexts.push(text); + return embedText(text); + }, embedBatch: async (texts: string[]) => { providerState.embedBatchCalls += 1; providerState.embeddedBatchTexts.push(...texts); @@ -548,6 +556,8 @@ export function createManagerIndexFixture(deps: { beforeEach(async () => { vi.useRealTimers(); clearRegistry(); + providerState.embedQueryCalls = 0; + providerState.embeddedQueryTexts = []; providerState.embedBatchCalls = 0; providerState.embeddedBatchTexts = []; providerState.embedBatchInputCalls = 0; diff --git a/extensions/memory-core/src/tools.real-manager.test.ts b/extensions/memory-core/src/tools.real-manager.test.ts new file mode 100644 index 000000000000..2baac3c4d9fb --- /dev/null +++ b/extensions/memory-core/src/tools.real-manager.test.ts @@ -0,0 +1,108 @@ +// Memory Core integration tests exercise the real SQLite search manager through tools. +import type { OpenClawConfig } from "openclaw/plugin-sdk/memory-core-host-runtime-core"; +import { describe, expect, it } from "vitest"; +import { + createManagerIndexFixture, + type ManagerIndexFixture, +} from "./memory/manager-index.test-support.js"; +import { createMemorySearchTool } from "./tools.js"; + +const { closeAllMemorySearchManagers, getMemorySearchManager } = await import("./memory/index.js"); + +describe("memory_search real manager", () => { + const fixture: ManagerIndexFixture = createManagerIndexFixture({ + getMemorySearchManager, + closeAllMemorySearchManagers, + }); + + it("backfills visible sessions with one bounded query embedding", async () => { + const baseConfig = fixture.createConfig({ + sources: ["sessions"], + sessionMemory: true, + minScore: 0, + vectorEnabled: false, + }); + const cfg = { + ...baseConfig, + memory: { ...baseConfig.memory, citations: "off" }, + tools: { ...baseConfig.tools, sessions: { visibility: "self" } }, + } satisfies OpenClawConfig; + const anchorSessionKey = "agent:main:telegram:direct:owner"; + + await fixture.seedSessionTranscript({ + sessionId: "current", + sessionKey: anchorSessionKey, + messages: [], + }); + for (const [sessionId, sessionKey, content] of [ + ["hidden-a", "agent:main:discord:group:hidden-a", "alpha alpha alpha hidden group a"], + ["hidden-b", "agent:main:discord:group:hidden-b", "alpha alpha alpha hidden group b"], + ["visible-a", "agent:main:telegram:direct:visible-a", "alpha beta visible private a"], + ["visible-b", "agent:main:telegram:direct:visible-b", "alpha beta visible private b"], + ] as const) { + await fixture.seedSessionTranscript({ + sessionId, + sessionKey, + messages: [{ role: "assistant", content, timestamp: "2026-08-17T12:00:00.000Z" }], + }); + } + + const manager = await fixture.getFreshManager(cfg); + await manager.sync({ reason: "test", force: true }); + expect(manager.status().sourceCounts).toEqual([{ source: "sessions", files: 5, chunks: 4 }]); + + const ranked = await manager.search("alpha", { + maxResults: 4, + minScore: 0, + sources: ["sessions"], + }); + expect(ranked.slice(0, 2).map((hit) => hit.path)).toEqual([ + expect.stringContaining("hidden-a"), + expect.stringContaining("hidden-b"), + ]); + fixture.provider.embedQueryCalls = 0; + fixture.provider.embeddedQueryTexts = []; + + const tool = createMemorySearchTool({ + config: cfg, + agentId: "main", + agentSessionKey: `${anchorSessionKey}:active-memory:abcdef123456`, + conversationRecall: { + anchorSessionKey, + scope: "same-agent-private", + corpus: "sessions", + }, + }); + if (!tool) { + throw new Error("memory_search tool missing"); + } + + const result = await tool.execute("real-manager-visible-backfill", { + query: "alpha", + corpus: "sessions", + maxResults: 2, + }); + const details = result.details as { + results: Array<{ snippet: string }>; + debug?: { + hits: number; + candidateHits: number; + withheldHits: number; + searchWindow: number; + }; + }; + + expect(details.results.map((hit) => hit.snippet)).toEqual([ + expect.stringContaining("visible private a"), + expect.stringContaining("visible private b"), + ]); + expect(fixture.provider.embedQueryCalls).toBe(1); + expect(fixture.provider.embeddedQueryTexts).toEqual(["alpha"]); + expect(details.debug).toMatchObject({ + hits: 2, + candidateHits: 4, + withheldHits: 2, + searchWindow: 4, + }); + }); +}); diff --git a/extensions/memory-core/src/tools.test.ts b/extensions/memory-core/src/tools.test.ts index 2cdc08bc1f01..f50e86b8d703 100644 --- a/extensions/memory-core/src/tools.test.ts +++ b/extensions/memory-core/src/tools.test.ts @@ -776,8 +776,10 @@ describe("memory_search corpus labels", () => { it("keeps ordinary memory_search on explicitly configured sources when recall indexing is enabled", async () => { let seenSources: readonly string[] | undefined; + let seenMaxResults: number | undefined; setMemorySearchImpl(async (opts) => { seenSources = opts?.sources; + seenMaxResults = opts?.maxResults; return []; }); const tool = createMemorySearchToolOrThrow({ @@ -795,9 +797,10 @@ describe("memory_search corpus labels", () => { agentSessionKey: "agent:main:main", }); - await tool.execute("ordinary-search", { query: "favorite food" }); + await tool.execute("ordinary-search", { query: "favorite food", maxResults: 3 }); expect(seenSources).toEqual(["memory"]); + expect(seenMaxResults).toBe(3); }); it("applies active-project ranking through the production memory_search tool", async () => { @@ -1061,6 +1064,7 @@ describe("memory_search corpus labels", () => { }); it("widens ranked candidates to fill the visible session result window", async () => { + const searchedLimits: Array = []; const ranked = [ { path: "sessions/missing-high-rank-a.jsonl", @@ -1096,6 +1100,7 @@ describe("memory_search corpus labels", () => { }, ]; setMemorySearchImpl(async (opts) => { + searchedLimits.push(opts?.maxResults); return ranked.slice(0, opts?.maxResults); }); setMemorySourceCounts([{ source: "sessions", files: 3, chunks: 4 }]); @@ -1140,12 +1145,42 @@ describe("memory_search corpus labels", () => { ]); expect(details.results).toHaveLength(2); expect(details.results.every((entry) => entry.path.startsWith("sessions/"))).toBe(true); + expect(searchedLimits).toEqual([4]); expect(details.debug).toMatchObject({ hits: 2, candidateHits: 4, withheldHits: 2, searchWindow: 4, }); + + searchedLimits.length = 0; + const boundedResult = await tool.execute("indexed-candidate-bound", { + query: "session result", + maxResults: 5, + }); + expect(searchedLimits).toEqual([4]); + expect(boundedResult.details).toMatchObject({ + results: expect.arrayContaining([ + expect.objectContaining({ snippet: "First visible session result" }), + expect.objectContaining({ snippet: "Second visible session result" }), + ]), + debug: { hits: 2, candidateHits: 4, withheldHits: 2, searchWindow: 4 }, + }); + + searchedLimits.length = 0; + setMemorySourceCounts([]); + const bootstrapResult = await tool.execute("bootstrap-candidate-window", { + query: "session result", + maxResults: 2, + }); + expect(searchedLimits).toEqual([200]); + expect(bootstrapResult.details).toMatchObject({ + results: expect.arrayContaining([ + expect.objectContaining({ snippet: "First visible session result" }), + expect.objectContaining({ snippet: "Second visible session result" }), + ]), + debug: { hits: 2, candidateHits: 4, withheldHits: 2, searchWindow: 200 }, + }); }); it("preserves source corpus labels for memory and session transcript hits", async () => { diff --git a/extensions/memory-core/src/tools.ts b/extensions/memory-core/src/tools.ts index 82a81ea3ce75..957a40a0b7f1 100644 --- a/extensions/memory-core/src/tools.ts +++ b/extensions/memory-core/src/tools.ts @@ -17,23 +17,22 @@ import { type MemoryCorpusSearchResult, type OpenClawConfig, } from "openclaw/plugin-sdk/memory-core-host-runtime-core"; -import type { - MemorySearchResult, - MemorySearchRuntimeDebug, -} from "openclaw/plugin-sdk/memory-core-host-runtime-files"; +import type { MemorySearchResult } from "openclaw/plugin-sdk/memory-core-host-runtime-files"; import { resolveMemoryDreamingConfig, resolveMemoryDeepDreamingConfig, } from "openclaw/plugin-sdk/memory-core-host-status"; import type { OpenClawPluginToolContext } from "openclaw/plugin-sdk/plugin-entry"; -import { asNullableRecord } from "openclaw/plugin-sdk/string-coerce-runtime"; +import { + buildPausedMemoryIndexUnavailableResult, + executeMemorySearchToolQuery, +} from "./memory-search-tool-query.js"; import type { MemoryCoreAcquireLocalService } from "./memory/embedding-local-service.js"; import { DEFAULT_MEMORY_SEARCH_TIMEOUT_MS, resolveMemorySearchAbortError, runMemorySearchWithDeadline, } from "./memory/search-deadline.js"; -import { filterMemorySearchHitsBySessionVisibility } from "./session-search-visibility.js"; import { recordShortTermRecalls } from "./short-term-promotion.js"; import { decorateCitations, @@ -56,12 +55,11 @@ type MemorySearchToolResult = | MemoryCorpusSearchResult; type MemoryManagerContext = Awaited>; type ActiveMemoryManagerContext = Extract; -type MemoryManagerSearchOptions = NonNullable< - Parameters[1] +type MemorySearchToolQueryDebug = NonNullable< + Awaited>["debug"] >; const MEMORY_SEARCH_TOOL_COOLDOWN_MS = 60_000; -const MEMORY_SEARCH_POST_FILTER_MAX_CANDIDATES = 200; const memorySearchToolCooldowns = new Map(); @@ -85,18 +83,6 @@ function readCorpusParam( throw new Error(`corpus must be one of: ${allowed.join(", ")}`); } -function mergeEmbeddingBootstrapRuntimeDebug( - entries: readonly MemorySearchRuntimeDebug[], -): MemorySearchRuntimeDebug["embeddingBootstrap"] | undefined { - let merged: MemorySearchRuntimeDebug["embeddingBootstrap"]; - for (const entry of entries) { - if (entry.embeddingBootstrap) { - merged = entry.embeddingBootstrap; - } - } - return merged; -} - function resolveMemorySearchToolCooldownKey(options: { agentId?: string; agentSessionKey?: string; @@ -156,28 +142,6 @@ async function closeMemoryManagers( } } -const PAUSED_MEMORY_INDEX_WARNING = - "Tell the user: memory search is paused because the memory index was built with a different embedding provider/model/settings."; -const PAUSED_MEMORY_INDEX_ACTION = - "Tell the user to run: openclaw memory status --index or openclaw memory index --force."; - -function resolvePausedMemoryIndexIdentityReason(status: { custom?: unknown }): string | undefined { - const indexIdentity = asNullableRecord(asNullableRecord(status.custom)?.indexIdentity); - if (indexIdentity?.status !== "mismatched" && indexIdentity?.status !== "missing") { - return undefined; - } - return typeof indexIdentity.reason === "string" && indexIdentity.reason.trim() - ? indexIdentity.reason.trim() - : "memory index identity is missing or mismatched"; -} - -function buildPausedMemoryIndexUnavailableResult(reason: string) { - return buildMemorySearchUnavailableResult(reason, { - warning: PAUSED_MEMORY_INDEX_WARNING, - action: PAUSED_MEMORY_INDEX_ACTION, - }); -} - function mergeRankedMemorySearchToolStreams( memoryResults: MemorySearchToolResult[], supplementResults: MemorySearchToolResult[], @@ -243,16 +207,6 @@ function mergeMemorySearchCorpusResults(params: { ).slice(0, params.maxResults); } -function isClosedMemoryStoreError(error: unknown): boolean { - const message = formatErrorMessage(error).toLowerCase(); - return ( - message.includes("database is not open") || - message.includes("database connection is not open") || - message.includes("database handle is closed") || - message.includes("memory search manager is closed") - ); -} - function buildRecallKey( result: Pick, ): string { @@ -539,93 +493,35 @@ export function createMemorySearchTool(options: { pluginConfig, cfg, }); - const searchStartedAt = Date.now(); let rawResults: MemorySearchResult[] = []; let surfacedMemoryResults: Array = []; - let provider: string | undefined; - let model: string | undefined; + let provider: string | undefined, model: string | undefined; let fallback: unknown; - let searchMode: string | undefined; - let pausedIndexIdentityReason: string | undefined; + let searchMode: string | undefined, pausedIndexIdentityReason: string | undefined; let staleness: | Exclude, null> | undefined; - let managerMs: number | undefined; let searchDebug: - | { - backend: string; - configuredMode?: string; - effectiveMode?: string; - fallback?: string; - toolMs?: number; - managerMs?: number; - outsideSearchMs?: number; - searchMs: number; - embeddingBootstrap?: MemorySearchRuntimeDebug["embeddingBootstrap"]; - hits: number; - candidateHits: number; - withheldHits: number; - searchWindow: number; - } + | (MemorySearchToolQueryDebug & { toolMs?: number; outsideSearchMs?: number }) | undefined; if (shouldQueryMemory && memorySetup && memory && !("error" in memory)) { await runUnavailablePhase("memory", async () => { - let activeMemory = memory; - const runtimeDebug: MemorySearchRuntimeDebug[] = []; const memorySearchConfig = resolveMemorySearchConfig(cfg, agentId); const defaultSearchSources = memorySearchConfig?.searchSources; - const trustedConfiguredRecall = options.conversationRecall?.corpus === "configured"; - const effectiveSearchSources = trustedConfiguredRecall - ? memorySearchConfig?.sources - : defaultSearchSources; - const trustedTranscriptRecall = options.conversationRecall !== undefined; - const configuredSessionSearch = defaultSearchSources?.includes("sessions") === true; - // Product recall may index transcripts without adding them to ordinary model search. - // Only trusted recall or explicit configuration may search those indexed transcripts. - const searchSources: MemorySource[] | undefined = - requestedCorpus === "sessions" - ? trustedTranscriptRecall || configuredSessionSearch - ? (["sessions"] as MemorySource[]) - : defaultSearchSources + const explicitSearchSources: MemorySource[] | undefined = + requestedCorpus === "sessions" && + (options.conversationRecall || defaultSearchSources?.includes("sessions")) + ? (["sessions"] as MemorySource[]) : requestedCorpus === "memory" ? (["memory"] as MemorySource[]) - : requestedCorpus == null || requestedCorpus === "all" - ? effectiveSearchSources - : undefined; + : undefined; const resultLimit = maxResults ?? memorySearchConfig?.query.maxResults ?? 10; - const createSearchOptions = (signal: AbortSignal, candidateLimit: number) => - ({ - maxResults: candidateLimit, - minScore, - sessionKey: options.agentSessionKey, - activeProjectKeys: options.activeProjectKeys - ? [...options.activeProjectKeys] - : undefined, - signal, - onDebug: (debug: MemorySearchRuntimeDebug) => { - runtimeDebug.push(debug); - }, - ...(searchSources ? { sources: searchSources } : {}), - }) satisfies MemoryManagerSearchOptions; - const searchActiveMemory = async ( - candidateLimit: number, - ): Promise => - await runWithDefaultDeadline( - async (signal) => - await activeMemory.manager.search( - query, - createSearchOptions(signal, candidateLimit), - ), - ); - const searchWithStoreRefresh = async ( - candidateLimit: number, - ): Promise => { - try { - return await searchActiveMemory(candidateLimit); - } catch (error) { - if (!isClosedMemoryStoreError(error)) { - throw error; - } + const executed = await executeMemorySearchToolQuery({ + initialManager: { + manager: memory.manager, + managerMs: memory.debug?.managerMs, + }, + refreshManager: async () => { const refreshed = await runWithDefaultDeadline(async () => trackMemoryManager( await getMemoryManagerContextWithPurpose({ @@ -637,66 +533,38 @@ export function createMemorySearchTool(options: { ), ); if ("error" in refreshed) { - throw error; + return null; } - managerMs = refreshed.debug?.managerMs; - activeMemory = refreshed; - return await searchActiveMemory(candidateLimit); - } - }; - const applyPostFilters = async ( - hits: MemorySearchResult[], - ): Promise => { - let filtered = await runWithDefaultDeadline( - async () => - await filterMemorySearchHitsBySessionVisibility({ - cfg, - agentId, - requesterSessionKey: options.agentSessionKey, - sandboxed: options.sandboxed === true, - hits, - conversationRecall: options.conversationRecall, - }), - ); - if (searchSources) { - const allowedSources = new Set(searchSources); - filtered = filtered.filter((hit) => allowedSources.has(hit.source)); - } - if (requestedCorpus === "sessions") { - filtered = filtered.filter((hit) => hit.source === "sessions"); - } else if (requestedCorpus === "memory") { - filtered = filtered.filter((hit) => hit.source === "memory"); - } - return filtered; - }; - managerMs = memory.debug?.managerMs; - let searchWindow = resultLimit; - let candidates = await searchWithStoreRefresh(searchWindow); - const statusBeforeRetry = activeMemory.manager.status(); - pausedIndexIdentityReason = - resolvePausedMemoryIndexIdentityReason(statusBeforeRetry); + return { + manager: refreshed.manager, + managerMs: refreshed.debug?.managerMs, + }; + }, + query: { + text: query, + resultLimit, + minScore, + explicitSources: explicitSearchSources, + defaultSources: defaultSearchSources, + indexedSources: memorySearchConfig?.sources, + requestedCorpus, + sessionKey: options.agentSessionKey, + activeProjectKeys: options.activeProjectKeys, + conversationRecall: options.conversationRecall, + }, + visibility: { + cfg, + agentId, + sandboxed: options.sandboxed === true, + }, + runWithDeadline: runWithDefaultDeadline, + }); + pausedIndexIdentityReason = executed.pausedIndexIdentityReason; if (pausedIndexIdentityReason) { return; } - rawResults = await applyPostFilters(candidates); - if (candidates.length === resultLimit && rawResults.length < resultLimit) { - const allowedSources = searchSources ? new Set(searchSources) : null; - const indexedCandidateCount = (statusBeforeRetry.sourceCounts ?? []) - .filter((entry) => !allowedSources || allowedSources.has(entry.source)) - .reduce((total, entry) => total + entry.chunks, 0); - const widenedLimit = Math.min( - MEMORY_SEARCH_POST_FILTER_MAX_CANDIDATES, - indexedCandidateCount, - ); - if (widenedLimit > searchWindow) { - searchWindow = widenedLimit; - candidates = await searchWithStoreRefresh(searchWindow); - rawResults = await applyPostFilters(candidates); - } - } - const postFilterHits = rawResults.length; - rawResults = rawResults.slice(0, resultLimit); - const status = activeMemory.manager.status(); + rawResults = executed.rawResults; + const status = executed.status; staleness = resolveMemorySearchStaleness(status, agentId) ?? undefined; const payloadResults = rawResults.map((result) => ({ ...result, @@ -720,23 +588,8 @@ export function createMemorySearchTool(options: { provider = status.provider; model = status.model; fallback = status.fallback; - const latestDebug = runtimeDebug.at(-1); - const embeddingBootstrap = mergeEmbeddingBootstrapRuntimeDebug(runtimeDebug); - searchMode = latestDebug?.effectiveMode; - const searchMs = Math.max(0, Date.now() - searchStartedAt); - searchDebug = { - backend: status.backend, - configuredMode: latestDebug?.configuredMode, - effectiveMode: "n/a", - fallback: latestDebug?.fallback, - managerMs, - searchMs, - embeddingBootstrap, - hits: rawResults.length, - candidateHits: candidates.length, - withheldHits: Math.max(0, candidates.length - postFilterHits), - searchWindow, - }; + searchMode = executed.searchMode; + searchDebug = executed.debug; }); if (pausedIndexIdentityReason) { return jsonResult( @@ -885,4 +738,3 @@ export function createMemoryGetTool(options: { }, }); } -/* oxlint-disable max-lines -- TODO: split this grandfathered oversized file. */ diff --git a/src/hooks/bundled/session-memory/handler.ts b/src/hooks/bundled/session-memory/handler.ts index 64d195d9c249..335c4cb28023 100644 --- a/src/hooks/bundled/session-memory/handler.ts +++ b/src/hooks/bundled/session-memory/handler.ts @@ -41,6 +41,10 @@ const SESSION_MEMORY_CAPTURE_MAX_BYTES = 8 * 1024 * 1024; const SESSION_MEMORY_CAPTURE_PAGE_MESSAGES = 256; const SESSION_MEMORY_CAPTURE_MAX_SCANNED_MESSAGES = 4_096; +type SessionMemoryTranscript = + | { status: "available"; content: string | null } + | { status: "unavailable"; reason: string }; + function pickDateTimePart( parts: Intl.DateTimeFormatPart[], type: Intl.DateTimeFormatPartTypes, @@ -274,31 +278,34 @@ async function saveSessionMemoryNow( : 15; let slug: string | null = null; - let sessionContent: string | null = null; - let transcriptUnavailableReason: string | null = null; + let transcript: SessionMemoryTranscript = { status: "available", content: null }; if (currentSessionId) { try { - sessionContent = await getRecentSqliteSessionContent( - { - agentId, - sessionId: currentSessionId, - sessionKey: event.sessionKey, - storePath: - contextStorePath ?? resolveSessionStorePathCore(cfg?.session?.store, { agentId }), - }, - messageCount, - capturedEvents, - ); + transcript = { + status: "available", + content: await getRecentSqliteSessionContent( + { + agentId, + sessionId: currentSessionId, + sessionKey: event.sessionKey, + storePath: + contextStorePath ?? resolveSessionStorePathCore(cfg?.session?.store, { agentId }), + }, + messageCount, + capturedEvents, + ), + }; } catch (error) { - transcriptUnavailableReason = formatHookErrorForLog(error); + const reason = formatHookErrorForLog(error); + transcript = { status: "unavailable", reason }; log.warn("Session transcript unavailable for memory capture", { sessionKey: event.sessionKey, - error: transcriptUnavailableReason, + error: reason, }); } log.debug("Session content loaded", { - length: sessionContent?.length ?? 0, + length: transcript.status === "available" ? (transcript.content?.length ?? 0) : 0, messageCount, }); @@ -306,11 +313,16 @@ async function saveSessionMemoryNow( const isTestEnv = isVitestRuntimeEnv(); const allowLlmSlug = !isTestEnv && hookConfig?.llmSlug === true; - if (sessionContent && cfg && allowLlmSlug) { + if (transcript.status === "available" && transcript.content && cfg && allowLlmSlug) { log.debug("Calling generateSlugViaLLM..."); // Use LLM to generate a descriptive slug const slugModel = typeof hookConfig?.model === "string" ? hookConfig.model : undefined; - slug = await generateSlugViaLLM({ sessionContent, cfg, agentId, model: slugModel }); + slug = await generateSlugViaLLM({ + sessionContent: transcript.content, + cfg, + agentId, + model: slugModel, + }); log.debug("Generated slug", { slug }); } } @@ -349,13 +361,13 @@ async function saveSessionMemoryNow( ]; // Include conversation content if available - if (sessionContent) { - entryParts.push("## Conversation Summary", "", sessionContent, ""); - } else if (transcriptUnavailableReason) { + if (transcript.status === "available" && transcript.content) { + entryParts.push("## Conversation Summary", "", transcript.content, ""); + } else if (transcript.status === "unavailable") { entryParts.push( "## Conversation Summary", "", - `> Transcript content was unavailable: ${JSON.stringify(transcriptUnavailableReason)}`, + `> Transcript content was unavailable: ${JSON.stringify(transcript.reason)}`, "", ); }