diff --git a/.changeset/queue-waits-for-interrupts.md b/.changeset/queue-waits-for-interrupts.md new file mode 100644 index 0000000000..6744236f2b --- /dev/null +++ b/.changeset/queue-waits-for-interrupts.md @@ -0,0 +1,5 @@ +--- +'@tanstack/ai-client': patch +--- + +Keep queued messages in the queue while a run is paused on an interrupt, for example a tool approval. Before, the queue sent the next message as soon as the paused run ended, so a new run started while the tool call was still unanswered. Now the queue sends after the run that answers the interrupt ends with success. diff --git a/docs/chat/queueing.md b/docs/chat/queueing.md index 3d28d6d38a..7dda75253d 100644 --- a/docs/chat/queueing.md +++ b/docs/chat/queueing.md @@ -174,5 +174,6 @@ export function useQueuedChat() { ### When the queue sends and when it empties - **Send.** The queue sends its items only after a stream ends with success. This includes the end of tool continuations. +- **Wait for a paused run.** A run can end paused on an interrupt, for example a tool approval. Then the queue waits. It sends after the run that answers the interrupt ends with success. - **Empty without sending.** The queue empties when the active run fails or stops (`stop()` or a real stream error). It also empties on `clear()`, `unsubscribe()`, and `reload()`. - **`interrupt` keeps the queue.** Items that already wait stay in the queue. They send after the interrupting turn succeeds. diff --git a/docs/config.json b/docs/config.json index b093c507f4..d1855bbf8b 100644 --- a/docs/config.json +++ b/docs/config.json @@ -224,7 +224,7 @@ "label": "Message Queue", "to": "chat/queueing", "addedAt": "2026-08-26", - "updatedAt": "2026-10-22" + "updatedAt": "2026-11-02" } ], "tab": "guides" diff --git a/packages/ai-client/src/chat-client.ts b/packages/ai-client/src/chat-client.ts index e413cd8db6..368dd7373f 100644 --- a/packages/ai-client/src/chat-client.ts +++ b/packages/ai-client/src/chat-client.ts @@ -3580,7 +3580,16 @@ export class ChatClient< * by control-flow analysis. */ private shouldAbortMessageQueueDrain(): boolean { - return this.isLoading || this.stopMessageQueueDrain + return ( + this.isLoading || this.stopMessageQueueDrain || this.hasOpenInterrupts() + ) + } + + // A paused run waits for its interrupts. A queued message sent now would + // start a new run with the paused tool calls unanswered. The run that + // answers them drains the queue when it ends. + private hasOpenInterrupts(): boolean { + return this.interruptManager.getInterrupts().length > 0 } /** @@ -3600,6 +3609,7 @@ export class ChatClient< if ( this.messageQueueDraining || this.isLoading || + this.hasOpenInterrupts() || this.messageQueue.length === 0 ) { return diff --git a/packages/ai-client/tests/chat-client-approval-followup.test.ts b/packages/ai-client/tests/chat-client-approval-followup.test.ts index 7ed195e5b8..86a569b35e 100644 --- a/packages/ai-client/tests/chat-client-approval-followup.test.ts +++ b/packages/ai-client/tests/chat-client-approval-followup.test.ts @@ -198,4 +198,84 @@ describe('tool-approval follow-up after resolve', () => { .join('\n') expect(followUpText).toContain('follow-up') }) + + it('keeps a queued message until the server approval is answered', async () => { + const approval = normalizeApprovalSchema(undefined, cartTool.inputSchema) + const input = { guitarId: '1', quantity: 1 } + const requests: Array = [] + const connection: ConnectConnectionAdapter = { + // eslint-disable-next-line @typescript-eslint/require-await + async *connect(messages, _d, _s, runContext) { + requests.push(JSON.stringify(messages)) + const runId = runContext?.runId ?? `r${requests.length}` + const threadId = runContext?.threadId ?? 't1' + const t = Date.now() + yield { type: EventType.RUN_STARTED, runId, threadId, timestamp: t } + yield { + type: EventType.RUN_FINISHED, + runId, + threadId, + timestamp: t, + outcome: + requests.length === 1 + ? { + type: 'interrupt', + interrupts: [ + { + id: 'approval_tc1', + reason: 'tool_call', + toolCallId: 'tc1', + responseSchema: approval.responseSchema, + metadata: { + kind: 'approval', + toolName: 'addToCart', + input, + 'tanstack:interruptBinding': { + kind: 'tool-approval', + interruptId: 'approval_tc1', + interruptedRunId: runId, + generation: 0, + toolName: 'addToCart', + toolCallId: 'tc1', + originalArgs: input, + inputSchemaHash: hashSchemaInput( + cartTool.inputSchema, + ), + approvalSchemaHash: approval.approvalSchemaHash, + responseSchemaHash: approval.responseSchemaHash, + }, + }, + }, + ], + } + : { type: 'success' }, + } as StreamChunk + }, + } + const client = new ChatClient({ + connection, + threadId: 't1', + tools: [cartTool], + }) + + // The second message waits in the queue while the first run streams. + await Promise.all([ + client.sendMessage('add the stratocaster'), + client.sendMessage('queued'), + ]) + await new Promise((resolve) => setTimeout(resolve, 20)) + + const interrupt = client.getInterrupts()[0] + if (interrupt?.kind !== 'tool-approval') { + throw new Error('expected tool-approval') + } + expect(requests).toHaveLength(1) + + interrupt.resolveInterrupt(true) + + // The resume run goes first. The queued message follows it. + await vi.waitFor(() => expect(requests).toHaveLength(3)) + expect(requests[1]).not.toContain('queued') + expect(requests[2]).toContain('queued') + }) }) 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 fa003021b4..00ca2dde61 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 @@ -178,4 +178,45 @@ describe('client tools from a foreign AG-UI server', () => { }, ) }) + + it('keeps a queued message until the approval is answered', async () => { + const execute = vi.fn(() => ({ reply: 'pong' })) + const ping = toolDefinition({ + name: 'ping', + description: 'Reply with pong', + inputSchema: z.object({}), + needsApproval: true, + }).client(execute) + const { connection, requests } = foreignServer() + 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 new Promise((resolve) => setTimeout(resolve, 20)) + + const interrupt = client.getInterrupts()[0] + if (interrupt?.kind !== 'tool-approval') { + throw new Error('expected a tool-approval interrupt') + } + expect(requests).toHaveLength(1) + + interrupt.resolveInterrupt(true) + + 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\\"}') + }) }) diff --git a/testing/e2e/src/routes/foreign-client-tool.tsx b/testing/e2e/src/routes/foreign-client-tool.tsx index e3553d5500..ab5e011c63 100644 --- a/testing/e2e/src/routes/foreign-client-tool.tsx +++ b/testing/e2e/src/routes/foreign-client-tool.tsx @@ -38,11 +38,13 @@ export const Route = createFileRoute('/foreign-client-tool')({ // The router writes the result back to the URL (`approval=true`) and // checks it again, so `true` must stay `true`. approval: search.approval === true || search.approval === 1, + // `?queue=1` queues a second message while the first run streams. + queue: search.queue === true || search.queue === 1, }), }) function ForeignClientToolPage() { - const { finish, approval } = Route.useSearch() + const { finish, approval, queue } = Route.useSearch() const { messages, sendMessage, interrupts } = useChat({ threadId: 'foreign-client-tool', connection: fetchServerSentEvents( @@ -70,6 +72,7 @@ function ForeignClientToolPage() { useEffect(() => { void sendMessage('ping me') + if (queue) void sendMessage('queued') // Fire the single run once on mount; the harness route ignores the content. // eslint-disable-next-line react-hooks/exhaustive-deps }, []) diff --git a/testing/e2e/tests/foreign-client-tool.spec.ts b/testing/e2e/tests/foreign-client-tool.spec.ts index bbe4458456..f80b0468d9 100644 --- a/testing/e2e/tests/foreign-client-tool.spec.ts +++ b/testing/e2e/tests/foreign-client-tool.spec.ts @@ -54,4 +54,32 @@ test.describe('frontend tool from a spec-compliant AG-UI server', () => { }) } } + + test('keeps a queued message until the approval is answered', async ({ + page, + }) => { + const posts: Array = [] + page.on('request', (request) => { + if ( + request.url().includes('/api/foreign-client-tool') && + request.method() === 'POST' + ) { + posts.push(request.postData() ?? '') + } + }) + + await page.goto('/foreign-client-tool?approval=1&queue=1') + const approve = page + .getByTestId('approval') + .getByRole('button', { name: 'Approve' }) + await expect(approve).toBeVisible() + expect(posts).toHaveLength(1) + + await approve.click() + + // The run with the tool result goes first. The queued message follows. + await expect.poll(() => posts.length).toBe(3) + expect(posts[1]).not.toContain('queued') + expect(posts[2]).toContain('queued') + }) })