@assistant-ui/react-a2a 0.2.35 → 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;
@@ -83,18 +103,23 @@ export class A2AThreadRuntimeCore {
   private readonly session = createMessageRepositorySession();
   private isRunningFlag = false;
   private abortController: AbortController | null = null;
+  private runGeneration = 0;
   private pendingError: Error | null = null;
 
   // 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;
   private _loadRequested = false;
   private _agentCardPromise: Promise<void> | undefined;
   private _agentCardRetryAfter = 0;
@@ -146,6 +171,8 @@ export class A2AThreadRuntimeCore {
   /** Thread-boundary reset: applyExternalMessages alone also serves branch
    * switches, deletes, and cancel resyncs, which must keep the live context. */
   resetContext(): void {
+    this._historyLoadGeneration++;
+    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
     // must not be discarded.
@@ -184,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;
   }
@@ -257,16 +332,19 @@ export class A2AThreadRuntimeCore {
 
     this._isLoading = true;
 
+    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);
           this.finalizeExternalApply();
         }
       })
       .catch((error) => {
+        if (generation !== this._historyLoadGeneration) return;
         invokeRuntimeCallback(
           "onError",
           this.onError,
@@ -274,6 +352,7 @@ export class A2AThreadRuntimeCore {
         );
       })
       .finally(() => {
+        if (generation !== this._historyLoadGeneration) return;
         this._isLoading = false;
         this.notifyUpdate();
       });
@@ -305,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);
   }
@@ -328,14 +415,26 @@ export class A2AThreadRuntimeCore {
   async cancel(): Promise<void> {
     if (!this.abortController) return;
 
+    // Read the server target before aborting: the abort listener runs the
+    // 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;
+
     // Abort locally first so the stream stops immediately
     this.abortController.abort();
 
     // Then try to cancel the task on the server
-    if (this.currentTask?.id) {
+    if (task?.id) {
       try {
-        const updated = await this.client.cancelTask(this.currentTask.id);
-        this.currentTask = updated;
+        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
+        // that case out.
+        if (this.currentTask === task && this.runGeneration === generation) {
+          this.currentTask = updated;
+        }
       } catch {
         // Server cancel failed; local abort already handled
       }
@@ -366,6 +465,8 @@ export class A2AThreadRuntimeCore {
     }
     this.currentTask = undefined;
     this.currentArtifacts = [];
+    this.a2uiState = new Map();
+    this.a2uiMessageIds.clear();
     this.notifyUpdate();
   }
 
@@ -405,6 +506,8 @@ 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();
@@ -429,6 +532,8 @@ export class A2AThreadRuntimeCore {
     }
 
     this.currentArtifacts = [];
+    this.a2uiState = new Map();
+    this.a2uiMessageIds.clear();
 
     const assistantParentId = userThreadMessage.id;
     const assistantId = this.insertAssistantPlaceholder(assistantParentId);
@@ -561,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);
@@ -591,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);
     }
@@ -601,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(
@@ -632,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",
@@ -646,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();
@@ -673,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);
@@ -718,10 +862,124 @@ 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) {
     const touched = this.session.updateMessage(messageId, (message) => {
       if (message.role !== "assistant") return message;
@@ -729,7 +987,7 @@ export class A2AThreadRuntimeCore {
     });
     if (touched) {
       this.notifyUpdate();
-      if (status.type === "complete" || status.type === "incomplete") {
+      if (this.isPersistableAssistantStatus(status)) {
         this.persistAssistantHistory(messageId);
       }
     }
@@ -741,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) {
@@ -757,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(
@@ -770,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;
   }
 }