@assistant-ui/react-ag-ui 0.1.0 → 0.1.2

package/src/runtime/adapter/conversions.content.test.ts

@@ -1,7 +1,10 @@
 import { describe, expect, expectTypeOf, it } from "vitest";
 import type { Message as AgUiWireMessage } from "@ag-ui/client";
 import { MessageSchema, UserMessageSchema } from "@ag-ui/core/schemas";
-import type { AppendMessage } from "@assistant-ui/core";
+import {
+  SimpleImageAttachmentAdapter,
+  type AppendMessage,
+} from "@assistant-ui/core";
 import {
   fromAgUiMessages,
   toAgUiMessages,
@@ -79,6 +82,27 @@ describe("toAgUiMessages data URLs", () => {
 });
 
 describe("toAgUiMessages content metadata", () => {
+  it("preserves an uploaded image's filename through a snapshot round trip", async () => {
+    const adapter = new SimpleImageAttachmentAdapter();
+    const pending = await adapter.add({
+      file: new File([new Uint8Array([137, 80, 78, 71])], "cat.png", {
+        type: "image/png",
+      }),
+    });
+    const attachment = await adapter.send(pending);
+    const sent = toAgUiMessages([userMessage([], [attachment])]);
+    const restored = fromAgUiMessages(sent);
+
+    expect(restored[0]?.attachments?.[0]?.name).toBe("cat.png");
+    expect(sent[0]?.content).toEqual([
+      {
+        type: "image",
+        source: { type: "data", value: "iVBORw==", mimeType: "image/png" },
+        metadata: { filename: "cat.png" },
+      },
+    ]);
+  });
+
   it.each(["image", "document"] as const)(
     "preserves a %s file source and its metadata through a snapshot round trip",
     (type) => {
@@ -114,6 +138,135 @@ describe("toAgUiMessages content metadata", () => {
     },
   );
 
+  it("restores each of two file sources that share a handle", () => {
+    const original = {
+      id: "u-shared",
+      role: "user",
+      content: [
+        {
+          type: "document",
+          source: {
+            type: "file",
+            value: "shared",
+            mimeType: "application/pdf",
+          },
+          metadata: { filename: "a.pdf" },
+        },
+        {
+          type: "document",
+          source: {
+            type: "file",
+            value: "shared",
+            provider: "openai",
+            mimeType: "application/pdf",
+          },
+          metadata: { filename: "b.pdf" },
+        },
+      ],
+    };
+
+    expect(toAgUiMessages(fromAgUiMessages([original]))).toEqual([original]);
+  });
+
+  it.each(["image", "audio", "video"])(
+    "exports a %s file source without a mime type once",
+    (type) => {
+      const original = {
+        id: "u-untyped",
+        role: "user",
+        content: [{ type, source: { type: "file", value: `${type}-handle` } }],
+      };
+
+      expect(toAgUiMessages(fromAgUiMessages([original]))).toEqual([original]);
+    },
+  );
+
+  it("restores a file source a message carries only in its import metadata", () => {
+    const original = {
+      id: "u-legacy",
+      role: "user",
+      content: [
+        { type: "text", text: "label" },
+        {
+          type: "image",
+          source: { type: "url", value: "https://example.com/other.png" },
+        },
+        {
+          type: "document",
+          source: { type: "file", value: "doc-handle", provider: "openai" },
+        },
+      ],
+    };
+    const [imported] = fromAgUiMessages([original]);
+    const withoutHandles = {
+      ...imported!,
+      attachments: imported!.attachments?.filter((attachment) =>
+        attachment.content?.every(
+          (part) => part.type !== "file" || part.sourceType !== "id",
+        ),
+      ),
+    };
+
+    expect(withoutHandles.attachments).toHaveLength(1);
+    expect(toAgUiMessages([withoutHandles])).toEqual([original]);
+  });
+
+  describe.each(["content", "attachment"])("%s file handles", (location) => {
+    it.each([
+      ["document", "application/pdf", "file-pdf"],
+      ["image", "image/png", "https://provider.example/files/image"],
+      ["audio", "audio/wav", "file-audio"],
+      ["video", "video/mp4", "file-video"],
+    ])(
+      "preserves a %s handle instead of sending bytes",
+      (type, mimeType, data) => {
+        const part = {
+          type: "file",
+          data,
+          mimeType,
+          sourceType: "id",
+          filename: "upload.bin",
+          providerMetadata: { agui: { trace: "kept" } },
+        };
+        const message =
+          location === "content"
+            ? userMessage([part])
+            : userMessage([], [{ name: "upload.bin", content: [part] }]);
+        const exported = toAgUiMessages([message]);
+        const expected = {
+          id: "u1",
+          role: "user",
+          content: [
+            {
+              type,
+              source: { type: "file", value: data, mimeType },
+              metadata: { trace: "kept", filename: "upload.bin" },
+            },
+          ],
+        };
+
+        const imported = fromAgUiMessages(exported);
+
+        expect(MessageSchema.parse(exported[0])).toEqual(expected);
+        expect(imported[0]?.attachments).toEqual([
+          expect.objectContaining({
+            content: [
+              {
+                type: "file",
+                data,
+                mimeType,
+                sourceType: "id",
+                filename: "upload.bin",
+                providerMetadata: { agui: { trace: "kept" } },
+              },
+            ],
+          }),
+        ]);
+        expect(toAgUiMessages(imported)).toEqual([expected]);
+      },
+    );
+  });
+
   it("emits AgUiMessage values assignable to AG-UI Message", () => {
     expectTypeOf<AgUiMessage>().toExtend<AgUiWireMessage>();
   });
@@ -405,3 +558,91 @@ describe("toAgUiMessages percent-encoded data URLs", () => {
     expect(toAgUiMessages(rebuilt as never)).toEqual(sent);
   });
 });
+
+describe("toAgUiMessages tool model content", () => {
+  const withModelContent = (modelContent: unknown): Message => ({
+    id: "a1",
+    role: "assistant",
+    content: [
+      {
+        type: "tool-call",
+        toolCallId: "call-1",
+        toolName: "download",
+        args: {},
+        argsText: "{}",
+        result: { private: "not model-visible" },
+        modelContent,
+      },
+    ],
+  });
+
+  it.each([
+    ["image", "image/png"],
+    ["document", "application/pdf"],
+    ["audio", "audio/wav"],
+    ["video", "video/mp4"],
+  ])(
+    "keeps %s model content through a snapshot round trip",
+    (type, mediaType) => {
+      const exported = toAgUiMessages([
+        withModelContent([
+          { type: "text", text: "Before" },
+          { type: "file", data: "YWJj", mediaType, filename: "result.bin" },
+          { type: "text", text: "After" },
+        ]),
+      ]);
+      const expected = {
+        id: "call-1:tool",
+        role: "tool",
+        toolCallId: "call-1",
+        content: [
+          { type: "text", text: "Before" },
+          {
+            type,
+            source: { type: "data", value: "YWJj", mimeType: mediaType },
+            metadata: { filename: "result.bin" },
+          },
+          { type: "text", text: "After" },
+        ],
+      };
+
+      expect(MessageSchema.parse(exported[1])).toEqual(expected);
+      expect(toAgUiMessages(fromAgUiMessages(exported))[1]).toEqual(expected);
+    },
+  );
+
+  it("sends a file-only model result instead of an empty string", () => {
+    const exported = toAgUiMessages([
+      withModelContent([
+        {
+          type: "file",
+          data: "https://example.com/report.pdf",
+          mediaType: "application/pdf",
+        },
+      ]),
+    ]);
+
+    expect(exported[1]?.content).toEqual([
+      {
+        type: "document",
+        source: {
+          type: "url",
+          value: "https://example.com/report.pdf",
+          mimeType: "application/pdf",
+        },
+      },
+    ]);
+  });
+
+  it.each([
+    { modelContent: [], content: "" },
+    { modelContent: [{ type: "text", text: "public" }], content: "public" },
+  ])(
+    "keeps text-only model content as $content",
+    ({ modelContent, content }) => {
+      expect(toAgUiMessages([withModelContent(modelContent)])[1]?.content).toBe(
+        content,
+      );
+    },
+  );
+});