@assistant-ui/react-a2a 0.2.34 → 0.2.36

package/src/useA2ARuntime.test.tsx

@@ -8,7 +8,10 @@ import {
   type PropsWithChildren,
 } from "react";
 import { afterEach, describe, expect, it, vi } from "vitest";
-import type { ThreadMessage } from "@assistant-ui/core";
+import type {
+  ExportedMessageRepository,
+  ThreadMessage,
+} from "@assistant-ui/core";
 import type { A2AClient } from "./A2AClient";
 import type { A2AStreamEvent } from "./types";
 import { useA2ARuntime } from "./useA2ARuntime";
@@ -74,7 +77,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" },
@@ -161,10 +178,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 () => {
@@ -341,6 +360,93 @@ describe("useA2ARuntime", () => {
     ]);
   });
 
+  it("keeps pending history when a run is cancelled in the same thread", async () => {
+    const { client, streamMessage } = createMockClient(true);
+    let resolve!: (repo: ExportedMessageRepository) => void;
+    const pending = new Promise<ExportedMessageRepository>((res) => {
+      resolve = res;
+    });
+    const { result } = renderHook(() =>
+      useA2ARuntime({
+        client,
+        adapters: { history: { load: () => pending, append: async () => {} } },
+      }),
+    );
+    act(() => {
+      result.current.thread.append("Hello");
+    });
+    await waitFor(() => expect(streamMessage).toHaveBeenCalledOnce());
+    await act(async () => {
+      result.current.thread.cancelRun();
+      await new Promise((done) => setTimeout(done, 0));
+    });
+    const wasLoading = result.current.thread.getState().isLoading;
+    const restored = createThreadMessage("restored");
+    await act(async () => {
+      resolve({
+        headId: restored.id,
+        messages: [{ parentId: null, message: restored }],
+      });
+      await pending;
+    });
+    expect(wasLoading).toBe(true);
+    expect(result.current.thread.getState().messages).toEqual([restored]);
+    expect(result.current.thread.getState().isLoading).toBe(false);
+  });
+
+  it.each(["existing", "new"])(
+    "keeps the selected %s thread when initial history finishes later",
+    async (target) => {
+      const { client } = createMockClient();
+      let resolveHistory!: (repo: ExportedMessageRepository) => void;
+      const pending = new Promise<ExportedMessageRepository>((resolve) => {
+        resolveHistory = resolve;
+      });
+      const history = { load: () => pending, append: async () => {} };
+      const { result } = renderHook(() => {
+        const [threadId, setThreadId] = useState("thread-a");
+        return useA2ARuntime({
+          client,
+          adapters: {
+            history,
+            threadList: {
+              threadId,
+              onSwitchToThread: async (id) => {
+                setThreadId(id);
+                return { messages: [createThreadMessage("message-b")] };
+              },
+              onSwitchToNewThread: async () => {
+                setThreadId("thread-new");
+              },
+            },
+          },
+        });
+      });
+
+      await act(async () => {
+        if (target === "existing")
+          await result.current.threads.switchToThread("thread-b");
+        else await result.current.threads.switchToNewThread();
+      });
+      const selectedId = result.current.threads.getState().mainThreadId;
+      const selectedMessages = result.current.thread.getState().messages;
+      await act(async () => {
+        resolveHistory({
+          headId: "message-a",
+          messages: [
+            { parentId: null, message: createThreadMessage("message-a") },
+          ],
+        });
+        await pending;
+      });
+      expect(result.current.threads.getState().mainThreadId).toBe(selectedId);
+      expect(result.current.thread.getState().messages).toEqual(
+        selectedMessages,
+      );
+      expect(result.current.thread.getState().isLoading).toBe(false);
+    },
+  );
+
   it("ignores a thread load superseded by a new thread", async () => {
     const { client } = createMockClient();
     let resolveLoad!: (value: { messages: ThreadMessage[] }) => void;