ai
Version:
AI SDK by Vercel - The AI Toolkit for TypeScript and JavaScript
1,696 lines (1,637 loc) • 240 kB
text/typescript
import { convertArrayToReadableStream } from '@ai-sdk/provider-utils/test';
import { UIMessageChunk } from '../ui-message-stream/ui-message-chunks';
import { consumeStream } from '../util/consume-stream';
import {
createStreamingUIMessageState,
processUIMessageStream,
StreamingUIMessageState,
} from './process-ui-message-stream';
import { InferUIMessageData, UIMessage } from './ui-messages';
import { beforeEach, describe, it, expect, vi } from 'vitest';
import { UIMessageStreamError } from '../error/ui-message-stream-error';
function createUIMessageStream(parts: UIMessageChunk[]) {
return convertArrayToReadableStream(parts);
}
describe('processUIMessageStream', () => {
let writeCalls: Array<{ message: UIMessage }> = [];
let state: StreamingUIMessageState<UIMessage> | undefined;
beforeEach(() => {
writeCalls = [];
state = undefined;
});
const runUpdateMessageJob = async (
job: (options: {
state: StreamingUIMessageState<UIMessage>;
write: () => void;
}) => Promise<void>,
) => {
await job({
state: state!,
write: () => {
writeCalls.push({ message: structuredClone(state!.message) });
},
});
};
describe('text', () => {
beforeEach(async () => {
const stream = createUIMessageStream([
{ type: 'start', messageId: 'msg-123' },
{ type: 'start-step' },
{ type: 'text-start', id: 'text-1' },
{ type: 'text-delta', id: 'text-1', delta: 'Hello, ' },
{ type: 'text-delta', id: 'text-1', delta: 'world!' },
{ type: 'text-end', id: 'text-1' },
{ type: 'finish-step' },
{ type: 'finish' },
]);
state = createStreamingUIMessageState({
messageId: 'msg-123',
lastMessage: undefined,
});
await consumeStream({
stream: processUIMessageStream({
stream,
runUpdateMessageJob,
onError: error => {
throw error;
},
}),
});
});
it('should call the update function with the correct arguments', async () => {
expect(writeCalls).toMatchInlineSnapshot(`
[
{
"message": {
"id": "msg-123",
"metadata": undefined,
"parts": [],
"role": "assistant",
},
},
{
"message": {
"id": "msg-123",
"metadata": undefined,
"parts": [
{
"type": "step-start",
},
{
"providerMetadata": undefined,
"state": "streaming",
"text": "",
"type": "text",
},
],
"role": "assistant",
},
},
{
"message": {
"id": "msg-123",
"metadata": undefined,
"parts": [
{
"type": "step-start",
},
{
"providerMetadata": undefined,
"state": "streaming",
"text": "Hello, ",
"type": "text",
},
],
"role": "assistant",
},
},
{
"message": {
"id": "msg-123",
"metadata": undefined,
"parts": [
{
"type": "step-start",
},
{
"providerMetadata": undefined,
"state": "streaming",
"text": "Hello, world!",
"type": "text",
},
],
"role": "assistant",
},
},
{
"message": {
"id": "msg-123",
"metadata": undefined,
"parts": [
{
"type": "step-start",
},
{
"providerMetadata": undefined,
"state": "done",
"text": "Hello, world!",
"type": "text",
},
],
"role": "assistant",
},
},
]
`);
});
it('should have the correct final message state', async () => {
expect(state!.message).toMatchInlineSnapshot(`
{
"id": "msg-123",
"metadata": undefined,
"parts": [
{
"type": "step-start",
},
{
"providerMetadata": undefined,
"state": "done",
"text": "Hello, world!",
"type": "text",
},
],
"role": "assistant",
}
`);
});
});
describe('errors', () => {
let errors: Array<unknown>;
beforeEach(async () => {
errors = [];
const stream = createUIMessageStream([
{ type: 'error', errorText: 'test error' },
]);
state = createStreamingUIMessageState({
messageId: 'msg-123',
lastMessage: undefined,
});
await consumeStream({
stream: processUIMessageStream({
stream,
runUpdateMessageJob,
onError: error => {
errors.push(error);
},
}),
});
});
it('should call the update function with the correct arguments', async () => {
expect(writeCalls).toMatchInlineSnapshot(`[]`);
});
it('should have the correct final message state', async () => {
expect(state!.message).toMatchInlineSnapshot(`
{
"id": "msg-123",
"metadata": undefined,
"parts": [],
"role": "assistant",
}
`);
});
it('should call the onError function with the correct arguments', async () => {
expect(errors).toMatchInlineSnapshot(`
[
[Error: test error],
]
`);
});
});
describe('malformed stream errors', () => {
it('should throw descriptive error when text-delta is received without text-start', async () => {
const stream = createUIMessageStream([
{ type: 'start', messageId: 'msg-123' },
{ type: 'start-step' },
{ type: 'text-delta', id: 'text-1', delta: 'Hello' },
]);
state = createStreamingUIMessageState({
messageId: 'msg-123',
lastMessage: undefined,
});
await expect(
consumeStream({
stream: processUIMessageStream({
stream,
runUpdateMessageJob,
onError: error => {
throw error;
},
}),
onError: error => {
throw error;
},
}),
).rejects.toThrow(
'Received text-delta for missing text part with ID "text-1". ' +
'Ensure a "text-start" chunk is sent before any "text-delta" chunks.',
);
});
it('should throw descriptive error when reasoning-delta is received without reasoning-start', async () => {
const stream = createUIMessageStream([
{ type: 'start', messageId: 'msg-123' },
{ type: 'start-step' },
{ type: 'reasoning-delta', id: 'reasoning-1', delta: 'Thinking...' },
]);
state = createStreamingUIMessageState({
messageId: 'msg-123',
lastMessage: undefined,
});
await expect(
consumeStream({
stream: processUIMessageStream({
stream,
runUpdateMessageJob,
onError: error => {
throw error;
},
}),
onError: error => {
throw error;
},
}),
).rejects.toThrow(
'Received reasoning-delta for missing reasoning part with ID "reasoning-1". ' +
'Ensure a "reasoning-start" chunk is sent before any "reasoning-delta" chunks.',
);
});
it('should throw descriptive error when tool-input-delta is received without tool-input-start', async () => {
const stream = createUIMessageStream([
{ type: 'start', messageId: 'msg-123' },
{ type: 'start-step' },
{
type: 'tool-input-delta',
toolCallId: 'tool-1',
inputTextDelta: '{"key":',
},
]);
state = createStreamingUIMessageState({
messageId: 'msg-123',
lastMessage: undefined,
});
await expect(
consumeStream({
stream: processUIMessageStream({
stream,
runUpdateMessageJob,
onError: error => {
throw error;
},
}),
onError: error => {
throw error;
},
}),
).rejects.toThrow(
'Received tool-input-delta for missing tool call with ID "tool-1". ' +
'Ensure a "tool-input-start" chunk is sent before any "tool-input-delta" chunks.',
);
});
it('should throw descriptive error when text-end is received without text-start', async () => {
const stream = createUIMessageStream([
{ type: 'start', messageId: 'msg-123' },
{ type: 'start-step' },
{ type: 'text-end', id: 'text-1' },
]);
state = createStreamingUIMessageState({
messageId: 'msg-123',
lastMessage: undefined,
});
await expect(
consumeStream({
stream: processUIMessageStream({
stream,
runUpdateMessageJob,
onError: error => {
throw error;
},
}),
onError: error => {
throw error;
},
}),
).rejects.toThrow(
'Received text-end for missing text part with ID "text-1". ' +
'Ensure a "text-start" chunk is sent before any "text-end" chunks.',
);
});
it('should throw descriptive error when reasoning-end is received without reasoning-start', async () => {
const stream = createUIMessageStream([
{ type: 'start', messageId: 'msg-123' },
{ type: 'start-step' },
{ type: 'reasoning-end', id: 'reasoning-1' },
]);
state = createStreamingUIMessageState({
messageId: 'msg-123',
lastMessage: undefined,
});
await expect(
consumeStream({
stream: processUIMessageStream({
stream,
runUpdateMessageJob,
onError: error => {
throw error;
},
}),
onError: error => {
throw error;
},
}),
).rejects.toThrow(
'Received reasoning-end for missing reasoning part with ID "reasoning-1". ' +
'Ensure a "reasoning-start" chunk is sent before any "reasoning-end" chunks.',
);
});
it('should throw UIMessageStreamError with correct properties for text-delta without text-start', async () => {
const stream = createUIMessageStream([
{ type: 'start', messageId: 'msg-123' },
{ type: 'start-step' },
{ type: 'text-delta', id: 'missing-id', delta: 'Hello' },
]);
state = createStreamingUIMessageState({
messageId: 'msg-123',
lastMessage: undefined,
});
let caughtError: unknown;
try {
await consumeStream({
stream: processUIMessageStream({
stream,
runUpdateMessageJob,
onError: error => {
throw error;
},
}),
onError: error => {
throw error;
},
});
} catch (error) {
caughtError = error;
}
expect(UIMessageStreamError.isInstance(caughtError)).toBe(true);
expect((caughtError as UIMessageStreamError).chunkType).toBe(
'text-delta',
);
expect((caughtError as UIMessageStreamError).chunkId).toBe('missing-id');
});
it('should throw UIMessageStreamError with correct properties for tool-input-delta without tool-input-start', async () => {
const stream = createUIMessageStream([
{ type: 'start', messageId: 'msg-123' },
{ type: 'start-step' },
{
type: 'tool-input-delta',
toolCallId: 'missing-tool-id',
inputTextDelta: '{"key":',
},
]);
state = createStreamingUIMessageState({
messageId: 'msg-123',
lastMessage: undefined,
});
let caughtError: unknown;
try {
await consumeStream({
stream: processUIMessageStream({
stream,
runUpdateMessageJob,
onError: error => {
throw error;
},
}),
onError: error => {
throw error;
},
});
} catch (error) {
caughtError = error;
}
expect(UIMessageStreamError.isInstance(caughtError)).toBe(true);
expect((caughtError as UIMessageStreamError).chunkType).toBe(
'tool-input-delta',
);
expect((caughtError as UIMessageStreamError).chunkId).toBe(
'missing-tool-id',
);
});
});
describe('server-side tool roundtrip', () => {
beforeEach(async () => {
const stream = createUIMessageStream([
{ type: 'start', messageId: 'msg-123' },
{ type: 'start-step' },
{
type: 'tool-input-available',
toolCallId: 'tool-call-id',
toolName: 'tool-name',
input: { city: 'London' },
},
{
type: 'tool-output-available',
toolCallId: 'tool-call-id',
output: { weather: 'sunny' },
},
{ type: 'finish-step' },
{ type: 'start-step' },
{ type: 'text-start', id: 'text-1' },
{
type: 'text-delta',
id: 'text-1',
delta: 'The weather in London is sunny.',
},
{ type: 'text-end', id: 'text-1' },
{ type: 'finish-step' },
{ type: 'finish' },
]);
state = createStreamingUIMessageState({
messageId: 'msg-123',
lastMessage: undefined,
});
await consumeStream({
stream: processUIMessageStream({
stream,
runUpdateMessageJob,
onError: error => {
throw error;
},
}),
});
});
it('should call the update function with the correct arguments', async () => {
expect(writeCalls).toMatchInlineSnapshot(`
[
{
"message": {
"id": "msg-123",
"metadata": undefined,
"parts": [],
"role": "assistant",
},
},
{
"message": {
"id": "msg-123",
"metadata": undefined,
"parts": [
{
"type": "step-start",
},
{
"errorText": undefined,
"input": {
"city": "London",
},
"output": undefined,
"preliminary": undefined,
"providerExecuted": undefined,
"rawInput": undefined,
"state": "input-available",
"title": undefined,
"toolCallId": "tool-call-id",
"type": "tool-tool-name",
},
],
"role": "assistant",
},
},
{
"message": {
"id": "msg-123",
"metadata": undefined,
"parts": [
{
"type": "step-start",
},
{
"errorText": undefined,
"input": {
"city": "London",
},
"output": {
"weather": "sunny",
},
"preliminary": undefined,
"providerExecuted": undefined,
"rawInput": undefined,
"state": "output-available",
"title": undefined,
"toolCallId": "tool-call-id",
"type": "tool-tool-name",
},
],
"role": "assistant",
},
},
{
"message": {
"id": "msg-123",
"metadata": undefined,
"parts": [
{
"type": "step-start",
},
{
"errorText": undefined,
"input": {
"city": "London",
},
"output": {
"weather": "sunny",
},
"preliminary": undefined,
"providerExecuted": undefined,
"rawInput": undefined,
"state": "output-available",
"title": undefined,
"toolCallId": "tool-call-id",
"type": "tool-tool-name",
},
{
"type": "step-start",
},
{
"providerMetadata": undefined,
"state": "streaming",
"text": "",
"type": "text",
},
],
"role": "assistant",
},
},
{
"message": {
"id": "msg-123",
"metadata": undefined,
"parts": [
{
"type": "step-start",
},
{
"errorText": undefined,
"input": {
"city": "London",
},
"output": {
"weather": "sunny",
},
"preliminary": undefined,
"providerExecuted": undefined,
"rawInput": undefined,
"state": "output-available",
"title": undefined,
"toolCallId": "tool-call-id",
"type": "tool-tool-name",
},
{
"type": "step-start",
},
{
"providerMetadata": undefined,
"state": "streaming",
"text": "The weather in London is sunny.",
"type": "text",
},
],
"role": "assistant",
},
},
{
"message": {
"id": "msg-123",
"metadata": undefined,
"parts": [
{
"type": "step-start",
},
{
"errorText": undefined,
"input": {
"city": "London",
},
"output": {
"weather": "sunny",
},
"preliminary": undefined,
"providerExecuted": undefined,
"rawInput": undefined,
"state": "output-available",
"title": undefined,
"toolCallId": "tool-call-id",
"type": "tool-tool-name",
},
{
"type": "step-start",
},
{
"providerMetadata": undefined,
"state": "done",
"text": "The weather in London is sunny.",
"type": "text",
},
],
"role": "assistant",
},
},
]
`);
});
it('should have the correct final message state', async () => {
expect(state!.message).toMatchInlineSnapshot(`
{
"id": "msg-123",
"metadata": undefined,
"parts": [
{
"type": "step-start",
},
{
"errorText": undefined,
"input": {
"city": "London",
},
"output": {
"weather": "sunny",
},
"preliminary": undefined,
"providerExecuted": undefined,
"rawInput": undefined,
"state": "output-available",
"title": undefined,
"toolCallId": "tool-call-id",
"type": "tool-tool-name",
},
{
"type": "step-start",
},
{
"providerMetadata": undefined,
"state": "done",
"text": "The weather in London is sunny.",
"type": "text",
},
],
"role": "assistant",
}
`);
});
});
describe('server-side tool roundtrip with existing assistant message', () => {
beforeEach(async () => {
const stream = createUIMessageStream([
{ type: 'start', messageId: 'msg-123' },
{ type: 'start-step' },
{
type: 'tool-input-available',
toolCallId: 'tool-call-id',
toolName: 'tool-name',
input: { city: 'London' },
},
{
type: 'tool-output-available',
toolCallId: 'tool-call-id',
output: { weather: 'sunny' },
},
{ type: 'finish-step' },
{ type: 'start-step' },
{ type: 'text-start', id: 'text-1' },
{
type: 'text-delta',
id: 'text-1',
delta: 'The weather in London is sunny.',
},
{ type: 'text-end', id: 'text-1' },
{ type: 'finish-step' },
{ type: 'finish' },
]);
state = createStreamingUIMessageState({
messageId: 'msg-123',
lastMessage: {
role: 'assistant',
id: 'original-id',
metadata: undefined,
parts: [
{
type: 'tool-tool-name-original',
toolCallId: 'tool-call-id-original',
state: 'output-available',
input: {},
output: { location: 'Berlin' },
},
],
},
});
await consumeStream({
stream: processUIMessageStream({
stream,
runUpdateMessageJob,
onError: error => {
throw error;
},
}),
});
});
it('should call the update function with the correct arguments', async () => {
expect(writeCalls).toMatchInlineSnapshot(`
[
{
"message": {
"id": "msg-123",
"metadata": undefined,
"parts": [
{
"input": {},
"output": {
"location": "Berlin",
},
"state": "output-available",
"toolCallId": "tool-call-id-original",
"type": "tool-tool-name-original",
},
],
"role": "assistant",
},
},
{
"message": {
"id": "msg-123",
"metadata": undefined,
"parts": [
{
"input": {},
"output": {
"location": "Berlin",
},
"state": "output-available",
"toolCallId": "tool-call-id-original",
"type": "tool-tool-name-original",
},
{
"type": "step-start",
},
{
"errorText": undefined,
"input": {
"city": "London",
},
"output": undefined,
"preliminary": undefined,
"providerExecuted": undefined,
"rawInput": undefined,
"state": "input-available",
"title": undefined,
"toolCallId": "tool-call-id",
"type": "tool-tool-name",
},
],
"role": "assistant",
},
},
{
"message": {
"id": "msg-123",
"metadata": undefined,
"parts": [
{
"input": {},
"output": {
"location": "Berlin",
},
"state": "output-available",
"toolCallId": "tool-call-id-original",
"type": "tool-tool-name-original",
},
{
"type": "step-start",
},
{
"errorText": undefined,
"input": {
"city": "London",
},
"output": {
"weather": "sunny",
},
"preliminary": undefined,
"providerExecuted": undefined,
"rawInput": undefined,
"state": "output-available",
"title": undefined,
"toolCallId": "tool-call-id",
"type": "tool-tool-name",
},
],
"role": "assistant",
},
},
{
"message": {
"id": "msg-123",
"metadata": undefined,
"parts": [
{
"input": {},
"output": {
"location": "Berlin",
},
"state": "output-available",
"toolCallId": "tool-call-id-original",
"type": "tool-tool-name-original",
},
{
"type": "step-start",
},
{
"errorText": undefined,
"input": {
"city": "London",
},
"output": {
"weather": "sunny",
},
"preliminary": undefined,
"providerExecuted": undefined,
"rawInput": undefined,
"state": "output-available",
"title": undefined,
"toolCallId": "tool-call-id",
"type": "tool-tool-name",
},
{
"type": "step-start",
},
{
"providerMetadata": undefined,
"state": "streaming",
"text": "",
"type": "text",
},
],
"role": "assistant",
},
},
{
"message": {
"id": "msg-123",
"metadata": undefined,
"parts": [
{
"input": {},
"output": {
"location": "Berlin",
},
"state": "output-available",
"toolCallId": "tool-call-id-original",
"type": "tool-tool-name-original",
},
{
"type": "step-start",
},
{
"errorText": undefined,
"input": {
"city": "London",
},
"output": {
"weather": "sunny",
},
"preliminary": undefined,
"providerExecuted": undefined,
"rawInput": undefined,
"state": "output-available",
"title": undefined,
"toolCallId": "tool-call-id",
"type": "tool-tool-name",
},
{
"type": "step-start",
},
{
"providerMetadata": undefined,
"state": "streaming",
"text": "The weather in London is sunny.",
"type": "text",
},
],
"role": "assistant",
},
},
{
"message": {
"id": "msg-123",
"metadata": undefined,
"parts": [
{
"input": {},
"output": {
"location": "Berlin",
},
"state": "output-available",
"toolCallId": "tool-call-id-original",
"type": "tool-tool-name-original",
},
{
"type": "step-start",
},
{
"errorText": undefined,
"input": {
"city": "London",
},
"output": {
"weather": "sunny",
},
"preliminary": undefined,
"providerExecuted": undefined,
"rawInput": undefined,
"state": "output-available",
"title": undefined,
"toolCallId": "tool-call-id",
"type": "tool-tool-name",
},
{
"type": "step-start",
},
{
"providerMetadata": undefined,
"state": "done",
"text": "The weather in London is sunny.",
"type": "text",
},
],
"role": "assistant",
},
},
]
`);
});
it('should have the correct final message state', async () => {
expect(state!.message).toMatchInlineSnapshot(`
{
"id": "msg-123",
"metadata": undefined,
"parts": [
{
"input": {},
"output": {
"location": "Berlin",
},
"state": "output-available",
"toolCallId": "tool-call-id-original",
"type": "tool-tool-name-original",
},
{
"type": "step-start",
},
{
"errorText": undefined,
"input": {
"city": "London",
},
"output": {
"weather": "sunny",
},
"preliminary": undefined,
"providerExecuted": undefined,
"rawInput": undefined,
"state": "output-available",
"title": undefined,
"toolCallId": "tool-call-id",
"type": "tool-tool-name",
},
{
"type": "step-start",
},
{
"providerMetadata": undefined,
"state": "done",
"text": "The weather in London is sunny.",
"type": "text",
},
],
"role": "assistant",
}
`);
});
});
describe('server-side tool roundtrip with multiple assistant texts', () => {
beforeEach(async () => {
const stream = createUIMessageStream([
{ type: 'start', messageId: 'msg-123' },
{ type: 'start-step' },
{ type: 'text-start', id: 'text-1' },
{ type: 'text-delta', id: 'text-1', delta: 'I will ' },
{
type: 'text-delta',
id: 'text-1',
delta: 'use a tool to get the weather in London.',
},
{ type: 'text-end', id: 'text-1' },
{
type: 'tool-input-available',
toolCallId: 'tool-call-id',
toolName: 'tool-name',
input: { city: 'London' },
},
{
type: 'tool-output-available',
toolCallId: 'tool-call-id',
output: { weather: 'sunny' },
},
{ type: 'finish-step' },
{ type: 'start-step' },
{ type: 'text-start', id: 'text-2' },
{ type: 'text-delta', id: 'text-2', delta: 'The weather in London ' },
{ type: 'text-delta', id: 'text-2', delta: 'is sunny.' },
{ type: 'text-end', id: 'text-2' },
{ type: 'finish-step' },
{ type: 'finish' },
]);
state = createStreamingUIMessageState({
messageId: 'msg-123',
lastMessage: undefined,
});
await consumeStream({
stream: processUIMessageStream({
stream,
runUpdateMessageJob,
onError: error => {
throw error;
},
}),
});
});
it('should call the update function with the correct arguments', async () => {
expect(writeCalls).toMatchInlineSnapshot(`
[
{
"message": {
"id": "msg-123",
"metadata": undefined,
"parts": [],
"role": "assistant",
},
},
{
"message": {
"id": "msg-123",
"metadata": undefined,
"parts": [
{
"type": "step-start",
},
{
"providerMetadata": undefined,
"state": "streaming",
"text": "",
"type": "text",
},
],
"role": "assistant",
},
},
{
"message": {
"id": "msg-123",
"metadata": undefined,
"parts": [
{
"type": "step-start",
},
{
"providerMetadata": undefined,
"state": "streaming",
"text": "I will ",
"type": "text",
},
],
"role": "assistant",
},
},
{
"message": {
"id": "msg-123",
"metadata": undefined,
"parts": [
{
"type": "step-start",
},
{
"providerMetadata": undefined,
"state": "streaming",
"text": "I will use a tool to get the weather in London.",
"type": "text",
},
],
"role": "assistant",
},
},
{
"message": {
"id": "msg-123",
"metadata": undefined,
"parts": [
{
"type": "step-start",
},
{
"providerMetadata": undefined,
"state": "done",
"text": "I will use a tool to get the weather in London.",
"type": "text",
},
],
"role": "assistant",
},
},
{
"message": {
"id": "msg-123",
"metadata": undefined,
"parts": [
{
"type": "step-start",
},
{
"providerMetadata": undefined,
"state": "done",
"text": "I will use a tool to get the weather in London.",
"type": "text",
},
{
"errorText": undefined,
"input": {
"city": "London",
},
"output": undefined,
"preliminary": undefined,
"providerExecuted": undefined,
"rawInput": undefined,
"state": "input-available",
"title": undefined,
"toolCallId": "tool-call-id",
"type": "tool-tool-name",
},
],
"role": "assistant",
},
},
{
"message": {
"id": "msg-123",
"metadata": undefined,
"parts": [
{
"type": "step-start",
},
{
"providerMetadata": undefined,
"state": "done",
"text": "I will use a tool to get the weather in London.",
"type": "text",
},
{
"errorText": undefined,
"input": {
"city": "London",
},
"output": {
"weather": "sunny",
},
"preliminary": undefined,
"providerExecuted": undefined,
"rawInput": undefined,
"state": "output-available",
"title": undefined,
"toolCallId": "tool-call-id",
"type": "tool-tool-name",
},
],
"role": "assistant",
},
},
{
"message": {
"id": "msg-123",
"metadata": undefined,
"parts": [
{
"type": "step-start",
},
{
"providerMetadata": undefined,
"state": "done",
"text": "I will use a tool to get the weather in London.",
"type": "text",
},
{
"errorText": undefined,
"input": {
"city": "London",
},
"output": {
"weather": "sunny",
},
"preliminary": undefined,
"providerExecuted": undefined,
"rawInput": undefined,
"state": "output-available",
"title": undefined,
"toolCallId": "tool-call-id",
"type": "tool-tool-name",
},
{
"type": "step-start",
},
{
"providerMetadata": undefined,
"state": "streaming",
"text": "",
"type": "text",
},
],
"role": "assistant",
},
},
{
"message": {
"id": "msg-123",
"metadata": undefined,
"parts": [
{
"type": "step-start",
},
{
"providerMetadata": undefined,
"state": "done",
"text": "I will use a tool to get the weather in London.",
"type": "text",
},
{
"errorText": undefined,
"input": {
"city": "London",
},
"output": {
"weather": "sunny",
},
"preliminary": undefined,
"providerExecuted": undefined,
"rawInput": undefined,
"state": "output-available",
"title": undefined,
"toolCallId": "tool-call-id",
"type": "tool-tool-name",
},
{
"type": "step-start",
},
{
"providerMetadata": undefined,
"state": "streaming",
"text": "The weather in London ",
"type": "text",
},
],
"role": "assistant",
},
},
{
"message": {
"id": "msg-123",
"metadata": undefined,
"parts": [
{
"type": "step-start",
},
{
"providerMetadata": undefined,
"state": "done",
"text": "I will use a tool to get the weather in London.",
"type": "text",
},
{
"errorText": undefined,
"input": {
"city": "London",
},
"output": {
"weather": "sunny",
},
"preliminary": undefined,
"providerExecuted": undefined,
"rawInput": undefined,
"state": "output-available",
"title": undefined,
"toolCallId": "tool-call-id",
"type": "tool-tool-name",
},
{
"type": "step-start",
},
{
"providerMetadata": undefined,
"state": "streaming",
"text": "The weather in London is sunny.",
"type": "text",
},
],
"role": "assistant",
},
},
{
"message": {
"id": "msg-123",
"metadata": undefined,
"parts": [
{
"type": "step-start",
},
{
"providerMetadata": undefined,
"state": "done",
"text": "I will use a tool to get the weather in London.",
"type": "text",
},
{
"errorText": undefined,
"input": {
"city": "London",
},
"output": {
"weather": "sunny",
},
"preliminary": undefined,
"providerExecuted": undefined,
"rawInput": undefined,
"state": "output-available",
"title": undefined,
"toolCallId": "tool-call-id",
"type": "tool-tool-name",
},
{
"type": "step-start",
},
{
"providerMetadata": undefined,
"state": "done",
"text": "The weather in London is sunny.",
"type": "text",
},
],
"role": "assistant",
},
},
]
`);
});
it('should have the correct final message state', async () => {
expect(state!.message).toMatchInlineSnapshot(`
{
"id": "msg-123",
"metadata": undefined,
"parts": [
{
"type": "step-start",
},
{
"providerMetadata": undefined,
"state": "done",
"text": "I will use a tool to get the weather in London.",
"type": "text",
},
{
"errorText": undefined,
"input": {
"city": "London",
},
"output": {
"weather": "sunny",
},
"preliminary": undefined,
"providerExecuted": undefined,
"rawInput": undefined,
"state": "output-available",
"title": undefined,
"toolCallId": "tool-call-id",
"type": "tool-tool-name",
},
{
"type": "step-start",
},
{
"providerMetadata": undefined,
"state": "done",
"text": "The weather in London is sunny.",
"type": "text",
},
],
"role": "assistant",
}
`);
});
});
describe('server-side tool roundtrip with multiple assistant reasoning', () => {
beforeEach(async () => {
const stream = createUIMessageStream([
{ type: 'start', messageId: 'msg-123' },
{ type: 'start-step' },
{ type: 'reasoning-start', id: 'reasoning-1' },
{
type: 'reasoning-delta',
id: 'reasoning-1',
delta: 'I will ',
providerMetadata: {
testProvider: { signature: '1234567890' },
},
},
{
type: 'reasoning-delta',
id: 'reasoning-1',
delta: 'use a tool to get the weather in London.',
},
{ type: 'reasoning-end', id: 'reasoning-1' },
{
type: 'tool-input-available',
toolCallId: 'tool-call-id',
toolName: 'tool-name',
input: { city: 'London' },
},
{
type: 'tool-output-available',
toolCallId: 'tool-call-id',
output: { weather: 'sunny' },
},
{ type: 'finish-step' },
{ type: 'start-step' },
{ type: 'reasoning-start', id: 'reasoning-2' },
{
type: 'reasoning-delta',
id: 'reasoning-2',
delta: 'I now know the weather in London.',
providerMetadata: {
testProvider: { signature: 'abc123' },
},
},
{ type: 'reasoning-end', id: 'reasoning-2' },
{ type: 'text-start', id: 'text-1' },
{
type: 'text-delta',
id: 'text-1',
delta: 'The weather in London is sunny.',
},
{ type: 'text-end', id: 'text-1' },
{ type: 'finish-step' },
{ type: 'finish' },
]);
state = createStreamingUIMessageState({
messageId: 'msg-123',
lastMessage: undefined,
});
await consumeStream({
stream: processUIMessageStream({
stream,
runUpdateMessageJob,
onError: error => {
throw error;
},
}),
});
});
it('should call the update function with the correct arguments', async () => {
expect(writeCalls).toMatchInlineSnapshot(`
[
{
"message": {
"id": "msg-123",
"metadata": undefined,
"parts": [],
"role": "assistant",
},
},
{
"message": {
"id": "msg-123",
"metadata": undefined,
"parts": [
{
"type": "step-start",
},
{
"providerMetadata": undefined,
"state": "streaming",
"text": "",
"type": "reasoning",
},
],
"role": "assistant",
},
},
{
"message": {
"id": "msg-123",
"metadata": undefined,
"parts": [
{
"type": "step-start",
},
{
"providerMetadata": {
"testProvider": {
"signature": "1234567890",
},
},
"state": "streaming",
"text": "I will ",
"type": "reasoning",
},
],
"role": "assistant",
},
},
{
"message": {
"id": "msg-123",
"metadata": undefined,
"parts": [
{
"type": "step-start",
},
{
"providerMetadata": {
"testProvider": {
"signature": "1234567890",
},
},
"state": "streaming",
"text": "I will use a tool to get the weather in London.",
"type": "reasoning",
},
],
"role": "assistant",
},
},
{
"message": {
"id": "msg-123",
"metadata": undefined,
"parts": [
{
"type": "step-start",
},
{
"providerMetadata": {
"testProvider": {
"signature": "1234567890",
},
},
"state": "done",
"text": "I will use a tool to get the weather in London.",
"type": "reasoning",
},
],
"role": "assistant",
},
},
{
"message": {
"id": "msg-123",
"metadata": undefined,
"parts": [
{
"type": "step-start",
},
{
"providerMetadata": {
"testProvider": {
"signature": "1234567890",
},
},
"state": "done",
"text": "I will use a tool to get the weather in London.",
"type": "reasoning",
},
{
"errorText": undefined,
"input": {
"city": "London",
},
"output": undefined,
"preliminary": undefined,
"providerExecuted": undefined,
"rawInput": undefined,
"state": "input-available",
"title": undefined,
"toolCallId": "tool-call-id",
"type": "tool-tool-name",
},