@assistant-ui/core 0.3.20 → 0.3.21

package/src/react/runtimes/cloud/auiV0DecodeSafely.test.ts

@@ -1,6 +1,10 @@
 import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
 import type { CloudMessage } from "assistant-cloud";
-import type { ThreadAssistantMessage, ThreadUserMessage } from "../../../types";
+import type {
+  MessageStatus,
+  ThreadAssistantMessage,
+  ThreadUserMessage,
+} from "../../../types";
 import { auiV0DecodeSafely, auiV0Encode } from "./auiV0";
 
 const storedRow = (content: unknown) =>
@@ -59,6 +63,44 @@ describe("auiV0DecodeSafely", () => {
     });
   });
 
+  it("normalizes malformed assistant status and step metadata", () => {
+    const item = auiV0DecodeSafely(
+      assistantRow([{ type: "text", text: "answer" }], {
+        status: { type: "unknown" },
+        metadata: {
+          custom: {},
+          steps: [
+            null,
+            {},
+            {
+              messageId: 42,
+              usage: { inputTokens: 7, outputTokens: 8 },
+            },
+            { usage: { inputTokens: 1, outputTokens: 2 } },
+            { usage: { inputTokens: 1 } },
+            { usage: { promptTokens: 3, completionTokens: 4 } },
+            { usage: { inputTokenDetails: { cacheReadTokens: 5 } } },
+            { usage: "invalid" },
+          ],
+        },
+      }),
+    );
+
+    expect(item?.message.status).toEqual({
+      type: "complete",
+      reason: "unknown",
+    });
+    expect(item?.message.metadata.steps).toEqual([
+      {},
+      { usage: { inputTokens: 7, outputTokens: 8 } },
+      { usage: { inputTokens: 1, outputTokens: 2 } },
+      { usage: { inputTokens: 1 } },
+      { usage: { promptTokens: 3, completionTokens: 4 } },
+      { usage: { inputTokenDetails: { cacheReadTokens: 5 } } },
+      {},
+    ]);
+  });
+
   it("keeps both one-sided tool call encodings", () => {
     const item = auiV0DecodeSafely(
       assistantRow([
@@ -83,6 +125,29 @@ describe("auiV0DecodeSafely", () => {
     ]);
   });
 
+  it("drops malformed tool-call interactions without dropping the message", () => {
+    const item = auiV0DecodeSafely(
+      assistantRow([
+        {
+          type: "tool-call",
+          toolCallId: "call-1",
+          toolName: "search",
+          args: {},
+          unstable_interactions: {
+            entries: [{ type: "unknown", occurredAt: 1, payload: null }],
+          },
+        },
+        { type: "text", text: "kept" },
+      ]),
+    );
+
+    expect(item?.message.content).toHaveLength(2);
+    expect(item?.message.content[0]).not.toHaveProperty(
+      "unstable_interactions",
+    );
+    expect(item?.message.content[1]).toEqual({ type: "text", text: "kept" });
+  });
+
   it("keeps a data prefixed part the decoder can still convert", () => {
     const item = auiV0DecodeSafely(
       assistantRow([{ type: "data-weather", data: { city: "Berlin" } }]),
@@ -416,6 +481,215 @@ describe("auiV0DecodeSafely against encoder output", () => {
       content: auiV0Encode(message),
     }) as unknown as CloudMessage & { format: "aui/v0" };
 
+  const withToolResult = (result: unknown): ThreadAssistantMessage => ({
+    id: "assistant-1",
+    role: "assistant",
+    status: { type: "complete", reason: "stop" },
+    createdAt: new Date(0),
+    metadata: {
+      unstable_state: null,
+      unstable_annotations: [],
+      unstable_data: [],
+      steps: [],
+      custom: {},
+    },
+    content: [
+      {
+        type: "tool-call",
+        toolCallId: "call-1",
+        toolName: "lookup",
+        args: {},
+        argsText: "{}",
+        result,
+      },
+    ],
+  });
+
+  const encodedResult = (result: unknown) => {
+    const warn = vi.spyOn(console, "warn").mockImplementation(() => {});
+    try {
+      const encoded = auiV0Encode(withToolResult(result));
+      const part = encoded.content[0] as { result?: unknown };
+      return { result: part.result, warned: warn.mock.calls.length };
+    } finally {
+      warn.mockRestore();
+    }
+  };
+
+  it("persists the JSON form of a result JSON would empty, and warns", () => {
+    expect(encodedResult(new Map([["answer", 42]]))).toEqual({
+      result: {},
+      warned: 1,
+    });
+    expect(encodedResult(new Set([1, 2]))).toEqual({ result: {}, warned: 1 });
+  });
+
+  it("keeps the siblings of a lossy field and warns once", () => {
+    expect(
+      encodedResult({ text: "long answer", index: new Map([["a", 1]]) }),
+    ).toEqual({ result: { text: "long answer", index: {} }, warned: 1 });
+    expect(encodedResult({ a: 1, b: undefined })).toEqual({
+      result: { a: 1 },
+      warned: 1,
+    });
+    expect(encodedResult({ label: "x", score: NaN })).toEqual({
+      result: { label: "x", score: null },
+      warned: 1,
+    });
+  });
+
+  it("keeps tool results JSON preserves", () => {
+    class Weather {
+      tempC = 21;
+      city = "Berlin";
+    }
+    const cases: [unknown, unknown][] = [
+      [{ createdAt: new Date(0) }, { createdAt: "1970-01-01T00:00:00.000Z" }],
+      [new Weather(), { tempC: 21, city: "Berlin" }],
+      [new URL("https://x.y/"), "https://x.y/"],
+      [{ nested: [{ items: new Map() }] }, { nested: [{ items: {} }] }],
+    ];
+    for (const [input, expected] of cases) {
+      expect(encodedResult(input)).toEqual({ result: expected, warned: 0 });
+    }
+  });
+
+  it("warns for a result that hides data behind non-string keys", () => {
+    const symbolKeyed = { [Symbol("hidden")]: 1, shown: 2 };
+    const sparse = Object.defineProperty([1, 2], "extra", {
+      value: 3,
+      enumerable: true,
+    });
+    expect(encodedResult(symbolKeyed)).toEqual({
+      result: { shown: 2 },
+      warned: 1,
+    });
+    expect(encodedResult(sparse)).toEqual({ result: [1, 2], warned: 1 });
+  });
+
+  it("omits a result JSON cannot serialize at all, and warns", () => {
+    const circular: { self?: unknown } = {};
+    circular.self = circular;
+    expect(encodedResult(circular)).toEqual({ result: undefined, warned: 1 });
+  });
+
+  it("takes an object literal's toJSON as its serialized form", () => {
+    expect(encodedResult({ toJSON: () => "summary" })).toEqual({
+      result: "summary",
+      warned: 0,
+    });
+  });
+
+  it("cannot see data behind prototype getters, matching main", () => {
+    class Wrapper {
+      #value = 1;
+      get value() {
+        return this.#value;
+      }
+    }
+    expect(encodedResult(new Wrapper())).toEqual({ result: {}, warned: 0 });
+  });
+
+  it("warns for opaque own state JSON cannot see", () => {
+    const empty = Object.assign(new Map(), { [Symbol("hidden")]: 1 });
+    expect(encodedResult(empty)).toEqual({ result: {}, warned: 1 });
+    expect(encodedResult(/x/)).toEqual({ result: {}, warned: 1 });
+  });
+
+  it("warns when a toJSON round trip yields nothing to store", () => {
+    expect(encodedResult({ toJSON: () => undefined })).toEqual({
+      result: undefined,
+      warned: 1,
+    });
+    expect(encodedResult({ toJSON: () => 1n })).toEqual({
+      result: undefined,
+      warned: 1,
+    });
+  });
+
+  it("reports a getter that throws on its second read instead of failing", () => {
+    let reads = 0;
+    const flaky = Object.defineProperty({}, "x", {
+      enumerable: true,
+      get() {
+        if (reads++ > 0) throw new Error("second read");
+        return 1;
+      },
+    });
+    expect(encodedResult(flaky)).toEqual({ result: { x: 1 }, warned: 1 });
+  });
+
+  it("keeps every assistant status the encoder writes", () => {
+    const statuses: MessageStatus[] = [
+      { type: "running" },
+      { type: "requires-action", reason: "tool-calls" },
+      { type: "requires-action", reason: "interrupt" },
+      { type: "complete", reason: "stop" },
+      { type: "complete", reason: "unknown" },
+      { type: "incomplete", reason: "cancelled" },
+      { type: "incomplete", reason: "tool-calls" },
+      { type: "incomplete", reason: "length" },
+      { type: "incomplete", reason: "content-filter" },
+      { type: "incomplete", reason: "other" },
+      { type: "incomplete", reason: "error", error: { message: "failed" } },
+      { type: "incomplete", reason: "max_tokens" } as unknown as MessageStatus,
+    ];
+
+    for (const status of statuses) {
+      const message: ThreadAssistantMessage = {
+        id: `assistant-${status.type}`,
+        role: "assistant",
+        status,
+        createdAt: new Date(0),
+        metadata: {
+          unstable_state: null,
+          unstable_annotations: [],
+          unstable_data: [],
+          steps: [
+            { messageId: "step-1", usage: { inputTokens: 1, outputTokens: 2 } },
+          ],
+          custom: {},
+        },
+        content: [{ type: "text", text: "answer" }],
+      };
+
+      const item = auiV0DecodeSafely(encodedRow(message));
+      const expectedStatus =
+        status.type === "running"
+          ? { type: "incomplete", reason: "cancelled" }
+          : status;
+
+      expect(item?.message.status).toEqual(expectedStatus);
+      expect(item?.message.metadata.steps).toEqual(message.metadata.steps);
+    }
+  });
+
+  it("keeps provider usage shapes the encoder writes", () => {
+    const providerSteps = [
+      { usage: { promptTokens: 3, completionTokens: 4 } },
+      { usage: { inputTokens: 0, cachedInputTokens: 5 } },
+      { usage: { inputTokenDetails: { cacheReadTokens: 5 } } },
+    ];
+    const message = {
+      id: "assistant-provider-usage",
+      role: "assistant",
+      status: { type: "complete", reason: "stop" },
+      createdAt: new Date(0),
+      metadata: {
+        unstable_state: null,
+        unstable_annotations: [],
+        unstable_data: [],
+        steps: providerSteps,
+        custom: {},
+      },
+      content: [{ type: "text", text: "answer" }],
+    } as unknown as ThreadAssistantMessage;
+
+    const item = auiV0DecodeSafely(encodedRow(message));
+
+    expect(item?.message.metadata.steps).toEqual(providerSteps);
+  });
+
   it("keeps every assistant part the encoder writes", () => {
     const message: ThreadAssistantMessage = {
       id: "assistant-1",
@@ -519,4 +793,81 @@ describe("auiV0DecodeSafely against encoder output", () => {
       ),
     ).toEqual(message.attachments[0]?.content.map((part) => part.type));
   });
+
+  it("keeps provider metadata on file and attachment parts", () => {
+    const providerMetadata = { openai: { fileId: "file-123" } };
+    const assistant: ThreadAssistantMessage = {
+      id: "assistant-file",
+      role: "assistant",
+      status: { type: "complete", reason: "stop" },
+      createdAt: new Date(0),
+      metadata: {
+        unstable_state: null,
+        unstable_annotations: [],
+        unstable_data: [],
+        steps: [],
+        custom: {},
+      },
+      content: [
+        {
+          type: "file",
+          data: "file-123",
+          mimeType: "application/pdf",
+          sourceType: "id",
+          providerMetadata,
+        },
+      ],
+    };
+    const user: ThreadUserMessage = {
+      id: "user-file",
+      role: "user",
+      createdAt: new Date(0),
+      metadata: { custom: {} },
+      content: [],
+      attachments: [
+        {
+          id: "attachment-1",
+          type: "file",
+          name: "report.pdf",
+          status: { type: "complete" },
+          content: [
+            {
+              type: "text",
+              text: "report",
+              providerMetadata,
+            },
+            {
+              type: "image",
+              image: "https://x.dev/preview.png",
+              providerMetadata,
+            },
+            {
+              type: "file",
+              data: "file-123",
+              mimeType: "application/pdf",
+              sourceType: "id",
+              providerMetadata,
+            },
+          ],
+        },
+      ],
+    };
+
+    expect(
+      auiV0DecodeSafely(encodedRow(assistant))?.message.content[0],
+    ).toMatchObject({
+      providerMetadata,
+    });
+    expect(auiV0DecodeSafely(encodedRow(user))?.message).toMatchObject({
+      attachments: [
+        {
+          content: [
+            { type: "text", providerMetadata },
+            { type: "image", providerMetadata },
+            { type: "file", providerMetadata },
+          ],
+        },
+      ],
+    });
+  });
 });

package/src/react/runtimes/createMessageConverter.ts

@@ -1,6 +1,6 @@
 "use client";
 import type { ThreadMessage } from "../../types/message";
-import type { ThreadState } from "../../runtime/api/thread-runtime";
+import type { ThreadRuntimeState } from "../../runtime/api/thread-runtime";
 import { useAui, useAuiState } from "@assistant-ui/store";
 import {
   useExternalMessageConverter,
@@ -40,14 +40,20 @@ export const createMessageConverter = <T extends object>(
       return convertExternalMessages(messages, callback, isRunning, metadata);
     },
     toOriginalMessages: (
-      input: ThreadState | ThreadMessage | ThreadMessage["content"][number],
+      input:
+        | ThreadRuntimeState
+        | ThreadMessage
+        | ThreadMessage["content"][number],
     ) => {
       const messages = getExternalStoreMessages(input);
       if (messages.length === 0) throw new Error("No original messages found");
       return messages;
     },
     toOriginalMessage: (
-      input: ThreadState | ThreadMessage | ThreadMessage["content"][number],
+      input:
+        | ThreadRuntimeState
+        | ThreadMessage
+        | ThreadMessage["content"][number],
     ) => {
       const messages = result.toOriginalMessages(input);
       return messages[0]!;

package/src/react/runtimes/external-store-history-copy.ts

@@ -0,0 +1,223 @@
+import type { ThreadHistoryAdapter } from "../../adapters/thread-history";
+import { getThreadRuntimeCoreIsRunning } from "../../runtime/api/thread-runtime";
+import type {
+  ThreadRuntimeCore,
+  Unstable_RecordToolInteractionOptions,
+} from "../../runtime/interfaces/thread-runtime-core";
+import { FALLBACK_ID_PREFIX } from "../../runtime/utils/external-store-message";
+import {
+  appendToolInteraction,
+  readToolInteractionLog,
+} from "../../runtime/utils/tool-interactions";
+import type {
+  ThreadMessage,
+  Unstable_ToolInteractionLog,
+} from "../../types/message";
+import { isErrorMessageId } from "../../utils/id";
+
+type Signature = string | ThreadMessage;
+
+export class ExternalStoreHistoryCopy {
+  private thread: ThreadRuntimeCore | undefined;
+  private history: ThreadHistoryAdapter | undefined;
+  private lastHistory: ThreadHistoryAdapter | undefined;
+  private session = 0;
+  private copied = new Map<string, Signature>();
+  private failedIds = new Set<string>();
+  private interactions = new Map<
+    string,
+    Map<string, Unstable_ToolInteractionLog>
+  >();
+  private overlaid = new WeakMap<
+    ThreadMessage,
+    {
+      logs: Map<string, Unstable_ToolInteractionLog>;
+      message: ThreadMessage;
+    }
+  >();
+  private pending = false;
+  private inFlight = false;
+  private timer: ReturnType<typeof setTimeout> | undefined;
+  private waiters: {
+    resolve: () => void;
+    reject: (error: unknown) => void;
+  }[] = [];
+  private warned = false;
+
+  private signature(message: ThreadMessage): Signature {
+    try {
+      return (
+        JSON.stringify({
+          role: message.role,
+          content: message.content,
+          status: message.role === "assistant" ? message.status : undefined,
+          metadata: message.metadata,
+        }) ?? message
+      );
+    } catch {
+      return message;
+    }
+  }
+
+  private branch(): ThreadMessage[] {
+    return (this.thread?.messages ?? [])
+      .filter(
+        (message) =>
+          !message.metadata.isOptimistic &&
+          !message.id.startsWith(FALLBACK_ID_PREFIX) &&
+          !isErrorMessageId(message.id) &&
+          !(message.role === "assistant" && message.status.type === "running"),
+      )
+      .map((message) => {
+        const logs = this.interactions.get(message.id);
+        if (!logs) return message;
+        const cached = this.overlaid.get(message);
+        if (
+          cached &&
+          cached.logs.size === logs.size &&
+          [...logs].every(([id, log]) => cached.logs.get(id) === log)
+        ) {
+          return cached.message;
+        }
+        const copy = {
+          ...message,
+          content: message.content.map((part) => {
+            if (part.type !== "tool-call") return part;
+            const log = logs.get(part.toolCallId);
+            return log ? { ...part, unstable_interactions: log } : part;
+          }),
+        } as ThreadMessage;
+        this.overlaid.set(message, { logs: new Map(logs), message: copy });
+        return copy;
+      });
+  }
+
+  private seed(): void {
+    for (const message of this.branch()) {
+      if (!this.copied.has(message.id) && !this.failedIds.has(message.id)) {
+        this.copied.set(message.id, this.signature(message));
+      }
+    }
+  }
+
+  public attach(thread: ThreadRuntimeCore, history: ThreadHistoryAdapter) {
+    if (history !== this.lastHistory) {
+      this.copied.clear();
+      this.failedIds.clear();
+      this.interactions.clear();
+      this.overlaid = new WeakMap();
+      this.lastHistory = history;
+      this.session += 1;
+    }
+    this.thread = thread;
+    this.history = history;
+    this.seed();
+    const offStart = thread.unstable_on("runStart", () => this.seed());
+    const offEnd = thread.unstable_on("runEnd", () => this.schedule());
+    return () => {
+      offStart();
+      offEnd();
+      if (this.timer !== undefined) clearTimeout(this.timer);
+      this.timer = undefined;
+      this.pending = false;
+      const detached = new Error("History copy was detached.");
+      for (const waiter of this.waiters.splice(0)) waiter.reject(detached);
+      this.thread = undefined;
+      this.history = undefined;
+    };
+  }
+
+  public recordInteraction = async (
+    options: Unstable_RecordToolInteractionOptions,
+  ): Promise<void> => {
+    const thread = this.thread;
+    const message = thread?.messages.find(
+      (item) => item.id === options.messageId,
+    );
+    const part = message?.content.find(
+      (item) =>
+        item.type === "tool-call" && item.toolCallId === options.toolCallId,
+    );
+    if (!thread || !message || part?.type !== "tool-call") {
+      throw new Error("Tool call is not available.");
+    }
+
+    const isRunning = getThreadRuntimeCoreIsRunning(thread);
+    if (!isRunning) this.seed();
+    let logs = this.interactions.get(message.id);
+    if (!logs) {
+      logs = new Map();
+      this.interactions.set(message.id, logs);
+    }
+    logs.set(
+      options.toolCallId,
+      appendToolInteraction(
+        logs.get(options.toolCallId) ??
+          readToolInteractionLog(part.unstable_interactions),
+        options.interaction,
+      ),
+    );
+
+    if (!isRunning) await this.schedule(true);
+  };
+
+  private schedule(wait = false): Promise<void> {
+    this.pending = true;
+    const result = wait
+      ? new Promise<void>((resolve, reject) => {
+          this.waiters.push({ resolve, reject });
+        })
+      : Promise.resolve();
+    if (!this.inFlight && this.timer === undefined) {
+      this.timer = setTimeout(() => {
+        this.timer = undefined;
+        void this.flush();
+      }, 0);
+    }
+    return result;
+  }
+
+  private async flush(): Promise<void> {
+    if (!this.pending || !this.history) return;
+    const history = this.history;
+    const session = this.session;
+    this.pending = false;
+    this.inFlight = true;
+    const waiters = this.waiters.splice(0);
+    const branch = this.branch();
+    const changed = branch
+      .map((message) => [message.id, this.signature(message)] as const)
+      .filter(([id, signature]) => this.copied.get(id) !== signature);
+
+    try {
+      if (changed.length > 0) {
+        await history.unstable_copy?.(
+          branch,
+          changed.map(([id]) => id),
+        );
+        if (session === this.session) {
+          for (const [id, signature] of changed) {
+            this.copied.set(id, signature);
+            this.failedIds.delete(id);
+          }
+        }
+      }
+      for (const waiter of waiters) waiter.resolve();
+    } catch (error) {
+      if (session === this.session) {
+        for (const [id] of changed) this.failedIds.add(id);
+      }
+      if (!this.warned) {
+        this.warned = true;
+        console.warn(
+          "[useExternalStoreRuntime] Failed to copy history.",
+          error,
+        );
+      }
+      for (const waiter of waiters) waiter.reject(error);
+    } finally {
+      this.inFlight = false;
+      if (this.pending) void this.schedule();
+    }
+  }
+}