@assistant-ui/ai-sdk 0.0.2 → 0.0.4

package/src/runtime/useExternalHistory.test.ts

@@ -392,6 +392,7 @@ describe("useExternalHistory persistence", () => {
     options?: {
       loadMessages?: MessageFormatRepository<InnerMessage>;
       toThreadMessages?: (messages: InnerMessage[]) => ThreadMessage[];
+      initialIsRunning?: boolean;
     },
   ) => {
     const append = vi.fn(
@@ -426,7 +427,7 @@ describe("useExternalHistory persistence", () => {
     };
 
     let listener: (() => void) | undefined;
-    let isRunning = false;
+    let isRunning = options?.initialIsRunning ?? false;
     let messages: ThreadMessage[] = [];
     const getState = vi.fn(() => ({ isRunning, messages }));
     const thread = {
@@ -497,6 +498,83 @@ describe("useExternalHistory persistence", () => {
     };
   };
 
+  it("persists a settled turn when the history adapter becomes active after it", async () => {
+    let listener: (() => void) | undefined;
+    let isRunning = false;
+    let messages: ThreadMessage[] = [];
+    const append = vi.fn(async () => {});
+    const formattedAdapter = {
+      load: vi.fn().mockResolvedValue({ messages: [] }),
+      append,
+    };
+    const historyAdapter: ThreadHistoryAdapter = {
+      load: vi.fn(),
+      append: vi.fn(),
+      withFormat: vi.fn().mockReturnValue(formattedAdapter),
+    };
+    let activeAdapter: ThreadHistoryAdapter | undefined;
+    const thread = {
+      subscribe: (next: () => void) => {
+        listener = next;
+        return () => {};
+      },
+      getState: () => ({ isRunning, messages }),
+      import: vi.fn(),
+      export: vi.fn(() => ({ headId: null, messages: [] })),
+    } as unknown as AssistantRuntime["thread"];
+    const persistenceRuntimeRef = {
+      current: { thread } as AssistantRuntime,
+    };
+    const message = createAssistantMessage(
+      { type: "complete", reason: "stop" },
+      [{ id: "inner-1", parts: ["answer"] }],
+    );
+
+    mocks.hasThreadListItem = true;
+    mocks.remoteId = "remote-thread";
+
+    const { rerender } = renderHook(() =>
+      useExternalHistory(
+        persistenceRuntimeRef,
+        activeAdapter,
+        () => [],
+        persistenceStorageFormat,
+        () => {},
+      ),
+    );
+
+    await act(async () => {
+      messages = [message];
+      isRunning = true;
+      listener?.();
+      isRunning = false;
+      listener?.();
+    });
+
+    activeAdapter = historyAdapter;
+    await act(async () => rerender());
+
+    await waitFor(() => expect(append).toHaveBeenCalledTimes(1));
+    expect(append).toHaveBeenCalledWith({
+      parentId: null,
+      message: { id: "inner-1", parts: ["answer"] },
+    });
+  });
+
+  it("persists a turn that is already running when the subscription starts", async () => {
+    const { append, step } = createPersistenceHarness(false, {
+      initialIsRunning: true,
+    });
+    const message = createAssistantMessage(
+      { type: "complete", reason: "stop" },
+      [{ id: "inner-1", parts: ["answer"] }],
+    );
+
+    await step({ isRunning: false, messages: [message] });
+
+    await waitFor(() => expect(append).toHaveBeenCalledTimes(1));
+  });
+
   it("retries a failed append on the next persistence pass", async () => {
     const consoleError = vi
       .spyOn(console, "error")
@@ -1102,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);
@@ -234,6 +232,19 @@ export const useExternalHistory = <TMessage>(
       }, 0);
     });
 
+    const initialThreadState = runtimeRef.current.thread.getState();
+    wasRunningRef.current = initialThreadState.isRunning;
+    if (initialThreadState.isRunning) {
+      if (runStartRef.current == null) {
+        runStartRef.current = Date.now();
+        stepBoundariesRef.current = [];
+        toolCallCountRef.current = 0;
+        adapter.pin?.();
+      }
+    } else if (initialThreadState.messages.length > 0) {
+      persistSettled(true);
+    }
+
     function persistSettled(ignoreRunning: boolean) {
       persistTimerRef.current = null;
       const latest = runtimeRef.current.thread.getState();
@@ -281,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;
@@ -330,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);
             }
@@ -344,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.
+                  }
                 }
               }
             }
@@ -365,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);
@@ -417,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]);
 }