ai 7.0.124 → 7.0.127

package/src/agent/pipe-agent-ui-stream-to-response.ts

@@ -15,8 +15,9 @@ import type { UIMessageStreamOptions } from '../generate-text/stream-text-result
 import type { TimeoutConfiguration } from '../prompt/request-options';
 import { pipeUIMessageStreamToResponse } from '../ui-message-stream';
 import type { UIMessageStreamResponseInit } from '../ui-message-stream/ui-message-stream-response-init';
-import type { InferUITools, UIMessage } from '../ui/ui-messages';
+import type { InferUITools, UIMessage, UIDataTypes } from '../ui/ui-messages';
 import type { Agent } from './agent';
+import type { convertToModelMessages } from '../ui/convert-to-model-messages';
 import { createAgentUIStream } from './create-agent-ui-stream';
 
 /**
@@ -25,6 +26,7 @@ import { createAgentUIStream } from './create-agent-ui-stream';
  * @param response - The Node.js ServerResponse object to pipe to.
  * @param agent - The agent to run.
  * @param uiMessages - The input UI messages.
+ * @param convertDataPart - Optional function to convert custom data parts to text or file model message parts.
  * @param abortSignal - Abort signal. Optional.
  * @param timeout - Timeout in milliseconds. Optional.
  * @param experimental_sandbox - The sandbox environment that is passed through to tool execution. Optional.
@@ -44,6 +46,11 @@ export async function pipeAgentUIStreamToResponse<
   RUNTIME_CONTEXT extends Context = Context,
   OUTPUT extends Output = never,
   MESSAGE_METADATA = unknown,
+  UI_MESSAGE extends UIMessage<
+    MESSAGE_METADATA,
+    UIDataTypes,
+    InferUITools<TOOLS>
+  > = UIMessage<MESSAGE_METADATA, UIDataTypes, InferUITools<TOOLS>>,
 >({
   response,
   headers,
@@ -56,6 +63,13 @@ export async function pipeAgentUIStreamToResponse<
   response: ServerResponse;
   agent: Agent<CALL_OPTIONS, TOOLS, RUNTIME_CONTEXT, OUTPUT>;
   uiMessages: unknown[];
+  /**
+   * Converts custom UI data parts to text or file model message parts.
+   * Data parts are ignored when omitted or when the callback returns undefined.
+   */
+  convertDataPart?: NonNullable<
+    Parameters<typeof convertToModelMessages<UI_MESSAGE>>[1]
+  >['convertDataPart'];
   abortSignal?: AbortSignal;
   timeout?: TimeoutConfiguration<TOOLS>;
   experimental_sandbox?: SandboxSession;
@@ -65,9 +79,7 @@ export async function pipeAgentUIStreamToResponse<
   /** @deprecated Use `onStepEnd` instead. */
   onStepFinish?: GenerateTextOnStepFinishCallback<TOOLS>;
 } & UIMessageStreamResponseInit &
-  UIMessageStreamOptions<
-    UIMessage<MESSAGE_METADATA, never, InferUITools<TOOLS>>
-  >): Promise<void> {
+  UIMessageStreamOptions<UI_MESSAGE>): Promise<void> {
   return pipeUIMessageStreamToResponse({
     response,
     headers,

package/src/generate-text/validate-tool-approvals.ts

@@ -117,9 +117,13 @@ export async function validateApprovedToolApprovals<
             refineToolInput,
           });
 
+          // Clone both inputs into the current realm before comparing them.
+          const revalidatedInput = structuredClone(revalidatedToolCall.input);
+          const approvedInput = structuredClone(toolCall.input);
+
           // Revalidation must never change the operation that was approved,
           // including when older or projected history omits the schema input.
-          if (!isDeepEqualData(revalidatedToolCall.input, toolCall.input)) {
+          if (!isDeepEqualData(revalidatedInput, approvedInput)) {
             validationError = new Error(
               'Approved tool input does not match the validated schema output.',
             );

package/src/tool-search/prepare-tool-search.ts

@@ -11,7 +11,11 @@ import {
 } from '../generate-text/tool-caller-configuration';
 import { resolveToolDescription } from '../prompt/prepare-tools';
 import { getOwn } from '../util/get-own';
-import { isToolSearch } from './tool-search';
+import {
+  getToolSearchFunction,
+  getToolSearchMaxResults,
+  isToolSearch,
+} from './tool-search';
 
 /** Create discovery state for one generation, never for a shared tool instance. */
 export function createToolSearchState({
@@ -90,41 +94,57 @@ export function createToolSearchState({
             searchName,
             {
               ...tool,
-              execute: ({ query }: { query: string }) => {
+              execute: async ({ query }: { query: string }) => {
+                const availableTools = candidates.map(([name, candidate]) => {
+                  const description = resolveToolDescription({
+                    tool: candidate,
+                    toolName: name,
+                    toolsContext,
+                    experimental_sandbox,
+                  });
+                  return {
+                    name,
+                    ...(description == null ? {} : { description }),
+                  };
+                });
+                const toolsByName = new Map(
+                  availableTools.map(tool => [tool.name, tool]),
+                );
+                const customSearch = getToolSearchFunction(tool);
                 const terms = [...new Set(tokenize(query))];
-                const matches = candidates
-                  .map(([name, candidate]) => {
-                    const description = resolveToolDescription({
-                      tool: candidate,
-                      toolName: name,
-                      toolsContext,
-                      experimental_sandbox,
-                    });
-                    const nameTerms = tokenize(name);
-                    const descriptionTerms = tokenize(description ?? '');
-                    const score = terms.reduce(
-                      (score, term) =>
-                        score +
-                        (nameTerms.includes(term) ? 2 : 0) +
-                        (descriptionTerms.includes(term) ? 1 : 0),
-                      0,
-                    );
-                    return { name, description, score };
+                const rankedNames = customSearch
+                  ? await customSearch({
+                      query,
+                      tools: availableTools.map(tool => ({ ...tool })),
+                    })
+                  : availableTools
+                      .map(({ name, description }) => {
+                        const nameTerms = tokenize(name);
+                        const descriptionTerms = tokenize(description ?? '');
+                        const score = terms.reduce(
+                          (score, term) =>
+                            score +
+                            (nameTerms.includes(term) ? 2 : 0) +
+                            (descriptionTerms.includes(term) ? 1 : 0),
+                          0,
+                        );
+                        return { name, score };
+                      })
+                      .filter(match => match.score > 0)
+                      .sort((a, b) => b.score - a.score)
+                      .map(match => match.name);
+                const matches = [...new Set(rankedNames)]
+                  .flatMap(name => {
+                    const match = toolsByName.get(name);
+                    return match == null ? [] : [match];
                   })
-                  .filter(match => match.score > 0)
-                  .sort((a, b) => b.score - a.score)
-                  .slice(0, 5);
+                  .slice(0, getToolSearchMaxResults(tool));
 
                 for (const { name } of matches) {
                   discovered.add(name);
                 }
 
-                return {
-                  tools: matches.map(({ name, description }) => ({
-                    name,
-                    ...(description == null ? {} : { description }),
-                  })),
-                };
+                return { tools: matches };
               },
             },
           ];

package/src/tool-search/tool-search.ts

@@ -1,6 +1,16 @@
 import { jsonSchema, tool, type Tool } from '@ai-sdk/provider-utils';
+import { InvalidArgumentError } from '../error/invalid-argument-error';
 
 const toolSearchSymbol = Symbol.for('vercel.ai.toolSearch');
+const toolSearchFunctionSymbol = Symbol.for('vercel.ai.toolSearch.search');
+const toolSearchMaxResultsSymbol = Symbol.for(
+  'vercel.ai.toolSearch.maxResults',
+);
+
+type ToolSearchFunction = (options: {
+  query: string;
+  tools: Array<{ name: string; description?: string }>;
+}) => string[] | PromiseLike<string[]>;
 
 type ToolSearchInput = { query: string };
 type ToolSearchOutput = {
@@ -12,13 +22,36 @@ type ToolSearchOutput = {
  * binds the search registry and makes matches available on the next step.
  * Use directly or through code mode with `toolDiscovery: 'conversation'`.
  */
-export function toolSearch(): Tool<ToolSearchInput, ToolSearchOutput> & {
+export function toolSearch({
+  search,
+  maxResults = 5,
+}: {
+  /**
+   * Search eligible deferred tools and return their names in ranked order.
+   * Supports synchronous and asynchronous callbacks. Unknown names and
+   * duplicates are ignored before applying the result limit.
+   */
+  search?: ToolSearchFunction;
+  /**
+   * Maximum number of matching tools returned per search.
+   *
+   * @default 5
+   */
+  maxResults?: number;
+} = {}): Tool<ToolSearchInput, ToolSearchOutput> & {
   type: 'function';
 } {
+  if (!Number.isSafeInteger(maxResults) || maxResults < 1) {
+    throw new InvalidArgumentError({
+      parameter: 'maxResults',
+      value: maxResults,
+      message: 'maxResults must be a positive safe integer.',
+    });
+  }
+
   return Object.assign(
     tool({
-      description:
-        'Search for tools by keywords in their names and descriptions. Returns up to five matching tools. Matches become available on the next model step, after this execution finishes. Wait for their tool definitions before calling the discovered tools. If no tools match, try different keywords.',
+      description: `${search == null ? 'Search for tools by keywords in their names and descriptions.' : 'Search for tools matching a query.'} Returns up to ${maxResults === 5 ? 'five' : maxResults} matching tools. Matches become available on the next model step, after this execution finishes. Wait for their tool definitions before calling the discovered tools. If no tools match, try different keywords.`,
       inputSchema: jsonSchema<ToolSearchInput>({
         type: 'object',
         properties: { query: { type: 'string', minLength: 1 } },
@@ -48,10 +81,31 @@ export function toolSearch(): Tool<ToolSearchInput, ToolSearchOutput> & {
         throw new Error('toolSearch must be bound by an AI SDK generation.');
       },
     }),
-    { type: 'function' as const, [toolSearchSymbol]: true },
+    {
+      type: 'function' as const,
+      [toolSearchSymbol]: true,
+      [toolSearchFunctionSymbol]: search,
+      [toolSearchMaxResultsSymbol]: maxResults,
+    },
   );
 }
 
+export function getToolSearchMaxResults(tool: Tool): number {
+  return (
+    (tool as Tool & { [toolSearchMaxResultsSymbol]?: number })[
+      toolSearchMaxResultsSymbol
+    ] ?? 5
+  );
+}
+
+export function getToolSearchFunction(
+  tool: Tool,
+): ToolSearchFunction | undefined {
+  return (tool as Tool & { [toolSearchFunctionSymbol]?: ToolSearchFunction })[
+    toolSearchFunctionSymbol
+  ];
+}
+
 export function isToolSearch(tool: Tool): boolean {
   return (
     (tool as Tool & { [toolSearchSymbol]?: boolean })[toolSearchSymbol] === true

package/src/ui/chat.ts

@@ -609,7 +609,7 @@ export abstract class AbstractChat<UI_MESSAGE extends UIMessage> {
     });
 
   addToolOutput: ChatAddToolOutputFunction<UI_MESSAGE> = async ({
-    state = 'output-available',
+    state,
     toolCallId,
     output,
     errorText,
@@ -621,10 +621,41 @@ export abstract class AbstractChat<UI_MESSAGE extends UIMessage> {
 
       const updatePart = (
         part: UIMessagePart<UIDataTypes, UITools>,
-      ): UIMessagePart<UIDataTypes, UITools> =>
-        isToolUIPart(part) && part.toolCallId === toolCallId
-          ? ({ ...part, state, output, errorText } as typeof part)
-          : part;
+      ): UIMessagePart<UIDataTypes, UITools> => {
+        if (!isToolUIPart(part) || part.toolCallId !== toolCallId) {
+          return part;
+        }
+
+        // Output states can only retain approvals that were granted.
+        const { approval: existingApproval, ...toolPart } = part;
+        const approval =
+          existingApproval?.approved === true
+            ? {
+                approval: {
+                  ...existingApproval,
+                  approved: existingApproval.approved,
+                },
+              }
+            : {};
+
+        return state === 'output-error'
+          ? {
+              ...toolPart,
+              ...approval,
+              state,
+              input: part.input,
+              output: undefined,
+              errorText,
+            }
+          : {
+              ...toolPart,
+              ...approval,
+              state: 'output-available',
+              input: part.input,
+              output,
+              errorText: undefined,
+            };
+      };
 
       // update the message to trigger an immediate UI update
       this.state.replaceMessage(messages.length - 1, {

package/src/ui-message-stream/create-ui-message-stream.ts

@@ -73,12 +73,25 @@ export function createUIMessageStream<UI_MESSAGE extends UIMessage>({
   >;
 
   const ongoingStreamPromises: Promise<void>[] = [];
+  const activeReaders = new Set<
+    ReadableStreamDefaultReader<InferUIMessageChunk<UI_MESSAGE>>
+  >();
+  let isCancelled = false;
+  let cancelReason: unknown;
   let outcome: UIMessageStreamOutcome = { status: 'unknown' };
 
   const stream = new ReadableStream({
     start(controllerArg) {
       controller = controllerArg;
     },
+    async cancel(reason) {
+      isCancelled = true;
+      cancelReason = reason;
+
+      await Promise.all(
+        Array.from(activeReaders, reader => reader.cancel(reason)),
+      );
+    },
   });
 
   function safeEnqueue(data: InferUIMessageChunk<UI_MESSAGE>) {
@@ -132,13 +145,25 @@ export function createUIMessageStream<UI_MESSAGE extends UIMessage>({
           safeEnqueue(part);
         },
         merge(streamArg) {
+          if (isCancelled) {
+            void streamArg.cancel(cancelReason).catch(() => {});
+            return;
+          }
+
           ongoingStreamPromises.push(
             (async () => {
               const reader = streamArg.getReader();
-              while (true) {
-                const { done, value } = await reader.read();
-                if (done) break;
-                safeEnqueue(value);
+              activeReaders.add(reader);
+
+              try {
+                while (true) {
+                  const { done, value } = await reader.read();
+                  if (done) break;
+                  safeEnqueue(value);
+                }
+              } finally {
+                activeReaders.delete(reader);
+                reader.releaseLock();
               }
             })().catch(error => {
               handleError(error);