@assistant-ui/react-a2a 0.2.47 → 0.2.48

package/src/{A2AThreadRuntimeCore.ts → A2AThreadController.ts}

@@ -15,15 +15,12 @@ import type {
 } from "@assistant-ui/core";
 import {
   appendToolInteraction,
+  createHistoryWriteChain,
   createMessageRepositorySession,
   invokeUserCallback,
   RunLeases,
 } from "@assistant-ui/core/internal";
-import {
-  applyA2uiOperations,
-  surfaceToPresentToolCall,
-  type A2uiState,
-} from "@assistant-ui/generative-ui/a2ui";
+import { surfaceToPresentToolCall } from "@assistant-ui/generative-ui/a2ui";
 import type { A2AClient } from "./A2AClient";
 import type {
   A2AArtifact,
@@ -38,11 +35,16 @@ import type {
 
 import {
   a2aMessageToContent,
-  a2uiPartsToOperations,
   isTerminalTaskState,
   threadMessageToA2AMessage,
   taskStateToMessageStatus,
 } from "./conversions";
+import {
+  createA2AThreadState,
+  reduceA2AThreadState,
+  type A2AThreadAction,
+  type A2AThreadState,
+} from "./a2aThreadState";
 
 const INITIAL_AGENT_CARD_RETRY_DELAY_MS = 5_000;
 const MAX_AGENT_CARD_RETRY_DELAY_MS = 5 * 60_000;
@@ -55,7 +57,7 @@ const storedArtifact = (artifact: A2AArtifact): A2AArtifact => ({
   parts: artifact.parts.map(({ raw: _raw, ...part }) => part),
 });
 
-export type A2AThreadRuntimeCoreOptions = {
+export type A2AThreadControllerOptions = {
   client: A2AClient;
   contextId?: string | undefined;
   configuration?: A2ASendMessageConfiguration | undefined;
@@ -80,16 +82,9 @@ const invokeRuntimeCallback = <TArgs extends readonly unknown[]>(
   void invokeUserCallback("react-a2a", name, callback, ...args);
 };
 
-function normalizeArtifact(artifact: A2AArtifact): A2AArtifact {
-  return {
-    ...artifact,
-    parts: Array.isArray(artifact.parts) ? artifact.parts : [],
-  };
-}
-
-export class A2AThreadRuntimeCore {
+export class A2AThreadController {
   private client: A2AClient;
-  private contextId: string | undefined;
+  private state: A2AThreadState;
   private configuration: A2ASendMessageConfiguration | undefined;
   private onError: ((error: Error) => void) | undefined;
   private onCancel: (() => void) | undefined;
@@ -109,26 +104,16 @@ export class A2AThreadRuntimeCore {
 
   private runtime: AssistantRuntime | undefined;
   private readonly session = createMessageRepositorySession();
-  private isRunningFlag = false;
   private abortController: AbortController | null = null;
   private readonly runLeases = new RunLeases();
   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 readonly historyWrites = createHistoryWriteChain();
   private _loadPromise: Promise<void> | undefined;
   private _historyLoadGeneration = 0;
-  private _loadRequested = false;
   private _agentCardPromise: Promise<void> | undefined;
   private _agentCardRetryAfter = 0;
   private _agentCardRetryDelay = INITIAL_AGENT_CARD_RETRY_DELAY_MS;
@@ -136,9 +121,9 @@ export class A2AThreadRuntimeCore {
 
   private lastOptionsContextId: string | undefined;
 
-  constructor(options: A2AThreadRuntimeCoreOptions) {
+  constructor(options: A2AThreadControllerOptions) {
     this.client = options.client;
-    this.contextId = options.contextId;
+    this.state = createA2AThreadState(options.contextId);
     this.lastOptionsContextId = options.contextId;
     this.configuration = options.configuration;
     this.onError = options.onError;
@@ -154,25 +139,32 @@ export class A2AThreadRuntimeCore {
 
   getSnapshot = () =>
     (this.snapshot ??= {
-      isLoading: this._isLoading,
-      isRunning: this.isRunningFlag,
+      isLoading: this.state.isLoading,
+      isRunning: this.state.isRunning,
       messageRepository: this.session.export(),
-      task: this.currentTask,
-      artifacts: this.currentArtifacts,
-      agentCard: this.agentCardValue,
+      task: this.state.currentTask,
+      artifacts: this.state.currentArtifacts,
+      agentCard: this.state.agentCard,
     });
 
+  private dispatch(action: A2AThreadAction): void {
+    this.state = reduceA2AThreadState(this.state, action);
+  }
+
   private publishUpdate(): void {
     this.snapshot = undefined;
     for (const listener of this.listeners) listener();
   }
 
-  updateOptions(options: A2AThreadRuntimeCoreOptions) {
+  updateOptions(options: A2AThreadControllerOptions) {
     this.client = options.client;
     // 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.dispatch({
+        type: "contextLearned",
+        contextId: options.contextId,
+      });
       this.lastOptionsContextId = options.contextId;
     }
     this.configuration = options.configuration;
@@ -183,7 +175,7 @@ export class A2AThreadRuntimeCore {
     this.history = options.history;
 
     if (
-      this._loadRequested &&
+      this.state.loadRequested &&
       !this._loadPromise &&
       !previousHistory &&
       options.history &&
@@ -197,13 +189,15 @@ 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;
+    const wasLoading = this.state.isLoading;
     // 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.
     const controller = this.abortController;
-    this.contextId = this.lastOptionsContextId;
+    this.dispatch({
+      type: "resetContext",
+      contextId: this.lastOptionsContextId,
+    });
     if (controller) {
       controller.abort();
       if (this.abortController === controller) {
@@ -290,23 +284,23 @@ export class A2AThreadRuntimeCore {
   }
 
   getTask(): A2ATask | undefined {
-    return this.currentTask;
+    return this.state.currentTask;
   }
 
   getArtifacts(): readonly A2AArtifact[] {
-    return this.currentArtifacts;
+    return this.state.currentArtifacts;
   }
 
   getAgentCard(): A2AAgentCard | undefined {
-    return this.agentCardValue;
+    return this.state.agentCard;
   }
 
   isRunning(): boolean {
-    return this.isRunningFlag;
+    return this.state.isRunning;
   }
 
   get isLoading(): boolean {
-    return this._isLoading;
+    return this.state.isLoading;
   }
 
   private loadAgentCard(): Promise<void> {
@@ -314,7 +308,10 @@ export class A2AThreadRuntimeCore {
 
     this._agentCardPromise ??= this.client.getAgentCard().then(
       (agentCard) => {
-        this.agentCardValue = agentCard;
+        this.dispatch({
+          type: "agentCardLoaded",
+          agentCard,
+        });
         this._agentCardRetryAfter = 0;
         this._agentCardRetryDelay = INITIAL_AGENT_CARD_RETRY_DELAY_MS;
         this._agentCardDiscoveryFailed = false;
@@ -354,13 +351,13 @@ export class A2AThreadRuntimeCore {
   }
 
   __internal_load(): Promise<void> {
-    this._loadRequested = true;
+    this.dispatch({ type: "loadRequested" });
     const agentCardPromise = this.loadAgentCard();
 
     if (this._loadPromise) return this._loadPromise;
     if (!this.history) return agentCardPromise;
 
-    this._isLoading = true;
+    this.dispatch({ type: "loadStarted" });
 
     const generation = this._historyLoadGeneration;
     const historyPromise = this.history.load();
@@ -383,7 +380,7 @@ export class A2AThreadRuntimeCore {
       })
       .finally(() => {
         if (generation !== this._historyLoadGeneration) return;
-        this._isLoading = false;
+        this.dispatch({ type: "loadFinished" });
         this.publishUpdate();
       });
 
@@ -448,7 +445,7 @@ export class A2AThreadRuntimeCore {
     // 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 task = this.state.currentTask;
     const lease = this.runLeases.current();
 
     // Abort locally first so the stream stops immediately
@@ -462,8 +459,11 @@ export class A2AThreadRuntimeCore {
         // or a cleared thread replaces the task object; a follow-up run that
         // has not emitted yet keeps it, so the run lease is what rules
         // that case out.
-        if (this.currentTask === task && lease.isCurrent()) {
-          this.currentTask = updated;
+        if (this.state.currentTask === task && lease.isCurrent()) {
+          this.dispatch({
+            type: "taskReplaced",
+            task: updated,
+          });
           this.publishUpdate();
         }
       } catch {
@@ -494,10 +494,7 @@ export class A2AThreadRuntimeCore {
     for (const { message } of this.getMessageRepository().messages) {
       this.recordedHistoryIds.add(message.id);
     }
-    this.currentTask = undefined;
-    this.currentArtifacts = [];
-    this.a2uiState = new Map();
-    this.a2uiMessageIds.clear();
+    this.dispatch({ type: "reset" });
     this.publishUpdate();
   }
 
@@ -552,25 +549,15 @@ export class A2AThreadRuntimeCore {
     }
 
     const a2aMessage = threadMessageToA2AMessage(userThreadMessage, {
-      contextId: this.contextId,
+      contextId: this.state.contextId,
       taskId:
-        this.currentTask?.id &&
-        !isTerminalTaskState(this.currentTask.status.state)
-          ? this.currentTask.id
+        this.state.currentTask?.id &&
+        !isTerminalTaskState(this.state.currentTask.status.state)
+          ? this.state.currentTask.id
           : undefined,
     });
 
-    // Clear task if previous task reached terminal state
-    if (
-      this.currentTask &&
-      isTerminalTaskState(this.currentTask.status.state)
-    ) {
-      this.currentTask = undefined;
-    }
-
-    this.currentArtifacts = [];
-    this.a2uiState = new Map();
-    this.a2uiMessageIds.clear();
+    this.dispatch({ type: "prepareRun" });
 
     const assistantParentId = userThreadMessage.id;
     const assistantId = this.insertAssistantPlaceholder(assistantParentId);
@@ -598,7 +585,7 @@ export class A2AThreadRuntimeCore {
       if (!(await this.waitForAgentCard(abortController.signal))) return;
 
       const supportsStreaming =
-        this.agentCardValue?.capabilities?.streaming !== false;
+        this.state.agentCard?.capabilities?.streaming !== false;
       if (supportsStreaming) {
         await this.runStreaming(a2aMessage, assistantId, abortController);
       } else {
@@ -718,22 +705,12 @@ export class A2AThreadRuntimeCore {
     assistantId: string,
     event: A2ATaskStatusUpdateEvent,
   ) {
-    if (!this.currentTask) {
-      this.currentTask = {
-        id: event.taskId,
-        contextId: event.contextId,
-        status: event.status,
-      };
-    } else {
-      this.currentTask = { ...this.currentTask, status: event.status };
-    }
-
-    if (event.contextId) {
-      this.contextId = event.contextId;
-    }
+    this.dispatch({
+      type: "statusUpdate",
+      event,
+    });
 
     if (event.status.message) {
-      this.applyA2uiMessage(event.status.message);
       const content = a2aMessageToContent(event.status.message);
       this.updateAssistantContent(assistantId, content);
     }
@@ -748,41 +725,17 @@ export class A2AThreadRuntimeCore {
     assistantId: string,
     event: A2ATaskArtifactUpdateEvent,
   ) {
-    const { append, lastChunk } = event;
-    const artifact = normalizeArtifact(event.artifact);
-    const existingIdx = this.currentArtifacts.findIndex(
-      (a) => a.artifactId === artifact.artifactId,
-    );
-
-    let updated: A2AArtifact;
-    if (existingIdx >= 0 && append) {
-      const existing = this.currentArtifacts[existingIdx]!;
-      updated = {
-        ...existing,
-        parts: [...existing.parts, ...artifact.parts],
-      };
-      this.currentArtifacts = [
-        ...this.currentArtifacts.slice(0, existingIdx),
-        updated,
-        ...this.currentArtifacts.slice(existingIdx + 1),
-      ];
-    } else if (existingIdx >= 0) {
-      updated = artifact;
-      this.currentArtifacts = [
-        ...this.currentArtifacts.slice(0, existingIdx),
-        updated,
-        ...this.currentArtifacts.slice(existingIdx + 1),
-      ];
-    } else {
-      updated = artifact;
-      this.currentArtifacts = [...this.currentArtifacts, updated];
-    }
-
-    this.applyA2uiParts(artifact.parts);
+    this.dispatch({
+      type: "artifactUpdate",
+      event,
+    });
+    const updated = this.state.currentArtifacts.find(
+      ({ artifactId }) => artifactId === event.artifact.artifactId,
+    )!;
     this.updateAssistantArtifacts(assistantId);
     this.rebuildAssistantA2uiSurfaces(assistantId);
 
-    if (lastChunk) {
+    if (event.lastChunk) {
       invokeRuntimeCallback(
         "onArtifactComplete",
         this.onArtifactComplete,
@@ -796,63 +749,26 @@ export class A2AThreadRuntimeCore {
   private handleMessage(assistantId: string, message: A2AMessage) {
     if (message.role !== "agent") return;
 
-    this.applyA2uiMessage(message);
+    this.dispatch({
+      type: "a2uiMessage",
+      message,
+    });
     const content = a2aMessageToContent(message);
     this.updateAssistantContent(assistantId, content);
     this.publishUpdate();
   }
 
   private handleTaskSnapshot(assistantId: string, task: A2ATask) {
-    const artifacts =
-      task.artifacts === undefined
-        ? undefined
-        : Array.isArray(task.artifacts)
-          ? task.artifacts.map(normalizeArtifact)
-          : [];
-    const history =
-      task.history === undefined
-        ? undefined
-        : Array.isArray(task.history)
-          ? task.history
-          : [];
-    this.currentTask = {
-      ...task,
-      ...(artifacts === undefined ? {} : { artifacts }),
-      ...(history === undefined ? {} : { history }),
-    };
-
-    if (task.contextId) {
-      this.contextId = task.contextId;
-    }
-    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);
-    }
+    const previousArtifacts = this.state.currentArtifacts;
+    this.dispatch({
+      type: "taskSnapshot",
+      task,
+    });
 
-    if (isCompleteSnapshot || artifactsToApply?.length) {
+    if (
+      (task.history !== undefined && task.artifacts !== undefined) ||
+      this.state.currentArtifacts.length > previousArtifacts.length
+    ) {
       this.updateAssistantArtifacts(assistantId);
     }
 
@@ -906,23 +822,9 @@ export class A2AThreadRuntimeCore {
     });
   }
 
-  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) {
+    for (const [surfaceId, surface] of this.state.a2uiState) {
       const { toolCall } = surfaceToPresentToolCall(surfaceId, surface);
       if (!toolCall) continue;
       parts.push({
@@ -1003,7 +905,7 @@ export class A2AThreadRuntimeCore {
             ...custom,
             a2a: {
               ...a2a,
-              artifacts: this.currentArtifacts.map(storedArtifact),
+              artifacts: this.state.currentArtifacts.map(storedArtifact),
             },
           },
         },
@@ -1042,7 +944,9 @@ export class A2AThreadRuntimeCore {
   // --- Lifecycle helpers ---
 
   private setRunning(running: boolean) {
-    this.isRunningFlag = running;
+    this.dispatch({
+      type: running ? "runStarted" : "runFinished",
+    });
     this.publishUpdate();
   }
 
@@ -1088,7 +992,7 @@ export class A2AThreadRuntimeCore {
     if (this.recordedHistoryIds.has(messageId)) {
       if (this.history.update) {
         const update = this.history.update.bind(this.history);
-        const write = this.chainHistoryWrite(messageId, () =>
+        const write = this.historyWrites.chain(messageId, () =>
           update({ parentId: resolvedParentId, message }),
         );
         if (!isPausing) {
@@ -1099,7 +1003,7 @@ export class A2AThreadRuntimeCore {
             this.recordedHistoryIds.add(messageId);
           },
           (error) => {
-            const pending = this.historyWrites.get(messageId);
+            const pending = this.historyWrites.pending(messageId);
             if (pending === undefined || pending === write) {
               this.assistantHistoryParents.set(messageId, resolvedParentId);
             }
@@ -1119,7 +1023,7 @@ export class A2AThreadRuntimeCore {
       this.assistantHistoryParents.delete(messageId);
     }
     void write.catch((error) => {
-      const pending = this.historyWrites.get(messageId);
+      const pending = this.historyWrites.pending(messageId);
       if (pending === undefined || pending === write) {
         this.assistantHistoryParents.set(messageId, resolvedParentId);
       }
@@ -1134,7 +1038,7 @@ export class A2AThreadRuntimeCore {
     if (!this.history || this.recordedHistoryIds.has(message.id)) return;
     this.recordedHistoryIds.add(message.id);
     const append = this.history.append.bind(this.history);
-    const write = this.chainHistoryWrite(message.id, () =>
+    const write = this.historyWrites.chain(message.id, () =>
       append({ parentId, message }),
     );
     void write.catch(() => {
@@ -1142,35 +1046,4 @@ export class A2AThreadRuntimeCore {
     });
     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;
-  }
 }