Fix chat retrieval payload guard

src/agent-harness/runtime.test.ts

@@ -98,6 +98,73 @@ describe("agent harness runtime", () => {
     ])
   })
 
+  it("returns only newly retrieved evidence in each retrieve tool result", async () => {
+    const query = vi
+      .fn<RetrievalCapability["query"]>()
+      .mockResolvedValueOnce(makeRetrievalResponse())
+      .mockResolvedValueOnce({
+        ...makeRetrievalResponse(),
+        query: "second query",
+        evidenceText: "Second evidence",
+        results: [
+          {
+            content: "Second retrieval evidence.",
+            chunkType: "text",
+            score: 0.8,
+            source: {
+              documentId: "doc_2",
+              sourceFileName: "second.pdf",
+              sectionPath: "Second",
+            },
+          },
+        ],
+      })
+    const state: {
+      intent?: IntentFrame
+      contextPolicy?: ContextPolicy
+      toolCalls?: HarnessToolCallTrace[]
+    } = {
+      intent: {
+        task: "answer",
+        dependsOnPreviousTurn: false,
+        retrievalNeeded: "yes",
+        targetModalities: ["text"],
+        constraints: {},
+        groundingPolicy: "must_use_sources",
+      },
+      contextPolicy: {
+        carryHistory: "none",
+        reason: "Self-contained request.",
+        activePriorTurnIds: [],
+      },
+    }
+    const ledger = createEvidenceLedger()
+    const tools = createHarnessTools({
+      state,
+      ledger,
+      retrieval: { query },
+      recentTurns: [],
+    })
+
+    const firstResult = await executeTool(tools.retrieve, { query: "first" })
+    const secondResult = await executeTool(tools.retrieve, { query: "second" })
+
+    expect(firstResult).toMatchObject({
+      retrievalCount: 1,
+      chunks: [{ ref: "r1:result:1" }],
+    })
+    expect(secondResult).toMatchObject({
+      retrievalCount: 2,
+      evidenceText: "Second evidence",
+      chunks: [{ ref: "r2:result:1" }],
+    })
+    expect(JSON.stringify(secondResult)).not.toContain("r1:result:1")
+    expect(ledger.snapshot().chunks.map((chunk) => chunk.ref)).toEqual([
+      "r1:result:1",
+      "r2:result:1",
+    ])
+  })
+
   it("blocks finalize until intent and context policy are declared", async () => {
     const state: {
       intent?: IntentFrame

src/agent-harness/runtime.ts

@@ -353,6 +353,7 @@ export function createHarnessTools(input: {
               }
             }
 
+            const beforeSnapshot = input.ledger.snapshot()
             const response = await input.retrieval.query({
               query: request.query,
               modalities: request.modalities as TargetModality[],
@@ -363,21 +364,27 @@ export function createHarnessTools(input: {
               threshold: request.threshold,
             })
             const snapshot = input.ledger.addRetrievalResponse(response)
+            const currentChunks = snapshot.chunks.slice(
+              beforeSnapshot.chunks.length,
+            )
+            const currentAssets = snapshot.assets.slice(
+              beforeSnapshot.assets.length,
+            )
             return {
               ok: true,
               retrievalCount: snapshot.retrievalCount,
               evidenceText: response.evidenceText ?? "",
               stopReason: response.stopReason ?? null,
               failureReason: response.failureReason ?? null,
-              chunks: snapshot.chunks.map((chunk) => ({
+              chunks: currentChunks.map((chunk) => ({
                 ref: chunk.ref,
                 kind: chunk.kind,
                 type: chunk.chunkType,
                 preview: chunk.contentPreview,
                 source: chunk.source,
                 assetRef: chunk.assetRef,
               })),
-              assets: snapshot.assets.map((asset) => ({
+              assets: currentAssets.map((asset) => ({
                 ref: asset.ref,
                 type: asset.type,
                 label: asset.label,

src/domains/chat/index.test.ts

@@ -8,6 +8,7 @@ import {
   answerQuestionWithRetrieval,
   generateAgenticOutputManifest,
   parseChatRequestBody,
+  type SearchSources,
 } from "."
 import type {
   HardenableRetrievalResult,
@@ -174,6 +175,91 @@ describe("answerQuestionWithRetrieval", () => {
     });
   });
 
+  it("bounds merged retrieval evidence before passing it to the answer agent", async () => {
+    const defaultResults = Array.from({ length: 40 }, (_, index) =>
+      makeRetrievalResult({
+        content: `Default namespace result ${index + 1}`,
+        source: {
+          documentId: `doc_default_${index + 1}`,
+          sourceFileName: "default.pdf",
+          sectionPath: `Default ${index + 1}`,
+        },
+      }),
+    );
+    const workspaceResults = Array.from({ length: 40 }, (_, index) =>
+      makeRetrievalResult({
+        content: `Workspace result ${index + 1}`,
+        source: {
+          documentId: `doc_workspace_${index + 1}`,
+          sourceFileName: "workspace.pdf",
+          sectionPath: `Workspace ${index + 1}`,
+        },
+      }),
+    );
+    const referencedChunks = Array.from({ length: 40 }, (_, index) => ({
+      chunkId: `chunk_${index + 1}`,
+      documentId: `doc_reference_${index + 1}`,
+      chunkType: "text" as const,
+      sectionPath: `Reference ${index + 1}`,
+    }));
+    const retrieval = {
+      query: vi
+        .fn()
+        .mockResolvedValueOnce({
+          results: defaultResults,
+          evidenceText: "Default evidence",
+          referencedChunks,
+          namespace: "default",
+          query: "large response",
+          routerUsed: "workflow_single_step",
+          answerText: null,
+          stopReason: "answer_done",
+          failureReason: null,
+        })
+        .mockResolvedValueOnce({
+          results: workspaceResults,
+          evidenceText: "Workspace evidence",
+          referencedChunks,
+          namespace: "notebook-workspace",
+          query: "large response",
+          routerUsed: "workflow_single_step",
+          answerText: null,
+          stopReason: "answer_done",
+          failureReason: null,
+        }),
+    };
+    const generateAnswer = vi.fn(
+      async ({ searchSources }: { searchSources: SearchSources }) => {
+        const response = await searchSources({
+          query: "large response",
+          topK: 3,
+        });
+        expect(response.results).toHaveLength(6);
+        expect(response.referencedChunks).toHaveLength(6);
+        expect(response.results.map((result) => result.content)).toEqual(
+          [
+            ...defaultResults.slice(0, 3),
+            ...workspaceResults.slice(0, 3),
+          ].map((result) => result.content),
+        );
+        return makeHarnessRunResult("The answer is grounded.");
+      },
+    );
+
+    await Effect.runPromise(
+      answerQuestionWithRetrieval({
+        question: "What does the document say?",
+        namespace: "notebook-workspace",
+        namespaces: ["default", "notebook-workspace"],
+        sources: [makeSource()],
+        excludedSourceIds: [],
+        retrieval,
+        generateAnswer,
+        messages: [],
+      }),
+    );
+  });
+
   it("does not hide a failed namespace query behind an empty namespace result", async () => {
     const retrievalError = new Error("Legacy namespace query failed.");
     const retrieval = {
@@ -235,7 +321,14 @@ describe("answerQuestionWithRetrieval", () => {
     });
     const retrieval = {
       query: vi.fn().mockResolvedValue({
-        results: [result],
+        results: [
+          result,
+          ...Array.from({ length: 30 }, (_, index) =>
+            makeRetrievalResult({
+              content: `extra result ${index + 1}`,
+            }),
+          ),
+        ],
         evidenceText: `Evidence https://blob.example/evidence.jpg ${"evidence ".repeat(
           80,
         )}`,
@@ -251,6 +344,12 @@ describe("answerQuestionWithRetrieval", () => {
             jobId: "job_1",
             assetUrl: "https://blob.example/id.jpg",
           },
+          ...Array.from({ length: 30 }, (_, index) => ({
+            chunkId: `chunk_extra_${index + 1}`,
+            documentId: "doc_identity",
+            chunkType: "text" as const,
+            sectionPath: `Extra referenced chunk ${index + 1}`,
+          })),
         ],
         namespace: "notebook-workspace",
         query: "冯荣洲 身份证 ID card",
@@ -284,19 +383,13 @@ describe("answerQuestionWithRetrieval", () => {
     const response = meta.response as KnowhereQueryResponseLogMeta;
     expect(response).toMatchObject({
       query: "冯荣洲 身份证 ID card",
-      resultCount: 1,
-      referencedChunkCount: 1,
-      results: [
-        {
-          chunkType: "image",
-        },
-      ],
-      referencedChunks: [
-        {
-          chunkType: "image",
-        },
-      ],
+      resultCount: 31,
+      referencedChunkCount: 31,
     });
+    expect(response.results).toHaveLength(20);
+    expect(response.referencedChunks).toHaveLength(20);
+    expect(response.results[0]).toMatchObject({ chunkType: "image" });
+    expect(response.referencedChunks[0]).toMatchObject({ chunkType: "image" });
     expect(response.answerText.length).toBeLessThanOrEqual(203);
     expect(response.evidenceText.length).toBeLessThanOrEqual(203);
     expect(response.results[0]?.content.length).toBeLessThanOrEqual(103);

src/domains/chat/index.ts

@@ -43,9 +43,13 @@ import type { HardenableRetrievalResult } from "./media-asset-hardening"
 
 const DEFAULT_TOP_K = 8
 const MAX_AGENTIC_TOP_K = 12
+const MAX_AGENTIC_MERGED_RESULT_COUNT = 24
+const MAX_AGENTIC_MERGED_REFERENCED_CHUNK_COUNT = 24
+const MAX_AGENTIC_MERGED_TEXT_CHARS = 12_000
 const MAX_CITATION_RESULTS = 20
 const KNOWHERE_RESPONSE_TEXT_LOG_LIMIT = 200
 const KNOWHERE_CHUNK_LOG_LIMIT = 100
+const KNOWHERE_RESPONSE_LOG_ITEM_LIMIT = 20
 const NO_RESULTS_ANSWER = "I couldn't find that in your sources."
 const HARNESS_VALIDATION_FAILURE_ANSWER =
   "I couldn't safely finish that response because the agent output did not pass Notebook's validation checks. Please try again."
@@ -88,6 +92,13 @@ type KnowhereReferencedChunkLog = {
   readonly summary: string
 }
 
+type AgenticMergedEvidenceLimits = {
+  readonly resultCountPerResponse: number
+  readonly referencedChunkCountPerResponse: number
+  readonly resultCount: number
+  readonly referencedChunkCount: number
+}
+
 export type {
   AnswerQuestionInput,
   AnswerQuestionResult,
@@ -186,7 +197,14 @@ export const answerQuestionWithRetrieval = (
       ) {
         throw queryFailures[0]
       }
-      return mergeRetrievalResponses(queryResponses, retrievalPlan)
+      return mergeRetrievalResponses(
+        queryResponses,
+        retrievalPlan,
+        getAgenticMergedEvidenceLimits({
+          namespaceCount: queryResponses.length,
+          topK: queryInput.topK,
+        }),
+      )
     }
 
     const generatedAnswer = yield* Effect.tryPromise(() =>
@@ -593,10 +611,12 @@ function formatKnowhereQueryResponseForLog(
       response.evidenceText ?? "",
       KNOWHERE_RESPONSE_TEXT_LOG_LIMIT,
     ),
-    results: response.results.map(formatKnowhereResultChunkForLog),
-    referencedChunks: response.referencedChunks.map(
-      formatKnowhereReferencedChunkForLog,
-    ),
+    results: response.results
+      .slice(0, KNOWHERE_RESPONSE_LOG_ITEM_LIMIT)
+      .map(formatKnowhereResultChunkForLog),
+    referencedChunks: response.referencedChunks
+      .slice(0, KNOWHERE_RESPONSE_LOG_ITEM_LIMIT)
+      .map(formatKnowhereReferencedChunkForLog),
   }
 }
 
@@ -654,23 +674,35 @@ function getRetrievalNamespaces(input: AnswerQuestionInput): readonly string[] {
 function mergeRetrievalResponses(
   responses: readonly RetrievalQueryResponse[],
   retrievalPlan: AgenticRetrievalPlan,
+  evidenceLimits: AgenticMergedEvidenceLimits,
 ): AgenticRetrievalResponse {
   const [first] = responses
   if (!first) {
     throw new Error("No retrieval responses to merge.")
   }
 
   const statusResponses = getRetrievalStatusResponses(responses)
-  const results = responses.flatMap((response) => response.results)
-  const referencedChunks = responses.flatMap(
-    (response) => response.referencedChunks,
-  )
+  const results = responses
+    .flatMap((response) =>
+      response.results.slice(0, evidenceLimits.resultCountPerResponse),
+    )
+    .slice(0, evidenceLimits.resultCount)
+  const referencedChunks = responses
+    .flatMap((response) =>
+      response.referencedChunks.slice(
+        0,
+        evidenceLimits.referencedChunkCountPerResponse,
+      ),
+    )
+    .slice(0, evidenceLimits.referencedChunkCount)
   const evidenceTexts = responses
     .map((response) => response.evidenceText)
     .filter((value): value is string => Boolean(value))
+    .map(truncateAgenticModelText)
   const answerTexts = responses
     .map((response) => response.answerText)
     .filter((value): value is string => Boolean(value))
+    .map(truncateAgenticModelText)
 
   return {
     ...first,
@@ -693,6 +725,30 @@ function mergeRetrievalResponses(
   }
 }
 
+function getAgenticMergedEvidenceLimits(input: {
+  readonly namespaceCount: number
+  readonly topK: number | undefined
+}): AgenticMergedEvidenceLimits {
+  const namespaceCount = Math.max(input.namespaceCount, 1)
+  const perResponseCount = normalizeTopK(input.topK)
+  const requestedCount = perResponseCount * namespaceCount
+  return {
+    resultCountPerResponse: perResponseCount,
+    referencedChunkCountPerResponse: perResponseCount,
+    resultCount: Math.min(requestedCount, MAX_AGENTIC_MERGED_RESULT_COUNT),
+    referencedChunkCount: Math.min(
+      requestedCount,
+      MAX_AGENTIC_MERGED_REFERENCED_CHUNK_COUNT,
+    ),
+  }
+}
+
+function truncateAgenticModelText(value: string): string {
+  const trimmed = value.trim()
+  if (trimmed.length <= MAX_AGENTIC_MERGED_TEXT_CHARS) return trimmed
+  return `${trimmed.slice(0, MAX_AGENTIC_MERGED_TEXT_CHARS)}\n...[truncated]`
+}
+
 function getRetrievalStatusResponses(
   responses: readonly RetrievalQueryResponse[],
 ): readonly RetrievalQueryResponse[] {