fix(memory): prevent silent context loss in resets and search (#125534)

* fix(memory): record transcript capture failures

* fix(memory): backfill filtered search results
This commit is contained in:
Peter Steinberger
2026-08-17 20:46:52 -07:00
committed by GitHub
parent bc9c3a4ad8
commit 7242074ecc
5 changed files with 272 additions and 72 deletions
@@ -33,6 +33,9 @@ type MemoryManagerParams = {
let workspaceDir = "/workspace";
let statusDirty = false;
let customStatus: Record<string, unknown> | undefined;
let sourceCounts: Array<{ source: MemorySource; files: number; chunks: number }> = [
{ source: "memory", files: 1, chunks: 1 },
];
let searchImpl: SearchImpl = async () => [];
let closeImpl: () => Promise<void> = 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 () => {};
+89
View File
@@ -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 () => [
{
+83 -41
View File
@@ -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<string, { until: number; error: string }>();
@@ -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<MemorySearchResult[]> =>
const searchActiveMemory = async (
candidateLimit: number,
): Promise<MemorySearchResult[]> =>
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<MemorySearchResult[]> => {
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<MemorySearchResult[]> => {
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<MemorySource>(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<MemorySource>(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) {
@@ -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<typeof import("../../../config/sessions/session-accessor.js")>();
return {
...actual,
loadTranscriptEvents: vi.fn(actual.loadTranscriptEvents),
readSessionTranscriptBoundedMessageTailPage: vi.fn(
actual.readSessionTranscriptBoundedMessageTailPage,
),
};
});
async function readFileTranscript(filePath: string, messageCount = 15): Promise<string | null> {
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");
+42 -29
View File
@@ -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<string | null> {
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");