fix(assistants): retain terminal message snapshots

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<Message[]> {
     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');
       }

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]);
     },
   );