From adac1069237061ad024c8b8af2496e253c770e5e Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Mon, 2 Nov 2026 13:26:07 +0100 Subject: [PATCH] fix(ai-client): keep queued messages until the open interrupts are answered A message queued during a run went out as soon as that run ended, also when the run ended paused on an interrupt such as a tool approval. The new run then started with the tool call unanswered. drainQueue() now waits while interrupts are open. The run that answers them ends with success and drains the queue as before. --- .changeset/queue-waits-for-interrupts.md | 5 ++ docs/chat/queueing.md | 1 + docs/config.json | 2 +- packages/ai-client/src/chat-client.ts | 12 ++- .../chat-client-approval-followup.test.ts | 80 +++++++++++++++++++ .../chat-client-foreign-client-tool.test.ts | 41 ++++++++++ .../e2e/src/routes/foreign-client-tool.tsx | 5 +- testing/e2e/tests/foreign-client-tool.spec.ts | 28 +++++++ 8 files changed, 171 insertions(+), 3 deletions(-) create mode 100644 .changeset/queue-waits-for-interrupts.md 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') + }) })