From 7242074ecc4027fc4e2abd26f3a514b0684d5dd6 Mon Sep 17 00:00:00 2001 From: Peter Steinberger Date: Mon, 17 Aug 2026 20:46:52 -0700 Subject: [PATCH] fix(memory): prevent silent context loss in resets and search (#125534) * fix(memory): record transcript capture failures * fix(memory): backfill filtered search results --- .../src/memory-tool-manager.test-mocks.ts | 12 +- extensions/memory-core/src/tools.test.ts | 89 +++++++++++++ extensions/memory-core/src/tools.ts | 124 ++++++++++++------ .../bundled/session-memory/handler.test.ts | 48 ++++++- src/hooks/bundled/session-memory/handler.ts | 71 ++++++---- 5 files changed, 272 insertions(+), 72 deletions(-) diff --git a/extensions/memory-core/src/memory-tool-manager.test-mocks.ts b/extensions/memory-core/src/memory-tool-manager.test-mocks.ts index 733b5f205e61..807416cdb4b8 100644 --- a/extensions/memory-core/src/memory-tool-manager.test-mocks.ts +++ b/extensions/memory-core/src/memory-tool-manager.test-mocks.ts @@ -33,6 +33,9 @@ type MemoryManagerParams = { let workspaceDir = "/workspace"; let statusDirty = false; let customStatus: Record | undefined; +let sourceCounts: Array<{ source: MemorySource; files: number; chunks: number }> = [ + { source: "memory", files: 1, chunks: 1 }, +]; let searchImpl: SearchImpl = async () => []; let closeImpl: () => Promise = async () => {}; let getManagerImpl: @@ -63,7 +66,7 @@ const stubManager = { model: "builtin", requestedProvider: "builtin", sources: ["memory" as const], - sourceCounts: [{ source: "memory" as const, files: 1, chunks: 1 }], + sourceCounts, custom: customStatus, }), sync: vi.fn(), @@ -96,6 +99,12 @@ export function setMemoryStatusDirty(next: boolean): void { statusDirty = next; } +export function setMemorySourceCounts( + next: Array<{ source: MemorySource; files: number; chunks: number }>, +): void { + sourceCounts = next; +} + export function setMemorySearchImpl(next: SearchImpl): void { searchImpl = next; } @@ -127,6 +136,7 @@ export function resetMemoryToolMockState(overrides?: { workspaceDir = "/workspace"; statusDirty = false; customStatus = undefined; + sourceCounts = [{ source: "memory", files: 1, chunks: 1 }]; getManagerImpl = undefined; searchImpl = overrides?.searchImpl ?? (async () => []); closeImpl = async () => {}; diff --git a/extensions/memory-core/src/tools.test.ts b/extensions/memory-core/src/tools.test.ts index cfe7bcc8fc48..2cdc08bc1f01 100644 --- a/extensions/memory-core/src/tools.test.ts +++ b/extensions/memory-core/src/tools.test.ts @@ -13,6 +13,7 @@ import { setMemoryCustomStatus, setMemorySearchImpl, setMemorySearchManagerImpl, + setMemorySourceCounts, setMemoryStatusDirty, } from "./memory-tool-manager.test-mocks.js"; import { applyProjectRanking } from "./memory/project-ranking.js"; @@ -1059,6 +1060,94 @@ describe("memory_search corpus labels", () => { ]); }); + it("widens ranked candidates to fill the visible session result window", async () => { + const ranked = [ + { + path: "sessions/missing-high-rank-a.jsonl", + startLine: 1, + endLine: 2, + score: 0.99, + snippet: "Invisible higher-ranked session", + source: "sessions" as const, + }, + { + path: "sessions/missing-high-rank-b.jsonl", + startLine: 3, + endLine: 4, + score: 0.98, + snippet: "Another invisible higher-ranked session", + source: "sessions" as const, + }, + { + path: "sessions/past-thread.jsonl", + startLine: 5, + endLine: 6, + score: 0.9, + snippet: "First visible session result", + source: "sessions" as const, + }, + { + path: "sessions/past-thread.jsonl", + startLine: 7, + endLine: 8, + score: 0.8, + snippet: "Second visible session result", + source: "sessions" as const, + }, + ]; + setMemorySearchImpl(async (opts) => { + return ranked.slice(0, opts?.maxResults); + }); + setMemorySourceCounts([{ source: "sessions", files: 3, chunks: 4 }]); + const tool = createMemorySearchToolOrThrow({ + config: { + agents: { list: [{ id: "main", default: true }] }, + memory: { + citations: "off", + search: { + sources: ["sessions"], + rememberAcrossConversations: true, + }, + }, + tools: { sessions: { visibility: "self" } }, + }, + agentSessionKey: "agent:main:main:active-memory:abcdef123456", + conversationRecall: { + anchorSessionKey: "agent:main:main", + scope: "same-agent-private", + corpus: "sessions", + }, + }); + + const result = await tool.execute("visible-backfill", { + query: "session result", + corpus: "memory", + maxResults: 2, + }); + const details = result.details as { + results: Array<{ path: string; snippet: string }>; + debug?: { + hits: number; + candidateHits: number; + withheldHits: number; + searchWindow: number; + }; + }; + + expect(details.results.map((entry) => entry.snippet)).toEqual([ + "First visible session result", + "Second visible session result", + ]); + expect(details.results).toHaveLength(2); + expect(details.results.every((entry) => entry.path.startsWith("sessions/"))).toBe(true); + expect(details.debug).toMatchObject({ + hits: 2, + candidateHits: 4, + withheldHits: 2, + searchWindow: 4, + }); + }); + it("preserves source corpus labels for memory and session transcript hits", async () => { setMemorySearchImpl(async () => [ { diff --git a/extensions/memory-core/src/tools.ts b/extensions/memory-core/src/tools.ts index bc3ffa76d79c..82a81ea3ce75 100644 --- a/extensions/memory-core/src/tools.ts +++ b/extensions/memory-core/src/tools.ts @@ -61,6 +61,7 @@ type MemoryManagerSearchOptions = NonNullable< >; const MEMORY_SEARCH_TOOL_COOLDOWN_MS = 60_000; +const MEMORY_SEARCH_POST_FILTER_MAX_CANDIDATES = 200; const memorySearchToolCooldowns = new Map(); @@ -562,6 +563,9 @@ export function createMemorySearchTool(options: { searchMs: number; embeddingBootstrap?: MemorySearchRuntimeDebug["embeddingBootstrap"]; hits: number; + candidateHits: number; + withheldHits: number; + searchWindow: number; } | undefined; if (shouldQueryMemory && memorySetup && memory && !("error" in memory)) { @@ -588,9 +592,10 @@ export function createMemorySearchTool(options: { : requestedCorpus == null || requestedCorpus === "all" ? effectiveSearchSources : undefined; - const createSearchOptions = (signal: AbortSignal) => + const resultLimit = maxResults ?? memorySearchConfig?.query.maxResults ?? 10; + const createSearchOptions = (signal: AbortSignal, candidateLimit: number) => ({ - maxResults, + maxResults: candidateLimit, minScore, sessionKey: options.agentSessionKey, activeProjectKeys: options.activeProjectKeys @@ -602,61 +607,95 @@ export function createMemorySearchTool(options: { }, ...(searchSources ? { sources: searchSources } : {}), }) satisfies MemoryManagerSearchOptions; - const searchActiveMemory = async (): Promise => + const searchActiveMemory = async ( + candidateLimit: number, + ): Promise => await runWithDefaultDeadline( async (signal) => - await activeMemory.manager.search(query, createSearchOptions(signal)), + await activeMemory.manager.search( + query, + createSearchOptions(signal, candidateLimit), + ), ); - managerMs = memory.debug?.managerMs; - try { - rawResults = await searchActiveMemory(); - } catch (error) { - if (!isClosedMemoryStoreError(error)) { - throw error; + const searchWithStoreRefresh = async ( + candidateLimit: number, + ): Promise => { + try { + return await searchActiveMemory(candidateLimit); + } catch (error) { + if (!isClosedMemoryStoreError(error)) { + throw error; + } + const refreshed = await runWithDefaultDeadline(async () => + trackMemoryManager( + await getMemoryManagerContextWithPurpose({ + cfg, + agentId, + purpose: memoryManagerPurpose, + acquireLocalService: options.acquireLocalService, + }), + ), + ); + if ("error" in refreshed) { + throw error; + } + managerMs = refreshed.debug?.managerMs; + activeMemory = refreshed; + return await searchActiveMemory(candidateLimit); } - const refreshed = await runWithDefaultDeadline(async () => - trackMemoryManager( - await getMemoryManagerContextWithPurpose({ + }; + const applyPostFilters = async ( + hits: MemorySearchResult[], + ): Promise => { + let filtered = await runWithDefaultDeadline( + async () => + await filterMemorySearchHitsBySessionVisibility({ cfg, agentId, - purpose: memoryManagerPurpose, - acquireLocalService: options.acquireLocalService, + requesterSessionKey: options.agentSessionKey, + sandboxed: options.sandboxed === true, + hits, + conversationRecall: options.conversationRecall, }), - ), ); - if ("error" in refreshed) { - throw error; + if (searchSources) { + const allowedSources = new Set(searchSources); + filtered = filtered.filter((hit) => allowedSources.has(hit.source)); } - managerMs = refreshed.debug?.managerMs; - activeMemory = refreshed; - rawResults = await searchActiveMemory(); - } + 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); if (pausedIndexIdentityReason) { return; } - rawResults = await runWithDefaultDeadline( - async () => - await filterMemorySearchHitsBySessionVisibility({ - cfg, - agentId, - requesterSessionKey: options.agentSessionKey, - sandboxed: options.sandboxed === true, - hits: rawResults, - conversationRecall: options.conversationRecall, - }), - ); - if (searchSources) { - const allowedSources = new Set(searchSources); - rawResults = rawResults.filter((hit) => allowedSources.has(hit.source)); - } - if (requestedCorpus === "sessions") { - rawResults = rawResults.filter((hit) => hit.source === "sessions"); - } else if (requestedCorpus === "memory") { - rawResults = rawResults.filter((hit) => hit.source === "memory"); + 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(); staleness = resolveMemorySearchStaleness(status, agentId) ?? undefined; const payloadResults = rawResults.map((result) => ({ @@ -694,6 +733,9 @@ export function createMemorySearchTool(options: { searchMs, embeddingBootstrap, hits: rawResults.length, + candidateHits: candidates.length, + withheldHits: Math.max(0, candidates.length - postFilterHits), + searchWindow, }; }); if (pausedIndexIdentityReason) { diff --git a/src/hooks/bundled/session-memory/handler.test.ts b/src/hooks/bundled/session-memory/handler.test.ts index dd417b0264dc..c9d12176cb3f 100644 --- a/src/hooks/bundled/session-memory/handler.test.ts +++ b/src/hooks/bundled/session-memory/handler.test.ts @@ -9,7 +9,11 @@ import { formatSqliteSessionFileMarker, parseSqliteSessionFileMarker, } from "../../../config/sessions/legacy-sqlite-marker.js"; -import { replaceTranscriptEvents } from "../../../config/sessions/session-accessor.js"; +import { + loadTranscriptEvents, + readSessionTranscriptBoundedMessageTailPage, + replaceTranscriptEvents, +} from "../../../config/sessions/session-accessor.js"; import { parseAgentSessionKey } from "../../../routing/session-key.js"; import { writeWorkspaceFile } from "../../../test-helpers/workspace.js"; import { withEnvAsync } from "../../../test-utils/env.js"; @@ -33,6 +37,18 @@ vi.mock("../../../logging/subsystem.js", () => ({ createSubsystemLogger: () => loggerMocks, })); +vi.mock("../../../config/sessions/session-accessor.js", async (importOriginal) => { + const actual = + await importOriginal(); + return { + ...actual, + loadTranscriptEvents: vi.fn(actual.loadTranscriptEvents), + readSessionTranscriptBoundedMessageTailPage: vi.fn( + actual.readSessionTranscriptBoundedMessageTailPage, + ), + }; +}); + async function readFileTranscript(filePath: string, messageCount = 15): Promise { try { const content = await fs.readFile(filePath, "utf8"); @@ -419,6 +435,36 @@ describe("session-memory hook", () => { expect(memoryContent).not.toContain("Inactive branch content"); }); + it("records and warns when transcript loading fails after reset capture", async () => { + const tempDir = await createCaseWorkspace("workspace"); + const sessionId = "unavailable-transcript"; + const sessionKey = "agent:main:main"; + const failure = new Error("transcript projection unavailable\nretry later"); + vi.mocked(readSessionTranscriptBoundedMessageTailPage).mockImplementationOnce(() => { + throw new Error("bounded capture unavailable"); + }); + vi.mocked(loadTranscriptEvents).mockRejectedValueOnce(failure); + loggerMocks.warn.mockClear(); + + const { memoryContent } = await runNewWithPreviousSessionEntry({ + tempDir, + sessionKey, + previousSessionEntry: { sessionId }, + }); + + expect(loggerMocks.warn).toHaveBeenCalledWith( + "Session transcript unavailable for memory capture", + { + sessionKey, + error: "transcript projection unavailable retry later", + }, + ); + expect(memoryContent).toContain("## Conversation Summary"); + expect(memoryContent).toContain( + '> Transcript content was unavailable: "transcript projection unavailable retry later"', + ); + }); + it("fills the configured memory window past ineligible tail messages", async () => { const tempDir = await createCaseWorkspace("workspace"); const storePath = path.join(tempDir, "sessions.json"); diff --git a/src/hooks/bundled/session-memory/handler.ts b/src/hooks/bundled/session-memory/handler.ts index ff3d0c30e04b..64d195d9c249 100644 --- a/src/hooks/bundled/session-memory/handler.ts +++ b/src/hooks/bundled/session-memory/handler.ts @@ -30,6 +30,7 @@ import { runWithGatewayIndependentRootWorkContinuation } from "../../../process/ import { parseAgentSessionKey, toAgentStoreSessionKey } from "../../../routing/session-key.js"; import { shortenHomePath } from "../../../utils.js"; import { resolveHookConfig } from "../../config.js"; +import { formatHookErrorForLog } from "../../fire-and-forget.js"; import type { HookHandler } from "../../hooks.js"; import { generateSlugViaLLM } from "../../llm-slug-generator.js"; import { isSessionAutoResetReason } from "../../session-auto-reset.js"; @@ -106,25 +107,21 @@ async function getRecentSqliteSessionContent( messageCount: number, capturedEvents?: TranscriptEvent[], ): Promise { - try { - const events = capturedEvents ?? (await loadTranscriptEvents({ ...scope })); - const latestResetIndex = capturedEvents - ? -1 - : events.findLastIndex( - (event) => - Boolean(event) && - typeof event === "object" && - !Array.isArray(event) && - (event as { type?: unknown }).type === "reset", - ); - const retiredEvents = latestResetIndex >= 0 ? events.slice(0, latestResetIndex) : events; - return getRecentSessionContentFromEvents( - selectVisibleTranscriptEvents(retiredEvents), - messageCount, - ); - } catch { - return null; - } + const events = capturedEvents ?? (await loadTranscriptEvents({ ...scope })); + const latestResetIndex = capturedEvents + ? -1 + : events.findLastIndex( + (event) => + Boolean(event) && + typeof event === "object" && + !Array.isArray(event) && + (event as { type?: unknown }).type === "reset", + ); + const retiredEvents = latestResetIndex >= 0 ? events.slice(0, latestResetIndex) : events; + return getRecentSessionContentFromEvents( + selectVisibleTranscriptEvents(retiredEvents), + messageCount, + ); } // The bounded reader already projects the active branch, but message pages @@ -278,19 +275,28 @@ async function saveSessionMemoryNow( let slug: string | null = null; let sessionContent: string | null = null; + let transcriptUnavailableReason: string | null = null; if (currentSessionId) { - sessionContent = await getRecentSqliteSessionContent( - { - agentId, - sessionId: currentSessionId, + try { + sessionContent = await getRecentSqliteSessionContent( + { + agentId, + sessionId: currentSessionId, + sessionKey: event.sessionKey, + storePath: + contextStorePath ?? resolveSessionStorePathCore(cfg?.session?.store, { agentId }), + }, + messageCount, + capturedEvents, + ); + } catch (error) { + transcriptUnavailableReason = formatHookErrorForLog(error); + log.warn("Session transcript unavailable for memory capture", { sessionKey: event.sessionKey, - storePath: - contextStorePath ?? resolveSessionStorePathCore(cfg?.session?.store, { agentId }), - }, - messageCount, - capturedEvents, - ); + error: transcriptUnavailableReason, + }); + } log.debug("Session content loaded", { length: sessionContent?.length ?? 0, messageCount, @@ -345,6 +351,13 @@ async function saveSessionMemoryNow( // Include conversation content if available if (sessionContent) { entryParts.push("## Conversation Summary", "", sessionContent, ""); + } else if (transcriptUnavailableReason) { + entryParts.push( + "## Conversation Summary", + "", + `> Transcript content was unavailable: ${JSON.stringify(transcriptUnavailableReason)}`, + "", + ); } const entry = entryParts.join("\n");