From a8b66e6d11cc5b6994dc12c54b566b2b51a73055 Mon Sep 17 00:00:00 2001
From: suguanYang <wangbinqi77@gmail.com>
Date: Mon, 6 Jul 2026 04:03:45 +0800
Subject: [PATCH] Fix chat harness finalization budget

---
 src/agent-harness/runtime.test.ts | 83 +++++++++++++++++++++++++++++++
 src/agent-harness/runtime.ts      | 65 ++++++++++++++++++++++--
 2 files changed, 143 insertions(+), 5 deletions(-)

diff --git a/src/agent-harness/runtime.test.ts b/src/agent-harness/runtime.test.ts
index ef7e5e0..f6abf7d 100644
--- a/src/agent-harness/runtime.test.ts
+++ b/src/agent-harness/runtime.test.ts
@@ -5,6 +5,7 @@ import {
   buildHarnessMessages,
   buildHarnessSystemPrompt,
   createHarnessTools,
+  prepareHarnessStep,
   sanitizeHarnessModelMessagesForStep,
 } from "./runtime"
 import { createEvidenceLedger } from "./ledger"
@@ -393,6 +394,88 @@ describe("agent harness runtime", () => {
       },
     ])
   })
+
+  it("keeps normal steps unconstrained before the finalization step", () => {
+    const result = prepareHarnessStep({
+      stepNumber: 11,
+      messages: [
+        {
+          role: "user",
+          content: "Find the penalty amount.",
+        },
+      ],
+    })
+
+    expect(result).toEqual({
+      messages: [
+        {
+          role: "user",
+          content: "Find the penalty amount.",
+        },
+      ],
+    })
+  })
+
+  it("forces finalize at step 12 using existing tool results", () => {
+    const result = prepareHarnessStep({
+      stepNumber: 12,
+      messages: [
+        {
+          role: "tool",
+          content: [
+            {
+              type: "tool-result",
+              toolCallId: "call_1",
+              toolName: "retrieve",
+              output: {
+                type: "json",
+                value: {
+                  ok: true,
+                  chunks: [{ ref: "r1:result:1" }],
+                },
+              },
+              providerOptions: {
+                google: {
+                  thoughtSignature: "signature-1",
+                },
+              },
+            },
+          ],
+        },
+      ],
+    })
+
+    expect(result.activeTools).toEqual(["finalize"])
+    expect(result.toolChoice).toEqual({
+      type: "tool",
+      toolName: "finalize",
+    })
+    expect(result.messages).toEqual([
+      {
+        role: "tool",
+        content: [
+          {
+            type: "tool-result",
+            toolCallId: "call_1",
+            toolName: "retrieve",
+            output: {
+              type: "json",
+              value: {
+                ok: true,
+                chunks: [{ ref: "r1:result:1" }],
+              },
+            },
+          },
+        ],
+      },
+      {
+        role: "user",
+        content: expect.stringContaining(
+          "Use only the evidence and tool results already available",
+        ),
+      },
+    ])
+  })
 })
 
 function executeTool(tool: unknown, input: unknown): Promise<unknown> {
diff --git a/src/agent-harness/runtime.ts b/src/agent-harness/runtime.ts
index fabc627..b565918 100644
--- a/src/agent-harness/runtime.ts
+++ b/src/agent-harness/runtime.ts
@@ -1,4 +1,5 @@
 import {
+  hasToolCall,
   stepCountIs,
   ToolLoopAgent,
   tool,
@@ -22,8 +23,9 @@ import type {
 } from "./types"
 import { validateOutputManifest } from "./validator"
 
-const defaultMaxSteps = 10
+const defaultMaxSteps = 13
 const defaultMaxRevisions = 1
+const forcedFinalizationStepNumber = 12
 
 type ToolLoopAgentSettings = ConstructorParameters<typeof ToolLoopAgent>[0]
 
@@ -50,6 +52,17 @@ type HarnessToolState = {
   toolCalls?: HarnessToolCallTrace[]
 }
 
+type HarnessTools = ReturnType<typeof createHarnessTools>
+
+type HarnessStepPreparation = {
+  messages: ModelMessage[]
+  activeTools?: Array<Extract<keyof HarnessTools, string>>
+  toolChoice?: {
+    type: "tool"
+    toolName: Extract<keyof HarnessTools, string>
+  }
+}
+
 const targetModalitySchema = z.enum(["text", "image", "table"])
 
 const intentFrameSchema = z.object({
@@ -153,10 +166,15 @@ export async function runAgentHarness(
     model: input.model,
     instructions: buildHarnessSystemPrompt(input.turn),
     tools,
-    prepareStep: ({ messages: stepMessages }) => ({
-      messages: sanitizeHarnessModelMessagesForStep(stepMessages),
-    }),
-    stopWhen: stepCountIs(input.maxSteps ?? defaultMaxSteps),
+    prepareStep: ({ messages: stepMessages, stepNumber }) =>
+      prepareHarnessStep({
+        messages: stepMessages,
+        stepNumber,
+      }),
+    stopWhen: [
+      hasToolCall("finalize"),
+      stepCountIs(input.maxSteps ?? defaultMaxSteps),
+    ],
   })
 
   const maxRevisions = input.maxRevisions ?? defaultMaxRevisions
@@ -214,6 +232,32 @@ export async function runAgentHarness(
   }
 }
 
+export function prepareHarnessStep(input: {
+  readonly stepNumber: number
+  readonly messages: readonly ModelMessage[]
+}): HarnessStepPreparation {
+  const messages = sanitizeHarnessModelMessagesForStep(input.messages)
+
+  if (input.stepNumber < forcedFinalizationStepNumber) {
+    return { messages }
+  }
+
+  return {
+    messages: [
+      ...messages,
+      {
+        role: "user",
+        content: buildForcedFinalizationFeedback(),
+      },
+    ],
+    activeTools: ["finalize"],
+    toolChoice: {
+      type: "tool",
+      toolName: "finalize",
+    },
+  }
+}
+
 export function sanitizeHarnessModelMessagesForStep(
   messages: readonly ModelMessage[],
 ): ModelMessage[] {
@@ -284,6 +328,17 @@ function buildRevisionFeedback(errors: readonly string[]): string {
   ].join("\n")
 }
 
+function buildForcedFinalizationFeedback(): string {
+  return [
+    "The retrieval step budget has been reached.",
+    "Do not search again or call any evidence-reading tools.",
+    "Use only the evidence and tool results already available in this turn.",
+    "Call finalize now with the best supported answer.",
+    "If the existing evidence is insufficient, explain the gap in unresolved",
+    "instead of making unsupported claims.",
+  ].join("\n")
+}
+
 export function createHarnessTools(input: {
   readonly state: HarnessToolState
   readonly ledger: ReturnType<typeof createEvidenceLedger>