Repository navigation
fix(ai): run onAbort when a chat() consumer stops iterating #2098
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,5 @@ | ||
| --- | ||
| '@tanstack/ai': patch | ||
| --- | ||
|
|
||
| A `chat()` consumer that stops iterating (`break` or `return()`) now runs the middleware `onAbort` hook, so every run gets exactly one terminal hook. Stopping at a client-tool or approval pause's `RUN_FINISHED` chunk runs `onFinish`, as a drained pause does, so the paused state (such as a `withSandbox` sandbox) is kept. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,153 @@ | ||
| import { describe, expect, it } from 'vitest' | ||
| import { chat } from '../src/activities/chat/index' | ||
| import { fakeText } from '../src/testing' | ||
| import { EventType } from '../src/types' | ||
| import { clientTool, createMockAdapter, ev, serverTool } from './test-utils' | ||
| import type { ChatMiddleware } from '../src/activities/chat/middleware/types' | ||
| import type { StreamChunk, Tool } from '../src/types' | ||
|
|
||
| function terminalRecorder() { | ||
| const hooks: Array<string> = [] | ||
| const middleware: ChatMiddleware = { | ||
| name: 'record-terminal', | ||
| onFinish: () => { | ||
| hooks.push('onFinish') | ||
| }, | ||
| onAbort: () => { | ||
| hooks.push('onAbort') | ||
| }, | ||
| onError: () => { | ||
| hooks.push('onError') | ||
| }, | ||
| } | ||
| return { hooks, middleware } | ||
| } | ||
|
|
||
| function run(middleware: ChatMiddleware) { | ||
| const adapter = fakeText() | ||
| adapter.setResponses([{ text: 'a long answer in many chunks' }]) | ||
| return chat({ | ||
| threadId: 'thread-1', | ||
| adapter, | ||
| messages: [{ role: 'user', content: 'Hi' }], | ||
| middleware: [middleware], | ||
| }) as AsyncIterable<StreamChunk> | ||
| } | ||
|
|
||
| function runWithTool(middleware: ChatMiddleware, tool: Tool) { | ||
| const { adapter, calls } = createMockAdapter({ | ||
| iterations: [ | ||
| [ | ||
| ev.runStarted(), | ||
| ev.toolStart('call-1', tool.name), | ||
| ev.toolArgs('call-1', '{}'), | ||
| ev.runFinished('tool_calls'), | ||
| ], | ||
| [ | ||
| ev.runStarted('run-2'), | ||
| ev.textStart(), | ||
| ev.textContent('done'), | ||
| ev.textEnd(), | ||
| ev.runFinished('stop', 'run-2'), | ||
| ], | ||
| ], | ||
| }) | ||
| const stream = chat({ | ||
| threadId: 'thread-1', | ||
| adapter, | ||
| messages: [{ role: 'user', content: 'Use the tool' }], | ||
| tools: [tool], | ||
| middleware: [middleware], | ||
| }) as AsyncIterable<StreamChunk> | ||
| return { stream, calls } | ||
| } | ||
|
|
||
| describe('a consumer that stops iterating chat() runs one terminal hook', () => { | ||
| it('break after the first TEXT_MESSAGE_CONTENT runs onAbort', async () => { | ||
| const { hooks, middleware } = terminalRecorder() | ||
| for await (const chunk of run(middleware)) { | ||
| if (chunk.type === EventType.TEXT_MESSAGE_CONTENT) break | ||
| } | ||
| expect(hooks).toEqual(['onAbort']) | ||
| }) | ||
|
|
||
| it('break after the final RUN_FINISHED runs onAbort', async () => { | ||
| const { hooks, middleware } = terminalRecorder() | ||
| for await (const chunk of run(middleware)) { | ||
| if (chunk.type === EventType.RUN_FINISHED) break | ||
| } | ||
| expect(hooks).toEqual(['onAbort']) | ||
| }) | ||
|
|
||
| it('break after a tool turn RUN_FINISHED runs onAbort and no next turn', async () => { | ||
| const { hooks, middleware } = terminalRecorder() | ||
| const { stream, calls } = runWithTool( | ||
| middleware, | ||
| serverTool('lookup', () => 'found'), | ||
| ) | ||
| for await (const chunk of stream) { | ||
| if (chunk.type === EventType.RUN_FINISHED) break | ||
| } | ||
| expect(hooks).toEqual(['onAbort']) | ||
| expect(calls).toHaveLength(1) | ||
| }) | ||
|
|
||
| it.each([ | ||
| ['a client tool', clientTool('ask_client')], | ||
| [ | ||
| 'an approval', | ||
| { ...serverTool('delete_file', () => 'deleted'), needsApproval: true }, | ||
| ], | ||
| ])('break after %s pause RUN_FINISHED runs onFinish', async (_kind, tool) => { | ||
| const { hooks, middleware } = terminalRecorder() | ||
| const { stream, calls } = runWithTool(middleware, tool) | ||
| let stoppedAt: StreamChunk | undefined | ||
| for await (const chunk of stream) { | ||
| if (chunk.type === EventType.RUN_FINISHED) { | ||
| stoppedAt = chunk | ||
| break | ||
| } | ||
| } | ||
| expect(stoppedAt).toMatchObject({ outcome: { type: 'interrupt' } }) | ||
| expect(hooks).toEqual(['onFinish']) | ||
| expect(calls).toHaveLength(1) | ||
| }) | ||
|
|
||
| it('break at a custom event from setup runs onAbort after setup ends', async () => { | ||
| const events: Array<string> = [] | ||
| const middleware: ChatMiddleware = { | ||
| name: 'slow-setup', | ||
| async setup(ctx) { | ||
| ctx.emitCustomEvent('setup:progress', {}) | ||
| await new Promise((resolve) => setTimeout(resolve, 20)) | ||
| events.push('setup done') | ||
| }, | ||
| onAbort: () => { | ||
| events.push('onAbort') | ||
| }, | ||
| } | ||
| for await (const chunk of run(middleware)) { | ||
| if (chunk.type === EventType.CUSTOM) break | ||
| } | ||
| expect(events).toEqual(['setup done', 'onAbort']) | ||
| }) | ||
|
|
||
| it('return() after the first TEXT_MESSAGE_CONTENT runs onAbort', async () => { | ||
| const { hooks, middleware } = terminalRecorder() | ||
| const iterator = run(middleware)[Symbol.asyncIterator]() | ||
| let next = await iterator.next() | ||
| while (!next.done && next.value.type !== EventType.TEXT_MESSAGE_CONTENT) { | ||
| next = await iterator.next() | ||
| } | ||
| await iterator.return?.() | ||
| expect(hooks).toEqual(['onAbort']) | ||
| }) | ||
|
|
||
| it('a full drain runs onFinish once', async () => { | ||
| const { hooks, middleware } = terminalRecorder() | ||
| for await (const _ of run(middleware)) { | ||
| // drain | ||
| } | ||
| expect(hooks).toEqual(['onFinish']) | ||
| }) | ||
| }) |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,43 @@ | ||
| import { createFileRoute } from '@tanstack/react-router' | ||
| import { chat } from '@tanstack/ai' | ||
| import { createTextAdapter } from '@/lib/providers' | ||
| import type { Provider } from '@/lib/types' | ||
|
|
||
| /** | ||
| * Iterates `chat()` and stops with `break` at the first text chunk, with no | ||
| * abort. Returns the terminal middleware hooks that ran. | ||
| */ | ||
| export const Route = createFileRoute('/api/chat-consumer-stop')({ | ||
| server: { | ||
| handlers: { | ||
| POST: async ({ request }) => { | ||
| const url = new URL(request.url) | ||
| const provider = url.searchParams.get('provider') as Provider | ||
| const testId = url.searchParams.get('testId') ?? undefined | ||
| const hooks: Array<string> = [] | ||
| for await (const event of chat({ | ||
| ...createTextAdapter(provider, undefined, undefined, testId), | ||
| threadId: 'chat-consumer-stop', | ||
| messages: [{ role: 'user', content: '[chat] recommend a guitar' }], | ||
| middleware: [ | ||
| { | ||
| name: 'record-terminal-hooks', | ||
| onFinish: () => { | ||
| hooks.push('onFinish') | ||
| }, | ||
| onAbort: () => { | ||
| hooks.push('onAbort') | ||
| }, | ||
| onError: () => { | ||
| hooks.push('onError') | ||
| }, | ||
| }, | ||
| ], | ||
| })) { | ||
| if (event.type === 'TEXT_MESSAGE_CONTENT') break | ||
| } | ||
| return Response.json({ hooks }) | ||
| }, | ||
| }, | ||
| }, | ||
| }) |
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Uh oh!
There was an error while loading. Please reload this page.