@assistant-ui/react-a2a 0.2.33 → 0.2.35

package/src/A2AThreadRuntimeCore.test.ts

@@ -1,7 +1,12 @@
 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 {
+  A2AAgentCard,
+  A2AMessage,
+  A2AStreamEvent,
+  A2ATask,
+} from "./types";
 import type { AppendMessage, ThreadMessage } from "@assistant-ui/core";
 
 // --- Mock client factory ---
@@ -791,6 +796,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",

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;
@@ -93,6 +97,9 @@ export class A2AThreadRuntimeCore {
   private _loadPromise: Promise<void> | undefined;
   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;
 
@@ -197,24 +204,62 @@ 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 historyPromise = this.history.load();
 
-    this._loadPromise = Promise.all([historyPromise, this._agentCardPromise])
+    this._loadPromise = Promise.all([historyPromise, agentCardPromise])
       .then(([repo]) => {
         if (repo) {
           this.session.applyExternalMessageRepository(repo);
@@ -407,11 +452,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 {

package/src/useA2ARuntime.test.tsx

@@ -1,8 +1,14 @@
 // @vitest-environment jsdom
 
 import { act, renderHook, waitFor } from "@testing-library/react";
-import { startTransition, Suspense, type PropsWithChildren } from "react";
+import {
+  startTransition,
+  Suspense,
+  useState,
+  type PropsWithChildren,
+} from "react";
 import { afterEach, describe, expect, it, vi } from "vitest";
+import type { ThreadMessage } from "@assistant-ui/core";
 import type { A2AClient } from "./A2AClient";
 import type { A2AStreamEvent } from "./types";
 import { useA2ARuntime } from "./useA2ARuntime";
@@ -68,7 +74,21 @@ const createFetchMock = () =>
           : input.url;
     if (url.endsWith("/.well-known/agent-card.json")) {
       return new Response(
-        JSON.stringify({ capabilities: { streaming: true } }),
+        JSON.stringify({
+          name: "Test Agent",
+          version: "1.0",
+          supported_interfaces: [
+            {
+              url: "https://agent.test",
+              protocol_binding: "HTTP+JSON",
+              protocol_version: "1.0",
+            },
+          ],
+          capabilities: { streaming: true },
+          default_input_modes: ["text"],
+          default_output_modes: ["text"],
+          skills: [],
+        }),
         {
           status: 200,
           headers: { "Content-Type": "application/json" },
@@ -84,6 +104,15 @@ const createFetchMock = () =>
     );
   });
 
+const createThreadMessage = (id: string): ThreadMessage => ({
+  id,
+  role: "user",
+  content: [{ type: "text", text: id }],
+  attachments: [],
+  createdAt: new Date(0),
+  metadata: { custom: {} },
+});
+
 afterEach(() => {
   vi.unstubAllGlobals();
 });
@@ -146,10 +175,12 @@ describe("useA2ARuntime", () => {
     rerender({ client: second.client });
 
     await waitFor(() => expect(second.getAgentCard).toHaveBeenCalledOnce());
-    await waitFor(() => expect(history.load).toHaveBeenCalledTimes(2));
-    expect(result.current.thread.getState().messages.map((m) => m.id)).toEqual([
-      "restored",
-    ]);
+    await waitFor(() =>
+      expect(
+        result.current.thread.getState().messages.map((m) => m.id),
+      ).toEqual(["restored"]),
+    );
+    expect(history.load).toHaveBeenCalledTimes(2);
   });
 
   it("switches provided clients and aborts the previous client run", async () => {
@@ -277,4 +308,92 @@ describe("useA2ARuntime", () => {
       Authorization: "Bearer workspace-a",
     });
   });
+
+  it("ignores an older thread load after a newer selection", async () => {
+    const { client } = createMockClient();
+    let resolveFirst!: (value: { messages: ThreadMessage[] }) => void;
+    let resolveSecond!: (value: { messages: ThreadMessage[] }) => void;
+    const first = new Promise<{ messages: ThreadMessage[] }>((resolve) => {
+      resolveFirst = resolve;
+    });
+    const second = new Promise<{ messages: ThreadMessage[] }>((resolve) => {
+      resolveSecond = resolve;
+    });
+    const { result } = renderHook(() => {
+      const [threadId, setThreadId] = useState("initial");
+      return useA2ARuntime({
+        client,
+        adapters: {
+          threadList: {
+            threadId,
+            onSwitchToThread: async (nextThreadId) => {
+              setThreadId(nextThreadId);
+              return nextThreadId === "thread-a" ? first : second;
+            },
+          },
+        },
+      });
+    });
+
+    let switchA!: Promise<void>;
+    let switchB!: Promise<void>;
+    act(() => {
+      switchA = result.current.threads.switchToThread("thread-a");
+      switchB = result.current.threads.switchToThread("thread-b");
+    });
+    expect(result.current.threads.getState().mainThreadId).toBe("thread-b");
+
+    await act(async () => {
+      resolveSecond({ messages: [createThreadMessage("thread-b")] });
+      await switchB;
+    });
+    await act(async () => {
+      resolveFirst({ messages: [createThreadMessage("thread-a")] });
+      await switchA;
+    });
+
+    expect(result.current.thread.getState().messages.map((m) => m.id)).toEqual([
+      "thread-b",
+    ]);
+  });
+
+  it("ignores a thread load superseded by a new thread", async () => {
+    const { client } = createMockClient();
+    let resolveLoad!: (value: { messages: ThreadMessage[] }) => void;
+    const load = new Promise<{ messages: ThreadMessage[] }>((resolve) => {
+      resolveLoad = resolve;
+    });
+    const { result } = renderHook(() => {
+      const [threadId, setThreadId] = useState("initial");
+      return useA2ARuntime({
+        client,
+        adapters: {
+          threadList: {
+            threadId,
+            onSwitchToThread: async (nextThreadId) => {
+              setThreadId(nextThreadId);
+              return load;
+            },
+            onSwitchToNewThread: async () => {
+              setThreadId("thread-new");
+            },
+          },
+        },
+      });
+    });
+
+    let staleSwitch!: Promise<void>;
+    act(() => {
+      staleSwitch = result.current.threads.switchToThread("thread-a");
+    });
+    await act(async () => {
+      await result.current.threads.switchToNewThread();
+    });
+    await act(async () => {
+      resolveLoad({ messages: [createThreadMessage("thread-a")] });
+      await staleSwitch;
+    });
+
+    expect(result.current.thread.getState().messages).toEqual([]);
+  });
 });

package/src/useA2ARuntime.ts

@@ -110,6 +110,7 @@ export function useA2ARuntime(options: UseA2ARuntimeOptions): AssistantRuntime {
   });
 
   // Thread list
+  const threadSwitchGenerationRef = useRef(0);
   const threadList = useMemo(() => {
     if (!threadListAdapter) return undefined;
 
@@ -119,7 +120,9 @@ export function useA2ARuntime(options: UseA2ARuntimeOptions): AssistantRuntime {
       threadId: threadListAdapter.threadId,
       onSwitchToNewThread: onSwitchToNewThread
         ? async () => {
+            const generation = ++threadSwitchGenerationRef.current;
             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.
@@ -129,7 +132,9 @@ export function useA2ARuntime(options: UseA2ARuntimeOptions): AssistantRuntime {
         : undefined,
       onSwitchToThread: onSwitchToThread
         ? async (threadId: string) => {
+            const generation = ++threadSwitchGenerationRef.current;
             const result = await onSwitchToThread(threadId);
+            if (generation !== threadSwitchGenerationRef.current) return;
             core.applyExternalMessages(result.messages);
             core.resetContext();
           }