@assistant-ui/core 0.2.18 → 0.2.20

package/src/runtime/utils/message-repository.ts

@@ -189,6 +189,16 @@ export class MessageRepository {
     return this.head?.current.id ?? null;
   }
 
+  get canonicalHeadId() {
+    // Optimistic messages are ephemeral, so persisted callers need the nearest
+    // non-optimistic ancestor rather than the raw head.
+    let head = this.head;
+    while (head?.current.metadata?.isOptimistic) {
+      head = head.prev;
+    }
+    return head?.current.id ?? null;
+  }
+
   getMessages(headId?: string) {
     if (headId === undefined || headId === this.head?.current.id) {
       return this._messages.value;
@@ -436,14 +446,8 @@ export class MessageRepository {
       });
     }
 
-    // The head may itself be optimistic; walk up to the nearest persisted ancestor.
-    let head = this.head;
-    while (head?.current.metadata?.isOptimistic) {
-      head = head.prev;
-    }
-
     return {
-      headId: head?.current.id ?? null,
+      headId: this.canonicalHeadId,
       messages: exportItems,
     };
   }

package/src/runtime/utils/streaming-timing.test.ts

@@ -0,0 +1,301 @@
+import { describe, expect, it } from "vitest";
+import {
+  stepStreamingTiming,
+  type StreamingTimingAccessors,
+} from "./streaming-timing";
+
+type TestMessage = {
+  readonly id: string;
+  readonly role: "assistant" | "user";
+  readonly text: string;
+  readonly toolCalls: number;
+};
+
+const accessors: StreamingTimingAccessors<TestMessage> = {
+  getAssistantMessageId: (messages) => {
+    for (let i = messages.length - 1; i >= 0; i--) {
+      if (messages[i]!.role === "assistant") return messages[i]!.id;
+    }
+    return undefined;
+  },
+  getTextLength: (messages, id) =>
+    messages.find((m) => m.id === id)?.text.length ?? 0,
+  getToolCallCount: (messages, id) =>
+    messages.find((m) => m.id === id)?.toolCalls ?? 0,
+};
+
+const makeClock = () => {
+  let t = 0;
+  return () => (t += 10);
+};
+
+const msg = (id: string, text: string, toolCalls = 0): TestMessage => ({
+  id,
+  role: "assistant",
+  text,
+  toolCalls,
+});
+
+describe("stepStreamingTiming", () => {
+  it("is idle when not running and there is no state", () => {
+    const now = makeClock();
+    expect(
+      stepStreamingTiming(null, [], false, accessors, undefined, now),
+    ).toEqual({ state: null, timings: {} });
+  });
+
+  it("starts tracking on the first running update and counts content growth as chunks", () => {
+    const now = makeClock();
+    const r1 = stepStreamingTiming(
+      null,
+      [msg("m1", "")],
+      true,
+      accessors,
+      undefined,
+      now,
+    );
+    // No growth yet (length 0), state is initialised but no chunks.
+    expect(r1.timings).toEqual({});
+    expect(r1.state).toMatchObject({ messageId: "m1", totalChunks: 0 });
+
+    const r2 = stepStreamingTiming(
+      r1.state,
+      [msg("m1", "Hello")],
+      true,
+      accessors,
+      undefined,
+      now,
+    );
+    expect(r2.state).toMatchObject({
+      messageId: "m1",
+      totalChunks: 1,
+      lastContentLength: 5,
+      firstTokenTime: 10,
+    });
+    expect(r2.timings).toEqual({});
+  });
+
+  it("finalizes timing when streaming stops and counts the final chunk", () => {
+    const now = makeClock();
+    const r1 = stepStreamingTiming(
+      null,
+      [msg("m1", "a")],
+      true,
+      accessors,
+      undefined,
+      now,
+    );
+    const r2 = stepStreamingTiming(
+      r1.state,
+      [msg("m1", "ab")],
+      true,
+      accessors,
+      undefined,
+      now,
+    );
+    // Final content delta arrives in the SAME update as isRunning -> false.
+    // The finalize re-reads the final length (3) rather than the last delta.
+    const r3 = stepStreamingTiming(
+      r2.state,
+      [msg("m1", "abc")],
+      false,
+      accessors,
+      undefined,
+      now,
+    );
+    expect(r3.state).toBeNull();
+    const timing = r3.timings["m1"]!;
+    expect(timing).toBeDefined();
+    // The final delta ("abc", arrived with the stop signal) is reconciled,
+    // so totalChunks and tokenCount count the same three-character content.
+    expect(timing.totalChunks).toBe(3);
+    expect(timing.tokenCount).toBe(Math.ceil(3 / 4));
+    expect(timing.toolCallCount).toBe(0);
+    expect(timing.tokensPerSecond).toBeCloseTo(
+      timing.tokenCount! / (timing.totalStreamTime! / 1000),
+    );
+  });
+
+  it("reconciles totalChunks and tokenCount when the only growth lands at finalize", () => {
+    const now = makeClock();
+    // Running the whole time with empty content; the text arrives only with
+    // the stop signal.
+    const r1 = stepStreamingTiming(
+      null,
+      [msg("m1", "")],
+      true,
+      accessors,
+      undefined,
+      now,
+    );
+    const r2 = stepStreamingTiming(
+      r1.state,
+      [msg("m1", "done")],
+      false,
+      accessors,
+      undefined,
+      now,
+    );
+    const timing = r2.timings["m1"]!;
+    expect(timing.totalChunks).toBe(1);
+    expect(timing.tokenCount).toBe(Math.ceil(4 / 4));
+    expect(timing.firstTokenTime).toBe(timing.totalStreamTime);
+  });
+
+  it("does not finalize while still running", () => {
+    const now = makeClock();
+    const r1 = stepStreamingTiming(
+      null,
+      [msg("m1", "content")],
+      true,
+      accessors,
+      undefined,
+      now,
+    );
+    const r2 = stepStreamingTiming(
+      r1.state,
+      [msg("m1", "content")],
+      true,
+      accessors,
+      undefined,
+      now,
+    );
+    expect(r2.timings).toEqual({});
+    expect(r2.state).not.toBeNull();
+  });
+
+  it("resets tracking when the assistant message id changes mid-stream", () => {
+    const now = makeClock();
+    const r1 = stepStreamingTiming(
+      null,
+      [msg("m1", "hi")],
+      true,
+      accessors,
+      undefined,
+      now,
+    );
+    const r2 = stepStreamingTiming(
+      r1.state,
+      [msg("m1", "hi"), msg("m2", "")],
+      true,
+      accessors,
+      undefined,
+      now,
+    );
+    expect(r2.state?.messageId).toBe("m2");
+    expect(r2.state?.totalChunks).toBe(0);
+    expect(r2.state?.lastContentLength).toBe(0);
+  });
+
+  it("records firstTokenTime only once, at the first content growth", () => {
+    const now = makeClock();
+    const r1 = stepStreamingTiming(
+      null,
+      [msg("m1", "")],
+      true,
+      accessors,
+      undefined,
+      now,
+    );
+    const r2 = stepStreamingTiming(
+      r1.state,
+      [msg("m1", "first")],
+      true,
+      accessors,
+      undefined,
+      now,
+    );
+    expect(r2.state?.firstTokenTime).toBe(10);
+    const r3 = stepStreamingTiming(
+      r2.state,
+      [msg("m1", "first second")],
+      true,
+      accessors,
+      undefined,
+      now,
+    );
+    // firstTokenTime unchanged on later growth.
+    expect(r3.state?.firstTokenTime).toBe(10);
+    expect(r3.state?.totalChunks).toBe(2);
+  });
+
+  it("counts tool calls from the final message", () => {
+    const now = makeClock();
+    const r1 = stepStreamingTiming(
+      null,
+      [msg("m1", "a", 2)],
+      true,
+      accessors,
+      undefined,
+      now,
+    );
+    const r2 = stepStreamingTiming(
+      r1.state,
+      [msg("m1", "ab", 2)],
+      false,
+      accessors,
+      undefined,
+      now,
+    );
+    expect(r2.timings["m1"]?.toolCallCount).toBe(2);
+  });
+
+  it("omits tokenCount and tokensPerSecond when the final text is empty", () => {
+    const now = makeClock();
+    const r1 = stepStreamingTiming(
+      null,
+      [msg("m1", "")],
+      true,
+      accessors,
+      undefined,
+      now,
+    );
+    const r2 = stepStreamingTiming(
+      r1.state,
+      [msg("m1", "")],
+      false,
+      accessors,
+      undefined,
+      now,
+    );
+    const timing = r2.timings["m1"]!;
+    expect(timing.tokenCount).toBeUndefined();
+    expect(timing.tokensPerSecond).toBeUndefined();
+    expect(timing.totalChunks).toBe(0);
+  });
+
+  it("respects a custom estimateTokens option", () => {
+    const now = makeClock();
+    const r1 = stepStreamingTiming(
+      null,
+      [msg("m1", "abcdefgh")],
+      true,
+      accessors,
+      { estimateTokens: (n) => Math.floor(n / 2) },
+      now,
+    );
+    const r2 = stepStreamingTiming(
+      r1.state,
+      [msg("m1", "abcdefgh")],
+      false,
+      accessors,
+      { estimateTokens: (n) => Math.floor(n / 2) },
+      now,
+    );
+    expect(r2.timings["m1"]?.tokenCount).toBe(4);
+  });
+
+  it("produces no timings when running but no assistant message is present", () => {
+    const now = makeClock();
+    const r = stepStreamingTiming(
+      null,
+      [{ id: "u1", role: "user", text: "hi", toolCalls: 0 }],
+      true,
+      accessors,
+      undefined,
+      now,
+    );
+    expect(r.state).toBeNull();
+    expect(r.timings).toEqual({});
+  });
+});

package/src/runtime/utils/streaming-timing.ts

@@ -0,0 +1,161 @@
+import type { MessageTiming } from "../../types/message";
+
+/**
+ * Shape-specific accessors that adapt {@link stepStreamingTiming} /
+ * {@link useStreamingTiming} to a runtime's message list. Each adapter
+ * provides three pure functions over its own message type; the timing state
+ * machine itself lives once in core.
+ */
+export type StreamingTimingAccessors<TMessage> = {
+  /** Resolve the id of the last assistant message, or `undefined` if none. */
+  readonly getAssistantMessageId: (
+    messages: readonly TMessage[],
+  ) => string | undefined;
+  /** Total text length (incl. reasoning/thinking) of the assistant message. */
+  readonly getTextLength: (
+    messages: readonly TMessage[],
+    messageId: string,
+  ) => number;
+  /** Number of tool calls recorded on the assistant message. */
+  readonly getToolCallCount: (
+    messages: readonly TMessage[],
+    messageId: string,
+  ) => number;
+};
+
+export type StreamingTimingOptions = {
+  /**
+   * Estimate token count from a text length. Defaults to `Math.ceil(n / 4)`.
+   * Adapters with real usage data can override this.
+   */
+  readonly estimateTokens?: (textLength: number) => number;
+};
+
+/**
+ * Mutable tracking state for a single in-flight assistant message. Kept
+ * between updates by the caller (a ref in the React hook); `null` means no
+ * message is being tracked.
+ *
+ * `totalChunks` counts content growth deltas observed while streaming. The
+ * final delta that lands in the same update as the `isRunning -> false`
+ * transition is reconciled at finalize, so `totalChunks` and the token
+ * estimate (derived from final length) count the same content.
+ */
+export type StreamingTimingState = {
+  readonly messageId: string;
+  readonly startTime: number;
+  readonly firstTokenTime?: number;
+  readonly lastContentLength: number;
+  readonly totalChunks: number;
+};
+
+const defaultEstimateTokens = (textLength: number): number =>
+  Math.ceil(textLength / 4);
+
+/**
+ * Apply one content growth delta to `state`, recording the first-token time
+ * on the first growth and bumping the chunk count. Returns `state` unchanged
+ * when `len` has not grown.
+ */
+const applyGrowth = (
+  state: StreamingTimingState,
+  len: number,
+  now: () => number,
+): StreamingTimingState => {
+  if (len <= state.lastContentLength) return state;
+  return {
+    ...state,
+    firstTokenTime: state.firstTokenTime ?? now() - state.startTime,
+    lastContentLength: len,
+    totalChunks: state.totalChunks + 1,
+  };
+};
+
+/**
+ * Advance the client-side streaming timing tracker by one update.
+ *
+ * Pure: given the previous state and the current `(messages, isRunning)`
+ * snapshot, returns the next state plus any `MessageTiming` finalized this
+ * update (empty when streaming is still in flight). The caller owns the
+ * state (e.g. a ref) and merges finalized timings into its store.
+ *
+ * At finalize the text length is recomputed from the final messages, and any
+ * final growth delta that landed in the same update as the
+ * `isRunning -> false` transition is reconciled, so the token estimate and
+ * `totalChunks` count the same content (the per-adapter hooks this replaces
+ * reuse a stale growth delta and miss that last chunk).
+ */
+export const stepStreamingTiming = <TMessage>(
+  state: StreamingTimingState | null,
+  messages: readonly TMessage[],
+  isRunning: boolean,
+  accessors: StreamingTimingAccessors<TMessage>,
+  options: StreamingTimingOptions | undefined,
+  now: () => number = Date.now,
+): {
+  readonly state: StreamingTimingState | null;
+  readonly timings: Record<string, MessageTiming>;
+} => {
+  const lastId = accessors.getAssistantMessageId(messages);
+  const estimateTokens = options?.estimateTokens ?? defaultEstimateTokens;
+
+  if (isRunning && lastId !== undefined) {
+    let next: StreamingTimingState;
+    if (state === null || state.messageId !== lastId) {
+      next = {
+        messageId: lastId,
+        startTime: now(),
+        lastContentLength: 0,
+        totalChunks: 0,
+      };
+    } else {
+      next = state;
+    }
+
+    return {
+      state: applyGrowth(
+        next,
+        accessors.getTextLength(messages, next.messageId),
+        now,
+      ),
+      timings: {},
+    };
+  }
+
+  if (!isRunning && state !== null) {
+    const nowMs = now();
+    // Reconcile any final growth that landed with the stop signal so the
+    // token estimate (derived from final length) and totalChunks agree.
+    const reconciled = applyGrowth(
+      state,
+      accessors.getTextLength(messages, state.messageId),
+      () => nowMs,
+    );
+
+    const totalStreamTime = nowMs - reconciled.startTime;
+    const tokenCount = estimateTokens(reconciled.lastContentLength);
+    const toolCallCount = accessors.getToolCallCount(
+      messages,
+      reconciled.messageId,
+    );
+
+    const timing: MessageTiming = {
+      streamStartTime: reconciled.startTime,
+      totalStreamTime,
+      totalChunks: reconciled.totalChunks,
+      toolCallCount,
+      ...(reconciled.firstTokenTime !== undefined && {
+        firstTokenTime: reconciled.firstTokenTime,
+      }),
+      ...(tokenCount > 0 && { tokenCount }),
+      ...(totalStreamTime > 0 &&
+        tokenCount > 0 && {
+          tokensPerSecond: tokenCount / (totalStreamTime / 1000),
+        }),
+    };
+
+    return { state: null, timings: { [reconciled.messageId]: timing } };
+  }
+
+  return { state, timings: {} };
+};

package/src/runtime/utils/thread-message-like.ts

@@ -127,6 +127,7 @@ export const fromThreadMessageLike = (
     image,
     ...rest
   }: ImageMessagePart): ImageMessagePart | null => {
+    if (typeof image !== "string") return null;
     const dataUri = image.match(
       /^data:image\/(png|jpeg|jpg|gif|webp|svg\+xml);base64,(.*)$/,
     );
@@ -160,7 +161,7 @@ export const fromThreadMessageLike = (
             switch (type) {
               case "text":
               case "reasoning":
-                if (part.text.trim().length === 0) return null;
+                if (!part.text?.trim()) return null;
                 return part;
 
               case "file":

package/src/runtimes/external-store/external-store-adapter.ts

@@ -64,6 +64,14 @@ export type ExternalStoreMessageConverter<T> = (
   idx: number,
 ) => ThreadMessageLike;
 
+/**
+ * @deprecated This API is still under active development and might change without notice.
+ */
+export type ExternalStoreBranchChange = {
+  headId: string | null;
+  visibleMessageIds: readonly string[];
+};
+
 type ExternalStoreMessageConverterAdapter<T> = {
   convertMessage: ExternalStoreMessageConverter<T>;
 };
@@ -102,6 +110,26 @@ type ExternalStoreAdapterBase<T> = {
   extras?: unknown;
 
   setMessages?: ((messages: readonly T[]) => void) | undefined;
+  /**
+   * Fires when the user explicitly switches branches via the runtime's
+   * `switchToBranch` action (e.g. a BranchPicker click). It does not fire on
+   * adapter resync, `append`, edit/regenerate, content-only updates, or while
+   * the thread is running. Consecutive switches that resolve to the same
+   * canonical head are de-duped.
+   *
+   * `headId` is the canonical (persisted) head of the now-visible branch —
+   * optimistic/transient ids are never surfaced. `visibleMessageIds` lists the
+   * visible path in order.
+   *
+   * This complements `setMessages` rather than replacing it: switching still
+   * requires `setMessages`, and this callback does not on its own enable branch
+   * switching.
+   *
+   * @deprecated This API is still under active development and might change without notice.
+   */
+  unstable_onBranchChange?:
+    | ((event: ExternalStoreBranchChange) => void)
+    | undefined;
   onImport?: ((messages: readonly ThreadMessage[]) => void) | undefined;
   onExportExternalState?: (() => any) | undefined;
   onLoadExternalState?: ((state: any) => void) | undefined;

package/src/runtimes/external-store/external-store-thread-runtime-core.ts

@@ -8,7 +8,10 @@ import type {
   ThreadSuggestion,
 } from "../../runtime/interfaces/thread-runtime-core";
 
-import type { ExternalStoreAdapter } from "./external-store-adapter";
+import type {
+  ExternalStoreAdapter,
+  ExternalStoreBranchChange,
+} from "./external-store-adapter";
 import {
   getExternalStoreMessages,
   bindExternalStoreMessage,
@@ -405,8 +408,37 @@ export class ExternalStoreThreadRuntimeCore
       return;
     }
 
+    const onBranchChange = this._store.unstable_onBranchChange;
+    const previousHeadId = onBranchChange
+      ? this.repository.canonicalHeadId
+      : null;
+
     this.repository.switchToBranch(branchId);
     this.updateMessages(this.repository.getMessages());
+    if (onBranchChange) {
+      this._notifyBranchChange(previousHeadId, onBranchChange);
+    }
+  }
+
+  /**
+   * Emit `unstable_onBranchChange` for an explicit branch switch. Reads the
+   * canonical head from the repository (which skips optimistic/transient
+   * messages) and de-dupes switches that leave the canonical head unchanged.
+   * Comparing against the head observed just before the switch — rather than the
+   * last emitted head — keeps a switch firing after an adapter resync moved the
+   * head elsewhere in the meantime.
+   */
+  private _notifyBranchChange(
+    previousHeadId: string | null,
+    onBranchChange: (event: ExternalStoreBranchChange) => void,
+  ): void {
+    const headId = this.repository.canonicalHeadId;
+    if (headId === previousHeadId) return;
+
+    onBranchChange({
+      headId,
+      visibleMessageIds: this.repository.getMessages().map((m) => m.id),
+    });
   }
 
   public async append(message: AppendMessage): Promise<void> {

package/src/runtimes/local/local-thread-runtime-core.ts

@@ -493,7 +493,10 @@ export class LocalThreadRuntimeCore
 
     try {
       this._lastRunConfig = runConfig ?? {};
-      const context = this.getModelContext();
+      // unstable_composerMetadata is composer-only (stamped onto the outgoing
+      // message); never expose it to the chat-model adapter's run context.
+      const { unstable_composerMetadata: _, ...context } =
+        this.getModelContext();
 
       runCallback =
         runCallback ??

package/src/runtimes/remote-thread-list/types.ts

@@ -77,6 +77,19 @@ export type RemoteThreadListOptions = {
    */
   threadId?: string | undefined;
 
+  /**
+   * Called whenever the active thread's canonical (remote) ID changes, so the
+   * value can be treated as a managed/controlled variable (e.g. synced to a
+   * URL query param). Together with `threadId` this forms the controlled
+   * pattern: `threadId` in, `onThreadIdChange` out.
+   *
+   * Only the settled remote ID is emitted: while a freshly created thread is
+   * still optimistic (no remote ID yet) the value is `undefined`, and the real
+   * ID is emitted once the thread is initialized. The transient local ID is
+   * never surfaced.
+   */
+  onThreadIdChange?: ((threadId: string | undefined) => void) | undefined;
+
   /**
    * When true, if this runtime is used inside another RemoteThreadListRuntime,
    * it becomes a no-op and simply calls the runtimeHook directly.

package/src/tests/MessageRepository.test.ts

@@ -302,6 +302,8 @@ describe("MessageRepository", () => {
       const exported = repository.export();
 
       expect(exported.messages.map((m) => m.message.id)).toEqual(["u"]);
+      expect(repository.headId).toBe("placeholder");
+      expect(repository.canonicalHeadId).toBe("u");
       // head was the optimistic placeholder; the exported head must fall back
       // to the nearest persisted ancestor so it always resolves on import.
       expect(exported.headId).toBe("u");

package/src/tests/composer-can-send.test.ts

@@ -15,6 +15,7 @@ const makeRuntimeStub = (
   const stub = {
     append: vi.fn(),
     cancelRun: vi.fn(),
+    getModelContext: () => ({}),
     subscribe: (cb: () => void) => {
       subscribers.add(cb);
       return () => subscribers.delete(cb);