ai 7.0.97 → 7.0.98

package/src/rerank/rerank.ts

@@ -1,12 +1,18 @@
-import type { JSONObject, RerankingModelV4CallOptions } from '@ai-sdk/provider';
+import {
+  InvalidResponseDataError,
+  type JSONObject,
+  type RerankingModelV4CallOptions,
+  type RerankingModelV4Result,
+} from '@ai-sdk/provider';
 import {
   createIdGenerator,
+  type Context,
   type ProviderOptions,
 } from '@ai-sdk/provider-utils';
 import { prepareRetries } from '../../src/util/prepare-retries';
 import { logWarnings } from '../logger/log-warnings';
 import { resolveRerankingModel } from '../model/resolve-model';
-import { createTelemetryDispatcher } from '../telemetry/create-telemetry-dispatcher';
+import { createRestrictedTelemetryDispatcher } from './restricted-telemetry-dispatcher';
 import type { TelemetryOptions } from '../telemetry/telemetry-options';
 import type { RerankingModel } from '../types';
 import type { Callback } from '../util/callback';
@@ -32,10 +38,14 @@ const originalGenerateCallId = createIdGenerator({
  * @param headers - Additional HTTP headers to be sent with the request. Only applicable for HTTP-based providers.
  * @param providerOptions - Additional provider-specific options.
  * @param telemetry - Optional telemetry configuration.
+ * @param runtimeContext - User-defined runtime context passed to callbacks and, when explicitly included, telemetry.
  *
  * @returns A result object that contains the reranked documents, the reranked indices, and additional information.
  */
-export async function rerank<VALUE extends JSONObject | string>({
+export async function rerank<
+  VALUE extends JSONObject | string,
+  RUNTIME_CONTEXT extends Context = Context,
+>({
   model: modelArg,
   documents,
   query,
@@ -46,6 +56,7 @@ export async function rerank<VALUE extends JSONObject | string>({
   providerOptions,
   experimental_telemetry,
   telemetry = experimental_telemetry,
+  runtimeContext = {} as RUNTIME_CONTEXT,
   onStart,
   experimental_onStart,
   onEnd,
@@ -93,14 +104,19 @@ export async function rerank<VALUE extends JSONObject | string>({
   /**
    * Optional telemetry configuration.
    */
-  telemetry?: TelemetryOptions;
+  telemetry?: TelemetryOptions<RUNTIME_CONTEXT>;
 
   /**
    * Optional telemetry configuration.
    *
    * @deprecated Use `telemetry` instead. This alias will be removed in a future major release.
    */
-  experimental_telemetry?: TelemetryOptions;
+  experimental_telemetry?: TelemetryOptions<RUNTIME_CONTEXT>;
+
+  /**
+   * User-defined runtime context. Treat runtime context as immutable.
+   */
+  runtimeContext?: RUNTIME_CONTEXT;
 
   /**
    * Additional provider-specific options. They are passed through
@@ -113,7 +129,7 @@ export async function rerank<VALUE extends JSONObject | string>({
    * Callback that is called when the rerank operation begins,
    * before the reranking model is called.
    */
-  onStart?: Callback<RerankStartEvent>;
+  onStart?: Callback<RerankStartEvent<RUNTIME_CONTEXT>>;
 
   /**
    * Callback that is called when the rerank operation begins,
@@ -121,13 +137,13 @@ export async function rerank<VALUE extends JSONObject | string>({
    *
    * @deprecated Use `onStart` instead.
    */
-  experimental_onStart?: Callback<RerankStartEvent>;
+  experimental_onStart?: Callback<RerankStartEvent<RUNTIME_CONTEXT>>;
 
   /**
    * Callback that is called when the rerank operation completes,
    * after the reranking model returns.
    */
-  onEnd?: Callback<RerankEndEvent>;
+  onEnd?: Callback<RerankEndEvent<RUNTIME_CONTEXT>>;
 
   /**
    * Callback that is called when the rerank operation completes,
@@ -135,7 +151,7 @@ export async function rerank<VALUE extends JSONObject | string>({
    *
    * @deprecated Use `onEnd` instead.
    */
-  experimental_onEnd?: Callback<RerankEndEvent>;
+  experimental_onEnd?: Callback<RerankEndEvent<RUNTIME_CONTEXT>>;
 
   /**
    * Internal. For test use only. May change without notice.
@@ -149,7 +165,7 @@ export async function rerank<VALUE extends JSONObject | string>({
   const resolvedOnStart = onStart ?? experimental_onStart;
   const resolvedOnEnd = onEnd ?? experimental_onEnd;
 
-  const telemetryDispatcher = createTelemetryDispatcher({
+  const telemetryDispatcher = createRestrictedTelemetryDispatcher({
     telemetry,
   });
 
@@ -163,6 +179,7 @@ export async function rerank<VALUE extends JSONObject | string>({
       event: {
         callId,
         operationId: 'ai.rerank',
+        runtimeContext,
         provider: model.provider,
         modelId: model.modelId,
         documents,
@@ -179,6 +196,7 @@ export async function rerank<VALUE extends JSONObject | string>({
       event: {
         callId,
         operationId: 'ai.rerank',
+        runtimeContext,
         provider: model.provider,
         modelId: model.modelId,
         documents,
@@ -218,6 +236,7 @@ export async function rerank<VALUE extends JSONObject | string>({
   const startEvent = {
     callId,
     operationId: 'ai.rerank',
+    runtimeContext,
     provider: model.provider,
     modelId: model.modelId,
     documents,
@@ -286,6 +305,8 @@ export async function rerank<VALUE extends JSONObject | string>({
           },
         );
 
+        validateRankingIndices({ ranking, documents });
+
         logWarnings({
           warnings: warnings ?? [],
           provider: model.provider,
@@ -296,14 +317,15 @@ export async function rerank<VALUE extends JSONObject | string>({
           event: {
             callId,
             operationId: 'ai.rerank',
+            runtimeContext,
             provider: model.provider,
             modelId: model.modelId,
             documents,
             query,
-            ranking: ranking.map(r => ({
-              originalIndex: r.index,
-              score: r.relevanceScore,
-              document: documents[r.index],
+            ranking: ranking.map(ranking => ({
+              originalIndex: ranking.index,
+              score: ranking.relevanceScore,
+              document: documents[ranking.index],
             })),
             warnings: warnings ?? [],
             providerMetadata,
@@ -342,6 +364,23 @@ export async function rerank<VALUE extends JSONObject | string>({
   });
 }
 
+function validateRankingIndices<VALUE>({
+  ranking,
+  documents,
+}: {
+  ranking: RerankingModelV4Result['ranking'];
+  documents: Array<VALUE>;
+}) {
+  for (const { index } of ranking) {
+    if (!Number.isInteger(index) || index < 0 || index >= documents.length) {
+      throw new InvalidResponseDataError({
+        data: ranking,
+        message: `Invalid ranking index ${index}. Expected an integer between 0 and ${documents.length - 1}.`,
+      });
+    }
+  }
+}
+
 class DefaultRerankResult<VALUE> implements RerankResult<VALUE> {
   readonly originalDocuments: RerankResult<VALUE>['originalDocuments'];
   readonly ranking: RerankResult<VALUE>['ranking'];

package/src/rerank/restricted-telemetry-dispatcher.ts

@@ -0,0 +1,44 @@
+import type { Context } from '@ai-sdk/provider-utils';
+import { createTelemetryDispatcher } from '../telemetry/create-telemetry-dispatcher';
+import { filterIncludedContext } from '../telemetry/filter-included-context';
+import type { TelemetryDispatcher } from '../telemetry/telemetry';
+import type { TelemetryOptions } from '../telemetry/telemetry-options';
+import type { Callback } from '../util/callback';
+import type { RerankStartEvent, RerankEndEvent } from './rerank-events';
+
+/**
+ * Filters runtime context before sending operation events to telemetry integrations.
+ * User callbacks continue to receive the original context.
+ */
+export function createRestrictedTelemetryDispatcher<
+  RUNTIME_CONTEXT extends Context,
+>({
+  telemetry,
+}: {
+  telemetry?: TelemetryOptions<RUNTIME_CONTEXT>;
+}): Omit<TelemetryDispatcher, 'onStart' | 'onEnd'> & {
+  onStart: Callback<RerankStartEvent<RUNTIME_CONTEXT>>;
+  onEnd: Callback<RerankEndEvent<RUNTIME_CONTEXT>>;
+} {
+  const dispatcher = createTelemetryDispatcher({ telemetry });
+
+  return {
+    ...dispatcher,
+    onStart: event =>
+      dispatcher.onStart?.({
+        ...event,
+        runtimeContext: filterIncludedContext({
+          context: event.runtimeContext,
+          includeContext: telemetry?.includeRuntimeContext,
+        }),
+      }),
+    onEnd: event =>
+      dispatcher.onEnd?.({
+        ...event,
+        runtimeContext: filterIncludedContext({
+          context: event.runtimeContext,
+          includeContext: telemetry?.includeRuntimeContext,
+        }),
+      }),
+  };
+}

package/src/telemetry/filter-included-context.ts

@@ -0,0 +1,24 @@
+import type { Context } from '@ai-sdk/provider-utils';
+import type { IncludedContext } from './telemetry-options';
+
+/**
+ * Returns a shallow copy of the runtime context with only top-level
+ * properties marked for telemetry inclusion.
+ */
+export function filterIncludedContext<CONTEXT extends Context>({
+  context,
+  includeContext,
+}: {
+  context: CONTEXT;
+  includeContext: IncludedContext<CONTEXT>;
+}): Context {
+  if (context == null) {
+    return {};
+  }
+
+  return Object.fromEntries(
+    Object.entries(context).filter(
+      ([key]) => includeContext?.[key as keyof CONTEXT] === true,
+    ),
+  );
+}

package/src/ui/index.ts

@@ -44,6 +44,7 @@ export {
   isReasoningUIPart,
   isStaticToolUIPart,
   isTextUIPart,
+  isToolOutputErrorUIPart,
   isToolUIPart,
   type CustomContentUIPart,
   type DataUIPart,
@@ -57,6 +58,7 @@ export {
   type SourceUrlUIPart,
   type StepStartUIPart,
   type TextUIPart,
+  type ToolOutputErrorUIPart,
   type ToolUIPart,
   type UIDataTypes,
   type UIMessage,

package/src/ui/ui-messages.ts

@@ -361,7 +361,7 @@ export type UIToolInvocation<TOOL extends UITool | Tool> = {
       };
     }
   | {
-      state: 'output-error'; // TODO AI SDK 6: change to 'error' state
+      state: 'output-error';
       input: asUITool<TOOL>['input'] | undefined;
       rawInput?: unknown; // TODO AI SDK 6: remove this field, input should be unknown
       output?: never;
@@ -489,7 +489,7 @@ export type DynamicToolUIPart = {
       };
     }
   | {
-      state: 'output-error'; // TODO AI SDK 6: change to 'error' state
+      state: 'output-error';
       input: unknown;
       output?: never;
       errorText: string;
@@ -523,6 +523,17 @@ export type DynamicToolUIPart = {
     }
 );
 
+/**
+ * A static or dynamic tool UI part whose execution failed.
+ *
+ * Use `isToolOutputErrorUIPart` to identify tool output errors without
+ * depending on the underlying tool state discriminator.
+ */
+export type ToolOutputErrorUIPart<TOOLS extends UITools = UITools> = Extract<
+  ToolUIPart<TOOLS> | DynamicToolUIPart,
+  { state: 'output-error' }
+>;
+
 /**
  * Type guard to check if a message part is a text part.
  */
@@ -603,6 +614,17 @@ export function isToolUIPart<TOOLS extends UITools>(
   return isStaticToolUIPart(part) || isDynamicToolUIPart(part);
 }
 
+/**
+ * Check if a message part is a tool output error part.
+ *
+ * This works for both static and dynamic tools.
+ */
+export function isToolOutputErrorUIPart<TOOLS extends UITools>(
+  part: UIMessagePart<UIDataTypes, TOOLS>,
+): part is ToolOutputErrorUIPart<TOOLS> {
+  return isToolUIPart(part) && part.state === 'output-error';
+}
+
 /**
  * Returns the name of the static tool.
  *

package/src/ui/validate-ui-messages.ts

@@ -438,10 +438,19 @@ type ValidateUIMessagesOptions<UI_MESSAGE extends UIMessage> = {
     >;
   };
   tools?: {
-    [NAME in keyof InferUIMessageTools<UI_MESSAGE> & string]?: Tool<
-      InferUIMessageTools<UI_MESSAGE>[NAME]['input'],
-      InferUIMessageTools<UI_MESSAGE>[NAME]['output']
-    >;
+    [NAME in keyof InferUIMessageTools<UI_MESSAGE> & string]?:
+      | Tool<
+          InferUIMessageTools<UI_MESSAGE>[NAME]['input'],
+          InferUIMessageTools<UI_MESSAGE>[NAME]['output']
+        >
+      | {
+          inputSchema: FlexibleSchema<
+            InferUIMessageTools<UI_MESSAGE>[NAME]['input']
+          >;
+          outputSchema?: FlexibleSchema<
+            InferUIMessageTools<UI_MESSAGE>[NAME]['output']
+          >;
+        };
   };
 };