@assistant-ui/react-a2a 0.2.41 → 0.2.42

package/src/A2AThreadRuntimeCore.ts

@@ -17,14 +17,13 @@ import {
   appendToolInteraction,
   createMessageRepositorySession,
   invokeUserCallback,
+  RunLeases,
 } from "@assistant-ui/core/internal";
-import type { ReadonlyJSONObject } from "assistant-stream/utils";
 import {
   applyA2uiOperations,
-  convertSurfaceToUISpec,
-  surfaceToOperations,
+  surfaceToPresentToolCall,
   type A2uiState,
-} from "@assistant-ui/react-generative-ui/a2ui";
+} from "@assistant-ui/generative-ui/a2ui";
 import type { A2AClient } from "./A2AClient";
 import type {
   A2AArtifact,
@@ -64,7 +63,6 @@ export type A2AThreadRuntimeCoreOptions = {
   onCancel?: (() => void) | undefined;
   onArtifactComplete?: ((artifact: A2AArtifact) => void) | undefined;
   history?: ThreadHistoryAdapter | undefined;
-  notifyUpdate: () => void;
 };
 
 const FALLBACK_USER_STATUS = {
@@ -97,13 +95,23 @@ export class A2AThreadRuntimeCore {
   private onCancel: (() => void) | undefined;
   private onArtifactComplete: ((artifact: A2AArtifact) => void) | undefined;
   private history: ThreadHistoryAdapter | undefined;
-  private readonly notifyUpdate: () => void;
+  private readonly listeners = new Set<() => void>();
+  private snapshot:
+    | {
+        isLoading: boolean;
+        isRunning: boolean;
+        messageRepository: ExportedMessageRepository;
+        task: A2ATask | undefined;
+        artifacts: readonly A2AArtifact[];
+        agentCard: A2AAgentCard | undefined;
+      }
+    | undefined;
 
   private runtime: AssistantRuntime | undefined;
   private readonly session = createMessageRepositorySession();
   private isRunningFlag = false;
   private abortController: AbortController | null = null;
-  private runGeneration = 0;
+  private readonly runLeases = new RunLeases();
   private pendingError: Error | null = null;
 
   // A2A-specific state
@@ -137,15 +145,32 @@ export class A2AThreadRuntimeCore {
     this.onCancel = options.onCancel;
     this.onArtifactComplete = options.onArtifactComplete;
     this.history = options.history;
-    this.notifyUpdate = options.notifyUpdate;
   }
 
-  updateOptions(options: Omit<A2AThreadRuntimeCoreOptions, "notifyUpdate">) {
+  subscribe = (listener: () => void): (() => void) => {
+    this.listeners.add(listener);
+    return () => this.listeners.delete(listener);
+  };
+
+  getSnapshot = () =>
+    (this.snapshot ??= {
+      isLoading: this._isLoading,
+      isRunning: this.isRunningFlag,
+      messageRepository: this.session.export(),
+      task: this.currentTask,
+      artifacts: this.currentArtifacts,
+      agentCard: this.agentCardValue,
+    });
+
+  private publishUpdate(): void {
+    this.snapshot = undefined;
+    for (const listener of this.listeners) listener();
+  }
+
+  updateOptions(options: A2AThreadRuntimeCoreOptions) {
     this.client = options.client;
-    // The hook re-applies options on every render, including renders caused
-    // by this core's own notifyUpdate. The option only seeds the context: a
-    // re-render with the same value must not clobber a server-assigned
-    // contextId learned from the stream.
+    // The option only seeds the context; a re-render with the same value must
+    // not clobber a server-assigned contextId learned from the stream.
     if (options.contextId !== this.lastOptionsContextId) {
       this.contextId = options.contextId;
       this.lastOptionsContextId = options.contextId;
@@ -172,6 +197,7 @@ export class A2AThreadRuntimeCore {
    * switches, deletes, and cancel resyncs, which must keep the live context. */
   resetContext(): void {
     this._historyLoadGeneration++;
+    const wasLoading = this._isLoading;
     this._isLoading = false;
     // Restore the seed before aborting: an onCancel callback that starts a
     // new run must not pick up the old thread's context, and its controller
@@ -184,6 +210,7 @@ export class A2AThreadRuntimeCore {
         this.abortController = null;
       }
     }
+    if (wasLoading) this.publishUpdate();
   }
 
   attachRuntime(runtime: AssistantRuntime) {
@@ -257,7 +284,7 @@ export class A2AThreadRuntimeCore {
       };
     });
     if (touched) {
-      this.notifyUpdate();
+      this.publishUpdate();
       this.persistAssistantHistory(messageId);
     }
   }
@@ -291,7 +318,7 @@ export class A2AThreadRuntimeCore {
         this._agentCardRetryAfter = 0;
         this._agentCardRetryDelay = INITIAL_AGENT_CARD_RETRY_DELAY_MS;
         this._agentCardDiscoveryFailed = false;
-        this.notifyUpdate();
+        this.publishUpdate();
       },
       () => {
         this._agentCardDiscoveryFailed = true;
@@ -357,10 +384,10 @@ export class A2AThreadRuntimeCore {
       .finally(() => {
         if (generation !== this._historyLoadGeneration) return;
         this._isLoading = false;
-        this.notifyUpdate();
+        this.publishUpdate();
       });
 
-    this.notifyUpdate();
+    this.publishUpdate();
     return this._loadPromise;
   }
 
@@ -380,7 +407,7 @@ export class A2AThreadRuntimeCore {
           : this.session.headId;
     this.session.addOrUpdateMessage(parentId, threadMessage);
     this.session.switchToBranch(threadMessage.id);
-    this.notifyUpdate();
+    this.publishUpdate();
     this.recordHistoryEntry(parentId, threadMessage);
 
     if (!startRun) return;
@@ -391,7 +418,7 @@ export class A2AThreadRuntimeCore {
     const parentId = this.session.headId;
     this.session.addOrUpdateMessage(parentId, message);
     this.session.switchToBranch(message.id);
-    this.notifyUpdate();
+    this.publishUpdate();
     this.recordHistoryEntry(parentId, message);
   }
 
@@ -422,7 +449,7 @@ export class A2AThreadRuntimeCore {
     // onCancel callback synchronously, which may clear the thread and with it
     // the task this cancellation is for, or start a new run.
     const task = this.currentTask;
-    const generation = this.runGeneration;
+    const lease = this.runLeases.current();
 
     // Abort locally first so the stream stops immediately
     this.abortController.abort();
@@ -433,11 +460,11 @@ export class A2AThreadRuntimeCore {
         const updated = await this.client.cancelTask(task.id);
         // Only apply the response while nothing newer exists. A newer snapshot
         // or a cleared thread replaces the task object; a follow-up run that
-        // has not emitted yet keeps it, so the run generation is what rules
+        // has not emitted yet keeps it, so the run lease is what rules
         // that case out.
-        if (this.currentTask === task && this.runGeneration === generation) {
+        if (this.currentTask === task && lease.isCurrent()) {
           this.currentTask = updated;
-          this.notifyUpdate();
+          this.publishUpdate();
         }
       } catch {
         // Server cancel failed; local abort already handled
@@ -471,7 +498,7 @@ export class A2AThreadRuntimeCore {
     this.currentArtifacts = [];
     this.a2uiState = new Map();
     this.a2uiMessageIds.clear();
-    this.notifyUpdate();
+    this.publishUpdate();
   }
 
   applyExternalMessages(messages: readonly ThreadMessage[]): void {
@@ -510,11 +537,17 @@ export class A2AThreadRuntimeCore {
   // --- Run logic ---
 
   private async startRun(userThreadMessage: ThreadMessage): Promise<void> {
-    this.runGeneration++;
-
-    // Cancel any in-progress run before starting a new one
-    if (this.abortController) {
-      this.abortController.abort();
+    this.runLeases.begin();
+
+    // Cancel any in-progress run before starting a new one. Its abort runs
+    // onCancel synchronously, and a run that callback starts keeps the thread.
+    // A listener that throws before finishRun leaves `previous` installed, so
+    // only a different, non-null controller counts as a replacement run.
+    const previous = this.abortController;
+    if (previous) {
+      previous.abort();
+      if (this.abortController !== previous && this.abortController !== null)
+        return;
       this.abortController = null;
     }
 
@@ -708,7 +741,7 @@ export class A2AThreadRuntimeCore {
     const status = taskStateToMessageStatus(event.status.state);
     this.updateAssistantStatus(assistantId, status);
 
-    this.notifyUpdate();
+    this.publishUpdate();
   }
 
   private handleArtifactUpdate(
@@ -757,7 +790,7 @@ export class A2AThreadRuntimeCore {
       );
     }
 
-    this.notifyUpdate();
+    this.publishUpdate();
   }
 
   private handleMessage(assistantId: string, message: A2AMessage) {
@@ -766,7 +799,7 @@ export class A2AThreadRuntimeCore {
     this.applyA2uiMessage(message);
     const content = a2aMessageToContent(message);
     this.updateAssistantContent(assistantId, content);
-    this.notifyUpdate();
+    this.publishUpdate();
   }
 
   private handleTaskSnapshot(assistantId: string, task: A2ATask) {
@@ -833,7 +866,7 @@ export class A2AThreadRuntimeCore {
     const status = taskStateToMessageStatus(task.status.state);
     this.updateAssistantStatus(assistantId, status);
 
-    this.notifyUpdate();
+    this.publishUpdate();
   }
 
   // --- Message helpers ---
@@ -856,7 +889,7 @@ export class A2AThreadRuntimeCore {
     };
     this.session.addOrUpdateMessage(parentId, assistant);
     this.session.switchToBranch(id);
-    this.notifyUpdate();
+    this.publishUpdate();
     return id;
   }
 
@@ -890,16 +923,11 @@ export class A2AThreadRuntimeCore {
   private a2uiSurfaceParts(): ThreadAssistantMessagePart[] {
     const parts: ThreadAssistantMessagePart[] = [];
     for (const [surfaceId, surface] of this.a2uiState) {
-      const { spec } = convertSurfaceToUISpec(surface);
-      if (!spec) continue;
+      const { toolCall } = surfaceToPresentToolCall(surfaceId, surface);
+      if (!toolCall) continue;
       parts.push({
         type: "tool-call",
-        toolCallId: `a2ui:${surfaceId}`,
-        toolName: "present",
-        args: spec as unknown as ReadonlyJSONObject,
-        argsText: JSON.stringify(spec),
-        result: {},
-        artifact: { a2ui: surfaceToOperations(surface) },
+        ...toolCall,
       });
     }
     return parts;
@@ -959,7 +987,7 @@ export class A2AThreadRuntimeCore {
         content: this.withA2uiSurfaces(message.content),
       };
     });
-    if (touched) this.notifyUpdate();
+    if (touched) this.publishUpdate();
   }
 
   private updateAssistantArtifacts(messageId: string) {
@@ -981,7 +1009,7 @@ export class A2AThreadRuntimeCore {
         },
       };
     });
-    if (touched) this.notifyUpdate();
+    if (touched) this.publishUpdate();
   }
 
   private updateAssistantStatus(messageId: string, status: MessageStatus) {
@@ -990,7 +1018,7 @@ export class A2AThreadRuntimeCore {
       return { ...message, status };
     });
     if (touched) {
-      this.notifyUpdate();
+      this.publishUpdate();
       if (this.isPersistableAssistantStatus(status)) {
         this.persistAssistantHistory(messageId);
       }
@@ -1015,7 +1043,7 @@ export class A2AThreadRuntimeCore {
 
   private setRunning(running: boolean) {
     this.isRunningFlag = running;
-    this.notifyUpdate();
+    this.publishUpdate();
   }
 
   private finishRun(controller: AbortController | null) {

package/src/conversions.test.ts

@@ -475,6 +475,19 @@ describe("isInterruptedTaskState", () => {
 });
 
 describe("contentPartsToA2AParts", () => {
+  it.each(["application/pdf", ""])(
+    "unwraps a media-less data URL while retaining the adapter's %j MIME fallback",
+    (mimeType) => {
+      expect(
+        contentPartsToA2AParts([
+          { type: "file", data: "data:;base64,SGVsbG8=", mimeType },
+        ]),
+      ).toEqual([
+        { raw: "SGVsbG8=", mediaType: mimeType || "application/octet-stream" },
+      ]);
+    },
+  );
+
   it("converts text parts", () => {
     const result = contentPartsToA2AParts([{ type: "text", text: "hi" }]);
     expect(result).toEqual([{ text: "hi" }]);
@@ -525,6 +538,14 @@ describe("contentPartsToA2AParts", () => {
     expect(result).toEqual([{ raw: "aGVsbG8=", mediaType: "image/png" }]);
   });
 
+  it("sniffs media-less image data URLs before sending their raw bytes", () => {
+    expect(
+      contentPartsToA2AParts([
+        { type: "image", image: "data:;base64,iVBORw0KGgo=" },
+      ]),
+    ).toEqual([{ raw: "iVBORw0KGgo=", mediaType: "image/png" }]);
+  });
+
   it("propagates image filenames", () => {
     const result = contentPartsToA2AParts([
       {

package/src/conversions.ts

@@ -8,6 +8,7 @@ import type {
 import {
   parseDataUrl,
   resolveFilePartSource,
+  resolveImageMediaType,
 } from "@assistant-ui/core/internal";
 import type { A2AMessage, A2APart, A2ATaskState } from "./types";
 
@@ -176,7 +177,7 @@ export function contentPartsToA2AParts(
           if (parsed) {
             return {
               raw: parsed.data,
-              mediaType: parsed.mimeType,
+              mediaType: resolveImageMediaType(part.image, fallbackMimeType),
               ...(part.filename && { filename: part.filename }),
             };
           }

package/src/useA2ARuntime.test.tsx

@@ -140,6 +140,42 @@ describe("useA2ARuntime", () => {
     append: vi.fn().mockResolvedValue(undefined),
   });
 
+  it("re-renders when the core publishes loaded history", async () => {
+    const { client, getAgentCard } = createMockClient();
+    getAgentCard.mockImplementation(() => new Promise(() => {}));
+    let resolveHistory!: (repo: ExportedMessageRepository) => void;
+    const pendingHistory = new Promise<ExportedMessageRepository>((resolve) => {
+      resolveHistory = resolve;
+    });
+    const history = {
+      load: vi.fn(() => pendingHistory),
+      append: async () => {},
+    };
+    const render = vi.fn();
+    const { result } = renderHook(() => {
+      render();
+      return useA2ARuntime({ client, adapters: { history } });
+    });
+
+    await waitFor(() => expect(history.load).toHaveBeenCalledOnce());
+    const rendersBeforeHistory = render.mock.calls.length;
+
+    await act(async () => {
+      resolveHistory({
+        headId: "restored",
+        messages: [
+          { parentId: null, message: createThreadMessage("restored") },
+        ],
+      });
+      await pendingHistory;
+    });
+
+    expect(render.mock.calls.length).toBeGreaterThan(rendersBeforeHistory);
+    expect(result.current.thread.getState().messages.map((m) => m.id)).toEqual([
+      "restored",
+    ]);
+  });
+
   it("loads a history adapter that arrives on a later render", async () => {
     const { client } = createMockClient();
     const history = createHistory();
@@ -466,6 +502,74 @@ describe("useA2ARuntime", () => {
     expect(result.current.thread.export().messages).toEqual([]);
   });
 
+  it.each(["existing", "new"] as const)(
+    "drops a send started by onCancel while switching to the %s thread",
+    async (destination) => {
+      const { client, streamMessage } = createMockClient(true);
+      let resolveSwitch!: () => void;
+      const pendingSwitch = new Promise<void>((resolve) => {
+        resolveSwitch = resolve;
+      });
+      let result!: { current: ReturnType<typeof useA2ARuntime> };
+      let cancelled = false;
+      ({ result } = renderHook(() => {
+        const [threadId, setThreadId] = useState("initial");
+        return useA2ARuntime({
+          client,
+          onCancel: () => {
+            if (cancelled) return;
+            cancelled = true;
+            void result.current.thread.append("started by onCancel");
+          },
+          adapters: {
+            threadList: {
+              threadId,
+              onSwitchToThread: async (id) => {
+                setThreadId(id);
+                await pendingSwitch;
+                return { messages: [createThreadMessage("loaded")] };
+              },
+              onSwitchToNewThread: async () => {
+                setThreadId("thread-new");
+                await pendingSwitch;
+              },
+            },
+          },
+        });
+      }));
+
+      act(() => {
+        void result.current.thread.append("old prompt");
+      });
+      await waitFor(() => expect(streamMessage).toHaveBeenCalledOnce());
+
+      let switching!: Promise<void>;
+      act(() => {
+        switching =
+          destination === "existing"
+            ? result.current.threads.switchToThread("thread-a")
+            : result.current.threads.switchToNewThread();
+      });
+      act(() => {
+        void result.current.thread.append("during switch");
+      });
+      await act(async () => {
+        await new Promise((resolve) => setTimeout(resolve, 0));
+      });
+      await act(async () => {
+        resolveSwitch();
+        await switching;
+      });
+
+      expect(cancelled).toBe(true);
+      expect(streamMessage).toHaveBeenCalledOnce();
+      expect(result.current.thread.getState().isRunning).toBe(false);
+      expect(
+        result.current.thread.getState().messages.map((m) => m.id),
+      ).toEqual(destination === "existing" ? ["loaded"] : []);
+    },
+  );
+
   it("leaves a new thread empty after an active run and failed creation", async () => {
     const { client, streamMessage } = createMockClient(true);
     let rejectNew!: (error: Error) => void;

package/src/useA2ARuntime.ts

@@ -7,6 +7,7 @@ import {
   useMemo,
   useRef,
   useState,
+  useSyncExternalStore,
 } from "react";
 import {
   useExternalStoreRuntime,
@@ -43,8 +44,6 @@ const serializeManagedClientOptions = (
 };
 
 export function useA2ARuntime(options: UseA2ARuntimeOptions): AssistantRuntime {
-  const [_version, setVersion] = useState(0);
-  const notifyUpdate = useCallback(() => setVersion((v) => v + 1), []);
   const runtimeAdapters = useRuntimeAdapters();
   const historyAdapter = options.adapters?.history ?? runtimeAdapters?.history;
   const threadListAdapter = options.adapters?.threadList;
@@ -118,7 +117,6 @@ export function useA2ARuntime(options: UseA2ARuntimeOptions): AssistantRuntime {
     new A2AThreadRuntimeCore({
       ...coreOptionsRef.current,
       client,
-      notifyUpdate,
     });
   const [pinnedCore, setPinnedCore] = useState(() => ({
     client,
@@ -130,6 +128,11 @@ export function useA2ARuntime(options: UseA2ARuntimeOptions): AssistantRuntime {
     setPinnedCore(currentCore);
   }
   const core = currentCore.core;
+  const snapshot = useSyncExternalStore(
+    core.subscribe,
+    core.getSnapshot,
+    core.getSnapshot,
+  );
 
   useEffect(() => {
     core.updateOptions(coreOptions);
@@ -137,41 +140,58 @@ export function useA2ARuntime(options: UseA2ARuntimeOptions): AssistantRuntime {
 
   // Thread list
   const threadSwitchGenerationRef = useRef(0);
+  const switchingGenerationRef = useRef<number | null>(null);
   const threadList = useMemo(() => {
     if (!threadListAdapter) return undefined;
 
     const { onSwitchToNewThread, onSwitchToThread } = threadListAdapter;
 
+    const releaseSwitch = (generation: number) => {
+      if (switchingGenerationRef.current === generation) {
+        switchingGenerationRef.current = null;
+      }
+    };
+
     return {
       threadId: threadListAdapter.threadId,
       onSwitchToNewThread: onSwitchToNewThread
         ? async () => {
             const generation = ++threadSwitchGenerationRef.current;
-            // Clear before the thread id flips, or the old messages leak
-            // into the new thread as a sibling branch.
-            core.applyExternalMessages([]);
-            core.resetContext();
-            await onSwitchToNewThread();
-            if (generation !== threadSwitchGenerationRef.current) return;
-            // Apply first so the abort inside resetContext finds an already
-            // cleared repository and cannot persist the old thread's partial
-            // assistant message.
-            core.applyExternalMessages([]);
-            core.resetContext();
+            switchingGenerationRef.current = generation;
+            try {
+              // Clear before the thread id flips, or the old messages leak
+              // into the new thread as a sibling branch.
+              core.applyExternalMessages([]);
+              core.resetContext();
+              await onSwitchToNewThread();
+              if (generation !== threadSwitchGenerationRef.current) return;
+              // Apply first so the abort inside resetContext finds an already
+              // cleared repository and cannot persist the old thread's partial
+              // assistant message.
+              core.applyExternalMessages([]);
+              core.resetContext();
+            } finally {
+              releaseSwitch(generation);
+            }
           }
         : undefined,
       onSwitchToThread: onSwitchToThread
         ? async (threadId: string) => {
             const generation = ++threadSwitchGenerationRef.current;
-            // Clear before the thread id flips, or the old messages leak
-            // into the new thread as a sibling branch.
-            core.applyExternalMessages([]);
-            core.resetContext();
-            const result = await onSwitchToThread(threadId);
-            if (generation !== threadSwitchGenerationRef.current) return;
-            core.applyExternalMessages([]);
-            core.applyExternalMessages(result.messages);
-            core.resetContext();
+            switchingGenerationRef.current = generation;
+            try {
+              // Clear before the thread id flips, or the old messages leak
+              // into the new thread as a sibling branch.
+              core.applyExternalMessages([]);
+              core.resetContext();
+              const result = await onSwitchToThread(threadId);
+              if (generation !== threadSwitchGenerationRef.current) return;
+              core.applyExternalMessages([]);
+              core.applyExternalMessages(result.messages);
+              core.resetContext();
+            } finally {
+              releaseSwitch(generation);
+            }
           }
         : undefined,
     };
@@ -194,24 +214,31 @@ export function useA2ARuntime(options: UseA2ARuntimeOptions): AssistantRuntime {
   // Build store adapter
   const shared = useExternalStoreSharedOptions(options);
   const store = useMemo(() => {
-    void _version;
-
     return {
       ...shared,
-      isLoading: core.isLoading,
-      messageRepository: core.getMessageRepository(),
-      isRunning: core.isRunning(),
+      isLoading: snapshot.isLoading,
+      messageRepository: snapshot.messageRepository,
+      isRunning: snapshot.isRunning,
       unstable_persistsHistory: true,
       extras: a2aExtras.provide({
-        task: core.getTask(),
-        artifacts: core.getArtifacts(),
-        agentCard: core.getAgentCard(),
+        task: snapshot.task,
+        artifacts: snapshot.artifacts,
+        agentCard: snapshot.agentCard,
       }),
-      onNew: (message: AppendMessage) => core.append(message),
+      onNew: (message: AppendMessage) =>
+        switchingGenerationRef.current === null
+          ? core.append(message)
+          : Promise.resolve(),
       onVoiceTranscript: (message: ThreadMessage) =>
         core.appendVoiceTranscript(message),
-      onEdit: (message: AppendMessage) => core.edit(message),
-      onReload: (parentId: string | null) => core.reload(parentId),
+      onEdit: (message: AppendMessage) =>
+        switchingGenerationRef.current === null
+          ? core.edit(message)
+          : Promise.resolve(),
+      onReload: (parentId: string | null) =>
+        switchingGenerationRef.current === null
+          ? core.reload(parentId)
+          : Promise.resolve(),
       onCancel: () => core.cancel(),
       unstable_onRecordToolInteraction: (options) =>
         core.recordToolInteraction(options),
@@ -221,7 +248,7 @@ export function useA2ARuntime(options: UseA2ARuntimeOptions): AssistantRuntime {
         core.applyExternalMessages(messages),
       adapters: adapterAdapters,
     } satisfies ExternalStoreAdapter<ThreadMessage>;
-  }, [adapterAdapters, core, _version, shared]);
+  }, [adapterAdapters, core, snapshot, shared]);
 
   const runtime = useExternalStoreRuntime(store);