@assistant-ui/react-ai-sdk 1.3.7 → 1.3.9

package/src/ui/use-chat/useExternalHistory.ts

@@ -0,0 +1,283 @@
+"use client";
+
+import {
+  AssistantRuntime,
+  ThreadHistoryAdapter,
+  ThreadMessage,
+  MessageFormatAdapter,
+  getExternalStoreMessages,
+  MessageFormatRepository,
+  ExportedMessageRepository,
+  INTERNAL,
+  useAui,
+} from "@assistant-ui/react";
+import { useRef, useEffect, useState, RefObject, useCallback } from "react";
+
+const { MessageRepository } = INTERNAL;
+
+export const toExportedMessageRepository = <TMessage>(
+  toThreadMessages: (messages: TMessage[]) => ThreadMessage[],
+  messages: MessageFormatRepository<TMessage>,
+): ExportedMessageRepository => {
+  return {
+    headId: messages.headId!,
+    messages: messages.messages.map((m) => {
+      const message = toThreadMessages([m.message])[0]!;
+      return {
+        ...m,
+        message,
+      };
+    }),
+  };
+};
+
+export const useExternalHistory = <TMessage>(
+  runtimeRef: RefObject<AssistantRuntime>,
+  historyAdapter: ThreadHistoryAdapter | undefined,
+  toThreadMessages: (messages: TMessage[]) => ThreadMessage[],
+  storageFormatAdapter: MessageFormatAdapter<TMessage, any>,
+  onSetMessages: (messages: TMessage[]) => void,
+) => {
+  const loadedRef = useRef(false);
+
+  const aui = useAui();
+  const optionalThreadListItem = useCallback(
+    () => (aui.threadListItem.source ? aui.threadListItem() : null),
+    [aui],
+  );
+
+  const [isLoading, setIsLoading] = useState(false);
+
+  const historyIds = useRef(new Set<string>());
+
+  const onSetMessagesRef = useRef(onSetMessages);
+  useEffect(() => {
+    onSetMessagesRef.current = onSetMessages;
+  });
+
+  useEffect(() => {
+    if (!historyAdapter || loadedRef.current) return;
+
+    const loadHistory = async () => {
+      setIsLoading(true);
+      try {
+        const repo = await historyAdapter
+          .withFormat?.(storageFormatAdapter)
+          .load();
+        if (repo && repo.messages.length > 0) {
+          const converted = toExportedMessageRepository(toThreadMessages, repo);
+          runtimeRef.current.thread.import(converted);
+
+          const tempRepo = new MessageRepository();
+          tempRepo.import(converted);
+          const messages = tempRepo.getMessages();
+
+          onSetMessagesRef.current(
+            messages.map(getExternalStoreMessages<TMessage>).flat(),
+          );
+
+          historyIds.current = new Set(
+            converted.messages.map((m) => m.message.id),
+          );
+        }
+      } catch (error) {
+        console.error("Failed to load message history:", error);
+      } finally {
+        setIsLoading(false);
+      }
+    };
+
+    loadedRef.current = true;
+
+    if (!optionalThreadListItem()?.getState().remoteId) {
+      setIsLoading(false);
+      return;
+    }
+
+    loadHistory();
+  }, [
+    historyAdapter,
+    storageFormatAdapter,
+    toThreadMessages,
+    runtimeRef,
+    optionalThreadListItem,
+  ]);
+
+  const runStartRef = useRef<number | null>(null);
+  const persistTimerRef = useRef<ReturnType<typeof setTimeout> | null>(null);
+  const stepBoundariesRef = useRef<number[]>([]);
+  const wasRunningRef = useRef(false);
+  const toolCallCountRef = useRef(0);
+
+  useEffect(() => {
+    const unsubscribe = runtimeRef.current.thread.subscribe(() => {
+      const { isRunning } = runtimeRef.current.thread.getState();
+      const wasRunning = wasRunningRef.current;
+      wasRunningRef.current = isRunning;
+
+      // Track step boundaries by content changes (more reliable than isRunning)
+      if (runStartRef.current != null) {
+        const lastMsg = runtimeRef.current.thread.getState().messages.at(-1);
+        if (lastMsg?.role === "assistant") {
+          const currentToolCallCount = lastMsg.content.filter(
+            (p) => p.type === "tool-call",
+          ).length;
+          while (toolCallCountRef.current < currentToolCallCount) {
+            stepBoundariesRef.current.push(Date.now() - runStartRef.current);
+            toolCallCountRef.current++;
+          }
+        }
+      }
+
+      if (isRunning) {
+        if (runStartRef.current == null) {
+          runStartRef.current = Date.now();
+          stepBoundariesRef.current = [];
+          toolCallCountRef.current = 0;
+        }
+        // Cancel any pending persist — isRunning went back to true
+        if (persistTimerRef.current) {
+          clearTimeout(persistTimerRef.current);
+          persistTimerRef.current = null;
+        }
+        return;
+      }
+
+      // Only act on the true→false transition
+      if (!wasRunning) return;
+
+      // Record step boundary offset (synchronous for accuracy)
+      if (runStartRef.current != null) {
+        stepBoundariesRef.current.push(Date.now() - runStartRef.current);
+      }
+
+      // Debounce: wait one macrotask so agentic step flickers are absorbed
+      if (persistTimerRef.current) clearTimeout(persistTimerRef.current);
+      persistTimerRef.current = setTimeout(async () => {
+        persistTimerRef.current = null;
+
+        // Re-read latest state — may have changed since the timeout was scheduled
+        const latest = runtimeRef.current.thread.getState();
+        if (latest.isRunning) return; // was just a flicker
+
+        // Derive durationMs from the last boundary (covers all steps)
+        const boundaries = stepBoundariesRef.current;
+        const durationMs =
+          boundaries.length > 0 ? boundaries.at(-1) : undefined;
+
+        // Fallback: if only 1 boundary but message has multiple steps, distribute evenly
+        if (boundaries.length === 1 && durationMs != null) {
+          const lastAssistant = latest.messages.findLast(
+            (m) => m.role === "assistant",
+          );
+          if (lastAssistant) {
+            const tcCount = lastAssistant.content.filter(
+              (p) => p.type === "tool-call",
+            ).length;
+            if (tcCount > 0) {
+              const totalSteps = tcCount + 1;
+              const stepDur = durationMs / totalSteps;
+              boundaries.length = 0;
+              for (let i = 0; i < totalSteps; i++) {
+                boundaries.push(Math.round((i + 1) * stepDur));
+              }
+            }
+          }
+        }
+
+        // Build per-step timestamps when there are multiple steps
+        const stepTimestamps =
+          boundaries.length > 1
+            ? boundaries.map((endMs, i) => ({
+                start_ms: i === 0 ? 0 : boundaries[i - 1]!,
+                end_ms: endMs,
+              }))
+            : undefined;
+
+        runStartRef.current = null;
+        stepBoundariesRef.current = [];
+
+        const telemetryOptions = {
+          ...(durationMs != null ? { durationMs } : undefined),
+          ...(stepTimestamps != null ? { stepTimestamps } : undefined),
+        };
+
+        const { messages } = latest;
+        let lastInnerMessageId: string | null = null;
+
+        const getLastInnerId = (msgs: TMessage[]): string | null =>
+          msgs.length > 0 ? storageFormatAdapter.getId(msgs.at(-1)!) : null;
+
+        const toBatchItems = (msgs: TMessage[]) =>
+          msgs.map((msg, idx) => ({
+            parentId:
+              idx === 0
+                ? lastInnerMessageId
+                : storageFormatAdapter.getId(msgs[idx - 1]!),
+            message: msg,
+          }));
+
+        for (const message of messages) {
+          const innerMessages = getExternalStoreMessages<TMessage>(message);
+
+          const isReady =
+            message.status === undefined ||
+            message.status.type === "complete" ||
+            message.status.type === "incomplete";
+
+          if (!isReady) {
+            lastInnerMessageId =
+              getLastInnerId(innerMessages) ?? lastInnerMessageId;
+            continue;
+          }
+
+          if (historyIds.current.has(message.id)) {
+            if (durationMs !== undefined) {
+              const formatAdapter =
+                historyAdapter?.withFormat?.(storageFormatAdapter);
+              let parentId = lastInnerMessageId;
+              for (const innerMessage of innerMessages) {
+                try {
+                  await formatAdapter?.update?.(
+                    { parentId, message: innerMessage },
+                    storageFormatAdapter.getId(innerMessage),
+                  );
+                } catch {
+                  // ignore update failures to avoid breaking the message processing loop
+                }
+                parentId = storageFormatAdapter.getId(innerMessage);
+              }
+            }
+            lastInnerMessageId =
+              getLastInnerId(innerMessages) ?? lastInnerMessageId;
+            continue;
+          }
+          historyIds.current.add(message.id);
+
+          const formatAdapter =
+            historyAdapter?.withFormat?.(storageFormatAdapter);
+
+          const batchItems = toBatchItems(innerMessages);
+          for (const item of batchItems) {
+            await formatAdapter?.append(item);
+          }
+
+          lastInnerMessageId =
+            getLastInnerId(innerMessages) ?? lastInnerMessageId;
+
+          formatAdapter?.reportTelemetry?.(batchItems, telemetryOptions);
+        }
+      }, 0);
+    });
+
+    return () => {
+      unsubscribe();
+      if (persistTimerRef.current) {
+        clearTimeout(persistTimerRef.current);
+        persistTimerRef.current = null;
+      }
+    };
+  }, [historyAdapter, storageFormatAdapter, runtimeRef]);
+
+  return isLoading;
+};

package/src/ui/use-chat/useStreamingTiming.ts

@@ -0,0 +1,102 @@
+"use client";
+
+import { useEffect, useRef, useState } from "react";
+import type { UIMessage } from "@ai-sdk/react";
+import { isToolUIPart } from "ai";
+import type { MessageTiming } from "@assistant-ui/react";
+
+type TrackingState = {
+  messageId: string;
+  startTime: number;
+  firstTokenTime?: number;
+  lastContentLength: number;
+  totalChunks: number;
+};
+
+function getTextLength(message: UIMessage | undefined): number {
+  if (!message?.parts) return 0;
+  let len = 0;
+  for (const part of message.parts) {
+    if (part.type === "text") len += part.text.length;
+  }
+  return len;
+}
+
+function getToolCallCount(message: UIMessage | undefined): number {
+  if (!message?.parts) return 0;
+  let count = 0;
+  for (const part of message.parts) {
+    if (isToolUIPart(part)) count++;
+  }
+  return count;
+}
+
+/**
+ * Tracks streaming timing for AI SDK messages client-side.
+ *
+ * Observes `isRunning` transitions and content changes to estimate
+ * timing metrics (TTFT, duration, tok/s). Timing is finalized when
+ * streaming ends and stored per message ID.
+ */
+export const useStreamingTiming = (
+  messages: UIMessage[],
+  isRunning: boolean,
+): Record<string, MessageTiming> => {
+  const [timings, setTimings] = useState<Record<string, MessageTiming>>({});
+  const trackRef = useRef<TrackingState | null>(null);
+
+  useEffect(() => {
+    const lastAssistant = messages.findLast((m) => m.role === "assistant");
+
+    if (isRunning && lastAssistant) {
+      // Start tracking if not already or if message changed
+      if (
+        !trackRef.current ||
+        trackRef.current.messageId !== lastAssistant.id
+      ) {
+        trackRef.current = {
+          messageId: lastAssistant.id,
+          startTime: Date.now(),
+          lastContentLength: 0,
+          totalChunks: 0,
+        };
+      }
+
+      // Track content changes
+      const t = trackRef.current;
+      const len = getTextLength(lastAssistant);
+      if (len > t.lastContentLength) {
+        if (t.firstTokenTime === undefined) {
+          t.firstTokenTime = Date.now() - t.startTime;
+        }
+        t.totalChunks++;
+        t.lastContentLength = len;
+      }
+    } else if (!isRunning && trackRef.current) {
+      // Streaming ended — finalize timing
+      const t = trackRef.current;
+      const totalStreamTime = Date.now() - t.startTime;
+      const tokenCount = Math.ceil(t.lastContentLength / 4);
+      const toolCallCount = getToolCallCount(lastAssistant);
+
+      const timing: MessageTiming = {
+        streamStartTime: t.startTime,
+        totalStreamTime,
+        totalChunks: t.totalChunks,
+        toolCallCount,
+        ...(t.firstTokenTime !== undefined && {
+          firstTokenTime: t.firstTokenTime,
+        }),
+        ...(tokenCount > 0 && { tokenCount }),
+        ...(totalStreamTime > 0 &&
+          tokenCount > 0 && {
+            tokensPerSecond: tokenCount / (totalStreamTime / 1000),
+          }),
+      };
+      setTimings((prev) => ({ ...prev, [t.messageId]: timing }));
+      trackRef.current = null;
+    }
+  }, [messages, isRunning]);
+
+  return timings;
+};

package/src/ui/utils/convertMessage.test.ts

@@ -0,0 +1,125 @@
+import { describe, expect, it } from "vitest";
+import { AISDKMessageConverter } from "./convertMessage";
+
+describe("AISDKMessageConverter", () => {
+  it("converts user files into attachments and keeps text content", () => {
+    const converted = AISDKMessageConverter.toThreadMessages([
+      {
+        id: "u1",
+        role: "user",
+        parts: [
+          { type: "text", text: "hello" },
+          {
+            type: "file",
+            mediaType: "image/png",
+            url: "https://cdn/img.png",
+            filename: "img.png",
+          },
+          {
+            type: "file",
+            mediaType: "application/pdf",
+            url: "https://cdn/file.pdf",
+            filename: "file.pdf",
+          },
+        ],
+      } as any,
+    ]);
+
+    expect(converted).toHaveLength(1);
+    expect(converted[0]?.role).toBe("user");
+    expect(converted[0]?.content).toHaveLength(1);
+    expect(converted[0]?.content[0]).toMatchObject({
+      type: "text",
+      text: "hello",
+    });
+    expect(converted[0]?.attachments).toHaveLength(2);
+    expect(converted[0]?.attachments?.[0]?.type).toBe("image");
+    expect(converted[0]?.attachments?.[1]?.type).toBe("file");
+  });
+
+  it("deduplicates tool calls by toolCallId and maps interrupt states", () => {
+    const converted = AISDKMessageConverter.toThreadMessages(
+      [
+        {
+          id: "a1",
+          role: "assistant",
+          parts: [
+            {
+              type: "tool-weather",
+              toolCallId: "tc-1",
+              state: "output-available",
+              input: { city: "NYC" },
+              output: { temp: 72 },
+            },
+            {
+              type: "tool-weather",
+              toolCallId: "tc-1",
+              state: "output-available",
+              input: { city: "NYC" },
+              output: { temp: 73 },
+            },
+            {
+              type: "tool-approve",
+              toolCallId: "tc-2",
+              state: "approval-requested",
+              input: { action: "deploy" },
+              approval: { reason: "need human review" },
+            },
+            {
+              type: "tool-human",
+              toolCallId: "tc-3",
+              state: "input-available",
+              input: { task: "confirm" },
+            },
+          ],
+        } as any,
+      ],
+      false,
+      {
+        toolStatuses: {
+          "tc-3": {
+            type: "interrupt",
+            payload: { type: "human", payload: { kind: "human" } },
+          },
+        },
+      },
+    );
+
+    const toolCalls = converted[0]?.content.filter(
+      (part): part is any => part.type === "tool-call",
+    );
+    expect(toolCalls).toHaveLength(3);
+
+    expect(toolCalls?.filter((p) => p.toolCallId === "tc-1")).toHaveLength(1);
+    expect(toolCalls?.find((p) => p.toolCallId === "tc-2")?.status).toEqual({
+      type: "requires-action",
+      reason: "interrupt",
+    });
+    expect(toolCalls?.find((p) => p.toolCallId === "tc-3")?.interrupt).toEqual({
+      type: "human",
+      payload: { kind: "human" },
+    });
+  });
+
+  it("strips closing delimiters from streaming tool argsText", () => {
+    const converted = AISDKMessageConverter.toThreadMessages([
+      {
+        id: "a1",
+        role: "assistant",
+        parts: [
+          {
+            type: "tool-weather",
+            toolCallId: "tc-1",
+            state: "input-streaming",
+            input: { city: "NYC" },
+          },
+        ],
+      } as any,
+    ]);
+
+    const toolCall = converted[0]?.content.find(
+      (part): part is any => part.type === "tool-call",
+    );
+    expect(toolCall?.argsText).toBe('{"city":"NYC');
+  });
+});