diff --git a/src/lib/AssistantStream.ts b/src/lib/AssistantStream.ts index ea9d9a194d..9cc64fa15d 100644 --- a/src/lib/AssistantStream.ts +++ b/src/lib/AssistantStream.ts @@ -322,7 +322,10 @@ export class AssistantStream return Object.values(this.#runStepSnapshots); } - /** Waits for successful completion and returns the final snapshot of every observed message. */ + /** + * Waits for successful completion and returns the final snapshot of every observed message. + * Terminal message events replace accumulated snapshots without mutating earlier snapshots. + */ async finalMessages(): Promise { await this.done(); @@ -479,6 +482,9 @@ export class AssistantStream case 'thread.message.delta': case 'thread.message.completed': case 'thread.message.incomplete': { + if (messageID !== undefined && this.#messageSnapshot) { + this.#reserveMessageAlias(this.#messageSnapshot, messageID); + } this.#handleMessage(stableEvent); if (messageID !== undefined) { this.#reserveMessageAlias(stableEvent.data, messageID); @@ -918,9 +924,9 @@ export class AssistantStream case 'thread.message.in_progress': case 'thread.message.completed': case 'thread.message.incomplete': { - //No changes on other thread events + //Terminal events provide the authoritative message snapshot. if (snapshot) { - return [snapshot, newContent]; + return [event.event === 'thread.message.in_progress' ? snapshot : event.data, newContent]; } throw new Error('Received thread message event with no existing snapshot'); } diff --git a/tests/lib/AssistantStream.test.ts b/tests/lib/AssistantStream.test.ts index 52c9ba825d..b3f65511de 100644 --- a/tests/lib/AssistantStream.test.ts +++ b/tests/lib/AssistantStream.test.ts @@ -2,6 +2,8 @@ import { vi } from 'vitest'; import { APIError, APIUserAbortError, OpenAIError } from 'openai/core/error'; import { AssistantStream } from 'openai/lib/AssistantStream'; import type { AssistantStreamEvent } from 'openai/resources/beta/assistants'; +import type { Message } from 'openai/resources/beta/threads/messages'; +import type { Run } from 'openai/resources/beta/threads/runs/runs'; import { Stream } from 'openai/streaming'; import { assistantStream, completedRun } from './assistant-stream-test-utils'; @@ -18,6 +20,65 @@ function iterableEvents(events: Event[], controller = new AbortController()) { }; } +function terminalMessageFixture(status: 'completed' | 'incomplete') { + const initialMessage: Message = { + id: 'msg_terminal', + assistant_id: 'asst_terminal', + attachments: [], + completed_at: null, + content: [], + created_at: 1_700_000_000, + incomplete_at: null, + incomplete_details: null, + metadata: {}, + object: 'thread.message', + role: 'assistant', + run_id: 'run_terminal', + status: 'in_progress', + thread_id: 'thread_terminal', + }; + const finalText = { value: 'Synthetic response.', annotations: [] }; + const finalMessage: Message = { + ...initialMessage, + status, + completed_at: status === 'completed' ? 1_700_000_001 : null, + incomplete_at: status === 'incomplete' ? 1_700_000_001 : null, + incomplete_details: status === 'incomplete' ? { reason: 'max_tokens' } : null, + content: [{ type: 'text', text: finalText }], + }; + const finalRun: Run = { + id: 'run_terminal', + assistant_id: 'asst_terminal', + cancelled_at: null, + completed_at: finalMessage.completed_at, + created_at: initialMessage.created_at, + expires_at: 1_700_000_600, + failed_at: null, + incomplete_details: status === 'incomplete' ? { reason: 'max_completion_tokens' } : null, + instructions: '', + last_error: null, + max_completion_tokens: 100, + max_prompt_tokens: null, + metadata: {}, + model: 'gpt-4o', + object: 'thread.run', + parallel_tool_calls: true, + required_action: null, + response_format: 'auto', + started_at: initialMessage.created_at, + status, + thread_id: initialMessage.thread_id, + tool_choice: 'auto', + tools: [], + truncation_strategy: { type: 'auto' }, + usage: + status === 'incomplete' + ? { prompt_tokens: 20, completion_tokens: 100, total_tokens: 120 } + : { prompt_tokens: 20, completion_tokens: 5, total_tokens: 25 }, + }; + return { initialMessage, finalMessage, finalText, finalRun }; +} + describe('AssistantStream delta accumulation', () => { test('accumulates numbers, nested records, and arrays of primitive values', () => { expect( @@ -268,6 +329,141 @@ describe('AssistantStream snapshots and message lifecycle', () => { expect(finalMessages[0]).toBe(message); }); + test.each(['completed', 'incomplete'] as const)( + 'adopts the terminal %s message snapshot without mutating retained delta snapshots', + async (status) => { + const { initialMessage, finalMessage, finalText, finalRun } = terminalMessageFixture(status); + const events: AssistantStreamEvent[] = [ + { event: 'thread.message.created', data: initialMessage }, + { event: 'thread.message.in_progress', data: initialMessage }, + { + event: 'thread.message.delta', + data: { + id: initialMessage.id, + object: 'thread.message.delta', + delta: { + content: [{ index: 0, type: 'text', text: { value: 'Synthetic response.', annotations: [] } }], + }, + }, + }, + { event: `thread.message.${status}`, data: finalMessage }, + { event: `thread.run.${status}`, data: finalRun }, + ]; + const runner = assistantStream(events); + const created = vi.fn(); + const textDone = vi.fn(); + const messageDone = vi.fn(); + const rawEvents: AssistantStreamEvent[] = []; + const order: string[] = []; + let deltaSnapshot: Message | undefined; + let retainedBeforeCompletion: Message | undefined; + let snapshotDuringRawTerminal: Message | undefined; + + runner.on('event', (event) => { + rawEvents.push(event); + order.push(event.event); + if (event.event === `thread.message.${status}`) { + snapshotDuringRawTerminal = runner.currentMessageSnapshot(); + } + }); + runner.on('messageCreated', created); + runner.on('messageDelta', (_delta, snapshot) => { + deltaSnapshot = snapshot; + retainedBeforeCompletion = structuredClone(snapshot); + }); + runner.on('textDone', (text, snapshot) => { + order.push('textDone'); + textDone(text, snapshot, runner.currentMessageSnapshot(), runner.currentEvent()); + }); + runner.on('messageDone', (message) => { + order.push('messageDone'); + messageDone(message, runner.currentMessageSnapshot(), runner.currentEvent()); + }); + + await runner.done(); + const finalMessages = await runner.finalMessages(); + const [terminalEvent, finalRunEvent] = rawEvents.slice(-2); + if ( + terminalEvent?.event !== 'thread.message.completed' && + terminalEvent?.event !== 'thread.message.incomplete' + ) { + throw new Error('Expected a terminal message event'); + } + const terminalMessage = terminalEvent.data; + + expect(finalMessages).toEqual([finalMessage]); + expect(finalMessages[0]).toBe(terminalMessage); + expect(finalMessages[0]?.content).toBe(terminalMessage.content); + expect(textDone).toHaveBeenCalledTimes(1); + expect(textDone).toHaveBeenCalledWith(finalText, terminalMessage, terminalMessage, terminalEvent); + expect(textDone.mock.calls[0]?.[1]).toBe(terminalMessage); + expect(textDone.mock.calls[0]?.[2]).toBe(terminalMessage); + expect(textDone.mock.calls[0]?.[3]).toBe(terminalEvent); + expect(messageDone).toHaveBeenCalledTimes(1); + expect(messageDone).toHaveBeenCalledWith(terminalMessage, terminalMessage, terminalEvent); + expect(messageDone.mock.calls[0]?.[0]).toBe(terminalMessage); + expect(messageDone.mock.calls[0]?.[1]).toBe(terminalMessage); + expect(messageDone.mock.calls[0]?.[2]).toBe(terminalEvent); + expect(created).toHaveBeenCalledTimes(1); + expect(created.mock.calls[0]?.[0]).toBe(deltaSnapshot); + expect(snapshotDuringRawTerminal).toBe(deltaSnapshot); + expect(finalMessages[0]).not.toBe(deltaSnapshot); + expect(deltaSnapshot).toEqual(retainedBeforeCompletion); + expect(deltaSnapshot?.status).toBe('in_progress'); + expect(runner.currentMessageSnapshot()).toBeUndefined(); + expect(runner.currentEvent()).toBe(finalRunEvent); + await expect(runner.finalRun()).resolves.toEqual(finalRun); + expect(order).toEqual([ + 'thread.message.created', + 'thread.message.in_progress', + 'thread.message.delta', + `thread.message.${status}`, + 'textDone', + 'messageDone', + `thread.run.${status}`, + ]); + }, + ); + + test.each(['completed', 'incomplete'] as const)( + 'reserves a retained message alias changed during the raw %s event', + async (status) => { + const { initialMessage, finalMessage, finalRun } = terminalMessageFixture(status); + const alias = 'msg_terminal_alias'; + const events: AssistantStreamEvent[] = [ + { event: 'thread.message.created', data: initialMessage }, + { event: `thread.message.${status}`, data: finalMessage }, + { event: 'thread.message.created', data: { ...initialMessage, id: alias } }, + { event: `thread.message.${status}`, data: { ...finalMessage, id: alias } }, + { event: `thread.run.${status}`, data: finalRun }, + ]; + const runner = assistantStream(events); + const created = vi.fn(); + const exposed: string[] = []; + let retainedSnapshot: Message | undefined; + runner.on('messageCreated', created); + runner.on('event', (event) => { + exposed.push(event.event); + if ( + (event.event === 'thread.message.completed' || event.event === 'thread.message.incomplete') && + event.data.id === initialMessage.id + ) { + retainedSnapshot = runner.currentMessageSnapshot(); + if (!retainedSnapshot) { + throw new Error('Expected the active message snapshot'); + } + retainedSnapshot.id = alias; + } + }); + + await expect(runner.done()).rejects.toThrow('already been created'); + expect(created).toHaveBeenCalledTimes(1); + expect(retainedSnapshot).toBe(created.mock.calls[0]?.[0]); + expect(retainedSnapshot?.id).toBe(alias); + expect(exposed).toEqual(['thread.message.created', `thread.message.${status}`]); + }, + ); + test('captures accessor-backed event data once before exposing or routing its message', async () => { const first = { id: 'msg_first', role: 'assistant', content: [] }; const second = { id: 'msg_second', role: 'assistant', content: [] }; @@ -897,8 +1093,12 @@ describe('AssistantStream snapshots and message lifecycle', () => { async (id) => { const message = { id, role: 'assistant', content: [] }; const runner = assistantStream([{ event: 'thread.message.created', data: message }, completedRun()]); + const created = vi.fn(); + runner.on('messageCreated', created); - await expect(runner.finalMessages()).resolves.toEqual([message]); + const finalMessages = await runner.finalMessages(); + expect(finalMessages).toEqual([message]); + expect(finalMessages[0]).toBe(created.mock.calls[0]?.[0]); }, );