@assistant-ui/react-a2a 0.2.21 → 0.2.23

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" }]);
+  });
+});