ai 5.0.0-canary.17 → 5.0.0-canary.18

package/dist/index.mjs

@@ -17,18 +17,6 @@ import {
 // core/util/process-chat-response.ts
 import { generateId as generateIdFunction } from "@ai-sdk/provider-utils";
 
-// core/types/duplicated/usage.ts
-function calculateLanguageModelUsage({
-  promptTokens,
-  completionTokens
-}) {
-  return {
-    promptTokens,
-    completionTokens,
-    totalTokens: promptTokens + completionTokens
-  };
-}
-
 // core/ui/get-tool-invocations.ts
 function getToolInvocations(message) {
   return message.parts.filter(
@@ -494,10 +482,13 @@ var finishMessageStreamPart = {
     const result = {
       finishReason: value.finishReason
     };
-    if ("usage" in value && value.usage != null && typeof value.usage === "object" && "promptTokens" in value.usage && "completionTokens" in value.usage) {
+    if ("usage" in value && value.usage != null && typeof value.usage === "object") {
       result.usage = {
-        promptTokens: typeof value.usage.promptTokens === "number" ? value.usage.promptTokens : Number.NaN,
-        completionTokens: typeof value.usage.completionTokens === "number" ? value.usage.completionTokens : Number.NaN
+        inputTokens: "inputTokens" in value.usage && typeof value.usage.inputTokens === "number" ? value.usage.inputTokens : void 0,
+        outputTokens: "outputTokens" in value.usage && typeof value.usage.outputTokens === "number" ? value.usage.outputTokens : void 0,
+        totalTokens: "totalTokens" in value.usage && typeof value.usage.totalTokens === "number" ? value.usage.totalTokens : void 0,
+        reasoningTokens: "reasoningTokens" in value.usage && typeof value.usage.reasoningTokens === "number" ? value.usage.reasoningTokens : void 0,
+        cachedInputTokens: "cachedInputTokens" in value.usage && typeof value.usage.cachedInputTokens === "number" ? value.usage.cachedInputTokens : void 0
       };
     }
     return {
@@ -519,10 +510,13 @@ var finishStepStreamPart = {
       finishReason: value.finishReason,
       isContinued: false
     };
-    if ("usage" in value && value.usage != null && typeof value.usage === "object" && "promptTokens" in value.usage && "completionTokens" in value.usage) {
+    if ("usage" in value && value.usage != null && typeof value.usage === "object") {
       result.usage = {
-        promptTokens: typeof value.usage.promptTokens === "number" ? value.usage.promptTokens : Number.NaN,
-        completionTokens: typeof value.usage.completionTokens === "number" ? value.usage.completionTokens : Number.NaN
+        inputTokens: "inputTokens" in value.usage && typeof value.usage.inputTokens === "number" ? value.usage.inputTokens : void 0,
+        outputTokens: "outputTokens" in value.usage && typeof value.usage.outputTokens === "number" ? value.usage.outputTokens : void 0,
+        totalTokens: "totalTokens" in value.usage && typeof value.usage.totalTokens === "number" ? value.usage.totalTokens : void 0,
+        reasoningTokens: "reasoningTokens" in value.usage && typeof value.usage.reasoningTokens === "number" ? value.usage.reasoningTokens : void 0,
+        cachedInputTokens: "cachedInputTokens" in value.usage && typeof value.usage.cachedInputTokens === "number" ? value.usage.cachedInputTokens : void 0
       };
     }
     if ("isContinued" in value && typeof value.isContinued === "boolean") {
@@ -795,9 +789,9 @@ async function processChatResponse({
   let messageAnnotations = replaceLastMessage ? lastMessage == null ? void 0 : lastMessage.annotations : void 0;
   const partialToolCalls = {};
   let usage = {
-    completionTokens: NaN,
-    promptTokens: NaN,
-    totalTokens: NaN
+    inputTokens: void 0,
+    outputTokens: void 0,
+    totalTokens: void 0
   };
   let finishReason = "unknown";
   function execUpdate() {
@@ -970,7 +964,7 @@ async function processChatResponse({
     onFinishMessagePart(value) {
       finishReason = value.finishReason;
       if (value.usage != null) {
-        usage = calculateLanguageModelUsage(value.usage);
+        usage = value.usage;
       }
     },
     onErrorPart(error) {
@@ -1027,7 +1021,11 @@ async function processChatTextResponse({
     }
   });
   onFinish == null ? void 0 : onFinish(resultMessage, {
-    usage: { completionTokens: NaN, promptTokens: NaN, totalTokens: NaN },
+    usage: {
+      inputTokens: void 0,
+      outputTokens: void 0,
+      totalTokens: void 0
+    },
     finishReason: "unknown"
   });
 }
@@ -1041,7 +1039,6 @@ async function callChatApi({
   credentials,
   headers,
   abortController,
-  restoreMessagesOnFailure,
   onResponse,
   onUpdate,
   onFinish,
@@ -1053,7 +1050,7 @@ async function callChatApi({
   requestType = "generate"
 }) {
   var _a17, _b, _c;
-  const request = requestType === "resume" ? fetch2(`${api}?chatId=${body.id}`, {
+  const response = requestType === "resume" ? await fetch2(`${api}?chatId=${body.id}`, {
     method: "GET",
     headers: {
       "Content-Type": "application/json",
@@ -1061,7 +1058,7 @@ async function callChatApi({
     },
     signal: (_a17 = abortController == null ? void 0 : abortController()) == null ? void 0 : _a17.signal,
     credentials
-  }) : fetch2(api, {
+  }) : await fetch2(api, {
     method: "POST",
     body: JSON.stringify(body),
     headers: {
@@ -1071,19 +1068,10 @@ async function callChatApi({
     signal: (_b = abortController == null ? void 0 : abortController()) == null ? void 0 : _b.signal,
     credentials
   });
-  const response = await request.catch((err) => {
-    restoreMessagesOnFailure();
-    throw err;
-  });
-  if (onResponse) {
-    try {
-      await onResponse(response);
-    } catch (err) {
-      throw err;
-    }
+  if (onResponse != null) {
+    await onResponse(response);
   }
   if (!response.ok) {
-    restoreMessagesOnFailure();
     throw new Error(
       (_c = await response.text()) != null ? _c : "Failed to fetch the chat response."
     );
@@ -1097,7 +1085,8 @@ async function callChatApi({
         stream: response.body,
         update: onUpdate,
         onFinish,
-        generateId: generateId3
+        generateId: generateId3,
+        getCurrentDate
       });
       return;
     }
@@ -1992,6 +1981,7 @@ function splitArray(array, chunkSize) {
 async function embedMany({
   model,
   values,
+  maxParallelCalls = Infinity,
   maxRetries: maxRetriesArg,
   abortSignal,
   headers,
@@ -2021,7 +2011,10 @@ async function embedMany({
     }),
     tracer,
     fn: async (span) => {
-      const maxEmbeddingsPerCall = await model.maxEmbeddingsPerCall;
+      const [maxEmbeddingsPerCall, supportsParallelCalls] = await Promise.all([
+        model.maxEmbeddingsPerCall,
+        model.supportsParallelCalls
+      ]);
       if (maxEmbeddingsPerCall == null || maxEmbeddingsPerCall === Infinity) {
         const { embeddings: embeddings2, usage, response } = await retry(() => {
           return recordSpan({
@@ -2092,61 +2085,69 @@ async function embedMany({
       const embeddings = [];
       const responses = [];
       let tokens = 0;
-      for (const chunk of valueChunks) {
-        const {
-          embeddings: responseEmbeddings,
-          usage,
-          response
-        } = await retry(() => {
-          return recordSpan({
-            name: "ai.embedMany.doEmbed",
-            attributes: selectTelemetryAttributes({
-              telemetry,
-              attributes: {
-                ...assembleOperationName({
-                  operationId: "ai.embedMany.doEmbed",
-                  telemetry
-                }),
-                ...baseTelemetryAttributes,
-                // specific settings that only make sense on the outer level:
-                "ai.values": {
-                  input: () => chunk.map((value) => JSON.stringify(value))
-                }
-              }
-            }),
-            tracer,
-            fn: async (doEmbedSpan) => {
-              var _a17;
-              const modelResponse = await model.doEmbed({
-                values: chunk,
-                abortSignal,
-                headers,
-                providerOptions
-              });
-              const embeddings2 = modelResponse.embeddings;
-              const usage2 = (_a17 = modelResponse.usage) != null ? _a17 : { tokens: NaN };
-              doEmbedSpan.setAttributes(
-                selectTelemetryAttributes({
+      const parallelChunks = splitArray(
+        valueChunks,
+        supportsParallelCalls ? maxParallelCalls : 1
+      );
+      for (const parallelChunk of parallelChunks) {
+        const results = await Promise.all(
+          parallelChunk.map((chunk) => {
+            return retry(() => {
+              return recordSpan({
+                name: "ai.embedMany.doEmbed",
+                attributes: selectTelemetryAttributes({
                   telemetry,
                   attributes: {
-                    "ai.embeddings": {
-                      output: () => embeddings2.map((embedding) => JSON.stringify(embedding))
-                    },
-                    "ai.usage.tokens": usage2.tokens
+                    ...assembleOperationName({
+                      operationId: "ai.embedMany.doEmbed",
+                      telemetry
+                    }),
+                    ...baseTelemetryAttributes,
+                    // specific settings that only make sense on the outer level:
+                    "ai.values": {
+                      input: () => chunk.map((value) => JSON.stringify(value))
+                    }
                   }
-                })
-              );
-              return {
-                embeddings: embeddings2,
-                usage: usage2,
-                response: modelResponse.response
-              };
-            }
-          });
-        });
-        embeddings.push(...responseEmbeddings);
-        responses.push(response);
-        tokens += usage.tokens;
+                }),
+                tracer,
+                fn: async (doEmbedSpan) => {
+                  var _a17;
+                  const modelResponse = await model.doEmbed({
+                    values: chunk,
+                    abortSignal,
+                    headers,
+                    providerOptions
+                  });
+                  const embeddings2 = modelResponse.embeddings;
+                  const usage = (_a17 = modelResponse.usage) != null ? _a17 : { tokens: NaN };
+                  doEmbedSpan.setAttributes(
+                    selectTelemetryAttributes({
+                      telemetry,
+                      attributes: {
+                        "ai.embeddings": {
+                          output: () => embeddings2.map(
+                            (embedding) => JSON.stringify(embedding)
+                          )
+                        },
+                        "ai.usage.tokens": usage.tokens
+                      }
+                    })
+                  );
+                  return {
+                    embeddings: embeddings2,
+                    usage,
+                    response: modelResponse.response
+                  };
+                }
+              });
+            });
+          })
+        );
+        for (const result of results) {
+          embeddings.push(...result.embeddings);
+          responses.push(result.response);
+          tokens += result.usage.tokens;
+        }
       }
       span.setAttributes(
         selectTelemetryAttributes({
@@ -3434,25 +3435,6 @@ async function standardizePrompt({
   throw new Error("unreachable");
 }
 
-// core/types/usage.ts
-function calculateLanguageModelUsage2({
-  inputTokens,
-  outputTokens
-}) {
-  return {
-    promptTokens: inputTokens != null ? inputTokens : NaN,
-    completionTokens: outputTokens != null ? outputTokens : NaN,
-    totalTokens: (inputTokens != null ? inputTokens : 0) + (outputTokens != null ? outputTokens : 0)
-  };
-}
-function addLanguageModelUsage(usage1, usage2) {
-  return {
-    promptTokens: usage1.promptTokens + usage2.promptTokens,
-    completionTokens: usage1.completionTokens + usage2.completionTokens,
-    totalTokens: usage1.totalTokens + usage2.totalTokens
-  };
-}
-
 // core/generate-object/output-strategy.ts
 import {
   isJSONArray,
@@ -3990,7 +3972,7 @@ async function generateObject(options) {
               throw new NoObjectGeneratedError({
                 message: "No object generated: the model did not return a response.",
                 response: responseData,
-                usage: calculateLanguageModelUsage2(result2.usage),
+                usage: result2.usage,
                 finishReason: result2.finishReason
               });
             }
@@ -4034,7 +4016,7 @@ async function generateObject(options) {
             cause: parseResult.error,
             text: result2,
             response,
-            usage: calculateLanguageModelUsage2(usage),
+            usage,
             finishReason
           });
         }
@@ -4043,7 +4025,7 @@ async function generateObject(options) {
           {
             text: result2,
             response,
-            usage: calculateLanguageModelUsage2(usage)
+            usage
           }
         );
         if (!validationResult.success) {
@@ -4052,7 +4034,7 @@ async function generateObject(options) {
             cause: validationResult.error,
             text: result2,
             response,
-            usage: calculateLanguageModelUsage2(usage),
+            usage,
             finishReason
           });
         }
@@ -4092,7 +4074,7 @@ async function generateObject(options) {
       return new DefaultGenerateObjectResult({
         object: object2,
         finishReason,
-        usage: calculateLanguageModelUsage2(usage),
+        usage,
         warnings,
         request,
         response,
@@ -4474,7 +4456,11 @@ var DefaultStreamObjectResult = class {
         );
         self.requestPromise.resolve(request != null ? request : {});
         let warnings;
-        let usage;
+        let usage = {
+          inputTokens: void 0,
+          outputTokens: void 0,
+          totalTokens: void 0
+        };
         let finishReason;
         let providerMetadata;
         let object2;
@@ -4554,7 +4540,7 @@ var DefaultStreamObjectResult = class {
                     controller.enqueue({ type: "text-delta", textDelta });
                   }
                   finishReason = chunk.finishReason;
-                  usage = calculateLanguageModelUsage2(chunk.usage);
+                  usage = chunk.usage;
                   providerMetadata = chunk.providerMetadata;
                   controller.enqueue({
                     ...chunk,
@@ -4616,14 +4602,17 @@ var DefaultStreamObjectResult = class {
                       "ai.response.id": fullResponse.id,
                       "ai.response.model": fullResponse.modelId,
                       "ai.response.timestamp": fullResponse.timestamp.toISOString(),
-                      "ai.usage.promptTokens": finalUsage.promptTokens,
-                      "ai.usage.completionTokens": finalUsage.completionTokens,
+                      "ai.usage.inputTokens": finalUsage.inputTokens,
+                      "ai.usage.outputTokens": finalUsage.outputTokens,
+                      "ai.usage.totalTokens": finalUsage.totalTokens,
+                      "ai.usage.reasoningTokens": finalUsage.reasoningTokens,
+                      "ai.usage.cachedInputTokens": finalUsage.cachedInputTokens,
                       // standardized gen-ai llm span attributes:
                       "gen_ai.response.finish_reasons": [finishReason],
                       "gen_ai.response.id": fullResponse.id,
                       "gen_ai.response.model": fullResponse.modelId,
-                      "gen_ai.usage.input_tokens": finalUsage.promptTokens,
-                      "gen_ai.usage.output_tokens": finalUsage.completionTokens
+                      "gen_ai.usage.input_tokens": finalUsage.inputTokens,
+                      "gen_ai.usage.output_tokens": finalUsage.outputTokens
                     }
                   })
                 );
@@ -4632,8 +4621,11 @@ var DefaultStreamObjectResult = class {
                   selectTelemetryAttributes({
                     telemetry,
                     attributes: {
-                      "ai.usage.promptTokens": finalUsage.promptTokens,
-                      "ai.usage.completionTokens": finalUsage.completionTokens,
+                      "ai.usage.inputTokens": finalUsage.inputTokens,
+                      "ai.usage.outputTokens": finalUsage.outputTokens,
+                      "ai.usage.totalTokens": finalUsage.totalTokens,
+                      "ai.usage.reasoningTokens": finalUsage.reasoningTokens,
+                      "ai.usage.cachedInputTokens": finalUsage.cachedInputTokens,
                       "ai.response.object": {
                         output: () => JSON.stringify(object2)
                       }
@@ -4863,6 +4855,26 @@ function prepareToolsAndToolChoice({
   };
 }
 
+// core/types/usage.ts
+function addLanguageModelUsage(usage1, usage2) {
+  return {
+    inputTokens: addTokenCounts(usage1.inputTokens, usage2.inputTokens),
+    outputTokens: addTokenCounts(usage1.outputTokens, usage2.outputTokens),
+    totalTokens: addTokenCounts(usage1.totalTokens, usage2.totalTokens),
+    reasoningTokens: addTokenCounts(
+      usage1.reasoningTokens,
+      usage2.reasoningTokens
+    ),
+    cachedInputTokens: addTokenCounts(
+      usage1.cachedInputTokens,
+      usage2.cachedInputTokens
+    )
+  };
+}
+function addTokenCounts(tokenCount1, tokenCount2) {
+  return tokenCount1 == null && tokenCount2 == null ? void 0 : (tokenCount1 != null ? tokenCount1 : 0) + (tokenCount2 != null ? tokenCount2 : 0);
+}
+
 // core/util/split-on-last-whitespace.ts
 var lastWhitespaceRegexp = /^([\s\S]*?)(\s+)(\S*)$/;
 function splitOnLastWhitespace(text2) {
@@ -5187,9 +5199,9 @@ async function generateText({
       const sources = [];
       const steps = [];
       let usage = {
-        completionTokens: 0,
-        promptTokens: 0,
-        totalTokens: 0
+        inputTokens: void 0,
+        outputTokens: void 0,
+        totalTokens: void 0
       };
       let stepType = "initial";
       do {
@@ -5330,10 +5342,7 @@ async function generateText({
           messages: stepInputMessages,
           abortSignal
         });
-        const currentUsage = calculateLanguageModelUsage2(
-          currentModelResponse.usage
-        );
-        usage = addLanguageModelUsage(usage, currentUsage);
+        usage = addLanguageModelUsage(usage, currentModelResponse.usage);
         let nextStepType = "done";
         if (++stepCount < maxSteps) {
           if (continueSteps && currentModelResponse.finishReason === "length" && // only use continue when there are no tool calls:
@@ -5398,7 +5407,7 @@ async function generateText({
           toolCalls: currentToolCalls,
           toolResults: currentToolResults,
           finishReason: currentModelResponse.finishReason,
-          usage: currentUsage,
+          usage: currentModelResponse.usage,
           warnings: currentModelResponse.warnings,
           request: (_e = currentModelResponse.request) != null ? _e : {},
           response: {
@@ -5936,6 +5945,7 @@ function runToolsTransformation({
       const chunkType = chunk.type;
       switch (chunkType) {
         case "stream-start":
+        case "finish":
         case "text":
         case "reasoning":
         case "reasoning-part-finish":
@@ -6056,15 +6066,6 @@ function runToolsTransformation({
           }
           break;
         }
-        case "finish": {
-          finishChunk = {
-            type: "finish",
-            finishReason: chunk.finishReason,
-            usage: calculateLanguageModelUsage2(chunk.usage),
-            providerMetadata: chunk.providerMetadata
-          };
-          break;
-        }
         default: {
           const _exhaustiveCheck = chunkType;
           throw new Error(`Unhandled chunk type: ${_exhaustiveCheck}`);
@@ -6427,9 +6428,9 @@ var DefaultStreamTextResult = class {
           self.reasoningDetailsPromise.resolve(lastStep.reasoning);
           const finishReason = recordedFinishReason != null ? recordedFinishReason : "unknown";
           const usage = recordedUsage != null ? recordedUsage : {
-            completionTokens: NaN,
-            promptTokens: NaN,
-            totalTokens: NaN
+            inputTokens: void 0,
+            outputTokens: void 0,
+            totalTokens: void 0
           };
           self.finishReasonPromise.resolve(finishReason);
           self.usagePromise.resolve(usage);
@@ -6465,8 +6466,11 @@ var DefaultStreamTextResult = class {
                     return ((_a18 = lastStep.toolCalls) == null ? void 0 : _a18.length) ? JSON.stringify(lastStep.toolCalls) : void 0;
                   }
                 },
-                "ai.usage.promptTokens": usage.promptTokens,
-                "ai.usage.completionTokens": usage.completionTokens
+                "ai.usage.inputTokens": usage.inputTokens,
+                "ai.usage.outputTokens": usage.outputTokens,
+                "ai.usage.totalTokens": usage.totalTokens,
+                "ai.usage.reasoningTokens": usage.reasoningTokens,
+                "ai.usage.cachedInputTokens": usage.cachedInputTokens
               }
             })
           );
@@ -6631,9 +6635,9 @@ var DefaultStreamTextResult = class {
           let activeReasoningPart2 = void 0;
           let stepFinishReason = "unknown";
           let stepUsage = {
-            promptTokens: 0,
-            completionTokens: 0,
-            totalTokens: 0
+            inputTokens: void 0,
+            outputTokens: void 0,
+            totalTokens: void 0
           };
           let stepProviderMetadata;
           let stepFirstChunk = true;
@@ -6662,7 +6666,7 @@ var DefaultStreamTextResult = class {
             transformedStream.pipeThrough(
               new TransformStream({
                 async transform(chunk, controller) {
-                  var _a17, _b, _c;
+                  var _a17, _b, _c, _d;
                   if (chunk.type === "stream-start") {
                     warnings = chunk.warnings;
                     return;
@@ -6758,7 +6762,7 @@ var DefaultStreamTextResult = class {
                       doStreamSpan.addEvent("ai.stream.finish");
                       doStreamSpan.setAttributes({
                         "ai.response.msToFinish": msToFinish,
-                        "ai.response.avgCompletionTokensPerSecond": 1e3 * stepUsage.completionTokens / msToFinish
+                        "ai.response.avgOutputTokensPerSecond": 1e3 * ((_d = stepUsage.outputTokens) != null ? _d : 0) / msToFinish
                       });
                       break;
                     }
@@ -6821,14 +6825,17 @@ var DefaultStreamTextResult = class {
                           "ai.response.id": stepResponse.id,
                           "ai.response.model": stepResponse.modelId,
                           "ai.response.timestamp": stepResponse.timestamp.toISOString(),
-                          "ai.usage.promptTokens": stepUsage.promptTokens,
-                          "ai.usage.completionTokens": stepUsage.completionTokens,
+                          "ai.usage.inputTokens": stepUsage.inputTokens,
+                          "ai.usage.outputTokens": stepUsage.outputTokens,
+                          "ai.usage.totalTokens": stepUsage.totalTokens,
+                          "ai.usage.reasoningTokens": stepUsage.reasoningTokens,
+                          "ai.usage.cachedInputTokens": stepUsage.cachedInputTokens,
                           // standardized gen-ai llm span attributes:
                           "gen_ai.response.finish_reasons": [stepFinishReason],
                           "gen_ai.response.id": stepResponse.id,
                           "gen_ai.response.model": stepResponse.modelId,
-                          "gen_ai.usage.input_tokens": stepUsage.promptTokens,
-                          "gen_ai.usage.output_tokens": stepUsage.completionTokens
+                          "gen_ai.usage.input_tokens": stepUsage.inputTokens,
+                          "gen_ai.usage.output_tokens": stepUsage.outputTokens
                         }
                       })
                     );
@@ -6910,9 +6917,9 @@ var DefaultStreamTextResult = class {
           currentStep: 0,
           responseMessages: [],
           usage: {
-            promptTokens: 0,
-            completionTokens: 0,
-            totalTokens: 0
+            inputTokens: void 0,
+            outputTokens: void 0,
+            totalTokens: void 0
           },
           previousStepText: "",
           stepType: "initial",
@@ -7141,8 +7148,11 @@ var DefaultStreamTextResult = class {
                 formatDataStreamPart("finish_step", {
                   finishReason: chunk.finishReason,
                   usage: sendUsage ? {
-                    promptTokens: chunk.usage.promptTokens,
-                    completionTokens: chunk.usage.completionTokens
+                    inputTokens: chunk.usage.inputTokens,
+                    outputTokens: chunk.usage.outputTokens,
+                    totalTokens: chunk.usage.totalTokens,
+                    reasoningTokens: chunk.usage.reasoningTokens,
+                    cachedInputTokens: chunk.usage.cachedInputTokens
                   } : void 0,
                   isContinued: chunk.isContinued
                 })
@@ -7155,8 +7165,11 @@ var DefaultStreamTextResult = class {
                   formatDataStreamPart("finish_message", {
                     finishReason: chunk.finishReason,
                     usage: sendUsage ? {
-                      promptTokens: chunk.usage.promptTokens,
-                      completionTokens: chunk.usage.completionTokens
+                      inputTokens: chunk.usage.inputTokens,
+                      outputTokens: chunk.usage.outputTokens,
+                      totalTokens: chunk.usage.totalTokens,
+                      reasoningTokens: chunk.usage.reasoningTokens,
+                      cachedInputTokens: chunk.usage.cachedInputTokens
                     } : void 0
                   })
                 );
@@ -7808,7 +7821,6 @@ function appendResponseMessages({
         if (isLastMessageAssistant) {
           const maxStep = extractMaxToolInvocationStep(
             getToolInvocations(lastMessage)
-            // TODO remove once Message is removed
           );
           (_a17 = lastMessage.parts) != null ? _a17 : lastMessage.parts = [];
           lastMessage.content = textContent;
@@ -7845,10 +7857,9 @@ function appendResponseMessages({
         }
         (_b = lastMessage.parts) != null ? _b : lastMessage.parts = [];
         for (const contentPart of message.content) {
-          const toolCall = getToolInvocations(
-            lastMessage
-            // TODO remove once Message is removed
-          ).find((call) => call.toolCallId === contentPart.toolCallId);
+          const toolCall = getToolInvocations(lastMessage).find(
+            (call) => call.toolCallId === contentPart.toolCallId
+          );
           const toolCallPart = lastMessage.parts.find(
             (part) => part.type === "tool-invocation" && part.toolInvocation.toolCallId === contentPart.toolCallId
           );