@assistant-ui/react-a2a 0.2.31 → 0.2.32

package/src/useA2ARuntime.test.tsx

@@ -1,6 +1,7 @@
 // @vitest-environment jsdom
 
 import { act, renderHook, waitFor } from "@testing-library/react";
+import { startTransition, Suspense, type PropsWithChildren } from "react";
 import { afterEach, describe, expect, it, vi } from "vitest";
 import type { A2AClient } from "./A2AClient";
 import type { A2AStreamEvent } from "./types";
@@ -57,11 +58,100 @@ const createMockClient = (waitForAbort = false) => {
   };
 };
 
+const createFetchMock = () =>
+  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(
+      'data: {"message":{"message_id":"response","role":"ROLE_AGENT","parts":[{"text":"Done"}]}}\n\n',
+      {
+        status: 200,
+        headers: { "Content-Type": "text/event-stream" },
+      },
+    );
+  });
+
 afterEach(() => {
   vi.unstubAllGlobals();
 });
 
 describe("useA2ARuntime", () => {
+  const createHistory = () => ({
+    load: vi.fn().mockResolvedValue({
+      headId: "restored",
+      messages: [
+        {
+          parentId: null,
+          message: {
+            id: "restored",
+            role: "user" as const,
+            content: [{ type: "text" as const, text: "hello" }],
+            createdAt: new Date(0),
+            metadata: { custom: {} },
+          },
+        },
+      ],
+    }),
+    append: vi.fn().mockResolvedValue(undefined),
+  });
+
+  it("loads a history adapter that arrives on a later render", async () => {
+    const { client } = createMockClient();
+    const history = createHistory();
+    const { result, rerender } = renderHook(
+      ({ history }: { history?: ReturnType<typeof createHistory> }) =>
+        useA2ARuntime({ client, adapters: history ? { history } : {} }),
+      { initialProps: {} },
+    );
+
+    await waitFor(() =>
+      expect(result.current.thread.getState().isLoading).toBe(false),
+    );
+    expect(history.load).not.toHaveBeenCalled();
+
+    rerender({ history });
+
+    await waitFor(() =>
+      expect(
+        result.current.thread.getState().messages.map((m) => m.id),
+      ).toEqual(["restored"]),
+    );
+    expect(history.load).toHaveBeenCalledOnce();
+  });
+
+  it("loads history through a swapped client's core", async () => {
+    const first = createMockClient();
+    const second = createMockClient();
+    const history = createHistory();
+    const { result, rerender } = renderHook(
+      ({ client }) => useA2ARuntime({ client, adapters: { history } }),
+      { initialProps: { client: first.client } },
+    );
+
+    await waitFor(() => expect(history.load).toHaveBeenCalledOnce());
+
+    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",
+    ]);
+  });
+
   it("switches provided clients and aborts the previous client run", async () => {
     const first = createMockClient(true);
     const second = createMockClient();
@@ -97,32 +187,7 @@ describe("useA2ARuntime", () => {
   });
 
   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(
-          'data: {"message":{"message_id":"response","role":"ROLE_AGENT","parts":[{"text":"Done"}]}}\n\n',
-          {
-            status: 200,
-            headers: { "Content-Type": "text/event-stream" },
-          },
-        );
-      },
-    );
+    const fetchMock = createFetchMock();
     vi.stubGlobal("fetch", fetchMock);
 
     const { result, rerender } = renderHook(
@@ -153,4 +218,63 @@ describe("useA2ARuntime", () => {
       Authorization: "Bearer second",
     });
   });
+
+  it("keeps managed headers scoped to committed renders", async () => {
+    const fetchMock = createFetchMock();
+    vi.stubGlobal("fetch", fetchMock);
+
+    const interruptedRender = vi.fn();
+    const renderedToken = vi.fn();
+    const pending = new Promise<never>(() => {});
+    let blocked = false;
+    const Blocker = ({ blocked }: { blocked: boolean }) => {
+      if (blocked) {
+        interruptedRender();
+        throw pending;
+      }
+      return null;
+    };
+    const Wrapper = ({ children }: PropsWithChildren) => (
+      <Suspense fallback={null}>
+        {children}
+        <Blocker blocked={blocked} />
+      </Suspense>
+    );
+
+    const { result, rerender } = renderHook(
+      ({ token }) => {
+        renderedToken(token);
+        return useA2ARuntime({
+          baseUrl: "https://agent.test",
+          headers: { Authorization: `Bearer ${token}` },
+        });
+      },
+      {
+        initialProps: { token: "workspace-a" },
+        wrapper: Wrapper,
+      },
+    );
+    await waitFor(() => expect(fetchMock).toHaveBeenCalledOnce());
+
+    act(() => {
+      blocked = true;
+      startTransition(() => {
+        rerender({ token: "workspace-b" });
+      });
+    });
+    expect(interruptedRender).toHaveBeenCalled();
+    expect(renderedToken).toHaveBeenCalledWith("workspace-b");
+
+    act(() => {
+      result.current.thread.append({
+        role: "user",
+        content: [{ type: "text", text: "hello" }],
+      });
+    });
+
+    await waitFor(() => expect(fetchMock).toHaveBeenCalledTimes(2));
+    expect(fetchMock.mock.calls[1]![1]?.headers).toMatchObject({
+      Authorization: "Bearer workspace-a",
+    });
+  });
 });

package/src/useA2ARuntime.ts

@@ -1,6 +1,13 @@
 "use client";
 
-import { useCallback, useEffect, useMemo, useRef, useState } from "react";
+import {
+  useCallback,
+  useEffect,
+  useInsertionEffect,
+  useMemo,
+  useRef,
+  useState,
+} from "react";
 import {
   useExternalStoreRuntime,
   useExternalStoreSharedOptions,
@@ -42,7 +49,9 @@ export function useA2ARuntime(options: UseA2ARuntimeOptions): AssistantRuntime {
   const threadListAdapter = options.adapters?.threadList;
 
   const headersRef = useRef(options.headers);
-  headersRef.current = options.headers;
+  useInsertionEffect(() => {
+    headersRef.current = options.headers;
+  });
   const resolveHeaders = useCallback(() => {
     const headers = headersRef.current;
     return typeof headers === "function" ? headers() : (headers ?? {});
@@ -72,16 +81,7 @@ export function useA2ARuntime(options: UseA2ARuntimeOptions): AssistantRuntime {
     });
   }, [managedClientOptionsKey, options.client, resolveHeaders]);
 
-  const core = useMemo(
-    () =>
-      new A2AThreadRuntimeCore({
-        client,
-        notifyUpdate,
-      }),
-    [client, notifyUpdate],
-  );
-
-  core.updateOptions({
+  const coreOptions = {
     client,
     contextId: options.contextId,
     configuration: options.configuration,
@@ -91,6 +91,22 @@ export function useA2ARuntime(options: UseA2ARuntimeOptions): AssistantRuntime {
       onArtifactComplete: options.onArtifactComplete,
     }),
     ...(historyAdapter && { history: historyAdapter }),
+  };
+  const coreOptionsRef = useRef(coreOptions);
+  coreOptionsRef.current = coreOptions;
+
+  const core = useMemo(
+    () =>
+      new A2AThreadRuntimeCore({
+        ...coreOptionsRef.current,
+        client,
+        notifyUpdate,
+      }),
+    [client, notifyUpdate],
+  );
+
+  useEffect(() => {
+    core.updateOptions(coreOptions);
   });
 
   // Thread list
@@ -104,13 +120,18 @@ export function useA2ARuntime(options: UseA2ARuntimeOptions): AssistantRuntime {
       onSwitchToNewThread: onSwitchToNewThread
         ? async () => {
             await onSwitchToNewThread();
+            // Apply first so the abort inside resetContext finds an already
+            // cleared repository and cannot persist the old thread's partial
+            // assistant message.
             core.applyExternalMessages([]);
+            core.resetContext();
           }
         : undefined,
       onSwitchToThread: onSwitchToThread
         ? async (threadId: string) => {
             const result = await onSwitchToThread(threadId);
             core.applyExternalMessages(result.messages);
+            core.resetContext();
           }
         : undefined,
     };