@assistant-ui/react-a2a 0.2.36 → 0.2.37

package/src/A2AThreadRuntimeCore.ts

@@ -7,13 +7,24 @@ import type {
   ExportedMessageRepository,
   MessageStatus,
   ThreadAssistantMessage,
+  ThreadAssistantMessagePart,
   ThreadHistoryAdapter,
   ThreadMessage,
+  ToolCallMessagePart,
+  Unstable_RecordToolInteractionOptions,
 } from "@assistant-ui/core";
 import {
+  appendToolInteraction,
   createMessageRepositorySession,
   invokeUserCallback,
 } from "@assistant-ui/core/internal";
+import type { ReadonlyJSONObject } from "assistant-stream/utils";
+import {
+  applyA2uiOperations,
+  convertSurfaceToUISpec,
+  surfaceToOperations,
+  type A2uiState,
+} from "@assistant-ui/react-generative-ui/a2ui";
 import type { A2AClient } from "./A2AClient";
 import type {
   A2AArtifact,
@@ -28,6 +39,7 @@ import type {
 
 import {
   a2aMessageToContent,
+  a2uiPartsToOperations,
   isTerminalTaskState,
   threadMessageToA2AMessage,
   taskStateToMessageStatus,
@@ -36,6 +48,14 @@ import {
 const INITIAL_AGENT_CARD_RETRY_DELAY_MS = 5_000;
 const MAX_AGENT_CARD_RETRY_DELAY_MS = 5 * 60_000;
 
+const isRecord = (value: unknown): value is Record<string, unknown> =>
+  typeof value === "object" && value !== null && !Array.isArray(value);
+
+const storedArtifact = (artifact: A2AArtifact): A2AArtifact => ({
+  ...artifact,
+  parts: artifact.parts.map(({ raw: _raw, ...part }) => part),
+});
+
 export type A2AThreadRuntimeCoreOptions = {
   client: A2AClient;
   contextId?: string | undefined;
@@ -89,11 +109,14 @@ export class A2AThreadRuntimeCore {
   // A2A-specific state
   private currentTask: A2ATask | undefined;
   private currentArtifacts: A2AArtifact[] = [];
+  private a2uiState: A2uiState = new Map();
+  private readonly a2uiMessageIds = new Set<string>();
   private agentCardValue: A2AAgentCard | undefined;
 
   // History tracking
   private readonly assistantHistoryParents = new Map<string, string | null>();
   private readonly recordedHistoryIds = new Set<string>();
+  private readonly historyWrites = new Map<string, Promise<void>>();
   private _isLoading = false;
   private _loadPromise: Promise<void> | undefined;
   private _historyLoadGeneration = 0;
@@ -188,6 +211,54 @@ export class A2AThreadRuntimeCore {
     return this.session.export();
   }
 
+  public async recordToolInteraction({
+    messageId,
+    toolCallId,
+    interaction,
+  }: Unstable_RecordToolInteractionOptions): Promise<void> {
+    const message = this.session.tryGetMessage(messageId)?.message;
+    if (!message) {
+      throw new Error(
+        "Tried to record a tool interaction on a non-existing message",
+      );
+    }
+    if (message.role !== "assistant") {
+      throw new Error(
+        "Tried to record a tool interaction on a non-assistant message",
+      );
+    }
+    const target = message.content.find(
+      (part) => part.type === "tool-call" && part.toolCallId === toolCallId,
+    );
+    if (!target || target.type !== "tool-call") {
+      throw new Error(
+        "Tried to record a tool interaction on a non-existing tool call",
+      );
+    }
+
+    const touched = this.session.updateMessage(messageId, (current) => {
+      if (current.role !== "assistant") return current;
+      return {
+        ...current,
+        content: current.content.map((part) =>
+          part.type === "tool-call" && part.toolCallId === toolCallId
+            ? {
+                ...part,
+                unstable_interactions: appendToolInteraction(
+                  part.unstable_interactions,
+                  interaction,
+                ),
+              }
+            : part,
+        ),
+      };
+    });
+    if (touched) {
+      this.notifyUpdate();
+      this.persistAssistantHistory(messageId);
+    }
+  }
+
   getTask(): A2ATask | undefined {
     return this.currentTask;
   }
@@ -264,8 +335,8 @@ export class A2AThreadRuntimeCore {
     const generation = this._historyLoadGeneration;
     const historyPromise = this.history.load();
 
-    this._loadPromise = Promise.all([historyPromise, agentCardPromise])
-      .then(([repo]) => {
+    this._loadPromise = historyPromise
+      .then((repo) => {
         if (generation !== this._historyLoadGeneration) return;
         if (repo) {
           this.session.applyExternalMessageRepository(repo);
@@ -313,6 +384,14 @@ export class A2AThreadRuntimeCore {
     await this.startRun(threadMessage);
   }
 
+  appendVoiceTranscript(message: ThreadMessage): void {
+    const parentId = this.session.headId;
+    this.session.addOrUpdateMessage(parentId, message);
+    this.session.switchToBranch(message.id);
+    this.notifyUpdate();
+    this.recordHistoryEntry(parentId, message);
+  }
+
   async edit(message: AppendMessage): Promise<void> {
     await this.append(message);
   }
@@ -386,6 +465,8 @@ export class A2AThreadRuntimeCore {
     }
     this.currentTask = undefined;
     this.currentArtifacts = [];
+    this.a2uiState = new Map();
+    this.a2uiMessageIds.clear();
     this.notifyUpdate();
   }
 
@@ -451,6 +532,8 @@ export class A2AThreadRuntimeCore {
     }
 
     this.currentArtifacts = [];
+    this.a2uiState = new Map();
+    this.a2uiMessageIds.clear();
 
     const assistantParentId = userThreadMessage.id;
     const assistantId = this.insertAssistantPlaceholder(assistantParentId);
@@ -583,7 +666,7 @@ export class A2AThreadRuntimeCore {
         this.handleStatusUpdate(assistantId, event.event);
         break;
       case "artifactUpdate":
-        this.handleArtifactUpdate(event.event);
+        this.handleArtifactUpdate(assistantId, event.event);
         break;
       case "message":
         this.handleMessage(assistantId, event.message);
@@ -613,6 +696,7 @@ export class A2AThreadRuntimeCore {
     }
 
     if (event.status.message) {
+      this.applyA2uiMessage(event.status.message);
       const content = a2aMessageToContent(event.status.message);
       this.updateAssistantContent(assistantId, content);
     }
@@ -623,7 +707,10 @@ export class A2AThreadRuntimeCore {
     this.notifyUpdate();
   }
 
-  private handleArtifactUpdate(event: A2ATaskArtifactUpdateEvent) {
+  private handleArtifactUpdate(
+    assistantId: string,
+    event: A2ATaskArtifactUpdateEvent,
+  ) {
     const { append, lastChunk } = event;
     const artifact = normalizeArtifact(event.artifact);
     const existingIdx = this.currentArtifacts.findIndex(
@@ -654,6 +741,10 @@ export class A2AThreadRuntimeCore {
       this.currentArtifacts = [...this.currentArtifacts, updated];
     }
 
+    this.applyA2uiParts(artifact.parts);
+    this.updateAssistantArtifacts(assistantId);
+    this.rebuildAssistantA2uiSurfaces(assistantId);
+
     if (lastChunk) {
       invokeRuntimeCallback(
         "onArtifactComplete",
@@ -668,6 +759,7 @@ export class A2AThreadRuntimeCore {
   private handleMessage(assistantId: string, message: A2AMessage) {
     if (message.role !== "agent") return;
 
+    this.applyA2uiMessage(message);
     const content = a2aMessageToContent(message);
     this.updateAssistantContent(assistantId, content);
     this.notifyUpdate();
@@ -695,13 +787,43 @@ export class A2AThreadRuntimeCore {
     if (task.contextId) {
       this.contextId = task.contextId;
     }
-    if (artifacts) {
+    const isCompleteSnapshot = history !== undefined && artifacts !== undefined;
+    const artifactsToApply = isCompleteSnapshot
+      ? artifacts
+      : artifacts?.filter(
+          (artifact) =>
+            !this.currentArtifacts.some(
+              ({ artifactId }) => artifactId === artifact.artifactId,
+            ),
+        );
+
+    if (isCompleteSnapshot) {
       this.currentArtifacts = artifacts;
+      this.a2uiState = new Map();
+      this.a2uiMessageIds.clear();
+    } else if (artifactsToApply) {
+      this.currentArtifacts = [...this.currentArtifacts, ...artifactsToApply];
+    }
+
+    for (const message of history ?? []) {
+      if (message.role === "agent") this.applyA2uiMessage(message, true);
+    }
+    for (const artifact of artifactsToApply ?? []) {
+      this.applyA2uiParts(artifact.parts);
+    }
+    if (task.status.message) {
+      this.applyA2uiMessage(task.status.message, true);
+    }
+
+    if (isCompleteSnapshot || artifactsToApply?.length) {
+      this.updateAssistantArtifacts(assistantId);
     }
 
     if (task.status.message) {
       const content = a2aMessageToContent(task.status.message);
       this.updateAssistantContent(assistantId, content);
+    } else {
+      this.rebuildAssistantA2uiSurfaces(assistantId);
     }
 
     const status = taskStateToMessageStatus(task.status.state);
@@ -740,8 +862,122 @@ export class A2AThreadRuntimeCore {
   ) {
     this.session.updateMessage(messageId, (message) => {
       if (message.role !== "assistant") return message;
-      return { ...message, content };
+      return {
+        ...message,
+        content: this.withA2uiSurfaces(content, message.content),
+      };
+    });
+  }
+
+  private applyA2uiMessage(message: A2AMessage, replay = false) {
+    if (message.messageId) {
+      if (replay && this.a2uiMessageIds.has(message.messageId)) return;
+      this.a2uiMessageIds.add(message.messageId);
+    }
+    this.applyA2uiParts(message.parts);
+  }
+
+  private applyA2uiParts(parts: readonly A2AMessage["parts"][number][]) {
+    const operations = a2uiPartsToOperations(parts);
+    if (operations.length === 0) return;
+    this.a2uiState = applyA2uiOperations(this.a2uiState, operations).state;
+  }
+
+  private a2uiSurfaceParts(): ThreadAssistantMessagePart[] {
+    const parts: ThreadAssistantMessagePart[] = [];
+    for (const [surfaceId, surface] of this.a2uiState) {
+      const { spec } = convertSurfaceToUISpec(surface);
+      if (!spec) 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) },
+      });
+    }
+    return parts;
+  }
+
+  private withA2uiSurfaces(
+    content: ThreadAssistantMessage["content"],
+    previousContent: ThreadAssistantMessage["content"] = content,
+  ): ThreadAssistantMessage["content"] {
+    const preserved = content.filter(
+      (part) =>
+        !(part.type === "tool-call" && part.toolCallId.startsWith("a2ui:")),
+    );
+    const interactionSource = [...previousContent, ...content];
+    return this.withPreservedToolInteractions(interactionSource, [
+      ...preserved,
+      ...this.a2uiSurfaceParts(),
+    ]);
+  }
+
+  private withPreservedToolInteractions(
+    previousContent: ThreadAssistantMessage["content"],
+    content: ThreadAssistantMessage["content"],
+  ): ThreadAssistantMessage["content"] {
+    const interactions = new Map<
+      string,
+      ToolCallMessagePart["unstable_interactions"]
+    >();
+    for (const part of previousContent) {
+      if (
+        part.type === "tool-call" &&
+        part.unstable_interactions !== undefined
+      ) {
+        interactions.set(part.toolCallId, part.unstable_interactions);
+      }
+    }
+    if (interactions.size === 0) return content;
+    return content.map((part) => {
+      if (
+        part.type !== "tool-call" ||
+        part.unstable_interactions !== undefined
+      ) {
+        return part;
+      }
+      const unstable_interactions = interactions.get(part.toolCallId);
+      return unstable_interactions === undefined
+        ? part
+        : { ...part, unstable_interactions };
+    });
+  }
+
+  private rebuildAssistantA2uiSurfaces(messageId: string) {
+    const touched = this.session.updateMessage(messageId, (message) => {
+      if (message.role !== "assistant") return message;
+      return {
+        ...message,
+        content: this.withA2uiSurfaces(message.content),
+      };
+    });
+    if (touched) this.notifyUpdate();
+  }
+
+  private updateAssistantArtifacts(messageId: string) {
+    const touched = this.session.updateMessage(messageId, (message) => {
+      if (message.role !== "assistant") return message;
+      const custom = message.metadata.custom;
+      const a2a = isRecord(custom.a2a) ? custom.a2a : {};
+      return {
+        ...message,
+        metadata: {
+          ...message.metadata,
+          custom: {
+            ...custom,
+            a2a: {
+              ...a2a,
+              artifacts: this.currentArtifacts.map(storedArtifact),
+            },
+          },
+        },
+      };
     });
+    if (touched) this.notifyUpdate();
   }
 
   private updateAssistantStatus(messageId: string, status: MessageStatus) {
@@ -751,7 +987,7 @@ export class A2AThreadRuntimeCore {
     });
     if (touched) {
       this.notifyUpdate();
-      if (status.type === "complete" || status.type === "incomplete") {
+      if (this.isPersistableAssistantStatus(status)) {
         this.persistAssistantHistory(messageId);
       }
     }
@@ -763,6 +999,14 @@ export class A2AThreadRuntimeCore {
     return msg.status;
   }
 
+  private isPersistableAssistantStatus(status: MessageStatus): boolean {
+    return (
+      status.type === "complete" ||
+      status.type === "incomplete" ||
+      (status.type === "requires-action" && status.reason === "interrupt")
+    );
+  }
+
   // --- Lifecycle helpers ---
 
   private setRunning(running: boolean) {
@@ -779,7 +1023,9 @@ export class A2AThreadRuntimeCore {
   // --- History persistence ---
 
   private recordHistoryEntry(parentId: string | null, message: ThreadMessage) {
-    this.appendHistoryItem(parentId, message);
+    void this.appendHistoryItem(parentId, message)?.catch((error) => {
+      console.error("[react-a2a] failed to append history entry", error);
+    });
   }
 
   private markPendingAssistantHistory(
@@ -792,24 +1038,107 @@ export class A2AThreadRuntimeCore {
 
   private persistAssistantHistory(messageId: string) {
     if (!this.history) return;
+    const messageData = this.session.tryGetMessage(messageId);
     const parentId = this.assistantHistoryParents.get(messageId);
-    if (parentId === undefined) return;
-    const message = this.session.tryGetMessage(messageId)?.message;
+    if (parentId === undefined && !this.recordedHistoryIds.has(messageId)) {
+      return;
+    }
+    const resolvedParentId =
+      parentId === undefined ? messageData?.parentId : parentId;
+    if (resolvedParentId === undefined) return;
+    const message = messageData?.message;
     if (!message || message.role !== "assistant") return;
-    if (
-      message.status?.type !== "complete" &&
-      message.status?.type !== "incomplete"
-    )
+    if (!this.isPersistableAssistantStatus(message.status)) return;
+    const isPausing = message.status.type === "requires-action";
+
+    if (isPausing && !this.history.update) return;
+
+    if (this.recordedHistoryIds.has(messageId)) {
+      if (this.history.update) {
+        const update = this.history.update.bind(this.history);
+        const write = this.chainHistoryWrite(messageId, () =>
+          update({ parentId: resolvedParentId, message }),
+        );
+        if (!isPausing) {
+          this.assistantHistoryParents.delete(messageId);
+        }
+        void write.then(
+          () => {
+            this.recordedHistoryIds.add(messageId);
+          },
+          (error) => {
+            const pending = this.historyWrites.get(messageId);
+            if (pending === undefined || pending === write) {
+              this.assistantHistoryParents.set(messageId, resolvedParentId);
+            }
+            console.error("[react-a2a] failed to update history entry", error);
+          },
+        );
+        return;
+      }
+      if (!isPausing) {
+        this.assistantHistoryParents.delete(messageId);
+      }
       return;
-    this.assistantHistoryParents.delete(messageId);
-    this.appendHistoryItem(parentId, message);
+    }
+    const write = this.appendHistoryItem(resolvedParentId, message);
+    if (!write) return;
+    if (!isPausing) {
+      this.assistantHistoryParents.delete(messageId);
+    }
+    void write.catch((error) => {
+      const pending = this.historyWrites.get(messageId);
+      if (pending === undefined || pending === write) {
+        this.assistantHistoryParents.set(messageId, resolvedParentId);
+      }
+      console.error("[react-a2a] failed to append history entry", error);
+    });
   }
 
-  private appendHistoryItem(parentId: string | null, message: ThreadMessage) {
+  private appendHistoryItem(
+    parentId: string | null,
+    message: ThreadMessage,
+  ): Promise<void> | undefined {
     if (!this.history || this.recordedHistoryIds.has(message.id)) return;
     this.recordedHistoryIds.add(message.id);
-    void this.history.append({ parentId, message }).catch(() => {
+    const append = this.history.append.bind(this.history);
+    const write = this.chainHistoryWrite(message.id, () =>
+      append({ parentId, message }),
+    );
+    void write.catch(() => {
       this.recordedHistoryIds.delete(message.id);
     });
+    return write;
+  }
+
+  private chainHistoryWrite(
+    id: string,
+    write: () => Promise<void>,
+  ): Promise<void> {
+    const pending = this.historyWrites.get(id);
+    let next: Promise<void>;
+    if (pending) {
+      next = pending.then(write, write);
+    } else {
+      try {
+        next = Promise.resolve(write());
+      } catch (error) {
+        next = Promise.reject(error);
+      }
+    }
+    this.historyWrites.set(id, next);
+    void next.then(
+      () => {
+        if (this.historyWrites.get(id) === next) {
+          this.historyWrites.delete(id);
+        }
+      },
+      () => {
+        if (this.historyWrites.get(id) === next) {
+          this.historyWrites.delete(id);
+        }
+      },
+    );
+    return next;
   }
 }

package/src/conversions.test.ts

@@ -3,6 +3,7 @@ import {
   a2aPartToContent,
   a2aPartsToContent,
   a2aMessageToContent,
+  isA2uiDataPart,
   taskStateToMessageStatus,
   contentPartsToA2AParts,
   isTerminalTaskState,
@@ -157,6 +158,82 @@ describe("a2aPartToContent", () => {
   });
 });
 
+describe("A2UI data parts", () => {
+  it.each([
+    [
+      "media type",
+      {
+        mediaType: "application/vnd.A2UI+json",
+        data: { value: "from media type" },
+      },
+    ],
+    [
+      "metadata",
+      {
+        metadata: { mimeType: "application/a2ui+json" },
+        data: { value: "from metadata" },
+      },
+    ],
+    [
+      "metadata media type",
+      {
+        metadata: { mediaType: "application/a2ui+json" },
+        data: { value: "from metadata media type" },
+      },
+    ],
+    [
+      "operation shape",
+      {
+        data: {
+          version: "v0.9",
+          createSurface: { surfaceId: "surface" },
+        },
+      },
+    ],
+    [
+      "operation array",
+      {
+        data: [
+          {
+            version: "v0.9",
+            createSurface: { surfaceId: "surface" },
+          },
+          {
+            version: "v0.9",
+            deleteSurface: { surfaceId: "surface" },
+          },
+        ],
+      },
+    ],
+  ] as const)("detects A2UI by %s", (_source, part) => {
+    expect(isA2uiDataPart(part)).toBe(true);
+    expect(a2aPartsToContent([part])).toEqual([]);
+  });
+
+  it("does not detect an operation array containing a non-operation", () => {
+    const part = {
+      data: [
+        {
+          version: "v0.9",
+          createSurface: { surfaceId: "surface" },
+        },
+        { value: "not an operation" },
+      ],
+    };
+
+    expect(isA2uiDataPart(part)).toBe(false);
+  });
+
+  it("keeps ordinary data parts as JSON text", () => {
+    expect(a2aPartsToContent([{ data: { value: "plain" } }])).toEqual([
+      {
+        type: "text",
+        text: '{\n  "value": "plain"\n}',
+      },
+    ]);
+  });
+});
+
 describe("inbound file part round trip", () => {
   const restoreFilePart = (part: A2APart) => {
     const restored = a2aPartToContent(part);
@@ -690,4 +767,43 @@ describe("threadMessageToA2AMessage", () => {
     expect(result.contextId).toBeUndefined();
     expect(result.taskId).toBeUndefined();
   });
+
+  it("keeps tool interactions out of outbound messages", () => {
+    const result = threadMessageToA2AMessage({
+      ...userMessage,
+      role: "user",
+      attachments: [],
+      content: [
+        { type: "text", text: "hello" },
+        {
+          type: "file",
+          data: "ZmlsZQ==",
+          mimeType: "text/plain",
+          filename: "file.txt",
+        },
+        {
+          type: "tool-call",
+          toolCallId: "tool-1",
+          toolName: "present",
+          args: {},
+          argsText: "{}",
+          result: {},
+          unstable_interactions: {
+            entries: [
+              {
+                type: "action",
+                occurredAt: 1_700_000_000_000,
+                payload: { value: "selected" },
+              },
+            ],
+          },
+        },
+      ],
+    } as any);
+
+    expect(result.parts).toEqual([
+      { text: "hello" },
+      { raw: "ZmlsZQ==", mediaType: "text/plain", filename: "file.txt" },
+    ]);
+  });
 });