diff --git a/src/agents/pi-embedded-runner.sanitize-session-history.test.ts b/src/agents/pi-embedded-runner.sanitize-session-history.test.ts index 339df520458..008ad26dbbd 100644 --- a/src/agents/pi-embedded-runner.sanitize-session-history.test.ts +++ b/src/agents/pi-embedded-runner.sanitize-session-history.test.ts @@ -1522,6 +1522,59 @@ describe("sanitizeSessionHistory", () => { ]); }); + it.each([ + { + provider: "anthropic", + modelApi: "anthropic-messages", + label: "anthropic", + }, + { + provider: "amazon-bedrock", + modelApi: "bedrock-converse-stream", + label: "bedrock", + }, + ])( + "preserves active tool-turn thinking signatures for $label even when a tool result follows", + async ({ provider, modelApi }) => { + setNonGoogleModelApi(); + + const messages = castAgentMessages([ + makeUserMessage("look up the answer"), + makeAssistantMessage([ + { + type: "thinking", + thinking: "call the tool", + signature: "", + } as unknown as ThinkingContent, + { type: "toolCall", id: "call_1", name: "lookup", arguments: {} }, + ]), + castAgentMessage({ + role: "toolResult", + toolCallId: "call_1", + toolName: "lookup", + content: [{ type: "text", text: "42" }], + isError: false, + }), + ]); + + const result = await sanitizeAnthropicHistory({ + provider, + modelApi, + messages, + modelId: "claude-sonnet-4-6", + }); + + expect((result[1] as Extract).content).toEqual([ + { + type: "thinking", + thinking: "call the tool", + signature: "", + }, + { type: "toolCall", id: "call_1", name: "lookup", arguments: {} }, + ]); + }, + ); + it.each([ { provider: "anthropic", diff --git a/src/agents/pi-embedded-runner/replay-history.ts b/src/agents/pi-embedded-runner/replay-history.ts index 8db634257ca..42c7d1232cf 100644 --- a/src/agents/pi-embedded-runner/replay-history.ts +++ b/src/agents/pi-embedded-runner/replay-history.ts @@ -50,6 +50,7 @@ import { isZeroUsageEmptyStopAssistantTurn } from "./empty-assistant-turn.js"; import { dropReasoningFromHistory, dropThinkingBlocks, + shouldPreserveLatestAssistantThinking, stripInvalidThinkingSignatures, } from "./thinking.js"; @@ -727,12 +728,9 @@ export async function sanitizeSessionHistory(params: { ...resolveImageSanitizationLimits(params.config), }, ); - const lastMessage = sanitizedImages[sanitizedImages.length - 1]; const preserveLatestAssistantThinking = params.preserveLatestAssistantThinking ?? - (!!lastMessage && - typeof lastMessage === "object" && - (lastMessage as { role?: unknown }).role === "assistant"); + shouldPreserveLatestAssistantThinking(sanitizedImages); // Some recovery paths supply a narrow policy with preserveSignatures disabled. // Native signed-thinking providers still cannot replay missing/blank // signatures once the assistant turn is no longer latest in the outbound diff --git a/src/agents/pi-embedded-runner/thinking.ts b/src/agents/pi-embedded-runner/thinking.ts index 43cd0e5f3da..b54376427ef 100644 --- a/src/agents/pi-embedded-runner/thinking.ts +++ b/src/agents/pi-embedded-runner/thinking.ts @@ -261,6 +261,31 @@ function shouldPreserveCurrentToolTurnReasoning( return false; } +export function shouldPreserveLatestAssistantThinking(messages: AgentMessage[]): boolean { + let latestAssistantIndex = -1; + for (let index = messages.length - 1; index >= 0; index -= 1) { + if (isAssistantMessageWithContent(messages[index])) { + latestAssistantIndex = index; + break; + } + } + if (latestAssistantIndex < 0) { + return false; + } + if (latestAssistantIndex === messages.length - 1) { + return true; + } + + let latestUserIndex = -1; + for (let index = messages.length - 1; index >= 0; index -= 1) { + if ((messages[index] as { role?: unknown })?.role === "user") { + latestUserIndex = index; + break; + } + } + return shouldPreserveCurrentToolTurnReasoning(messages, latestAssistantIndex, latestUserIndex); +} + function stripAllThinkingBlocks(messages: AgentMessage[]): AgentMessage[] { let touched = false; const out: AgentMessage[] = [];