diff --git a/src/agents/embedded-agent-runner/run/attempt.spawn-workspace.context-engine.test.ts b/src/agents/embedded-agent-runner/run/attempt.spawn-workspace.context-engine.test.ts index fa0cd4171656..ea0bde0776d3 100644 --- a/src/agents/embedded-agent-runner/run/attempt.spawn-workspace.context-engine.test.ts +++ b/src/agents/embedded-agent-runner/run/attempt.spawn-workspace.context-engine.test.ts @@ -538,6 +538,193 @@ describe("runEmbeddedAttempt context engine sessionKey forwarding", () => { ]); }); + it("keeps newly generated thinking after repairing rejected Anthropic replay", async () => { + const { SessionManager: ActualSessionManager } = + await vi.importActual("../../sessions/index.js"); + const staleAssistant = { + role: "assistant", + content: [ + { + type: "thinking", + thinking: "historical stale thinking", + thinkingSignature: "stale-signature", + }, + { type: "text", text: "historical answer" }, + ], + stopReason: "stop", + api: "anthropic-messages", + provider: "anthropic", + model: "claude-sonnet-4-6", + timestamp: 2, + } as AgentMessage; + const sessionMessages = [ + { role: "user", content: "historical question", timestamp: 1 } as AgentMessage, + staleAssistant, + ]; + const sessionManager = ActualSessionManager.inMemory(); + const appendSessionMessage = (message: AgentMessage) => + sessionManager.appendMessage(message as Parameters[0]); + for (const message of sessionMessages) { + appendSessionMessage(message); + } + const retryAssistant = { + role: "assistant", + content: [ + { + type: "thinking", + thinking: "fresh valid retry thinking", + thinkingSignature: "fresh-valid-signature", + }, + { type: "text", text: "retry answer" }, + ], + stopReason: "stop", + api: "anthropic-messages", + provider: "anthropic", + model: "claude-sonnet-4-6", + timestamp: 4, + } as AgentMessage; + const providerContexts: AgentMessage[][] = []; + const afterTurn = vi.fn(async (_params: { messages: AgentMessage[] }) => {}); + + hoisted.sessionManagerOpenMock.mockReturnValue(sessionManager); + + await createContextEngineAttemptRunner({ + contextEngine: { + ...createContextEngineBootstrapAndAssemble(), + afterTurn, + }, + sessionKey, + tempPaths, + sessionMessages, + attemptOverrides: { + provider: "anthropic", + modelId: "claude-sonnet-4-6", + model: { + api: "anthropic-messages", + provider: "anthropic", + id: "claude-sonnet-4-6", + contextWindow: 128_000, + input: ["text"], + } as never, + runtimePlan: { + prompt: { + resolveSystemPromptContribution: () => undefined, + }, + transcript: { + resolvePolicy: () => ({ + sanitizeMode: "full", + sanitizeToolCallIds: true, + preserveNativeAnthropicToolUseIds: false, + repairToolUseResultPairing: true, + preserveSignatures: true, + sanitizeThinkingSignatures: false, + dropThinkingBlocks: false, + dropReasoningFromHistory: false, + applyGoogleTurnOrdering: false, + validateGeminiTurns: false, + validateAnthropicTurns: false, + allowSyntheticToolResults: false, + }), + }, + transport: { + extraParams: {}, + resolveExtraParams: () => ({}), + }, + tools: { + normalize: (tools: unknown[]) => tools, + logDiagnostics: () => {}, + }, + auth: { + providerForAuth: "anthropic", + authProfileProviderForAuth: "", + forwardedAuthProfileId: undefined, + }, + delivery: { + isSilentPayload: () => false, + resolveFollowupRoute: () => undefined, + }, + outcome: { + classifyRunResult: () => undefined, + }, + observability: { + resolvedRef: "anthropic/claude-sonnet-4-6", + provider: "anthropic", + modelId: "claude-sonnet-4-6", + modelApi: "anthropic-messages", + }, + } as never, + }, + createSession: () => { + const session = createDefaultEmbeddedSession({ initialMessages: sessionMessages }); + let streamCalls = 0; + session.agent.streamFn = async (_model, context) => { + streamCalls += 1; + providerContexts.push([ + ...((context as { messages?: AgentMessage[] } | undefined)?.messages ?? []), + ]); + if (streamCalls === 1) { + throw new Error("invalid signature in thinking block"); + } + return { + async result() { + return retryAssistant; + }, + [Symbol.asyncIterator]() { + return (async function* () {})(); + }, + }; + }; + session.prompt = async (prompt, options) => { + options?.preflightResult?.(true); + const userMessage = { + role: "user", + content: [{ type: "text", text: prompt }], + timestamp: 3, + } as AgentMessage; + session.messages = [...session.messages, userMessage]; + appendSessionMessage(userMessage); + const stream = await session.agent.streamFn?.( + {} as never, + { messages: session.messages } as never, + {}, + ); + const assistantMessage = await ( + stream as { result: () => Promise } + ).result(); + session.messages = [...session.messages, assistantMessage]; + appendSessionMessage(assistantMessage); + }; + return session; + }, + }); + + const firstProviderContext = providerContexts[0] ?? []; + const retryProviderContext = providerContexts[1] ?? []; + expect(JSON.stringify(firstProviderContext)).toContain("stale-signature"); + expect(JSON.stringify(retryProviderContext)).not.toContain("stale-signature"); + + const finalMessages = sessionManager.buildSessionContext().messages; + expect(JSON.stringify(finalMessages[1])).not.toContain("historical stale thinking"); + expect(JSON.stringify(finalMessages.at(-1))).toContain("fresh valid retry thinking"); + expect(JSON.stringify(finalMessages.at(-1))).toContain("fresh-valid-signature"); + expect(afterTurn).toHaveBeenCalledWith( + expect.objectContaining({ + messages: expect.arrayContaining([ + expect.objectContaining({ + role: "assistant", + content: expect.arrayContaining([ + expect.objectContaining({ + type: "thinking", + thinking: "fresh valid retry thinking", + thinkingSignature: "fresh-valid-signature", + }), + ]), + }), + ]), + }), + ); + }); + it("sends transcriptPrompt visibly and keeps runtime context out of transcript messages", async () => { const seen: { prompt?: string; messages?: unknown[]; systemPrompt?: string } = {}; diff --git a/src/agents/embedded-agent-runner/run/attempt.spawn-workspace.test-support.ts b/src/agents/embedded-agent-runner/run/attempt.spawn-workspace.test-support.ts index 88cfa4af3600..4bd3f3c50ec2 100644 --- a/src/agents/embedded-agent-runner/run/attempt.spawn-workspace.test-support.ts +++ b/src/agents/embedded-agent-runner/run/attempt.spawn-workspace.test-support.ts @@ -793,10 +793,14 @@ vi.mock("../sandbox-info.js", () => ({ resolveEmbeddedSandboxInfoExecPolicy: () => ({}), })); -vi.mock("../thinking.js", () => ({ - dropReasoningFromHistory: (messages: T) => messages, - dropThinkingBlocks: (messages: T) => messages, -})); +vi.mock("../thinking.js", async (importOriginal) => { + const actual = await importOriginal(); + return { + ...actual, + dropReasoningFromHistory: (messages: T) => messages, + dropThinkingBlocks: (messages: T) => messages, + }; +}); vi.mock("../tool-name-allowlist.js", async (importOriginal) => { const actual = await importOriginal(); diff --git a/src/agents/embedded-agent-runner/run/attempt.ts b/src/agents/embedded-agent-runner/run/attempt.ts index 48631185b3b5..34f2e8f02daa 100644 --- a/src/agents/embedded-agent-runner/run/attempt.ts +++ b/src/agents/embedded-agent-runner/run/attempt.ts @@ -277,6 +277,7 @@ import { resolveEmbeddedAgentStreamFn, } from "../stream-resolution.js"; import { applySystemPromptToSession } from "../system-prompt.js"; +import { repairRejectedThinkingReplayInSessionManager } from "../thinking-replay-repair.js"; import { dropReasoningFromHistory, dropThinkingBlocks, @@ -1990,6 +1991,7 @@ export async function runEmbeddedAttempt( let trajectoryEndRecorded = false; let buildAbortSettlePromise: () => Promise | null = () => null; let cleanupYieldAborted = false; + let repairedRejectedThinkingReplay = false; try { await repairSessionFileIfNeeded({ sessionFile: params.sessionFile, @@ -2744,6 +2746,30 @@ export async function runEmbeddedAttempt( activeSession.agent.streamFn, { id: activeSession.sessionId, + onRecoveredAnthropicThinking: () => { + if (!sessionManager) { + log.warn( + `[session-recovery] unable to repair rejected thinking replay: session manager unavailable sessionId=${activeSession.sessionId}`, + ); + return; + } + const repair = repairRejectedThinkingReplayInSessionManager({ + sessionManager, + sessionFile: params.sessionFile, + sessionId: params.sessionId, + sessionKey: params.sessionKey, + agentId: sessionAgentId, + }); + if (repair.repaired) { + repairedRejectedThinkingReplay = true; + sessionLockController.refreshAfterOwnedSessionWrite(); + return; + } + log.warn( + `[session-recovery] rejected thinking replay retry succeeded but transcript repair made no changes: ` + + `sessionId=${activeSession.sessionId} reason=${repair.reason ?? "unknown"}`, + ); + }, }, ); } @@ -4538,6 +4564,9 @@ export async function runEmbeddedAttempt( await sessionLockController.waitForSessionEvents(activeSession); await waitForPendingEvents(); + if (repairedRejectedThinkingReplay) { + activeSession.agent.state.messages = activeSessionManager.buildSessionContext().messages; + } await sessionLockController.releaseForPrompt(); if ( diff --git a/src/agents/embedded-agent-runner/thinking-replay-repair.test.ts b/src/agents/embedded-agent-runner/thinking-replay-repair.test.ts new file mode 100644 index 000000000000..9c64614c0185 --- /dev/null +++ b/src/agents/embedded-agent-runner/thinking-replay-repair.test.ts @@ -0,0 +1,139 @@ +// Provider thinking replay repair tests cover durable transcript cleanup after +// Anthropic/Bedrock proves a signed thinking block invalid. +import type { AgentMessage } from "openclaw/plugin-sdk/agent-core"; +import { SessionManager } from "openclaw/plugin-sdk/agent-sessions"; +import { describe, expect, it } from "vitest"; +import { repairRejectedThinkingReplayInSessionManager } from "./thinking-replay-repair.js"; + +type AppendMessage = Parameters[0]; + +function asAppendMessage(message: unknown): AppendMessage { + return message as AppendMessage; +} + +function branchMessages(sessionManager: SessionManager): AgentMessage[] { + return sessionManager + .getBranch() + .filter((entry) => entry.type === "message") + .map((entry) => entry.message); +} + +function branchAssistantContents(sessionManager: SessionManager): unknown[] { + return branchMessages(sessionManager) + .filter((message): message is Extract => { + return message.role === "assistant"; + }) + .map((message) => message.content); +} + +describe("repairRejectedThinkingReplayInSessionManager", () => { + it("strips thinking blocks from active-branch assistant messages and preserves visible content", () => { + const sessionManager = SessionManager.inMemory(); + sessionManager.appendMessage(asAppendMessage({ role: "user", content: "first", timestamp: 1 })); + sessionManager.appendMessage( + asAppendMessage({ + role: "assistant", + content: [ + { type: "thinking", thinking: "private", thinkingSignature: "sig_bad" }, + { type: "text", text: "visible answer" }, + ], + timestamp: 2, + }), + ); + sessionManager.appendMessage( + asAppendMessage({ role: "user", content: "second", timestamp: 3 }), + ); + + const result = repairRejectedThinkingReplayInSessionManager({ sessionManager }); + + expect(result).toMatchObject({ repaired: true, repairedCount: 1 }); + expect(branchMessages(sessionManager).map((message) => message.role)).toEqual([ + "user", + "assistant", + "user", + ]); + expect(branchAssistantContents(sessionManager)).toEqual([ + [{ type: "text", text: "visible answer" }], + ]); + }); + + it("keeps thinking-only assistant turns as omitted-reasoning placeholders", () => { + const sessionManager = SessionManager.inMemory(); + sessionManager.appendMessage(asAppendMessage({ role: "user", content: "first", timestamp: 1 })); + sessionManager.appendMessage( + asAppendMessage({ + role: "assistant", + content: [{ type: "thinking", thinking: "private", thinkingSignature: "sig_bad" }], + timestamp: 2, + }), + ); + + const result = repairRejectedThinkingReplayInSessionManager({ sessionManager }); + + expect(result).toMatchObject({ repaired: true, repairedCount: 1 }); + expect(branchAssistantContents(sessionManager)).toEqual([ + [{ type: "text", text: "[assistant reasoning omitted]" }], + ]); + }); + + it("preserves downstream branch suffix entries after rewriting the first repaired assistant", () => { + const sessionManager = SessionManager.inMemory(); + sessionManager.appendMessage(asAppendMessage({ role: "user", content: "first", timestamp: 1 })); + sessionManager.appendMessage( + asAppendMessage({ + role: "assistant", + content: [ + { type: "thinking", thinking: "private", thinkingSignature: "sig_bad" }, + { type: "text", text: "first answer" }, + ], + timestamp: 2, + }), + ); + sessionManager.appendMessage( + asAppendMessage({ role: "user", content: "follow-up", timestamp: 3 }), + ); + sessionManager.appendMessage( + asAppendMessage({ + role: "assistant", + content: [{ type: "text", text: "follow-up answer" }], + timestamp: 4, + }), + ); + + const result = repairRejectedThinkingReplayInSessionManager({ sessionManager }); + + expect(result).toMatchObject({ repaired: true, repairedCount: 1 }); + expect(branchMessages(sessionManager).map((message) => message.role)).toEqual([ + "user", + "assistant", + "user", + "assistant", + ]); + expect(branchAssistantContents(sessionManager)).toEqual([ + [{ type: "text", text: "first answer" }], + [{ type: "text", text: "follow-up answer" }], + ]); + }); + + it("does not rewrite sessions without active-branch thinking blocks", () => { + const sessionManager = SessionManager.inMemory(); + sessionManager.appendMessage(asAppendMessage({ role: "user", content: "first", timestamp: 1 })); + sessionManager.appendMessage( + asAppendMessage({ + role: "assistant", + content: [{ type: "text", text: "visible answer" }], + timestamp: 2, + }), + ); + + const beforeLeafId = sessionManager.getLeafId(); + const result = repairRejectedThinkingReplayInSessionManager({ sessionManager }); + + expect(result).toMatchObject({ + repaired: false, + repairedCount: 0, + reason: "no thinking blocks on active branch", + }); + expect(sessionManager.getLeafId()).toBe(beforeLeafId); + }); +}); diff --git a/src/agents/embedded-agent-runner/thinking-replay-repair.ts b/src/agents/embedded-agent-runner/thinking-replay-repair.ts new file mode 100644 index 000000000000..beea522cdf15 --- /dev/null +++ b/src/agents/embedded-agent-runner/thinking-replay-repair.ts @@ -0,0 +1,70 @@ +/** + * Repairs persisted signed-thinking replay state after provider-confirmed rejection. + */ +import { emitSessionTranscriptUpdate } from "../../sessions/transcript-events.js"; +import type { AgentMessage } from "../runtime/index.js"; +import { log } from "./logger.js"; +import { stripThinkingBlocksFromMessage } from "./thinking.js"; +import { rewriteTranscriptEntriesInSessionManager } from "./transcript-rewrite.js"; + +type RewritableSessionManager = Parameters< + typeof rewriteTranscriptEntriesInSessionManager +>[0]["sessionManager"]; + +export function repairRejectedThinkingReplayInSessionManager(params: { + sessionManager: RewritableSessionManager; + sessionFile?: string; + sessionId?: string; + sessionKey?: string; + agentId?: string; +}): { repaired: boolean; repairedCount: number; reason?: string } { + const replacements: Array<{ entryId: string; message: AgentMessage }> = []; + for (const entry of params.sessionManager.getBranch()) { + if (entry.type !== "message") { + continue; + } + const replacement = stripThinkingBlocksFromMessage(entry.message); + if (replacement === entry.message) { + continue; + } + replacements.push({ entryId: entry.id, message: replacement }); + } + + if (replacements.length === 0) { + return { + repaired: false, + repairedCount: 0, + reason: "no thinking blocks on active branch", + }; + } + + const rewriteResult = rewriteTranscriptEntriesInSessionManager({ + sessionManager: params.sessionManager, + replacements, + }); + if (!rewriteResult.changed) { + return { + repaired: false, + repairedCount: 0, + reason: rewriteResult.reason, + }; + } + + if (params.sessionFile) { + emitSessionTranscriptUpdate({ + sessionFile: params.sessionFile, + sessionKey: params.sessionKey, + ...(params.agentId ? { agentId: params.agentId } : {}), + }); + } + log.warn( + `[session-recovery] stripped thinking blocks after provider rejected replay: ` + + `repaired=${rewriteResult.rewrittenEntries} sessionKey=${params.sessionKey ?? params.sessionId ?? "unknown"}`, + ); + + return { + repaired: true, + repairedCount: rewriteResult.rewrittenEntries, + reason: rewriteResult.reason, + }; +} diff --git a/src/agents/embedded-agent-runner/thinking.test.ts b/src/agents/embedded-agent-runner/thinking.test.ts index 38fbc6e11789..dac06116b5e4 100644 --- a/src/agents/embedded-agent-runner/thinking.test.ts +++ b/src/agents/embedded-agent-runner/thinking.test.ts @@ -2,7 +2,7 @@ // recovery behavior for provider transcripts and active assistant turns. import type { AgentMessage } from "openclaw/plugin-sdk/agent-core"; import { createAssistantMessageEventStream } from "openclaw/plugin-sdk/llm"; -import { describe, expect, it } from "vitest"; +import { describe, expect, it, vi } from "vitest"; import { castAgentMessage, castAgentMessages } from "../test-helpers/agent-message-fixtures.js"; import { OMITTED_ASSISTANT_REASONING_TEXT, @@ -594,6 +594,139 @@ describe("wrapAnthropicStreamWithRecovery", () => { expect(retryMessage.content).toEqual([{ type: "text", text: "visible answer" }]); }); + it("notifies recovery only after a rejected request retry succeeds", async () => { + let callCount = 0; + const recovered = vi.fn(); + const finalMessage = createTestAssistantMessage({ + content: [{ type: "text", text: "recovered" }], + stopReason: "stop", + }); + const originalMessages = castAgentMessages([ + { + role: "assistant", + content: [ + { type: "thinking", thinking: "secret", thinkingSignature: "sig" }, + { type: "text", text: "visible answer" }, + ], + }, + ]); + const wrapped = wrapAnthropicStreamWithRecovery( + (() => { + callCount += 1; + if (callCount === 1) { + return Promise.reject(anthropicThinkingError); + } + const stream = createAssistantMessageEventStream(); + queueMicrotask(() => { + stream.push({ type: "done", reason: "stop", message: finalMessage }); + stream.end(); + }); + return stream; + }) as Parameters[0], + { id: "test-session", onRecoveredAnthropicThinking: recovered }, + ); + + const response = (await wrapped( + {} as never, + { + messages: originalMessages, + } as never, + {} as never, + )) as { result: () => Promise } & AsyncIterable; + for await (const event of response) { + void event; + // Drain the retry stream before reading result(). + } + + await expect(response.result()).resolves.toEqual(finalMessage); + expect(callCount).toBe(2); + expect(recovered).toHaveBeenCalledTimes(1); + expect(recovered).toHaveBeenCalledWith({ + originalMessages, + cleanedMessages: [ + { + ...originalMessages[0], + content: [{ type: "text", text: "visible answer" }], + }, + ], + }); + }); + + it("does not notify recovery when the stripped-thinking retry also fails", async () => { + const recovered = vi.fn(); + let callCount = 0; + const retryError = new Error("retry failed"); + const wrapped = wrapAnthropicStreamWithRecovery( + (() => { + callCount += 1; + return Promise.reject(callCount === 1 ? anthropicThinkingError : retryError); + }) as Parameters[0], + { id: "test-session", onRecoveredAnthropicThinking: recovered }, + ); + + await expect( + wrapped( + {} as never, + { + messages: castAgentMessages([ + { + role: "assistant", + content: [{ type: "thinking", thinking: "secret", thinkingSignature: "sig" }], + }, + ]), + } as never, + {} as never, + ), + ).rejects.toBe(retryError); + expect(recovered).not.toHaveBeenCalled(); + }); + + it("does not notify recovery when the stripped-thinking retry resolves to a stream error", async () => { + const recovered = vi.fn(); + let callCount = 0; + const errorMessage = createTestStreamErrorMessage("retry stream failed"); + const wrapped = wrapAnthropicStreamWithRecovery( + (() => { + callCount += 1; + if (callCount === 1) { + return Promise.reject(anthropicThinkingError); + } + const stream = createAssistantMessageEventStream(); + queueMicrotask(() => { + stream.push({ + type: "error", + reason: "error", + error: errorMessage, + }); + stream.end(); + }); + return stream; + }) as Parameters[0], + { id: "test-session", onRecoveredAnthropicThinking: recovered }, + ); + + const response = (await wrapped( + {} as never, + { + messages: castAgentMessages([ + { + role: "assistant", + content: [{ type: "thinking", thinking: "secret", thinkingSignature: "sig" }], + }, + ]), + } as never, + {} as never, + )) as { result: () => Promise } & AsyncIterable; + for await (const event of response) { + void event; + // Drain the retry stream before reading result(). + } + + await expect(response.result()).resolves.toEqual(errorMessage); + expect(callCount).toBe(2); + expect(recovered).not.toHaveBeenCalled(); + }); + it("retries Bedrock-style invalid thinking signature errors", async () => { let callCount = 0; const bedrockThinkingError = new Error( diff --git a/src/agents/embedded-agent-runner/thinking.ts b/src/agents/embedded-agent-runner/thinking.ts index 50c3202e3090..02e5b5589bb9 100644 --- a/src/agents/embedded-agent-runner/thinking.ts +++ b/src/agents/embedded-agent-runner/thinking.ts @@ -10,7 +10,15 @@ import { log } from "./logger.js"; type AssistantContentBlock = Extract["content"][number]; type AssistantMessage = Extract; type RecoveryAssessment = "valid" | "incomplete-thinking" | "incomplete-text"; -type RecoverySessionMeta = { id: string; recoveredAnthropicThinking?: boolean }; +export type AnthropicThinkingRecovery = { + originalMessages: AgentMessage[]; + cleanedMessages: AgentMessage[]; +}; +type RecoverySessionMeta = { + id: string; + recoveredAnthropicThinking?: boolean; + onRecoveredAnthropicThinking?: (recovery: AnthropicThinkingRecovery) => void | Promise; +}; const THINKING_BLOCK_ERROR_PATTERN = /(?:thinking|redacted_thinking).*?(?:cannot be modified|signature|invalid|missing|empty|blank)|(?:signature|invalid|missing|empty|blank).*?(?:thinking|redacted_thinking)/i; @@ -417,26 +425,31 @@ export function shouldPreserveLatestAssistantThinking(messages: AgentMessage[]): return shouldPreserveCurrentToolTurnReasoning(messages, latestAssistantIndex, latestUserIndex); } +export function stripThinkingBlocksFromMessage(message: AgentMessage): AgentMessage { + if (!isAssistantMessageWithContent(message)) { + return message; + } + const nextContent = message.content.filter((block) => !isThinkingBlock(block)); + if (nextContent.length === message.content.length) { + return message; + } + return { + ...message, + content: nextContent.length > 0 ? nextContent : buildOmittedAssistantReasoningContent(), + }; +} + function stripAllThinkingBlocks(messages: AgentMessage[]): AgentMessage[] { let touched = false; const out: AgentMessage[] = []; for (const message of messages) { - if (!isAssistantMessageWithContent(message)) { - out.push(message); + const stripped = stripThinkingBlocksFromMessage(message); + if (stripped === message) { + out.push(stripped); continue; } - - const nextContent = message.content.filter((block) => !isThinkingBlock(block)); - if (nextContent.length === message.content.length) { - out.push(message); - continue; - } - touched = true; - out.push({ - ...message, - content: nextContent.length > 0 ? nextContent : buildOmittedAssistantReasoningContent(), - }); + out.push(stripped); } return touched ? out : messages; } @@ -592,9 +605,58 @@ function getAssistantMessageErrorText( return typeof errorMessage === "string" ? errorMessage : ""; } +async function notifyRecoveredAnthropicThinking( + sessionMeta: RecoverySessionMeta, + recovery: AnthropicThinkingRecovery, +): Promise { + try { + await sessionMeta.onRecoveredAnthropicThinking?.(recovery); + } catch (error: unknown) { + log.warn( + `[session-recovery] Anthropic thinking transcript repair hook failed: sessionId=${sessionMeta.id} error=${formatErrorMessage(error)}`, + ); + } +} + +function isSuccessfulRecoveryRetryResult(message: AssistantMessage | undefined): boolean { + if (!message) { + return false; + } + return message.stopReason !== "error" && message.stopReason !== "aborted"; +} + +function wrapRetryStreamWithRecoveryNotification( + retryStream: ReturnType, + notify: () => Promise, +): ReturnType { + if (retryStream instanceof Promise) { + return retryStream.then((resolved) => + wrapRetryStreamWithRecoveryNotification(resolved as ReturnType, notify), + ) as ReturnType; + } + const streamWithResult = retryStream as unknown as { + result?: () => Promise; + }; + if (typeof streamWithResult.result !== "function") { + return retryStream; + } + const result = streamWithResult.result.bind(streamWithResult); + let notified = false; + streamWithResult.result = async () => { + const message = await result(); + if (!notified && isSuccessfulRecoveryRetryResult(message)) { + notified = true; + await notify(); + } + return message; + }; + return retryStream; +} + async function retryStreamWithoutThinking( outer: ReturnType, retry: () => ReturnType, + notify: () => Promise, ): Promise { const retryStream = retry(); const resolvedRetry = retryStream instanceof Promise ? await retryStream : retryStream; @@ -602,6 +664,9 @@ async function retryStreamWithoutThinking( outer.push(chunk as Parameters[0]); } const result = await (resolvedRetry as { result?: () => Promise }).result?.(); + if (isSuccessfulRecoveryRetryResult(result)) { + await notify(); + } return result as AssistantMessage; } @@ -610,6 +675,7 @@ async function pumpStreamWithRecovery( stream: ReturnType, sessionMeta: RecoverySessionMeta, retry: () => ReturnType, + notify: () => Promise, ): Promise { let yieldedOutput = false; try { @@ -631,7 +697,7 @@ async function pumpStreamWithRecovery( log.warn( `[session-recovery] Anthropic thinking stream error; retrying once without thinking blocks: sessionId=${sessionMeta.id}`, ); - return retryStreamWithoutThinking(outer, retry); + return retryStreamWithoutThinking(outer, retry, notify); } } } else { @@ -655,7 +721,7 @@ async function pumpStreamWithRecovery( log.warn( `[session-recovery] Anthropic thinking error during stream; retrying once without thinking blocks: sessionId=${sessionMeta.id}`, ); - return retryStreamWithoutThinking(outer, retry); + return retryStreamWithoutThinking(outer, retry, notify); } } @@ -664,7 +730,10 @@ export function wrapAnthropicStreamWithRecovery( sessionMeta: RecoverySessionMeta, ): StreamFn { return (model, context, options) => { - const requestMeta: RecoverySessionMeta = { id: sessionMeta.id }; + const requestMeta: RecoverySessionMeta = { + id: sessionMeta.id, + onRecoveredAnthropicThinking: sessionMeta.onRecoveredAnthropicThinking, + }; const contextRecord = context as unknown as { messages?: unknown }; const originalMessages = Array.isArray(contextRecord.messages) ? (contextRecord.messages as AgentMessage[]) @@ -677,6 +746,11 @@ export function wrapAnthropicStreamWithRecovery( } as typeof context; return innerStreamFn(model, nextContext, options); }; + const notify = () => + notifyRecoveredAnthropicThinking(requestMeta, { + originalMessages, + cleanedMessages: stripAllThinkingBlocks(originalMessages), + }); const stream = innerStreamFn(model, context, options); if (stream instanceof Promise) { @@ -688,15 +762,19 @@ export function wrapAnthropicStreamWithRecovery( log.warn( `[session-recovery] Anthropic thinking request rejected; retrying once without thinking blocks: sessionId=${requestMeta.id}`, ); - return retry(); + return wrapRetryStreamWithRecoveryNotification(retry(), notify); }) as ReturnType; } const outer = createAssistantMessageEventStream(); - const finalResultPromise = pumpStreamWithRecovery(outer, stream, requestMeta, retry).finally( - () => { - outer.end(); - }, - ); + const finalResultPromise = pumpStreamWithRecovery( + outer, + stream, + requestMeta, + retry, + notify, + ).finally(() => { + outer.end(); + }); outer.result = () => finalResultPromise; return outer as unknown as ReturnType; };