Skip to content

Commit cb3b9c9

Browse files
lgrammellgrammel
authored
fix (ai): catch errors in ui message stream (#6898)
## Background Adding error parts to the ui message stream leads to processing failures. ## Summary * add `ErrorHandler` type * add default error handler with console logging to `Chat` * use error handler in stream processing ## Verification <!-- For features & bugfixes. Please explain how you *manually* verified that the change works end-to-end as expected (independent of automated tests). Remove the section if it's not needed (e.g. for docs). --> ## Tasks - [x] add test for handleUIMessageStreamFinish error case - [x] changeset ## Related Issues #6879
1 parent bbf3128 commit cb3b9c9

File tree

14 files changed

+147
-14
lines changed

14 files changed

+147
-14
lines changed

‎.changeset/spotty-countries-appear.md

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
---
2+
'ai': patch
3+
---
4+
5+
fix (ai): catch errors in ui message stream

‎examples/next-openai-kasada-bot-protection/app/page.tsx

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@ export default function Chat() {
1313
});
1414

1515
return (
16-
<div className="flex flex-col w-full max-w-md py-24 mx-auto stretch">
16+
<div className="flex flex-col py-24 mx-auto w-full max-w-md stretch">
1717
{messages.length > 0
1818
? messages.map(m => (
1919
<div key={m.id} className="whitespace-pre-wrap">
@@ -33,7 +33,7 @@ export default function Chat() {
3333
}}
3434
>
3535
<input
36-
className="fixed bottom-0 w-full max-w-md p-2 mb-8 border border-gray-300 rounded shadow-xl"
36+
className="fixed bottom-0 p-2 mb-8 w-full max-w-md rounded border border-gray-300 shadow-xl"
3737
value={input}
3838
placeholder="Say something..."
3939
onChange={e => setInput(e.target.value)}

‎examples/next-openai-upstash-rate-limits/app/page.tsx

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,7 @@ export default function Chat() {
1111
});
1212

1313
return (
14-
<div className="flex flex-col w-full max-w-md py-24 mx-auto stretch">
14+
<div className="flex flex-col py-24 mx-auto w-full max-w-md stretch">
1515
{messages.length > 0
1616
? messages.map(m => (
1717
<div key={m.id} className="whitespace-pre-wrap">
@@ -31,7 +31,7 @@ export default function Chat() {
3131
}}
3232
>
3333
<input
34-
className="fixed bottom-0 w-full max-w-md p-2 mb-8 border border-gray-300 rounded shadow-xl"
34+
className="fixed bottom-0 p-2 mb-8 w-full max-w-md rounded border border-gray-300 shadow-xl"
3535
value={input}
3636
placeholder="Say something..."
3737
onChange={e => setInput(e.target.value)}

‎packages/ai/core/generate-text/stream-text-result.ts

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@ import { StepResult } from './step-result';
2020
import { ToolCallUnion } from './tool-call';
2121
import { ToolErrorUnion, ToolResultUnion } from './tool-output';
2222
import { ToolSet } from './tool-set';
23+
import { ErrorHandler } from '../../src/util/error-handler';
2324

2425
export type UIMessageStreamOptions<UI_MESSAGE extends UIMessage> = {
2526
/**
@@ -98,7 +99,7 @@ export type UIMessageStreamOptions<UI_MESSAGE extends UIMessage> = {
9899
};
99100

100101
export type ConsumeStreamOptions = {
101-
onError?: (error: unknown) => void;
102+
onError?: ErrorHandler;
102103
};
103104

104105
/**

‎packages/ai/core/generate-text/stream-text.ts

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1803,6 +1803,7 @@ However, the LLM results are expected to be small enough to not cause issues.
18031803
messageId: responseMessageId ?? this.generateId(),
18041804
originalMessages,
18051805
onFinish,
1806+
onError,
18061807
});
18071808
}
18081809

‎packages/ai/src/ui-message-stream/create-ui-message-stream.ts

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -135,5 +135,6 @@ export function createUIMessageStream<UI_MESSAGE extends UIMessage>({
135135
messageId: generateId(),
136136
originalMessages,
137137
onFinish,
138+
onError,
138139
});
139140
}

‎packages/ai/src/ui-message-stream/handle-ui-message-stream-finish.ts

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@ import {
44
StreamingUIMessageState,
55
} from '../ui/process-ui-message-stream';
66
import { UIMessage } from '../ui/ui-messages';
7+
import { ErrorHandler } from '../util/error-handler';
78
import {
89
InferUIMessageStreamPart,
910
UIMessageStreamPart,
@@ -13,6 +14,7 @@ export function handleUIMessageStreamFinish<UI_MESSAGE extends UIMessage>({
1314
messageId,
1415
originalMessages = [],
1516
onFinish,
17+
onError,
1618
stream,
1719
}: {
1820
stream: ReadableStream<InferUIMessageStreamPart<UI_MESSAGE>>;
@@ -24,6 +26,8 @@ export function handleUIMessageStreamFinish<UI_MESSAGE extends UIMessage>({
2426
*/
2527
originalMessages?: UI_MESSAGE[];
2628

29+
onError: ErrorHandler;
30+
2731
onFinish?: (options: {
2832
/**
2933
* The updates list of UI messages.
@@ -87,6 +91,7 @@ export function handleUIMessageStreamFinish<UI_MESSAGE extends UIMessage>({
8791
}),
8892
),
8993
runUpdateMessageJob,
94+
onError,
9095
}).pipeThrough(
9196
new TransformStream<
9297
InferUIMessageStreamPart<UI_MESSAGE>,

‎packages/ai/src/ui-message-stream/ui-message-stream-writer.ts

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
import { UIMessage } from '../ui';
2+
import { ErrorHandler } from '../util/error-handler';
23
import { InferUIMessageStreamPart } from './ui-message-stream-parts';
34

45
export interface UIMessageStreamWriter<
@@ -19,5 +20,5 @@ export interface UIMessageStreamWriter<
1920
* This is intended for forwarding when merging streams
2021
* to prevent duplicated error masking.
2122
*/
22-
onError: ((error: unknown) => string) | undefined;
23+
onError: ErrorHandler | undefined;
2324
}

‎packages/ai/src/ui/chat.ts

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -496,6 +496,9 @@ export abstract class AbstractChat<UI_MESSAGE extends UIMessage> {
496496
messageMetadataSchema: this.messageMetadataSchema,
497497
dataPartSchemas: this.dataPartSchemas,
498498
runUpdateMessageJob,
499+
onError: error => {
500+
throw error;
501+
},
499502
}),
500503
onError: error => {
501504
throw error;
@@ -506,8 +509,6 @@ export abstract class AbstractChat<UI_MESSAGE extends UIMessage> {
506509

507510
this.setStatus({ status: 'ready' });
508511
} catch (err) {
509-
console.error(err);
510-
511512
// Ignore abort errors as they are expected.
512513
if ((err as any).name === 'AbortError') {
513514
this.setStatus({ status: 'ready' });

‎packages/ai/src/ui/process-ui-message-stream.test.ts

Lines changed: 110 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -62,6 +62,9 @@ describe('processUIMessageStream', () => {
6262
stream: processUIMessageStream({
6363
stream,
6464
runUpdateMessageJob,
65+
onError: error => {
66+
throw error;
67+
},
6568
}),
6669
});
6770
});
@@ -170,6 +173,56 @@ describe('processUIMessageStream', () => {
170173
});
171174
});
172175

176+
describe('errors', () => {
177+
let errors: Array<unknown>;
178+
179+
beforeEach(async () => {
180+
errors = [];
181+
182+
const stream = createUIMessageStream([
183+
{ type: 'error', errorText: 'test error' },
184+
]);
185+
186+
state = createStreamingUIMessageState({
187+
messageId: 'msg-123',
188+
lastMessage: undefined,
189+
});
190+
191+
await consumeStream({
192+
stream: processUIMessageStream({
193+
stream,
194+
runUpdateMessageJob,
195+
onError: error => {
196+
errors.push(error);
197+
},
198+
}),
199+
});
200+
});
201+
202+
it('should call the update function with the correct arguments', async () => {
203+
expect(writeCalls).toMatchInlineSnapshot(`[]`);
204+
});
205+
206+
it('should have the correct final message state', async () => {
207+
expect(state!.message).toMatchInlineSnapshot(`
208+
{
209+
"id": "msg-123",
210+
"metadata": undefined,
211+
"parts": [],
212+
"role": "assistant",
213+
}
214+
`);
215+
});
216+
217+
it('should call the onError function with the correct arguments', async () => {
218+
expect(errors).toMatchInlineSnapshot(`
219+
[
220+
[Error: test error],
221+
]
222+
`);
223+
});
224+
});
225+
173226
describe('server-side tool roundtrip', () => {
174227
beforeEach(async () => {
175228
const stream = createUIMessageStream([
@@ -208,6 +261,9 @@ describe('processUIMessageStream', () => {
208261
stream: processUIMessageStream({
209262
stream,
210263
runUpdateMessageJob,
264+
onError: error => {
265+
throw error;
266+
},
211267
}),
212268
});
213269
});
@@ -462,6 +518,9 @@ describe('processUIMessageStream', () => {
462518
stream: processUIMessageStream({
463519
stream,
464520
runUpdateMessageJob,
521+
onError: error => {
522+
throw error;
523+
},
465524
}),
466525
});
467526
});
@@ -772,6 +831,9 @@ describe('processUIMessageStream', () => {
772831
stream: processUIMessageStream({
773832
stream,
774833
runUpdateMessageJob,
834+
onError: error => {
835+
throw error;
836+
},
775837
}),
776838
});
777839
});
@@ -1174,6 +1236,9 @@ describe('processUIMessageStream', () => {
11741236
stream: processUIMessageStream({
11751237
stream,
11761238
runUpdateMessageJob,
1239+
onError: error => {
1240+
throw error;
1241+
},
11771242
}),
11781243
});
11791244
});
@@ -1739,6 +1804,9 @@ describe('processUIMessageStream', () => {
17391804
stream: processUIMessageStream({
17401805
stream,
17411806
runUpdateMessageJob,
1807+
onError: error => {
1808+
throw error;
1809+
},
17421810
}),
17431811
});
17441812
});
@@ -1979,6 +2047,9 @@ describe('processUIMessageStream', () => {
19792047
stream: processUIMessageStream({
19802048
stream,
19812049
runUpdateMessageJob,
2050+
onError: error => {
2051+
throw error;
2052+
},
19822053
}),
19832054
});
19842055
});
@@ -2205,6 +2276,9 @@ describe('processUIMessageStream', () => {
22052276
stream: processUIMessageStream({
22062277
stream,
22072278
runUpdateMessageJob,
2279+
onError: error => {
2280+
throw error;
2281+
},
22082282
}),
22092283
});
22102284
});
@@ -2353,6 +2427,9 @@ describe('processUIMessageStream', () => {
23532427
stream: processUIMessageStream({
23542428
stream,
23552429
runUpdateMessageJob,
2430+
onError: error => {
2431+
throw error;
2432+
},
23562433
}),
23572434
});
23582435
});
@@ -2508,6 +2585,9 @@ describe('processUIMessageStream', () => {
25082585
stream: processUIMessageStream({
25092586
stream,
25102587
runUpdateMessageJob,
2588+
onError: error => {
2589+
throw error;
2590+
},
25112591
}),
25122592
});
25132593
});
@@ -2689,6 +2769,9 @@ describe('processUIMessageStream', () => {
26892769
stream: processUIMessageStream({
26902770
stream,
26912771
runUpdateMessageJob,
2772+
onError: error => {
2773+
throw error;
2774+
},
26922775
}),
26932776
});
26942777
});
@@ -2858,6 +2941,9 @@ describe('processUIMessageStream', () => {
28582941
stream: processUIMessageStream({
28592942
stream,
28602943
runUpdateMessageJob,
2944+
onError: error => {
2945+
throw error;
2946+
},
28612947
}),
28622948
});
28632949
});
@@ -3426,6 +3512,9 @@ describe('processUIMessageStream', () => {
34263512
stream,
34273513
runUpdateMessageJob,
34283514
onToolCall: vi.fn().mockResolvedValue('test-result'),
3515+
onError: error => {
3516+
throw error;
3517+
},
34293518
}),
34303519
});
34313520
});
@@ -3549,6 +3638,9 @@ describe('processUIMessageStream', () => {
35493638
stream: processUIMessageStream({
35503639
stream,
35513640
runUpdateMessageJob,
3641+
onError: error => {
3642+
throw error;
3643+
},
35523644
}),
35533645
});
35543646
});
@@ -3705,6 +3797,9 @@ describe('processUIMessageStream', () => {
37053797
stream: processUIMessageStream({
37063798
stream,
37073799
runUpdateMessageJob,
3800+
onError: error => {
3801+
throw error;
3802+
},
37083803
}),
37093804
});
37103805
});
@@ -3968,6 +4063,9 @@ describe('processUIMessageStream', () => {
39684063
stream: processUIMessageStream({
39694064
stream,
39704065
runUpdateMessageJob,
4066+
onError: error => {
4067+
throw error;
4068+
},
39714069
}),
39724070
});
39734071
});
@@ -4051,6 +4149,9 @@ describe('processUIMessageStream', () => {
40514149
stream: processUIMessageStream({
40524150
stream,
40534151
runUpdateMessageJob,
4152+
onError: error => {
4153+
throw error;
4154+
},
40544155
}),
40554156
});
40564157
});
@@ -4159,6 +4260,9 @@ describe('processUIMessageStream', () => {
41594260
stream: processUIMessageStream({
41604261
stream,
41614262
runUpdateMessageJob,
4263+
onError: error => {
4264+
throw error;
4265+
},
41624266
}),
41634267
});
41644268
});
@@ -4307,6 +4411,9 @@ describe('processUIMessageStream', () => {
43074411
return 'client-result';
43084412
},
43094413
runUpdateMessageJob,
4414+
onError: error => {
4415+
throw error;
4416+
},
43104417
}),
43114418
});
43124419
});
@@ -4557,6 +4664,9 @@ describe('processUIMessageStream', () => {
45574664
return 'client-result';
45584665
},
45594666
runUpdateMessageJob,
4667+
onError: error => {
4668+
throw error;
4669+
},
45604670
}),
45614671
});
45624672

0 commit comments

Comments
 (0)