@assistant-ui/react-ai-sdk 1.3.21 → 1.3.25

package/src/ui/utils/convertMessage.test.ts

@@ -1,4 +1,5 @@
 import { describe, expect, it } from "vitest";
+import type { ReadonlyJSONObject } from "assistant-stream/utils";
 import { AISDKMessageConverter } from "./convertMessage";
 
 describe("AISDKMessageConverter", () => {
@@ -37,6 +38,94 @@ describe("AISDKMessageConverter", () => {
     expect(converted[0]?.attachments?.[1]?.type).toBe("file");
   });
 
+  it("converts source-document parts into document sources", () => {
+    const converted = AISDKMessageConverter.toThreadMessages([
+      {
+        id: "a1",
+        role: "assistant",
+        parts: [
+          {
+            type: "source-document",
+            sourceId: "doc_123",
+            title: "proposal.pdf",
+            mediaType: "application/pdf",
+            filename: "proposal.pdf",
+            providerMetadata: {
+              openai: {
+                type: "file_citation",
+                fileId: "file_123",
+                index: 0,
+              },
+            },
+          },
+        ],
+      } as any,
+    ]);
+
+    expect(converted).toHaveLength(1);
+    expect(converted[0]?.role).toBe("assistant");
+
+    const sourcePart = converted[0]?.content.find(
+      (part): part is any => part.type === "source",
+    );
+
+    expect(sourcePart).toMatchObject({
+      type: "source",
+      sourceType: "document",
+      id: "doc_123",
+      title: "proposal.pdf",
+      mediaType: "application/pdf",
+      filename: "proposal.pdf",
+      providerMetadata: {
+        openai: {
+          type: "file_citation",
+          fileId: "file_123",
+          index: 0,
+        },
+      },
+    });
+  });
+
+  it("converts source-url parts without synthesizing missing optional fields", () => {
+    const converted = AISDKMessageConverter.toThreadMessages([
+      {
+        id: "a1",
+        role: "assistant",
+        parts: [
+          {
+            type: "source-url",
+            sourceId: "url_123",
+            url: "https://example.com/report",
+            providerMetadata: {
+              openai: {
+                type: "url_citation",
+                index: 1,
+              },
+            },
+          },
+        ],
+      } as any,
+    ]);
+
+    const sourcePart = converted[0]?.content.find(
+      (part): part is any => part.type === "source",
+    );
+
+    expect(sourcePart).toMatchObject({
+      type: "source",
+      sourceType: "url",
+      id: "url_123",
+      url: "https://example.com/report",
+      providerMetadata: {
+        openai: {
+          type: "url_citation",
+          index: 1,
+        },
+      },
+    });
+    expect(sourcePart).not.toHaveProperty("title");
+  });
+
   it("converts assistant image file parts into file content", () => {
     const converted = AISDKMessageConverter.toThreadMessages([
       {
@@ -307,4 +396,246 @@ describe("AISDKMessageConverter", () => {
       limit: 5,
     });
   });
+
+  it("preserves last good input when AI SDK briefly emits null input", () => {
+    const metadata = {
+      toolArgsKeyOrderCache: new Map<string, Map<string, string[]>>(),
+      toolLastInputCache: new Map<string, ReadonlyJSONObject>(),
+    };
+
+    const convertWithInput = (input: unknown) =>
+      AISDKMessageConverter.toThreadMessages(
+        [
+          {
+            id: "a1",
+            role: "assistant",
+            parts: [
+              {
+                type: "tool-weather",
+                toolCallId: "tc-1",
+                state: "input-streaming",
+                input,
+              },
+            ],
+          } as any,
+        ],
+        false,
+        metadata,
+      )[0]?.content.find((part): part is any => part.type === "tool-call");
+
+    const first = convertWithInput({ city: "NYC" });
+    expect(first?.argsText).toBe('{"city":"NYC');
+    expect(first?.args).toEqual({ city: "NYC" });
+
+    const dropped = convertWithInput(null);
+    expect(dropped?.argsText).toBe('{"city":"NYC');
+    expect(dropped?.args).toEqual({ city: "NYC" });
+
+    const undef = convertWithInput(undefined);
+    expect(undef?.argsText).toBe('{"city":"NYC');
+    expect(undef?.args).toEqual({ city: "NYC" });
+
+    const grown = convertWithInput({ city: "NYC", units: "F" });
+    expect(grown?.argsText).toBe('{"city":"NYC","units":"F');
+    expect(grown?.args).toEqual({ city: "NYC", units: "F" });
+  });
+
+  it("preserves last good input across terminal state transitions", () => {
+    const metadata = {
+      toolArgsKeyOrderCache: new Map<string, Map<string, string[]>>(),
+      toolLastInputCache: new Map<string, ReadonlyJSONObject>(),
+    };
+
+    AISDKMessageConverter.toThreadMessages(
+      [
+        {
+          id: "a1",
+          role: "assistant",
+          parts: [
+            {
+              type: "tool-weather",
+              toolCallId: "tc-1",
+              state: "input-available",
+              input: { city: "NYC" },
+            },
+          ],
+        } as any,
+      ],
+      false,
+      metadata,
+    );
+
+    const terminal = AISDKMessageConverter.toThreadMessages(
+      [
+        {
+          id: "a1",
+          role: "assistant",
+          parts: [
+            {
+              type: "tool-weather",
+              toolCallId: "tc-1",
+              state: "output-available",
+              input: null,
+              output: { temp: 70 },
+            },
+          ],
+        } as any,
+      ],
+      false,
+      metadata,
+    );
+
+    const call = terminal[0]?.content.find(
+      (part): part is any => part.type === "tool-call",
+    );
+    expect(call?.args).toEqual({ city: "NYC" });
+    expect(call?.result).toEqual({ temp: 70 });
+  });
+
+  it("unwraps the modelContent envelope produced by frontend tool execution", () => {
+    const converted = AISDKMessageConverter.toThreadMessages([
+      {
+        id: "a1",
+        role: "assistant",
+        parts: [
+          {
+            type: "tool-readPdf",
+            toolCallId: "tc-pdf",
+            state: "output-available",
+            input: {},
+            output: {
+              __aui_modelContent: [
+                { type: "text", text: "PDF contents:" },
+                {
+                  type: "file",
+                  data: "JVBERi0xLjQK",
+                  mediaType: "application/pdf",
+                },
+              ],
+              value: { mediaType: "application/pdf", base64: "JVBERi0xLjQK" },
+            },
+          },
+        ],
+      } as any,
+    ]);
+
+    const call = converted[0]?.content.find(
+      (part): part is any => part.type === "tool-call",
+    );
+    expect(call?.result).toEqual({
+      mediaType: "application/pdf",
+      base64: "JVBERi0xLjQK",
+    });
+    expect(call?.modelContent).toEqual([
+      { type: "text", text: "PDF contents:" },
+      {
+        type: "file",
+        data: "JVBERi0xLjQK",
+        mediaType: "application/pdf",
+      },
+    ]);
+  });
+
+  it("leaves a plain output untouched when no envelope is present", () => {
+    const converted = AISDKMessageConverter.toThreadMessages([
+      {
+        id: "a1",
+        role: "assistant",
+        parts: [
+          {
+            type: "tool-weather",
+            toolCallId: "tc-1",
+            state: "output-available",
+            input: { city: "NYC" },
+            output: { temp: 72 },
+          },
+        ],
+      } as any,
+    ]);
+
+    const call = converted[0]?.content.find(
+      (part): part is any => part.type === "tool-call",
+    );
+    expect(call?.result).toEqual({ temp: 72 });
+    expect(call?.modelContent).toBeUndefined();
+  });
+
+  it("forwards callProviderMetadata.mcp.app onto ToolCallMessagePart.mcp.app", () => {
+    const converted = AISDKMessageConverter.toThreadMessages([
+      {
+        id: "a1",
+        role: "assistant",
+        parts: [
+          {
+            type: "tool-search",
+            toolCallId: "tc-1",
+            state: "output-available",
+            input: { query: "hi" },
+            output: { results: [] },
+            callProviderMetadata: {
+              mcp: {
+                app: {
+                  resourceUri: "ui://example/search",
+                  mimeType: "text/html;profile=mcp-app",
+                  visibility: ["app", "model", "bogus"],
+                },
+              },
+            },
+          },
+        ],
+      } as any,
+    ]);
+
+    const call = converted[0]?.content.find(
+      (part): part is any => part.type === "tool-call",
+    );
+    expect(call?.mcp?.app).toEqual({
+      resourceUri: "ui://example/search",
+      mimeType: "text/html;profile=mcp-app",
+      visibility: ["app", "model"],
+    });
+  });
+
+  it("memoizes MCP app metadata across conversions by resourceUri", () => {
+    const metadata = {
+      mcpAppMetadataCache: new Map(),
+    };
+
+    const buildMessage = (id: string) => ({
+      id,
+      role: "assistant" as const,
+      parts: [
+        {
+          type: "tool-search",
+          toolCallId: `${id}-call`,
+          state: "output-available",
+          input: { q: "hi" },
+          output: {},
+          callProviderMetadata: {
+            mcp: { app: { resourceUri: "ui://example/search" } },
+          },
+        } as any,
+      ],
+    });
+
+    const first = AISDKMessageConverter.toThreadMessages(
+      [buildMessage("a1")],
+      false,
+      metadata,
+    );
+    const second = AISDKMessageConverter.toThreadMessages(
+      [buildMessage("a2")],
+      false,
+      metadata,
+    );
+
+    const firstApp = first[0]?.content.find(
+      (p): p is any => p.type === "tool-call",
+    )?.mcp?.app;
+    const secondApp = second[0]?.content.find(
+      (p): p is any => p.type === "tool-call",
+    )?.mcp?.app;
+    expect(firstApp).toBeDefined();
+    expect(firstApp).toBe(secondApp);
+  });
 });

package/src/ui/utils/convertMessage.ts

@@ -3,27 +3,75 @@ import {
   createMessageConverter as unstable_createMessageConverter,
   type useExternalMessageConverter,
 } from "@assistant-ui/core/react";
-import type {
-  ReasoningMessagePart,
-  ToolCallMessagePart,
-  TextMessagePart,
-  DataMessagePart,
-  SourceMessagePart,
-  FileMessagePart,
-  ThreadMessageLike,
+import {
+  isMcpAppUri,
+  type ReasoningMessagePart,
+  type ToolCallMessagePart,
+  type TextMessagePart,
+  type DataMessagePart,
+  type SourceMessagePart,
+  type SourceProviderMetadata,
+  type FileMessagePart,
+  type ThreadMessageLike,
+  type McpAppMetadata,
 } from "@assistant-ui/core";
 import type { ReadonlyJSONObject } from "assistant-stream/utils";
+import { unwrapModelContentEnvelope } from "../../modelContentEnvelope";
 
 type MessageMetadata = ThreadMessageLike["metadata"];
 export type AISDKMessageConverterMetadata =
   useExternalMessageConverter.Metadata & {
     toolArgsKeyOrderCache?: Map<string, Map<string, string[]>>;
+    toolLastInputCache?: Map<string, ReadonlyJSONObject>;
+    mcpAppMetadataCache?: Map<string, McpAppMetadata>;
   };
 
 function stripClosingDelimiters(json: string): string {
   return json.replace(/[}\]"]+$/, "");
 }
 
+const MCP_APP_METADATA_CACHE_MAX = 100;
+
+function extractMcpAppMetadata(
+  part: unknown,
+  cache: Map<string, McpAppMetadata> | undefined,
+): McpAppMetadata | undefined {
+  if (!part || typeof part !== "object") return undefined;
+  const meta = (part as { callProviderMetadata?: unknown })
+    .callProviderMetadata;
+  if (!meta || typeof meta !== "object") return undefined;
+  const mcp = (meta as { mcp?: unknown }).mcp;
+  if (!mcp || typeof mcp !== "object") return undefined;
+  const app = (mcp as { app?: unknown }).app;
+  if (!app || typeof app !== "object") return undefined;
+  const a = app as Record<string, unknown>;
+  if (typeof a["resourceUri"] !== "string") return undefined;
+  if (!isMcpAppUri(a["resourceUri"])) return undefined;
+  const cached = cache?.get(a["resourceUri"]);
+  if (cached) {
+    cache!.delete(a["resourceUri"]);
+    cache!.set(a["resourceUri"], cached);
+    return cached;
+  }
+  const out: { -readonly [K in keyof McpAppMetadata]: McpAppMetadata[K] } = {
+    resourceUri: a["resourceUri"],
+  };
+  if (typeof a["mimeType"] === "string") out.mimeType = a["mimeType"];
+  if (Array.isArray(a["visibility"])) {
+    out.visibility = a["visibility"].filter(
+      (v): v is "model" | "app" => v === "model" || v === "app",
+    );
+  }
+  if (cache) {
+    if (cache.size >= MCP_APP_METADATA_CACHE_MAX) {
+      const oldest = cache.keys().next().value;
+      if (oldest !== undefined) cache.delete(oldest);
+    }
+    cache.set(a["resourceUri"], out);
+  }
+  return out;
+}
+
 const hasOwn = (value: object, key: string) => Object.hasOwn(value, key);
 
 const stabilizeToolArgsValue = (
@@ -147,14 +195,28 @@ function convertParts(
         const toolName = getToolName(part);
         const toolCallId = part.toolCallId;
         const argsKeyOrderCacheKey = `${message.id}:${toolCallId}`;
-        const args: ReadonlyJSONObject =
-          (part.input as ReadonlyJSONObject) || {};
+
+        const rawInput = part.input as ReadonlyJSONObject | null | undefined;
+        let args: ReadonlyJSONObject;
+        if (
+          rawInput != null &&
+          typeof rawInput === "object" &&
+          !Array.isArray(rawInput)
+        ) {
+          args = rawInput;
+          metadata.toolLastInputCache?.set(argsKeyOrderCacheKey, args);
+        } else {
+          args = metadata.toolLastInputCache?.get(argsKeyOrderCacheKey) ?? {};
+        }
 
         let result: unknown;
+        let modelContent: ToolCallMessagePart["modelContent"];
         let isError = false;
 
         if (part.state === "output-available") {
-          result = part.output;
+          const unwrapped = unwrapModelContentEnvelope(part.output);
+          result = unwrapped.result;
+          modelContent = unwrapped.modelContent;
         } else if (part.state === "output-error") {
           isError = true;
           result = { error: part.errorText };
@@ -177,9 +239,20 @@ function convertParts(
           argsText = stripClosingDelimiters(argsText);
         } else {
           metadata.toolArgsKeyOrderCache?.delete(argsKeyOrderCacheKey);
+          if (
+            part.state === "output-available" ||
+            part.state === "output-error" ||
+            part.state === "output-denied"
+          ) {
+            metadata.toolLastInputCache?.delete(argsKeyOrderCacheKey);
+          }
         }
 
         const toolStatus = metadata.toolStatuses?.[toolCallId];
+        const mcpApp = extractMcpAppMetadata(
+          part,
+          metadata.mcpAppMetadataCache,
+        );
         return {
           type: "tool-call",
           toolName,
@@ -188,6 +261,8 @@ function convertParts(
           args,
           result,
           isError,
+          ...(modelContent !== undefined && { modelContent }),
+          ...(mcpApp && { mcp: { app: mcpApp } }),
           ...getToolInterrupt(part, toolStatus),
         } satisfies ToolCallMessagePart;
       }
@@ -198,7 +273,13 @@ function convertParts(
           sourceType: "url",
           id: part.sourceId,
           url: part.url,
-          title: part.title || "",
+          ...(part.title != null ? { title: part.title } : undefined),
+          ...(part.providerMetadata != null
+            ? {
+                providerMetadata:
+                  part.providerMetadata as SourceProviderMetadata,
+              }
+            : undefined),
         } satisfies SourceMessagePart;
       }
 
@@ -212,10 +293,20 @@ function convertParts(
       }
 
       if (part.type === "source-document") {
-        console.warn(
-          "Source document parts are not yet supported in conversion",
-        );
-        return null;
+        return {
+          type: "source",
+          sourceType: "document",
+          id: part.sourceId,
+          title: part.title,
+          mediaType: part.mediaType,
+          ...(part.filename != null ? { filename: part.filename } : undefined),
+          ...(part.providerMetadata != null
+            ? {
+                providerMetadata:
+                  part.providerMetadata as SourceProviderMetadata,
+              }
+            : undefined),
+        } satisfies SourceMessagePart;
       }
 
       if (part.type.startsWith("data-")) {