Repository navigation
fix(ai): run onAbort when a chat() consumer stops iterating #2098
+166
−1
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. |
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,71 @@ | ||
| import { describe, expect, it } from 'vitest' | ||
| import { chat } from '../src/activities/chat/index' | ||
| import { fakeText } from '../src/testing' | ||
| import { EventType } from '../src/types' | ||
| import type { ChatMiddleware } from '../src/activities/chat/middleware/types' | ||
| import type { StreamChunk } 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> | ||
| } | ||
|
|
||
| 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 RUN_FINISHED runs one terminal hook', async () => { | ||
| const { hooks, middleware } = terminalRecorder() | ||
| for await (const chunk of run(middleware)) { | ||
| if (chunk.type === EventType.RUN_FINISHED) break | ||
| } | ||
| expect(hooks).toHaveLength(1) | ||
| }) | ||
|
|
||
| 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 }) | ||
| }, | ||
| }, | ||
| }, | ||
| }) |
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
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.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Script executed:
Repository: TanStack/ai
Length of output: 11932
Script executed:
Repository: TanStack/ai
Length of output: 18089
Script executed:
Repository: TanStack/ai
Length of output: 33067
Script executed:
Repository: TanStack/ai
Length of output: 23266
Clarify which hook runs at
RUN_FINISHED.For normal completion and client-tool or approval pauses, draining the stream lets
onFinishrun. Breaking on theRUN_FINISHEDchunk can instead runonAbort. Clarify this distinction so readers do not expect the pause alone to triggeronFinish.LanguageTool
[style] ~788-~788: Consider using “who” when you are referring to a person instead of an object.
Context: ...om
onBeforeToolCall, or by a consumer that stopped iterating (breakor `return()...(THAT_WHO)