From b6f1544754a533b3b9a5273237e8676f3c227b28 Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Mon, 2 Nov 2026 12:26:58 +0100 Subject: [PATCH 1/2] fix(ai, ai-client): run client tools when a custom server ends on tool_calls A custom AG-UI server can end a run with finishReason 'tool_calls' and close the stream. The stream processor skipped the tool handover for that finish, because chat() starts a next run in the same stream. So the client tool did not run and the agent loop stopped. The processor now marks that finish and hands the unanswered calls to the client when the stream ends. A new RUN_STARTED clears the mark, so chat() works as before. ChatClient applies its approval check to this handover too, so a needsApproval tool still does not run without approval. --- .changeset/client-tool-tool-calls-finish.md | 6 +++ docs/config.json | 2 +- docs/tools/client-tools.md | 2 + packages/ai-client/src/chat-client.ts | 11 ++++- .../chat-client-foreign-client-tool.test.ts | 47 ++++++++++++++---- .../src/activities/chat/stream/processor.ts | 11 +++++ packages/ai/tests/stream-processor.test.ts | 26 ++++++++++ .../e2e/src/routes/api.foreign-client-tool.ts | 6 ++- .../e2e/src/routes/foreign-client-tool.tsx | 13 ++++- testing/e2e/tests/foreign-client-tool.spec.ts | 49 ++++++++++--------- 10 files changed, 137 insertions(+), 36 deletions(-) create mode 100644 .changeset/client-tool-tool-calls-finish.md diff --git a/.changeset/client-tool-tool-calls-finish.md b/.changeset/client-tool-tool-calls-finish.md new file mode 100644 index 0000000000..72c44b9e74 --- /dev/null +++ b/.changeset/client-tool-tool-calls-finish.md @@ -0,0 +1,6 @@ +--- +'@tanstack/ai': patch +'@tanstack/ai-client': patch +--- + +Run client tools when a custom AG-UI server ends a run with `finishReason: 'tool_calls'` and closes the stream. Before, the client waited for a next run that never came, so the tool did not run and the agent loop stopped. The client now runs the unanswered calls when the stream ends. A tool that needs approval still does not run there, the same as for a success outcome. diff --git a/docs/config.json b/docs/config.json index 5ebdf1256b..b093c507f4 100644 --- a/docs/config.json +++ b/docs/config.json @@ -356,7 +356,7 @@ "label": "Client Tools", "to": "tools/client-tools", "addedAt": "2026-04-15", - "updatedAt": "2026-10-22" + "updatedAt": "2026-11-02" }, { "label": "WebMCP Tools", diff --git a/docs/tools/client-tools.md b/docs/tools/client-tools.md index 2045f04031..cae9d8675e 100644 --- a/docs/tools/client-tools.md +++ b/docs/tools/client-tools.md @@ -402,6 +402,8 @@ The AG-UI spec ends such a run with `RUN_FINISHED` and a success outcome. The to - If the server names the calls in `outcome.pendingToolCallIds`, the client runs only those of them that this run started. - The next request carries each result as a `role: "tool"` message. +Some servers end the run with `finishReason: "tool_calls"` instead, then close the stream. The client does the same steps when the stream closes. + These calls do not run: - A call with a name that no client tool has. diff --git a/packages/ai-client/src/chat-client.ts b/packages/ai-client/src/chat-client.ts index 3085d2f526..24c0b5d62a 100644 --- a/packages/ai-client/src/chat-client.ts +++ b/packages/ai-client/src/chat-client.ts @@ -2814,8 +2814,15 @@ export class ChatClient< await Promise.all(this.pendingToolExecutions.values()) } - // Finalize (idempotent — may already be done by RUN_FINISHED handler) - this.processor.finalizeStream() + // Finalize (idempotent — may already be done by RUN_FINISHED handler). + // It can hand over the calls of a run that ended with 'tool_calls', so + // the approval check in onToolCall applies here too. + this.handingOverPendingToolCalls = true + try { + this.processor.finalizeStream() + } finally { + this.handingOverPendingToolCalls = false + } streamCompletedSuccessfully = true } catch (err: unknown) { const error = err instanceof Error ? err : new Error(String(err)) diff --git a/packages/ai-client/tests/chat-client-foreign-client-tool.test.ts b/packages/ai-client/tests/chat-client-foreign-client-tool.test.ts index 2998b8e7fb..6ff2c297cb 100644 --- a/packages/ai-client/tests/chat-client-foreign-client-tool.test.ts +++ b/packages/ai-client/tests/chat-client-foreign-client-tool.test.ts @@ -9,7 +9,11 @@ import type { StreamChunk } from '@tanstack/ai/client' * An AG-UI server outside TanStack AI ends a run that calls a frontend tool * with a success outcome and no TOOL_CALL_RESULT (issue #1464). */ -function foreignServer() { +function foreignServer( + finished: Record = { + outcome: { type: 'success' }, + }, +) { const requests: Array = [] const connection: ConnectConnectionAdapter = { async *connect(messages) { @@ -26,12 +30,7 @@ function foreignServer() { }, { type: 'TOOL_CALL_ARGS', toolCallId: 'tc-1', delta: '{}' }, { type: 'TOOL_CALL_END', toolCallId: 'tc-1' }, - { - type: 'RUN_FINISHED', - threadId: 't', - runId, - outcome: { type: 'success' }, - }, + { type: 'RUN_FINISHED', threadId: 't', runId, ...finished }, ] : [ { type: 'RUN_STARTED', threadId: 't', runId }, @@ -67,7 +66,37 @@ describe('client tools from a foreign AG-UI server', () => { expect(JSON.stringify(requests[1])).toContain('{\\"reply\\":\\"pong\\"}') }) - it('does not run a tool that needs approval', async () => { + it('runs the tool when the run ends with finishReason tool_calls', async () => { + const execute = vi.fn(async () => { + await new Promise((resolve) => setTimeout(resolve, 5)) + return { reply: 'pong' } + }) + const ping = toolDefinition({ + name: 'ping', + description: 'Reply with pong', + inputSchema: z.object({}), + }).client(execute) + // A custom server that forwards the provider finish reason and closes. + const { connection, requests } = foreignServer({ + finishReason: 'tool_calls', + }) + const client = new ChatClient({ + threadId: 'thread-1', + connection, + tools: [ping], + }) + + await client.sendMessage('ping me') + + await vi.waitFor(() => expect(requests).toHaveLength(2)) + expect(execute).toHaveBeenCalledTimes(1) + expect(JSON.stringify(requests[1])).toContain('{\\"reply\\":\\"pong\\"}') + }) + + it.each([ + ['a success outcome', { outcome: { type: 'success' } }], + ['finishReason tool_calls', { finishReason: 'tool_calls' }], + ])('does not run a tool that needs approval (%s)', async (_, finished) => { const warn = vi.spyOn(console, 'warn').mockImplementation(() => {}) const execute = vi.fn(() => ({ reply: 'pong' })) const ping = toolDefinition({ @@ -76,7 +105,7 @@ describe('client tools from a foreign AG-UI server', () => { inputSchema: z.object({}), needsApproval: true, }).client(execute) - const { connection, requests } = foreignServer() + const { connection, requests } = foreignServer(finished) const client = new ChatClient({ threadId: 'thread-1', connection, diff --git a/packages/ai/src/activities/chat/stream/processor.ts b/packages/ai/src/activities/chat/stream/processor.ts index 89e0862cf3..4ac1dde7b0 100644 --- a/packages/ai/src/activities/chat/stream/processor.ts +++ b/packages/ai/src/activities/chat/stream/processor.ts @@ -279,6 +279,10 @@ export class StreamProcessor { // Tool calls started since a RUN_STARTED that found no other run active, in // order. A success RUN_FINISHED hands the unanswered ones to the client. private runToolCallIds: Array = [] + // The last run ended with finishReason 'tool_calls'. chat() starts a new + // run after it, but a custom server can close the stream there. Then the + // stream end hands the unanswered calls to the client. + private handOverAtStreamEnd = false // Tool calls that an interrupt already handed to onToolCall since that same // reset, so the success path does not run them twice. private readonly toolCallsSentToClient = new Set() @@ -2660,6 +2664,7 @@ export class StreamProcessor { if (this.activeRuns.size === 0) { this.runToolCallIds = [] this.toolCallsSentToClient.clear() + this.handOverAtStreamEnd = false // A new run owes its own stream end. A client that only watches never // resets between runs, so without this only its first run would end. this.streamEndEmitted = false @@ -2723,6 +2728,7 @@ export class StreamProcessor { // The next model call is a new assistant message. Leave this // call's messages closed so its reasoning cannot land here. for (const id of ids) this.sealedMessageIds.add(id) + this.handOverAtStreamEnd = true return } if (!chunk.outcome || chunk.outcome.type === 'success') { @@ -3444,6 +3450,10 @@ export class StreamProcessor { * @see docs/chat-architecture.md#single-shot-text-response — Finalization step */ finalizeStream(): void { + if (this.handOverAtStreamEnd) { + this.handOverAtStreamEnd = false + if (!this.hasError) this.handlePendingToolCalls() + } this.isDone = true for (const event of this.closeChunk()) this.dispatchChunk(event) this.closeChildChunkLanes() @@ -3632,6 +3642,7 @@ export class StreamProcessor { this.thinkingOwners.clear() this.openChunk = null this.finishReason = null + this.handOverAtStreamEnd = false this.hasError = false this.isDone = false this.streamEndEmitted = false diff --git a/packages/ai/tests/stream-processor.test.ts b/packages/ai/tests/stream-processor.test.ts index 4abcbdaa45..40e04e1541 100644 --- a/packages/ai/tests/stream-processor.test.ts +++ b/packages/ai/tests/stream-processor.test.ts @@ -1961,6 +1961,32 @@ describe('StreamProcessor', () => { expect(events.onToolCall).not.toHaveBeenCalled() }) + it('fires at the stream end when the last run finished with tool_calls', () => { + const { events, processor } = run( + ...toolCall('tc-1'), + ev.runFinished('tool_calls'), + ) + expect(events.onToolCall).not.toHaveBeenCalled() + + processor.finalizeStream() + + expect(events.onToolCall).toHaveBeenCalledTimes(1) + expect(events.onToolCall.mock.calls[0]![0].toolCallId).toBe('tc-1') + }) + + it('does not fire for a tool_calls run that the stream continues', () => { + const { events, processor } = run( + ...toolCall('tc-1'), + ev.runFinished('tool_calls'), + ev.runStarted('run-2'), + ev.toolResult('tc-1', '"pong"'), + ev.runFinished('stop', 'run-2'), + ) + processor.finalizeStream() + + expect(events.onToolCall).not.toHaveBeenCalled() + }) + it('skips calls the run answered with TOOL_CALL_RESULT', () => { const { events } = run( ...toolCall('tc-1'), diff --git a/testing/e2e/src/routes/api.foreign-client-tool.ts b/testing/e2e/src/routes/api.foreign-client-tool.ts index f4df742385..2916e4044f 100644 --- a/testing/e2e/src/routes/api.foreign-client-tool.ts +++ b/testing/e2e/src/routes/api.foreign-client-tool.ts @@ -48,6 +48,8 @@ export const Route = createFileRoute('/api/foreign-client-tool')({ const body: unknown = await request.json() const threadId = 'foreign-client-tool' const result = toolResultContent(body) + const toolCalls = + new URL(request.url).searchParams.get('finish') === 'tool_calls' if (result === undefined) { const runId = 'run-1' @@ -64,7 +66,9 @@ export const Route = createFileRoute('/api/foreign-client-tool')({ type: 'RUN_FINISHED', threadId, runId, - outcome: { type: 'success' }, + ...(toolCalls + ? { finishReason: 'tool_calls' } + : { outcome: { type: 'success' } }), }, ]) } diff --git a/testing/e2e/src/routes/foreign-client-tool.tsx b/testing/e2e/src/routes/foreign-client-tool.tsx index aac5c72885..e365531b0c 100644 --- a/testing/e2e/src/routes/foreign-client-tool.tsx +++ b/testing/e2e/src/routes/foreign-client-tool.tsx @@ -19,12 +19,23 @@ const ping = toolDefinition({ export const Route = createFileRoute('/foreign-client-tool')({ component: ForeignClientToolPage, + // `?finish=tool_calls` ends turn 1 with finishReason 'tool_calls' instead + // of a success outcome, the way a custom server forwards it. + validateSearch: (search: Record) => ({ + finish: + search.finish === 'tool_calls' + ? ('tool_calls' as const) + : ('success' as const), + }), }) function ForeignClientToolPage() { + const { finish } = Route.useSearch() const { messages, sendMessage } = useChat({ threadId: 'foreign-client-tool', - connection: fetchServerSentEvents('/api/foreign-client-tool'), + connection: fetchServerSentEvents( + `/api/foreign-client-tool?finish=${finish}`, + ), tools: [ping], }) diff --git a/testing/e2e/tests/foreign-client-tool.spec.ts b/testing/e2e/tests/foreign-client-tool.spec.ts index c0ba0edf5b..0bbd2e9b25 100644 --- a/testing/e2e/tests/foreign-client-tool.spec.ts +++ b/testing/e2e/tests/foreign-client-tool.spec.ts @@ -2,30 +2,35 @@ import { expect, test } from '@playwright/test' /** * An AG-UI producer that is not TanStack AI ends a run that calls a frontend - * tool with a success outcome and leaves the call unanswered (issue #1464). - * The client runs the tool, then starts one more run with the tool result in - * its history. + * tool and leaves the call unanswered. It ends the run with a success outcome + * (issue #1464), or with finishReason 'tool_calls', as a custom server that + * forwards the provider finish reason does. The client runs the tool, then + * starts one more run with the tool result in its history. */ test.describe('frontend tool from a spec-compliant AG-UI server', () => { - test('runs the tool and continues with its result', async ({ page }) => { - const posts: Array = [] - page.on('request', (request) => { - if ( - request.url().includes('/api/foreign-client-tool') && - request.method() === 'POST' - ) { - posts.push(request.postDataJSON()) - } - }) + for (const finish of ['success', 'tool_calls']) { + test(`runs the tool and continues with its result (${finish})`, async ({ + page, + }) => { + const posts: Array = [] + page.on('request', (request) => { + if ( + request.url().includes('/api/foreign-client-tool') && + request.method() === 'POST' + ) { + posts.push(request.postDataJSON()) + } + }) - await page.goto('/foreign-client-tool') + await page.goto(`/foreign-client-tool?finish=${finish}`) - await expect(page.getByTestId('ping-outputs')).toHaveText( - '{"reply":"pong"}', - ) - await expect(page.getByTestId('assistant-text')).toHaveText( - 'ping returned {"reply":"pong"}', - ) - expect(posts).toHaveLength(2) - }) + await expect(page.getByTestId('ping-outputs')).toHaveText( + '{"reply":"pong"}', + ) + await expect(page.getByTestId('assistant-text')).toHaveText( + 'ping returned {"reply":"pong"}', + ) + expect(posts).toHaveLength(2) + }) + } }) From aa2f67b0f41c6b589b49b6ac920fb3cdcd2caee8 Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Mon, 2 Nov 2026 12:46:35 +0100 Subject: [PATCH 2/2] fix(ai, ai-client): keep the tool_calls handover to the top-level stream Two review findings on the stream-end handover: - A subagent child that ended a run with tool_calls handed its calls to the client at SUBAGENT_FINISHED or SUBAGENT_ERROR. A subagent run comes from chat(), which goes on with its next run itself, so only the top-level stream hands over at its end now. - ChatClient did not wait for a tool that finalizeStream() started. A queued user message could go out before the tool result. It now waits, as the success path already does. --- packages/ai-client/src/chat-client.ts | 4 ++ .../chat-client-foreign-client-tool.test.ts | 42 +++++++++++++++++++ .../src/activities/chat/stream/processor.ts | 4 +- .../tests/stream-processor-subagents.test.ts | 32 ++++++++++++++ 4 files changed, 81 insertions(+), 1 deletion(-) diff --git a/packages/ai-client/src/chat-client.ts b/packages/ai-client/src/chat-client.ts index 24c0b5d62a..3ac64b7cad 100644 --- a/packages/ai-client/src/chat-client.ts +++ b/packages/ai-client/src/chat-client.ts @@ -2823,6 +2823,10 @@ export class ChatClient< } finally { this.handingOverPendingToolCalls = false } + // Wait for the tools that finalizeStream() started, as above. + if (this.pendingToolExecutions.size > 0) { + await Promise.all(this.pendingToolExecutions.values()) + } streamCompletedSuccessfully = true } catch (err: unknown) { const error = err instanceof Error ? err : new Error(String(err)) diff --git a/packages/ai-client/tests/chat-client-foreign-client-tool.test.ts b/packages/ai-client/tests/chat-client-foreign-client-tool.test.ts index 6ff2c297cb..447c361e1d 100644 --- a/packages/ai-client/tests/chat-client-foreign-client-tool.test.ts +++ b/packages/ai-client/tests/chat-client-foreign-client-tool.test.ts @@ -93,6 +93,48 @@ describe('client tools from a foreign AG-UI server', () => { expect(JSON.stringify(requests[1])).toContain('{\\"reply\\":\\"pong\\"}') }) + it.each([ + ['a success outcome', { outcome: { type: 'success' } }], + ['finishReason tool_calls', { finishReason: 'tool_calls' }], + ])( + 'sends the tool result before a queued message (%s)', + async (_, finished) => { + const execute = vi.fn(async () => { + await new Promise((resolve) => setTimeout(resolve, 20)) + return { reply: 'pong' } + }) + const ping = toolDefinition({ + name: 'ping', + description: 'Reply with pong', + inputSchema: z.object({}), + }).client(execute) + const { connection, requests } = foreignServer(finished) + const client = new ChatClient({ + threadId: 'thread-1', + connection, + tools: [ping], + }) + + // The second message waits in the queue while the first run streams. + await Promise.all([ + client.sendMessage('ping me'), + client.sendMessage('queued'), + ]) + + await vi.waitFor(() => + expect( + requests.some((request) => + JSON.stringify(request).includes('queued'), + ), + ).toBe(true), + ) + const queued = requests.find((request) => + JSON.stringify(request).includes('queued'), + ) + expect(JSON.stringify(queued)).toContain('{\\"reply\\":\\"pong\\"}') + }, + ) + it.each([ ['a success outcome', { outcome: { type: 'success' } }], ['finishReason tool_calls', { finishReason: 'tool_calls' }], diff --git a/packages/ai/src/activities/chat/stream/processor.ts b/packages/ai/src/activities/chat/stream/processor.ts index 4ac1dde7b0..917406236f 100644 --- a/packages/ai/src/activities/chat/stream/processor.ts +++ b/packages/ai/src/activities/chat/stream/processor.ts @@ -2728,7 +2728,9 @@ export class StreamProcessor { // The next model call is a new assistant message. Leave this // call's messages closed so its reasoning cannot land here. for (const id of ids) this.sealedMessageIds.add(id) - this.handOverAtStreamEnd = true + // Only the top-level stream can end here. A subagent run comes from + // chat(), which goes on with its next run itself. + this.handOverAtStreamEnd = this.subagentRunId === undefined return } if (!chunk.outcome || chunk.outcome.type === 'success') { diff --git a/packages/ai/tests/stream-processor-subagents.test.ts b/packages/ai/tests/stream-processor-subagents.test.ts index 6b9b499938..1ae9380ad8 100644 --- a/packages/ai/tests/stream-processor-subagents.test.ts +++ b/packages/ai/tests/stream-processor-subagents.test.ts @@ -265,4 +265,36 @@ describe('StreamProcessor subagent parts', () => { warn.mockRestore() }) }) + + it.each([EventType.SUBAGENT_FINISHED, EventType.SUBAGENT_ERROR])( + 'does not hand over a child tool_calls run at %s', + (type) => { + const onToolCall = vi.fn() + const processor = new StreamProcessor({ events: { onToolCall } }) + processor.processChunk(ev.runStarted()) + processor.processChunk({ + type: EventType.SUBAGENT_STARTED, + subagentRunId: 'sub-1', + name: 'researcher', + timestamp: Date.now(), + }) + for (const chunk of [ + ev.runStarted('child-run'), + ev.toolStart('tc-1', 'ping'), + ev.toolArgs('tc-1', '{}'), + ev.toolEnd('tc-1'), + ev.runFinished('tool_calls', 'child-run'), + ]) { + processor.processChunk(attributed(chunk, 'sub-1')) + } + processor.processChunk({ + type, + subagentRunId: 'sub-1', + message: 'Stopped', + timestamp: Date.now(), + } as StreamChunk) + + expect(onToolCall).not.toHaveBeenCalled() + }, + ) })