@assistant-ui/react-a2a 0.2.34 → 0.2.36

package/src/A2AThreadRuntimeCore.test.ts

@@ -1,8 +1,17 @@
 import { describe, it, expect, vi, beforeEach, afterEach } from "vitest";
 import { A2AThreadRuntimeCore } from "./A2AThreadRuntimeCore";
 import type { A2AClient } from "./A2AClient";
-import type { A2AMessage, A2AStreamEvent, A2ATask } from "./types";
-import type { AppendMessage, ThreadMessage } from "@assistant-ui/core";
+import type {
+  A2AAgentCard,
+  A2AMessage,
+  A2AStreamEvent,
+  A2ATask,
+} from "./types";
+import type {
+  AppendMessage,
+  ExportedMessageRepository,
+  ThreadMessage,
+} from "@assistant-ui/core";
 
 // --- Mock client factory ---
 
@@ -264,6 +273,118 @@ describe("A2AThreadRuntimeCore", () => {
   });
 
   describe("history loading", () => {
+    it("keeps initial history loading across same-thread message resyncs", async () => {
+      let resolve!: (repo: ExportedMessageRepository) => void;
+      const pending = new Promise<ExportedMessageRepository>((res) => {
+        resolve = res;
+      });
+      const core = createCore(
+        {},
+        { history: { load: () => pending, append: async () => {} } },
+      );
+      const loading = core.__internal_load();
+      core.applyExternalMessages([]);
+      const wasLoading = core.isLoading;
+      const restored = createHistoryMessage(
+        "restored",
+        "user",
+        "Saved history",
+      );
+      resolve({
+        headId: restored.id,
+        messages: [{ parentId: null, message: restored }],
+      });
+      await loading;
+      expect(wasLoading).toBe(true);
+      expect(core.getMessages()).toEqual([restored]);
+      expect(core.isLoading).toBe(false);
+    });
+
+    it.each(["resolve", "reject"])(
+      "ignores an initial history load that later %s after a thread switch",
+      async (outcome) => {
+        let resolve!: (repo: ExportedMessageRepository) => void;
+        let reject!: (error: Error) => void;
+        const pending = new Promise<ExportedMessageRepository>((res, rej) => {
+          resolve = res;
+          reject = rej;
+        });
+        const onError = vi.fn();
+        const history = {
+          load: vi.fn(() => pending),
+          append: vi.fn().mockResolvedValue(undefined),
+        };
+        const core = createCore({}, { history, onError });
+        const loading = core.__internal_load();
+        const replacement = createHistoryMessage(
+          "selected",
+          "user",
+          "Selected thread",
+        );
+        core.applyExternalMessages([replacement]);
+        core.resetContext();
+
+        if (outcome === "resolve") {
+          const previous = createHistoryMessage(
+            "previous",
+            "user",
+            "Previous thread",
+          );
+          resolve({
+            headId: previous.id,
+            messages: [{ parentId: null, message: previous }],
+          });
+        } else {
+          reject(new Error("Old history failed"));
+        }
+        await loading;
+
+        expect(core.getMessages()).toEqual([replacement]);
+        expect(core.isLoading).toBe(false);
+        expect(onError).not.toHaveBeenCalled();
+        await core.__internal_load();
+        expect(history.load).toHaveBeenCalledTimes(1);
+      },
+    );
+
+    it("ends initial loading immediately when switching threads", async () => {
+      let resolve!: (repo: ExportedMessageRepository) => void;
+      const pending = new Promise<ExportedMessageRepository>((res) => {
+        resolve = res;
+      });
+      const core = createCore(
+        {},
+        { history: { load: () => pending, append: async () => {} } },
+      );
+      const loading = core.__internal_load();
+      expect(core.isLoading).toBe(true);
+      core.applyExternalMessages([]);
+      core.resetContext();
+      const updatesAfterReplacement = notifyUpdate.mock.calls.length;
+      expect(core.isLoading).toBe(false);
+      resolve({ messages: [] });
+      await loading;
+      expect(notifyUpdate).toHaveBeenCalledTimes(updatesAfterReplacement);
+    });
+
+    it("does not import old history when agent-card discovery settles after a thread switch", async () => {
+      let resolveCard!: (card: A2AAgentCard) => void;
+      const card = new Promise<A2AAgentCard>((resolve) => {
+        resolveCard = resolve;
+      });
+      const { history } = createBranchedHistory();
+      const core = createCore({ getAgentCard: () => card }, { history });
+      const loading = core.__internal_load();
+      await Promise.resolve();
+      core.applyExternalMessages([]);
+      core.resetContext();
+      resolveCard({ name: "Agent" } as A2AAgentCard);
+      await loading;
+      expect(core.getMessages()).toEqual([]);
+      expect(core.getAgentCard()?.name).toBe("Agent");
+      expect(core.isLoading).toBe(false);
+    });
+
     it("preserves sibling branches and selects the persisted head", async () => {
       const { user, firstAssistant, secondAssistant, history } =
         createBranchedHistory();
@@ -791,6 +912,147 @@ describe("A2AThreadRuntimeCore", () => {
   // --- Sync (non-streaming) fallback ---
 
   describe("sync fallback", () => {
+    it("retries agent card discovery after a transient failure", async () => {
+      let now = 1_000;
+      vi.spyOn(Date, "now").mockImplementation(() => now);
+      let resolveRecovery!: (card: A2AAgentCard) => void;
+      const transportOrder: string[] = [];
+      const getAgentCard = vi
+        .fn()
+        .mockRejectedValueOnce(new Error("temporary failure"))
+        .mockImplementationOnce(
+          () =>
+            new Promise<A2AAgentCard>((resolve) => {
+              resolveRecovery = resolve;
+            }),
+        );
+      const sendMessage = vi.fn().mockImplementation(async () => {
+        transportOrder.push("sync");
+        return {
+          id: "t1",
+          status: { state: "completed" },
+        } satisfies A2ATask;
+      });
+      const streamMessage = vi.fn().mockImplementation(async function* () {
+        transportOrder.push("stream");
+        yield statusUpdateEvent("completed");
+      });
+      const core = createCore({ getAgentCard, sendMessage, streamMessage });
+
+      await core.append(createUserAppendMessage("First"));
+      now += 5_000;
+      await core.append(createUserAppendMessage("Second"));
+      resolveRecovery({
+        name: "Agent",
+        capabilities: { streaming: false },
+      } as A2AAgentCard);
+      await vi.waitFor(() => expect(core.getAgentCard()).toBeDefined());
+      await core.append(createUserAppendMessage("Third"));
+
+      expect(getAgentCard).toHaveBeenCalledTimes(2);
+      expect(streamMessage).toHaveBeenCalledTimes(2);
+      expect(sendMessage).toHaveBeenCalledOnce();
+      expect(transportOrder).toEqual(["stream", "stream", "sync"]);
+    });
+
+    it("does not repeat failed discovery during the retry delay", async () => {
+      vi.spyOn(Date, "now").mockReturnValue(1_000);
+      const getAgentCard = vi.fn().mockRejectedValue(new Error("unavailable"));
+      const streamMessage = vi.fn().mockImplementation(async function* () {
+        yield statusUpdateEvent("completed");
+      });
+      const core = createCore({ getAgentCard, streamMessage });
+
+      await core.append(createUserAppendMessage("First"));
+      await core.append(createUserAppendMessage("Second"));
+
+      expect(getAgentCard).toHaveBeenCalledOnce();
+      expect(streamMessage).toHaveBeenCalledTimes(2);
+    });
+
+    it("backs off persistent failures without blocking later sends", async () => {
+      let now = 1_000;
+      vi.spyOn(Date, "now").mockImplementation(() => now);
+      let rejectSecond!: (error: Error) => void;
+      const getAgentCard = vi
+        .fn()
+        .mockRejectedValueOnce(new Error("first failure"))
+        .mockImplementationOnce(
+          () =>
+            new Promise<A2AAgentCard>((_resolve, reject) => {
+              rejectSecond = reject;
+            }),
+        )
+        .mockRejectedValue(new Error("still unavailable"));
+      const streamMessage = vi.fn().mockImplementation(async function* () {
+        yield statusUpdateEvent("completed");
+      });
+      const core = createCore({ getAgentCard, streamMessage });
+
+      await core.append(createUserAppendMessage("First"));
+      now = 6_000;
+      await core.append(createUserAppendMessage("Second"));
+      expect(streamMessage).toHaveBeenCalledTimes(2);
+      rejectSecond(new Error("second failure"));
+      await vi.waitFor(() => expect(getAgentCard).toHaveBeenCalledTimes(2));
+
+      now = 15_999;
+      await core.append(createUserAppendMessage("Third"));
+      expect(getAgentCard).toHaveBeenCalledTimes(2);
+
+      now = 16_000;
+      await core.append(createUserAppendMessage("Fourth"));
+      expect(getAgentCard).toHaveBeenCalledTimes(3);
+      expect(streamMessage).toHaveBeenCalledTimes(4);
+    });
+
+    it("waits for agent capabilities before choosing the first send method", async () => {
+      let resolveAgentCard!: (value: A2AAgentCard) => void;
+      const getAgentCard = vi.fn(
+        () =>
+          new Promise<A2AAgentCard>((resolve) => {
+            resolveAgentCard = resolve;
+          }),
+      );
+      const sendMessage = vi.fn().mockResolvedValue({
+        id: "t1",
+        status: { state: "completed" },
+      } satisfies A2ATask);
+      const streamMessage = vi.fn();
+      const core = createCore({ getAgentCard, sendMessage, streamMessage });
+
+      const run = core.append(createUserAppendMessage("Hello"));
+      expect(sendMessage).not.toHaveBeenCalled();
+      expect(streamMessage).not.toHaveBeenCalled();
+
+      resolveAgentCard({
+        name: "Agent",
+        capabilities: { streaming: false },
+      } as A2AAgentCard);
+      await run;
+
+      expect(getAgentCard).toHaveBeenCalledOnce();
+      expect(sendMessage).toHaveBeenCalledOnce();
+      expect(streamMessage).not.toHaveBeenCalled();
+    });
+
+    it("stops waiting for agent capabilities when the run is cancelled", async () => {
+      const getAgentCard = vi.fn(() => new Promise<A2AAgentCard>(() => {}));
+      const sendMessage = vi.fn();
+      const streamMessage = vi.fn();
+      const core = createCore({ getAgentCard, sendMessage, streamMessage });
+
+      const run = core.append(createUserAppendMessage("Hello"));
+      expect(core.isRunning()).toBe(true);
+
+      await core.cancel();
+      await run;
+
+      expect(core.isRunning()).toBe(false);
+      expect(sendMessage).not.toHaveBeenCalled();
+      expect(streamMessage).not.toHaveBeenCalled();
+    });
+
     it("uses sendMessage when streaming is false in agent card", async () => {
       const sendMessage = vi.fn().mockResolvedValue({
         id: "t1",
@@ -1256,6 +1518,136 @@ describe("A2AThreadRuntimeCore", () => {
       expect(cancelTask).not.toHaveBeenCalled();
     });
 
+    it("ignores a cancellation response that resolves after a newer task snapshot", async () => {
+      let resolveCancel!: (task: A2ATask) => void;
+      const cancelTask = vi.fn().mockReturnValue(
+        new Promise<A2ATask>((resolve) => {
+          resolveCancel = resolve;
+        }),
+      );
+      const streamMessage = vi.fn().mockImplementation(async function* (
+        _msg: any,
+        _cfg: any,
+        _meta: any,
+        signal: AbortSignal,
+      ) {
+        if (streamMessage.mock.calls.length === 1) {
+          yield statusUpdateEvent("working");
+          await new Promise<void>((resolve) => {
+            if (signal.aborted) resolve();
+            else
+              signal.addEventListener("abort", () => resolve(), {
+                once: true,
+              });
+          });
+          return;
+        }
+        yield statusUpdateEvent("working", "Second answer");
+      });
+      const core = createCore({ cancelTask, streamMessage });
+
+      const firstRun = core.append(createUserAppendMessage("First"));
+      await vi.waitFor(() => expect(core.getTask()?.id).toBe("t1"));
+
+      const cancelPromise = core.cancel();
+      await firstRun;
+
+      await core.append(createUserAppendMessage("Second"));
+      expect(core.getTask()?.status.state).toBe("working");
+
+      resolveCancel({ id: "t1", status: { state: "canceled" } });
+      await cancelPromise;
+
+      expect(cancelTask).toHaveBeenCalledExactlyOnceWith("t1");
+      expect(core.getTask()).toMatchObject({
+        id: "t1",
+        status: { state: "working" },
+      });
+    });
+
+    it("ignores a cancellation response that resolves after a newer run starts", async () => {
+      let resolveCancel!: (task: A2ATask) => void;
+      const cancelTask = vi.fn().mockReturnValue(
+        new Promise<A2ATask>((resolve) => {
+          resolveCancel = resolve;
+        }),
+      );
+      // The follow-up run never emits, so it leaves the non-terminal task
+      // object from the first run in place: only the run itself is newer.
+      const streamMessage = vi.fn().mockImplementation(async function* (
+        _msg: any,
+        _cfg: any,
+        _meta: any,
+        signal: AbortSignal,
+      ) {
+        if (streamMessage.mock.calls.length === 1) {
+          yield statusUpdateEvent("working");
+        }
+        await new Promise<void>((resolve) => {
+          if (signal.aborted) resolve();
+          else
+            signal.addEventListener("abort", () => resolve(), { once: true });
+        });
+      });
+      const core = createCore({ cancelTask, streamMessage });
+
+      const firstRun = core.append(createUserAppendMessage("First"));
+      await vi.waitFor(() => expect(core.getTask()?.id).toBe("t1"));
+
+      const cancelPromise = core.cancel();
+      await firstRun;
+
+      void core.append(createUserAppendMessage("Second"));
+      await vi.waitFor(() => expect(streamMessage).toHaveBeenCalledTimes(2));
+      expect(core.isRunning()).toBe(true);
+
+      resolveCancel({ id: "t1", status: { state: "canceled" } });
+      await cancelPromise;
+
+      expect(core.getTask()?.status.state).toBe("working");
+    });
+
+    it("still cancels the server task when onCancel clears the thread", async () => {
+      let resolveCancel!: (task: A2ATask) => void;
+      const cancelTask = vi.fn().mockReturnValue(
+        new Promise<A2ATask>((resolve) => {
+          resolveCancel = resolve;
+        }),
+      );
+      const streamMessage = vi.fn().mockImplementation(async function* (
+        _msg: any,
+        _cfg: any,
+        _meta: any,
+        signal: AbortSignal,
+      ) {
+        yield statusUpdateEvent("working");
+        await new Promise<void>((resolve) => {
+          if (signal.aborted) resolve();
+          else
+            signal.addEventListener("abort", () => resolve(), { once: true });
+        });
+      });
+      let core!: A2AThreadRuntimeCore;
+      core = createCore(
+        { cancelTask, streamMessage },
+        { onCancel: () => core.applyExternalMessages([]) },
+      );
+
+      const runPromise = core.append(createUserAppendMessage("Go"));
+      await vi.waitFor(() => expect(core.getTask()?.id).toBe("t1"));
+
+      const cancelPromise = core.cancel();
+      await runPromise;
+
+      expect(cancelTask).toHaveBeenCalledExactlyOnceWith("t1");
+
+      resolveCancel({ id: "t1", status: { state: "canceled" } });
+      await cancelPromise;
+
+      expect(core.getTask()).toBeUndefined();
+      expect(core.getMessages()).toEqual([]);
+    });
+
     it.each(["throws", "rejects"] as const)(
       "isolates onCancel callbacks that %s",
       async (failureMode) => {

package/src/A2AThreadRuntimeCore.ts

@@ -25,6 +25,7 @@ import type {
   A2ATaskArtifactUpdateEvent,
   A2ATaskStatusUpdateEvent,
 } from "./types";
+
 import {
   a2aMessageToContent,
   isTerminalTaskState,
@@ -32,6 +33,9 @@ import {
   taskStateToMessageStatus,
 } from "./conversions";
 
+const INITIAL_AGENT_CARD_RETRY_DELAY_MS = 5_000;
+const MAX_AGENT_CARD_RETRY_DELAY_MS = 5 * 60_000;
+
 export type A2AThreadRuntimeCoreOptions = {
   client: A2AClient;
   contextId?: string | undefined;
@@ -79,6 +83,7 @@ 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
@@ -91,8 +96,12 @@ export class A2AThreadRuntimeCore {
   private readonly recordedHistoryIds = new Set<string>();
   private _isLoading = false;
   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;
+  private _agentCardDiscoveryFailed = false;
 
   private lastOptionsContextId: string | undefined;
 
@@ -139,6 +148,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.
@@ -197,31 +208,72 @@ export class A2AThreadRuntimeCore {
     return this._isLoading;
   }
 
-  __internal_load(): Promise<void> {
-    this._loadRequested = true;
-    this._agentCardPromise ??= this.client
-      .getAgentCard()
-      .then((agentCard) => {
+  private loadAgentCard(): Promise<void> {
+    if (Date.now() < this._agentCardRetryAfter) return Promise.resolve();
+
+    this._agentCardPromise ??= this.client.getAgentCard().then(
+      (agentCard) => {
         this.agentCardValue = agentCard;
+        this._agentCardRetryAfter = 0;
+        this._agentCardRetryDelay = INITIAL_AGENT_CARD_RETRY_DELAY_MS;
+        this._agentCardDiscoveryFailed = false;
         this.notifyUpdate();
-      })
-      .catch(() => undefined);
+      },
+      () => {
+        this._agentCardDiscoveryFailed = true;
+        this._agentCardRetryAfter = Date.now() + this._agentCardRetryDelay;
+        this._agentCardRetryDelay = Math.min(
+          this._agentCardRetryDelay * 2,
+          MAX_AGENT_CARD_RETRY_DELAY_MS,
+        );
+        this._agentCardPromise = undefined;
+      },
+    );
+    return this._agentCardPromise;
+  }
+
+  private async waitForAgentCard(signal: AbortSignal): Promise<boolean> {
+    const shouldWait = !this._agentCardDiscoveryFailed;
+    const load = this.loadAgentCard();
+    if (signal.aborted) return false;
+    if (!shouldWait) return true;
+
+    let onAbort!: () => void;
+    const abort = new Promise<void>((resolve) => {
+      onAbort = resolve;
+      signal.addEventListener("abort", onAbort, { once: true });
+    });
+
+    try {
+      await Promise.race([load, abort]);
+    } finally {
+      signal.removeEventListener("abort", onAbort);
+    }
+    return !signal.aborted;
+  }
+
+  __internal_load(): Promise<void> {
+    this._loadRequested = true;
+    const agentCardPromise = this.loadAgentCard();
 
     if (this._loadPromise) return this._loadPromise;
-    if (!this.history) return this._agentCardPromise;
+    if (!this.history) return agentCardPromise;
 
     this._isLoading = true;
 
+    const generation = this._historyLoadGeneration;
     const historyPromise = this.history.load();
 
-    this._loadPromise = Promise.all([historyPromise, this._agentCardPromise])
+    this._loadPromise = Promise.all([historyPromise, agentCardPromise])
       .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,
@@ -229,6 +281,7 @@ export class A2AThreadRuntimeCore {
         );
       })
       .finally(() => {
+        if (generation !== this._historyLoadGeneration) return;
         this._isLoading = false;
         this.notifyUpdate();
       });
@@ -283,14 +336,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
       }
@@ -360,6 +425,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();
@@ -407,11 +474,11 @@ export class A2AThreadRuntimeCore {
 
     this.setRunning(true);
 
-    // Check if agent supports streaming; fall back to sync sendMessage if not
-    const supportsStreaming =
-      this.agentCardValue?.capabilities?.streaming !== false;
-
     try {
+      if (!(await this.waitForAgentCard(abortController.signal))) return;
+
+      const supportsStreaming =
+        this.agentCardValue?.capabilities?.streaming !== false;
       if (supportsStreaming) {
         await this.runStreaming(a2aMessage, assistantId, abortController);
       } else {