pass tool values through more places

Author: ianmacartney

src/UIMessages.ts

@@ -381,10 +381,9 @@ function createAssistantUIMessage<
   const allParts: UIMessage<METADATA, DATA_PARTS, TOOLS>["parts"] = [];
 
   for (const message of group) {
-    const coreMessage = message.message && toModelMessage(message.message);
-    if (!coreMessage) continue;
+    if (!message.message) continue;
 
-    const content = coreMessage.content;
+    const content = message.message.content;
     const nonStringContent =
       content && typeof content !== "string" ? content : [];
     const text = extractTextFromMessageDoc(message);
@@ -441,12 +440,11 @@ function createAssistantUIMessage<
             type: "step-start",
           } satisfies StepStartUIPart);
           const toolPart: ToolUIPart<TOOLS> = {
+            ...omit(contentPart, ["args", "type"]),
             type: `tool-${contentPart.toolName as keyof TOOLS & string}`,
-            toolCallId: contentPart.toolCallId,
-            input: contentPart.input as DeepPartial<
+            input: contentPart.args as DeepPartial<
               TOOLS[keyof TOOLS & string]["input"]
             >,
-            providerExecuted: contentPart.providerExecuted,
             ...(message.streaming
               ? { state: "input-streaming" }
               : {
@@ -478,25 +476,22 @@ function createAssistantUIMessage<
               call.output = output;
             }
           } else {
-            console.warn(
-              "Tool result without preceding tool call.. adding anyways",
-              contentPart,
-            );
             if (message.error) {
               allParts.push({
+                input: contentPart.args,
+                ...omit(contentPart, ["args", "type", "output"]),
                 type: `tool-${contentPart.toolName}`,
-                toolCallId: contentPart.toolCallId,
                 state: "output-error",
-                input: undefined,
+                output,
                 errorText: message.error,
                 callProviderMetadata: message.providerMetadata,
               } satisfies ToolUIPart<TOOLS>);
             } else {
               allParts.push({
+                input: contentPart.args,
+                ...omit(contentPart, ["args", "type", "output"]),
                 type: `tool-${contentPart.toolName}`,
-                toolCallId: contentPart.toolCallId,
                 state: "output-available",
-                input: undefined,
                 output,
                 callProviderMetadata: message.providerMetadata,
               } satisfies ToolUIPart<TOOLS>);

src/mapping.ts

@@ -54,7 +54,7 @@ import {
   getProviderName,
   type ModelOrMetadata,
 } from "./shared.js";
-import { pick } from "convex-helpers";
+import { omit, pick } from "convex-helpers";
 export type AIMessageWithoutId = Omit<AIMessage, "id">;
 
 export type SerializeUrlsAndUint8Arrays<T> = T extends URL
@@ -102,9 +102,7 @@ export function fromModelMessage(message: ModelMessage): Message {
   return {
     role: message.role,
     content,
-    ...(message.providerOptions
-      ? { providerOptions: message.providerOptions }
-      : {}),
+    ...pick(message, ["providerOptions"]),
   } as SerializedMessage;
 }
 
@@ -121,9 +119,7 @@ export async function serializeOrThrow(
   return {
     role: message.role,
     content,
-    ...(message.providerOptions
-      ? { providerOptions: message.providerOptions }
-      : {}),
+    ...pick(message, ["providerOptions"]),
   } as SerializedMessage;
 }
 
@@ -613,20 +609,20 @@ export function toModelMessageContent(
           mediaType: getMimeOrMediaType(part)!,
           ...metadata,
         } satisfies FilePart;
-      case "tool-call": {
-        const input = "input" in part ? part.input : part.args;
+      case "tool-call":
         return {
-          type: part.type,
-          input: input ?? null,
-          toolCallId: part.toolCallId,
-          toolName: part.toolName,
-          providerExecuted: part.providerExecuted,
-          ...metadata,
+          input: ("input" in part ? part.input : part.args) ?? null,
+          ...omit(part as Infer<typeof vToolCallPart>, ["args"]),
         } satisfies ToolCallPart;
-      }
-      case "tool-result": {
-        return normalizeToolResult(part, metadata);
-      }
+      case "tool-result":
+        return {
+          input: (part as Infer<typeof vToolResultPart>).args,
+          output: normalizeToolOutput(
+            (part as Infer<typeof vToolResultPart>).result,
+          ),
+          ...omit(part as Infer<typeof vToolResultPart>, ["result", "args"]),
+          ...metadata,
+        } satisfies ToolResultPart & { input: unknown };
       case "reasoning":
         return {
           type: part.type,
@@ -684,16 +680,13 @@ function normalizeToolResult(
     providerOptions?: ProviderOptions;
     providerMetadata?: ProviderMetadata;
   },
-): ToolResultPart & Infer<typeof vToolResultPart> {
+): ToolResultPart & Infer<typeof vToolResultPart> & { input: unknown } {
   return {
-    type: part.type,
-    output:
-      part.output ??
-      normalizeToolOutput("result" in part ? part.result : undefined),
-    toolCallId: part.toolCallId,
-    toolName: part.toolName,
+    input: (part as Infer<typeof vToolResultPart>).args,
+    output: normalizeToolOutput("result" in part ? part.result : undefined),
+    ...omit(part as Infer<typeof vToolResultPart>, ["result", "args"]),
     ...metadata,
-  } satisfies ToolResultPart;
+  } satisfies ToolResultPart & { input: unknown };
 }
 
 /**
@@ -804,16 +797,24 @@ export function toModelMessageDataOrUrl(
   return urlOrString;
 }
 
-export function toUIFilePart(part: ImagePart | FilePart): FileUIPart {
+export function toUIFilePart(
+  part:
+    | ImagePart
+    | FilePart
+    | Infer<typeof vImagePart>
+    | Infer<typeof vFilePart>,
+): FileUIPart {
   const dataOrUrl = part.type === "image" ? part.image : part.data;
   const url =
     dataOrUrl instanceof ArrayBuffer
       ? convertUint8ArrayToBase64(new Uint8Array(dataOrUrl))
       : dataOrUrl.toString();
 
+  const mediaType = getMimeOrMediaType(part);
+
   return {
     type: "file",
-    mediaType: part.mediaType!,
+    mediaType: mediaType!,
     filename: part.type === "file" ? part.filename : undefined,
     url,
     providerMetadata: part.providerOptions,