From bee1d44bf9f1b7cdf9a8925926c278bc552e2a64 Mon Sep 17 00:00:00 2001 From: roboomp Date: Fri, 14 Aug 2026 14:34:27 +0000 Subject: [PATCH] fix(ai): preserved anthropic tool-search replay blocks - Retained tool-search server calls and opaque results in signed assistant history across direct streams, gateways, and custom-endpoint projection. - Added replay regressions for interleaved thinking and client tool continuations. Fixes #8559 --- packages/agent/src/compaction/compaction.ts | 4 +- packages/ai/CHANGELOG.md | 4 + .../providers/anthropic-messages-server.ts | 10 +- packages/ai/src/providers/anthropic-wire.ts | 36 +++++- packages/ai/src/providers/anthropic.ts | 9 +- packages/ai/src/types.ts | 9 +- .../ai/src/utils/leaked-thinking-stream.ts | 4 +- .../ai/test/anthropic-stream-envelope.test.ts | 116 +++++++++++++++++- .../auth-gateway-anthropic-messages.test.ts | 43 +++++++ .../ai/test/leaked-thinking-stream.test.ts | 52 ++++++++ 10 files changed, 265 insertions(+), 22 deletions(-) diff --git a/packages/agent/src/compaction/compaction.ts b/packages/agent/src/compaction/compaction.ts index c9c64b09a..7ce34fb39 100644 --- a/packages/agent/src/compaction/compaction.ts +++ b/packages/agent/src/compaction/compaction.ts @@ -470,8 +470,8 @@ function computeMessageTokens(message: AgentMessage, options?: { excludeEncrypte if (!options?.excludeEncryptedReasoning) fragments.push(block.data); } else if (block.type === "anthropicServerTool") { // Native Anthropic server-tool call/result replayed verbatim on the - // wire (server_tool_use input, web_search_tool_result - // encrypted_content). Opaque provider-replay state the provider still + // wire (server_tool_use input and opaque result content). This opaque + // provider-replay state the provider still // bills for on same-provider replay; excluded from the compaction // floor like other encrypted reasoning because its local byte size // diverges from provider billing. diff --git a/packages/ai/CHANGELOG.md b/packages/ai/CHANGELOG.md index 08c36e86c..3da4ed69e 100644 --- a/packages/ai/CHANGELOG.md +++ b/packages/ai/CHANGELOG.md @@ -2,6 +2,10 @@ ## [Unreleased] +### Fixed + +- Fixed Anthropic custom signing-proxy continuations dropping native tool-search call/result blocks from signed assistant history, preserving interleaved thinking for byte-identical replay ([#8559](https://github.com/can1357/oh-my-pi/issues/8559)). + ## [17.3.4] - 2026-08-14 ### Fixed diff --git a/packages/ai/src/providers/anthropic-messages-server.ts b/packages/ai/src/providers/anthropic-messages-server.ts index 6bb79d992..dc2e74a38 100644 --- a/packages/ai/src/providers/anthropic-messages-server.ts +++ b/packages/ai/src/providers/anthropic-messages-server.ts @@ -27,7 +27,7 @@ import { type AnthropicUserContentBlock, anthropicMessagesRequestSchema, } from "./anthropic-messages-server-schema"; -import { isAnthropicWebSearchHistoryBlock } from "./anthropic-wire"; +import { isAnthropicServerToolHistoryBlock } from "./anthropic-wire"; /** * Anthropic Messages API (https://docs.anthropic.com/en/api/messages) ↔ pi-ai @@ -216,10 +216,10 @@ function walkAssistantContent( break; case "server_tool_use": case "web_search_tool_result": - if (isAnthropicWebSearchHistoryBlock(block)) { - // Native web-search call/result. Anthropic requires these - // replayed verbatim (encrypted_content included), so retain - // the block instead of flattening it to text. + case "tool_search_tool_result": + if (isAnthropicServerToolHistoryBlock(block)) { + // Anthropic requires supported server-tool call/results replayed + // verbatim, so retain each opaque block instead of flattening it. out.push({ type: "anthropicServerTool", block: { ...block } }); } else { // Other server tools use distinct result block types that omp diff --git a/packages/ai/src/providers/anthropic-wire.ts b/packages/ai/src/providers/anthropic-wire.ts index 11665e0d5..bca3c24e7 100644 --- a/packages/ai/src/providers/anthropic-wire.ts +++ b/packages/ai/src/providers/anthropic-wire.ts @@ -77,6 +77,11 @@ export type ServerToolUseBlockParam = { /** Web-search server-tool call whose matching result is replayable by omp. */ export type WebSearchServerToolUseBlockParam = ServerToolUseBlockParam & { name: "web_search" }; +/** Tool-search server-tool call whose matching result is replayable by omp. */ +export type ToolSearchServerToolUseBlockParam = ServerToolUseBlockParam & { + name: "tool_search_tool_regex" | "tool_search_tool_bm25"; +}; + /** Native web-search result replayed inside an assistant turn. */ export type WebSearchToolResultBlockParam = { type: "web_search_tool_result"; @@ -85,18 +90,37 @@ export type WebSearchToolResultBlockParam = { [key: string]: unknown; }; -/** True for the complete native web-search history variants omp can replay. */ -export function isAnthropicWebSearchHistoryBlock(block: { +/** Native tool-search result replayed inside an assistant turn. */ +export type ToolSearchToolResultBlockParam = { + type: "tool_search_tool_result"; + tool_use_id: string; + content: unknown; + [key: string]: unknown; +}; + +/** Anthropic server-tool history variants omp can replay atomically. */ +export type AnthropicServerToolHistoryBlockParam = + | WebSearchServerToolUseBlockParam + | WebSearchToolResultBlockParam + | ToolSearchServerToolUseBlockParam + | ToolSearchToolResultBlockParam; + +/** True when a block is complete Anthropic server-tool history omp can replay. */ +export function isAnthropicServerToolHistoryBlock(block: { type: string; name?: unknown; id?: unknown; tool_use_id?: unknown; content?: unknown; -}): block is WebSearchServerToolUseBlockParam | WebSearchToolResultBlockParam { +}): block is AnthropicServerToolHistoryBlockParam { if (block.type === "server_tool_use") { - return block.name === "web_search" && typeof block.id === "string" && block.id.length > 0; + const supportedName = + block.name === "web_search" || + block.name === "tool_search_tool_regex" || + block.name === "tool_search_tool_bm25"; + return supportedName && typeof block.id === "string" && block.id.length > 0; } - if (block.type === "web_search_tool_result") { + if (block.type === "web_search_tool_result" || block.type === "tool_search_tool_result") { return typeof block.tool_use_id === "string" && block.tool_use_id.length > 0 && Object.hasOwn(block, "content"); } return false; @@ -132,6 +156,7 @@ export type ContentBlockParam = | ToolResultBlockParam | ServerToolUseBlockParam | WebSearchToolResultBlockParam + | ToolSearchToolResultBlockParam | ThinkingBlockParam | RedactedThinkingBlockParam | FallbackBlockParam; @@ -319,6 +344,7 @@ export type ResponseContentBlock = | { type: "tool_use"; id: string; name: string; input?: Record | null } | ServerToolUseBlockParam | WebSearchToolResultBlockParam + | ToolSearchToolResultBlockParam | { type: "fallback"; from: { model: string }; to: { model: string } }; export type ContentBlockDelta = diff --git a/packages/ai/src/providers/anthropic.ts b/packages/ai/src/providers/anthropic.ts index 2d3a49170..6e4fda1d6 100644 --- a/packages/ai/src/providers/anthropic.ts +++ b/packages/ai/src/providers/anthropic.ts @@ -79,7 +79,7 @@ import { type Usage as AnthropicWireUsage, type ContentBlockParam, type FallbackParam, - isAnthropicWebSearchHistoryBlock, + isAnthropicServerToolHistoryBlock, type MessageCreateParams, type MessageCreateParamsStreaming, type MessageParam, @@ -2429,8 +2429,11 @@ const streamAnthropicOnce = ( kind: "redactedThinking", }); } else if ( - isAnthropicWebSearchHistoryBlock(event.content_block) && - umansGatewayWebSearchHeader === undefined + isAnthropicServerToolHistoryBlock(event.content_block) && + (umansGatewayWebSearchHeader === undefined || + (event.content_block.type === "server_tool_use" + ? event.content_block.name !== "web_search" + : event.content_block.type !== "web_search_tool_result")) ) { streamedReplayUnsafeContent = true; const block: Block = { diff --git a/packages/ai/src/types.ts b/packages/ai/src/types.ts index 43ae36965..48452b5c8 100644 --- a/packages/ai/src/types.ts +++ b/packages/ai/src/types.ts @@ -716,8 +716,9 @@ export interface AnthropicFallbackContent { } /** - * Verbatim Anthropic web-search call/result retained for same-provider - * history replay. Other providers discard it in `transformMessages`. + * Verbatim Anthropic web-search or tool-search call/result retained for + * same-provider history replay. Other providers discard it in + * `transformMessages`. */ export interface AnthropicServerToolContent { type: "anthropicServerTool"; @@ -725,12 +726,12 @@ export interface AnthropicServerToolContent { | { type: "server_tool_use"; id: string; - name: "web_search"; + name: "web_search" | "tool_search_tool_regex" | "tool_search_tool_bm25"; input?: Record | null; [key: string]: unknown; } | { - type: "web_search_tool_result"; + type: "web_search_tool_result" | "tool_search_tool_result"; tool_use_id: string; content: unknown; [key: string]: unknown; diff --git a/packages/ai/src/utils/leaked-thinking-stream.ts b/packages/ai/src/utils/leaked-thinking-stream.ts index 93729a82d..79077ef04 100644 --- a/packages/ai/src/utils/leaked-thinking-stream.ts +++ b/packages/ai/src/utils/leaked-thinking-stream.ts @@ -25,7 +25,7 @@ * events are forwarded verbatim. */ -import { isAnthropicWebSearchHistoryBlock } from "../providers/anthropic-wire"; +import { isAnthropicServerToolHistoryBlock } from "../providers/anthropic-wire"; import type { AnthropicServerToolContent, AssistantMessage, @@ -430,7 +430,7 @@ class LeakedThinkingProjector { const pairedIndexes = new Set(); for (let srcIndex = 0; srcIndex < message.content.length; srcIndex++) { const content = message.content[srcIndex]; - if (content?.type !== "anthropicServerTool" || !isAnthropicWebSearchHistoryBlock(content.block)) continue; + if (content?.type !== "anthropicServerTool" || !isAnthropicServerToolHistoryBlock(content.block)) continue; if (content.block.type === "server_tool_use") { pendingCalls.set(content.block.id, srcIndex); continue; diff --git a/packages/ai/test/anthropic-stream-envelope.test.ts b/packages/ai/test/anthropic-stream-envelope.test.ts index f6c2f9364..7184a8b3a 100644 --- a/packages/ai/test/anthropic-stream-envelope.test.ts +++ b/packages/ai/test/anthropic-stream-envelope.test.ts @@ -6,7 +6,10 @@ import { type AnthropicMessagesClientLike, type AnthropicRequestOptions, } from "@oh-my-pi/pi-ai/providers/anthropic-client"; -import type { WebSearchToolResultBlockParam } from "@oh-my-pi/pi-ai/providers/anthropic-wire"; +import type { + ToolSearchToolResultBlockParam, + WebSearchToolResultBlockParam, +} from "@oh-my-pi/pi-ai/providers/anthropic-wire"; import type { AssistantMessageEvent, Context, Model, ModelSpec, ProviderSessionState } from "@oh-my-pi/pi-ai/types"; import { buildModel } from "@oh-my-pi/pi-catalog/build"; import { structuredCloneJSON } from "@oh-my-pi/pi-utils"; @@ -676,6 +679,117 @@ describe("anthropic stream envelope handling", () => { ]); }); + it("replays tool-search server blocks between signed thinking and a client tool call", async () => { + const toolSearchResult: ToolSearchToolResultBlockParam = { + type: "tool_search_tool_result", + tool_use_id: "srvtoolu_search", + content: { + type: "tool_search_tool_search_result", + tool_references: [{ type: "tool_reference", tool_name: "_read" }], + }, + }; + vi.spyOn(AnthropicMessages.prototype, "create").mockImplementation( + () => + createMockRequest([ + { + type: "message_start", + message: { + id: "msg_tool_search", + usage: { + input_tokens: 12, + output_tokens: 0, + cache_read_input_tokens: 0, + cache_creation_input_tokens: 0, + }, + }, + }, + { type: "content_block_start", index: 0, content_block: { type: "thinking", thinking: "" } }, + { type: "content_block_delta", index: 0, delta: { type: "thinking_delta", thinking: "Find read." } }, + { type: "content_block_delta", index: 0, delta: { type: "signature_delta", signature: "sig-1" } }, + { type: "content_block_stop", index: 0 }, + { + type: "content_block_start", + index: 1, + content_block: { + type: "server_tool_use", + id: "srvtoolu_search", + name: "tool_search_tool_regex", + }, + }, + { + type: "content_block_delta", + index: 1, + delta: { type: "input_json_delta", partial_json: '{"pattern":"read"}' }, + }, + { type: "content_block_stop", index: 1 }, + { type: "content_block_start", index: 2, content_block: toolSearchResult }, + { type: "content_block_stop", index: 2 }, + { type: "content_block_start", index: 3, content_block: { type: "thinking", thinking: "" } }, + { type: "content_block_delta", index: 3, delta: { type: "thinking_delta", thinking: "Use read." } }, + { type: "content_block_delta", index: 3, delta: { type: "signature_delta", signature: "sig-2" } }, + { type: "content_block_stop", index: 3 }, + { + type: "content_block_start", + index: 4, + content_block: { type: "tool_use", id: "tool_1", name: "_read", input: {} }, + }, + { + type: "content_block_delta", + index: 4, + delta: { type: "input_json_delta", partial_json: '{"path":"notes.txt"}' }, + }, + { type: "content_block_stop", index: 4 }, + { + type: "message_delta", + delta: { stop_reason: "tool_use" }, + usage: { + input_tokens: 12, + output_tokens: 20, + cache_read_input_tokens: 0, + cache_creation_input_tokens: 0, + }, + }, + { type: "message_stop" }, + ]) as never, + ); + + const stream = streamAnthropic(model, context, { apiKey: "sk-ant-test" }); + for await (const _ of stream) { + // drain stream + } + const result = structuredCloneJSON(await stream.result()); + const replay = convertAnthropicMessages( + [ + context.messages[0], + result, + { + role: "toolResult", + toolCallId: "tool_1", + toolName: "_read", + content: [{ type: "text", text: "notes" }], + isError: false, + timestamp: 2, + }, + ], + model, + false, + ); + const assistant = replay.find(message => message.role === "assistant"); + + expect(assistant?.content).toEqual([ + { type: "thinking", thinking: "Find read.", signature: "sig-1" }, + { + type: "server_tool_use", + id: "srvtoolu_search", + name: "tool_search_tool_regex", + input: { pattern: "read" }, + }, + toolSearchResult, + { type: "thinking", thinking: "Use read.", signature: "sig-2" }, + { type: "tool_use", id: "tool_1", name: "_read", input: { path: "notes.txt" } }, + ]); + }); + it("does not persist a code-execution call without its unsupported result block", async () => { vi.spyOn(AnthropicMessages.prototype, "create").mockImplementation( () => diff --git a/packages/ai/test/auth-gateway-anthropic-messages.test.ts b/packages/ai/test/auth-gateway-anthropic-messages.test.ts index 538df9197..5cc34d2fc 100644 --- a/packages/ai/test/auth-gateway-anthropic-messages.test.ts +++ b/packages/ai/test/auth-gateway-anthropic-messages.test.ts @@ -1,6 +1,8 @@ import { describe, expect, it } from "bun:test"; import { encodeResponse, encodeStream, parseRequest } from "@oh-my-pi/pi-ai/providers/anthropic-messages-server"; import type { + ToolSearchServerToolUseBlockParam, + ToolSearchToolResultBlockParam, WebSearchServerToolUseBlockParam, WebSearchToolResultBlockParam, } from "@oh-my-pi/pi-ai/providers/anthropic-wire"; @@ -344,6 +346,47 @@ describe("anthropic-messages parseRequest", () => { ]); }); + it("preserves inbound assistant tool-search call/result blocks verbatim", () => { + const serverToolUse: ToolSearchServerToolUseBlockParam = { + type: "server_tool_use", + id: "srvtoolu_search", + name: "tool_search_tool_regex", + input: { pattern: "read" }, + }; + const searchResult: ToolSearchToolResultBlockParam = { + type: "tool_search_tool_result", + tool_use_id: "srvtoolu_search", + content: { + type: "tool_search_tool_search_result", + tool_references: [{ type: "tool_reference", tool_name: "_read" }], + }, + }; + const parsed = parseRequest({ + model: "claude-opus-4-7", + max_tokens: 8, + messages: [ + { role: "user", content: "read notes" }, + { + role: "assistant", + content: [ + { type: "thinking", thinking: "find read", signature: "sig-1" }, + serverToolUse, + searchResult, + { type: "text", text: "tool loaded" }, + ], + }, + ], + }); + const assistant = parsed.context.messages.find(message => message.role === "assistant"); + + expect(assistant?.content).toEqual([ + { type: "thinking", thinking: "find read", thinkingSignature: "sig-1" }, + { type: "anthropicServerTool", block: serverToolUse }, + { type: "anthropicServerTool", block: searchResult }, + { type: "text", text: "tool loaded" }, + ]); + }); + it("flattens malformed web-search history blocks instead of preserving invalid replay state", () => { const parsed = parseRequest({ model: "claude-opus-4-7", diff --git a/packages/ai/test/leaked-thinking-stream.test.ts b/packages/ai/test/leaked-thinking-stream.test.ts index bd0564d22..e699acf3e 100644 --- a/packages/ai/test/leaked-thinking-stream.test.ts +++ b/packages/ai/test/leaked-thinking-stream.test.ts @@ -419,6 +419,58 @@ describe("wrapLeakedThinkingStream", () => { expect(result.content.slice(1, 3)).toEqual(serverBlocks); }); + it("preserves complete Anthropic tool-search history through the custom-endpoint projector", async () => { + const firstThinking: ThinkingContent = { + type: "thinking", + thinking: "find the deferred tool", + thinkingSignature: "sig-1", + }; + const serverBlocks: AnthropicServerToolContent[] = [ + { + type: "anthropicServerTool", + block: { + type: "server_tool_use", + id: "srvtoolu_search", + name: "tool_search_tool_regex", + input: { pattern: "read" }, + }, + }, + { + type: "anthropicServerTool", + block: { + type: "tool_search_tool_result", + tool_use_id: "srvtoolu_search", + content: { + type: "tool_search_tool_search_result", + tool_references: [{ type: "tool_reference", tool_name: "_read" }], + }, + }, + }, + ]; + const secondThinking: ThinkingContent = { + type: "thinking", + thinking: "use the discovered tool", + thinkingSignature: "sig-2", + }; + const call: ToolCall = { + type: "toolCall", + id: "toolu_read", + name: "_read", + arguments: { path: "notes.txt" }, + }; + const content: AssistantMessage["content"] = [firstThinking, ...serverBlocks, secondThinking, call]; + const terminal = msg({ content, stopReason: "toolUse" }); + + const { result } = await runWrapper(inner => { + inner.push({ type: "start", partial: msg() }); + inner.push({ type: "toolcall_start", contentIndex: 4, partial: terminal }); + inner.push({ type: "toolcall_end", contentIndex: 4, toolCall: call, partial: terminal }); + inner.push({ type: "done", reason: "toolUse", message: terminal }); + }); + + expect(result.content).toEqual(content); + }); + it("drops incomplete Anthropic web-search history instead of replaying orphan blocks", async () => { const content: AssistantMessage["content"] = [ {