diff --git a/.changeset/tfy-fine-grained-tool-streaming-header.md b/.changeset/tfy-fine-grained-tool-streaming-header.md new file mode 100644 index 000000000..9af6dac41 --- /dev/null +++ b/.changeset/tfy-fine-grained-tool-streaming-header.md @@ -0,0 +1,5 @@ +--- +'@truefoundry/trueforge-core': patch +--- + +Enable fine-grained tool streaming via TrueFoundry `anthropic-beta` header and Anthropic `toolStreaming`, and report invalid tool-arg JSON as a distinct parse error diff --git a/packages/trueforge-core/src/core/llm/VercelAILLM.ts b/packages/trueforge-core/src/core/llm/VercelAILLM.ts index ec126ba88..71cfdc8d1 100644 --- a/packages/trueforge-core/src/core/llm/VercelAILLM.ts +++ b/packages/trueforge-core/src/core/llm/VercelAILLM.ts @@ -132,6 +132,12 @@ function compatibleModel(config: VercelAIProviderConfig): LanguageModel { if (baseUrl === undefined) { throw new Error(`Provider "${provider.type}" requires a baseUrl`); } + // More straightforward than mutating tool definitions for fine-grained tool streaming. + // TODO: fine-grained-tool-streaming-2025-05-14 is considered deprecated. + const outboundHeaders = + provider.type === 'truefoundry' + ? { ...headers, 'anthropic-beta': 'fine-grained-tool-streaming-2025-05-14' } + : headers; const client = createOpenAICompatible({ name: provider.type, baseURL: baseUrl, @@ -141,7 +147,7 @@ function compatibleModel(config: VercelAIProviderConfig): LanguageModel { supportsStructuredOutputs: true, // These endpoints omit token counts from streamed responses unless asked. includeUsage: true, - ...(Object.keys(headers).length > 0 ? { headers } : {}), + ...(Object.keys(outboundHeaders).length > 0 ? { headers: outboundHeaders } : {}), }); return client(model.id); } @@ -390,6 +396,9 @@ function openaiProviderOptions({ * * Prompt caching defaults on: Anthropic bills a write once, then cheaper reads on stable prefixes * (system / tools / early messages). Callers can override via `cache_control` in the request body. + * + * `toolStreaming` is pinned on so every function tool gets `eager_input_streaming` on the Anthropic + * wire (call-level default). */ function anthropicProviderOptions(rawBody: unknown): JSONObject | undefined { const cacheControl = readBodyField({ rawBody, key: 'cache_control' }) ?? { type: 'ephemeral' }; @@ -398,6 +407,7 @@ function anthropicProviderOptions(rawBody: unknown): JSONObject | undefined { const effort = readBodyField({ rawBody, key: 'effort' }); const opts: JSONObject = { cacheControl, + toolStreaming: true, ...(disableParallelToolUse !== undefined ? { disableParallelToolUse } : {}), ...(thinking !== undefined ? { thinking } : {}), ...(effort !== undefined ? { effort } : {}), diff --git a/packages/trueforge-core/src/core/mcp/executeToolCalls.ts b/packages/trueforge-core/src/core/mcp/executeToolCalls.ts index f8a8e77ab..80625f11c 100644 --- a/packages/trueforge-core/src/core/mcp/executeToolCalls.ts +++ b/packages/trueforge-core/src/core/mcp/executeToolCalls.ts @@ -37,6 +37,36 @@ export interface ExecuteToolCallsResult { events: RegisteredPassthroughEvent[]; } +function isJsonObject(value: unknown): value is Record { + return typeof value === 'object' && value !== null && !Array.isArray(value); +} + +/** + * Parse tool-call `function.arguments`. Fine-grained streaming can leave incomplete JSON; + * failures are reported as tool errors so the model can retry rather than treating them as + * generic execution failures. + */ +function parseToolCallArguments( + rawArguments: string, +): { ok: true; args: Record } | { ok: false; parseError: string } { + let parsed: unknown; + try { + parsed = JSON.parse(rawArguments || '{}'); + } catch (error) { + return { + ok: false, + parseError: error instanceof Error ? error.message : 'JSON parse failed', + }; + } + if (!isJsonObject(parsed)) { + return { + ok: false, + parseError: 'Tool call arguments must be a JSON object', + }; + } + return { ok: true, args: parsed }; +} + export async function executeToolCalls({ assistantMessage, toolMapping, @@ -82,12 +112,28 @@ export async function executeToolCalls({ }; } + const parsedArgs = parseToolCallArguments(toolCall.function.arguments); + if (!parsedArgs.ok) { + return { + toolCall, + toolInfo, + response: toolResultResponse({ + text: JSON.stringify({ + error: 'Failed to parse tool call arguments as JSON', + parse_error: parsedArgs.parseError, + }), + isError: true, + }), + failure: true, + completedAt: new Date().toISOString(), + }; + } + try { - const args: Record = JSON.parse(toolCall.function.arguments || '{}') as Record; const response = await toolInfo.toolSet.callTool( { name: toolInfo.originalToolName, - arguments: args, + arguments: parsedArgs.args, }, approvalDecisions.get(toolCall.id), ); diff --git a/packages/trueforge-core/tests/core/llm/VercelAILLM.output.test.ts b/packages/trueforge-core/tests/core/llm/VercelAILLM.output.test.ts index 6c9a784f3..955d2aca1 100644 --- a/packages/trueforge-core/tests/core/llm/VercelAILLM.output.test.ts +++ b/packages/trueforge-core/tests/core/llm/VercelAILLM.output.test.ts @@ -249,7 +249,9 @@ describe('buildProviderOptions', () => { structuredOutputSpec: textSpec, rawBody: {}, }); - expect(opts).toEqual({ anthropic: { cacheControl: { type: 'ephemeral' } } }); + expect(opts).toEqual({ + anthropic: { cacheControl: { type: 'ephemeral' }, toolStreaming: true }, + }); } }); @@ -263,6 +265,7 @@ describe('buildProviderOptions', () => { }); expect(opts['anthropic']).toEqual({ cacheControl: { type: 'ephemeral' }, + toolStreaming: true, thinking: { type: 'adaptive', display: 'summarized' }, effort: 'max', }); @@ -275,7 +278,9 @@ describe('buildProviderOptions', () => { structuredOutputSpec: schemaSpecStrict, rawBody: {}, }); - expect(opts).toEqual({ anthropic: { cacheControl: { type: 'ephemeral' } } }); + expect(opts).toEqual({ + anthropic: { cacheControl: { type: 'ephemeral' }, toolStreaming: true }, + }); }); it('defaults cacheControl to ephemeral when cache_control is omitted', () => { @@ -285,7 +290,7 @@ describe('buildProviderOptions', () => { structuredOutputSpec: textSpec, rawBody: {}, }); - expect(opts['anthropic']).toEqual({ cacheControl: { type: 'ephemeral' } }); + expect(opts['anthropic']).toEqual({ cacheControl: { type: 'ephemeral' }, toolStreaming: true }); }); it('forwards cache_control and disable_parallel_tool_use from rawBody', () => { @@ -298,6 +303,7 @@ describe('buildProviderOptions', () => { expect(opts['anthropic']).toMatchObject({ cacheControl: { type: 'ephemeral', ttl: '1h' }, disableParallelToolUse: true, + toolStreaming: true, }); }); }); diff --git a/packages/trueforge-core/tests/core/mcp/executeToolCalls.test.ts b/packages/trueforge-core/tests/core/mcp/executeToolCalls.test.ts new file mode 100644 index 000000000..251387a6d --- /dev/null +++ b/packages/trueforge-core/tests/core/mcp/executeToolCalls.test.ts @@ -0,0 +1,38 @@ +import type { MappedMCPTool } from '../../../src/core/mcp/convertMCPServers'; +import { executeToolCalls } from '../../../src/core/mcp/executeToolCalls'; +import { makeMockIMCPServer } from '../harnessMocks'; + +describe('executeToolCalls', () => { + it('reports invalid tool-arg JSON as a parse failure without calling the tool', async () => { + const toolSet = makeMockIMCPServer({ name: 'srv', preload: true }); + const toolMapping = new Map([['search', { toolSet, originalToolName: 'search' }]]); + + const { toolCallResults } = await executeToolCalls({ + assistantMessage: { + role: 'assistant', + content: null, + tool_calls: [ + { + id: 'call_1', + type: 'function', + function: { name: 'search', arguments: '{"q":' }, + tool_info: { + type: 'mcp', + mcp_server_id: 'srv', + mcp_server_name: 'srv', + original_tool_name: 'search', + }, + }, + ], + }, + toolMapping, + threadId: 'thread-1', + approvalDecisions: new Map(), + }); + + expect(toolSet.callTool).not.toHaveBeenCalled(); + expect(toolCallResults).toHaveLength(1); + expect(toolCallResults[0]?.failure).toBe(true); + expect(toolCallResults[0]?.message.content).toContain('Failed to parse tool call arguments as JSON'); + }); +});