@assistant-ui/ai-sdk 0.0.3 → 0.0.4

package/src/runtime/useChatRuntime.integration.test.tsx

@@ -8,6 +8,7 @@ import { StrictMode, useState } from "react";
 import { describe, expect, it } from "vitest";
 import { AssistantChatTransport } from "../transport/AssistantChatTransport";
 import { useChatRuntime } from "./useChatRuntime";
+import { useThreadTokenUsage } from "../usage";
 
 const messages: UIMessage[] = [
   {
@@ -65,3 +66,48 @@ describe("useChatRuntime integration", () => {
     });
   });
 });
+
+const UsageProbe = () => {
+  const usage = useThreadTokenUsage();
+  return (
+    <output data-testid="total-tokens">{usage?.totalTokens ?? "none"}</output>
+  );
+};
+
+const UsageApp = () => {
+  const [transport] = useState(
+    () => new AssistantChatTransport({ api: "/api/chat" }),
+  );
+  const runtime = useChatRuntime({
+    messages: [
+      ...messages,
+      {
+        id: "assistant-with-usage",
+        role: "assistant",
+        parts: [{ type: "text", text: "Hi" }],
+        metadata: { usage: { inputTokens: 40, outputTokens: 2 } },
+      },
+    ],
+    transport,
+  });
+
+  return (
+    <AssistantRuntimeProvider runtime={runtime}>
+      <UsageProbe />
+    </AssistantRuntimeProvider>
+  );
+};
+
+describe("useThreadTokenUsage through useChatRuntime", () => {
+  it("reads usage from the message metadata a server attached", async () => {
+    render(
+      <StrictMode>
+        <UsageApp />
+      </StrictMode>,
+    );
+
+    await waitFor(() => {
+      expect(screen.getByTestId("total-tokens").textContent).toBe("42");
+    });
+  });
+});

package/src/runtime/useChatThread.ts

@@ -1,6 +1,7 @@
 "use client";
 
 import { useChat, type Chat, type UIMessage } from "@ai-sdk/react";
+import type { MessageRepository } from "@assistant-ui/core/internal";
 import {
   pickExternalStoreSharedOptions,
   type AssistantRuntime,
@@ -60,6 +61,12 @@ export type ChatThreadEnvironment<UI_MESSAGE extends UIMessage = UIMessage> = {
    * from the instance.
    */
   chat?: Chat<UI_MESSAGE> | undefined;
+  /**
+   * An externally owned per-thread message repository. Hosts that route
+   * multiple threads through one mounting pass a distinct instance per
+   * thread so histories and branches stay isolated.
+   */
+  messageRepositoryInstance?: MessageRepository | undefined;
 };
 
 const useDynamicChatTransport = <UI_MESSAGE extends UIMessage = UIMessage>(
@@ -183,6 +190,7 @@ export const useChatThread = <UI_MESSAGE extends UIMessage = UIMessage>(
     getThreadListItem,
     stopOnClientDestroy = false,
     chat: externalChat,
+    messageRepositoryInstance,
   } = env;
 
   const defaultTransport = useMemo(() => new AssistantChatTransport(), []);
@@ -209,6 +217,9 @@ export const useChatThread = <UI_MESSAGE extends UIMessage = UIMessage>(
     ...(onResumeToolCall && { onResumeToolCall }),
     ...(joinStrategy && { joinStrategy }),
     ...(messageRepository && { messageRepository }),
+    ...(messageRepositoryInstance && {
+      unstable_messageRepositoryInstance: messageRepositoryInstance,
+    }),
     ...(unstable_onBranchChange && { unstable_onBranchChange }),
   });
 

package/src/runtime/useExternalHistory.test.ts

@@ -1180,6 +1180,70 @@ describe("useExternalHistory persistence", () => {
     expect(reportTelemetry).not.toHaveBeenCalled();
   });
 
+  it("skips updates for a persisted message restored by a branch switch", async () => {
+    const { append, update, runCycle, flush } = createPersistenceHarness(true);
+    const completeStatus: ThreadAssistantMessage["status"] = {
+      type: "complete",
+      reason: "stop",
+    };
+    const createUserMessage = (
+      id: string,
+      inner: InnerMessage,
+    ): ThreadMessage => {
+      const message: ThreadMessage = {
+        id,
+        role: "user",
+        content: [{ type: "text", text: "hi" }],
+        attachments: [],
+        createdAt: new Date(),
+        metadata: { custom: {} },
+      };
+      bindExternalStoreMessage(message, [inner]);
+      return message;
+    };
+    const branchA = () => [
+      createUserMessage("user-a", { id: "inner-user-a", parts: ["question"] }),
+      createAssistantMessage(
+        completeStatus,
+        [{ id: "inner-assistant-a", parts: ["answer"] }],
+        "assistant-a",
+      ),
+    ];
+
+    await runCycle(branchA());
+    await waitFor(() => expect(append).toHaveBeenCalledTimes(2));
+
+    await runCycle([
+      createUserMessage("user-b", { id: "inner-user-b", parts: ["edited"] }),
+      createAssistantMessage(
+        completeStatus,
+        [{ id: "inner-assistant-b", parts: ["answer"] }],
+        "assistant-b",
+      ),
+    ]);
+    await waitFor(() => expect(append).toHaveBeenCalledTimes(4));
+
+    append.mockClear();
+    update.mockClear();
+
+    await runCycle([
+      ...branchA(),
+      createUserMessage("user-c", { id: "inner-user-c", parts: ["follow-up"] }),
+      createAssistantMessage(
+        completeStatus,
+        [{ id: "inner-assistant-c", parts: ["answer"] }],
+        "assistant-c",
+      ),
+    ]);
+    await flush();
+
+    expect(update).not.toHaveBeenCalled();
+    expect(append.mock.calls.map(([item]) => item.message.id)).toEqual([
+      "inner-user-c",
+      "inner-assistant-c",
+    ]);
+  });
+
   it("absorbs agentic flickers without losing change detection", async () => {
     const { append, update, reportTelemetry, runCycle, flush, step } =
       createPersistenceHarness(true);

package/src/runtime/useExternalHistory.ts

@@ -5,6 +5,7 @@ import type {
   ThreadHistoryAdapter,
   ThreadMessage,
   MessageFormatAdapter,
+  MessageFormatItem,
   MessageFormatRepository,
   ExportedMessageRepository,
 } from "@assistant-ui/core";
@@ -49,15 +50,10 @@ const isAwaitingToolApproval = (message: ThreadMessage) =>
   message.status?.type === "requires-action" &&
   message.status.reason === "tool-calls";
 
-const snapshotExternalMessages = <TMessage>(
-  messages: readonly ThreadMessage[],
-) =>
-  new Map<string, TMessage[]>(
-    messages.map((message) => [
-      message.id,
-      [...getExternalStoreMessages<TMessage>(message)],
-    ]),
-  );
+const encodeContent = <TMessage>(
+  storageFormatAdapter: MessageFormatAdapter<TMessage, any>,
+  item: MessageFormatItem<TMessage>,
+) => JSON.stringify(storageFormatAdapter.encode(item));
 
 export const useExternalHistory = <TMessage>(
   runtimeRef: RefObject<AssistantRuntime>,
@@ -78,9 +74,11 @@ export const useExternalHistory = <TMessage>(
   const [hasLoaded, setHasLoaded] = useState(false);
 
   const historyIds = useRef(new Set<string>());
-  const persistedInnerIds = useRef(new Set<string>());
   const deferredTelemetryIds = useRef(new Set<string>());
-  const persistedExternalMessages = useRef(new Map<string, TMessage[]>());
+  // `content` is a snapshot taken at write time rather than a re-encode of `source`, because a retained message object can be mutated in place by the runtime that produced it.
+  const persistedInnerMessages = useRef(
+    new Map<string, { source: TMessage; content: string }>(),
+  );
 
   const onSetMessagesRef = useRef(onSetMessages);
   useEffect(() => {
@@ -107,8 +105,12 @@ export const useExternalHistory = <TMessage>(
         const repo = await formatAdapter.load();
         if (repo && repo.messages.length > 0) {
           for (const m of repo.messages) {
-            persistedInnerIds.current.add(
+            persistedInnerMessages.current.set(
               storageFormatAdapter.getId(m.message),
+              {
+                source: m.message,
+                content: encodeContent(storageFormatAdapter, m),
+              },
             );
           }
           const converted = toExportedMessageRepository(toThreadMessages, repo);
@@ -129,10 +131,6 @@ export const useExternalHistory = <TMessage>(
               deferredTelemetryIds.current.add(m.message.id);
             }
           }
-          persistedExternalMessages.current =
-            snapshotExternalMessages<TMessage>(
-              converted.messages.map((m) => m.message),
-            );
         }
       } catch (error) {
         console.error("Failed to load message history:", error);
@@ -294,23 +292,8 @@ export const useExternalHistory = <TMessage>(
 
       persistInFlightRef.current = persistInFlightRef.current
         .then(async () => {
-          const changedRunMessageIds = new Set<string>();
-          for (const message of latest.messages) {
-            const externalMessages =
-              getExternalStoreMessages<TMessage>(message);
-            const previous = persistedExternalMessages.current.get(message.id);
-            if (
-              previous === undefined ||
-              previous.length !== externalMessages.length ||
-              externalMessages.some((item, index) => item !== previous[index])
-            ) {
-              changedRunMessageIds.add(message.id);
-            }
-          }
-
           const { messages } = latest;
           let lastInnerMessageId: string | null = null;
-          const failedUpdateIds = new Set<string>();
 
           const getLastInnerId = (msgs: TMessage[]): string | null =>
             msgs.length > 0 ? storageFormatAdapter.getId(msgs.at(-1)!) : null;
@@ -343,13 +326,7 @@ export const useExternalHistory = <TMessage>(
               continue;
             }
 
-            const isPersistedMessage = historyIds.current.has(message.id);
-            if (isPersistedMessage && !changedRunMessageIds.has(message.id)) {
-              lastInnerMessageId =
-                getLastInnerId(innerMessages) ?? lastInnerMessageId;
-              continue;
-            }
-            if (!isPersistedMessage) {
+            if (!historyIds.current.has(message.id)) {
               historyIds.current.add(message.id);
               deferredTelemetryIds.current.add(message.id);
             }
@@ -357,15 +334,31 @@ export const useExternalHistory = <TMessage>(
             const batchItems = toBatchItems(innerMessages);
             for (const item of batchItems) {
               const innerId = storageFormatAdapter.getId(item.message);
-              if (!persistedInnerIds.current.has(innerId)) {
+              const persisted = persistedInnerMessages.current.get(innerId);
+              if (!persisted) {
                 await adapter.append(item);
-                persistedInnerIds.current.add(innerId);
-              } else if (durationMs !== undefined) {
-                try {
-                  await adapter.update?.(item, innerId);
-                } catch {
-                  // A failed update drops the message from the refreshed baseline so it retries on the next run stop.
-                  failedUpdateIds.add(message.id);
+                persistedInnerMessages.current.set(innerId, {
+                  source: item.message,
+                  content: encodeContent(storageFormatAdapter, item),
+                });
+              } else if (
+                persisted.source !== item.message &&
+                durationMs !== undefined &&
+                adapter.update
+              ) {
+                const content = encodeContent(storageFormatAdapter, item);
+                if (content === persisted.content) {
+                  persisted.source = item.message;
+                } else {
+                  try {
+                    await adapter.update(item, innerId);
+                    persistedInnerMessages.current.set(innerId, {
+                      source: item.message,
+                      content,
+                    });
+                  } catch {
+                    // A failed update leaves the stale baseline behind so the next run stop retries it.
+                  }
                 }
               }
             }
@@ -378,14 +371,6 @@ export const useExternalHistory = <TMessage>(
               adapter.reportTelemetry?.(batchItems, telemetryOptions);
             }
           }
-
-          const nextSnapshot = snapshotExternalMessages<TMessage>(
-            latest.messages,
-          );
-          for (const id of failedUpdateIds) {
-            nextSnapshot.delete(id);
-          }
-          persistedExternalMessages.current = nextSnapshot;
         })
         .catch((error) => {
           console.error("Failed to persist message history:", error);
@@ -430,9 +415,8 @@ export const useExternalHistory = <TMessage>(
 
         historyIds.current.delete(messageId);
         deferredTelemetryIds.current.delete(messageId);
-        persistedExternalMessages.current.delete(messageId);
         for (const item of itemsToDelete) {
-          persistedInnerIds.current.delete(
+          persistedInnerMessages.current.delete(
             storageFormatAdapter.getId(item.message),
           );
         }

package/src/usage.test.ts

@@ -1,5 +1,5 @@
 import { describe, expect, it } from "vitest";
-import { getLatestThreadTokenUsage, getThreadMessageTokenUsage } from "./usage";
+import { getThreadMessageTokenUsage } from "./usage";
 
 function msg(metadata: unknown): { role: "assistant"; metadata: unknown } {
   return {
@@ -9,6 +9,22 @@ function msg(metadata: unknown): { role: "assistant"; metadata: unknown } {
 }
 
 describe("getThreadMessageTokenUsage", () => {
+  it("reads usage from custom.usage", () => {
+    const usage = getThreadMessageTokenUsage(
+      msg({
+        custom: {
+          usage: { inputTokens: 4, outputTokens: 6 },
+        },
+      }),
+    );
+
+    expect(usage).toEqual({
+      totalTokens: 10,
+      inputTokens: 4,
+      outputTokens: 6,
+    });
+  });
+
   it("does not double-count reasoning/cached in fallback totalTokens", () => {
     const usage = getThreadMessageTokenUsage(
       msg({
@@ -157,25 +173,27 @@ describe("getThreadMessageTokenUsage", () => {
   });
 });
 
-describe("getLatestThreadTokenUsage", () => {
-  it("falls back to the latest assistant message with usage", () => {
-    const usage = getLatestThreadTokenUsage([
+describe("getThreadMessageTokenUsage", () => {
+  it("returns token usage from an earlier assistant message", () => {
+    const messages = [
       { role: "assistant", metadata: { usage: { totalTokens: 100 } } },
       { role: "user", metadata: {} },
       { role: "assistant", metadata: {} },
-    ]);
+    ];
+    const usage = getThreadMessageTokenUsage(messages[0]);
 
     expect(usage).toEqual({ totalTokens: 100 });
   });
 
-  it("prefers the newest assistant message when it has usage", () => {
-    const usage = getLatestThreadTokenUsage([
+  it("returns token usage from the newest assistant message", () => {
+    const messages = [
       { role: "assistant", metadata: { usage: { totalTokens: 100 } } },
       {
         role: "assistant",
         metadata: { usage: { inputTokens: 40, outputTokens: 2 } },
       },
-    ]);
+    ];
+    const usage = getThreadMessageTokenUsage(messages[1]);
 
     expect(usage).toEqual({
       totalTokens: 42,

package/src/usage.ts

@@ -1,4 +1,5 @@
 /// <reference types="@assistant-ui/core/react" />
+import { useMemo } from "react";
 import { useAuiState } from "@assistant-ui/store";
 
 export type ThreadTokenUsage = {
@@ -141,18 +142,12 @@ export function getThreadMessageTokenUsage(
   const topLevelUsage = normalizeUsage(metadata.usage);
   if (topLevelUsage) return withComputedTotal(topLevelUsage);
 
-  const legacyUsage = normalizeUsage(asRecord(metadata.custom)?.usage);
-  if (legacyUsage) return withComputedTotal(legacyUsage);
+  const customUsage = normalizeUsage(asRecord(metadata.custom)?.usage);
+  if (customUsage) return withComputedTotal(customUsage);
 
   return usageFromSteps(metadata.steps);
 }
 
-export function getLatestThreadTokenUsage(
-  messages: readonly TokenUsageExtractableMessage[] | undefined,
-): ThreadTokenUsage | undefined {
-  return getThreadMessageTokenUsage(findLatestMessageWithUsage(messages));
-}
-
 function findLatestMessageWithUsage(
   messages: readonly TokenUsageExtractableMessage[] | undefined,
 ): TokenUsageExtractableMessage | undefined {
@@ -170,5 +165,5 @@ function findLatestMessageWithUsage(
 
 export function useThreadTokenUsage(): ThreadTokenUsage | undefined {
   const msg = useAuiState((s) => findLatestMessageWithUsage(s.thread.messages));
-  return getThreadMessageTokenUsage(msg);
+  return useMemo(() => getThreadMessageTokenUsage(msg), [msg]);
 }