@assistant-ui/react-a2a 0.2.22 → 0.2.24

package/src/A2AClient.ts

@@ -1,4 +1,5 @@
 import { SSEEventDecoder, type SSEEvent } from "assistant-stream/utils";
+import { isRecord } from "@assistant-ui/core/internal";
 import type {
   A2AAgentCard,
   A2AErrorInfo,
@@ -153,6 +154,114 @@ function discriminateStreamResponse(
   return null;
 }
 
+const TASK_STATES: ReadonlySet<string> = new Set(
+  Object.keys({
+    unspecified: true,
+    submitted: true,
+    working: true,
+    completed: true,
+    failed: true,
+    canceled: true,
+    input_required: true,
+    rejected: true,
+    auth_required: true,
+  } satisfies Record<A2ATaskState, true>),
+);
+
+const isTaskState = (value: unknown): value is A2ATaskState =>
+  typeof value === "string" && TASK_STATES.has(value);
+
+const isTask = (value: unknown): value is A2ATask =>
+  isRecord(value) &&
+  typeof value.id === "string" &&
+  value.id.length > 0 &&
+  isRecord(value.status) &&
+  isTaskState(value.status.state);
+
+const isMessage = (value: unknown): value is A2AMessage =>
+  isRecord(value) &&
+  typeof value.messageId === "string" &&
+  value.messageId.length > 0 &&
+  typeof value.role === "string" &&
+  value.role.length > 0 &&
+  Array.isArray(value.parts) &&
+  value.parts.every(isRecord);
+
+const isStringArray = (value: unknown): value is string[] =>
+  Array.isArray(value) && value.every((item) => typeof item === "string");
+
+const invalidAgentCard = (): never => {
+  throw new Error(
+    "Invalid A2A agent card response: expected a valid agent card payload.",
+  );
+};
+
+const parseCardString = (value: unknown): string =>
+  value == null ? "" : typeof value === "string" ? value : invalidAgentCard();
+
+const parseCardStringArray = (value: unknown): string[] =>
+  value == null ? [] : isStringArray(value) ? value : invalidAgentCard();
+
+const parseCardRecordArray = (value: unknown): Record<string, unknown>[] =>
+  value == null
+    ? []
+    : Array.isArray(value) && value.every(isRecord)
+      ? (value as Record<string, unknown>[])
+      : invalidAgentCard();
+
+const parseCardRecord = (value: unknown): Record<string, unknown> =>
+  value == null ? {} : isRecord(value) ? value : invalidAgentCard();
+
+// Proto3 JSON parsing treats omitted and null fields as defaults, so a valid
+// card may arrive without its empty lists, strings, or capabilities. Fill
+// those per the proto3 JSON mapping rules; a payload without a name or with a
+// present field of the wrong type rejects.
+const parseAgentCardResponse = (value: unknown): A2AAgentCard => {
+  if (
+    !isRecord(value) ||
+    typeof value.name !== "string" ||
+    value.name.length === 0
+  ) {
+    return invalidAgentCard();
+  }
+
+  return {
+    ...value,
+    name: value.name,
+    description: parseCardString(value.description),
+    version: parseCardString(value.version),
+    supportedInterfaces: parseCardRecordArray(value.supportedInterfaces).map(
+      (entry) => ({
+        ...entry,
+        url: parseCardString(entry.url),
+        protocolBinding: parseCardString(entry.protocolBinding),
+        protocolVersion: parseCardString(entry.protocolVersion),
+      }),
+    ),
+    capabilities: parseCardRecord(value.capabilities),
+    defaultInputModes: parseCardStringArray(value.defaultInputModes),
+    defaultOutputModes: parseCardStringArray(value.defaultOutputModes),
+    skills: parseCardRecordArray(value.skills).map((entry) => ({
+      ...entry,
+      id: parseCardString(entry.id),
+      name: parseCardString(entry.name),
+      description: parseCardString(entry.description),
+      tags: parseCardStringArray(entry.tags),
+    })),
+  } as A2AAgentCard;
+};
+
+const parseSendMessageResponse = (value: unknown): A2ATask | A2AMessage => {
+  if (isRecord(value)) {
+    const candidate = value.task ?? value.message ?? value;
+    if (isTask(candidate) || isMessage(candidate)) return candidate;
+  }
+
+  throw new Error(
+    "Invalid A2A message:send response: expected a valid task or message payload.",
+  );
+};
+
 function signalInit(signal?: AbortSignal): RequestInit {
   return signal ? { signal } : {};
 }
@@ -273,14 +382,15 @@ export class A2AClient {
       await this.throwResponseError(response);
     }
     const json = await response.json();
-    return normalizeKeys(json) as A2AAgentCard;
+    return parseAgentCardResponse(normalizeKeys(json));
   }
 
   async getExtendedAgentCard(signal?: AbortSignal): Promise<A2AAgentCard> {
-    return this.fetchJSON<A2AAgentCard>(
+    const result = await this.fetchJSON<unknown>(
       `${this.getBasePath()}/extendedAgentCard`,
       signalInit(signal),
     );
+    return parseAgentCardResponse(result);
   }
 
   // --- Message ---
@@ -297,7 +407,7 @@ export class A2AClient {
     if (configuration) body.configuration = configuration;
     if (metadata) body.metadata = metadata;
 
-    const result = await this.fetchJSON<Record<string, unknown>>(
+    const result = await this.fetchJSON<unknown>(
       `${this.getBasePath()}/message:send`,
       {
         method: "POST",
@@ -306,11 +416,7 @@ export class A2AClient {
       },
     );
 
-    // Unwrap SendMessageResponse: {task: Task} | {message: Message}
-    if ("task" in result && result.task) return result.task as A2ATask;
-    if ("message" in result && result.message)
-      return result.message as A2AMessage;
-    return result as unknown as A2ATask | A2AMessage;
+    return parseSendMessageResponse(result);
   }
 
   async *streamMessage(

package/src/A2AThreadRuntimeCore.test.ts

@@ -40,6 +40,57 @@ function createUserAppendMessage(text: string): AppendMessage {
   } as unknown as AppendMessage;
 }
 
+function createHistoryMessage(
+  id: string,
+  role: "user" | "assistant",
+  text: string,
+): ThreadMessage {
+  return {
+    id,
+    role,
+    createdAt: new Date(),
+    content: [{ type: "text", text }],
+    status: { type: "complete", reason: "stop" },
+    ...(role === "assistant"
+      ? {
+          metadata: {
+            unstable_state: null,
+            unstable_annotations: [],
+            unstable_data: [],
+            steps: [],
+            custom: {},
+          },
+        }
+      : {}),
+  } as ThreadMessage;
+}
+
+function createBranchedHistory() {
+  const user = createHistoryMessage("user", "user", "Question");
+  const firstAssistant = createHistoryMessage(
+    "assistant-a",
+    "assistant",
+    "First answer",
+  );
+  const secondAssistant = createHistoryMessage(
+    "assistant-b",
+    "assistant",
+    "Second answer",
+  );
+  const history = {
+    load: vi.fn().mockResolvedValue({
+      headId: secondAssistant.id,
+      messages: [
+        { parentId: null, message: user },
+        { parentId: user.id, message: firstAssistant },
+        { parentId: user.id, message: secondAssistant },
+      ],
+    }),
+    append: vi.fn().mockResolvedValue(undefined),
+  };
+  return { user, firstAssistant, secondAssistant, history };
+}
+
 function statusUpdateEvent(state: string, text?: string): A2AStreamEvent {
   return {
     type: "statusUpdate",
@@ -119,6 +170,203 @@ describe("A2AThreadRuntimeCore", () => {
     });
   });
 
+  describe("history loading", () => {
+    it("preserves sibling branches and selects the persisted head", async () => {
+      const { user, firstAssistant, secondAssistant, history } =
+        createBranchedHistory();
+      const core = createCore({}, { history });
+
+      await core.__internal_load();
+
+      expect(core.getMessages().map((message) => message.id)).toEqual([
+        user.id,
+        secondAssistant.id,
+      ]);
+      expect(
+        core.getMessageRepository().messages.map(({ message, parentId }) => ({
+          id: message.id,
+          parentId,
+        })),
+      ).toEqual([
+        { id: user.id, parentId: null },
+        { id: firstAssistant.id, parentId: user.id },
+        { id: secondAssistant.id, parentId: user.id },
+      ]);
+    });
+
+    it("falls back to linear history when stored parents are invalid", async () => {
+      const user = createHistoryMessage("user", "user", "Question");
+      const assistant = createHistoryMessage(
+        "assistant",
+        "assistant",
+        "Answer",
+      );
+      const history = {
+        load: vi.fn().mockResolvedValue({
+          headId: assistant.id,
+          messages: [
+            { parentId: "missing", message: user },
+            { parentId: user.id, message: assistant },
+          ],
+        }),
+        append: vi.fn().mockResolvedValue(undefined),
+      };
+      const core = createCore({}, { history });
+
+      await expect(core.__internal_load()).resolves.toBeUndefined();
+
+      expect(core.getMessages().map((message) => message.id)).toEqual([
+        user.id,
+        assistant.id,
+      ]);
+      expect(
+        core.getMessageRepository().messages.map(({ message, parentId }) => ({
+          id: message.id,
+          parentId,
+        })),
+      ).toEqual([
+        { id: user.id, parentId: null },
+        { id: assistant.id, parentId: user.id },
+      ]);
+    });
+
+    it("falls back to linear history when stored ids are duplicated", async () => {
+      const first = createHistoryMessage("duplicate", "user", "First");
+      const replacement = createHistoryMessage(
+        "duplicate",
+        "assistant",
+        "Replacement",
+      );
+      const tail = createHistoryMessage("tail", "user", "Tail");
+      const history = {
+        load: vi.fn().mockResolvedValue({
+          headId: tail.id,
+          messages: [
+            { parentId: null, message: first },
+            { parentId: null, message: replacement },
+            { parentId: replacement.id, message: tail },
+          ],
+        }),
+        append: vi.fn().mockResolvedValue(undefined),
+      };
+      const core = createCore({}, { history });
+
+      await expect(core.__internal_load()).resolves.toBeUndefined();
+
+      const repository = core.getMessageRepository();
+      expect(
+        repository.messages.map(({ message, parentId }) => ({
+          id: message.id,
+          parentId,
+        })),
+      ).toEqual([
+        { id: replacement.id, parentId: null },
+        { id: tail.id, parentId: replacement.id },
+      ]);
+      expect(repository.messages[0]!.message.content).toEqual(
+        replacement.content,
+      );
+    });
+
+    it("falls back to linear history when stored parents form a cycle", async () => {
+      const first = createHistoryMessage("first", "user", "First");
+      const second = createHistoryMessage("second", "assistant", "Second");
+      const history = {
+        load: vi.fn().mockResolvedValue({
+          headId: second.id,
+          messages: [
+            { parentId: second.id, message: first },
+            { parentId: first.id, message: second },
+          ],
+        }),
+        append: vi.fn().mockResolvedValue(undefined),
+      };
+      const core = createCore({}, { history });
+
+      await expect(core.__internal_load()).resolves.toBeUndefined();
+
+      expect(
+        core.getMessageRepository().messages.map(({ message, parentId }) => ({
+          id: message.id,
+          parentId,
+        })),
+      ).toEqual([
+        { id: first.id, parentId: null },
+        { id: second.id, parentId: first.id },
+      ]);
+    });
+
+    it("keeps hidden siblings when the visible branch changes", async () => {
+      const { user, firstAssistant, secondAssistant, history } =
+        createBranchedHistory();
+      const core = createCore({}, { history });
+
+      await core.__internal_load();
+      core.applyExternalMessages([user, firstAssistant]);
+
+      expect(core.getMessages().map((message) => message.id)).toEqual([
+        user.id,
+        firstAssistant.id,
+      ]);
+      expect(core.getMessageRepository().headId).toBe(firstAssistant.id);
+      expect(
+        core.getMessageRepository().messages.map(({ message }) => message.id),
+      ).toEqual([user.id, firstAssistant.id, secondAssistant.id]);
+    });
+
+    it("replaces stored branches when external messages change parentage", async () => {
+      const { secondAssistant, history } = createBranchedHistory();
+      const core = createCore({}, { history });
+
+      await core.__internal_load();
+      core.applyExternalMessages([secondAssistant]);
+
+      expect(core.getMessages().map((message) => message.id)).toEqual([
+        secondAssistant.id,
+      ]);
+      expect(
+        core.getMessageRepository().messages.map(({ message, parentId }) => ({
+          id: message.id,
+          parentId,
+        })),
+      ).toEqual([{ id: secondAssistant.id, parentId: null }]);
+    });
+
+    it("adds regenerated responses without dropping loaded siblings", async () => {
+      const { user, firstAssistant, secondAssistant, history } =
+        createBranchedHistory();
+      const core = createCore(
+        {
+          streamMessage: vi.fn().mockImplementation(async function* () {
+            yield statusUpdateEvent("completed", "Regenerated answer");
+          }),
+        },
+        { history },
+      );
+
+      await core.__internal_load();
+      await core.reload(user.id);
+
+      const visibleMessages = core.getMessages();
+      const regenerated = visibleMessages[1]!;
+      expect(visibleMessages.map((message) => message.id)).toEqual([
+        user.id,
+        regenerated.id,
+      ]);
+      expect(regenerated.content).toEqual([
+        { type: "text", text: "Regenerated answer" },
+      ]);
+      expect(
+        core.getMessageRepository().messages.map(({ message }) => message.id),
+      ).toEqual([
+        user.id,
+        firstAssistant.id,
+        secondAssistant.id,
+        regenerated.id,
+      ]);
+    });
+  });
+
   // --- Edit & Reload ---
 
   describe("edit", () => {
@@ -136,6 +384,80 @@ describe("A2AThreadRuntimeCore", () => {
       expect(messages[0]!.role).toBe("user");
       expect(messages[1]!.role).toBe("assistant");
     });
+
+    it("keeps the replaced message subtree as a sibling branch", async () => {
+      const root = createHistoryMessage("root", "user", "First question");
+      const parent = createHistoryMessage(
+        "parent",
+        "assistant",
+        "First answer",
+      );
+      const source = createHistoryMessage(
+        "source",
+        "user",
+        "Original follow-up",
+      );
+      const child = createHistoryMessage(
+        "child",
+        "assistant",
+        "Original response",
+      );
+      const sibling = createHistoryMessage(
+        "sibling",
+        "user",
+        "Alternate follow-up",
+      );
+      const history = {
+        load: vi.fn().mockResolvedValue({
+          headId: child.id,
+          messages: [
+            { parentId: null, message: root },
+            { parentId: root.id, message: parent },
+            { parentId: parent.id, message: source },
+            { parentId: source.id, message: child },
+            { parentId: parent.id, message: sibling },
+          ],
+        }),
+        append: vi.fn().mockResolvedValue(undefined),
+      };
+      const core = createCore({}, { history });
+
+      await core.__internal_load();
+      await core.edit({
+        ...createUserAppendMessage("Edited follow-up"),
+        parentId: parent.id,
+        sourceId: source.id,
+        startRun: false,
+      });
+
+      const edited = core.getMessages().at(-1)!;
+      expect(core.getMessages().map((message) => message.id)).toEqual([
+        root.id,
+        parent.id,
+        edited.id,
+      ]);
+      expect(
+        core.getMessageRepository().messages.map(({ message, parentId }) => ({
+          id: message.id,
+          parentId,
+        })),
+      ).toEqual([
+        { id: root.id, parentId: null },
+        { id: parent.id, parentId: root.id },
+        { id: source.id, parentId: parent.id },
+        { id: child.id, parentId: source.id },
+        { id: sibling.id, parentId: parent.id },
+        { id: edited.id, parentId: parent.id },
+      ]);
+
+      core.applyExternalMessages([root, parent, source, child]);
+      expect(core.getMessages().map((message) => message.id)).toEqual([
+        root.id,
+        parent.id,
+        source.id,
+        child.id,
+      ]);
+    });
   });
 
   describe("reload", () => {
@@ -690,3 +1012,139 @@ describe("A2AThreadRuntimeCore", () => {
     });
   });
 });
+
+describe("outbound message conversion", () => {
+  function createCoreWithStream() {
+    const streamMessage = vi.fn().mockImplementation(async function* () {});
+    const core = new A2AThreadRuntimeCore({
+      client: createMockClient({ streamMessage }),
+      notifyUpdate: vi.fn() as unknown as () => void,
+    });
+    return { core, streamMessage };
+  }
+
+  it("forwards attachment content to the wire with the attachment MIME type", async () => {
+    const { core, streamMessage } = createCoreWithStream();
+
+    await core.append({
+      parentId: null,
+      role: "user",
+      content: [{ type: "text", text: "See attached" }],
+      attachments: [
+        {
+          id: "att-1",
+          type: "document",
+          name: "doc.pdf",
+          contentType: "application/pdf",
+          status: { type: "complete" },
+          content: [
+            {
+              type: "file",
+              data: "https://files.com/doc.pdf",
+              mimeType: "",
+              filename: "doc.pdf",
+            },
+          ],
+        },
+      ],
+    } as unknown as AppendMessage);
+
+    const sent = streamMessage.mock.calls[0]![0] as A2AMessage;
+    expect(sent.parts).toEqual([
+      { text: "See attached" },
+      {
+        url: "https://files.com/doc.pdf",
+        mediaType: "application/pdf",
+        filename: "doc.pdf",
+      },
+    ]);
+  });
+
+  it("forwards file parts in message content to the wire", async () => {
+    const { core, streamMessage } = createCoreWithStream();
+
+    await core.append({
+      parentId: null,
+      role: "user",
+      content: [
+        {
+          type: "file",
+          data: "data:text/csv;base64,ZmlsZQ==",
+          mimeType: "text/csv",
+        },
+      ],
+    } as unknown as AppendMessage);
+
+    const sent = streamMessage.mock.calls[0]![0] as A2AMessage;
+    expect(sent.parts).toEqual([{ raw: "ZmlsZQ==", mediaType: "text/csv" }]);
+  });
+
+  it("does not throw for user messages missing attachments at runtime", async () => {
+    const { core, streamMessage } = createCoreWithStream();
+    const message = {
+      id: "u1",
+      role: "user",
+      createdAt: new Date(),
+      content: [{ type: "text", text: "Loaded" }],
+      status: { type: "complete", reason: "stop" },
+    } as unknown as ThreadMessage;
+
+    core.applyExternalMessages([message]);
+    await core.reload("u1");
+
+    const sent = streamMessage.mock.calls[0]![0] as A2AMessage;
+    expect(sent.parts).toEqual([{ text: "Loaded" }]);
+  });
+
+  it("forwards data URL image attachments as raw bytes", async () => {
+    const { core, streamMessage } = createCoreWithStream();
+
+    await core.append({
+      parentId: null,
+      role: "user",
+      content: [{ type: "text", text: "See image" }],
+      attachments: [
+        {
+          id: "att-2",
+          type: "image",
+          name: "a.png",
+          contentType: "image/png",
+          status: { type: "complete" },
+          content: [{ type: "image", image: "data:image/png;base64,aGVsbG8=" }],
+        },
+      ],
+    } as unknown as AppendMessage);
+
+    const sent = streamMessage.mock.calls[0]![0] as A2AMessage;
+    expect(sent.parts).toEqual([
+      { text: "See image" },
+      { raw: "aGVsbG8=", mediaType: "image/png" },
+    ]);
+  });
+
+  it("does not throw for attachments missing content at runtime", async () => {
+    const { core, streamMessage } = createCoreWithStream();
+    const message = {
+      id: "u2",
+      role: "user",
+      createdAt: new Date(),
+      content: [{ type: "text", text: "Loaded" }],
+      attachments: [
+        {
+          id: "att-3",
+          type: "file",
+          name: "x.txt",
+          contentType: "text/plain",
+          status: { type: "complete" },
+        },
+      ],
+      status: { type: "complete", reason: "stop" },
+    } as unknown as ThreadMessage;
+
+    core.applyExternalMessages([message]);
+    await core.reload("u2");
+
+    const sent = streamMessage.mock.calls[0]![0] as A2AMessage;
+    expect(sent.parts).toEqual([{ text: "Loaded" }]);
+  });
+});