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..3ac64b7cad 100644 --- a/packages/ai-client/src/chat-client.ts +++ b/packages/ai-client/src/chat-client.ts @@ -2814,8 +2814,19 @@ 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 + } + // 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 2998b8e7fb..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 @@ -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,79 @@ 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' }], + ])( + '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' }], + ])('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 +147,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..917406236f 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,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) + // 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') { @@ -3444,6 +3452,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 +3644,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-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() + }, + ) }) 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) + }) + } })