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

* fix(memory): record transcript capture failures * fix(memory): backfill filtered search results
2026-08-27 04:47:03 -06:00 · 2026-08-17 20:46:52 -07:00
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 () => {};
@@ -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 () => [
      {
@@ -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) {