@assistant-ui/react-ag-ui 0.0.56 → 0.0.57

package/src/runtime/adapter/conversions.ts

@@ -1,6 +1,7 @@
 "use client";
 
 import type { InputContent, RunAgentParameters } from "@ag-ui/client";
+import { generateId } from "@assistant-ui/core";
 import type {
   ThreadMessageLike as CoreThreadMessageLike,
   PartProviderMetadata,
@@ -9,8 +10,8 @@ import type {
 } from "@assistant-ui/core";
 import {
   getAutoStatus,
-  httpUrlPattern,
   parseDataUrl,
+  resolveFilePartSource,
 } from "@assistant-ui/core/internal";
 import { type Tool, toToolsJSONSchema } from "assistant-stream";
 import type { ReadonlyJSONObject } from "assistant-stream/utils";
@@ -177,13 +178,6 @@ function parseJSONText(value: string): unknown {
   }
 }
 
-function generateId(): string {
-  return (
-    (globalThis.crypto as { randomUUID?: () => string })?.randomUUID?.() ??
-    Math.random().toString(36).slice(2)
-  );
-}
-
 function normalizeToolCall(part: ToolCallPart): {
   id: string;
   call: AgUiToolCall;
@@ -242,17 +236,20 @@ function buildInputSource(
   declaredMimeType: string | undefined,
   sourceType?: string,
 ): InputContentSource {
-  if (sourceType === "url" || httpUrlPattern.test(value)) {
+  const source = resolveFilePartSource({
+    data: value,
+    mimeType: declaredMimeType ?? "application/octet-stream",
+    sourceType,
+  });
+  if (source.kind === "url") {
     return declaredMimeType !== undefined
-      ? { type: "url", value, mimeType: declaredMimeType }
-      : { type: "url", value };
+      ? { type: "url", value: source.url, mimeType: declaredMimeType }
+      : { type: "url", value: source.url };
   }
-  const parsed = parseDataUrl(value);
   return {
     type: "data",
-    value: parsed?.data ?? value,
-    mimeType:
-      parsed?.mimeType ?? declaredMimeType ?? "application/octet-stream",
+    value: source.data,
+    mimeType: source.mimeType,
   };
 }
 
@@ -662,7 +659,7 @@ function toAssistantSnapshotMessage(
 }
 
 function toUserOrSystemSnapshotMessage(
-  role: "user" | "system",
+  role: "user" | "system" | "developer",
   rawMessage: Record<string, unknown>,
 ): CoreThreadMessageLike {
   const messageName = getString(rawMessage, "name");
@@ -670,10 +667,23 @@ function toUserOrSystemSnapshotMessage(
     role === "user" ? toSnapshotAttachments(rawMessage.content) : [];
   return {
     id: getString(rawMessage, "id") ?? generateId(),
-    role,
+    // The internal message model has no developer role; it rides as a system
+    // message with its wire role kept in metadata so the export restores it.
+    role: role === "developer" ? "system" : role,
     content: extractText(rawMessage.content),
     ...(messageName !== undefined ? { name: messageName } : {}),
     ...(attachments.length > 0 ? { attachments } : {}),
+    ...(role === "developer"
+      ? {
+          metadata: {
+            custom: {
+              [AG_UI_METADATA_NAMESPACE]: {
+                role: "developer",
+              } satisfies AgUiCustomMetadata,
+            },
+          },
+        }
+      : {}),
   };
 }
 
@@ -939,7 +949,7 @@ export function fromAgUiMessages(
       continue;
     }
 
-    if (role === "user" || role === "system") {
+    if (role === "user" || role === "system" || role === "developer") {
       converted.push(toUserOrSystemSnapshotMessage(role, rawMessage));
     }
   }
@@ -1137,9 +1147,21 @@ export function toAgUiMessages(
     }
 
     if (message.role === "system" || message.role === "developer") {
+      const custom = isObject(message.metadata)
+        ? message.metadata.custom
+        : undefined;
+      const namespaced = isObject(custom)
+        ? custom[AG_UI_METADATA_NAMESPACE]
+        : undefined;
+      const wireRole =
+        message.role === "system" &&
+        isObject(namespaced) &&
+        namespaced.role === "developer"
+          ? ("developer" as const)
+          : message.role;
       converted.push({
         id: message.id,
-        role: message.role,
+        role: wireRole,
         content: extractText(message.content),
         ...(message.name ? { name: message.name } : {}),
       });

package/src/runtime/adapter/run-aggregator.ts

@@ -28,6 +28,9 @@ export type AgUiOpaqueReasoning = {
 };
 
 export type AgUiCustomMetadata = {
+  /** Wire role restored on export for messages the internal model cannot
+   * represent (a developer record rides as a system message). */
+  role?: "developer";
   interrupts?: AgUiInterrupt[];
   opaqueReasoning?: AgUiOpaqueReasoning[];
 };
@@ -63,6 +66,7 @@ export type RunAggregatorOptions = {
   logger: Logger;
   emit: Emit;
   onServerMessageId?: (messageId: string) => void;
+  onTextMessageStart?: (messageId: string) => void;
 };
 
 /**
@@ -80,16 +84,29 @@ export class RunAggregator {
   private readonly showThinking: boolean;
   private readonly logger: Logger;
   private readonly onServerMessageId: ((messageId: string) => void) | undefined;
+  private readonly onTextMessageStart:
+    | ((messageId: string) => void)
+    | undefined;
 
   private status: ChatModelRunResult["status"] | undefined;
   private interrupts: AgUiInterrupt[] | undefined;
-  private readonly textParts = new Map<
-    string,
-    { buffer: string; touched: boolean }
-  >();
+  private readonly textParts = new Map<string, { buffer: string }>();
   private activeTextMessageId: string | undefined;
   private readonly reasoningParts = new Map<string, string>(); // key → buffer
   private readonly reasoningSignatures = new Map<string, string>();
+  private readonly reasoningSignatureIds = new Map<string, string>();
+  // Signatures captured while thinking is hidden have no block to live on;
+  // they are transport state that must survive to the next run input. The
+  // claim rules mirror the visible path: an id-carrying block, or an open
+  // anonymous one.
+  private readonly hiddenSignatures = new Map<string, string>();
+  private readonly hiddenSignatureAnchors = new Map<string, number>();
+  private readonly hiddenBlockAnchors = new Map<string, number>();
+  private hiddenAnonymousAnchor: number | undefined;
+  private readonly hiddenReasoningIds = new Set<string>();
+  private hiddenActiveReasoning: "none" | "anonymous" | "identified" = "none";
+  private hasEmittedOpaqueReasoning = false;
+  private readonly loggedDroppedOpaqueIds = new Set<string>();
   private readonly reasoningMessageIds = new Map<string, string>();
   private readonly anonymousReasoningKeys = new Set<string>();
   private activeReasoningKey: string | undefined;
@@ -106,6 +123,8 @@ export class RunAggregator {
   )[] = [];
   private textPartCounter = 0;
   private serverMessageIdReported = false;
+  private reportedServerMessageId: string | undefined;
+  private lastTextMessageId: string | undefined;
 
   private streamStartTime: number | undefined;
   private firstTokenTime: number | undefined;
@@ -116,6 +135,7 @@ export class RunAggregator {
     this.showThinking = options.showThinking;
     this.logger = options.logger;
     this.onServerMessageId = options.onServerMessageId;
+    this.onTextMessageStart = options.onTextMessageStart;
   }
 
   hasToolCall(toolCallId: string): boolean {
@@ -125,22 +145,11 @@ export class RunAggregator {
   handle(event: AgUiEvent): void {
     switch (event.type) {
       case "RUN_STARTED": {
-        this.clearTextParts();
-        this.reasoningParts.clear();
-        this.reasoningSignatures.clear();
-        this.reasoningMessageIds.clear();
-        this.anonymousReasoningKeys.clear();
-        this.activeReasoningKey = undefined;
-        this.reasoningPartCounter = 0;
-        this.toolCalls.clear();
-        this.a2uiBuckets.clear();
-        this.a2uiToolCallIds.clear();
-        this.lastResolvedToolCallId = undefined;
-        this.partOrder.length = 0;
-        this.textPartCounter = 0;
-        this.activeTextMessageId = undefined;
+        this.resetMessageParts();
         this.interrupts = undefined;
         this.serverMessageIdReported = false;
+        this.lastTextMessageId = undefined;
+        this.reportedServerMessageId = undefined;
         this.streamStartTime = Date.now();
         this.firstTokenTime = undefined;
         this.totalChunks = 0;
@@ -184,18 +193,19 @@ export class RunAggregator {
       }
 
       case "TEXT_MESSAGE_START": {
+        this.beginDistinctTextMessage(event.messageId);
         this.reportServerMessageId(event.messageId);
-        const id = this.startTextMessage(event.messageId);
-        if (id) {
-          this.markTextPartTouched(id);
-        }
+        this.startTextMessage(event.messageId);
+        if (event.messageId) this.lastTextMessageId = event.messageId;
         this.emit();
         break;
       }
       case "TEXT_MESSAGE_CONTENT":
       case "TEXT_MESSAGE_CHUNK": {
         const incomingId = "messageId" in event ? event.messageId : undefined;
+        this.beginDistinctTextMessage(incomingId);
         this.reportServerMessageId(incomingId);
+        if (incomingId) this.lastTextMessageId = incomingId;
         if (!event.delta) break;
         this.recordFirstToken();
         const id = this.resolveTextMessageId(incomingId);
@@ -238,13 +248,31 @@ export class RunAggregator {
           // entityId names any message, not necessarily a reasoning one, so an
           // unmatched id may only claim a block that has no id to contradict it.
           const active = this.activeReasoningKey;
-          const key = this.reasoningParts.has(event.entityId)
-            ? event.entityId
-            : active !== undefined && this.anonymousReasoningKeys.has(active)
-              ? active
-              : undefined;
+          const key = this.showThinking
+            ? this.reasoningParts.has(event.entityId)
+              ? event.entityId
+              : active !== undefined && this.anonymousReasoningKeys.has(active)
+                ? active
+                : undefined
+            : undefined;
           if (key !== undefined) {
             this.reasoningSignatures.set(key, event.encryptedValue);
+            this.reasoningSignatureIds.set(key, event.entityId);
+            this.emit();
+          } else if (
+            !this.showThinking &&
+            event.entityId.trim().length > 0 &&
+            event.encryptedValue.trim().length > 0 &&
+            (this.hiddenReasoningIds.has(event.entityId) ||
+              this.hiddenActiveReasoning === "anonymous")
+          ) {
+            this.hiddenSignatures.set(event.entityId, event.encryptedValue);
+            this.hiddenSignatureAnchors.set(
+              event.entityId,
+              this.hiddenBlockAnchors.get(event.entityId) ??
+                this.hiddenAnonymousAnchor ??
+                this.partOrder.length,
+            );
             this.emit();
           }
         }
@@ -436,8 +464,13 @@ export class RunAggregator {
   }
 
   private reportServerMessageId(messageId: string | undefined): void {
-    if (this.serverMessageIdReported || !messageId) return;
+    if (!messageId) return;
+    if (this.lastTextMessageId === undefined) {
+      this.lastTextMessageId = messageId;
+    }
+    if (this.serverMessageIdReported) return;
     this.serverMessageIdReported = true;
+    this.reportedServerMessageId = messageId;
     this.onServerMessageId?.(messageId);
   }
 
@@ -445,6 +478,45 @@ export class RunAggregator {
     this.textParts.clear();
   }
 
+  private resetMessageParts(): void {
+    this.clearTextParts();
+    this.reasoningParts.clear();
+    this.reasoningSignatures.clear();
+    this.reasoningSignatureIds.clear();
+    this.reasoningMessageIds.clear();
+    this.hiddenSignatures.clear();
+    this.hiddenSignatureAnchors.clear();
+    this.hiddenBlockAnchors.clear();
+    this.hiddenAnonymousAnchor = undefined;
+    this.hiddenReasoningIds.clear();
+    this.hiddenActiveReasoning = "none";
+    this.loggedDroppedOpaqueIds.clear();
+    this.anonymousReasoningKeys.clear();
+    this.activeReasoningKey = undefined;
+    this.reasoningPartCounter = 0;
+    this.toolCalls.clear();
+    this.a2uiBuckets.clear();
+    this.a2uiToolCallIds.clear();
+    this.lastResolvedToolCallId = undefined;
+    this.partOrder.length = 0;
+    this.textPartCounter = 0;
+    this.activeTextMessageId = undefined;
+    this.reportedServerMessageId = undefined;
+  }
+
+  private beginDistinctTextMessage(messageId: string | undefined): void {
+    if (
+      !messageId ||
+      this.lastTextMessageId === undefined ||
+      this.lastTextMessageId === messageId ||
+      !this.onTextMessageStart
+    ) {
+      return;
+    }
+    this.resetMessageParts();
+    this.onTextMessageStart(messageId);
+  }
+
   private generateTextKey(): string {
     this.textPartCounter += 1;
     return `text-${this.textPartCounter}`;
@@ -476,7 +548,7 @@ export class RunAggregator {
 
   private ensureTextPart(id: string): void {
     if (!this.textParts.has(id)) {
-      this.textParts.set(id, { buffer: "", touched: false });
+      this.textParts.set(id, { buffer: "" });
       if (
         !this.partOrder.some((part) => part.kind === "text" && part.key === id)
       ) {
@@ -485,18 +557,11 @@ export class RunAggregator {
     }
   }
 
-  private markTextPartTouched(id: string): void {
-    const entry = this.textParts.get(id);
-    if (!entry) return;
-    entry.touched = true;
-  }
-
   private appendText(id: string, delta: string): void {
     this.ensureTextPart(id);
     const entry = this.textParts.get(id);
     if (!entry) return;
     entry.buffer += delta;
-    entry.touched = true;
   }
 
   private startToolCall(
@@ -640,7 +705,22 @@ export class RunAggregator {
       new Set(this.toolCalls.keys()),
     );
 
-    for (const part of this.partOrder) {
+    const opaqueCandidates: (AgUiOpaqueReasoning & { anchor: number })[] =
+      Array.from(this.hiddenSignatures, ([id, encryptedValue]) => ({
+        id,
+        encryptedValue,
+        anchor: this.hiddenSignatureAnchors.get(id)!,
+      }));
+
+    let currentIndex = -1;
+    let lastMaterializedIndex = -1;
+    const pushSnapshotPart = (part: (typeof snapshot)[number]) => {
+      snapshot.push(part);
+      lastMaterializedIndex = currentIndex;
+    };
+
+    for (const [index, part] of this.partOrder.entries()) {
+      currentIndex = index;
       if (part.kind === "reasoning") {
         if (this.showThinking) {
           const buffer = this.reasoningParts.get(part.key) ?? "";
@@ -651,13 +731,22 @@ export class RunAggregator {
               ...(reasoningId !== undefined ? { reasoningId } : {}),
               ...(encryptedValue !== undefined ? { encryptedValue } : {}),
             };
-            snapshot.push({
+            pushSnapshotPart({
               type: "reasoning",
               text: buffer,
               ...(Object.keys(meta).length > 0
                 ? { providerMetadata: { [AG_UI_METADATA_NAMESPACE]: meta } }
                 : {}),
             } as const);
+          } else {
+            // A retracted empty block still carries transport state: without
+            // a part to live on, its signature rides the message metadata,
+            // matching the shape a snapshot reload produces.
+            const encryptedValue = this.reasoningSignatures.get(part.key);
+            const id = this.reasoningSignatureIds.get(part.key);
+            if (id?.trim() && encryptedValue?.trim()) {
+              opaqueCandidates.push({ id, encryptedValue, anchor: index + 1 });
+            }
           }
         }
         continue;
@@ -665,14 +754,14 @@ export class RunAggregator {
 
       if (part.kind === "text") {
         const entry = this.textParts.get(part.key);
-        if (entry?.touched) {
-          snapshot.push({ type: "text", text: entry.buffer } as const);
+        if (entry && entry.buffer.trim().length > 0) {
+          pushSnapshotPart({ type: "text", text: entry.buffer } as const);
         }
         continue;
       }
 
       if (part.kind === "data") {
-        snapshot.push({
+        pushSnapshotPart({
           type: "data",
           name: part.name,
           data: part.value,
@@ -712,17 +801,57 @@ export class RunAggregator {
           ? { unstable_toolMessageId: entry.toolMessageId }
           : {}),
       } as ToolCallMessagePart & { unstable_toolMessageId?: string };
-      snapshot.push(toolPart);
+      pushSnapshotPart(toolPart);
     }
 
+    // An anonymous claim promotes the signature's entityId to a wire message
+    // id; when that id actually names a text message or the adopted assistant
+    // message id (which can also arrive via TOOL_CALL_START.parentMessageId),
+    // replaying it as its own record would put two wire records under one id,
+    // so such entries are dropped instead — the same outcome the visible
+    // path's claim guard produces for a signature that is "not for this
+    // block". Entries with no materialized part after their anchor trail the
+    // assistant record, matching the snapshot path's anchor/after
+    // bookkeeping.
+    const toolMessageIds = new Set<string>();
+    for (const call of this.toolCalls.values()) {
+      if (call.toolMessageId) toolMessageIds.add(call.toolMessageId);
+    }
+    opaqueCandidates.sort((a, b) => a.anchor - b.anchor);
+    const publishableOpaqueReasoning = opaqueCandidates
+      .filter((entry) => {
+        const collides =
+          this.textParts.has(entry.id) ||
+          entry.id === this.reportedServerMessageId ||
+          toolMessageIds.has(entry.id);
+        if (collides && !this.loggedDroppedOpaqueIds.has(entry.id)) {
+          this.loggedDroppedOpaqueIds.add(entry.id);
+          this.logger.debug?.(
+            "[agui] aggregator dropped opaque reasoning signature: id collides with a message id",
+            entry.id,
+          );
+        }
+        return !collides;
+      })
+      .map(({ anchor, ...entry }) =>
+        lastMaterializedIndex < anchor ? { ...entry, after: true } : entry,
+      );
+    if (publishableOpaqueReasoning.length > 0)
+      this.hasEmittedOpaqueReasoning = true;
+    // Once an opaque entry has been published this run, the key keeps being
+    // emitted (empty when withdrawn) so the namespace merge can retract it.
+    const includeOpaque = this.hasEmittedOpaqueReasoning;
     const timing = this.getTiming();
     const metadata = {
       ...(timing ? { timing } : {}),
-      ...(this.interrupts
+      ...(this.interrupts || includeOpaque
         ? {
             custom: {
               [AG_UI_METADATA_NAMESPACE]: {
-                interrupts: this.interrupts,
+                ...(this.interrupts ? { interrupts: this.interrupts } : {}),
+                ...(includeOpaque
+                  ? { opaqueReasoning: publishableOpaqueReasoning }
+                  : {}),
               } satisfies AgUiCustomMetadata,
             },
           }
@@ -778,7 +907,18 @@ export class RunAggregator {
   }
 
   private handleReasoningStart(messageId?: string, isMessageId = false): void {
-    if (!this.showThinking) return;
+    if (!this.showThinking) {
+      const anchor = this.partOrder.length;
+      if (messageId === undefined) {
+        this.hiddenActiveReasoning = "anonymous";
+        this.hiddenAnonymousAnchor = anchor;
+      } else {
+        this.hiddenReasoningIds.add(messageId);
+        this.hiddenBlockAnchors.set(messageId, anchor);
+        this.hiddenActiveReasoning = "identified";
+      }
+      return;
+    }
     // A reasoning block acts as a boundary: anonymous text arriving after it
     // should be a new part, not appended to any pre-reasoning text.
     this.activeTextMessageId = undefined;
@@ -804,7 +944,15 @@ export class RunAggregator {
     messageId?: string,
     isMessageId = false,
   ): void {
-    if (!this.showThinking || !delta) return;
+    if (!delta) return;
+    if (!this.showThinking) {
+      // Content without a preceding START still names the block (the visible
+      // path opens it lazily); register it so a signature can claim it.
+      if (this.hiddenActiveReasoning === "none") {
+        this.handleReasoningStart(messageId, isMessageId);
+      }
+      return;
+    }
     if (!this.activeReasoningKey) {
       // Content arrived without a preceding START — create the slot lazily.
       this.handleReasoningStart(messageId, isMessageId);
@@ -816,7 +964,10 @@ export class RunAggregator {
   }
 
   private handleReasoningEnd(): void {
-    if (!this.showThinking) return;
+    if (!this.showThinking) {
+      this.hiddenActiveReasoning = "none";
+      return;
+    }
     this.activeReasoningKey = undefined;
     this.emit();
   }

package/src/useAgUiRuntime.history.test.tsx

@@ -0,0 +1,76 @@
+// @vitest-environment jsdom
+
+import { describe, it, expect, vi, afterEach } from "vitest";
+import { cleanup, renderHook, waitFor } from "@testing-library/react";
+import type { ThreadHistoryAdapter } from "@assistant-ui/core";
+import type { HttpAgent } from "@ag-ui/client";
+import { useAgUiRuntime } from "./useAgUiRuntime";
+
+const agent = { runAgent: vi.fn(), abortRun: vi.fn() } as unknown as HttpAgent;
+
+function createHistory(): ThreadHistoryAdapter {
+  return {
+    load: vi.fn().mockResolvedValue({
+      headId: "restored",
+      messages: [
+        {
+          parentId: null,
+          message: {
+            id: "restored",
+            role: "user" as const,
+            content: [{ type: "text" as const, text: "hello" }],
+            createdAt: new Date(0),
+            metadata: { custom: {} },
+          },
+        },
+      ],
+    }),
+    append: vi.fn().mockResolvedValue(undefined),
+  };
+}
+
+afterEach(() => {
+  cleanup();
+});
+
+describe("useAgUiRuntime history", () => {
+  it("loads a history adapter that arrives on a later render", async () => {
+    const history = createHistory();
+    const { result, rerender } = renderHook(
+      ({ history }: { history?: ThreadHistoryAdapter }) =>
+        useAgUiRuntime({ agent, adapters: history ? { history } : {} }),
+      { initialProps: {} },
+    );
+
+    await waitFor(() =>
+      expect(result.current.thread.getState().isLoading).toBe(false),
+    );
+    expect(history.load).not.toHaveBeenCalled();
+    expect(result.current.thread.getState().messages).toEqual([]);
+
+    rerender({ history });
+
+    await waitFor(() =>
+      expect(
+        result.current.thread.getState().messages.map((m) => m.id),
+      ).toEqual(["restored"]),
+    );
+    expect(history.load).toHaveBeenCalledOnce();
+  });
+
+  it("loads once when the adapter is present from the first render", async () => {
+    const history = createHistory();
+    const { result, rerender } = renderHook(() =>
+      useAgUiRuntime({ agent, adapters: { history } }),
+    );
+
+    await waitFor(() =>
+      expect(
+        result.current.thread.getState().messages.map((m) => m.id),
+      ).toEqual(["restored"]),
+    );
+    rerender();
+    rerender();
+    expect(history.load).toHaveBeenCalledOnce();
+  });
+});

package/src/useAgUiRuntime.ts

@@ -81,14 +81,16 @@ export function useAgUiRuntime(
   }
 
   const core = coreRef.current;
-  core.updateOptions({
-    agent: options.agent,
-    logger,
-    showThinking: options.showThinking ?? true,
-    autoCancelPendingToolCalls: options.autoCancelPendingToolCalls,
-    ...(options.onError && { onError: options.onError }),
-    ...(options.onCancel && { onCancel: options.onCancel }),
-    ...(historyAdapter && { history: historyAdapter }),
+  useEffect(() => {
+    core.updateOptions({
+      agent: options.agent,
+      logger,
+      showThinking: options.showThinking ?? true,
+      autoCancelPendingToolCalls: options.autoCancelPendingToolCalls,
+      ...(options.onError && { onError: options.onError }),
+      ...(options.onCancel && { onCancel: options.onCancel }),
+      ...(historyAdapter && { history: historyAdapter }),
+    });
   });
 
   const [toolStatuses, setToolStatuses] = useState<
@@ -228,7 +230,7 @@ export function useAgUiRuntime(
         isLoading: core.isLoading,
         messageRepository: core.getMessageRepository(),
         state: core.getState(),
-        isRunning,
+        isRunning: core.isRunning(),
         extras: agUiExtras.provide({
           interrupts:
             core.getPendingInterrupts()?.interrupts ?? EMPTY_INTERRUPTS,