@assistant-ui/react-a2a 0.2.21 → 0.2.23

package/src/useA2ARuntime.test.tsx

@@ -0,0 +1,142 @@
+// @vitest-environment jsdom
+
+import { act, renderHook, waitFor } from "@testing-library/react";
+import { afterEach, describe, expect, it, vi } from "vitest";
+import type { A2AClient } from "./A2AClient";
+import { useA2ARuntime } from "./useA2ARuntime";
+
+const createMockClient = (waitForAbort = false) => {
+  let streamSignal: AbortSignal | undefined;
+  const getAgentCard = vi.fn().mockResolvedValue(undefined);
+  const streamMessage = vi.fn(
+    (
+      _message: unknown,
+      _configuration: unknown,
+      _metadata: unknown,
+      signal?: AbortSignal,
+    ): AsyncIterable<never> => {
+      streamSignal = signal;
+      return {
+        async *[Symbol.asyncIterator]() {
+          if (!waitForAbort || !signal) return;
+          await new Promise<void>((resolve) => {
+            if (signal.aborted) {
+              resolve();
+            } else {
+              signal.addEventListener("abort", () => resolve(), {
+                once: true,
+              });
+            }
+          });
+        },
+      };
+    },
+  );
+
+  return {
+    client: {
+      getAgentCard,
+      streamMessage,
+    } as unknown as A2AClient,
+    getAgentCard,
+    streamMessage,
+    get streamSignal() {
+      return streamSignal;
+    },
+  };
+};
+
+afterEach(() => {
+  vi.unstubAllGlobals();
+});
+
+describe("useA2ARuntime", () => {
+  it("switches provided clients and aborts the previous client run", async () => {
+    const first = createMockClient(true);
+    const second = createMockClient();
+    const { result, rerender } = renderHook(
+      ({ client }) => useA2ARuntime({ client }),
+      { initialProps: { client: first.client } },
+    );
+
+    await waitFor(() => expect(first.getAgentCard).toHaveBeenCalledOnce());
+
+    act(() => {
+      result.current.thread.append({
+        role: "user",
+        content: [{ type: "text", text: "first" }],
+      });
+    });
+    await waitFor(() => expect(first.streamMessage).toHaveBeenCalledOnce());
+
+    rerender({ client: second.client });
+
+    await waitFor(() => expect(first.streamSignal?.aborted).toBe(true));
+    await waitFor(() => expect(second.getAgentCard).toHaveBeenCalledOnce());
+
+    act(() => {
+      result.current.thread.append({
+        role: "user",
+        content: [{ type: "text", text: "second" }],
+      });
+    });
+
+    await waitFor(() => expect(second.streamMessage).toHaveBeenCalledOnce());
+    expect(first.streamMessage).toHaveBeenCalledOnce();
+  });
+
+  it("uses current headers without recreating the managed client", async () => {
+    const fetchMock = vi.fn(
+      async (input: string | URL | Request, _init?: RequestInit) => {
+        const url =
+          typeof input === "string"
+            ? input
+            : input instanceof URL
+              ? input.href
+              : input.url;
+        if (url.endsWith("/.well-known/agent-card.json")) {
+          return new Response(
+            JSON.stringify({ capabilities: { streaming: true } }),
+            {
+              status: 200,
+              headers: { "Content-Type": "application/json" },
+            },
+          );
+        }
+        return new Response("", {
+          status: 200,
+          headers: { "Content-Type": "text/event-stream" },
+        });
+      },
+    );
+    vi.stubGlobal("fetch", fetchMock);
+
+    const { result, rerender } = renderHook(
+      ({ token }) =>
+        useA2ARuntime({
+          baseUrl: "https://agent.test",
+          headers: { Authorization: `Bearer ${token}` },
+          extensions: ["urn:example"],
+          fetchOptions: { credentials: "include" },
+        }),
+      { initialProps: { token: "first" } },
+    );
+
+    await waitFor(() => expect(fetchMock).toHaveBeenCalledOnce());
+
+    rerender({ token: "second" });
+
+    act(() => {
+      result.current.thread.append({
+        role: "user",
+        content: [{ type: "text", text: "hello" }],
+      });
+    });
+
+    await waitFor(() => expect(fetchMock).toHaveBeenCalledTimes(2));
+    const streamRequest = fetchMock.mock.calls[1]!;
+    expect(streamRequest[1]?.headers).toMatchObject({
+      Authorization: "Bearer second",
+    });
+  });
+});

package/src/useA2ARuntime.ts

@@ -12,11 +12,28 @@ import type {
   ExternalStoreAdapter,
   ThreadMessage,
 } from "@assistant-ui/core";
-import { A2AClient } from "./A2AClient";
+import { A2AClient, type A2AClientOptions } from "./A2AClient";
 import { A2AThreadRuntimeCore } from "./A2AThreadRuntimeCore";
 import { a2aExtras } from "./a2aExtras";
 import type { UseA2ARuntimeOptions } from "./types";
 
+type ManagedA2AClientOptions = Omit<A2AClientOptions, "headers">;
+
+const serializeManagedClientOptions = (
+  options: ManagedA2AClientOptions,
+): string => {
+  const fetchOptions = Object.fromEntries(
+    Object.entries(options.fetchOptions ?? {}).sort(([a], [b]) =>
+      a.localeCompare(b),
+    ),
+  );
+
+  return JSON.stringify({
+    ...options,
+    fetchOptions,
+  });
+};
+
 export function useA2ARuntime(options: UseA2ARuntimeOptions): AssistantRuntime {
   const [_version, setVersion] = useState(0);
   const notifyUpdate = useCallback(() => setVersion((v) => v + 1), []);
@@ -24,44 +41,46 @@ export function useA2ARuntime(options: UseA2ARuntimeOptions): AssistantRuntime {
   const historyAdapter = options.adapters?.history ?? runtimeAdapters?.history;
   const threadListAdapter = options.adapters?.threadList;
 
-  // Create or reuse client
-  const clientRef = useRef<A2AClient | null>(null);
-  if (!clientRef.current) {
-    if (options.client) {
-      clientRef.current = options.client;
-    } else if (options.baseUrl) {
-      clientRef.current = new A2AClient({
-        baseUrl: options.baseUrl,
-        basePath: options.basePath,
-        tenant: options.tenant,
-        headers: options.headers,
-        extensions: options.extensions,
-        fetchOptions: options.fetchOptions,
-      });
-    } else {
+  const headersRef = useRef(options.headers);
+  headersRef.current = options.headers;
+  const resolveHeaders = useCallback(() => {
+    const headers = headersRef.current;
+    return typeof headers === "function" ? headers() : (headers ?? {});
+  }, []);
+
+  const managedClientOptionsKey = options.client
+    ? null
+    : options.baseUrl
+      ? serializeManagedClientOptions({
+          baseUrl: options.baseUrl,
+          basePath: options.basePath,
+          tenant: options.tenant,
+          extensions: options.extensions,
+          fetchOptions: options.fetchOptions,
+        })
+      : null;
+
+  const client = useMemo(() => {
+    if (options.client) return options.client;
+    if (!managedClientOptionsKey) {
       throw new Error("useA2ARuntime requires either `client` or `baseUrl`");
     }
-  }
-  const client = clientRef.current;
-
-  // Create or reuse core
-  const coreRef = useRef<A2AThreadRuntimeCore | null>(null);
-  if (!coreRef.current) {
-    coreRef.current = new A2AThreadRuntimeCore({
-      client,
-      contextId: options.contextId,
-      configuration: options.configuration,
-      ...(options.onError && { onError: options.onError }),
-      ...(options.onCancel && { onCancel: options.onCancel }),
-      ...(options.onArtifactComplete && {
-        onArtifactComplete: options.onArtifactComplete,
-      }),
-      ...(historyAdapter && { history: historyAdapter }),
-      notifyUpdate,
+
+    return new A2AClient({
+      ...(JSON.parse(managedClientOptionsKey) as ManagedA2AClientOptions),
+      headers: resolveHeaders,
     });
-  }
+  }, [managedClientOptionsKey, options.client, resolveHeaders]);
+
+  const core = useMemo(
+    () =>
+      new A2AThreadRuntimeCore({
+        client,
+        notifyUpdate,
+      }),
+    [client, notifyUpdate],
+  );
 
-  const core = coreRef.current;
   core.updateOptions({
     client,
     contextId: options.contextId,
@@ -119,7 +138,7 @@ export function useA2ARuntime(options: UseA2ARuntimeOptions): AssistantRuntime {
     return {
       ...shared,
       isLoading: core.isLoading,
-      messages: core.getMessages(),
+      messageRepository: core.getMessageRepository(),
       isRunning: core.isRunning(),
       extras: a2aExtras.provide({
         task: core.getTask(),