diff --git a/.changeset/quiet-append-drain.md b/.changeset/quiet-append-drain.md new file mode 100644 index 0000000000..7272da3dae --- /dev/null +++ b/.changeset/quiet-append-drain.md @@ -0,0 +1,6 @@ +--- +'@tanstack/ai': patch +'@tanstack/ai-client': patch +--- + +Keep `append()` pending until the HTTP response is fully processed, including later `RUN_FINISHED` events in the same agent loop. diff --git a/docs/api/ai-client.md b/docs/api/ai-client.md index 16456962b5..59af3902b1 100644 --- a/docs/api/ai-client.md +++ b/docs/api/ai-client.md @@ -176,6 +176,8 @@ export async function POST(request: Request) { Appends a message to the conversation. If you pass a `UIMessage`, `append` copies `uiMessage.metadata` onto the stored message. +`append()` resolves after the full HTTP response is processed when this call starts the stream. If a stream is already in progress, this call queues the send and the returned promise can resolve before that queued response is processed. A `RUN_FINISHED` with `finishReason: "tool_calls"` does not end the wait when the agent loop continues in that response. + ```typescript import { client } from "./client"; import type { UIMessage } from "@tanstack/ai-client"; diff --git a/packages/ai-client/src/chat-client.ts b/packages/ai-client/src/chat-client.ts index 0d320cdd84..fc0299c69c 100644 --- a/packages/ai-client/src/chat-client.ts +++ b/packages/ai-client/src/chat-client.ts @@ -164,6 +164,20 @@ function resolveTransport(transport: { throw new Error('ChatClient: either `connection` or `fetcher` is required.') } +function connectionDrainsOnSend(connection: ConnectionAdapter): boolean { + return 'connect' in connection +} + +function isIntermediateToolTurn(chunk: StreamChunk): boolean { + if (chunk.type !== 'RUN_FINISHED') return false + if (chunk.outcome?.type === 'interrupt') return false + const extra = chunk as StreamChunk & { finishReason?: unknown } + if (extra.finishReason !== undefined) { + return extra.finishReason === 'tool_calls' + } + return tanstackMetadata(chunk)?.finishReason === 'tool_calls' +} + export interface NormalizedQueueConfig { whenBusy: WhenBusy drain: 'fifo' | 'batch' @@ -416,6 +430,12 @@ export class ChatClient< private continuationPending = false private subscriptionAbortController: AbortController | null = null private processingResolve: (() => void) | null = null + /** + * `connect()` adapters push the full HTTP body into the subscribe queue, then + * wait until that queue is idle. After `send()` returns, every chunk from this + * request has been processed. Subscribe/send sockets do not drain that way. + */ + private connectionDrainsOnSend = false private errorReportedGeneration: number | null = null private streamGeneration = 0 private continuationGeneration = 0 @@ -518,7 +538,9 @@ export class ChatClient< this.byokProvider = options.byokProvider this.context = options.context this.queueConfig = normalizeQueueOption(options.queue) - this.connection = normalizeConnectionAdapter(resolveTransport(options)) + const transport = resolveTransport(options) + this.connectionDrainsOnSend = connectionDrainsOnSend(transport) + this.connection = normalizeConnectionAdapter(transport) // Build client tools map this.clientToolsRef = { current: new Map() } @@ -1140,7 +1162,9 @@ export class ChatClient< this.clearedStreamTracker.onSessionRunError() } this.setSessionGenerating(this.activeRunIds.size > 0) - if (options?.resolveProcessing !== false) { + const skipProcessingResolve = + chunk.type === 'RUN_FINISHED' && isIntermediateToolTurn(chunk) + if (options?.resolveProcessing !== false && !skipProcessingResolve) { this.resolveProcessing() } } @@ -2344,6 +2368,14 @@ export class ChatClient< return false } + // connect() send() already waited until the subscribe queue was idle. + // Kick the processing wait so a stream that ends on tool_calls (no + // interrupt / stop) cannot hang. Subscribe/send sockets still wait for + // a request-ending terminal below. + if (this.connectionDrainsOnSend) { + this.resolveProcessing() + } + // Wait for subscription loop to finish processing all chunks await processingComplete @@ -3045,12 +3077,12 @@ export class ChatClient< this.resetSessionGenerating() this.setIsSubscribed(false) this.setConnectionStatus('disconnected') - this.connection = normalizeConnectionAdapter( - resolveTransport({ - connection: options.connection, - fetcher: options.fetcher, - }), - ) + const transport = resolveTransport({ + connection: options.connection, + fetcher: options.fetcher, + }) + this.connectionDrainsOnSend = connectionDrainsOnSend(transport) + this.connection = normalizeConnectionAdapter(transport) if (wasSubscribed) { this.subscribe() diff --git a/packages/ai-client/src/connection-adapters.ts b/packages/ai-client/src/connection-adapters.ts index 63541ebe0c..6f6c30a658 100644 --- a/packages/ai-client/src/connection-adapters.ts +++ b/packages/ai-client/src/connection-adapters.ts @@ -1055,6 +1055,28 @@ export function normalizeConnectionAdapter( } } + async function waitUntilSubscriberIdle( + abortSignal?: AbortSignal, + ): Promise { + // Idle means the subscriber is waiting for the next chunk, so the + // previous chunk has left processIncomingChunk. Empty waiters with an + // empty buffer is in-flight delivery, not idle. + const idle = () => + activeBuffer.length === 0 && + (activeWaiters.length > 0 || abortSignal?.aborted) + for (let i = 0; i < 16 && !abortSignal?.aborted; i++) { + if (idle()) return + await Promise.resolve() + } + let macrotaskWaits = 0 + while (!abortSignal?.aborted) { + if (idle()) return + await new Promise((resolve) => setTimeout(resolve, 0)) + macrotaskWaits++ + if (activeWaiters.length === 0 && macrotaskWaits >= 32) return + } + } + return { subscribe(abortSignal?: AbortSignal): AsyncIterable { // Transfer ownership to the latest subscriber so only one active @@ -1162,6 +1184,7 @@ export function normalizeConnectionAdapter( } throw err } + await waitUntilSubscriberIdle(abortSignal) }, // Expose joinRun only when the underlying connection is resumable. Require // a real function — `'joinRun' in connection` is true for diff --git a/packages/ai-client/tests/chat-client.test.ts b/packages/ai-client/tests/chat-client.test.ts index 67429010b6..44a8ce4ce0 100644 --- a/packages/ai-client/tests/chat-client.test.ts +++ b/packages/ai-client/tests/chat-client.test.ts @@ -2306,6 +2306,113 @@ describe('ChatClient', () => { expect(messages[0]?.id).toBeTruthy() expect(messages[0]?.createdAt).toBeInstanceOf(Date) }) + + it('keeps append pending through an intermediate tool_calls RUN_FINISHED until the interrupt', async () => { + const adapter: ConnectConnectionAdapter = { + async *connect(_messages, _data, _signal, ctx) { + const runId = ctx?.runId ?? 'run-1' + const threadId = ctx?.threadId ?? 'thread-1' + yield { + type: EventType.RUN_STARTED, + runId, + threadId, + timestamp: Date.now(), + } + yield { + type: EventType.RUN_FINISHED, + runId, + threadId, + timestamp: Date.now(), + metadata: { tanstack: { finishReason: 'tool_calls' } }, + } + yield { + type: EventType.RUN_STARTED, + runId: 'provider-2', + threadId, + timestamp: Date.now(), + } + yield { + type: EventType.RUN_FINISHED, + runId: 'provider-2', + threadId, + timestamp: Date.now(), + outcome: { + type: 'interrupt', + interrupts: [{ id: 'interrupt-1', reason: 'client_tool_input' }], + }, + } + }, + } + const client = new ChatClient({ + connection: adapter, + threadId: 'thread-1', + }) + + await client.append({ + role: 'user', + content: 'Notify me', + }) + + expect(client.getPendingInterrupts()).toEqual([ + expect.objectContaining({ id: 'interrupt-1' }), + ]) + expect(client.getResumeState()?.runId).toBeTruthy() + }) + + it('keeps append pending when intermediate tool_calls is a direct finishReason', async () => { + const adapter: ConnectConnectionAdapter = { + async *connect(_messages, _data, _signal, ctx) { + const runId = ctx?.runId ?? 'run-1' + const threadId = ctx?.threadId ?? 'thread-1' + yield { + type: EventType.RUN_STARTED, + runId, + threadId, + timestamp: Date.now(), + } + yield { + type: EventType.RUN_FINISHED, + runId, + threadId, + timestamp: Date.now(), + finishReason: 'tool_calls', + } + await Promise.resolve() + yield { + type: EventType.RUN_STARTED, + runId: 'provider-2', + threadId, + timestamp: Date.now(), + } + yield { + type: EventType.RUN_FINISHED, + runId: 'provider-2', + threadId, + timestamp: Date.now(), + outcome: { + type: 'interrupt', + interrupts: [ + { id: 'interrupt-direct', reason: 'client_tool_input' }, + ], + }, + } + }, + } + const client = new ChatClient({ + connection: adapter, + threadId: 'thread-1', + }) + + await client.append({ + role: 'user', + content: 'Notify me', + }) + + expect(client.getPendingInterrupts()).toEqual([ + expect.objectContaining({ id: 'interrupt-direct' }), + ]) + expect(client.getResumeState()?.runId).toBeTruthy() + }) }) describe('reload', () => { diff --git a/packages/ai/src/activities/chat/stream/processor.ts b/packages/ai/src/activities/chat/stream/processor.ts index 8b6f5bc371..decc142833 100644 --- a/packages/ai/src/activities/chat/stream/processor.ts +++ b/packages/ai/src/activities/chat/stream/processor.ts @@ -212,6 +212,7 @@ export class StreamProcessor { private finishReason: string | null = null private hasError = false private isDone = false + private streamEndEmitted = false // Recording private recording: ChunkRecording | null = null @@ -729,9 +730,26 @@ export class StreamProcessor { return id } } + // finalizeStream() clears activeMessageIds but keeps messageStates. + // Leftover reasoning after an early RUN_FINISHED must resume that + // assistant. A new user turn calls prepareAssistantMessage(), which + // clears messageStates first. + for (const [id, state] of [...this.messageStates].reverse()) { + if (state.role === 'assistant') { + return id + } + } return null } + private resumeAssistantState(id: string, state: MessageStreamState): void { + this.activeMessageIds.add(id) + if (state.isComplete || this.isDone) { + state.isComplete = false + this.isDone = false + } + } + /** * Ensure an active assistant message exists, creating one if needed. * Used for backward compat when events arrive without prior TEXT_MESSAGE_START. @@ -748,14 +766,20 @@ export class StreamProcessor { // Try to find state by preferred ID if (preferredId) { const state = this.getMessageState(preferredId) - if (state) return { messageId: preferredId, state } + if (state) { + this.resumeAssistantState(preferredId, state) + return { messageId: preferredId, state } + } } // Try active assistant message const activeId = this.getActiveAssistantMessageId() if (activeId) { const state = this.getMessageState(activeId) - if (state) return { messageId: activeId, state } + if (state) { + this.resumeAssistantState(activeId, state) + return { messageId: activeId, state } + } } // Check if a message with preferredId already exists (reconnect/resume case). @@ -1647,8 +1671,14 @@ export class StreamProcessor { } if (this.activeRuns.size === 0) { - this.isDone = true this.completeAllToolCalls() + const isIntermediateToolTurn = + this.finishReason === 'tool_calls' && + chunk.outcome?.type !== 'interrupt' + if (isIntermediateToolTurn) { + return + } + this.isDone = true this.finalizeStream() } } @@ -2344,6 +2374,7 @@ export class StreamProcessor { * @see docs/chat-architecture.md#single-shot-text-response — Finalization step */ finalizeStream(): void { + this.isDone = true let lastAssistantMessage: UIMessage | undefined // Finalize ALL active messages @@ -2407,7 +2438,8 @@ export class StreamProcessor { } // Emit stream end for the last assistant message - if (lastAssistantMessage) { + if (lastAssistantMessage && !this.streamEndEmitted) { + this.streamEndEmitted = true this.events.onStreamEnd?.(lastAssistantMessage) } } @@ -2526,6 +2558,7 @@ export class StreamProcessor { this.finishReason = null this.hasError = false this.isDone = false + this.streamEndEmitted = false this.chunkStrategy.reset?.() } diff --git a/packages/ai/tests/stream-processor.test.ts b/packages/ai/tests/stream-processor.test.ts index 61a37de65b..c9b6a3057f 100644 --- a/packages/ai/tests/stream-processor.test.ts +++ b/packages/ai/tests/stream-processor.test.ts @@ -2976,7 +2976,9 @@ describe('StreamProcessor', () => { expect(state.toolCalls.size).toBe(1) expect(state.toolCallOrder).toEqual(['tc-1']) expect(state.finishReason).toBe('tool_calls') - expect(state.done).toBe(true) + expect(state.done).toBe(false) + processor.finalizeStream() + expect(processor.getState().done).toBe(true) }) it('should return independent copies (mutations do not affect internal state)', () => { @@ -4603,6 +4605,84 @@ describe('StreamProcessor', () => { expect(processor.getState().done).toBe(true) }) + it('appends leftover reasoning after a stop RUN_FINISHED onto the same thinking part', () => { + const processor = new StreamProcessor() + processor.prepareAssistantMessage() + + processor.processChunk(ev.runStarted()) + processor.processChunk(ev.stepStarted('step-1')) + processor.processChunk( + ev.reasoningContent( + 'The user is asking for a beginner guitar recommen.', + ), + ) + processor.processChunk(ev.runFinished('stop')) + processor.processChunk(ev.reasoningContent(' $1,299.')) + + const messages = processor.getMessages() + expect(messages).toHaveLength(1) + const thinkingParts = messages[0]!.parts.filter( + (part) => part.type === 'thinking', + ) + expect(thinkingParts).toHaveLength(1) + const thinkingPart = thinkingParts[0] + if (thinkingPart?.type !== 'thinking') { + throw new Error('expected a thinking part') + } + expect(thinkingPart.content).toBe( + 'The user is asking for a beginner guitar recommen. $1,299.', + ) + }) + + it('flushes leftover text after a stop RUN_FINISHED and fires onStreamEnd once', () => { + const events = spyEvents() + const processor = new StreamProcessor({ events }) + processor.prepareAssistantMessage() + + processor.processChunk(ev.runStarted()) + processor.processChunk(ev.textStart('msg-1')) + processor.processChunk(ev.textContent('Hello', 'msg-1')) + processor.processChunk(ev.runFinished('stop')) + + expect(events.onStreamEnd).toHaveBeenCalledTimes(1) + expect(processor.getState().done).toBe(true) + + processor.processChunk(ev.textContent(' world', 'msg-1')) + expect(processor.getState().done).toBe(false) + processor.processChunk(ev.textEnd('msg-1')) + processor.finalizeStream() + + const textPart = processor + .getMessages()[0]! + .parts.find((part) => part.type === 'text') + if (textPart?.type !== 'text') { + throw new Error('expected a text part') + } + expect(textPart.content).toBe('Hello world') + expect(events.onStreamEnd).toHaveBeenCalledTimes(1) + expect(processor.getState().done).toBe(true) + }) + + it('does not fire onStreamEnd on a sequential tool_calls terminal', () => { + const events = spyEvents() + const processor = new StreamProcessor({ events }) + + processor.processChunk(ev.runStarted('run-1')) + processor.processChunk(ev.textStart('msg-1')) + processor.processChunk(ev.textContent('calling', 'msg-1')) + processor.processChunk(ev.runFinished('tool_calls', 'run-1')) + + expect(events.onStreamEnd).not.toHaveBeenCalled() + expect(processor.getState().done).toBe(false) + + processor.processChunk(ev.runStarted('run-2')) + processor.processChunk(ev.textContent(' done', 'msg-1')) + processor.processChunk(ev.runFinished('stop', 'run-2')) + + expect(events.onStreamEnd).toHaveBeenCalledTimes(1) + expect(processor.getState().done).toBe(true) + }) + it('single run should finalize normally (backward compat)', () => { const events = spyEvents() const processor = new StreamProcessor({ events }) @@ -4656,6 +4736,8 @@ describe('StreamProcessor', () => { expect(processor.getState().toolCalls.get('tc-a')?.state).toBe( 'input-complete', ) + expect(processor.getState().done).toBe(false) + processor.finalizeStream() expect(processor.getState().done).toBe(true) }) diff --git a/testing/e2e/src/lib/tools-test-tools.ts b/testing/e2e/src/lib/tools-test-tools.ts index 8cbc8c381b..91a83f5a6a 100644 --- a/testing/e2e/src/lib/tools-test-tools.ts +++ b/testing/e2e/src/lib/tools-test-tools.ts @@ -253,6 +253,11 @@ export const SCENARIO_LIST = [ label: 'Client Tool Input Error', category: 'basic', }, + { + id: 'invalid-client-tool-retry', + label: 'Invalid Client Tool Retry (Regression #1192)', + category: 'race', + }, // Race condition / event flow scenarios { id: 'sequential-client-tools', @@ -312,6 +317,7 @@ export function getToolsForScenario(scenario: string) { case 'client-tool-reasoning': case 'client-tool-stop': case 'client-tool-input-error': + case 'invalid-client-tool-retry': return [clientToolDefinitions.show_notification] case 'server-context': diff --git a/testing/e2e/src/routes/api.tools-test.ts b/testing/e2e/src/routes/api.tools-test.ts index b6d1e9d446..5b1daaae6d 100644 --- a/testing/e2e/src/routes/api.tools-test.ts +++ b/testing/e2e/src/routes/api.tools-test.ts @@ -20,6 +20,7 @@ const providerFreeScenarios = new Set([ 'client-server-context', 'client-tool-stop', 'client-tool-input-error', + 'invalid-client-tool-retry', 'malformed-tool-arguments', 'provider-rejected-tool-call', ]) @@ -49,13 +50,20 @@ function createProviderFreeAdapter(scenario: string): AnyTextAdapter { state: undefined, toolName: 'check_status', } - : scenario === 'client-tool-input-error' + : scenario === 'client-tool-input-error' || + scenario === 'invalid-client-tool-retry' ? { arguments: '{"message":42,"type":"info"}', initialText: 'Showing a notification.', input: { message: 42, type: 'info' }, - name: 'client-tool-input-error-test', - responseText: 'Unexpected client continuation.', + name: + scenario === 'invalid-client-tool-retry' + ? 'invalid-client-tool-retry-test' + : 'client-tool-input-error-test', + responseText: + scenario === 'invalid-client-tool-retry' + ? 'Recovered after client tool input retry.' + : 'Unexpected client continuation.', result: undefined, state: undefined, toolName: 'show_notification', @@ -104,9 +112,12 @@ function createProviderFreeAdapter(scenario: string): AnyTextAdapter { const runId = options.runId ?? 'runtime-context-run' const threadId = options.threadId ?? 'runtime-context-thread' const messageId = `${runId}-message` - const hasToolResult = options.messages.some( + const toolResultCount = options.messages.filter( (message) => message.role === 'tool', - ) + ).length + const hasToolResult = toolResultCount > 0 + const retryClientTool = + scenario === 'invalid-client-tool-retry' && toolResultCount === 1 yield { type: EventType.RUN_STARTED, @@ -116,8 +127,17 @@ function createProviderFreeAdapter(scenario: string): AnyTextAdapter { timestamp: Date.now(), } - if (!hasToolResult) { - const toolCallId = `${scenario}-tool-call` + if (!hasToolResult || retryClientTool) { + const toolCallId = + scenario === 'invalid-client-tool-retry' + ? `${scenario}-tool-call-${toolResultCount + 1}` + : `${scenario}-tool-call` + const toolArguments = retryClientTool + ? '{"message":"done","type":"info"}' + : config.arguments + const toolInput = retryClientTool + ? { message: 'done', type: 'info' } + : config.input yield { type: EventType.TEXT_MESSAGE_START, @@ -150,7 +170,7 @@ function createProviderFreeAdapter(scenario: string): AnyTextAdapter { yield { type: EventType.TOOL_CALL_ARGS, toolCallId, - delta: config.arguments, + delta: toolArguments, model, timestamp: Date.now(), } @@ -159,7 +179,7 @@ function createProviderFreeAdapter(scenario: string): AnyTextAdapter { toolCallId, toolCallName: config.toolName, toolName: config.toolName, - ...(config.input === undefined ? {} : { input: config.input }), + ...(toolInput === undefined ? {} : { input: toolInput }), ...(config.result === undefined ? {} : { result: config.result, state: config.state }), diff --git a/testing/e2e/tests/tools-test/client-tool.spec.ts b/testing/e2e/tests/tools-test/client-tool.spec.ts index 5d47f98ad5..7734c43dec 100644 --- a/testing/e2e/tests/tools-test/client-tool.spec.ts +++ b/testing/e2e/tests/tools-test/client-tool.spec.ts @@ -20,6 +20,48 @@ import { */ test.describe('Client Tool E2E Tests', () => { + test('invalid client-tool input retries then runs the interrupt', async ({ + page, + testId, + aimockPort, + }) => { + const requests: Array = [] + page.on('request', (request) => { + if ( + request.method() === 'POST' && + request.url().includes('/api/tools-test') + ) { + requests.push(request.url()) + } + }) + + await selectScenario(page, 'invalid-client-tool-retry', testId, aimockPort) + await runTest(page) + await waitForTestComplete(page, 15000, 2) + + const metadata = await getMetadata(page) + expect(requests).toHaveLength(2) + expect(metadata.executionCompleteCount).toBe('1') + + const toolCalls = await getToolCalls(page) + expect(toolCalls).toEqual( + expect.arrayContaining([ + expect.objectContaining({ state: 'error' }), + expect.objectContaining({ + name: 'show_notification', + state: 'complete', + }), + ]), + ) + + const responseText = (await getMessages(page)) + .flatMap((message) => message.parts) + .filter((part) => part.type === 'text') + .map((part) => part.content) + .join(' ') + expect(responseText).toContain('Recovered after client tool input retry.') + }) + test('single client tool executes and completes', async ({ page, testId,