From 1febce0e487abec5da64d67b0d57ebeff7f970a5 Mon Sep 17 00:00:00 2001
From: Lars Grammel <lars.grammel@gmail.com>
Date: Thu, 5 Mar 2026 11:34:38 +0100
Subject: [PATCH 01/41] 1

---
 .../ai/src/generate-text/generate-text.ts     | 22 +++++++++----------
 packages/ai/src/generate-text/prepare-step.ts | 12 +++++++---
 packages/provider-utils/src/types/context.ts  |  5 +++++
 packages/provider-utils/src/types/index.ts    |  3 ++-
 .../provider-utils/src/types/tool.test-d.ts   | 12 +++++-----
 packages/provider-utils/src/types/tool.ts     | 22 ++++++++++++++-----
 6 files changed, 49 insertions(+), 27 deletions(-)
 create mode 100644 packages/provider-utils/src/types/context.ts

diff --git a/packages/ai/src/generate-text/generate-text.ts b/packages/ai/src/generate-text/generate-text.ts
index 355310668877..72612e04f2cf 100644
--- a/packages/ai/src/generate-text/generate-text.ts
+++ b/packages/ai/src/generate-text/generate-text.ts
@@ -2,20 +2,20 @@ import {
   LanguageModelV4,
   LanguageModelV4Content,
   LanguageModelV4ToolCall,
-  LanguageModelV4ToolChoice,
 } from '@ai-sdk/provider';
 import {
+  Context,
+  ContextRegistry,
   createIdGenerator,
   getErrorMessage,
   IdGenerator,
   ProviderOptions,
-  SystemModelMessage,
   ToolApprovalResponse,
   withUserAgentSuffix,
 } from '@ai-sdk/provider-utils';
 import { Tracer } from '@opentelemetry/api';
 import { NoOutputGeneratedError } from '../error';
-import { notify } from '../util/notify';
+import { ToolCallNotFoundForApprovalError } from '../error/tool-call-not-found-for-approval-error';
 import { logWarnings } from '../logger/log-warnings';
 import { resolveLanguageModel } from '../model/resolve-model';
 import { ModelMessage } from '../prompt';
@@ -23,7 +23,6 @@ import {
   CallSettings,
   getStepTimeoutMs,
   getTotalTimeoutMs,
-  TimeoutConfiguration,
 } from '../prompt/call-settings';
 import { convertToLanguageModelPrompt } from '../prompt/convert-to-language-model-prompt';
 import { createToolModelOutput } from '../prompt/create-tool-model-output';
@@ -32,14 +31,13 @@ import { prepareToolsAndToolChoice } from '../prompt/prepare-tools-and-tool-choi
 import { Prompt } from '../prompt/prompt';
 import { standardizePrompt } from '../prompt/standardize-prompt';
 import { wrapGatewayError } from '../prompt/wrap-gateway-error';
-import { ToolCallNotFoundForApprovalError } from '../error/tool-call-not-found-for-approval-error';
 import { assembleOperationName } from '../telemetry/assemble-operation-name';
 import { getBaseTelemetryAttributes } from '../telemetry/get-base-telemetry-attributes';
+import { getGlobalTelemetryIntegration } from '../telemetry/get-global-telemetry-integration';
 import { getTracer } from '../telemetry/get-tracer';
 import { recordSpan } from '../telemetry/record-span';
 import { selectTelemetryAttributes } from '../telemetry/select-telemetry-attributes';
 import { stringifyForTelemetry } from '../telemetry/stringify-for-telemetry';
-import { getGlobalTelemetryIntegration } from '../telemetry/get-global-telemetry-integration';
 import { TelemetrySettings } from '../telemetry/telemetry-settings';
 import {
   LanguageModel,
@@ -53,7 +51,9 @@ import {
 } from '../types/usage';
 import { asArray } from '../util/as-array';
 import { DownloadFunction } from '../util/download/download-function';
+import { mergeAbortSignals } from '../util/merge-abort-signals';
 import { mergeObjects } from '../util/merge-objects';
+import { notify } from '../util/notify';
 import { prepareRetries } from '../util/prepare-retries';
 import { VERSION } from '../version';
 import type {
@@ -91,7 +91,6 @@ import { TypedToolError } from './tool-error';
 import { ToolOutput } from './tool-output';
 import { TypedToolResult } from './tool-result';
 import { ToolSet } from './tool-set';
-import { mergeAbortSignals } from '../util/merge-abort-signals';
 
 const originalGenerateId = createIdGenerator({
   prefix: 'aitxt',
@@ -247,6 +246,7 @@ export type GenerateTextOnFinishCallback<TOOLS extends ToolSet> = (
  */
 export async function generateText<
   TOOLS extends ToolSet,
+  CONTEXT extends Partial<ContextRegistry> = ContextRegistry,
   OUTPUT extends Output = Output<string, string>,
 >({
   model: modelArg,
@@ -270,7 +270,7 @@ export async function generateText<
   prepareStep = experimental_prepareStep,
   experimental_repairToolCall: repairToolCall,
   experimental_download: download,
-  experimental_context,
+  experimental_context = {} as Context<CONTEXT>,
   experimental_include: include,
   _internal: { generateId = originalGenerateId } = {},
   experimental_onStart: onStart,
@@ -352,12 +352,12 @@ export async function generateText<
     /**
      * @deprecated Use `prepareStep` instead.
      */
-    experimental_prepareStep?: PrepareStepFunction<NoInfer<TOOLS>>;
+    experimental_prepareStep?: PrepareStepFunction<NoInfer<TOOLS>, CONTEXT>;
 
     /**
      * Optional function that you can use to provide different settings for a step.
      */
-    prepareStep?: PrepareStepFunction<NoInfer<TOOLS>>;
+    prepareStep?: PrepareStepFunction<NoInfer<TOOLS>, CONTEXT>;
 
     /**
      * A function that attempts to repair a tool call that failed to parse.
@@ -410,7 +410,7 @@ export async function generateText<
      *
      * @default undefined
      */
-    experimental_context?: unknown;
+    experimental_context?: Context<CONTEXT>;
 
     /**
      * Settings for controlling what data is included in step results.
diff --git a/packages/ai/src/generate-text/prepare-step.ts b/packages/ai/src/generate-text/prepare-step.ts
index f3f367094aee..334cdfaf4246 100644
--- a/packages/ai/src/generate-text/prepare-step.ts
+++ b/packages/ai/src/generate-text/prepare-step.ts
@@ -1,4 +1,6 @@
 import {
+  Context,
+  ContextRegistry,
   ModelMessage,
   ProviderOptions,
   SystemModelMessage,
@@ -22,6 +24,7 @@ import { StepResult } from './step-result';
  */
 export type PrepareStepFunction<
   TOOLS extends Record<string, Tool> = Record<string, Tool>,
+  CONTEXT extends Partial<ContextRegistry> = ContextRegistry,
 > = (options: {
   /**
    * The steps that have been executed so far.
@@ -46,8 +49,10 @@ export type PrepareStepFunction<
   /**
    * The context passed via the experimental_context setting (experimental).
    */
-  experimental_context: unknown;
-}) => PromiseLike<PrepareStepResult<TOOLS>> | PrepareStepResult<TOOLS>;
+  experimental_context: Context<CONTEXT>;
+}) =>
+  | PromiseLike<PrepareStepResult<TOOLS, CONTEXT>>
+  | PrepareStepResult<TOOLS, CONTEXT>;
 
 /**
  * The result type returned by a {@link PrepareStepFunction},
@@ -55,6 +60,7 @@ export type PrepareStepFunction<
  */
 export type PrepareStepResult<
   TOOLS extends Record<string, Tool> = Record<string, Tool>,
+  CONTEXT extends Partial<ContextRegistry> = ContextRegistry,
 > =
   | {
       /**
@@ -90,7 +96,7 @@ export type PrepareStepResult<
        * Changing the context will affect the context in this step
        * and all subsequent steps.
        */
-      experimental_context?: unknown;
+      experimental_context?: Context<CONTEXT>;
 
       /**
        * Additional provider-specific options for this step.
diff --git a/packages/provider-utils/src/types/context.ts b/packages/provider-utils/src/types/context.ts
new file mode 100644
index 000000000000..92927870b874
--- /dev/null
+++ b/packages/provider-utils/src/types/context.ts
@@ -0,0 +1,5 @@
+export interface ContextRegistry {} // open via declaration merging
+
+export type Context<T extends Partial<ContextRegistry> = ContextRegistry> = {
+  [K in keyof T]: T[K];
+};
diff --git a/packages/provider-utils/src/types/index.ts b/packages/provider-utils/src/types/index.ts
index 7cb7fb0e3459..eb3ead960a9d 100644
--- a/packages/provider-utils/src/types/index.ts
+++ b/packages/provider-utils/src/types/index.ts
@@ -11,6 +11,7 @@ export type {
   ToolResultOutput,
   ToolResultPart,
 } from './content-part';
+export type { Context, ContextRegistry } from './context';
 export type { DataContent } from './data-content';
 export { executeTool } from './execute-tool';
 export type { ModelMessage } from './model-message';
@@ -22,8 +23,8 @@ export {
   type InferToolInput,
   type InferToolOutput,
   type Tool,
-  type ToolExecutionOptions,
   type ToolExecuteFunction,
+  type ToolExecutionOptions,
   type ToolNeedsApprovalFunction,
 } from './tool';
 export type { ToolApprovalRequest } from './tool-approval-request';
diff --git a/packages/provider-utils/src/types/tool.test-d.ts b/packages/provider-utils/src/types/tool.test-d.ts
index 5b6a6904d96f..b09178639534 100644
--- a/packages/provider-utils/src/types/tool.test-d.ts
+++ b/packages/provider-utils/src/types/tool.test-d.ts
@@ -1,10 +1,10 @@
-import { LanguageModelV3ToolResultPart } from '@ai-sdk/provider';
 import { describe, expectTypeOf, it } from 'vitest';
 import { z } from 'zod/v4';
 import { FlexibleSchema } from '../schema';
+import { ToolResultOutput } from './content-part';
+import { Context, ContextRegistry } from './context';
 import { ModelMessage } from './model-message';
 import { Tool, tool, ToolExecuteFunction } from './tool';
-import { ToolResultOutput } from './content-part';
 
 describe('tool type', () => {
   describe('input type', () => {
@@ -176,7 +176,7 @@ describe('tool type', () => {
           expectTypeOf(options).toEqualTypeOf<{
             toolCallId: string;
             messages: ModelMessage[];
-            experimental_context?: unknown;
+            experimental_context?: Context<ContextRegistry> | undefined;
           }>();
           return true;
         },
@@ -189,7 +189,7 @@ describe('tool type', () => {
             options: {
               toolCallId: string;
               messages: ModelMessage[];
-              experimental_context: unknown;
+              experimental_context?: Context<ContextRegistry> | undefined;
             },
           ) => boolean | PromiseLike<boolean>)
         | undefined
@@ -205,7 +205,7 @@ describe('tool type', () => {
           expectTypeOf(options).toEqualTypeOf<{
             toolCallId: string;
             messages: ModelMessage[];
-            experimental_context?: unknown;
+            experimental_context?: Context<ContextRegistry> | undefined;
           }>();
           return true;
         },
@@ -218,7 +218,7 @@ describe('tool type', () => {
             options: {
               toolCallId: string;
               messages: ModelMessage[];
-              experimental_context: unknown;
+              experimental_context?: Context<ContextRegistry> | undefined;
             },
           ) => boolean | PromiseLike<boolean>)
         | undefined
diff --git a/packages/provider-utils/src/types/tool.ts b/packages/provider-utils/src/types/tool.ts
index b07089fe24ba..68c97f2e73a0 100644
--- a/packages/provider-utils/src/types/tool.ts
+++ b/packages/provider-utils/src/types/tool.ts
@@ -3,11 +3,14 @@ import { FlexibleSchema } from '../schema';
 import { ToolResultOutput } from './content-part';
 import { ModelMessage } from './model-message';
 import { ProviderOptions } from './provider-options';
+import { Context, ContextRegistry } from './context';
 
 /**
  * Additional options that are sent into each tool call.
  */
-export interface ToolExecutionOptions {
+export interface ToolExecutionOptions<
+  CONTEXT extends Partial<ContextRegistry> = ContextRegistry,
+> {
   /**
    * The ID of the tool call. You can use it e.g. when sending tool-call related information with stream data.
    */
@@ -36,13 +39,16 @@ export interface ToolExecutionOptions {
    *
    * Experimental (can break in patch releases).
    */
-  experimental_context?: unknown;
+  experimental_context?: Context<CONTEXT>;
 }
 
 /**
  * Function that is called to determine if the tool needs approval before it can be executed.
  */
-export type ToolNeedsApprovalFunction<INPUT> = (
+export type ToolNeedsApprovalFunction<
+  INPUT,
+  CONTEXT extends Partial<ContextRegistry> = ContextRegistry,
+> = (
   input: INPUT,
   options: {
     /**
@@ -61,13 +67,17 @@ export type ToolNeedsApprovalFunction<INPUT> = (
      *
      * Experimental (can break in patch releases).
      */
-    experimental_context?: unknown;
+    experimental_context?: Context<CONTEXT>;
   },
 ) => boolean | PromiseLike<boolean>;
 
-export type ToolExecuteFunction<INPUT, OUTPUT> = (
+export type ToolExecuteFunction<
+  INPUT,
+  OUTPUT,
+  CONTEXT extends Partial<ContextRegistry> = ContextRegistry,
+> = (
   input: INPUT,
-  options: ToolExecutionOptions,
+  options: ToolExecutionOptions<CONTEXT>,
 ) => AsyncIterable<OUTPUT> | PromiseLike<OUTPUT> | OUTPUT;
 
 // 0 extends 1 & N checks for any

From 672f5edf4837027b2579ee86782898d8eba8c268 Mon Sep 17 00:00:00 2001
From: Lars Grammel <lars.grammel@gmail.com>
Date: Thu, 5 Mar 2026 12:08:21 +0100
Subject: [PATCH 02/41] 2

---
 .../openai/tool-call-with-context.ts          | 29 ++++++-
 .../ai/src/generate-text/callback-events.ts   | 63 +++++++++-------
 packages/ai/src/generate-text/content-part.ts | 14 ++--
 .../ai/src/generate-text/generate-text.ts     | 27 ++++---
 packages/ai/src/generate-text/prepare-step.ts | 10 +--
 packages/ai/src/generate-text/step-result.ts  | 13 ++--
 .../ai/src/generate-text/stop-condition.ts    |  8 +-
 .../tool-approval-request-output.ts           |  8 +-
 packages/ai/src/generate-text/tool-call.ts    | 18 +++--
 packages/ai/src/generate-text/tool-error.ts   | 16 ++--
 packages/ai/src/generate-text/tool-result.ts  | 22 ++++--
 packages/ai/src/generate-text/tool-set.ts     | 15 +++-
 .../src/provider-tool-factory.ts              | 75 +++++++++++--------
 .../provider-utils/src/types/execute-tool.ts  | 11 ++-
 .../provider-utils/src/types/tool.test-d.ts   | 20 +++--
 packages/provider-utils/src/types/tool.ts     | 51 ++++++++-----
 16 files changed, 257 insertions(+), 143 deletions(-)

diff --git a/examples/ai-functions/src/generate-text/openai/tool-call-with-context.ts b/examples/ai-functions/src/generate-text/openai/tool-call-with-context.ts
index 743ca0c8b363..cae53ba9f0e1 100644
--- a/examples/ai-functions/src/generate-text/openai/tool-call-with-context.ts
+++ b/examples/ai-functions/src/generate-text/openai/tool-call-with-context.ts
@@ -12,10 +12,13 @@ run(async () => {
         inputSchema: z.object({
           location: z.string().describe('The location to get the weather for'),
         }),
+        contextSchema: z.object({
+          weatherApiKey: z.string().describe('The API key for the weather API'),
+        }),
         execute: async ({ location }, { experimental_context: context }) => {
-          const typedContext = context as { weatherApiKey: string }; // or use type validation library
+          console.log(context);
 
-          console.log(typedContext);
+          context satisfies { weatherApiKey: string };
 
           return {
             location,
@@ -23,8 +26,28 @@ run(async () => {
           };
         },
       }),
+      calculator: tool({
+        description: 'Calculate mathematical expressions',
+        inputSchema: z.object({
+          expression: z
+            .string()
+            .describe('The mathematical expression to calculate'),
+        }),
+        contextSchema: z.object({
+          calculatorApiKey: z
+            .string()
+            .describe('The API key for the calculator API'),
+        }),
+        execute: async ({ expression }, { experimental_context: context }) => {
+          console.log(context);
+          return {
+            expression,
+            result: eval(expression),
+          };
+        },
+      }),
     },
-    experimental_context: { weatherApiKey: '123' },
+    experimental_context: { weatherApiKey: '123', calculatorApiKey: '456' },
     prompt: 'What is the weather in San Francisco?',
   });
 
diff --git a/packages/ai/src/generate-text/callback-events.ts b/packages/ai/src/generate-text/callback-events.ts
index 454be6f96a9a..afb35b9282fb 100644
--- a/packages/ai/src/generate-text/callback-events.ts
+++ b/packages/ai/src/generate-text/callback-events.ts
@@ -1,5 +1,6 @@
 import type { LanguageModelV4ToolChoice } from '@ai-sdk/provider';
 import type {
+  ContextRegistry,
   ModelMessage,
   ProviderOptions,
   SystemModelMessage,
@@ -129,7 +130,8 @@ export interface OnStartEvent<
  * Each step represents a single LLM invocation.
  */
 export interface OnStepStartEvent<
-  TOOLS extends ToolSet = ToolSet,
+  CONTEXT extends Partial<ContextRegistry>,
+  TOOLS extends ToolSet<CONTEXT> = ToolSet<CONTEXT>,
   OUTPUT extends Output = Output,
   INCLUDE = { requestBody?: boolean; responseBody?: boolean },
 > {
@@ -184,8 +186,8 @@ export interface OnStepStartEvent<
    * When the condition is an array, any of the conditions can be met to stop.
    */
   readonly stopWhen:
-    | StopCondition<TOOLS>
-    | Array<StopCondition<TOOLS>>
+    | StopCondition<CONTEXT, TOOLS>
+    | Array<StopCondition<CONTEXT, TOOLS>>
     | undefined;
 
   /** The output specification for structured outputs, if configured. */
@@ -216,7 +218,10 @@ export interface OnStepStartEvent<
  *
  * Called when a tool execution begins, before the tool's `execute` function is invoked.
  */
-export interface OnToolCallStartEvent<TOOLS extends ToolSet = ToolSet> {
+export interface OnToolCallStartEvent<
+  CONTEXT extends Partial<ContextRegistry>,
+  TOOLS extends ToolSet<CONTEXT> = ToolSet<CONTEXT>,
+> {
   /** Zero-based index of the current step where this tool call occurs. */
   readonly stepNumber: number | undefined;
 
@@ -224,7 +229,7 @@ export interface OnToolCallStartEvent<TOOLS extends ToolSet = ToolSet> {
   readonly model: CallbackModelInfo | undefined;
 
   /** The full tool call object. */
-  readonly toolCall: TypedToolCall<TOOLS>;
+  readonly toolCall: TypedToolCall<CONTEXT, TOOLS>;
 
   /** The conversation messages available at tool execution time. */
   readonly messages: Array<ModelMessage>;
@@ -307,26 +312,28 @@ export type OnStepFinishEvent<TOOLS extends ToolSet = ToolSet> =
  * Called when the entire generation completes (all steps finished).
  * Includes the final step's result along with aggregated data from all steps.
  */
-export type OnFinishEvent<TOOLS extends ToolSet = ToolSet> =
-  StepResult<TOOLS> & {
-    /** Array containing results from all steps in the generation. */
-    readonly steps: StepResult<TOOLS>[];
-
-    /** Aggregated token usage across all steps. */
-    readonly totalUsage: LanguageModelUsage;
-
-    /**
-     * The final state of the user-defined context object.
-     *
-     * Experimental (can break in patch releases).
-     *
-     * @default undefined
-     */
-    experimental_context: unknown;
-
-    /** Identifier from telemetry settings for grouping related operations. */
-    readonly functionId: string | undefined;
-
-    /** Additional metadata from telemetry settings. */
-    readonly metadata: Record<string, unknown> | undefined;
-  };
+export type OnFinishEvent<
+  CONTEXT extends Partial<ContextRegistry>,
+  TOOLS extends ToolSet<CONTEXT> = ToolSet<CONTEXT>,
+> = StepResult<CONTEXT, TOOLS> & {
+  /** Array containing results from all steps in the generation. */
+  readonly steps: StepResult<CONTEXT, TOOLS>[];
+
+  /** Aggregated token usage across all steps. */
+  readonly totalUsage: LanguageModelUsage;
+
+  /**
+   * The final state of the user-defined context object.
+   *
+   * Experimental (can break in patch releases).
+   *
+   * @default undefined
+   */
+  experimental_context: unknown;
+
+  /** Identifier from telemetry settings for grouping related operations. */
+  readonly functionId: string | undefined;
+
+  /** Additional metadata from telemetry settings. */
+  readonly metadata: Record<string, unknown> | undefined;
+};
diff --git a/packages/ai/src/generate-text/content-part.ts b/packages/ai/src/generate-text/content-part.ts
index 925d1a3663e8..df1e9b1eb90a 100644
--- a/packages/ai/src/generate-text/content-part.ts
+++ b/packages/ai/src/generate-text/content-part.ts
@@ -1,3 +1,4 @@
+import { ContextRegistry } from '@ai-sdk/provider-utils';
 import { ProviderMetadata } from '../types';
 import { Source } from '../types/language-model';
 import { GeneratedFile } from './generated-file';
@@ -8,18 +9,21 @@ import { TypedToolError } from './tool-error';
 import { TypedToolResult } from './tool-result';
 import { ToolSet } from './tool-set';
 
-export type ContentPart<TOOLS extends ToolSet> =
+export type ContentPart<
+  CONTEXT extends Partial<ContextRegistry>,
+  TOOLS extends ToolSet<CONTEXT> = ToolSet<CONTEXT>,
+> =
   | { type: 'text'; text: string; providerMetadata?: ProviderMetadata }
   | ReasoningOutput
   | ({ type: 'source' } & Source)
   | { type: 'file'; file: GeneratedFile; providerMetadata?: ProviderMetadata } // different because of GeneratedFile object
-  | ({ type: 'tool-call' } & TypedToolCall<TOOLS> & {
+  | ({ type: 'tool-call' } & TypedToolCall<CONTEXT, TOOLS> & {
         providerMetadata?: ProviderMetadata;
       })
-  | ({ type: 'tool-result' } & TypedToolResult<TOOLS> & {
+  | ({ type: 'tool-result' } & TypedToolResult<CONTEXT, TOOLS> & {
         providerMetadata?: ProviderMetadata;
       })
-  | ({ type: 'tool-error' } & TypedToolError<TOOLS> & {
+  | ({ type: 'tool-error' } & TypedToolError<CONTEXT, TOOLS> & {
         providerMetadata?: ProviderMetadata;
       })
-  | ToolApprovalRequestOutput<TOOLS>;
+  | ToolApprovalRequestOutput<CONTEXT, TOOLS>;
diff --git a/packages/ai/src/generate-text/generate-text.ts b/packages/ai/src/generate-text/generate-text.ts
index 72612e04f2cf..e15f8e2113c3 100644
--- a/packages/ai/src/generate-text/generate-text.ts
+++ b/packages/ai/src/generate-text/generate-text.ts
@@ -131,10 +131,11 @@ export type GenerateTextOnStartCallback<
  * @param event - The event object containing step configuration.
  */
 export type GenerateTextOnStepStartCallback<
-  TOOLS extends ToolSet = ToolSet,
+  CONTEXT extends Partial<ContextRegistry>,
+  TOOLS extends ToolSet<CONTEXT> = ToolSet<CONTEXT>,
   OUTPUT extends Output = Output,
 > = (
-  event: OnStepStartEvent<TOOLS, OUTPUT, GenerateTextIncludeSettings>,
+  event: OnStepStartEvent<CONTEXT, TOOLS, OUTPUT, GenerateTextIncludeSettings>,
 ) => PromiseLike<void> | void;
 
 /**
@@ -146,8 +147,9 @@ export type GenerateTextOnStepStartCallback<
  * @param event - The event object containing tool call information.
  */
 export type GenerateTextOnToolCallStartCallback<
-  TOOLS extends ToolSet = ToolSet,
-> = (event: OnToolCallStartEvent<TOOLS>) => PromiseLike<void> | void;
+  CONTEXT extends Partial<ContextRegistry>,
+  TOOLS extends ToolSet<CONTEXT> = ToolSet<CONTEXT>,
+> = (event: OnToolCallStartEvent<CONTEXT, TOOLS>) => PromiseLike<void> | void;
 
 /**
  * Callback that is set using the `experimental_onToolCallFinish` option.
@@ -186,9 +188,10 @@ export type GenerateTextOnStepFinishCallback<TOOLS extends ToolSet> = (
  *
  * @param event - The final result along with aggregated step data.
  */
-export type GenerateTextOnFinishCallback<TOOLS extends ToolSet> = (
-  event: OnFinishEvent<TOOLS>,
-) => PromiseLike<void> | void;
+export type GenerateTextOnFinishCallback<
+  CONTEXT extends Partial<ContextRegistry>,
+  TOOLS extends ToolSet<CONTEXT> = ToolSet<CONTEXT>,
+> = (event: OnFinishEvent<CONTEXT, TOOLS>) => PromiseLike<void> | void;
 
 /**
  * Generate a text and call tools for a given prompt using a language model.
@@ -245,8 +248,8 @@ export type GenerateTextOnFinishCallback<TOOLS extends ToolSet> = (
  * A result object that contains the generated text, the results of the tool calls, and additional information.
  */
 export async function generateText<
-  TOOLS extends ToolSet,
-  CONTEXT extends Partial<ContextRegistry> = ContextRegistry,
+  CONTEXT extends Partial<ContextRegistry>,
+  TOOLS extends ToolSet<CONTEXT>,
   OUTPUT extends Output = Output<string, string>,
 >({
   model: modelArg,
@@ -352,12 +355,12 @@ export async function generateText<
     /**
      * @deprecated Use `prepareStep` instead.
      */
-    experimental_prepareStep?: PrepareStepFunction<NoInfer<TOOLS>, CONTEXT>;
+    experimental_prepareStep?: PrepareStepFunction<CONTEXT, NoInfer<TOOLS>>;
 
     /**
      * Optional function that you can use to provide different settings for a step.
      */
-    prepareStep?: PrepareStepFunction<NoInfer<TOOLS>, CONTEXT>;
+    prepareStep?: PrepareStepFunction<CONTEXT, NoInfer<TOOLS>>;
 
     /**
      * A function that attempts to repair a tool call that failed to parse.
@@ -1191,7 +1194,7 @@ export async function generateText<
             onFinish,
             globalTelemetry.onFinish as
               | undefined
-              | GenerateTextOnFinishCallback<TOOLS>,
+              | GenerateTextOnFinishCallback<CONTEXT, TOOLS>,
           ],
         });
 
diff --git a/packages/ai/src/generate-text/prepare-step.ts b/packages/ai/src/generate-text/prepare-step.ts
index 334cdfaf4246..8ba9901b631d 100644
--- a/packages/ai/src/generate-text/prepare-step.ts
+++ b/packages/ai/src/generate-text/prepare-step.ts
@@ -4,10 +4,10 @@ import {
   ModelMessage,
   ProviderOptions,
   SystemModelMessage,
-  Tool,
 } from '@ai-sdk/provider-utils';
 import { LanguageModel, ToolChoice } from '../types/language-model';
 import { StepResult } from './step-result';
+import { ToolSet } from './tool-set';
 
 /**
  * Function that you can use to provide different settings for a step.
@@ -23,8 +23,8 @@ import { StepResult } from './step-result';
  * If you return undefined (or for undefined settings), the settings from the outer level will be used.
  */
 export type PrepareStepFunction<
-  TOOLS extends Record<string, Tool> = Record<string, Tool>,
   CONTEXT extends Partial<ContextRegistry> = ContextRegistry,
+  TOOLS extends ToolSet<CONTEXT> = ToolSet<CONTEXT>,
 > = (options: {
   /**
    * The steps that have been executed so far.
@@ -51,16 +51,16 @@ export type PrepareStepFunction<
    */
   experimental_context: Context<CONTEXT>;
 }) =>
-  | PromiseLike<PrepareStepResult<TOOLS, CONTEXT>>
-  | PrepareStepResult<TOOLS, CONTEXT>;
+  | PromiseLike<PrepareStepResult<CONTEXT, TOOLS>>
+  | PrepareStepResult<CONTEXT, TOOLS>;
 
 /**
  * The result type returned by a {@link PrepareStepFunction},
  * allowing per-step overrides of model, tools, or messages.
  */
 export type PrepareStepResult<
-  TOOLS extends Record<string, Tool> = Record<string, Tool>,
   CONTEXT extends Partial<ContextRegistry> = ContextRegistry,
+  TOOLS extends ToolSet<CONTEXT> = ToolSet<CONTEXT>,
 > =
   | {
       /**
diff --git a/packages/ai/src/generate-text/step-result.ts b/packages/ai/src/generate-text/step-result.ts
index 73b53f084c67..94d7a9e92fd2 100644
--- a/packages/ai/src/generate-text/step-result.ts
+++ b/packages/ai/src/generate-text/step-result.ts
@@ -1,4 +1,4 @@
-import { ReasoningPart } from '@ai-sdk/provider-utils';
+import { ContextRegistry, ReasoningPart } from '@ai-sdk/provider-utils';
 import {
   CallWarning,
   FinishReason,
@@ -22,7 +22,10 @@ import { ToolSet } from './tool-set';
 /**
  * The result of a single step in the generation process.
  */
-export type StepResult<TOOLS extends ToolSet> = {
+export type StepResult<
+  CONTEXT extends Partial<ContextRegistry>,
+  TOOLS extends ToolSet<CONTEXT> = ToolSet<CONTEXT>,
+> = {
   /**
    * Zero-based index of this step.
    */
@@ -58,7 +61,7 @@ export type StepResult<TOOLS extends ToolSet> = {
   /**
    * The content that was generated in the last step.
    */
-  readonly content: Array<ContentPart<TOOLS>>;
+  readonly content: Array<ContentPart<CONTEXT, TOOLS>>;
 
   /**
    * The generated text.
@@ -88,7 +91,7 @@ export type StepResult<TOOLS extends ToolSet> = {
   /**
    * The tool calls that were made during the generation.
    */
-  readonly toolCalls: Array<TypedToolCall<TOOLS>>;
+  readonly toolCalls: Array<TypedToolCall<CONTEXT, TOOLS>>;
 
   /**
    * The static tool calls that were made in the last step.
@@ -103,7 +106,7 @@ export type StepResult<TOOLS extends ToolSet> = {
   /**
    * The results of the tool calls.
    */
-  readonly toolResults: Array<TypedToolResult<TOOLS>>;
+  readonly toolResults: Array<TypedToolResult<CONTEXT, TOOLS>>;
 
   /**
    * The static tool results that were made in the last step.
diff --git a/packages/ai/src/generate-text/stop-condition.ts b/packages/ai/src/generate-text/stop-condition.ts
index 4a2bd316ee49..37ce347927f1 100644
--- a/packages/ai/src/generate-text/stop-condition.ts
+++ b/packages/ai/src/generate-text/stop-condition.ts
@@ -1,8 +1,12 @@
+import { ContextRegistry } from '@ai-sdk/provider-utils';
 import { StepResult } from './step-result';
 import { ToolSet } from './tool-set';
 
-export type StopCondition<TOOLS extends ToolSet> = (options: {
-  steps: Array<StepResult<TOOLS>>;
+export type StopCondition<
+  CONTEXT extends Partial<ContextRegistry>,
+  TOOLS extends ToolSet<CONTEXT> = ToolSet<CONTEXT>,
+> = (options: {
+  steps: Array<StepResult<CONTEXT, TOOLS>>;
 }) => PromiseLike<boolean> | boolean;
 
 export function stepCountIs(stepCount: number): StopCondition<any> {
diff --git a/packages/ai/src/generate-text/tool-approval-request-output.ts b/packages/ai/src/generate-text/tool-approval-request-output.ts
index 2880b67e0a78..b43ac9f11bc2 100644
--- a/packages/ai/src/generate-text/tool-approval-request-output.ts
+++ b/packages/ai/src/generate-text/tool-approval-request-output.ts
@@ -1,3 +1,4 @@
+import { ContextRegistry } from '@ai-sdk/provider-utils';
 import { TypedToolCall } from './tool-call';
 import { ToolSet } from './tool-set';
 
@@ -6,7 +7,10 @@ import { ToolSet } from './tool-set';
  *
  * The tool approval request can be approved or denied in the next tool message.
  */
-export type ToolApprovalRequestOutput<TOOLS extends ToolSet> = {
+export type ToolApprovalRequestOutput<
+  CONTEXT extends Partial<ContextRegistry>,
+  TOOLS extends ToolSet<CONTEXT> = ToolSet<CONTEXT>,
+> = {
   type: 'tool-approval-request';
 
   /**
@@ -17,5 +21,5 @@ export type ToolApprovalRequestOutput<TOOLS extends ToolSet> = {
   /**
    * Tool call that the approval request is for.
    */
-  toolCall: TypedToolCall<TOOLS>;
+  toolCall: TypedToolCall<CONTEXT, TOOLS>;
 };
diff --git a/packages/ai/src/generate-text/tool-call.ts b/packages/ai/src/generate-text/tool-call.ts
index b6f087a675f9..b214f62b0405 100644
--- a/packages/ai/src/generate-text/tool-call.ts
+++ b/packages/ai/src/generate-text/tool-call.ts
@@ -1,4 +1,4 @@
-import { Tool } from '@ai-sdk/provider-utils';
+import { ContextRegistry, Tool } from '@ai-sdk/provider-utils';
 import { ProviderMetadata } from '../types';
 import { ValueOf } from '../util/value-of';
 import { ToolSet } from './tool-set';
@@ -10,10 +10,15 @@ type BaseToolCall = {
   providerMetadata?: ProviderMetadata;
 };
 
-export type StaticToolCall<TOOLS extends ToolSet> = ValueOf<{
+export type StaticToolCall<
+  CONTEXT extends Partial<ContextRegistry>,
+  TOOLS extends ToolSet<CONTEXT> = ToolSet<CONTEXT>,
+> = ValueOf<{
   [NAME in keyof TOOLS]: BaseToolCall & {
     toolName: NAME & string;
-    input: TOOLS[NAME] extends Tool<infer PARAMETERS> ? PARAMETERS : never;
+    input: TOOLS[NAME] extends Tool<CONTEXT, infer PARAMETERS>
+      ? PARAMETERS
+      : never;
     dynamic?: false | undefined;
     invalid?: false | undefined;
     error?: never;
@@ -42,6 +47,7 @@ export type DynamicToolCall = BaseToolCall & {
   error?: unknown;
 };
 
-export type TypedToolCall<TOOLS extends ToolSet> =
-  | StaticToolCall<TOOLS>
-  | DynamicToolCall;
+export type TypedToolCall<
+  CONTEXT extends Partial<ContextRegistry>,
+  TOOLS extends ToolSet<CONTEXT> = ToolSet<CONTEXT>,
+> = StaticToolCall<CONTEXT, TOOLS> | DynamicToolCall;
diff --git a/packages/ai/src/generate-text/tool-error.ts b/packages/ai/src/generate-text/tool-error.ts
index 12d0540ef8f7..da2bca8d2d00 100644
--- a/packages/ai/src/generate-text/tool-error.ts
+++ b/packages/ai/src/generate-text/tool-error.ts
@@ -1,14 +1,17 @@
-import { InferToolInput } from '@ai-sdk/provider-utils';
+import { ContextRegistry, InferToolInput } from '@ai-sdk/provider-utils';
 import { ProviderMetadata } from '../types';
 import { ValueOf } from '../util/value-of';
 import { ToolSet } from './tool-set';
 
-export type StaticToolError<TOOLS extends ToolSet> = ValueOf<{
+export type StaticToolError<
+  CONTEXT extends Partial<ContextRegistry>,
+  TOOLS extends ToolSet<CONTEXT> = ToolSet<CONTEXT>,
+> = ValueOf<{
   [NAME in keyof TOOLS]: {
     type: 'tool-error';
     toolCallId: string;
     toolName: NAME & string;
-    input: InferToolInput<TOOLS[NAME]>;
+    input: InferToolInput<CONTEXT, TOOLS[NAME]>;
     error: unknown;
     providerExecuted?: boolean;
     providerMetadata?: ProviderMetadata;
@@ -29,6 +32,7 @@ export type DynamicToolError = {
   title?: string;
 };
 
-export type TypedToolError<TOOLS extends ToolSet> =
-  | StaticToolError<TOOLS>
-  | DynamicToolError;
+export type TypedToolError<
+  CONTEXT extends Partial<ContextRegistry>,
+  TOOLS extends ToolSet<CONTEXT> = ToolSet<CONTEXT>,
+> = StaticToolError<TOOLS> | DynamicToolError;
diff --git a/packages/ai/src/generate-text/tool-result.ts b/packages/ai/src/generate-text/tool-result.ts
index 406deef0231e..1f2943f644d0 100644
--- a/packages/ai/src/generate-text/tool-result.ts
+++ b/packages/ai/src/generate-text/tool-result.ts
@@ -1,15 +1,22 @@
-import { InferToolInput, InferToolOutput } from '@ai-sdk/provider-utils';
+import {
+  ContextRegistry,
+  InferToolInput,
+  InferToolOutput,
+} from '@ai-sdk/provider-utils';
 import { ProviderMetadata } from '../types';
 import { ValueOf } from '../../src/util/value-of';
 import { ToolSet } from './tool-set';
 
-export type StaticToolResult<TOOLS extends ToolSet> = ValueOf<{
+export type StaticToolResult<
+  CONTEXT extends Partial<ContextRegistry>,
+  TOOLS extends ToolSet<CONTEXT> = ToolSet<CONTEXT>,
+> = ValueOf<{
   [NAME in keyof TOOLS]: {
     type: 'tool-result';
     toolCallId: string;
     toolName: NAME & string;
-    input: InferToolInput<TOOLS[NAME]>;
-    output: InferToolOutput<TOOLS[NAME]>;
+    input: InferToolInput<CONTEXT, TOOLS[NAME]>;
+    output: InferToolOutput<CONTEXT, TOOLS[NAME]>;
     providerExecuted?: boolean;
     providerMetadata?: ProviderMetadata;
     dynamic?: false | undefined;
@@ -31,6 +38,7 @@ export type DynamicToolResult = {
   title?: string;
 };
 
-export type TypedToolResult<TOOLS extends ToolSet> =
-  | StaticToolResult<TOOLS>
-  | DynamicToolResult;
+export type TypedToolResult<
+  CONTEXT extends Partial<ContextRegistry>,
+  TOOLS extends ToolSet<CONTEXT> = ToolSet<CONTEXT>,
+> = StaticToolResult<CONTEXT, TOOLS> | DynamicToolResult;
diff --git a/packages/ai/src/generate-text/tool-set.ts b/packages/ai/src/generate-text/tool-set.ts
index 09e64774d48e..68d308fb5a25 100644
--- a/packages/ai/src/generate-text/tool-set.ts
+++ b/packages/ai/src/generate-text/tool-set.ts
@@ -1,10 +1,17 @@
-import { Tool } from '@ai-sdk/provider-utils';
+import { ContextRegistry, Tool } from '@ai-sdk/provider-utils';
 
-export type ToolSet = Record<
+export type ToolSet<
+  CONTEXT extends Partial<ContextRegistry> = ContextRegistry,
+> = Record<
   string,
-  (Tool<never, never> | Tool<any, any> | Tool<any, never> | Tool<never, any>) &
+  (
+    | Tool<CONTEXT, never, never>
+    | Tool<CONTEXT, any, any>
+    | Tool<CONTEXT, any, never>
+    | Tool<CONTEXT, never, any>
+  ) &
     Pick<
-      Tool<any, any>,
+      Tool<CONTEXT, any, any>,
       | 'execute'
       | 'onInputAvailable'
       | 'onInputStart'
diff --git a/packages/provider-utils/src/provider-tool-factory.ts b/packages/provider-utils/src/provider-tool-factory.ts
index 2c29130e0a51..7595d518c8d9 100644
--- a/packages/provider-utils/src/provider-tool-factory.ts
+++ b/packages/provider-utils/src/provider-tool-factory.ts
@@ -1,24 +1,33 @@
 import { tool, Tool, ToolExecuteFunction } from './types/tool';
 import { FlexibleSchema } from './schema';
+import { ContextRegistry } from './types/context';
 
-export type ProviderToolFactory<INPUT, ARGS extends object> = <OUTPUT>(
+export type ProviderToolFactory<
+  CONTEXT extends Partial<ContextRegistry>,
+  INPUT,
+  ARGS extends object,
+> = <OUTPUT>(
   options: ARGS & {
-    execute?: ToolExecuteFunction<INPUT, OUTPUT>;
-    needsApproval?: Tool<INPUT, OUTPUT>['needsApproval'];
-    toModelOutput?: Tool<INPUT, OUTPUT>['toModelOutput'];
-    onInputStart?: Tool<INPUT, OUTPUT>['onInputStart'];
-    onInputDelta?: Tool<INPUT, OUTPUT>['onInputDelta'];
-    onInputAvailable?: Tool<INPUT, OUTPUT>['onInputAvailable'];
+    execute?: ToolExecuteFunction<CONTEXT, INPUT, OUTPUT>;
+    needsApproval?: Tool<CONTEXT, INPUT, OUTPUT>['needsApproval'];
+    toModelOutput?: Tool<CONTEXT, INPUT, OUTPUT>['toModelOutput'];
+    onInputStart?: Tool<CONTEXT, INPUT, OUTPUT>['onInputStart'];
+    onInputDelta?: Tool<CONTEXT, INPUT, OUTPUT>['onInputDelta'];
+    onInputAvailable?: Tool<CONTEXT, INPUT, OUTPUT>['onInputAvailable'];
   },
-) => Tool<INPUT, OUTPUT>;
+) => Tool<CONTEXT, INPUT, OUTPUT>;
 
-export function createProviderToolFactory<INPUT, ARGS extends object>({
+export function createProviderToolFactory<
+  CONTEXT extends Partial<ContextRegistry>,
+  INPUT,
+  ARGS extends object,
+>({
   id,
   inputSchema,
 }: {
   id: `${string}.${string}`;
   inputSchema: FlexibleSchema<INPUT>;
-}): ProviderToolFactory<INPUT, ARGS> {
+}): ProviderToolFactory<CONTEXT, INPUT, ARGS> {
   return <OUTPUT>({
     execute,
     outputSchema,
@@ -29,14 +38,14 @@ export function createProviderToolFactory<INPUT, ARGS extends object>({
     onInputAvailable,
     ...args
   }: ARGS & {
-    execute?: ToolExecuteFunction<INPUT, OUTPUT>;
+    execute?: ToolExecuteFunction<CONTEXT, INPUT, OUTPUT>;
     outputSchema?: FlexibleSchema<OUTPUT>;
-    needsApproval?: Tool<INPUT, OUTPUT>['needsApproval'];
-    toModelOutput?: Tool<INPUT, OUTPUT>['toModelOutput'];
-    onInputStart?: Tool<INPUT, OUTPUT>['onInputStart'];
-    onInputDelta?: Tool<INPUT, OUTPUT>['onInputDelta'];
-    onInputAvailable?: Tool<INPUT, OUTPUT>['onInputAvailable'];
-  }): Tool<INPUT, OUTPUT> =>
+    needsApproval?: Tool<CONTEXT, INPUT, OUTPUT>['needsApproval'];
+    toModelOutput?: Tool<CONTEXT, INPUT, OUTPUT>['toModelOutput'];
+    onInputStart?: Tool<CONTEXT, INPUT, OUTPUT>['onInputStart'];
+    onInputDelta?: Tool<CONTEXT, INPUT, OUTPUT>['onInputDelta'];
+    onInputAvailable?: Tool<CONTEXT, INPUT, OUTPUT>['onInputAvailable'];
+  }): Tool<CONTEXT, INPUT, OUTPUT> =>
     tool({
       type: 'provider',
       id,
@@ -53,21 +62,23 @@ export function createProviderToolFactory<INPUT, ARGS extends object>({
 }
 
 export type ProviderToolFactoryWithOutputSchema<
+  CONTEXT extends Partial<ContextRegistry>,
   INPUT,
   OUTPUT,
   ARGS extends object,
 > = (
   options: ARGS & {
-    execute?: ToolExecuteFunction<INPUT, OUTPUT>;
-    needsApproval?: Tool<INPUT, OUTPUT>['needsApproval'];
-    toModelOutput?: Tool<INPUT, OUTPUT>['toModelOutput'];
-    onInputStart?: Tool<INPUT, OUTPUT>['onInputStart'];
-    onInputDelta?: Tool<INPUT, OUTPUT>['onInputDelta'];
-    onInputAvailable?: Tool<INPUT, OUTPUT>['onInputAvailable'];
+    execute?: ToolExecuteFunction<CONTEXT, INPUT, OUTPUT>;
+    needsApproval?: Tool<CONTEXT, INPUT, OUTPUT>['needsApproval'];
+    toModelOutput?: Tool<CONTEXT, INPUT, OUTPUT>['toModelOutput'];
+    onInputStart?: Tool<CONTEXT, INPUT, OUTPUT>['onInputStart'];
+    onInputDelta?: Tool<CONTEXT, INPUT, OUTPUT>['onInputDelta'];
+    onInputAvailable?: Tool<CONTEXT, INPUT, OUTPUT>['onInputAvailable'];
   },
-) => Tool<INPUT, OUTPUT>;
+) => Tool<CONTEXT, INPUT, OUTPUT>;
 
 export function createProviderToolFactoryWithOutputSchema<
+  CONTEXT extends Partial<ContextRegistry>,
   INPUT,
   OUTPUT,
   ARGS extends object,
@@ -91,7 +102,7 @@ export function createProviderToolFactoryWithOutputSchema<
    * @default false
    */
   supportsDeferredResults?: boolean;
-}): ProviderToolFactoryWithOutputSchema<INPUT, OUTPUT, ARGS> {
+}): ProviderToolFactoryWithOutputSchema<CONTEXT, INPUT, OUTPUT, ARGS> {
   return ({
     execute,
     needsApproval,
@@ -101,13 +112,13 @@ export function createProviderToolFactoryWithOutputSchema<
     onInputAvailable,
     ...args
   }: ARGS & {
-    execute?: ToolExecuteFunction<INPUT, OUTPUT>;
-    needsApproval?: Tool<INPUT, OUTPUT>['needsApproval'];
-    toModelOutput?: Tool<INPUT, OUTPUT>['toModelOutput'];
-    onInputStart?: Tool<INPUT, OUTPUT>['onInputStart'];
-    onInputDelta?: Tool<INPUT, OUTPUT>['onInputDelta'];
-    onInputAvailable?: Tool<INPUT, OUTPUT>['onInputAvailable'];
-  }): Tool<INPUT, OUTPUT> =>
+    execute?: ToolExecuteFunction<CONTEXT, INPUT, OUTPUT>;
+    needsApproval?: Tool<CONTEXT, INPUT, OUTPUT>['needsApproval'];
+    toModelOutput?: Tool<CONTEXT, INPUT, OUTPUT>['toModelOutput'];
+    onInputStart?: Tool<CONTEXT, INPUT, OUTPUT>['onInputStart'];
+    onInputDelta?: Tool<CONTEXT, INPUT, OUTPUT>['onInputDelta'];
+    onInputAvailable?: Tool<CONTEXT, INPUT, OUTPUT>['onInputAvailable'];
+  }): Tool<CONTEXT, INPUT, OUTPUT> =>
     tool({
       type: 'provider',
       id,
diff --git a/packages/provider-utils/src/types/execute-tool.ts b/packages/provider-utils/src/types/execute-tool.ts
index 7780d582f4ce..4220dbd43744 100644
--- a/packages/provider-utils/src/types/execute-tool.ts
+++ b/packages/provider-utils/src/types/execute-tool.ts
@@ -1,14 +1,19 @@
 import { isAsyncIterable } from '../is-async-iterable';
+import { ContextRegistry } from './context';
 import { ToolExecutionOptions, ToolExecuteFunction } from './tool';
 
-export async function* executeTool<INPUT, OUTPUT>({
+export async function* executeTool<
+  CONTEXT extends Partial<ContextRegistry>,
+  INPUT,
+  OUTPUT,
+>({
   execute,
   input,
   options,
 }: {
-  execute: ToolExecuteFunction<INPUT, OUTPUT>;
+  execute: ToolExecuteFunction<CONTEXT, INPUT, OUTPUT>;
   input: INPUT;
-  options: ToolExecutionOptions;
+  options: ToolExecutionOptions<CONTEXT>;
 }): AsyncGenerator<
   { type: 'preliminary'; output: OUTPUT } | { type: 'final'; output: OUTPUT }
 > {
diff --git a/packages/provider-utils/src/types/tool.test-d.ts b/packages/provider-utils/src/types/tool.test-d.ts
index b09178639534..76318d7d804d 100644
--- a/packages/provider-utils/src/types/tool.test-d.ts
+++ b/packages/provider-utils/src/types/tool.test-d.ts
@@ -13,7 +13,9 @@ describe('tool type', () => {
         inputSchema: z.object({ number: z.number() }),
       });
 
-      expectTypeOf(aTool).toEqualTypeOf<Tool<{ number: number }, never>>();
+      expectTypeOf(aTool).toEqualTypeOf<
+        Tool<ContextRegistry, { number: number }, never>
+      >();
       expectTypeOf(aTool.execute).toEqualTypeOf<undefined>();
       expectTypeOf(aTool.execute).not.toEqualTypeOf<Function>();
       expectTypeOf(aTool.inputSchema).toEqualTypeOf<
@@ -26,7 +28,7 @@ describe('tool type', () => {
         inputSchema: null as unknown as FlexibleSchema<T>,
       });
 
-      expectTypeOf(aTool).toEqualTypeOf<Tool<T, never>>();
+      expectTypeOf(aTool).toEqualTypeOf<Tool<ContextRegistry, T, never>>();
       expectTypeOf(aTool.execute).toEqualTypeOf<undefined>();
       expectTypeOf(aTool.execute).not.toEqualTypeOf<Function>();
       expectTypeOf(aTool.inputSchema).toEqualTypeOf<FlexibleSchema<T>>();
@@ -78,9 +80,12 @@ describe('tool type', () => {
         },
       });
 
-      expectTypeOf(aTool).toEqualTypeOf<Tool<{ number: number }, 'test'>>();
+      expectTypeOf(aTool).toEqualTypeOf<
+        Tool<ContextRegistry, { number: number }, 'test'>
+      >();
       expectTypeOf(aTool.execute).toMatchTypeOf<
-        ToolExecuteFunction<{ number: number }, 'test'> | undefined
+        | ToolExecuteFunction<ContextRegistry, { number: number }, 'test'>
+        | undefined
       >();
       expectTypeOf(aTool.execute).not.toEqualTypeOf<undefined>();
       expectTypeOf(aTool.inputSchema).toEqualTypeOf<
@@ -96,9 +101,12 @@ describe('tool type', () => {
         },
       });
 
-      expectTypeOf(aTool).toEqualTypeOf<Tool<{ number: number }, 'test'>>();
+      expectTypeOf(aTool).toEqualTypeOf<
+        Tool<ContextRegistry, { number: number }, 'test'>
+      >();
       expectTypeOf(aTool.execute).toEqualTypeOf<
-        ToolExecuteFunction<{ number: number }, 'test'> | undefined
+        | ToolExecuteFunction<ContextRegistry, { number: number }, 'test'>
+        | undefined
       >();
       expectTypeOf(aTool.inputSchema).toEqualTypeOf<
         FlexibleSchema<{ number: number }>
diff --git a/packages/provider-utils/src/types/tool.ts b/packages/provider-utils/src/types/tool.ts
index 68c97f2e73a0..625b0df8af9f 100644
--- a/packages/provider-utils/src/types/tool.ts
+++ b/packages/provider-utils/src/types/tool.ts
@@ -39,7 +39,7 @@ export interface ToolExecutionOptions<
    *
    * Experimental (can break in patch releases).
    */
-  experimental_context?: Context<CONTEXT>;
+  experimental_context: Context<CONTEXT>;
 }
 
 /**
@@ -72,9 +72,9 @@ export type ToolNeedsApprovalFunction<
 ) => boolean | PromiseLike<boolean>;
 
 export type ToolExecuteFunction<
+  CONTEXT extends Partial<ContextRegistry>,
   INPUT,
   OUTPUT,
-  CONTEXT extends Partial<ContextRegistry> = ContextRegistry,
 > = (
   input: INPUT,
   options: ToolExecutionOptions<CONTEXT>,
@@ -88,7 +88,11 @@ type NeverOptional<N, T> = 0 extends 1 & N
     ? Partial<Record<keyof T, undefined>>
     : T;
 
-type ToolOutputProperties<INPUT, OUTPUT> = NeverOptional<
+type ToolOutputProperties<
+  CONTEXT extends Partial<ContextRegistry>,
+  INPUT,
+  OUTPUT,
+> = NeverOptional<
   OUTPUT,
   | {
       /**
@@ -98,7 +102,7 @@ type ToolOutputProperties<INPUT, OUTPUT> = NeverOptional<
        * @args is the input of the tool call.
        * @options.abortSignal is a signal that can be used to abort the tool call.
        */
-      execute: ToolExecuteFunction<INPUT, OUTPUT>;
+      execute: ToolExecuteFunction<CONTEXT, INPUT, OUTPUT>;
 
       outputSchema?: FlexibleSchema<OUTPUT>;
     }
@@ -116,6 +120,7 @@ type ToolOutputProperties<INPUT, OUTPUT> = NeverOptional<
  * The tool can also contain an optional execute function for the actual execution function of the tool.
  */
 export type Tool<
+  CONTEXT extends Partial<ContextRegistry> = ContextRegistry,
   INPUT extends JSONValue | unknown | never = any,
   OUTPUT extends JSONValue | unknown | never = any,
 > = {
@@ -153,6 +158,8 @@ export type Tool<
    */
   inputExamples?: Array<{ input: NoInfer<INPUT> }>;
 
+  contextSchema?: FlexibleSchema<CONTEXT>;
+
   /**
    * Whether the tool needs approval before it can be executed.
    */
@@ -192,7 +199,7 @@ export type Tool<
       input: [INPUT] extends [never] ? unknown : INPUT;
     } & ToolExecutionOptions,
   ) => void | PromiseLike<void>;
-} & ToolOutputProperties<INPUT, OUTPUT> & {
+} & ToolOutputProperties<CONTEXT, INPUT, OUTPUT> & {
     /**
      * Optional conversion function that maps the tool result to an output that can be used by the language model.
      *
@@ -268,25 +275,35 @@ export type Tool<
 /**
  * Infer the input type of a tool.
  */
-export type InferToolInput<TOOL extends Tool> =
-  TOOL extends Tool<infer INPUT, any> ? INPUT : never;
+export type InferToolInput<
+  CONTEXT extends Partial<ContextRegistry>,
+  TOOL extends Tool<CONTEXT>,
+> = TOOL extends Tool<CONTEXT, infer INPUT, any> ? INPUT : never;
 
 /**
  * Infer the output type of a tool.
  */
-export type InferToolOutput<TOOL extends Tool> =
-  TOOL extends Tool<any, infer OUTPUT> ? OUTPUT : never;
+export type InferToolOutput<
+  CONTEXT extends Partial<ContextRegistry>,
+  TOOL extends Tool<CONTEXT>,
+> = TOOL extends Tool<CONTEXT, any, infer OUTPUT> ? OUTPUT : never;
 
 /**
  * Helper function for inferring the execute args of a tool.
  */
 // Note: overload order is important for auto-completion
-export function tool<INPUT, OUTPUT>(
-  tool: Tool<INPUT, OUTPUT>,
-): Tool<INPUT, OUTPUT>;
-export function tool<INPUT>(tool: Tool<INPUT, never>): Tool<INPUT, never>;
-export function tool<OUTPUT>(tool: Tool<never, OUTPUT>): Tool<never, OUTPUT>;
-export function tool(tool: Tool<never, never>): Tool<never, never>;
+export function tool<CONTEXT extends Partial<ContextRegistry>, INPUT, OUTPUT>(
+  tool: Tool<CONTEXT, INPUT, OUTPUT>,
+): Tool<CONTEXT, INPUT, OUTPUT>;
+export function tool<CONTEXT extends Partial<ContextRegistry>, INPUT>(
+  tool: Tool<CONTEXT, INPUT, never>,
+): Tool<CONTEXT, INPUT, never>;
+export function tool<CONTEXT extends Partial<ContextRegistry>, OUTPUT>(
+  tool: Tool<CONTEXT, never, OUTPUT>,
+): Tool<CONTEXT, never, OUTPUT>;
+export function tool<CONTEXT extends Partial<ContextRegistry>>(
+  tool: Tool<CONTEXT, never, never>,
+): Tool<CONTEXT, never, never>;
 export function tool(tool: any): any {
   return tool;
 }
@@ -299,7 +316,7 @@ export function dynamicTool(tool: {
   title?: string;
   providerOptions?: ProviderOptions;
   inputSchema: FlexibleSchema<unknown>;
-  execute: ToolExecuteFunction<unknown, unknown>;
+  execute: ToolExecuteFunction<ContextRegistry, unknown, unknown>;
 
   /**
    * Optional conversion function that maps the tool result to an output that can be used by the language model.
@@ -327,7 +344,7 @@ export function dynamicTool(tool: {
    * Whether the tool needs approval before it can be executed.
    */
   needsApproval?: boolean | ToolNeedsApprovalFunction<unknown>;
-}): Tool<unknown, unknown> & {
+}): Tool<ContextRegistry, unknown, unknown> & {
   type: 'dynamic';
 } {
   return { ...tool, type: 'dynamic' };

From eaf33e252b2c143592e2dc1f2204dec13052213f Mon Sep 17 00:00:00 2001
From: Lars Grammel <lars.grammel@gmail.com>
Date: Thu, 5 Mar 2026 15:40:10 +0100
Subject: [PATCH 03/41] proof of concept

---
 .../ai-functions/src/test/typed-context.ts    | 42 +++++++++++++++++++
 1 file changed, 42 insertions(+)
 create mode 100644 examples/ai-functions/src/test/typed-context.ts

diff --git a/examples/ai-functions/src/test/typed-context.ts b/examples/ai-functions/src/test/typed-context.ts
new file mode 100644
index 000000000000..dd53f7cd0bf2
--- /dev/null
+++ b/examples/ai-functions/src/test/typed-context.ts
@@ -0,0 +1,42 @@
+import { run } from '../lib/run';
+
+interface Tool<CONTEXT> {
+  execute: (context: CONTEXT) => Promise<void>;
+}
+
+export type ToolSet<CONTEXT> = Record<string, Tool<CONTEXT>>;
+
+function executeTool<CONTEXT>({
+  tools,
+  toolName,
+  context,
+}: {
+  tools: ToolSet<CONTEXT>;
+  toolName: keyof ToolSet<CONTEXT>;
+  context: CONTEXT;
+}) {
+  return tools[toolName].execute(context);
+}
+
+run(async () => {
+  const tool1: Tool<{ name: string }> = {
+    execute: async context => {
+      console.log(context);
+    },
+  };
+
+  const tool2: Tool<{ age: number }> = {
+    execute: async context => {
+      console.log(context);
+    },
+  };
+
+  executeTool({
+    tools: {
+      tool1,
+      tool2,
+    },
+    toolName: 'tool1',
+    context: { name: 'John', age: 30 },
+  });
+});

From d67d969125424768ba882ec7dc164753486bf1de Mon Sep 17 00:00:00 2001
From: Lars Grammel <lars.grammel@gmail.com>
Date: Thu, 5 Mar 2026 15:45:36 +0100
Subject: [PATCH 04/41] s

---
 .../src/provider-tool-factory.ts              | 10 ++---
 packages/provider-utils/src/types/context.ts  |  6 +--
 .../provider-utils/src/types/execute-tool.ts  | 10 ++---
 packages/provider-utils/src/types/index.ts    |  2 +-
 .../provider-utils/src/types/tool.test-d.ts   | 24 ++++++------
 packages/provider-utils/src/types/tool.ts     | 38 ++++++++-----------
 6 files changed, 37 insertions(+), 53 deletions(-)

diff --git a/packages/provider-utils/src/provider-tool-factory.ts b/packages/provider-utils/src/provider-tool-factory.ts
index 7595d518c8d9..73590059e9e9 100644
--- a/packages/provider-utils/src/provider-tool-factory.ts
+++ b/packages/provider-utils/src/provider-tool-factory.ts
@@ -1,9 +1,9 @@
 import { tool, Tool, ToolExecuteFunction } from './types/tool';
 import { FlexibleSchema } from './schema';
-import { ContextRegistry } from './types/context';
+import { Context } from './types/context';
 
 export type ProviderToolFactory<
-  CONTEXT extends Partial<ContextRegistry>,
+  CONTEXT extends Context,
   INPUT,
   ARGS extends object,
 > = <OUTPUT>(
@@ -18,7 +18,7 @@ export type ProviderToolFactory<
 ) => Tool<CONTEXT, INPUT, OUTPUT>;
 
 export function createProviderToolFactory<
-  CONTEXT extends Partial<ContextRegistry>,
+  CONTEXT extends Context,
   INPUT,
   ARGS extends object,
 >({
@@ -62,7 +62,7 @@ export function createProviderToolFactory<
 }
 
 export type ProviderToolFactoryWithOutputSchema<
-  CONTEXT extends Partial<ContextRegistry>,
+  CONTEXT extends Context,
   INPUT,
   OUTPUT,
   ARGS extends object,
@@ -78,7 +78,7 @@ export type ProviderToolFactoryWithOutputSchema<
 ) => Tool<CONTEXT, INPUT, OUTPUT>;
 
 export function createProviderToolFactoryWithOutputSchema<
-  CONTEXT extends Partial<ContextRegistry>,
+  CONTEXT extends Context,
   INPUT,
   OUTPUT,
   ARGS extends object,
diff --git a/packages/provider-utils/src/types/context.ts b/packages/provider-utils/src/types/context.ts
index 92927870b874..9a1d29a311fd 100644
--- a/packages/provider-utils/src/types/context.ts
+++ b/packages/provider-utils/src/types/context.ts
@@ -1,5 +1 @@
-export interface ContextRegistry {} // open via declaration merging
-
-export type Context<T extends Partial<ContextRegistry> = ContextRegistry> = {
-  [K in keyof T]: T[K];
-};
+export type Context = Record<string, unknown>;
diff --git a/packages/provider-utils/src/types/execute-tool.ts b/packages/provider-utils/src/types/execute-tool.ts
index 4220dbd43744..68961ccb9c34 100644
--- a/packages/provider-utils/src/types/execute-tool.ts
+++ b/packages/provider-utils/src/types/execute-tool.ts
@@ -1,12 +1,8 @@
 import { isAsyncIterable } from '../is-async-iterable';
-import { ContextRegistry } from './context';
-import { ToolExecutionOptions, ToolExecuteFunction } from './tool';
+import { Context } from './context';
+import { ToolExecuteFunction, ToolExecutionOptions } from './tool';
 
-export async function* executeTool<
-  CONTEXT extends Partial<ContextRegistry>,
-  INPUT,
-  OUTPUT,
->({
+export async function* executeTool<CONTEXT extends Context, INPUT, OUTPUT>({
   execute,
   input,
   options,
diff --git a/packages/provider-utils/src/types/index.ts b/packages/provider-utils/src/types/index.ts
index eb3ead960a9d..b1fca4b85201 100644
--- a/packages/provider-utils/src/types/index.ts
+++ b/packages/provider-utils/src/types/index.ts
@@ -11,7 +11,7 @@ export type {
   ToolResultOutput,
   ToolResultPart,
 } from './content-part';
-export type { Context, ContextRegistry } from './context';
+export type { Context } from './context';
 export type { DataContent } from './data-content';
 export { executeTool } from './execute-tool';
 export type { ModelMessage } from './model-message';
diff --git a/packages/provider-utils/src/types/tool.test-d.ts b/packages/provider-utils/src/types/tool.test-d.ts
index 76318d7d804d..8a3e9d1fba96 100644
--- a/packages/provider-utils/src/types/tool.test-d.ts
+++ b/packages/provider-utils/src/types/tool.test-d.ts
@@ -2,7 +2,7 @@ import { describe, expectTypeOf, it } from 'vitest';
 import { z } from 'zod/v4';
 import { FlexibleSchema } from '../schema';
 import { ToolResultOutput } from './content-part';
-import { Context, ContextRegistry } from './context';
+import { Context } from './context';
 import { ModelMessage } from './model-message';
 import { Tool, tool, ToolExecuteFunction } from './tool';
 
@@ -14,7 +14,7 @@ describe('tool type', () => {
       });
 
       expectTypeOf(aTool).toEqualTypeOf<
-        Tool<ContextRegistry, { number: number }, never>
+        Tool<Context, { number: number }, never>
       >();
       expectTypeOf(aTool.execute).toEqualTypeOf<undefined>();
       expectTypeOf(aTool.execute).not.toEqualTypeOf<Function>();
@@ -28,7 +28,7 @@ describe('tool type', () => {
         inputSchema: null as unknown as FlexibleSchema<T>,
       });
 
-      expectTypeOf(aTool).toEqualTypeOf<Tool<ContextRegistry, T, never>>();
+      expectTypeOf(aTool).toEqualTypeOf<Tool<Context, T, never>>();
       expectTypeOf(aTool.execute).toEqualTypeOf<undefined>();
       expectTypeOf(aTool.execute).not.toEqualTypeOf<Function>();
       expectTypeOf(aTool.inputSchema).toEqualTypeOf<FlexibleSchema<T>>();
@@ -81,11 +81,10 @@ describe('tool type', () => {
       });
 
       expectTypeOf(aTool).toEqualTypeOf<
-        Tool<ContextRegistry, { number: number }, 'test'>
+        Tool<Context, { number: number }, 'test'>
       >();
       expectTypeOf(aTool.execute).toMatchTypeOf<
-        | ToolExecuteFunction<ContextRegistry, { number: number }, 'test'>
-        | undefined
+        ToolExecuteFunction<Context, { number: number }, 'test'> | undefined
       >();
       expectTypeOf(aTool.execute).not.toEqualTypeOf<undefined>();
       expectTypeOf(aTool.inputSchema).toEqualTypeOf<
@@ -102,11 +101,10 @@ describe('tool type', () => {
       });
 
       expectTypeOf(aTool).toEqualTypeOf<
-        Tool<ContextRegistry, { number: number }, 'test'>
+        Tool<Context, { number: number }, 'test'>
       >();
       expectTypeOf(aTool.execute).toEqualTypeOf<
-        | ToolExecuteFunction<ContextRegistry, { number: number }, 'test'>
-        | undefined
+        ToolExecuteFunction<Context, { number: number }, 'test'> | undefined
       >();
       expectTypeOf(aTool.inputSchema).toEqualTypeOf<
         FlexibleSchema<{ number: number }>
@@ -184,7 +182,7 @@ describe('tool type', () => {
           expectTypeOf(options).toEqualTypeOf<{
             toolCallId: string;
             messages: ModelMessage[];
-            experimental_context?: Context<ContextRegistry> | undefined;
+            experimental_context?: Context | undefined;
           }>();
           return true;
         },
@@ -197,7 +195,7 @@ describe('tool type', () => {
             options: {
               toolCallId: string;
               messages: ModelMessage[];
-              experimental_context?: Context<ContextRegistry> | undefined;
+              experimental_context?: Context | undefined;
             },
           ) => boolean | PromiseLike<boolean>)
         | undefined
@@ -213,7 +211,7 @@ describe('tool type', () => {
           expectTypeOf(options).toEqualTypeOf<{
             toolCallId: string;
             messages: ModelMessage[];
-            experimental_context?: Context<ContextRegistry> | undefined;
+            experimental_context?: Context | undefined;
           }>();
           return true;
         },
@@ -226,7 +224,7 @@ describe('tool type', () => {
             options: {
               toolCallId: string;
               messages: ModelMessage[];
-              experimental_context?: Context<ContextRegistry> | undefined;
+              experimental_context?: Context | undefined;
             },
           ) => boolean | PromiseLike<boolean>)
         | undefined
diff --git a/packages/provider-utils/src/types/tool.ts b/packages/provider-utils/src/types/tool.ts
index 625b0df8af9f..5785670b5332 100644
--- a/packages/provider-utils/src/types/tool.ts
+++ b/packages/provider-utils/src/types/tool.ts
@@ -3,14 +3,12 @@ import { FlexibleSchema } from '../schema';
 import { ToolResultOutput } from './content-part';
 import { ModelMessage } from './model-message';
 import { ProviderOptions } from './provider-options';
-import { Context, ContextRegistry } from './context';
+import { Context } from './context';
 
 /**
  * Additional options that are sent into each tool call.
  */
-export interface ToolExecutionOptions<
-  CONTEXT extends Partial<ContextRegistry> = ContextRegistry,
-> {
+export interface ToolExecutionOptions<CONTEXT extends Context = Context> {
   /**
    * The ID of the tool call. You can use it e.g. when sending tool-call related information with stream data.
    */
@@ -39,7 +37,7 @@ export interface ToolExecutionOptions<
    *
    * Experimental (can break in patch releases).
    */
-  experimental_context: Context<CONTEXT>;
+  experimental_context: CONTEXT;
 }
 
 /**
@@ -47,7 +45,7 @@ export interface ToolExecutionOptions<
  */
 export type ToolNeedsApprovalFunction<
   INPUT,
-  CONTEXT extends Partial<ContextRegistry> = ContextRegistry,
+  CONTEXT extends Context = Context,
 > = (
   input: INPUT,
   options: {
@@ -67,15 +65,11 @@ export type ToolNeedsApprovalFunction<
      *
      * Experimental (can break in patch releases).
      */
-    experimental_context?: Context<CONTEXT>;
+    experimental_context?: CONTEXT;
   },
 ) => boolean | PromiseLike<boolean>;
 
-export type ToolExecuteFunction<
-  CONTEXT extends Partial<ContextRegistry>,
-  INPUT,
-  OUTPUT,
-> = (
+export type ToolExecuteFunction<CONTEXT extends Context, INPUT, OUTPUT> = (
   input: INPUT,
   options: ToolExecutionOptions<CONTEXT>,
 ) => AsyncIterable<OUTPUT> | PromiseLike<OUTPUT> | OUTPUT;
@@ -89,7 +83,7 @@ type NeverOptional<N, T> = 0 extends 1 & N
     : T;
 
 type ToolOutputProperties<
-  CONTEXT extends Partial<ContextRegistry>,
+  CONTEXT extends Context,
   INPUT,
   OUTPUT,
 > = NeverOptional<
@@ -120,7 +114,7 @@ type ToolOutputProperties<
  * The tool can also contain an optional execute function for the actual execution function of the tool.
  */
 export type Tool<
-  CONTEXT extends Partial<ContextRegistry> = ContextRegistry,
+  CONTEXT extends Context,
   INPUT extends JSONValue | unknown | never = any,
   OUTPUT extends JSONValue | unknown | never = any,
 > = {
@@ -276,7 +270,7 @@ export type Tool<
  * Infer the input type of a tool.
  */
 export type InferToolInput<
-  CONTEXT extends Partial<ContextRegistry>,
+  CONTEXT extends Context,
   TOOL extends Tool<CONTEXT>,
 > = TOOL extends Tool<CONTEXT, infer INPUT, any> ? INPUT : never;
 
@@ -284,7 +278,7 @@ export type InferToolInput<
  * Infer the output type of a tool.
  */
 export type InferToolOutput<
-  CONTEXT extends Partial<ContextRegistry>,
+  CONTEXT extends Context,
   TOOL extends Tool<CONTEXT>,
 > = TOOL extends Tool<CONTEXT, any, infer OUTPUT> ? OUTPUT : never;
 
@@ -292,16 +286,16 @@ export type InferToolOutput<
  * Helper function for inferring the execute args of a tool.
  */
 // Note: overload order is important for auto-completion
-export function tool<CONTEXT extends Partial<ContextRegistry>, INPUT, OUTPUT>(
+export function tool<CONTEXT extends Context, INPUT, OUTPUT>(
   tool: Tool<CONTEXT, INPUT, OUTPUT>,
 ): Tool<CONTEXT, INPUT, OUTPUT>;
-export function tool<CONTEXT extends Partial<ContextRegistry>, INPUT>(
+export function tool<CONTEXT extends Context, INPUT>(
   tool: Tool<CONTEXT, INPUT, never>,
 ): Tool<CONTEXT, INPUT, never>;
-export function tool<CONTEXT extends Partial<ContextRegistry>, OUTPUT>(
+export function tool<CONTEXT extends Context, OUTPUT>(
   tool: Tool<CONTEXT, never, OUTPUT>,
 ): Tool<CONTEXT, never, OUTPUT>;
-export function tool<CONTEXT extends Partial<ContextRegistry>>(
+export function tool<CONTEXT extends Context>(
   tool: Tool<CONTEXT, never, never>,
 ): Tool<CONTEXT, never, never>;
 export function tool(tool: any): any {
@@ -316,7 +310,7 @@ export function dynamicTool(tool: {
   title?: string;
   providerOptions?: ProviderOptions;
   inputSchema: FlexibleSchema<unknown>;
-  execute: ToolExecuteFunction<ContextRegistry, unknown, unknown>;
+  execute: ToolExecuteFunction<Context, unknown, unknown>;
 
   /**
    * Optional conversion function that maps the tool result to an output that can be used by the language model.
@@ -344,7 +338,7 @@ export function dynamicTool(tool: {
    * Whether the tool needs approval before it can be executed.
    */
   needsApproval?: boolean | ToolNeedsApprovalFunction<unknown>;
-}): Tool<ContextRegistry, unknown, unknown> & {
+}): Tool<Context, unknown, unknown> & {
   type: 'dynamic';
 } {
   return { ...tool, type: 'dynamic' };

From 9f48c008d8251bf34f0d129c534f58e4b9acf0ba Mon Sep 17 00:00:00 2001
From: Lars Grammel <lars.grammel@gmail.com>
Date: Thu, 5 Mar 2026 16:00:54 +0100
Subject: [PATCH 05/41] x

---
 .../ai/src/generate-text/callback-events.ts   | 63 +++++++++----------
 .../ai/src/generate-text/execute-tool-call.ts |  9 ++-
 .../ai/src/generate-text/generate-text.ts     | 32 +++++-----
 .../ai/src/generate-text/stop-condition.ts    |  8 +--
 packages/ai/src/generate-text/tool-call.ts    | 18 ++----
 packages/ai/src/generate-text/tool-result.ts  | 22 +++----
 packages/ai/src/generate-text/tool-set.ts     |  6 +-
 packages/provider-utils/src/types/index.ts    |  5 +-
 packages/provider-utils/src/types/tool.ts     | 36 +++++------
 9 files changed, 86 insertions(+), 113 deletions(-)

diff --git a/packages/ai/src/generate-text/callback-events.ts b/packages/ai/src/generate-text/callback-events.ts
index afb35b9282fb..454be6f96a9a 100644
--- a/packages/ai/src/generate-text/callback-events.ts
+++ b/packages/ai/src/generate-text/callback-events.ts
@@ -1,6 +1,5 @@
 import type { LanguageModelV4ToolChoice } from '@ai-sdk/provider';
 import type {
-  ContextRegistry,
   ModelMessage,
   ProviderOptions,
   SystemModelMessage,
@@ -130,8 +129,7 @@ export interface OnStartEvent<
  * Each step represents a single LLM invocation.
  */
 export interface OnStepStartEvent<
-  CONTEXT extends Partial<ContextRegistry>,
-  TOOLS extends ToolSet<CONTEXT> = ToolSet<CONTEXT>,
+  TOOLS extends ToolSet = ToolSet,
   OUTPUT extends Output = Output,
   INCLUDE = { requestBody?: boolean; responseBody?: boolean },
 > {
@@ -186,8 +184,8 @@ export interface OnStepStartEvent<
    * When the condition is an array, any of the conditions can be met to stop.
    */
   readonly stopWhen:
-    | StopCondition<CONTEXT, TOOLS>
-    | Array<StopCondition<CONTEXT, TOOLS>>
+    | StopCondition<TOOLS>
+    | Array<StopCondition<TOOLS>>
     | undefined;
 
   /** The output specification for structured outputs, if configured. */
@@ -218,10 +216,7 @@ export interface OnStepStartEvent<
  *
  * Called when a tool execution begins, before the tool's `execute` function is invoked.
  */
-export interface OnToolCallStartEvent<
-  CONTEXT extends Partial<ContextRegistry>,
-  TOOLS extends ToolSet<CONTEXT> = ToolSet<CONTEXT>,
-> {
+export interface OnToolCallStartEvent<TOOLS extends ToolSet = ToolSet> {
   /** Zero-based index of the current step where this tool call occurs. */
   readonly stepNumber: number | undefined;
 
@@ -229,7 +224,7 @@ export interface OnToolCallStartEvent<
   readonly model: CallbackModelInfo | undefined;
 
   /** The full tool call object. */
-  readonly toolCall: TypedToolCall<CONTEXT, TOOLS>;
+  readonly toolCall: TypedToolCall<TOOLS>;
 
   /** The conversation messages available at tool execution time. */
   readonly messages: Array<ModelMessage>;
@@ -312,28 +307,26 @@ export type OnStepFinishEvent<TOOLS extends ToolSet = ToolSet> =
  * Called when the entire generation completes (all steps finished).
  * Includes the final step's result along with aggregated data from all steps.
  */
-export type OnFinishEvent<
-  CONTEXT extends Partial<ContextRegistry>,
-  TOOLS extends ToolSet<CONTEXT> = ToolSet<CONTEXT>,
-> = StepResult<CONTEXT, TOOLS> & {
-  /** Array containing results from all steps in the generation. */
-  readonly steps: StepResult<CONTEXT, TOOLS>[];
-
-  /** Aggregated token usage across all steps. */
-  readonly totalUsage: LanguageModelUsage;
-
-  /**
-   * The final state of the user-defined context object.
-   *
-   * Experimental (can break in patch releases).
-   *
-   * @default undefined
-   */
-  experimental_context: unknown;
-
-  /** Identifier from telemetry settings for grouping related operations. */
-  readonly functionId: string | undefined;
-
-  /** Additional metadata from telemetry settings. */
-  readonly metadata: Record<string, unknown> | undefined;
-};
+export type OnFinishEvent<TOOLS extends ToolSet = ToolSet> =
+  StepResult<TOOLS> & {
+    /** Array containing results from all steps in the generation. */
+    readonly steps: StepResult<TOOLS>[];
+
+    /** Aggregated token usage across all steps. */
+    readonly totalUsage: LanguageModelUsage;
+
+    /**
+     * The final state of the user-defined context object.
+     *
+     * Experimental (can break in patch releases).
+     *
+     * @default undefined
+     */
+    experimental_context: unknown;
+
+    /** Identifier from telemetry settings for grouping related operations. */
+    readonly functionId: string | undefined;
+
+    /** Additional metadata from telemetry settings. */
+    readonly metadata: Record<string, unknown> | undefined;
+  };
diff --git a/packages/ai/src/generate-text/execute-tool-call.ts b/packages/ai/src/generate-text/execute-tool-call.ts
index 55644265c277..80a311f0c90d 100644
--- a/packages/ai/src/generate-text/execute-tool-call.ts
+++ b/packages/ai/src/generate-text/execute-tool-call.ts
@@ -1,4 +1,4 @@
-import { executeTool, ModelMessage } from '@ai-sdk/provider-utils';
+import { Context, executeTool, ModelMessage } from '@ai-sdk/provider-utils';
 import { Tracer } from '@opentelemetry/api';
 import { notify } from '../util/notify';
 import { assembleOperationName } from '../telemetry/assemble-operation-name';
@@ -27,7 +27,10 @@ import { TypedToolError } from './tool-error';
  *
  * @returns The tool output (result or error), or undefined if the tool has no execute function.
  */
-export async function executeToolCall<TOOLS extends ToolSet>({
+export async function executeToolCall<
+  CONTEXT extends Context,
+  TOOLS extends ToolSet<CONTEXT>,
+>({
   toolCall,
   tools,
   tracer,
@@ -47,7 +50,7 @@ export async function executeToolCall<TOOLS extends ToolSet>({
   telemetry: TelemetrySettings | undefined;
   messages: ModelMessage[];
   abortSignal: AbortSignal | undefined;
-  experimental_context: unknown;
+  experimental_context: CONTEXT;
   stepNumber?: number;
   model?: { provider: string; modelId: string };
   onPreliminaryToolResult?: (result: TypedToolResult<TOOLS>) => void;
diff --git a/packages/ai/src/generate-text/generate-text.ts b/packages/ai/src/generate-text/generate-text.ts
index e15f8e2113c3..8d9f82ddc4fd 100644
--- a/packages/ai/src/generate-text/generate-text.ts
+++ b/packages/ai/src/generate-text/generate-text.ts
@@ -5,7 +5,6 @@ import {
 } from '@ai-sdk/provider';
 import {
   Context,
-  ContextRegistry,
   createIdGenerator,
   getErrorMessage,
   IdGenerator,
@@ -131,11 +130,10 @@ export type GenerateTextOnStartCallback<
  * @param event - The event object containing step configuration.
  */
 export type GenerateTextOnStepStartCallback<
-  CONTEXT extends Partial<ContextRegistry>,
-  TOOLS extends ToolSet<CONTEXT> = ToolSet<CONTEXT>,
+  TOOLS extends ToolSet = ToolSet,
   OUTPUT extends Output = Output,
 > = (
-  event: OnStepStartEvent<CONTEXT, TOOLS, OUTPUT, GenerateTextIncludeSettings>,
+  event: OnStepStartEvent<TOOLS, OUTPUT, GenerateTextIncludeSettings>,
 ) => PromiseLike<void> | void;
 
 /**
@@ -147,9 +145,8 @@ export type GenerateTextOnStepStartCallback<
  * @param event - The event object containing tool call information.
  */
 export type GenerateTextOnToolCallStartCallback<
-  CONTEXT extends Partial<ContextRegistry>,
-  TOOLS extends ToolSet<CONTEXT> = ToolSet<CONTEXT>,
-> = (event: OnToolCallStartEvent<CONTEXT, TOOLS>) => PromiseLike<void> | void;
+  TOOLS extends ToolSet = ToolSet,
+> = (event: OnToolCallStartEvent<TOOLS>) => PromiseLike<void> | void;
 
 /**
  * Callback that is set using the `experimental_onToolCallFinish` option.
@@ -188,10 +185,9 @@ export type GenerateTextOnStepFinishCallback<TOOLS extends ToolSet> = (
  *
  * @param event - The final result along with aggregated step data.
  */
-export type GenerateTextOnFinishCallback<
-  CONTEXT extends Partial<ContextRegistry>,
-  TOOLS extends ToolSet<CONTEXT> = ToolSet<CONTEXT>,
-> = (event: OnFinishEvent<CONTEXT, TOOLS>) => PromiseLike<void> | void;
+export type GenerateTextOnFinishCallback<TOOLS extends ToolSet = ToolSet> = (
+  event: OnFinishEvent<TOOLS>,
+) => PromiseLike<void> | void;
 
 /**
  * Generate a text and call tools for a given prompt using a language model.
@@ -248,8 +244,8 @@ export type GenerateTextOnFinishCallback<
  * A result object that contains the generated text, the results of the tool calls, and additional information.
  */
 export async function generateText<
-  CONTEXT extends Partial<ContextRegistry>,
-  TOOLS extends ToolSet<CONTEXT>,
+  CONTEXT extends Context,
+  TOOLS extends ToolSet = ToolSet,
   OUTPUT extends Output = Output<string, string>,
 >({
   model: modelArg,
@@ -273,7 +269,7 @@ export async function generateText<
   prepareStep = experimental_prepareStep,
   experimental_repairToolCall: repairToolCall,
   experimental_download: download,
-  experimental_context = {} as Context<CONTEXT>,
+  experimental_context = {} as CONTEXT,
   experimental_include: include,
   _internal: { generateId = originalGenerateId } = {},
   experimental_onStart: onStart,
@@ -413,7 +409,7 @@ export async function generateText<
      *
      * @default undefined
      */
-    experimental_context?: Context<CONTEXT>;
+    experimental_context?: CONTEXT;
 
     /**
      * Settings for controlling what data is included in step results.
@@ -1194,7 +1190,7 @@ export async function generateText<
             onFinish,
             globalTelemetry.onFinish as
               | undefined
-              | GenerateTextOnFinishCallback<CONTEXT, TOOLS>,
+              | GenerateTextOnFinishCallback<TOOLS>,
           ],
         });
 
@@ -1224,7 +1220,7 @@ export async function generateText<
   }
 }
 
-async function executeTools<TOOLS extends ToolSet>({
+async function executeTools<CONTEXT extends Context, TOOLS extends ToolSet>({
   toolCalls,
   tools,
   tracer,
@@ -1243,7 +1239,7 @@ async function executeTools<TOOLS extends ToolSet>({
   telemetry: TelemetrySettings | undefined;
   messages: ModelMessage[];
   abortSignal: AbortSignal | undefined;
-  experimental_context: unknown;
+  experimental_context: CONTEXT;
   stepNumber: number;
   model: { provider: string; modelId: string };
   onToolCallStart:
diff --git a/packages/ai/src/generate-text/stop-condition.ts b/packages/ai/src/generate-text/stop-condition.ts
index 37ce347927f1..a9438b923fc1 100644
--- a/packages/ai/src/generate-text/stop-condition.ts
+++ b/packages/ai/src/generate-text/stop-condition.ts
@@ -1,12 +1,8 @@
-import { ContextRegistry } from '@ai-sdk/provider-utils';
 import { StepResult } from './step-result';
 import { ToolSet } from './tool-set';
 
-export type StopCondition<
-  CONTEXT extends Partial<ContextRegistry>,
-  TOOLS extends ToolSet<CONTEXT> = ToolSet<CONTEXT>,
-> = (options: {
-  steps: Array<StepResult<CONTEXT, TOOLS>>;
+export type StopCondition<TOOLS extends ToolSet = ToolSet> = (options: {
+  steps: Array<StepResult<TOOLS>>;
 }) => PromiseLike<boolean> | boolean;
 
 export function stepCountIs(stepCount: number): StopCondition<any> {
diff --git a/packages/ai/src/generate-text/tool-call.ts b/packages/ai/src/generate-text/tool-call.ts
index b214f62b0405..59698ce0b550 100644
--- a/packages/ai/src/generate-text/tool-call.ts
+++ b/packages/ai/src/generate-text/tool-call.ts
@@ -1,4 +1,4 @@
-import { ContextRegistry, Tool } from '@ai-sdk/provider-utils';
+import { Tool } from '@ai-sdk/provider-utils';
 import { ProviderMetadata } from '../types';
 import { ValueOf } from '../util/value-of';
 import { ToolSet } from './tool-set';
@@ -10,15 +10,10 @@ type BaseToolCall = {
   providerMetadata?: ProviderMetadata;
 };
 
-export type StaticToolCall<
-  CONTEXT extends Partial<ContextRegistry>,
-  TOOLS extends ToolSet<CONTEXT> = ToolSet<CONTEXT>,
-> = ValueOf<{
+export type StaticToolCall<TOOLS extends ToolSet = ToolSet> = ValueOf<{
   [NAME in keyof TOOLS]: BaseToolCall & {
     toolName: NAME & string;
-    input: TOOLS[NAME] extends Tool<CONTEXT, infer PARAMETERS>
-      ? PARAMETERS
-      : never;
+    input: TOOLS[NAME] extends Tool<any, infer PARAMETERS> ? PARAMETERS : never;
     dynamic?: false | undefined;
     invalid?: false | undefined;
     error?: never;
@@ -47,7 +42,6 @@ export type DynamicToolCall = BaseToolCall & {
   error?: unknown;
 };
 
-export type TypedToolCall<
-  CONTEXT extends Partial<ContextRegistry>,
-  TOOLS extends ToolSet<CONTEXT> = ToolSet<CONTEXT>,
-> = StaticToolCall<CONTEXT, TOOLS> | DynamicToolCall;
+export type TypedToolCall<TOOLS extends ToolSet = ToolSet> =
+  | StaticToolCall<TOOLS>
+  | DynamicToolCall;
diff --git a/packages/ai/src/generate-text/tool-result.ts b/packages/ai/src/generate-text/tool-result.ts
index 1f2943f644d0..422f3dafb283 100644
--- a/packages/ai/src/generate-text/tool-result.ts
+++ b/packages/ai/src/generate-text/tool-result.ts
@@ -1,22 +1,15 @@
-import {
-  ContextRegistry,
-  InferToolInput,
-  InferToolOutput,
-} from '@ai-sdk/provider-utils';
+import { InferToolInput, InferToolOutput } from '@ai-sdk/provider-utils';
 import { ProviderMetadata } from '../types';
 import { ValueOf } from '../../src/util/value-of';
 import { ToolSet } from './tool-set';
 
-export type StaticToolResult<
-  CONTEXT extends Partial<ContextRegistry>,
-  TOOLS extends ToolSet<CONTEXT> = ToolSet<CONTEXT>,
-> = ValueOf<{
+export type StaticToolResult<TOOLS extends ToolSet = ToolSet> = ValueOf<{
   [NAME in keyof TOOLS]: {
     type: 'tool-result';
     toolCallId: string;
     toolName: NAME & string;
-    input: InferToolInput<CONTEXT, TOOLS[NAME]>;
-    output: InferToolOutput<CONTEXT, TOOLS[NAME]>;
+    input: InferToolInput<TOOLS[NAME]>;
+    output: InferToolOutput<TOOLS[NAME]>;
     providerExecuted?: boolean;
     providerMetadata?: ProviderMetadata;
     dynamic?: false | undefined;
@@ -38,7 +31,6 @@ export type DynamicToolResult = {
   title?: string;
 };
 
-export type TypedToolResult<
-  CONTEXT extends Partial<ContextRegistry>,
-  TOOLS extends ToolSet<CONTEXT> = ToolSet<CONTEXT>,
-> = StaticToolResult<CONTEXT, TOOLS> | DynamicToolResult;
+export type TypedToolResult<TOOLS extends ToolSet = ToolSet> =
+  | StaticToolResult<TOOLS>
+  | DynamicToolResult;
diff --git a/packages/ai/src/generate-text/tool-set.ts b/packages/ai/src/generate-text/tool-set.ts
index 68d308fb5a25..061410448768 100644
--- a/packages/ai/src/generate-text/tool-set.ts
+++ b/packages/ai/src/generate-text/tool-set.ts
@@ -1,8 +1,6 @@
-import { ContextRegistry, Tool } from '@ai-sdk/provider-utils';
+import { Context, Tool } from '@ai-sdk/provider-utils';
 
-export type ToolSet<
-  CONTEXT extends Partial<ContextRegistry> = ContextRegistry,
-> = Record<
+export type ToolSet<CONTEXT extends Context> = Record<
   string,
   (
     | Tool<CONTEXT, never, never>
diff --git a/packages/provider-utils/src/types/index.ts b/packages/provider-utils/src/types/index.ts
index b1fca4b85201..c70878137ece 100644
--- a/packages/provider-utils/src/types/index.ts
+++ b/packages/provider-utils/src/types/index.ts
@@ -33,9 +33,12 @@ export type { ToolCall } from './tool-call';
 export type { ToolContent, ToolModelMessage } from './tool-model-message';
 export type { ToolResult } from './tool-result';
 export type { UserContent, UserModelMessage } from './user-model-message';
+
 import type { ToolExecutionOptions } from './tool';
+import type { Context } from './context';
 
 /**
  * @deprecated Use ToolExecutionOptions instead.
  */
-export type ToolCallOptions = ToolExecutionOptions;
+export type ToolCallOptions<CONTEXT extends Context> =
+  ToolExecutionOptions<CONTEXT>;
diff --git a/packages/provider-utils/src/types/tool.ts b/packages/provider-utils/src/types/tool.ts
index 5785670b5332..d7745e029767 100644
--- a/packages/provider-utils/src/types/tool.ts
+++ b/packages/provider-utils/src/types/tool.ts
@@ -8,7 +8,7 @@ import { Context } from './context';
 /**
  * Additional options that are sent into each tool call.
  */
-export interface ToolExecutionOptions<CONTEXT extends Context = Context> {
+export interface ToolExecutionOptions<CONTEXT extends Context> {
   /**
    * The ID of the tool call. You can use it e.g. when sending tool-call related information with stream data.
    */
@@ -43,10 +43,7 @@ export interface ToolExecutionOptions<CONTEXT extends Context = Context> {
 /**
  * Function that is called to determine if the tool needs approval before it can be executed.
  */
-export type ToolNeedsApprovalFunction<
-  INPUT,
-  CONTEXT extends Context = Context,
-> = (
+export type ToolNeedsApprovalFunction<INPUT, CONTEXT extends Context> = (
   input: INPUT,
   options: {
     /**
@@ -114,7 +111,7 @@ type ToolOutputProperties<
  * The tool can also contain an optional execute function for the actual execution function of the tool.
  */
 export type Tool<
-  CONTEXT extends Context,
+  CONTEXT extends Context = any,
   INPUT extends JSONValue | unknown | never = any,
   OUTPUT extends JSONValue | unknown | never = any,
 > = {
@@ -159,7 +156,10 @@ export type Tool<
    */
   needsApproval?:
     | boolean
-    | ToolNeedsApprovalFunction<[INPUT] extends [never] ? unknown : INPUT>;
+    | ToolNeedsApprovalFunction<
+        [INPUT] extends [never] ? unknown : INPUT,
+        CONTEXT
+      >;
 
   /**
    * Strict mode setting for the tool.
@@ -174,14 +174,16 @@ export type Tool<
    * Optional function that is called when the argument streaming starts.
    * Only called when the tool is used in a streaming context.
    */
-  onInputStart?: (options: ToolExecutionOptions) => void | PromiseLike<void>;
+  onInputStart?: (
+    options: ToolExecutionOptions<CONTEXT>,
+  ) => void | PromiseLike<void>;
 
   /**
    * Optional function that is called when an argument streaming delta is available.
    * Only called when the tool is used in a streaming context.
    */
   onInputDelta?: (
-    options: { inputTextDelta: string } & ToolExecutionOptions,
+    options: { inputTextDelta: string } & ToolExecutionOptions<CONTEXT>,
   ) => void | PromiseLike<void>;
 
   /**
@@ -191,7 +193,7 @@ export type Tool<
   onInputAvailable?: (
     options: {
       input: [INPUT] extends [never] ? unknown : INPUT;
-    } & ToolExecutionOptions,
+    } & ToolExecutionOptions<CONTEXT>,
   ) => void | PromiseLike<void>;
 } & ToolOutputProperties<CONTEXT, INPUT, OUTPUT> & {
     /**
@@ -269,18 +271,14 @@ export type Tool<
 /**
  * Infer the input type of a tool.
  */
-export type InferToolInput<
-  CONTEXT extends Context,
-  TOOL extends Tool<CONTEXT>,
-> = TOOL extends Tool<CONTEXT, infer INPUT, any> ? INPUT : never;
+export type InferToolInput<TOOL extends Tool<any>> =
+  TOOL extends Tool<any, infer INPUT, any> ? INPUT : never;
 
 /**
  * Infer the output type of a tool.
  */
-export type InferToolOutput<
-  CONTEXT extends Context,
-  TOOL extends Tool<CONTEXT>,
-> = TOOL extends Tool<CONTEXT, any, infer OUTPUT> ? OUTPUT : never;
+export type InferToolOutput<TOOL extends Tool<any>> =
+  TOOL extends Tool<any, any, infer OUTPUT> ? OUTPUT : never;
 
 /**
  * Helper function for inferring the execute args of a tool.
@@ -337,7 +335,7 @@ export function dynamicTool(tool: {
   /**
    * Whether the tool needs approval before it can be executed.
    */
-  needsApproval?: boolean | ToolNeedsApprovalFunction<unknown>;
+  needsApproval?: boolean | ToolNeedsApprovalFunction<unknown, Context>;
 }): Tool<Context, unknown, unknown> & {
   type: 'dynamic';
 } {

From 37d2b0a2883f90f0ddc12a4bb59e6f3b89610b88 Mon Sep 17 00:00:00 2001
From: Lars Grammel <lars.grammel@gmail.com>
Date: Thu, 5 Mar 2026 16:12:13 +0100
Subject: [PATCH 06/41] xx

---
 .../ai/src/generate-text/callback-events.ts   | 71 +++++++++++--------
 .../src/generate-text/generate-text-result.ts |  4 +-
 .../ai/src/generate-text/generate-text.ts     | 65 ++++++++++-------
 packages/ai/src/generate-text/prepare-step.ts |  9 ++-
 packages/ai/src/generate-text/step-result.ts  | 17 +++--
 .../ai/src/generate-text/stop-condition.ts    | 13 +++-
 .../tool-call-repair-function.ts              |  6 +-
 packages/ai/src/generate-text/tool-call.ts    | 25 ++++---
 8 files changed, 127 insertions(+), 83 deletions(-)

diff --git a/packages/ai/src/generate-text/callback-events.ts b/packages/ai/src/generate-text/callback-events.ts
index 454be6f96a9a..3df6de677145 100644
--- a/packages/ai/src/generate-text/callback-events.ts
+++ b/packages/ai/src/generate-text/callback-events.ts
@@ -1,5 +1,6 @@
 import type { LanguageModelV4ToolChoice } from '@ai-sdk/provider';
 import type {
+  Context,
   ModelMessage,
   ProviderOptions,
   SystemModelMessage,
@@ -29,7 +30,8 @@ export interface CallbackModelInfo {
  * Called when the generation operation begins, before any LLM calls.
  */
 export interface OnStartEvent<
-  TOOLS extends ToolSet = ToolSet,
+  CONTEXT extends Context,
+  TOOLS extends ToolSet<CONTEXT> = ToolSet<CONTEXT>,
   OUTPUT extends Output = Output,
   INCLUDE = { requestBody?: boolean; responseBody?: boolean },
 > {
@@ -129,7 +131,8 @@ export interface OnStartEvent<
  * Each step represents a single LLM invocation.
  */
 export interface OnStepStartEvent<
-  TOOLS extends ToolSet = ToolSet,
+  CONTEXT extends Context,
+  TOOLS extends ToolSet<CONTEXT> = ToolSet<CONTEXT>,
   OUTPUT extends Output = Output,
   INCLUDE = { requestBody?: boolean; responseBody?: boolean },
 > {
@@ -216,7 +219,10 @@ export interface OnStepStartEvent<
  *
  * Called when a tool execution begins, before the tool's `execute` function is invoked.
  */
-export interface OnToolCallStartEvent<TOOLS extends ToolSet = ToolSet> {
+export interface OnToolCallStartEvent<
+  CONTEXT extends Context,
+  TOOLS extends ToolSet<CONTEXT> = ToolSet<CONTEXT>,
+> {
   /** Zero-based index of the current step where this tool call occurs. */
   readonly stepNumber: number | undefined;
 
@@ -248,7 +254,10 @@ export interface OnToolCallStartEvent<TOOLS extends ToolSet = ToolSet> {
  * Called when a tool execution completes, either successfully or with an error.
  * Uses a discriminated union on the `success` field.
  */
-export type OnToolCallFinishEvent<TOOLS extends ToolSet = ToolSet> = {
+export type OnToolCallFinishEvent<
+  CONTEXT extends Context,
+  TOOLS extends ToolSet<CONTEXT> = ToolSet<CONTEXT>,
+> = {
   /** Zero-based index of the current step where this tool call occurred. */
   readonly stepNumber: number | undefined;
 
@@ -298,8 +307,10 @@ export type OnToolCallFinishEvent<TOOLS extends ToolSet = ToolSet> = {
  * Called when a step (LLM call) completes.
  * This is simply the StepResult for that step.
  */
-export type OnStepFinishEvent<TOOLS extends ToolSet = ToolSet> =
-  StepResult<TOOLS>;
+export type OnStepFinishEvent<
+  CONTEXT extends Context,
+  TOOLS extends ToolSet<CONTEXT> = ToolSet<CONTEXT>,
+> = StepResult<TOOLS>;
 
 /**
  * Event passed to the `onFinish` callback.
@@ -307,26 +318,28 @@ export type OnStepFinishEvent<TOOLS extends ToolSet = ToolSet> =
  * Called when the entire generation completes (all steps finished).
  * Includes the final step's result along with aggregated data from all steps.
  */
-export type OnFinishEvent<TOOLS extends ToolSet = ToolSet> =
-  StepResult<TOOLS> & {
-    /** Array containing results from all steps in the generation. */
-    readonly steps: StepResult<TOOLS>[];
-
-    /** Aggregated token usage across all steps. */
-    readonly totalUsage: LanguageModelUsage;
-
-    /**
-     * The final state of the user-defined context object.
-     *
-     * Experimental (can break in patch releases).
-     *
-     * @default undefined
-     */
-    experimental_context: unknown;
-
-    /** Identifier from telemetry settings for grouping related operations. */
-    readonly functionId: string | undefined;
-
-    /** Additional metadata from telemetry settings. */
-    readonly metadata: Record<string, unknown> | undefined;
-  };
+export type OnFinishEvent<
+  CONTEXT extends Context,
+  TOOLS extends ToolSet<CONTEXT> = ToolSet<CONTEXT>,
+> = StepResult<TOOLS> & {
+  /** Array containing results from all steps in the generation. */
+  readonly steps: StepResult<TOOLS>[];
+
+  /** Aggregated token usage across all steps. */
+  readonly totalUsage: LanguageModelUsage;
+
+  /**
+   * The final state of the user-defined context object.
+   *
+   * Experimental (can break in patch releases).
+   *
+   * @default undefined
+   */
+  experimental_context: unknown;
+
+  /** Identifier from telemetry settings for grouping related operations. */
+  readonly functionId: string | undefined;
+
+  /** Additional metadata from telemetry settings. */
+  readonly metadata: Record<string, unknown> | undefined;
+};
diff --git a/packages/ai/src/generate-text/generate-text-result.ts b/packages/ai/src/generate-text/generate-text-result.ts
index 6f37f8d69eaf..f2311ba68ab9 100644
--- a/packages/ai/src/generate-text/generate-text-result.ts
+++ b/packages/ai/src/generate-text/generate-text-result.ts
@@ -1,3 +1,4 @@
+import { Context } from '@ai-sdk/provider-utils';
 import { CallWarning, FinishReason, ProviderMetadata } from '../types';
 import { Source } from '../types/language-model';
 import { LanguageModelRequestMetadata } from '../types/language-model-request-metadata';
@@ -23,7 +24,8 @@ import { ToolSet } from './tool-set';
  * It contains the generated text, the tool calls that were made during the generation, and the results of the tool calls.
  */
 export interface GenerateTextResult<
-  TOOLS extends ToolSet,
+  CONTEXT extends Context,
+  TOOLS extends ToolSet<CONTEXT>,
   OUTPUT extends Output,
 > {
   /**
diff --git a/packages/ai/src/generate-text/generate-text.ts b/packages/ai/src/generate-text/generate-text.ts
index 8d9f82ddc4fd..494ef59459bf 100644
--- a/packages/ai/src/generate-text/generate-text.ts
+++ b/packages/ai/src/generate-text/generate-text.ts
@@ -114,10 +114,11 @@ type GenerateTextIncludeSettings = {
  * @param event - The event object containing generation configuration.
  */
 export type GenerateTextOnStartCallback<
-  TOOLS extends ToolSet = ToolSet,
+  CONTEXT extends Context,
+  TOOLS extends ToolSet<CONTEXT> = ToolSet<CONTEXT>,
   OUTPUT extends Output = Output,
 > = (
-  event: OnStartEvent<TOOLS, OUTPUT, GenerateTextIncludeSettings>,
+  event: OnStartEvent<CONTEXT, TOOLS, OUTPUT, GenerateTextIncludeSettings>,
 ) => PromiseLike<void> | void;
 
 /**
@@ -130,10 +131,11 @@ export type GenerateTextOnStartCallback<
  * @param event - The event object containing step configuration.
  */
 export type GenerateTextOnStepStartCallback<
-  TOOLS extends ToolSet = ToolSet,
+  CONTEXT extends Context,
+  TOOLS extends ToolSet<CONTEXT> = ToolSet<CONTEXT>,
   OUTPUT extends Output = Output,
 > = (
-  event: OnStepStartEvent<TOOLS, OUTPUT, GenerateTextIncludeSettings>,
+  event: OnStepStartEvent<CONTEXT, TOOLS, OUTPUT, GenerateTextIncludeSettings>,
 ) => PromiseLike<void> | void;
 
 /**
@@ -145,7 +147,8 @@ export type GenerateTextOnStepStartCallback<
  * @param event - The event object containing tool call information.
  */
 export type GenerateTextOnToolCallStartCallback<
-  TOOLS extends ToolSet = ToolSet,
+  CONTEXT extends Context,
+  TOOLS extends ToolSet<CONTEXT> = ToolSet<CONTEXT>,
 > = (event: OnToolCallStartEvent<TOOLS>) => PromiseLike<void> | void;
 
 /**
@@ -161,7 +164,8 @@ export type GenerateTextOnToolCallStartCallback<
  * @param event - The event object containing tool call result information.
  */
 export type GenerateTextOnToolCallFinishCallback<
-  TOOLS extends ToolSet = ToolSet,
+  CONTEXT extends Context,
+  TOOLS extends ToolSet<CONTEXT> = ToolSet<CONTEXT>,
 > = (event: OnToolCallFinishEvent<TOOLS>) => PromiseLike<void> | void;
 
 /**
@@ -172,9 +176,10 @@ export type GenerateTextOnToolCallFinishCallback<
  *
  * @param stepResult - The result of the step.
  */
-export type GenerateTextOnStepFinishCallback<TOOLS extends ToolSet> = (
-  event: OnStepFinishEvent<TOOLS>,
-) => Promise<void> | void;
+export type GenerateTextOnStepFinishCallback<
+  CONTEXT extends Context,
+  TOOLS extends ToolSet<CONTEXT> = ToolSet<CONTEXT>,
+> = (event: OnStepFinishEvent<TOOLS>) => Promise<void> | void;
 
 /**
  * Callback that is set using the `onFinish` option.
@@ -185,9 +190,10 @@ export type GenerateTextOnStepFinishCallback<TOOLS extends ToolSet> = (
  *
  * @param event - The final result along with aggregated step data.
  */
-export type GenerateTextOnFinishCallback<TOOLS extends ToolSet = ToolSet> = (
-  event: OnFinishEvent<TOOLS>,
-) => PromiseLike<void> | void;
+export type GenerateTextOnFinishCallback<
+  CONTEXT extends Context,
+  TOOLS extends ToolSet<CONTEXT> = ToolSet<CONTEXT>,
+> = (event: OnFinishEvent<TOOLS>) => PromiseLike<void> | void;
 
 /**
  * Generate a text and call tools for a given prompt using a language model.
@@ -245,7 +251,7 @@ export type GenerateTextOnFinishCallback<TOOLS extends ToolSet = ToolSet> = (
  */
 export async function generateText<
   CONTEXT extends Context,
-  TOOLS extends ToolSet = ToolSet,
+  TOOLS extends ToolSet<CONTEXT> = ToolSet<CONTEXT>,
   OUTPUT extends Output = Output<string, string>,
 >({
   model: modelArg,
@@ -367,13 +373,18 @@ export async function generateText<
      * Callback that is called when the generateText operation begins,
      * before any LLM calls are made.
      */
-    experimental_onStart?: GenerateTextOnStartCallback<NoInfer<TOOLS>, OUTPUT>;
+    experimental_onStart?: GenerateTextOnStartCallback<
+      CONTEXT,
+      NoInfer<TOOLS>,
+      OUTPUT
+    >;
 
     /**
      * Callback that is called when a step (LLM call) begins,
      * before the provider is called.
      */
     experimental_onStepStart?: GenerateTextOnStepStartCallback<
+      CONTEXT,
       NoInfer<TOOLS>,
       OUTPUT
     >;
@@ -439,7 +450,7 @@ export async function generateText<
     _internal?: {
       generateId?: IdGenerator;
     };
-  }): Promise<GenerateTextResult<TOOLS, OUTPUT>> {
+  }): Promise<GenerateTextResult<CONTEXT, TOOLS, OUTPUT>> {
   const model = resolveLanguageModel(modelArg);
   const createGlobalTelemetry = getGlobalTelemetryIntegration<TOOLS, OUTPUT>();
   const stopConditions = asArray(stopWhen);
@@ -516,7 +527,7 @@ export async function generateText<
       onStart,
       globalTelemetry.onStart as
         | undefined
-        | GenerateTextOnStartCallback<TOOLS, OUTPUT>,
+        | GenerateTextOnStartCallback<CONTEXT, TOOLS, OUTPUT>,
     ],
   });
 
@@ -661,7 +672,7 @@ export async function generateText<
         > & { response: { id: string; timestamp: Date; modelId: string } };
         let clientToolCalls: Array<TypedToolCall<TOOLS>> = [];
         let clientToolOutputs: Array<ToolOutput<TOOLS>> = [];
-        const steps: GenerateTextResult<TOOLS, OUTPUT>['steps'] = [];
+        const steps: StepResult<CONTEXT, TOOLS>[] = [];
 
         // Track provider-executed tool calls that support deferred results
         // (e.g., code_execution in programmatic tool calling scenarios).
@@ -757,7 +768,7 @@ export async function generateText<
                 onStepStart,
                 globalTelemetry.onStepStart as
                   | undefined
-                  | GenerateTextOnStepStartCallback<TOOLS, OUTPUT>,
+                  | GenerateTextOnStepStartCallback<CONTEXT, TOOLS, OUTPUT>,
               ],
             });
 
@@ -1220,7 +1231,10 @@ export async function generateText<
   }
 }
 
-async function executeTools<CONTEXT extends Context, TOOLS extends ToolSet>({
+async function executeTools<
+  CONTEXT extends Context,
+  TOOLS extends ToolSet<CONTEXT> = ToolSet<CONTEXT>,
+>({
   toolCalls,
   tools,
   tracer,
@@ -1274,15 +1288,18 @@ async function executeTools<CONTEXT extends Context, TOOLS extends ToolSet>({
   );
 }
 
-class DefaultGenerateTextResult<TOOLS extends ToolSet, OUTPUT extends Output>
-  implements GenerateTextResult<TOOLS, OUTPUT>
+class DefaultGenerateTextResult<
+  CONTEXT extends Context,
+  TOOLS extends ToolSet<CONTEXT>,
+  OUTPUT extends Output,
+> implements GenerateTextResult<CONTEXT, TOOLS, OUTPUT>
 {
-  readonly steps: GenerateTextResult<TOOLS, OUTPUT>['steps'];
+  readonly steps: GenerateTextResult<CONTEXT, TOOLS, OUTPUT>['steps'];
   readonly totalUsage: LanguageModelUsage;
   private readonly _output: InferCompleteOutput<OUTPUT> | undefined;
 
   constructor(options: {
-    steps: GenerateTextResult<TOOLS, OUTPUT>['steps'];
+    steps: GenerateTextResult<CONTEXT, TOOLS, OUTPUT>['steps'];
     output: InferCompleteOutput<OUTPUT> | undefined;
     totalUsage: LanguageModelUsage;
   }) {
@@ -1400,7 +1417,7 @@ function asToolCalls(content: Array<LanguageModelV4Content>) {
   }));
 }
 
-function asContent<TOOLS extends ToolSet>({
+function asContent<TOOLS extends ToolSet<any>>({
   content,
   toolCalls,
   toolOutputs,
diff --git a/packages/ai/src/generate-text/prepare-step.ts b/packages/ai/src/generate-text/prepare-step.ts
index 8ba9901b631d..3dcab84207c0 100644
--- a/packages/ai/src/generate-text/prepare-step.ts
+++ b/packages/ai/src/generate-text/prepare-step.ts
@@ -1,6 +1,5 @@
 import {
   Context,
-  ContextRegistry,
   ModelMessage,
   ProviderOptions,
   SystemModelMessage,
@@ -23,7 +22,7 @@ import { ToolSet } from './tool-set';
  * If you return undefined (or for undefined settings), the settings from the outer level will be used.
  */
 export type PrepareStepFunction<
-  CONTEXT extends Partial<ContextRegistry> = ContextRegistry,
+  CONTEXT extends Context,
   TOOLS extends ToolSet<CONTEXT> = ToolSet<CONTEXT>,
 > = (options: {
   /**
@@ -49,7 +48,7 @@ export type PrepareStepFunction<
   /**
    * The context passed via the experimental_context setting (experimental).
    */
-  experimental_context: Context<CONTEXT>;
+  experimental_context: CONTEXT;
 }) =>
   | PromiseLike<PrepareStepResult<CONTEXT, TOOLS>>
   | PrepareStepResult<CONTEXT, TOOLS>;
@@ -59,7 +58,7 @@ export type PrepareStepFunction<
  * allowing per-step overrides of model, tools, or messages.
  */
 export type PrepareStepResult<
-  CONTEXT extends Partial<ContextRegistry> = ContextRegistry,
+  CONTEXT extends Context,
   TOOLS extends ToolSet<CONTEXT> = ToolSet<CONTEXT>,
 > =
   | {
@@ -96,7 +95,7 @@ export type PrepareStepResult<
        * Changing the context will affect the context in this step
        * and all subsequent steps.
        */
-      experimental_context?: Context<CONTEXT>;
+      experimental_context?: CONTEXT;
 
       /**
        * Additional provider-specific options for this step.
diff --git a/packages/ai/src/generate-text/step-result.ts b/packages/ai/src/generate-text/step-result.ts
index 94d7a9e92fd2..9b8bf2c09c1e 100644
--- a/packages/ai/src/generate-text/step-result.ts
+++ b/packages/ai/src/generate-text/step-result.ts
@@ -1,4 +1,4 @@
-import { ContextRegistry, ReasoningPart } from '@ai-sdk/provider-utils';
+import { Context, ReasoningPart } from '@ai-sdk/provider-utils';
 import {
   CallWarning,
   FinishReason,
@@ -22,10 +22,7 @@ import { ToolSet } from './tool-set';
 /**
  * The result of a single step in the generation process.
  */
-export type StepResult<
-  CONTEXT extends Partial<ContextRegistry>,
-  TOOLS extends ToolSet<CONTEXT> = ToolSet<CONTEXT>,
-> = {
+export type StepResult<TOOLS extends ToolSet = ToolSet> = {
   /**
    * Zero-based index of this step.
    */
@@ -91,7 +88,7 @@ export type StepResult<
   /**
    * The tool calls that were made during the generation.
    */
-  readonly toolCalls: Array<TypedToolCall<CONTEXT, TOOLS>>;
+  readonly toolCalls: Array<TypedToolCall<TOOLS>>;
 
   /**
    * The static tool calls that were made in the last step.
@@ -106,7 +103,7 @@ export type StepResult<
   /**
    * The results of the tool calls.
    */
-  readonly toolResults: Array<TypedToolResult<CONTEXT, TOOLS>>;
+  readonly toolResults: Array<TypedToolResult<TOOLS>>;
 
   /**
    * The static tool results that were made in the last step.
@@ -168,8 +165,10 @@ export type StepResult<
   readonly providerMetadata: ProviderMetadata | undefined;
 };
 
-export class DefaultStepResult<TOOLS extends ToolSet>
-  implements StepResult<TOOLS>
+export class DefaultStepResult<
+  CONTEXT extends Context,
+  TOOLS extends ToolSet<CONTEXT>,
+> implements StepResult<CONTEXT, TOOLS>
 {
   readonly stepNumber: StepResult<TOOLS>['stepNumber'];
   readonly model: StepResult<TOOLS>['model'];
diff --git a/packages/ai/src/generate-text/stop-condition.ts b/packages/ai/src/generate-text/stop-condition.ts
index a9438b923fc1..cdcb27734501 100644
--- a/packages/ai/src/generate-text/stop-condition.ts
+++ b/packages/ai/src/generate-text/stop-condition.ts
@@ -1,7 +1,11 @@
+import { Context } from '@ai-sdk/provider-utils';
 import { StepResult } from './step-result';
 import { ToolSet } from './tool-set';
 
-export type StopCondition<TOOLS extends ToolSet = ToolSet> = (options: {
+export type StopCondition<
+  CONTEXT extends Context,
+  TOOLS extends ToolSet<CONTEXT> = ToolSet<CONTEXT>,
+> = (options: {
   steps: Array<StepResult<TOOLS>>;
 }) => PromiseLike<boolean> | boolean;
 
@@ -16,11 +20,14 @@ export function hasToolCall(toolName: string): StopCondition<any> {
     ) ?? false;
 }
 
-export async function isStopConditionMet<TOOLS extends ToolSet>({
+export async function isStopConditionMet<
+  CONTEXT extends Context,
+  TOOLS extends ToolSet<CONTEXT>,
+>({
   stopConditions,
   steps,
 }: {
-  stopConditions: Array<StopCondition<TOOLS>>;
+  stopConditions: Array<StopCondition<CONTEXT, TOOLS>>;
   steps: Array<StepResult<TOOLS>>;
 }): Promise<boolean> {
   return (
diff --git a/packages/ai/src/generate-text/tool-call-repair-function.ts b/packages/ai/src/generate-text/tool-call-repair-function.ts
index d747916c6ff0..49f6f7e62983 100644
--- a/packages/ai/src/generate-text/tool-call-repair-function.ts
+++ b/packages/ai/src/generate-text/tool-call-repair-function.ts
@@ -1,4 +1,5 @@
 import { JSONSchema7, LanguageModelV4ToolCall } from '@ai-sdk/provider';
+import { Context } from '@ai-sdk/provider-utils';
 import { InvalidToolInputError } from '../error/invalid-tool-input-error';
 import { NoSuchToolError } from '../error/no-such-tool-error';
 import { ModelMessage, SystemModelMessage } from '../prompt';
@@ -17,7 +18,10 @@ import { ToolSet } from './tool-set';
  * @param options.inputSchema - A function that returns the JSON Schema for a tool.
  * @param options.error - The error that occurred while parsing the tool call.
  */
-export type ToolCallRepairFunction<TOOLS extends ToolSet> = (options: {
+export type ToolCallRepairFunction<
+  CONTEXT extends Context,
+  TOOLS extends ToolSet<CONTEXT> = ToolSet<CONTEXT>,
+> = (options: {
   system: string | SystemModelMessage | Array<SystemModelMessage> | undefined;
   messages: ModelMessage[];
   toolCall: LanguageModelV4ToolCall;
diff --git a/packages/ai/src/generate-text/tool-call.ts b/packages/ai/src/generate-text/tool-call.ts
index 59698ce0b550..831f93e3db25 100644
--- a/packages/ai/src/generate-text/tool-call.ts
+++ b/packages/ai/src/generate-text/tool-call.ts
@@ -10,16 +10,19 @@ type BaseToolCall = {
   providerMetadata?: ProviderMetadata;
 };
 
-export type StaticToolCall<TOOLS extends ToolSet = ToolSet> = ValueOf<{
-  [NAME in keyof TOOLS]: BaseToolCall & {
-    toolName: NAME & string;
-    input: TOOLS[NAME] extends Tool<any, infer PARAMETERS> ? PARAMETERS : never;
-    dynamic?: false | undefined;
-    invalid?: false | undefined;
-    error?: never;
-    title?: string;
-  };
-}>;
+export type StaticToolCall<TOOLS extends ToolSet<any> = ToolSet<any>> =
+  ValueOf<{
+    [NAME in keyof TOOLS]: BaseToolCall & {
+      toolName: NAME & string;
+      input: TOOLS[NAME] extends Tool<any, infer PARAMETERS>
+        ? PARAMETERS
+        : never;
+      dynamic?: false | undefined;
+      invalid?: false | undefined;
+      error?: never;
+      title?: string;
+    };
+  }>;
 
 export type DynamicToolCall = BaseToolCall & {
   toolName: string;
@@ -42,6 +45,6 @@ export type DynamicToolCall = BaseToolCall & {
   error?: unknown;
 };
 
-export type TypedToolCall<TOOLS extends ToolSet = ToolSet> =
+export type TypedToolCall<TOOLS extends ToolSet<any> = ToolSet<any>> =
   | StaticToolCall<TOOLS>
   | DynamicToolCall;

From 7ad6fd7fdfdf149e86bf56276361eeb785dcf5a2 Mon Sep 17 00:00:00 2001
From: Lars Grammel <lars.grammel@gmail.com>
Date: Thu, 5 Mar 2026 16:40:31 +0100
Subject: [PATCH 07/41] y

---
 .../ai-functions/src/test/typed-context.ts    |  6 +-
 .../ai/src/generate-text/callback-events.ts   | 29 +++-----
 packages/ai/src/generate-text/content-part.ts | 16 ++---
 .../src/generate-text/generate-text-result.ts |  4 +-
 .../ai/src/generate-text/generate-text.ts     | 68 +++++++++----------
 packages/ai/src/generate-text/step-result.ts  | 64 +++++++++--------
 .../ai/src/generate-text/stop-condition.ts    | 12 +---
 packages/ai/src/generate-text/tool-call.ts    | 25 +++----
 packages/ai/src/generate-text/tool-set.ts     |  5 +-
 packages/provider-utils/src/types/tool.ts     |  2 +-
 10 files changed, 107 insertions(+), 124 deletions(-)

diff --git a/examples/ai-functions/src/test/typed-context.ts b/examples/ai-functions/src/test/typed-context.ts
index dd53f7cd0bf2..4446294f8c8d 100644
--- a/examples/ai-functions/src/test/typed-context.ts
+++ b/examples/ai-functions/src/test/typed-context.ts
@@ -6,13 +6,13 @@ interface Tool<CONTEXT> {
 
 export type ToolSet<CONTEXT> = Record<string, Tool<CONTEXT>>;
 
-function executeTool<CONTEXT>({
+function executeTool<CONTEXT, TOOLS extends ToolSet<CONTEXT>>({
   tools,
   toolName,
   context,
 }: {
-  tools: ToolSet<CONTEXT>;
-  toolName: keyof ToolSet<CONTEXT>;
+  tools: TOOLS;
+  toolName: keyof TOOLS;
   context: CONTEXT;
 }) {
   return tools[toolName].execute(context);
diff --git a/packages/ai/src/generate-text/callback-events.ts b/packages/ai/src/generate-text/callback-events.ts
index 3df6de677145..697a711002d3 100644
--- a/packages/ai/src/generate-text/callback-events.ts
+++ b/packages/ai/src/generate-text/callback-events.ts
@@ -30,8 +30,7 @@ export interface CallbackModelInfo {
  * Called when the generation operation begins, before any LLM calls.
  */
 export interface OnStartEvent<
-  CONTEXT extends Context,
-  TOOLS extends ToolSet<CONTEXT> = ToolSet<CONTEXT>,
+  TOOLS extends ToolSet,
   OUTPUT extends Output = Output,
   INCLUDE = { requestBody?: boolean; responseBody?: boolean },
 > {
@@ -132,7 +131,7 @@ export interface OnStartEvent<
  */
 export interface OnStepStartEvent<
   CONTEXT extends Context,
-  TOOLS extends ToolSet<CONTEXT> = ToolSet<CONTEXT>,
+  TOOLS extends ToolSet<CONTEXT>,
   OUTPUT extends Output = Output,
   INCLUDE = { requestBody?: boolean; responseBody?: boolean },
 > {
@@ -168,7 +167,7 @@ export interface OnStepStartEvent<
   readonly activeTools: Array<keyof TOOLS> | undefined;
 
   /** Array of results from previous steps (empty for first step). */
-  readonly steps: ReadonlyArray<StepResult<TOOLS>>;
+  readonly steps: ReadonlyArray<StepResult<CONTEXT, TOOLS>>;
 
   /** Additional provider-specific options for this step. */
   readonly providerOptions: ProviderOptions | undefined;
@@ -219,10 +218,7 @@ export interface OnStepStartEvent<
  *
  * Called when a tool execution begins, before the tool's `execute` function is invoked.
  */
-export interface OnToolCallStartEvent<
-  CONTEXT extends Context,
-  TOOLS extends ToolSet<CONTEXT> = ToolSet<CONTEXT>,
-> {
+export interface OnToolCallStartEvent<TOOLS extends ToolSet> {
   /** Zero-based index of the current step where this tool call occurs. */
   readonly stepNumber: number | undefined;
 
@@ -254,10 +250,7 @@ export interface OnToolCallStartEvent<
  * Called when a tool execution completes, either successfully or with an error.
  * Uses a discriminated union on the `success` field.
  */
-export type OnToolCallFinishEvent<
-  CONTEXT extends Context,
-  TOOLS extends ToolSet<CONTEXT> = ToolSet<CONTEXT>,
-> = {
+export type OnToolCallFinishEvent<TOOLS extends ToolSet> = {
   /** Zero-based index of the current step where this tool call occurred. */
   readonly stepNumber: number | undefined;
 
@@ -309,8 +302,8 @@ export type OnToolCallFinishEvent<
  */
 export type OnStepFinishEvent<
   CONTEXT extends Context,
-  TOOLS extends ToolSet<CONTEXT> = ToolSet<CONTEXT>,
-> = StepResult<TOOLS>;
+  TOOLS extends ToolSet<CONTEXT>,
+> = StepResult<CONTEXT, TOOLS>;
 
 /**
  * Event passed to the `onFinish` callback.
@@ -320,10 +313,10 @@ export type OnStepFinishEvent<
  */
 export type OnFinishEvent<
   CONTEXT extends Context,
-  TOOLS extends ToolSet<CONTEXT> = ToolSet<CONTEXT>,
-> = StepResult<TOOLS> & {
+  TOOLS extends ToolSet<CONTEXT>,
+> = StepResult<CONTEXT, TOOLS> & {
   /** Array containing results from all steps in the generation. */
-  readonly steps: StepResult<TOOLS>[];
+  readonly steps: StepResult<CONTEXT, TOOLS>[];
 
   /** Aggregated token usage across all steps. */
   readonly totalUsage: LanguageModelUsage;
@@ -335,7 +328,7 @@ export type OnFinishEvent<
    *
    * @default undefined
    */
-  experimental_context: unknown;
+  experimental_context: CONTEXT;
 
   /** Identifier from telemetry settings for grouping related operations. */
   readonly functionId: string | undefined;
diff --git a/packages/ai/src/generate-text/content-part.ts b/packages/ai/src/generate-text/content-part.ts
index df1e9b1eb90a..8575387a179b 100644
--- a/packages/ai/src/generate-text/content-part.ts
+++ b/packages/ai/src/generate-text/content-part.ts
@@ -1,29 +1,25 @@
-import { ContextRegistry } from '@ai-sdk/provider-utils';
 import { ProviderMetadata } from '../types';
 import { Source } from '../types/language-model';
 import { GeneratedFile } from './generated-file';
-import { ToolApprovalRequestOutput } from './tool-approval-request-output';
 import { ReasoningOutput } from './reasoning-output';
+import { ToolApprovalRequestOutput } from './tool-approval-request-output';
 import { TypedToolCall } from './tool-call';
 import { TypedToolError } from './tool-error';
 import { TypedToolResult } from './tool-result';
 import { ToolSet } from './tool-set';
 
-export type ContentPart<
-  CONTEXT extends Partial<ContextRegistry>,
-  TOOLS extends ToolSet<CONTEXT> = ToolSet<CONTEXT>,
-> =
+export type ContentPart<TOOLS extends ToolSet = ToolSet> =
   | { type: 'text'; text: string; providerMetadata?: ProviderMetadata }
   | ReasoningOutput
   | ({ type: 'source' } & Source)
   | { type: 'file'; file: GeneratedFile; providerMetadata?: ProviderMetadata } // different because of GeneratedFile object
-  | ({ type: 'tool-call' } & TypedToolCall<CONTEXT, TOOLS> & {
+  | ({ type: 'tool-call' } & TypedToolCall<TOOLS> & {
         providerMetadata?: ProviderMetadata;
       })
-  | ({ type: 'tool-result' } & TypedToolResult<CONTEXT, TOOLS> & {
+  | ({ type: 'tool-result' } & TypedToolResult<TOOLS> & {
         providerMetadata?: ProviderMetadata;
       })
-  | ({ type: 'tool-error' } & TypedToolError<CONTEXT, TOOLS> & {
+  | ({ type: 'tool-error' } & TypedToolError<TOOLS> & {
         providerMetadata?: ProviderMetadata;
       })
-  | ToolApprovalRequestOutput<CONTEXT, TOOLS>;
+  | ToolApprovalRequestOutput<TOOLS>;
diff --git a/packages/ai/src/generate-text/generate-text-result.ts b/packages/ai/src/generate-text/generate-text-result.ts
index f2311ba68ab9..6f37f8d69eaf 100644
--- a/packages/ai/src/generate-text/generate-text-result.ts
+++ b/packages/ai/src/generate-text/generate-text-result.ts
@@ -1,4 +1,3 @@
-import { Context } from '@ai-sdk/provider-utils';
 import { CallWarning, FinishReason, ProviderMetadata } from '../types';
 import { Source } from '../types/language-model';
 import { LanguageModelRequestMetadata } from '../types/language-model-request-metadata';
@@ -24,8 +23,7 @@ import { ToolSet } from './tool-set';
  * It contains the generated text, the tool calls that were made during the generation, and the results of the tool calls.
  */
 export interface GenerateTextResult<
-  CONTEXT extends Context,
-  TOOLS extends ToolSet<CONTEXT>,
+  TOOLS extends ToolSet,
   OUTPUT extends Output,
 > {
   /**
diff --git a/packages/ai/src/generate-text/generate-text.ts b/packages/ai/src/generate-text/generate-text.ts
index 494ef59459bf..3a312235ac43 100644
--- a/packages/ai/src/generate-text/generate-text.ts
+++ b/packages/ai/src/generate-text/generate-text.ts
@@ -114,11 +114,10 @@ type GenerateTextIncludeSettings = {
  * @param event - The event object containing generation configuration.
  */
 export type GenerateTextOnStartCallback<
-  CONTEXT extends Context,
-  TOOLS extends ToolSet<CONTEXT> = ToolSet<CONTEXT>,
+  TOOLS extends ToolSet,
   OUTPUT extends Output = Output,
 > = (
-  event: OnStartEvent<CONTEXT, TOOLS, OUTPUT, GenerateTextIncludeSettings>,
+  event: OnStartEvent<TOOLS, OUTPUT, GenerateTextIncludeSettings>,
 ) => PromiseLike<void> | void;
 
 /**
@@ -179,7 +178,7 @@ export type GenerateTextOnToolCallFinishCallback<
 export type GenerateTextOnStepFinishCallback<
   CONTEXT extends Context,
   TOOLS extends ToolSet<CONTEXT> = ToolSet<CONTEXT>,
-> = (event: OnStepFinishEvent<TOOLS>) => Promise<void> | void;
+> = (event: OnStepFinishEvent<CONTEXT, TOOLS>) => Promise<void> | void;
 
 /**
  * Callback that is set using the `onFinish` option.
@@ -193,7 +192,7 @@ export type GenerateTextOnStepFinishCallback<
 export type GenerateTextOnFinishCallback<
   CONTEXT extends Context,
   TOOLS extends ToolSet<CONTEXT> = ToolSet<CONTEXT>,
-> = (event: OnFinishEvent<TOOLS>) => PromiseLike<void> | void;
+> = (event: OnFinishEvent<CONTEXT, TOOLS>) => PromiseLike<void> | void;
 
 /**
  * Generate a text and call tools for a given prompt using a language model.
@@ -373,11 +372,7 @@ export async function generateText<
      * Callback that is called when the generateText operation begins,
      * before any LLM calls are made.
      */
-    experimental_onStart?: GenerateTextOnStartCallback<
-      CONTEXT,
-      NoInfer<TOOLS>,
-      OUTPUT
-    >;
+    experimental_onStart?: GenerateTextOnStartCallback<NoInfer<TOOLS>, OUTPUT>;
 
     /**
      * Callback that is called when a step (LLM call) begins,
@@ -406,12 +401,12 @@ export async function generateText<
     /**
      * Callback that is called when each step (LLM call) is finished, including intermediate steps.
      */
-    onStepFinish?: GenerateTextOnStepFinishCallback<NoInfer<TOOLS>>;
+    onStepFinish?: GenerateTextOnStepFinishCallback<CONTEXT, NoInfer<TOOLS>>;
 
     /**
      * Callback that is called when all steps are finished and the response is complete.
      */
-    onFinish?: GenerateTextOnFinishCallback<NoInfer<TOOLS>>;
+    onFinish?: GenerateTextOnFinishCallback<CONTEXT, NoInfer<TOOLS>>;
 
     /**
      * Context that is passed into tool execution.
@@ -450,9 +445,9 @@ export async function generateText<
     _internal?: {
       generateId?: IdGenerator;
     };
-  }): Promise<GenerateTextResult<CONTEXT, TOOLS, OUTPUT>> {
+  }): Promise<GenerateTextResult<TOOLS, OUTPUT>> {
   const model = resolveLanguageModel(modelArg);
-  const createGlobalTelemetry = getGlobalTelemetryIntegration<TOOLS, OUTPUT>();
+  const createGlobalTelemetry = getGlobalTelemetryIntegration<any, OUTPUT>();
   const stopConditions = asArray(stopWhen);
 
   const totalTimeoutMs = getTotalTimeoutMs(timeout);
@@ -527,7 +522,7 @@ export async function generateText<
       onStart,
       globalTelemetry.onStart as
         | undefined
-        | GenerateTextOnStartCallback<CONTEXT, TOOLS, OUTPUT>,
+        | GenerateTextOnStartCallback<TOOLS, OUTPUT>,
     ],
   });
 
@@ -1072,23 +1067,24 @@ export async function generateText<
 
             const stepNumber = steps.length;
 
-            const currentStepResult: StepResult<TOOLS> = new DefaultStepResult({
-              stepNumber,
-              model: stepModelInfo,
-              functionId: telemetry?.functionId,
-              metadata: telemetry?.metadata as
-                | Record<string, unknown>
-                | undefined,
-              experimental_context,
-              content: stepContent,
-              finishReason: currentModelResponse.finishReason.unified,
-              rawFinishReason: currentModelResponse.finishReason.raw,
-              usage: asLanguageModelUsage(currentModelResponse.usage),
-              warnings: currentModelResponse.warnings,
-              providerMetadata: currentModelResponse.providerMetadata,
-              request: stepRequest,
-              response: stepResponse,
-            });
+            const currentStepResult: StepResult<CONTEXT, TOOLS> =
+              new DefaultStepResult({
+                stepNumber,
+                model: stepModelInfo,
+                functionId: telemetry?.functionId,
+                metadata: telemetry?.metadata as
+                  | Record<string, unknown>
+                  | undefined,
+                experimental_context,
+                content: stepContent,
+                finishReason: currentModelResponse.finishReason.unified,
+                rawFinishReason: currentModelResponse.finishReason.raw,
+                usage: asLanguageModelUsage(currentModelResponse.usage),
+                warnings: currentModelResponse.warnings,
+                providerMetadata: currentModelResponse.providerMetadata,
+                request: stepRequest,
+                response: stepResponse,
+              });
 
             logWarnings({
               warnings: currentModelResponse.warnings ?? [],
@@ -1201,7 +1197,7 @@ export async function generateText<
             onFinish,
             globalTelemetry.onFinish as
               | undefined
-              | GenerateTextOnFinishCallback<TOOLS>,
+              | GenerateTextOnFinishCallback<CONTEXT, TOOLS>,
           ],
         });
 
@@ -1292,14 +1288,14 @@ class DefaultGenerateTextResult<
   CONTEXT extends Context,
   TOOLS extends ToolSet<CONTEXT>,
   OUTPUT extends Output,
-> implements GenerateTextResult<CONTEXT, TOOLS, OUTPUT>
+> implements GenerateTextResult<TOOLS, OUTPUT>
 {
-  readonly steps: GenerateTextResult<CONTEXT, TOOLS, OUTPUT>['steps'];
+  readonly steps: StepResult<CONTEXT, TOOLS>[];
   readonly totalUsage: LanguageModelUsage;
   private readonly _output: InferCompleteOutput<OUTPUT> | undefined;
 
   constructor(options: {
-    steps: GenerateTextResult<CONTEXT, TOOLS, OUTPUT>['steps'];
+    steps: StepResult<CONTEXT, TOOLS>[];
     output: InferCompleteOutput<OUTPUT> | undefined;
     totalUsage: LanguageModelUsage;
   }) {
diff --git a/packages/ai/src/generate-text/step-result.ts b/packages/ai/src/generate-text/step-result.ts
index 9b8bf2c09c1e..6574053e66a5 100644
--- a/packages/ai/src/generate-text/step-result.ts
+++ b/packages/ai/src/generate-text/step-result.ts
@@ -22,7 +22,10 @@ import { ToolSet } from './tool-set';
 /**
  * The result of a single step in the generation process.
  */
-export type StepResult<TOOLS extends ToolSet = ToolSet> = {
+export type StepResult<
+  CONTEXT extends Context,
+  TOOLS extends ToolSet<CONTEXT>,
+> = {
   /**
    * Zero-based index of this step.
    */
@@ -53,12 +56,12 @@ export type StepResult<TOOLS extends ToolSet = ToolSet> = {
    *
    * Experimental (can break in patch releases).
    */
-  readonly experimental_context: unknown;
+  readonly experimental_context: CONTEXT;
 
   /**
    * The content that was generated in the last step.
    */
-  readonly content: Array<ContentPart<CONTEXT, TOOLS>>;
+  readonly content: Array<ContentPart<TOOLS>>;
 
   /**
    * The generated text.
@@ -170,19 +173,22 @@ export class DefaultStepResult<
   TOOLS extends ToolSet<CONTEXT>,
 > implements StepResult<CONTEXT, TOOLS>
 {
-  readonly stepNumber: StepResult<TOOLS>['stepNumber'];
-  readonly model: StepResult<TOOLS>['model'];
-  readonly functionId: StepResult<TOOLS>['functionId'];
-  readonly metadata: StepResult<TOOLS>['metadata'];
-  readonly experimental_context: StepResult<TOOLS>['experimental_context'];
-  readonly content: StepResult<TOOLS>['content'];
-  readonly finishReason: StepResult<TOOLS>['finishReason'];
-  readonly rawFinishReason: StepResult<TOOLS>['rawFinishReason'];
-  readonly usage: StepResult<TOOLS>['usage'];
-  readonly warnings: StepResult<TOOLS>['warnings'];
-  readonly request: StepResult<TOOLS>['request'];
-  readonly response: StepResult<TOOLS>['response'];
-  readonly providerMetadata: StepResult<TOOLS>['providerMetadata'];
+  readonly stepNumber: StepResult<CONTEXT, TOOLS>['stepNumber'];
+  readonly model: StepResult<CONTEXT, TOOLS>['model'];
+  readonly functionId: StepResult<CONTEXT, TOOLS>['functionId'];
+  readonly metadata: StepResult<CONTEXT, TOOLS>['metadata'];
+  readonly experimental_context: StepResult<
+    CONTEXT,
+    TOOLS
+  >['experimental_context'];
+  readonly content: StepResult<CONTEXT, TOOLS>['content'];
+  readonly finishReason: StepResult<CONTEXT, TOOLS>['finishReason'];
+  readonly rawFinishReason: StepResult<CONTEXT, TOOLS>['rawFinishReason'];
+  readonly usage: StepResult<CONTEXT, TOOLS>['usage'];
+  readonly warnings: StepResult<CONTEXT, TOOLS>['warnings'];
+  readonly request: StepResult<CONTEXT, TOOLS>['request'];
+  readonly response: StepResult<CONTEXT, TOOLS>['response'];
+  readonly providerMetadata: StepResult<CONTEXT, TOOLS>['providerMetadata'];
 
   constructor({
     stepNumber,
@@ -199,19 +205,19 @@ export class DefaultStepResult<
     response,
     providerMetadata,
   }: {
-    stepNumber: StepResult<TOOLS>['stepNumber'];
-    model: StepResult<TOOLS>['model'];
-    functionId: StepResult<TOOLS>['functionId'];
-    metadata: StepResult<TOOLS>['metadata'];
-    experimental_context: StepResult<TOOLS>['experimental_context'];
-    content: StepResult<TOOLS>['content'];
-    finishReason: StepResult<TOOLS>['finishReason'];
-    rawFinishReason: StepResult<TOOLS>['rawFinishReason'];
-    usage: StepResult<TOOLS>['usage'];
-    warnings: StepResult<TOOLS>['warnings'];
-    request: StepResult<TOOLS>['request'];
-    response: StepResult<TOOLS>['response'];
-    providerMetadata: StepResult<TOOLS>['providerMetadata'];
+    stepNumber: StepResult<CONTEXT, TOOLS>['stepNumber'];
+    model: StepResult<CONTEXT, TOOLS>['model'];
+    functionId: StepResult<CONTEXT, TOOLS>['functionId'];
+    metadata: StepResult<CONTEXT, TOOLS>['metadata'];
+    experimental_context: StepResult<CONTEXT, TOOLS>['experimental_context'];
+    content: StepResult<CONTEXT, TOOLS>['content'];
+    finishReason: StepResult<CONTEXT, TOOLS>['finishReason'];
+    rawFinishReason: StepResult<CONTEXT, TOOLS>['rawFinishReason'];
+    usage: StepResult<CONTEXT, TOOLS>['usage'];
+    warnings: StepResult<CONTEXT, TOOLS>['warnings'];
+    request: StepResult<CONTEXT, TOOLS>['request'];
+    response: StepResult<CONTEXT, TOOLS>['response'];
+    providerMetadata: StepResult<CONTEXT, TOOLS>['providerMetadata'];
   }) {
     this.stepNumber = stepNumber;
     this.model = model;
diff --git a/packages/ai/src/generate-text/stop-condition.ts b/packages/ai/src/generate-text/stop-condition.ts
index cdcb27734501..239e8d77ff63 100644
--- a/packages/ai/src/generate-text/stop-condition.ts
+++ b/packages/ai/src/generate-text/stop-condition.ts
@@ -2,10 +2,7 @@ import { Context } from '@ai-sdk/provider-utils';
 import { StepResult } from './step-result';
 import { ToolSet } from './tool-set';
 
-export type StopCondition<
-  CONTEXT extends Context,
-  TOOLS extends ToolSet<CONTEXT> = ToolSet<CONTEXT>,
-> = (options: {
+export type StopCondition<TOOLS extends ToolSet> = (options: {
   steps: Array<StepResult<TOOLS>>;
 }) => PromiseLike<boolean> | boolean;
 
@@ -20,14 +17,11 @@ export function hasToolCall(toolName: string): StopCondition<any> {
     ) ?? false;
 }
 
-export async function isStopConditionMet<
-  CONTEXT extends Context,
-  TOOLS extends ToolSet<CONTEXT>,
->({
+export async function isStopConditionMet<TOOLS extends ToolSet>({
   stopConditions,
   steps,
 }: {
-  stopConditions: Array<StopCondition<CONTEXT, TOOLS>>;
+  stopConditions: Array<StopCondition<TOOLS>>;
   steps: Array<StepResult<TOOLS>>;
 }): Promise<boolean> {
   return (
diff --git a/packages/ai/src/generate-text/tool-call.ts b/packages/ai/src/generate-text/tool-call.ts
index 831f93e3db25..59698ce0b550 100644
--- a/packages/ai/src/generate-text/tool-call.ts
+++ b/packages/ai/src/generate-text/tool-call.ts
@@ -10,19 +10,16 @@ type BaseToolCall = {
   providerMetadata?: ProviderMetadata;
 };
 
-export type StaticToolCall<TOOLS extends ToolSet<any> = ToolSet<any>> =
-  ValueOf<{
-    [NAME in keyof TOOLS]: BaseToolCall & {
-      toolName: NAME & string;
-      input: TOOLS[NAME] extends Tool<any, infer PARAMETERS>
-        ? PARAMETERS
-        : never;
-      dynamic?: false | undefined;
-      invalid?: false | undefined;
-      error?: never;
-      title?: string;
-    };
-  }>;
+export type StaticToolCall<TOOLS extends ToolSet = ToolSet> = ValueOf<{
+  [NAME in keyof TOOLS]: BaseToolCall & {
+    toolName: NAME & string;
+    input: TOOLS[NAME] extends Tool<any, infer PARAMETERS> ? PARAMETERS : never;
+    dynamic?: false | undefined;
+    invalid?: false | undefined;
+    error?: never;
+    title?: string;
+  };
+}>;
 
 export type DynamicToolCall = BaseToolCall & {
   toolName: string;
@@ -45,6 +42,6 @@ export type DynamicToolCall = BaseToolCall & {
   error?: unknown;
 };
 
-export type TypedToolCall<TOOLS extends ToolSet<any> = ToolSet<any>> =
+export type TypedToolCall<TOOLS extends ToolSet = ToolSet> =
   | StaticToolCall<TOOLS>
   | DynamicToolCall;
diff --git a/packages/ai/src/generate-text/tool-set.ts b/packages/ai/src/generate-text/tool-set.ts
index 061410448768..3d7b8a9acd42 100644
--- a/packages/ai/src/generate-text/tool-set.ts
+++ b/packages/ai/src/generate-text/tool-set.ts
@@ -1,6 +1,6 @@
 import { Context, Tool } from '@ai-sdk/provider-utils';
 
-export type ToolSet<CONTEXT extends Context> = Record<
+export type ToolSet<CONTEXT extends Context = any> = Record<
   string,
   (
     | Tool<CONTEXT, never, never>
@@ -17,3 +17,6 @@ export type ToolSet<CONTEXT extends Context> = Record<
       | 'needsApproval'
     >
 >;
+
+export type InferContextFromToolSet<TOOLS extends ToolSet> =
+  TOOLS extends ToolSet<infer CONTEXT> ? CONTEXT : never;
diff --git a/packages/provider-utils/src/types/tool.ts b/packages/provider-utils/src/types/tool.ts
index d7745e029767..bb07f386ab39 100644
--- a/packages/provider-utils/src/types/tool.ts
+++ b/packages/provider-utils/src/types/tool.ts
@@ -62,7 +62,7 @@ export type ToolNeedsApprovalFunction<INPUT, CONTEXT extends Context> = (
      *
      * Experimental (can break in patch releases).
      */
-    experimental_context?: CONTEXT;
+    experimental_context: CONTEXT;
   },
 ) => boolean | PromiseLike<boolean>;
 

From 9c2d1fd153e28b1e83f49d19c98111f30ca81457 Mon Sep 17 00:00:00 2001
From: Lars Grammel <lars.grammel@gmail.com>
Date: Thu, 5 Mar 2026 16:57:42 +0100
Subject: [PATCH 08/41] x

---
 .../ai/src/generate-text/generate-text.ts     |  2 +-
 packages/ai/src/generate-text/tool-set.ts     | 10 ++--
 .../ai/src/telemetry/telemetry-integration.ts |  7 ++-
 .../src/provider-tool-factory.ts              | 56 +++++++++----------
 .../provider-utils/src/types/execute-tool.ts  |  4 +-
 .../provider-utils/src/types/tool.test-d.ts   | 20 +++----
 packages/provider-utils/src/types/tool.ts     | 46 +++++++--------
 7 files changed, 74 insertions(+), 71 deletions(-)

diff --git a/packages/ai/src/generate-text/generate-text.ts b/packages/ai/src/generate-text/generate-text.ts
index 3a312235ac43..85f74a9df44e 100644
--- a/packages/ai/src/generate-text/generate-text.ts
+++ b/packages/ai/src/generate-text/generate-text.ts
@@ -250,7 +250,7 @@ export type GenerateTextOnFinishCallback<
  */
 export async function generateText<
   CONTEXT extends Context,
-  TOOLS extends ToolSet<CONTEXT> = ToolSet<CONTEXT>,
+  TOOLS extends ToolSet<CONTEXT>,
   OUTPUT extends Output = Output<string, string>,
 >({
   model: modelArg,
diff --git a/packages/ai/src/generate-text/tool-set.ts b/packages/ai/src/generate-text/tool-set.ts
index 3d7b8a9acd42..89a2658f862f 100644
--- a/packages/ai/src/generate-text/tool-set.ts
+++ b/packages/ai/src/generate-text/tool-set.ts
@@ -3,13 +3,13 @@ import { Context, Tool } from '@ai-sdk/provider-utils';
 export type ToolSet<CONTEXT extends Context = any> = Record<
   string,
   (
-    | Tool<CONTEXT, never, never>
-    | Tool<CONTEXT, any, any>
-    | Tool<CONTEXT, any, never>
-    | Tool<CONTEXT, never, any>
+    | Tool<never, never, CONTEXT>
+    | Tool<any, any, CONTEXT>
+    | Tool<any, never, CONTEXT>
+    | Tool<never, any, CONTEXT>
   ) &
     Pick<
-      Tool<CONTEXT, any, any>,
+      Tool<any, any, CONTEXT>,
       | 'execute'
       | 'onInputAvailable'
       | 'onInputStart'
diff --git a/packages/ai/src/telemetry/telemetry-integration.ts b/packages/ai/src/telemetry/telemetry-integration.ts
index 11d4eff8026b..3958728d9bdd 100644
--- a/packages/ai/src/telemetry/telemetry-integration.ts
+++ b/packages/ai/src/telemetry/telemetry-integration.ts
@@ -1,3 +1,4 @@
+import { Context } from '@ai-sdk/provider-utils';
 import type {
   OnFinishEvent,
   OnStartEvent,
@@ -16,9 +17,9 @@ import { Listener } from '../util/notify';
  */
 export interface TelemetryIntegration {
   onStart?: Listener<OnStartEvent<ToolSet, Output>>;
-  onStepStart?: Listener<OnStepStartEvent<ToolSet, Output>>;
+  onStepStart?: Listener<OnStepStartEvent<Context, ToolSet, Output>>;
   onToolCallStart?: Listener<OnToolCallStartEvent<ToolSet>>;
   onToolCallFinish?: Listener<OnToolCallFinishEvent<ToolSet>>;
-  onStepFinish?: Listener<OnStepFinishEvent<ToolSet>>;
-  onFinish?: Listener<OnFinishEvent<ToolSet>>;
+  onStepFinish?: Listener<OnStepFinishEvent<Context, ToolSet>>;
+  onFinish?: Listener<OnFinishEvent<Context, ToolSet>>;
 }
diff --git a/packages/provider-utils/src/provider-tool-factory.ts b/packages/provider-utils/src/provider-tool-factory.ts
index 73590059e9e9..9766b664a530 100644
--- a/packages/provider-utils/src/provider-tool-factory.ts
+++ b/packages/provider-utils/src/provider-tool-factory.ts
@@ -8,14 +8,14 @@ export type ProviderToolFactory<
   ARGS extends object,
 > = <OUTPUT>(
   options: ARGS & {
-    execute?: ToolExecuteFunction<CONTEXT, INPUT, OUTPUT>;
-    needsApproval?: Tool<CONTEXT, INPUT, OUTPUT>['needsApproval'];
-    toModelOutput?: Tool<CONTEXT, INPUT, OUTPUT>['toModelOutput'];
-    onInputStart?: Tool<CONTEXT, INPUT, OUTPUT>['onInputStart'];
-    onInputDelta?: Tool<CONTEXT, INPUT, OUTPUT>['onInputDelta'];
-    onInputAvailable?: Tool<CONTEXT, INPUT, OUTPUT>['onInputAvailable'];
+    execute?: ToolExecuteFunction<INPUT, OUTPUT, CONTEXT>;
+    needsApproval?: Tool<INPUT, OUTPUT, CONTEXT>['needsApproval'];
+    toModelOutput?: Tool<INPUT, OUTPUT, CONTEXT>['toModelOutput'];
+    onInputStart?: Tool<INPUT, OUTPUT, CONTEXT>['onInputStart'];
+    onInputDelta?: Tool<INPUT, OUTPUT, CONTEXT>['onInputDelta'];
+    onInputAvailable?: Tool<INPUT, OUTPUT, CONTEXT>['onInputAvailable'];
   },
-) => Tool<CONTEXT, INPUT, OUTPUT>;
+) => Tool<INPUT, OUTPUT, CONTEXT>;
 
 export function createProviderToolFactory<
   CONTEXT extends Context,
@@ -38,14 +38,14 @@ export function createProviderToolFactory<
     onInputAvailable,
     ...args
   }: ARGS & {
-    execute?: ToolExecuteFunction<CONTEXT, INPUT, OUTPUT>;
+    execute?: ToolExecuteFunction<INPUT, OUTPUT, CONTEXT>;
     outputSchema?: FlexibleSchema<OUTPUT>;
-    needsApproval?: Tool<CONTEXT, INPUT, OUTPUT>['needsApproval'];
-    toModelOutput?: Tool<CONTEXT, INPUT, OUTPUT>['toModelOutput'];
-    onInputStart?: Tool<CONTEXT, INPUT, OUTPUT>['onInputStart'];
-    onInputDelta?: Tool<CONTEXT, INPUT, OUTPUT>['onInputDelta'];
-    onInputAvailable?: Tool<CONTEXT, INPUT, OUTPUT>['onInputAvailable'];
-  }): Tool<CONTEXT, INPUT, OUTPUT> =>
+    needsApproval?: Tool<INPUT, OUTPUT, CONTEXT>['needsApproval'];
+    toModelOutput?: Tool<INPUT, OUTPUT, CONTEXT>['toModelOutput'];
+    onInputStart?: Tool<INPUT, OUTPUT, CONTEXT>['onInputStart'];
+    onInputDelta?: Tool<INPUT, OUTPUT, CONTEXT>['onInputDelta'];
+    onInputAvailable?: Tool<INPUT, OUTPUT, CONTEXT>['onInputAvailable'];
+  }): Tool<INPUT, OUTPUT, CONTEXT> =>
     tool({
       type: 'provider',
       id,
@@ -68,14 +68,14 @@ export type ProviderToolFactoryWithOutputSchema<
   ARGS extends object,
 > = (
   options: ARGS & {
-    execute?: ToolExecuteFunction<CONTEXT, INPUT, OUTPUT>;
-    needsApproval?: Tool<CONTEXT, INPUT, OUTPUT>['needsApproval'];
-    toModelOutput?: Tool<CONTEXT, INPUT, OUTPUT>['toModelOutput'];
-    onInputStart?: Tool<CONTEXT, INPUT, OUTPUT>['onInputStart'];
-    onInputDelta?: Tool<CONTEXT, INPUT, OUTPUT>['onInputDelta'];
-    onInputAvailable?: Tool<CONTEXT, INPUT, OUTPUT>['onInputAvailable'];
+    execute?: ToolExecuteFunction<INPUT, OUTPUT, CONTEXT>;
+    needsApproval?: Tool<INPUT, OUTPUT, CONTEXT>['needsApproval'];
+    toModelOutput?: Tool<INPUT, OUTPUT, CONTEXT>['toModelOutput'];
+    onInputStart?: Tool<INPUT, OUTPUT, CONTEXT>['onInputStart'];
+    onInputDelta?: Tool<INPUT, OUTPUT, CONTEXT>['onInputDelta'];
+    onInputAvailable?: Tool<INPUT, OUTPUT, CONTEXT>['onInputAvailable'];
   },
-) => Tool<CONTEXT, INPUT, OUTPUT>;
+) => Tool<INPUT, OUTPUT, CONTEXT>;
 
 export function createProviderToolFactoryWithOutputSchema<
   CONTEXT extends Context,
@@ -112,13 +112,13 @@ export function createProviderToolFactoryWithOutputSchema<
     onInputAvailable,
     ...args
   }: ARGS & {
-    execute?: ToolExecuteFunction<CONTEXT, INPUT, OUTPUT>;
-    needsApproval?: Tool<CONTEXT, INPUT, OUTPUT>['needsApproval'];
-    toModelOutput?: Tool<CONTEXT, INPUT, OUTPUT>['toModelOutput'];
-    onInputStart?: Tool<CONTEXT, INPUT, OUTPUT>['onInputStart'];
-    onInputDelta?: Tool<CONTEXT, INPUT, OUTPUT>['onInputDelta'];
-    onInputAvailable?: Tool<CONTEXT, INPUT, OUTPUT>['onInputAvailable'];
-  }): Tool<CONTEXT, INPUT, OUTPUT> =>
+    execute?: ToolExecuteFunction<INPUT, OUTPUT, CONTEXT>;
+    needsApproval?: Tool<INPUT, OUTPUT, CONTEXT>['needsApproval'];
+    toModelOutput?: Tool<INPUT, OUTPUT, CONTEXT>['toModelOutput'];
+    onInputStart?: Tool<INPUT, OUTPUT, CONTEXT>['onInputStart'];
+    onInputDelta?: Tool<INPUT, OUTPUT, CONTEXT>['onInputDelta'];
+    onInputAvailable?: Tool<INPUT, OUTPUT, CONTEXT>['onInputAvailable'];
+  }): Tool<INPUT, OUTPUT, CONTEXT> =>
     tool({
       type: 'provider',
       id,
diff --git a/packages/provider-utils/src/types/execute-tool.ts b/packages/provider-utils/src/types/execute-tool.ts
index 68961ccb9c34..8d9add4183a2 100644
--- a/packages/provider-utils/src/types/execute-tool.ts
+++ b/packages/provider-utils/src/types/execute-tool.ts
@@ -2,12 +2,12 @@ import { isAsyncIterable } from '../is-async-iterable';
 import { Context } from './context';
 import { ToolExecuteFunction, ToolExecutionOptions } from './tool';
 
-export async function* executeTool<CONTEXT extends Context, INPUT, OUTPUT>({
+export async function* executeTool<INPUT, OUTPUT, CONTEXT extends Context>({
   execute,
   input,
   options,
 }: {
-  execute: ToolExecuteFunction<CONTEXT, INPUT, OUTPUT>;
+  execute: ToolExecuteFunction<INPUT, OUTPUT, CONTEXT>;
   input: INPUT;
   options: ToolExecutionOptions<CONTEXT>;
 }): AsyncGenerator<
diff --git a/packages/provider-utils/src/types/tool.test-d.ts b/packages/provider-utils/src/types/tool.test-d.ts
index 8a3e9d1fba96..cc80184f7c7d 100644
--- a/packages/provider-utils/src/types/tool.test-d.ts
+++ b/packages/provider-utils/src/types/tool.test-d.ts
@@ -14,7 +14,7 @@ describe('tool type', () => {
       });
 
       expectTypeOf(aTool).toEqualTypeOf<
-        Tool<Context, { number: number }, never>
+        Tool<{ number: number }, never, Context>
       >();
       expectTypeOf(aTool.execute).toEqualTypeOf<undefined>();
       expectTypeOf(aTool.execute).not.toEqualTypeOf<Function>();
@@ -28,7 +28,7 @@ describe('tool type', () => {
         inputSchema: null as unknown as FlexibleSchema<T>,
       });
 
-      expectTypeOf(aTool).toEqualTypeOf<Tool<Context, T, never>>();
+      expectTypeOf(aTool).toEqualTypeOf<Tool<T, never, Context>>();
       expectTypeOf(aTool.execute).toEqualTypeOf<undefined>();
       expectTypeOf(aTool.execute).not.toEqualTypeOf<Function>();
       expectTypeOf(aTool.inputSchema).toEqualTypeOf<FlexibleSchema<T>>();
@@ -81,10 +81,10 @@ describe('tool type', () => {
       });
 
       expectTypeOf(aTool).toEqualTypeOf<
-        Tool<Context, { number: number }, 'test'>
+        Tool<{ number: number }, 'test', Context>
       >();
       expectTypeOf(aTool.execute).toMatchTypeOf<
-        ToolExecuteFunction<Context, { number: number }, 'test'> | undefined
+        ToolExecuteFunction<{ number: number }, 'test', Context> | undefined
       >();
       expectTypeOf(aTool.execute).not.toEqualTypeOf<undefined>();
       expectTypeOf(aTool.inputSchema).toEqualTypeOf<
@@ -101,10 +101,10 @@ describe('tool type', () => {
       });
 
       expectTypeOf(aTool).toEqualTypeOf<
-        Tool<Context, { number: number }, 'test'>
+        Tool<{ number: number }, 'test', Context>
       >();
       expectTypeOf(aTool.execute).toEqualTypeOf<
-        ToolExecuteFunction<Context, { number: number }, 'test'> | undefined
+        ToolExecuteFunction<{ number: number }, 'test', Context> | undefined
       >();
       expectTypeOf(aTool.inputSchema).toEqualTypeOf<
         FlexibleSchema<{ number: number }>
@@ -182,7 +182,7 @@ describe('tool type', () => {
           expectTypeOf(options).toEqualTypeOf<{
             toolCallId: string;
             messages: ModelMessage[];
-            experimental_context?: Context | undefined;
+            experimental_context: Context;
           }>();
           return true;
         },
@@ -195,7 +195,7 @@ describe('tool type', () => {
             options: {
               toolCallId: string;
               messages: ModelMessage[];
-              experimental_context?: Context | undefined;
+              experimental_context: Context;
             },
           ) => boolean | PromiseLike<boolean>)
         | undefined
@@ -211,7 +211,7 @@ describe('tool type', () => {
           expectTypeOf(options).toEqualTypeOf<{
             toolCallId: string;
             messages: ModelMessage[];
-            experimental_context?: Context | undefined;
+            experimental_context: Context;
           }>();
           return true;
         },
@@ -224,7 +224,7 @@ describe('tool type', () => {
             options: {
               toolCallId: string;
               messages: ModelMessage[];
-              experimental_context?: Context | undefined;
+              experimental_context: Context;
             },
           ) => boolean | PromiseLike<boolean>)
         | undefined
diff --git a/packages/provider-utils/src/types/tool.ts b/packages/provider-utils/src/types/tool.ts
index bb07f386ab39..ac3b9fa65e67 100644
--- a/packages/provider-utils/src/types/tool.ts
+++ b/packages/provider-utils/src/types/tool.ts
@@ -66,7 +66,7 @@ export type ToolNeedsApprovalFunction<INPUT, CONTEXT extends Context> = (
   },
 ) => boolean | PromiseLike<boolean>;
 
-export type ToolExecuteFunction<CONTEXT extends Context, INPUT, OUTPUT> = (
+export type ToolExecuteFunction<INPUT, OUTPUT, CONTEXT extends Context> = (
   input: INPUT,
   options: ToolExecutionOptions<CONTEXT>,
 ) => AsyncIterable<OUTPUT> | PromiseLike<OUTPUT> | OUTPUT;
@@ -80,9 +80,9 @@ type NeverOptional<N, T> = 0 extends 1 & N
     : T;
 
 type ToolOutputProperties<
-  CONTEXT extends Context,
   INPUT,
   OUTPUT,
+  CONTEXT extends Context,
 > = NeverOptional<
   OUTPUT,
   | {
@@ -93,7 +93,7 @@ type ToolOutputProperties<
        * @args is the input of the tool call.
        * @options.abortSignal is a signal that can be used to abort the tool call.
        */
-      execute: ToolExecuteFunction<CONTEXT, INPUT, OUTPUT>;
+      execute: ToolExecuteFunction<INPUT, OUTPUT, CONTEXT>;
 
       outputSchema?: FlexibleSchema<OUTPUT>;
     }
@@ -111,9 +111,9 @@ type ToolOutputProperties<
  * The tool can also contain an optional execute function for the actual execution function of the tool.
  */
 export type Tool<
-  CONTEXT extends Context = any,
   INPUT extends JSONValue | unknown | never = any,
   OUTPUT extends JSONValue | unknown | never = any,
+  CONTEXT extends Context = any,
 > = {
   /**
    * An optional description of what the tool does.
@@ -158,7 +158,7 @@ export type Tool<
     | boolean
     | ToolNeedsApprovalFunction<
         [INPUT] extends [never] ? unknown : INPUT,
-        CONTEXT
+        NoInfer<CONTEXT>
       >;
 
   /**
@@ -175,7 +175,7 @@ export type Tool<
    * Only called when the tool is used in a streaming context.
    */
   onInputStart?: (
-    options: ToolExecutionOptions<CONTEXT>,
+    options: ToolExecutionOptions<NoInfer<CONTEXT>>,
   ) => void | PromiseLike<void>;
 
   /**
@@ -183,7 +183,9 @@ export type Tool<
    * Only called when the tool is used in a streaming context.
    */
   onInputDelta?: (
-    options: { inputTextDelta: string } & ToolExecutionOptions<CONTEXT>,
+    options: { inputTextDelta: string } & ToolExecutionOptions<
+      NoInfer<CONTEXT>
+    >,
   ) => void | PromiseLike<void>;
 
   /**
@@ -193,9 +195,9 @@ export type Tool<
   onInputAvailable?: (
     options: {
       input: [INPUT] extends [never] ? unknown : INPUT;
-    } & ToolExecutionOptions<CONTEXT>,
+    } & ToolExecutionOptions<NoInfer<CONTEXT>>,
   ) => void | PromiseLike<void>;
-} & ToolOutputProperties<CONTEXT, INPUT, OUTPUT> & {
+} & ToolOutputProperties<INPUT, OUTPUT, NoInfer<CONTEXT>> & {
     /**
      * Optional conversion function that maps the tool result to an output that can be used by the language model.
      *
@@ -284,18 +286,18 @@ export type InferToolOutput<TOOL extends Tool<any>> =
  * Helper function for inferring the execute args of a tool.
  */
 // Note: overload order is important for auto-completion
-export function tool<CONTEXT extends Context, INPUT, OUTPUT>(
-  tool: Tool<CONTEXT, INPUT, OUTPUT>,
-): Tool<CONTEXT, INPUT, OUTPUT>;
-export function tool<CONTEXT extends Context, INPUT>(
-  tool: Tool<CONTEXT, INPUT, never>,
-): Tool<CONTEXT, INPUT, never>;
-export function tool<CONTEXT extends Context, OUTPUT>(
-  tool: Tool<CONTEXT, never, OUTPUT>,
-): Tool<CONTEXT, never, OUTPUT>;
+export function tool<INPUT, OUTPUT, CONTEXT extends Context>(
+  tool: Tool<INPUT, OUTPUT, CONTEXT>,
+): Tool<INPUT, OUTPUT, CONTEXT>;
+export function tool<INPUT, CONTEXT extends Context>(
+  tool: Tool<INPUT, never, CONTEXT>,
+): Tool<INPUT, never, CONTEXT>;
+export function tool<OUTPUT, CONTEXT extends Context>(
+  tool: Tool<never, OUTPUT, CONTEXT>,
+): Tool<never, OUTPUT, CONTEXT>;
 export function tool<CONTEXT extends Context>(
-  tool: Tool<CONTEXT, never, never>,
-): Tool<CONTEXT, never, never>;
+  tool: Tool<never, never, CONTEXT>,
+): Tool<never, never, CONTEXT>;
 export function tool(tool: any): any {
   return tool;
 }
@@ -308,7 +310,7 @@ export function dynamicTool(tool: {
   title?: string;
   providerOptions?: ProviderOptions;
   inputSchema: FlexibleSchema<unknown>;
-  execute: ToolExecuteFunction<Context, unknown, unknown>;
+  execute: ToolExecuteFunction<unknown, unknown, Context>;
 
   /**
    * Optional conversion function that maps the tool result to an output that can be used by the language model.
@@ -336,7 +338,7 @@ export function dynamicTool(tool: {
    * Whether the tool needs approval before it can be executed.
    */
   needsApproval?: boolean | ToolNeedsApprovalFunction<unknown, Context>;
-}): Tool<Context, unknown, unknown> & {
+}): Tool<unknown, unknown, Context> & {
   type: 'dynamic';
 } {
   return { ...tool, type: 'dynamic' };

From aabce40976dfcc223fbfcb7796ed101c4ab1aa4d Mon Sep 17 00:00:00 2001
From: Lars Grammel <lars.grammel@gmail.com>
Date: Thu, 5 Mar 2026 17:00:48 +0100
Subject: [PATCH 09/41] a

---
 packages/provider-utils/src/types/execute-tool.ts | 2 +-
 packages/provider-utils/src/types/tool.ts         | 8 ++++----
 2 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/packages/provider-utils/src/types/execute-tool.ts b/packages/provider-utils/src/types/execute-tool.ts
index 8d9add4183a2..41f83de3721c 100644
--- a/packages/provider-utils/src/types/execute-tool.ts
+++ b/packages/provider-utils/src/types/execute-tool.ts
@@ -9,7 +9,7 @@ export async function* executeTool<INPUT, OUTPUT, CONTEXT extends Context>({
 }: {
   execute: ToolExecuteFunction<INPUT, OUTPUT, CONTEXT>;
   input: INPUT;
-  options: ToolExecutionOptions<CONTEXT>;
+  options: ToolExecutionOptions<NoInfer<CONTEXT>>;
 }): AsyncGenerator<
   { type: 'preliminary'; output: OUTPUT } | { type: 'final'; output: OUTPUT }
 > {
diff --git a/packages/provider-utils/src/types/tool.ts b/packages/provider-utils/src/types/tool.ts
index ac3b9fa65e67..f6bd56d9dd4f 100644
--- a/packages/provider-utils/src/types/tool.ts
+++ b/packages/provider-utils/src/types/tool.ts
@@ -273,14 +273,14 @@ export type Tool<
 /**
  * Infer the input type of a tool.
  */
-export type InferToolInput<TOOL extends Tool<any>> =
-  TOOL extends Tool<any, infer INPUT, any> ? INPUT : never;
+export type InferToolInput<TOOL extends Tool> =
+  TOOL extends Tool<infer INPUT, any, any> ? INPUT : never;
 
 /**
  * Infer the output type of a tool.
  */
-export type InferToolOutput<TOOL extends Tool<any>> =
-  TOOL extends Tool<any, any, infer OUTPUT> ? OUTPUT : never;
+export type InferToolOutput<TOOL extends Tool> =
+  TOOL extends Tool<any, infer OUTPUT, any> ? OUTPUT : never;
 
 /**
  * Helper function for inferring the execute args of a tool.

From 50c741f5cd21fa2c7484d8e7be91e74fe9d4d0e8 Mon Sep 17 00:00:00 2001
From: Lars Grammel <lars.grammel@gmail.com>
Date: Thu, 5 Mar 2026 17:13:42 +0100
Subject: [PATCH 10/41] x

---
 .../ai/src/generate-text/callback-events.ts   |  6 +--
 .../ai/src/generate-text/execute-tool-call.ts |  2 +-
 .../ai/src/generate-text/generate-text.ts     | 48 +++++++++++--------
 packages/ai/src/generate-text/tool-set.ts     | 17 ++++---
 4 files changed, 41 insertions(+), 32 deletions(-)

diff --git a/packages/ai/src/generate-text/callback-events.ts b/packages/ai/src/generate-text/callback-events.ts
index 697a711002d3..c0b7c9d79f61 100644
--- a/packages/ai/src/generate-text/callback-events.ts
+++ b/packages/ai/src/generate-text/callback-events.ts
@@ -131,7 +131,7 @@ export interface OnStartEvent<
  */
 export interface OnStepStartEvent<
   CONTEXT extends Context,
-  TOOLS extends ToolSet<CONTEXT>,
+  TOOLS extends ToolSet,
   OUTPUT extends Output = Output,
   INCLUDE = { requestBody?: boolean; responseBody?: boolean },
 > {
@@ -302,7 +302,7 @@ export type OnToolCallFinishEvent<TOOLS extends ToolSet> = {
  */
 export type OnStepFinishEvent<
   CONTEXT extends Context,
-  TOOLS extends ToolSet<CONTEXT>,
+  TOOLS extends ToolSet,
 > = StepResult<CONTEXT, TOOLS>;
 
 /**
@@ -313,7 +313,7 @@ export type OnStepFinishEvent<
  */
 export type OnFinishEvent<
   CONTEXT extends Context,
-  TOOLS extends ToolSet<CONTEXT>,
+  TOOLS extends ToolSet,
 > = StepResult<CONTEXT, TOOLS> & {
   /** Array containing results from all steps in the generation. */
   readonly steps: StepResult<CONTEXT, TOOLS>[];
diff --git a/packages/ai/src/generate-text/execute-tool-call.ts b/packages/ai/src/generate-text/execute-tool-call.ts
index 80a311f0c90d..a52391a49356 100644
--- a/packages/ai/src/generate-text/execute-tool-call.ts
+++ b/packages/ai/src/generate-text/execute-tool-call.ts
@@ -29,7 +29,7 @@ import { TypedToolError } from './tool-error';
  */
 export async function executeToolCall<
   CONTEXT extends Context,
-  TOOLS extends ToolSet<CONTEXT>,
+  TOOLS extends ToolSet,
 >({
   toolCall,
   tools,
diff --git a/packages/ai/src/generate-text/generate-text.ts b/packages/ai/src/generate-text/generate-text.ts
index 85f74a9df44e..b356c8b11e03 100644
--- a/packages/ai/src/generate-text/generate-text.ts
+++ b/packages/ai/src/generate-text/generate-text.ts
@@ -131,10 +131,15 @@ export type GenerateTextOnStartCallback<
  */
 export type GenerateTextOnStepStartCallback<
   CONTEXT extends Context,
-  TOOLS extends ToolSet<CONTEXT> = ToolSet<CONTEXT>,
+  TOOLS extends ToolSet = ToolSet,
   OUTPUT extends Output = Output,
 > = (
-  event: OnStepStartEvent<CONTEXT, TOOLS, OUTPUT, GenerateTextIncludeSettings>,
+  event: OnStepStartEvent<
+    NoInfer<CONTEXT>,
+    TOOLS,
+    OUTPUT,
+    GenerateTextIncludeSettings
+  >,
 ) => PromiseLike<void> | void;
 
 /**
@@ -146,8 +151,7 @@ export type GenerateTextOnStepStartCallback<
  * @param event - The event object containing tool call information.
  */
 export type GenerateTextOnToolCallStartCallback<
-  CONTEXT extends Context,
-  TOOLS extends ToolSet<CONTEXT> = ToolSet<CONTEXT>,
+  TOOLS extends ToolSet = ToolSet,
 > = (event: OnToolCallStartEvent<TOOLS>) => PromiseLike<void> | void;
 
 /**
@@ -163,8 +167,7 @@ export type GenerateTextOnToolCallStartCallback<
  * @param event - The event object containing tool call result information.
  */
 export type GenerateTextOnToolCallFinishCallback<
-  CONTEXT extends Context,
-  TOOLS extends ToolSet<CONTEXT> = ToolSet<CONTEXT>,
+  TOOLS extends ToolSet = ToolSet,
 > = (event: OnToolCallFinishEvent<TOOLS>) => PromiseLike<void> | void;
 
 /**
@@ -177,7 +180,7 @@ export type GenerateTextOnToolCallFinishCallback<
  */
 export type GenerateTextOnStepFinishCallback<
   CONTEXT extends Context,
-  TOOLS extends ToolSet<CONTEXT> = ToolSet<CONTEXT>,
+  TOOLS extends ToolSet,
 > = (event: OnStepFinishEvent<CONTEXT, TOOLS>) => Promise<void> | void;
 
 /**
@@ -191,7 +194,7 @@ export type GenerateTextOnStepFinishCallback<
  */
 export type GenerateTextOnFinishCallback<
   CONTEXT extends Context,
-  TOOLS extends ToolSet<CONTEXT> = ToolSet<CONTEXT>,
+  TOOLS extends ToolSet = ToolSet,
 > = (event: OnFinishEvent<CONTEXT, TOOLS>) => PromiseLike<void> | void;
 
 /**
@@ -250,7 +253,7 @@ export type GenerateTextOnFinishCallback<
  */
 export async function generateText<
   CONTEXT extends Context,
-  TOOLS extends ToolSet<CONTEXT>,
+  TOOLS extends ToolSet,
   OUTPUT extends Output = Output<string, string>,
 >({
   model: modelArg,
@@ -274,7 +277,7 @@ export async function generateText<
   prepareStep = experimental_prepareStep,
   experimental_repairToolCall: repairToolCall,
   experimental_download: download,
-  experimental_context = {} as CONTEXT,
+  experimental_context = {} as NoInfer<CONTEXT>,
   experimental_include: include,
   _internal: { generateId = originalGenerateId } = {},
   experimental_onStart: onStart,
@@ -356,12 +359,15 @@ export async function generateText<
     /**
      * @deprecated Use `prepareStep` instead.
      */
-    experimental_prepareStep?: PrepareStepFunction<CONTEXT, NoInfer<TOOLS>>;
+    experimental_prepareStep?: PrepareStepFunction<
+      NoInfer<CONTEXT>,
+      NoInfer<TOOLS>
+    >;
 
     /**
      * Optional function that you can use to provide different settings for a step.
      */
-    prepareStep?: PrepareStepFunction<CONTEXT, NoInfer<TOOLS>>;
+    prepareStep?: PrepareStepFunction<NoInfer<CONTEXT>, NoInfer<TOOLS>>;
 
     /**
      * A function that attempts to repair a tool call that failed to parse.
@@ -379,7 +385,7 @@ export async function generateText<
      * before the provider is called.
      */
     experimental_onStepStart?: GenerateTextOnStepStartCallback<
-      CONTEXT,
+      NoInfer<CONTEXT>,
       NoInfer<TOOLS>,
       OUTPUT
     >;
@@ -401,12 +407,15 @@ export async function generateText<
     /**
      * Callback that is called when each step (LLM call) is finished, including intermediate steps.
      */
-    onStepFinish?: GenerateTextOnStepFinishCallback<CONTEXT, NoInfer<TOOLS>>;
+    onStepFinish?: GenerateTextOnStepFinishCallback<
+      NoInfer<CONTEXT>,
+      NoInfer<TOOLS>
+    >;
 
     /**
      * Callback that is called when all steps are finished and the response is complete.
      */
-    onFinish?: GenerateTextOnFinishCallback<CONTEXT, NoInfer<TOOLS>>;
+    onFinish?: GenerateTextOnFinishCallback<NoInfer<CONTEXT>, NoInfer<TOOLS>>;
 
     /**
      * Context that is passed into tool execution.
@@ -1227,10 +1236,7 @@ export async function generateText<
   }
 }
 
-async function executeTools<
-  CONTEXT extends Context,
-  TOOLS extends ToolSet<CONTEXT> = ToolSet<CONTEXT>,
->({
+async function executeTools<CONTEXT extends Context, TOOLS extends ToolSet>({
   toolCalls,
   tools,
   tracer,
@@ -1286,7 +1292,7 @@ async function executeTools<
 
 class DefaultGenerateTextResult<
   CONTEXT extends Context,
-  TOOLS extends ToolSet<CONTEXT>,
+  TOOLS extends ToolSet,
   OUTPUT extends Output,
 > implements GenerateTextResult<TOOLS, OUTPUT>
 {
@@ -1413,7 +1419,7 @@ function asToolCalls(content: Array<LanguageModelV4Content>) {
   }));
 }
 
-function asContent<TOOLS extends ToolSet<any>>({
+function asContent<TOOLS extends ToolSet>({
   content,
   toolCalls,
   toolOutputs,
diff --git a/packages/ai/src/generate-text/tool-set.ts b/packages/ai/src/generate-text/tool-set.ts
index 89a2658f862f..3abaa47fd0c0 100644
--- a/packages/ai/src/generate-text/tool-set.ts
+++ b/packages/ai/src/generate-text/tool-set.ts
@@ -1,15 +1,15 @@
 import { Context, Tool } from '@ai-sdk/provider-utils';
 
-export type ToolSet<CONTEXT extends Context = any> = Record<
+export type ToolSet = Record<
   string,
   (
-    | Tool<never, never, CONTEXT>
-    | Tool<any, any, CONTEXT>
-    | Tool<any, never, CONTEXT>
-    | Tool<never, any, CONTEXT>
+    | Tool<never, never, any>
+    | Tool<any, any, any>
+    | Tool<any, never, any>
+    | Tool<never, any, any>
   ) &
     Pick<
-      Tool<any, any, CONTEXT>,
+      Tool<any, any, any>,
       | 'execute'
       | 'onInputAvailable'
       | 'onInputStart'
@@ -18,5 +18,8 @@ export type ToolSet<CONTEXT extends Context = any> = Record<
     >
 >;
 
+export type InferContextFromTool<TOOL extends Tool> =
+  TOOL extends Tool<any, any, infer CONTEXT> ? CONTEXT : never;
+
 export type InferContextFromToolSet<TOOLS extends ToolSet> =
-  TOOLS extends ToolSet<infer CONTEXT> ? CONTEXT : never;
+  TOOLS extends ToolSet ? InferContextFromTool<TOOLS[keyof TOOLS]> : never;

From 1427dcc008ff1ed91d6df135584c0c6f16354d1a Mon Sep 17 00:00:00 2001
From: Lars Grammel <lars.grammel@gmail.com>
Date: Thu, 5 Mar 2026 17:18:32 +0100
Subject: [PATCH 11/41] Revert "x"

This reverts commit 50c741f5cd21fa2c7484d8e7be91e74fe9d4d0e8.
---
 .../ai/src/generate-text/callback-events.ts   |  6 +--
 .../ai/src/generate-text/execute-tool-call.ts |  2 +-
 .../ai/src/generate-text/generate-text.ts     | 48 ++++++++-----------
 packages/ai/src/generate-text/tool-set.ts     | 17 +++----
 4 files changed, 32 insertions(+), 41 deletions(-)

diff --git a/packages/ai/src/generate-text/callback-events.ts b/packages/ai/src/generate-text/callback-events.ts
index c0b7c9d79f61..697a711002d3 100644
--- a/packages/ai/src/generate-text/callback-events.ts
+++ b/packages/ai/src/generate-text/callback-events.ts
@@ -131,7 +131,7 @@ export interface OnStartEvent<
  */
 export interface OnStepStartEvent<
   CONTEXT extends Context,
-  TOOLS extends ToolSet,
+  TOOLS extends ToolSet<CONTEXT>,
   OUTPUT extends Output = Output,
   INCLUDE = { requestBody?: boolean; responseBody?: boolean },
 > {
@@ -302,7 +302,7 @@ export type OnToolCallFinishEvent<TOOLS extends ToolSet> = {
  */
 export type OnStepFinishEvent<
   CONTEXT extends Context,
-  TOOLS extends ToolSet,
+  TOOLS extends ToolSet<CONTEXT>,
 > = StepResult<CONTEXT, TOOLS>;
 
 /**
@@ -313,7 +313,7 @@ export type OnStepFinishEvent<
  */
 export type OnFinishEvent<
   CONTEXT extends Context,
-  TOOLS extends ToolSet,
+  TOOLS extends ToolSet<CONTEXT>,
 > = StepResult<CONTEXT, TOOLS> & {
   /** Array containing results from all steps in the generation. */
   readonly steps: StepResult<CONTEXT, TOOLS>[];
diff --git a/packages/ai/src/generate-text/execute-tool-call.ts b/packages/ai/src/generate-text/execute-tool-call.ts
index a52391a49356..80a311f0c90d 100644
--- a/packages/ai/src/generate-text/execute-tool-call.ts
+++ b/packages/ai/src/generate-text/execute-tool-call.ts
@@ -29,7 +29,7 @@ import { TypedToolError } from './tool-error';
  */
 export async function executeToolCall<
   CONTEXT extends Context,
-  TOOLS extends ToolSet,
+  TOOLS extends ToolSet<CONTEXT>,
 >({
   toolCall,
   tools,
diff --git a/packages/ai/src/generate-text/generate-text.ts b/packages/ai/src/generate-text/generate-text.ts
index b356c8b11e03..85f74a9df44e 100644
--- a/packages/ai/src/generate-text/generate-text.ts
+++ b/packages/ai/src/generate-text/generate-text.ts
@@ -131,15 +131,10 @@ export type GenerateTextOnStartCallback<
  */
 export type GenerateTextOnStepStartCallback<
   CONTEXT extends Context,
-  TOOLS extends ToolSet = ToolSet,
+  TOOLS extends ToolSet<CONTEXT> = ToolSet<CONTEXT>,
   OUTPUT extends Output = Output,
 > = (
-  event: OnStepStartEvent<
-    NoInfer<CONTEXT>,
-    TOOLS,
-    OUTPUT,
-    GenerateTextIncludeSettings
-  >,
+  event: OnStepStartEvent<CONTEXT, TOOLS, OUTPUT, GenerateTextIncludeSettings>,
 ) => PromiseLike<void> | void;
 
 /**
@@ -151,7 +146,8 @@ export type GenerateTextOnStepStartCallback<
  * @param event - The event object containing tool call information.
  */
 export type GenerateTextOnToolCallStartCallback<
-  TOOLS extends ToolSet = ToolSet,
+  CONTEXT extends Context,
+  TOOLS extends ToolSet<CONTEXT> = ToolSet<CONTEXT>,
 > = (event: OnToolCallStartEvent<TOOLS>) => PromiseLike<void> | void;
 
 /**
@@ -167,7 +163,8 @@ export type GenerateTextOnToolCallStartCallback<
  * @param event - The event object containing tool call result information.
  */
 export type GenerateTextOnToolCallFinishCallback<
-  TOOLS extends ToolSet = ToolSet,
+  CONTEXT extends Context,
+  TOOLS extends ToolSet<CONTEXT> = ToolSet<CONTEXT>,
 > = (event: OnToolCallFinishEvent<TOOLS>) => PromiseLike<void> | void;
 
 /**
@@ -180,7 +177,7 @@ export type GenerateTextOnToolCallFinishCallback<
  */
 export type GenerateTextOnStepFinishCallback<
   CONTEXT extends Context,
-  TOOLS extends ToolSet,
+  TOOLS extends ToolSet<CONTEXT> = ToolSet<CONTEXT>,
 > = (event: OnStepFinishEvent<CONTEXT, TOOLS>) => Promise<void> | void;
 
 /**
@@ -194,7 +191,7 @@ export type GenerateTextOnStepFinishCallback<
  */
 export type GenerateTextOnFinishCallback<
   CONTEXT extends Context,
-  TOOLS extends ToolSet = ToolSet,
+  TOOLS extends ToolSet<CONTEXT> = ToolSet<CONTEXT>,
 > = (event: OnFinishEvent<CONTEXT, TOOLS>) => PromiseLike<void> | void;
 
 /**
@@ -253,7 +250,7 @@ export type GenerateTextOnFinishCallback<
  */
 export async function generateText<
   CONTEXT extends Context,
-  TOOLS extends ToolSet,
+  TOOLS extends ToolSet<CONTEXT>,
   OUTPUT extends Output = Output<string, string>,
 >({
   model: modelArg,
@@ -277,7 +274,7 @@ export async function generateText<
   prepareStep = experimental_prepareStep,
   experimental_repairToolCall: repairToolCall,
   experimental_download: download,
-  experimental_context = {} as NoInfer<CONTEXT>,
+  experimental_context = {} as CONTEXT,
   experimental_include: include,
   _internal: { generateId = originalGenerateId } = {},
   experimental_onStart: onStart,
@@ -359,15 +356,12 @@ export async function generateText<
     /**
      * @deprecated Use `prepareStep` instead.
      */
-    experimental_prepareStep?: PrepareStepFunction<
-      NoInfer<CONTEXT>,
-      NoInfer<TOOLS>
-    >;
+    experimental_prepareStep?: PrepareStepFunction<CONTEXT, NoInfer<TOOLS>>;
 
     /**
      * Optional function that you can use to provide different settings for a step.
      */
-    prepareStep?: PrepareStepFunction<NoInfer<CONTEXT>, NoInfer<TOOLS>>;
+    prepareStep?: PrepareStepFunction<CONTEXT, NoInfer<TOOLS>>;
 
     /**
      * A function that attempts to repair a tool call that failed to parse.
@@ -385,7 +379,7 @@ export async function generateText<
      * before the provider is called.
      */
     experimental_onStepStart?: GenerateTextOnStepStartCallback<
-      NoInfer<CONTEXT>,
+      CONTEXT,
       NoInfer<TOOLS>,
       OUTPUT
     >;
@@ -407,15 +401,12 @@ export async function generateText<
     /**
      * Callback that is called when each step (LLM call) is finished, including intermediate steps.
      */
-    onStepFinish?: GenerateTextOnStepFinishCallback<
-      NoInfer<CONTEXT>,
-      NoInfer<TOOLS>
-    >;
+    onStepFinish?: GenerateTextOnStepFinishCallback<CONTEXT, NoInfer<TOOLS>>;
 
     /**
      * Callback that is called when all steps are finished and the response is complete.
      */
-    onFinish?: GenerateTextOnFinishCallback<NoInfer<CONTEXT>, NoInfer<TOOLS>>;
+    onFinish?: GenerateTextOnFinishCallback<CONTEXT, NoInfer<TOOLS>>;
 
     /**
      * Context that is passed into tool execution.
@@ -1236,7 +1227,10 @@ export async function generateText<
   }
 }
 
-async function executeTools<CONTEXT extends Context, TOOLS extends ToolSet>({
+async function executeTools<
+  CONTEXT extends Context,
+  TOOLS extends ToolSet<CONTEXT> = ToolSet<CONTEXT>,
+>({
   toolCalls,
   tools,
   tracer,
@@ -1292,7 +1286,7 @@ async function executeTools<CONTEXT extends Context, TOOLS extends ToolSet>({
 
 class DefaultGenerateTextResult<
   CONTEXT extends Context,
-  TOOLS extends ToolSet,
+  TOOLS extends ToolSet<CONTEXT>,
   OUTPUT extends Output,
 > implements GenerateTextResult<TOOLS, OUTPUT>
 {
@@ -1419,7 +1413,7 @@ function asToolCalls(content: Array<LanguageModelV4Content>) {
   }));
 }
 
-function asContent<TOOLS extends ToolSet>({
+function asContent<TOOLS extends ToolSet<any>>({
   content,
   toolCalls,
   toolOutputs,
diff --git a/packages/ai/src/generate-text/tool-set.ts b/packages/ai/src/generate-text/tool-set.ts
index 3abaa47fd0c0..89a2658f862f 100644
--- a/packages/ai/src/generate-text/tool-set.ts
+++ b/packages/ai/src/generate-text/tool-set.ts
@@ -1,15 +1,15 @@
 import { Context, Tool } from '@ai-sdk/provider-utils';
 
-export type ToolSet = Record<
+export type ToolSet<CONTEXT extends Context = any> = Record<
   string,
   (
-    | Tool<never, never, any>
-    | Tool<any, any, any>
-    | Tool<any, never, any>
-    | Tool<never, any, any>
+    | Tool<never, never, CONTEXT>
+    | Tool<any, any, CONTEXT>
+    | Tool<any, never, CONTEXT>
+    | Tool<never, any, CONTEXT>
   ) &
     Pick<
-      Tool<any, any, any>,
+      Tool<any, any, CONTEXT>,
       | 'execute'
       | 'onInputAvailable'
       | 'onInputStart'
@@ -18,8 +18,5 @@ export type ToolSet = Record<
     >
 >;
 
-export type InferContextFromTool<TOOL extends Tool> =
-  TOOL extends Tool<any, any, infer CONTEXT> ? CONTEXT : never;
-
 export type InferContextFromToolSet<TOOLS extends ToolSet> =
-  TOOLS extends ToolSet ? InferContextFromTool<TOOLS[keyof TOOLS]> : never;
+  TOOLS extends ToolSet<infer CONTEXT> ? CONTEXT : never;

From 920adb5aac8ffde76ff381fbfb1525656dfea30e Mon Sep 17 00:00:00 2001
From: Lars Grammel <lars.grammel@gmail.com>
Date: Fri, 6 Mar 2026 15:50:43 +0100
Subject: [PATCH 12/41] x

---
 .../ai-functions/src/test/typed-context-2.ts  | 72 +++++++++++++++++++
 1 file changed, 72 insertions(+)
 create mode 100644 examples/ai-functions/src/test/typed-context-2.ts

diff --git a/examples/ai-functions/src/test/typed-context-2.ts b/examples/ai-functions/src/test/typed-context-2.ts
new file mode 100644
index 000000000000..25b0f3c429b6
--- /dev/null
+++ b/examples/ai-functions/src/test/typed-context-2.ts
@@ -0,0 +1,72 @@
+import { FlexibleSchema } from 'ai';
+import { run } from '../lib/run';
+import { z } from 'zod';
+
+interface Tool<INPUT = any, CONTEXT = any> {
+  inputSchema: FlexibleSchema<INPUT>;
+  contextSchema: FlexibleSchema<CONTEXT>;
+  execute: (input: NoInfer<INPUT>, context: NoInfer<CONTEXT>) => unknown;
+}
+
+function tool<INPUT, CONTEXT>(options: {
+  inputSchema: FlexibleSchema<INPUT>;
+  contextSchema: FlexibleSchema<CONTEXT>;
+  execute: (input: NoInfer<INPUT>, context: NoInfer<CONTEXT>) => unknown;
+}) {
+  return options;
+}
+
+type InferToolInput<TOOL extends Tool> =
+  TOOL extends Tool<infer INPUT, any> ? INPUT : never;
+type InferToolContext<TOOL extends Tool> =
+  TOOL extends Tool<any, infer CONTEXT> ? CONTEXT : never;
+
+export type ToolSet = Record<string, Tool<any, any>>;
+
+// should be a union of all the context types of the tools
+type InferToolSetContext<TOOLS extends ToolSet> = {
+  [K in keyof TOOLS]: InferToolContext<TOOLS[K]>;
+}[keyof TOOLS];
+
+function executeTool<CONTEXT, TOOLS extends ToolSet>({
+  tools,
+  toolName,
+  input,
+  context,
+}: {
+  tools: TOOLS;
+  toolName: keyof TOOLS;
+  input: InferToolInput<TOOLS[typeof toolName]>;
+  context: InferToolSetContext<TOOLS>;
+}) {
+  const tool = tools[toolName];
+  return tool.execute(input, context);
+}
+
+run(async () => {
+  const tool1 = tool({
+    inputSchema: z.object({ input1: z.string() }),
+    contextSchema: z.object({ context1: z.number() }),
+    execute: async ({ input1 }, { context1 }) => {
+      console.log(input1, context1);
+    },
+  });
+
+  const tool2 = tool({
+    inputSchema: z.object({ input2: z.number() }),
+    contextSchema: z.object({ context2: z.string() }),
+    execute: async ({ input2 }, { context2 }) => {
+      console.log(input2, context2);
+    },
+  });
+
+  executeTool({
+    tools: {
+      tool1,
+      tool2,
+    },
+    toolName: 'tool1',
+    input: { input1: 'Hello' },
+    context: { context1: 1, context2: 'Hello' },
+  });
+});

From ba939e152a06f5f65eef64fb93467c4e085083d9 Mon Sep 17 00:00:00 2001
From: Lars Grammel <lars.grammel@gmail.com>
Date: Fri, 6 Mar 2026 15:57:46 +0100
Subject: [PATCH 13/41] perf

---
 .../ai-functions/src/test/typed-context-2.ts   | 18 +++++++++++++-----
 1 file changed, 13 insertions(+), 5 deletions(-)

diff --git a/examples/ai-functions/src/test/typed-context-2.ts b/examples/ai-functions/src/test/typed-context-2.ts
index 25b0f3c429b6..25d9e4502afe 100644
--- a/examples/ai-functions/src/test/typed-context-2.ts
+++ b/examples/ai-functions/src/test/typed-context-2.ts
@@ -23,10 +23,18 @@ type InferToolContext<TOOL extends Tool> =
 
 export type ToolSet = Record<string, Tool<any, any>>;
 
+type UnionToIntersection<U> = (
+  U extends unknown ? (arg: U) => void : never
+) extends (arg: infer I) => void
+  ? I
+  : never;
+
 // should be a union of all the context types of the tools
-type InferToolSetContext<TOOLS extends ToolSet> = {
-  [K in keyof TOOLS]: InferToolContext<TOOLS[K]>;
-}[keyof TOOLS];
+type InferToolSetContext<TOOLS extends ToolSet> = UnionToIntersection<
+  {
+    [K in keyof TOOLS]: InferToolContext<TOOLS[K]>;
+  }[keyof TOOLS]
+>;
 
 function executeTool<CONTEXT, TOOLS extends ToolSet>({
   tools,
@@ -37,7 +45,7 @@ function executeTool<CONTEXT, TOOLS extends ToolSet>({
   tools: TOOLS;
   toolName: keyof TOOLS;
   input: InferToolInput<TOOLS[typeof toolName]>;
-  context: InferToolSetContext<TOOLS>;
+  context: InferToolSetContext<TOOLS> & Record<string, unknown>;
 }) {
   const tool = tools[toolName];
   return tool.execute(input, context);
@@ -67,6 +75,6 @@ run(async () => {
     },
     toolName: 'tool1',
     input: { input1: 'Hello' },
-    context: { context1: 1, context2: 'Hello' },
+    context: { context1: 1, context2: 'world', somethingElse: 'context' },
   });
 });

From 509417d363105fb99463f59a5a9d7df3cb878e1b Mon Sep 17 00:00:00 2001
From: Lars Grammel <lars.grammel@gmail.com>
Date: Mon, 9 Mar 2026 14:01:52 +0100
Subject: [PATCH 14/41] a

---
 examples/ai-functions/src/test/typed-context-2.ts | 12 ++++++++++--
 1 file changed, 10 insertions(+), 2 deletions(-)

diff --git a/examples/ai-functions/src/test/typed-context-2.ts b/examples/ai-functions/src/test/typed-context-2.ts
index 25d9e4502afe..818a8ea1966a 100644
--- a/examples/ai-functions/src/test/typed-context-2.ts
+++ b/examples/ai-functions/src/test/typed-context-2.ts
@@ -36,7 +36,11 @@ type InferToolSetContext<TOOLS extends ToolSet> = UnionToIntersection<
   }[keyof TOOLS]
 >;
 
-function executeTool<CONTEXT, TOOLS extends ToolSet>({
+// TODO prepareStep
+function executeTool<
+  TOOLS extends ToolSet,
+  CONTEXT extends InferToolSetContext<TOOLS> & Record<string, unknown>,
+>({
   tools,
   toolName,
   input,
@@ -45,7 +49,8 @@ function executeTool<CONTEXT, TOOLS extends ToolSet>({
   tools: TOOLS;
   toolName: keyof TOOLS;
   input: InferToolInput<TOOLS[typeof toolName]>;
-  context: InferToolSetContext<TOOLS> & Record<string, unknown>;
+  context: CONTEXT;
+  prepareStep: (context: CONTEXT) => void;
 }) {
   const tool = tools[toolName];
   return tool.execute(input, context);
@@ -76,5 +81,8 @@ run(async () => {
     toolName: 'tool1',
     input: { input1: 'Hello' },
     context: { context1: 1, context2: 'world', somethingElse: 'context' },
+    prepareStep: context => {
+      console.log(context);
+    },
   });
 });

From 8e273175a68e047cc9f2cd5396fa0cec4f5146c9 Mon Sep 17 00:00:00 2001
From: Lars Grammel <lars.grammel@gmail.com>
Date: Mon, 9 Mar 2026 14:02:15 +0100
Subject: [PATCH 15/41] x

---
 examples/ai-functions/src/test/typed-context-2.ts | 1 -
 1 file changed, 1 deletion(-)

diff --git a/examples/ai-functions/src/test/typed-context-2.ts b/examples/ai-functions/src/test/typed-context-2.ts
index 818a8ea1966a..143a92f9285f 100644
--- a/examples/ai-functions/src/test/typed-context-2.ts
+++ b/examples/ai-functions/src/test/typed-context-2.ts
@@ -36,7 +36,6 @@ type InferToolSetContext<TOOLS extends ToolSet> = UnionToIntersection<
   }[keyof TOOLS]
 >;
 
-// TODO prepareStep
 function executeTool<
   TOOLS extends ToolSet,
   CONTEXT extends InferToolSetContext<TOOLS> & Record<string, unknown>,

From 7e9fbd9c5cc9a0ea3b18a8599f5dd8e283d48003 Mon Sep 17 00:00:00 2001
From: Lars Grammel <lars.grammel@gmail.com>
Date: Mon, 9 Mar 2026 14:05:02 +0100
Subject: [PATCH 16/41] x

---
 examples/ai-functions/src/test/typed-context-2.ts | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/examples/ai-functions/src/test/typed-context-2.ts b/examples/ai-functions/src/test/typed-context-2.ts
index 143a92f9285f..339238bab09b 100644
--- a/examples/ai-functions/src/test/typed-context-2.ts
+++ b/examples/ai-functions/src/test/typed-context-2.ts
@@ -36,9 +36,12 @@ type InferToolSetContext<TOOLS extends ToolSet> = UnionToIntersection<
   }[keyof TOOLS]
 >;
 
+type ExpandedContext<TOOLS extends ToolSet> = InferToolSetContext<TOOLS> &
+  Record<string, unknown>;
+
 function executeTool<
   TOOLS extends ToolSet,
-  CONTEXT extends InferToolSetContext<TOOLS> & Record<string, unknown>,
+  CONTEXT extends ExpandedContext<TOOLS>,
 >({
   tools,
   toolName,

From c5bc3810d70328e97c767a101e284c3f9b079187 Mon Sep 17 00:00:00 2001
From: Lars Grammel <lars.grammel@gmail.com>
Date: Mon, 9 Mar 2026 14:06:45 +0100
Subject: [PATCH 17/41] x

---
 examples/ai-functions/src/test/typed-context-2.ts | 8 +++++---
 1 file changed, 5 insertions(+), 3 deletions(-)

diff --git a/examples/ai-functions/src/test/typed-context-2.ts b/examples/ai-functions/src/test/typed-context-2.ts
index 339238bab09b..e34507f0d7ea 100644
--- a/examples/ai-functions/src/test/typed-context-2.ts
+++ b/examples/ai-functions/src/test/typed-context-2.ts
@@ -2,13 +2,15 @@ import { FlexibleSchema } from 'ai';
 import { run } from '../lib/run';
 import { z } from 'zod';
 
-interface Tool<INPUT = any, CONTEXT = any> {
+type Context = Record<string, unknown>;
+
+interface Tool<INPUT = any, CONTEXT extends Context = Context> {
   inputSchema: FlexibleSchema<INPUT>;
   contextSchema: FlexibleSchema<CONTEXT>;
   execute: (input: NoInfer<INPUT>, context: NoInfer<CONTEXT>) => unknown;
 }
 
-function tool<INPUT, CONTEXT>(options: {
+function tool<INPUT, CONTEXT extends Context>(options: {
   inputSchema: FlexibleSchema<INPUT>;
   contextSchema: FlexibleSchema<CONTEXT>;
   execute: (input: NoInfer<INPUT>, context: NoInfer<CONTEXT>) => unknown;
@@ -37,7 +39,7 @@ type InferToolSetContext<TOOLS extends ToolSet> = UnionToIntersection<
 >;
 
 type ExpandedContext<TOOLS extends ToolSet> = InferToolSetContext<TOOLS> &
-  Record<string, unknown>;
+  Context;
 
 function executeTool<
   TOOLS extends ToolSet,

From c41b21cb5875fc2f0a66fea02a5e991d23317309 Mon Sep 17 00:00:00 2001
From: Lars Grammel <lars.grammel@gmail.com>
Date: Mon, 9 Mar 2026 14:31:29 +0100
Subject: [PATCH 18/41] x

---
 .../ai/src/generate-text/callback-events.ts   | 23 +++---
 .../ai/src/generate-text/execute-tool-call.ts |  6 +-
 .../ai/src/generate-text/generate-text.ts     | 77 ++++++++++---------
 packages/ai/src/generate-text/step-result.ts  | 66 ++++++++--------
 .../ai/src/generate-text/stop-condition.ts    | 22 ++++--
 .../tool-approval-request-output.ts           |  8 +-
 .../tool-call-repair-function.ts              |  6 +-
 packages/ai/src/generate-text/tool-error.ts   | 16 ++--
 packages/ai/src/generate-text/tool-set.ts     | 25 +++---
 .../ai/src/telemetry/telemetry-integration.ts |  6 +-
 packages/ai/src/util/union-to-intersection.ts | 17 ++++
 packages/provider-utils/src/types/index.ts    |  3 +-
 packages/provider-utils/src/types/tool.ts     |  8 +-
 13 files changed, 155 insertions(+), 128 deletions(-)
 create mode 100644 packages/ai/src/util/union-to-intersection.ts

diff --git a/packages/ai/src/generate-text/callback-events.ts b/packages/ai/src/generate-text/callback-events.ts
index 697a711002d3..0968b784ee8f 100644
--- a/packages/ai/src/generate-text/callback-events.ts
+++ b/packages/ai/src/generate-text/callback-events.ts
@@ -1,6 +1,5 @@
 import type { LanguageModelV4ToolChoice } from '@ai-sdk/provider';
 import type {
-  Context,
   ModelMessage,
   ProviderOptions,
   SystemModelMessage,
@@ -12,7 +11,7 @@ import type { Output } from './output';
 import type { StepResult } from './step-result';
 import type { StopCondition } from './stop-condition';
 import type { TypedToolCall } from './tool-call';
-import type { ToolSet } from './tool-set';
+import type { ExpandedContext, ToolSet } from './tool-set';
 
 /**
  * Common model information used across callback events.
@@ -130,8 +129,8 @@ export interface OnStartEvent<
  * Each step represents a single LLM invocation.
  */
 export interface OnStepStartEvent<
-  CONTEXT extends Context,
-  TOOLS extends ToolSet<CONTEXT>,
+  TOOLS extends ToolSet,
+  CONTEXT extends ExpandedContext<TOOLS>,
   OUTPUT extends Output = Output,
   INCLUDE = { requestBody?: boolean; responseBody?: boolean },
 > {
@@ -167,7 +166,7 @@ export interface OnStepStartEvent<
   readonly activeTools: Array<keyof TOOLS> | undefined;
 
   /** Array of results from previous steps (empty for first step). */
-  readonly steps: ReadonlyArray<StepResult<CONTEXT, TOOLS>>;
+  readonly steps: ReadonlyArray<StepResult<TOOLS, CONTEXT>>;
 
   /** Additional provider-specific options for this step. */
   readonly providerOptions: ProviderOptions | undefined;
@@ -301,9 +300,9 @@ export type OnToolCallFinishEvent<TOOLS extends ToolSet> = {
  * This is simply the StepResult for that step.
  */
 export type OnStepFinishEvent<
-  CONTEXT extends Context,
-  TOOLS extends ToolSet<CONTEXT>,
-> = StepResult<CONTEXT, TOOLS>;
+  TOOLS extends ToolSet,
+  CONTEXT extends ExpandedContext<TOOLS>,
+> = StepResult<TOOLS, CONTEXT>;
 
 /**
  * Event passed to the `onFinish` callback.
@@ -312,11 +311,11 @@ export type OnStepFinishEvent<
  * Includes the final step's result along with aggregated data from all steps.
  */
 export type OnFinishEvent<
-  CONTEXT extends Context,
-  TOOLS extends ToolSet<CONTEXT>,
-> = StepResult<CONTEXT, TOOLS> & {
+  TOOLS extends ToolSet,
+  CONTEXT extends ExpandedContext<TOOLS>,
+> = StepResult<TOOLS, CONTEXT> & {
   /** Array containing results from all steps in the generation. */
-  readonly steps: StepResult<CONTEXT, TOOLS>[];
+  readonly steps: StepResult<TOOLS, CONTEXT>[];
 
   /** Aggregated token usage across all steps. */
   readonly totalUsage: LanguageModelUsage;
diff --git a/packages/ai/src/generate-text/execute-tool-call.ts b/packages/ai/src/generate-text/execute-tool-call.ts
index 80a311f0c90d..ce9f5ef4ce80 100644
--- a/packages/ai/src/generate-text/execute-tool-call.ts
+++ b/packages/ai/src/generate-text/execute-tool-call.ts
@@ -12,7 +12,7 @@ import {
 } from './generate-text';
 import { TypedToolCall } from './tool-call';
 import { ToolOutput } from './tool-output';
-import { ToolSet } from './tool-set';
+import { ExpandedContext, ToolSet } from './tool-set';
 import { TypedToolResult } from './tool-result';
 import { TypedToolError } from './tool-error';
 
@@ -28,8 +28,8 @@ import { TypedToolError } from './tool-error';
  * @returns The tool output (result or error), or undefined if the tool has no execute function.
  */
 export async function executeToolCall<
-  CONTEXT extends Context,
-  TOOLS extends ToolSet<CONTEXT>,
+  TOOLS extends ToolSet,
+  CONTEXT extends ExpandedContext<TOOLS>,
 >({
   toolCall,
   tools,
diff --git a/packages/ai/src/generate-text/generate-text.ts b/packages/ai/src/generate-text/generate-text.ts
index 63c575dc75a1..bbe46c70147c 100644
--- a/packages/ai/src/generate-text/generate-text.ts
+++ b/packages/ai/src/generate-text/generate-text.ts
@@ -89,7 +89,7 @@ import { ToolCallRepairFunction } from './tool-call-repair-function';
 import { TypedToolError } from './tool-error';
 import { ToolOutput } from './tool-output';
 import { TypedToolResult } from './tool-result';
-import { ToolSet } from './tool-set';
+import { ExpandedContext, InferToolSetContext, ToolSet } from './tool-set';
 
 const originalGenerateId = createIdGenerator({
   prefix: 'aitxt',
@@ -130,11 +130,11 @@ export type GenerateTextOnStartCallback<
  * @param event - The event object containing step configuration.
  */
 export type GenerateTextOnStepStartCallback<
-  CONTEXT extends Context,
-  TOOLS extends ToolSet<CONTEXT> = ToolSet<CONTEXT>,
+  TOOLS extends ToolSet,
+  CONTEXT extends ExpandedContext<TOOLS>,
   OUTPUT extends Output = Output,
 > = (
-  event: OnStepStartEvent<CONTEXT, TOOLS, OUTPUT, GenerateTextIncludeSettings>,
+  event: OnStepStartEvent<TOOLS, CONTEXT, OUTPUT, GenerateTextIncludeSettings>,
 ) => PromiseLike<void> | void;
 
 /**
@@ -145,10 +145,9 @@ export type GenerateTextOnStepStartCallback<
  *
  * @param event - The event object containing tool call information.
  */
-export type GenerateTextOnToolCallStartCallback<
-  CONTEXT extends Context,
-  TOOLS extends ToolSet<CONTEXT> = ToolSet<CONTEXT>,
-> = (event: OnToolCallStartEvent<TOOLS>) => PromiseLike<void> | void;
+export type GenerateTextOnToolCallStartCallback<TOOLS extends ToolSet> = (
+  event: OnToolCallStartEvent<TOOLS>,
+) => PromiseLike<void> | void;
 
 /**
  * Callback that is set using the `experimental_onToolCallFinish` option.
@@ -162,10 +161,9 @@ export type GenerateTextOnToolCallStartCallback<
  *
  * @param event - The event object containing tool call result information.
  */
-export type GenerateTextOnToolCallFinishCallback<
-  CONTEXT extends Context,
-  TOOLS extends ToolSet<CONTEXT> = ToolSet<CONTEXT>,
-> = (event: OnToolCallFinishEvent<TOOLS>) => PromiseLike<void> | void;
+export type GenerateTextOnToolCallFinishCallback<TOOLS extends ToolSet> = (
+  event: OnToolCallFinishEvent<TOOLS>,
+) => PromiseLike<void> | void;
 
 /**
  * Callback that is set using the `onStepFinish` option.
@@ -176,9 +174,9 @@ export type GenerateTextOnToolCallFinishCallback<
  * @param stepResult - The result of the step.
  */
 export type GenerateTextOnStepFinishCallback<
-  CONTEXT extends Context,
-  TOOLS extends ToolSet<CONTEXT> = ToolSet<CONTEXT>,
-> = (event: OnStepFinishEvent<CONTEXT, TOOLS>) => Promise<void> | void;
+  TOOLS extends ToolSet,
+  CONTEXT extends ExpandedContext<TOOLS>,
+> = (event: OnStepFinishEvent<TOOLS, CONTEXT>) => Promise<void> | void;
 
 /**
  * Callback that is set using the `onFinish` option.
@@ -190,9 +188,9 @@ export type GenerateTextOnStepFinishCallback<
  * @param event - The final result along with aggregated step data.
  */
 export type GenerateTextOnFinishCallback<
-  CONTEXT extends Context,
-  TOOLS extends ToolSet<CONTEXT> = ToolSet<CONTEXT>,
-> = (event: OnFinishEvent<CONTEXT, TOOLS>) => PromiseLike<void> | void;
+  TOOLS extends ToolSet,
+  CONTEXT extends ExpandedContext<TOOLS>,
+> = (event: OnFinishEvent<TOOLS, CONTEXT>) => PromiseLike<void> | void;
 
 /**
  * Generate a text and call tools for a given prompt using a language model.
@@ -249,8 +247,8 @@ export type GenerateTextOnFinishCallback<
  * A result object that contains the generated text, the results of the tool calls, and additional information.
  */
 export async function generateText<
-  CONTEXT extends Context,
-  TOOLS extends ToolSet<CONTEXT>,
+  TOOLS extends ToolSet,
+  CONTEXT extends ExpandedContext<TOOLS>,
   OUTPUT extends Output = Output<string, string>,
 >({
   model: modelArg,
@@ -308,8 +306,8 @@ export async function generateText<
      * @default stepCountIs(1)
      */
     stopWhen?:
-      | StopCondition<NoInfer<TOOLS>>
-      | Array<StopCondition<NoInfer<TOOLS>>>;
+      | StopCondition<NoInfer<TOOLS>, CONTEXT>
+      | Array<StopCondition<NoInfer<TOOLS>, CONTEXT>>;
 
     /**
      * Optional telemetry configuration (experimental).
@@ -379,8 +377,8 @@ export async function generateText<
      * before the provider is called.
      */
     experimental_onStepStart?: GenerateTextOnStepStartCallback<
+      TOOLS,
       CONTEXT,
-      NoInfer<TOOLS>,
       OUTPUT
     >;
 
@@ -401,12 +399,12 @@ export async function generateText<
     /**
      * Callback that is called when each step (LLM call) is finished, including intermediate steps.
      */
-    onStepFinish?: GenerateTextOnStepFinishCallback<CONTEXT, NoInfer<TOOLS>>;
+    onStepFinish?: GenerateTextOnStepFinishCallback<TOOLS, CONTEXT>;
 
     /**
      * Callback that is called when all steps are finished and the response is complete.
      */
-    onFinish?: GenerateTextOnFinishCallback<CONTEXT, NoInfer<TOOLS>>;
+    onFinish?: GenerateTextOnFinishCallback<TOOLS, CONTEXT>;
 
     /**
      * Context that is passed into tool execution.
@@ -667,7 +665,7 @@ export async function generateText<
         > & { response: { id: string; timestamp: Date; modelId: string } };
         let clientToolCalls: Array<TypedToolCall<TOOLS>> = [];
         let clientToolOutputs: Array<ToolOutput<TOOLS>> = [];
-        const steps: StepResult<CONTEXT, TOOLS>[] = [];
+        const steps: StepResult<TOOLS, CONTEXT>[] = [];
 
         // Track provider-executed tool calls that support deferred results
         // (e.g., code_execution in programmatic tool calling scenarios).
@@ -763,7 +761,7 @@ export async function generateText<
                 onStepStart,
                 globalTelemetry.onStepStart as
                   | undefined
-                  | GenerateTextOnStepStartCallback<CONTEXT, TOOLS, OUTPUT>,
+                  | GenerateTextOnStepStartCallback<TOOLS, CONTEXT, OUTPUT>,
               ],
             });
 
@@ -1082,7 +1080,7 @@ export async function generateText<
 
             const stepNumber = steps.length;
 
-            const currentStepResult: StepResult<CONTEXT, TOOLS> =
+            const currentStepResult: StepResult<TOOLS, CONTEXT> =
               new DefaultStepResult({
                 stepNumber,
                 model: stepModelInfo,
@@ -1111,7 +1109,12 @@ export async function generateText<
 
             await notify({
               event: currentStepResult,
-              callbacks: [onStepFinish, globalTelemetry.onStepFinish],
+              callbacks: [
+                onStepFinish,
+                globalTelemetry.onStepFinish as
+                  | undefined
+                  | GenerateTextOnStepFinishCallback<TOOLS, CONTEXT>,
+              ],
             });
           } finally {
             if (stepTimeoutId != null) {
@@ -1231,7 +1234,7 @@ export async function generateText<
             onFinish,
             globalTelemetry.onFinish as
               | undefined
-              | GenerateTextOnFinishCallback<CONTEXT, TOOLS>,
+              | GenerateTextOnFinishCallback<TOOLS, CONTEXT>,
           ],
         });
 
@@ -1262,8 +1265,8 @@ export async function generateText<
 }
 
 async function executeTools<
-  CONTEXT extends Context,
-  TOOLS extends ToolSet<CONTEXT> = ToolSet<CONTEXT>,
+  TOOLS extends ToolSet,
+  CONTEXT extends ExpandedContext<TOOLS>,
 >({
   toolCalls,
   tools,
@@ -1319,17 +1322,17 @@ async function executeTools<
 }
 
 class DefaultGenerateTextResult<
-  CONTEXT extends Context,
-  TOOLS extends ToolSet<CONTEXT>,
+  TOOLS extends ToolSet,
+  CONTEXT extends ExpandedContext<TOOLS>,
   OUTPUT extends Output,
 > implements GenerateTextResult<TOOLS, OUTPUT>
 {
-  readonly steps: StepResult<CONTEXT, TOOLS>[];
+  readonly steps: StepResult<TOOLS, CONTEXT>[];
   readonly totalUsage: LanguageModelUsage;
   private readonly _output: InferCompleteOutput<OUTPUT> | undefined;
 
   constructor(options: {
-    steps: StepResult<CONTEXT, TOOLS>[];
+    steps: StepResult<TOOLS, CONTEXT>[];
     output: InferCompleteOutput<OUTPUT> | undefined;
     totalUsage: LanguageModelUsage;
   }) {
@@ -1447,7 +1450,7 @@ function asToolCalls(content: Array<LanguageModelV4Content>) {
   }));
 }
 
-function asContent<TOOLS extends ToolSet<any>>({
+function asContent<TOOLS extends ToolSet>({
   content,
   toolCalls,
   toolOutputs,
diff --git a/packages/ai/src/generate-text/step-result.ts b/packages/ai/src/generate-text/step-result.ts
index 6574053e66a5..dee7eeba39ab 100644
--- a/packages/ai/src/generate-text/step-result.ts
+++ b/packages/ai/src/generate-text/step-result.ts
@@ -17,14 +17,14 @@ import {
   StaticToolResult,
   TypedToolResult,
 } from './tool-result';
-import { ToolSet } from './tool-set';
+import { ExpandedContext, ToolSet } from './tool-set';
 
 /**
  * The result of a single step in the generation process.
  */
 export type StepResult<
-  CONTEXT extends Context,
-  TOOLS extends ToolSet<CONTEXT>,
+  TOOLS extends ToolSet,
+  CONTEXT extends ExpandedContext<TOOLS>,
 > = {
   /**
    * Zero-based index of this step.
@@ -169,26 +169,26 @@ export type StepResult<
 };
 
 export class DefaultStepResult<
-  CONTEXT extends Context,
-  TOOLS extends ToolSet<CONTEXT>,
-> implements StepResult<CONTEXT, TOOLS>
+  TOOLS extends ToolSet,
+  CONTEXT extends ExpandedContext<TOOLS>,
+> implements StepResult<TOOLS, CONTEXT>
 {
-  readonly stepNumber: StepResult<CONTEXT, TOOLS>['stepNumber'];
-  readonly model: StepResult<CONTEXT, TOOLS>['model'];
-  readonly functionId: StepResult<CONTEXT, TOOLS>['functionId'];
-  readonly metadata: StepResult<CONTEXT, TOOLS>['metadata'];
+  readonly stepNumber: StepResult<TOOLS, CONTEXT>['stepNumber'];
+  readonly model: StepResult<TOOLS, CONTEXT>['model'];
+  readonly functionId: StepResult<TOOLS, CONTEXT>['functionId'];
+  readonly metadata: StepResult<TOOLS, CONTEXT>['metadata'];
   readonly experimental_context: StepResult<
-    CONTEXT,
-    TOOLS
+    TOOLS,
+    CONTEXT
   >['experimental_context'];
-  readonly content: StepResult<CONTEXT, TOOLS>['content'];
-  readonly finishReason: StepResult<CONTEXT, TOOLS>['finishReason'];
-  readonly rawFinishReason: StepResult<CONTEXT, TOOLS>['rawFinishReason'];
-  readonly usage: StepResult<CONTEXT, TOOLS>['usage'];
-  readonly warnings: StepResult<CONTEXT, TOOLS>['warnings'];
-  readonly request: StepResult<CONTEXT, TOOLS>['request'];
-  readonly response: StepResult<CONTEXT, TOOLS>['response'];
-  readonly providerMetadata: StepResult<CONTEXT, TOOLS>['providerMetadata'];
+  readonly content: StepResult<TOOLS, CONTEXT>['content'];
+  readonly finishReason: StepResult<TOOLS, CONTEXT>['finishReason'];
+  readonly rawFinishReason: StepResult<TOOLS, CONTEXT>['rawFinishReason'];
+  readonly usage: StepResult<TOOLS, CONTEXT>['usage'];
+  readonly warnings: StepResult<TOOLS, CONTEXT>['warnings'];
+  readonly request: StepResult<TOOLS, CONTEXT>['request'];
+  readonly response: StepResult<TOOLS, CONTEXT>['response'];
+  readonly providerMetadata: StepResult<TOOLS, CONTEXT>['providerMetadata'];
 
   constructor({
     stepNumber,
@@ -205,19 +205,19 @@ export class DefaultStepResult<
     response,
     providerMetadata,
   }: {
-    stepNumber: StepResult<CONTEXT, TOOLS>['stepNumber'];
-    model: StepResult<CONTEXT, TOOLS>['model'];
-    functionId: StepResult<CONTEXT, TOOLS>['functionId'];
-    metadata: StepResult<CONTEXT, TOOLS>['metadata'];
-    experimental_context: StepResult<CONTEXT, TOOLS>['experimental_context'];
-    content: StepResult<CONTEXT, TOOLS>['content'];
-    finishReason: StepResult<CONTEXT, TOOLS>['finishReason'];
-    rawFinishReason: StepResult<CONTEXT, TOOLS>['rawFinishReason'];
-    usage: StepResult<CONTEXT, TOOLS>['usage'];
-    warnings: StepResult<CONTEXT, TOOLS>['warnings'];
-    request: StepResult<CONTEXT, TOOLS>['request'];
-    response: StepResult<CONTEXT, TOOLS>['response'];
-    providerMetadata: StepResult<CONTEXT, TOOLS>['providerMetadata'];
+    stepNumber: StepResult<TOOLS, CONTEXT>['stepNumber'];
+    model: StepResult<TOOLS, CONTEXT>['model'];
+    functionId: StepResult<TOOLS, CONTEXT>['functionId'];
+    metadata: StepResult<TOOLS, CONTEXT>['metadata'];
+    experimental_context: StepResult<TOOLS, CONTEXT>['experimental_context'];
+    content: StepResult<TOOLS, CONTEXT>['content'];
+    finishReason: StepResult<TOOLS, CONTEXT>['finishReason'];
+    rawFinishReason: StepResult<TOOLS, CONTEXT>['rawFinishReason'];
+    usage: StepResult<TOOLS, CONTEXT>['usage'];
+    warnings: StepResult<TOOLS, CONTEXT>['warnings'];
+    request: StepResult<TOOLS, CONTEXT>['request'];
+    response: StepResult<TOOLS, CONTEXT>['response'];
+    providerMetadata: StepResult<TOOLS, CONTEXT>['providerMetadata'];
   }) {
     this.stepNumber = stepNumber;
     this.model = model;
diff --git a/packages/ai/src/generate-text/stop-condition.ts b/packages/ai/src/generate-text/stop-condition.ts
index 239e8d77ff63..8973618a279b 100644
--- a/packages/ai/src/generate-text/stop-condition.ts
+++ b/packages/ai/src/generate-text/stop-condition.ts
@@ -1,28 +1,34 @@
 import { Context } from '@ai-sdk/provider-utils';
 import { StepResult } from './step-result';
-import { ToolSet } from './tool-set';
+import { ExpandedContext, ToolSet } from './tool-set';
 
-export type StopCondition<TOOLS extends ToolSet> = (options: {
-  steps: Array<StepResult<TOOLS>>;
+export type StopCondition<
+  TOOLS extends ToolSet,
+  CONTEXT extends ExpandedContext<TOOLS>,
+> = (options: {
+  steps: Array<StepResult<TOOLS, CONTEXT>>;
 }) => PromiseLike<boolean> | boolean;
 
-export function stepCountIs(stepCount: number): StopCondition<any> {
+export function stepCountIs(stepCount: number): StopCondition<any, any> {
   return ({ steps }) => steps.length === stepCount;
 }
 
-export function hasToolCall(toolName: string): StopCondition<any> {
+export function hasToolCall(toolName: string): StopCondition<any, any> {
   return ({ steps }) =>
     steps[steps.length - 1]?.toolCalls?.some(
       toolCall => toolCall.toolName === toolName,
     ) ?? false;
 }
 
-export async function isStopConditionMet<TOOLS extends ToolSet>({
+export async function isStopConditionMet<
+  TOOLS extends ToolSet,
+  CONTEXT extends ExpandedContext<TOOLS>,
+>({
   stopConditions,
   steps,
 }: {
-  stopConditions: Array<StopCondition<TOOLS>>;
-  steps: Array<StepResult<TOOLS>>;
+  stopConditions: Array<StopCondition<TOOLS, CONTEXT>>;
+  steps: Array<StepResult<TOOLS, CONTEXT>>;
 }): Promise<boolean> {
   return (
     await Promise.all(stopConditions.map(condition => condition({ steps })))
diff --git a/packages/ai/src/generate-text/tool-approval-request-output.ts b/packages/ai/src/generate-text/tool-approval-request-output.ts
index b43ac9f11bc2..2880b67e0a78 100644
--- a/packages/ai/src/generate-text/tool-approval-request-output.ts
+++ b/packages/ai/src/generate-text/tool-approval-request-output.ts
@@ -1,4 +1,3 @@
-import { ContextRegistry } from '@ai-sdk/provider-utils';
 import { TypedToolCall } from './tool-call';
 import { ToolSet } from './tool-set';
 
@@ -7,10 +6,7 @@ import { ToolSet } from './tool-set';
  *
  * The tool approval request can be approved or denied in the next tool message.
  */
-export type ToolApprovalRequestOutput<
-  CONTEXT extends Partial<ContextRegistry>,
-  TOOLS extends ToolSet<CONTEXT> = ToolSet<CONTEXT>,
-> = {
+export type ToolApprovalRequestOutput<TOOLS extends ToolSet> = {
   type: 'tool-approval-request';
 
   /**
@@ -21,5 +17,5 @@ export type ToolApprovalRequestOutput<
   /**
    * Tool call that the approval request is for.
    */
-  toolCall: TypedToolCall<CONTEXT, TOOLS>;
+  toolCall: TypedToolCall<TOOLS>;
 };
diff --git a/packages/ai/src/generate-text/tool-call-repair-function.ts b/packages/ai/src/generate-text/tool-call-repair-function.ts
index 49f6f7e62983..d747916c6ff0 100644
--- a/packages/ai/src/generate-text/tool-call-repair-function.ts
+++ b/packages/ai/src/generate-text/tool-call-repair-function.ts
@@ -1,5 +1,4 @@
 import { JSONSchema7, LanguageModelV4ToolCall } from '@ai-sdk/provider';
-import { Context } from '@ai-sdk/provider-utils';
 import { InvalidToolInputError } from '../error/invalid-tool-input-error';
 import { NoSuchToolError } from '../error/no-such-tool-error';
 import { ModelMessage, SystemModelMessage } from '../prompt';
@@ -18,10 +17,7 @@ import { ToolSet } from './tool-set';
  * @param options.inputSchema - A function that returns the JSON Schema for a tool.
  * @param options.error - The error that occurred while parsing the tool call.
  */
-export type ToolCallRepairFunction<
-  CONTEXT extends Context,
-  TOOLS extends ToolSet<CONTEXT> = ToolSet<CONTEXT>,
-> = (options: {
+export type ToolCallRepairFunction<TOOLS extends ToolSet> = (options: {
   system: string | SystemModelMessage | Array<SystemModelMessage> | undefined;
   messages: ModelMessage[];
   toolCall: LanguageModelV4ToolCall;
diff --git a/packages/ai/src/generate-text/tool-error.ts b/packages/ai/src/generate-text/tool-error.ts
index da2bca8d2d00..12d0540ef8f7 100644
--- a/packages/ai/src/generate-text/tool-error.ts
+++ b/packages/ai/src/generate-text/tool-error.ts
@@ -1,17 +1,14 @@
-import { ContextRegistry, InferToolInput } from '@ai-sdk/provider-utils';
+import { InferToolInput } from '@ai-sdk/provider-utils';
 import { ProviderMetadata } from '../types';
 import { ValueOf } from '../util/value-of';
 import { ToolSet } from './tool-set';
 
-export type StaticToolError<
-  CONTEXT extends Partial<ContextRegistry>,
-  TOOLS extends ToolSet<CONTEXT> = ToolSet<CONTEXT>,
-> = ValueOf<{
+export type StaticToolError<TOOLS extends ToolSet> = ValueOf<{
   [NAME in keyof TOOLS]: {
     type: 'tool-error';
     toolCallId: string;
     toolName: NAME & string;
-    input: InferToolInput<CONTEXT, TOOLS[NAME]>;
+    input: InferToolInput<TOOLS[NAME]>;
     error: unknown;
     providerExecuted?: boolean;
     providerMetadata?: ProviderMetadata;
@@ -32,7 +29,6 @@ export type DynamicToolError = {
   title?: string;
 };
 
-export type TypedToolError<
-  CONTEXT extends Partial<ContextRegistry>,
-  TOOLS extends ToolSet<CONTEXT> = ToolSet<CONTEXT>,
-> = StaticToolError<TOOLS> | DynamicToolError;
+export type TypedToolError<TOOLS extends ToolSet> =
+  | StaticToolError<TOOLS>
+  | DynamicToolError;
diff --git a/packages/ai/src/generate-text/tool-set.ts b/packages/ai/src/generate-text/tool-set.ts
index 89a2658f862f..e6b039a9ddb5 100644
--- a/packages/ai/src/generate-text/tool-set.ts
+++ b/packages/ai/src/generate-text/tool-set.ts
@@ -1,15 +1,16 @@
-import { Context, Tool } from '@ai-sdk/provider-utils';
+import { Context, InferToolContext, Tool } from '@ai-sdk/provider-utils';
+import { UnionToIntersection } from '../util/union-to-intersection';
 
-export type ToolSet<CONTEXT extends Context = any> = Record<
+export type ToolSet = Record<
   string,
   (
-    | Tool<never, never, CONTEXT>
-    | Tool<any, any, CONTEXT>
-    | Tool<any, never, CONTEXT>
-    | Tool<never, any, CONTEXT>
+    | Tool<never, never, Context>
+    | Tool<any, any, Context>
+    | Tool<any, never, Context>
+    | Tool<never, any, Context>
   ) &
     Pick<
-      Tool<any, any, CONTEXT>,
+      Tool<any, any, Context>,
       | 'execute'
       | 'onInputAvailable'
       | 'onInputStart'
@@ -18,5 +19,11 @@ export type ToolSet<CONTEXT extends Context = any> = Record<
     >
 >;
 
-export type InferContextFromToolSet<TOOLS extends ToolSet> =
-  TOOLS extends ToolSet<infer CONTEXT> ? CONTEXT : never;
+export type InferToolSetContext<TOOLS extends ToolSet> = UnionToIntersection<
+  {
+    [K in keyof TOOLS]: InferToolContext<TOOLS[K]>;
+  }[keyof TOOLS]
+>;
+
+export type ExpandedContext<TOOLS extends ToolSet> =
+  InferToolSetContext<TOOLS> & Context;
diff --git a/packages/ai/src/telemetry/telemetry-integration.ts b/packages/ai/src/telemetry/telemetry-integration.ts
index 3958728d9bdd..fe4e710581f7 100644
--- a/packages/ai/src/telemetry/telemetry-integration.ts
+++ b/packages/ai/src/telemetry/telemetry-integration.ts
@@ -17,9 +17,9 @@ import { Listener } from '../util/notify';
  */
 export interface TelemetryIntegration {
   onStart?: Listener<OnStartEvent<ToolSet, Output>>;
-  onStepStart?: Listener<OnStepStartEvent<Context, ToolSet, Output>>;
+  onStepStart?: Listener<OnStepStartEvent<ToolSet, Context, Output>>;
   onToolCallStart?: Listener<OnToolCallStartEvent<ToolSet>>;
   onToolCallFinish?: Listener<OnToolCallFinishEvent<ToolSet>>;
-  onStepFinish?: Listener<OnStepFinishEvent<Context, ToolSet>>;
-  onFinish?: Listener<OnFinishEvent<Context, ToolSet>>;
+  onStepFinish?: Listener<OnStepFinishEvent<ToolSet, Context>>;
+  onFinish?: Listener<OnFinishEvent<ToolSet, Context>>;
 }
diff --git a/packages/ai/src/util/union-to-intersection.ts b/packages/ai/src/util/union-to-intersection.ts
new file mode 100644
index 000000000000..63186a2bde79
--- /dev/null
+++ b/packages/ai/src/util/union-to-intersection.ts
@@ -0,0 +1,17 @@
+/**
+ * Converts a union type `U` into an intersection type.
+ *
+ * For example:
+ *   type A = { a: number };
+ *   type B = { b: string };
+ *   type Union = A | B;
+ *   type Intersection = UnionToIntersection<Union>;
+ *   // Intersection is: { a: number } & { b: string }
+ *
+ * This is useful when you have a union of object types and need a type with all possible properties.
+ */
+export type UnionToIntersection<U> = (
+  U extends unknown ? (arg: U) => void : never
+) extends (arg: infer I) => void
+  ? I
+  : never;
diff --git a/packages/provider-utils/src/types/index.ts b/packages/provider-utils/src/types/index.ts
index c70878137ece..f2b0ac372cc9 100644
--- a/packages/provider-utils/src/types/index.ts
+++ b/packages/provider-utils/src/types/index.ts
@@ -20,6 +20,7 @@ export type { SystemModelMessage } from './system-model-message';
 export {
   dynamicTool,
   tool,
+  type InferToolContext,
   type InferToolInput,
   type InferToolOutput,
   type Tool,
@@ -34,8 +35,8 @@ export type { ToolContent, ToolModelMessage } from './tool-model-message';
 export type { ToolResult } from './tool-result';
 export type { UserContent, UserModelMessage } from './user-model-message';
 
-import type { ToolExecutionOptions } from './tool';
 import type { Context } from './context';
+import type { ToolExecutionOptions } from './tool';
 
 /**
  * @deprecated Use ToolExecutionOptions instead.
diff --git a/packages/provider-utils/src/types/tool.ts b/packages/provider-utils/src/types/tool.ts
index f6bd56d9dd4f..67b028cfb93b 100644
--- a/packages/provider-utils/src/types/tool.ts
+++ b/packages/provider-utils/src/types/tool.ts
@@ -113,7 +113,7 @@ type ToolOutputProperties<
 export type Tool<
   INPUT extends JSONValue | unknown | never = any,
   OUTPUT extends JSONValue | unknown | never = any,
-  CONTEXT extends Context = any,
+  CONTEXT extends Context = Context,
 > = {
   /**
    * An optional description of what the tool does.
@@ -282,6 +282,12 @@ export type InferToolInput<TOOL extends Tool> =
 export type InferToolOutput<TOOL extends Tool> =
   TOOL extends Tool<any, infer OUTPUT, any> ? OUTPUT : never;
 
+/**
+ * Infer the context type of a tool.
+ */
+export type InferToolContext<TOOL extends Tool> =
+  TOOL extends Tool<any, any, infer CONTEXT> ? CONTEXT : never;
+
 /**
  * Helper function for inferring the execute args of a tool.
  */

From c8bac431be0f166bc4d8b77b8f57acf38489b91a Mon Sep 17 00:00:00 2001
From: Lars Grammel <lars.grammel@gmail.com>
Date: Mon, 9 Mar 2026 14:35:30 +0100
Subject: [PATCH 19/41] x

---
 packages/ai/src/generate-text/callback-events.ts      |  9 +++++----
 packages/ai/src/generate-text/generate-text.ts        | 11 ++++++++---
 .../ai/src/generate-text/run-tools-transformation.ts  |  9 ++++++---
 packages/ai/src/telemetry/telemetry-integration.ts    |  2 +-
 4 files changed, 20 insertions(+), 11 deletions(-)

diff --git a/packages/ai/src/generate-text/callback-events.ts b/packages/ai/src/generate-text/callback-events.ts
index 0968b784ee8f..e3f45081fc8a 100644
--- a/packages/ai/src/generate-text/callback-events.ts
+++ b/packages/ai/src/generate-text/callback-events.ts
@@ -30,6 +30,7 @@ export interface CallbackModelInfo {
  */
 export interface OnStartEvent<
   TOOLS extends ToolSet,
+  CONTEXT extends ExpandedContext<TOOLS>,
   OUTPUT extends Output = Output,
   INCLUDE = { requestBody?: boolean; responseBody?: boolean },
 > {
@@ -94,8 +95,8 @@ export interface OnStartEvent<
    * When the condition is an array, any of the conditions can be met to stop.
    */
   readonly stopWhen:
-    | StopCondition<TOOLS>
-    | Array<StopCondition<TOOLS>>
+    | StopCondition<TOOLS, CONTEXT>
+    | Array<StopCondition<TOOLS, CONTEXT>>
     | undefined;
 
   /** The output specification for structured outputs, if configured. */
@@ -185,8 +186,8 @@ export interface OnStepStartEvent<
    * When the condition is an array, any of the conditions can be met to stop.
    */
   readonly stopWhen:
-    | StopCondition<TOOLS>
-    | Array<StopCondition<TOOLS>>
+    | StopCondition<TOOLS, CONTEXT>
+    | Array<StopCondition<TOOLS, CONTEXT>>
     | undefined;
 
   /** The output specification for structured outputs, if configured. */
diff --git a/packages/ai/src/generate-text/generate-text.ts b/packages/ai/src/generate-text/generate-text.ts
index bbe46c70147c..b9656f15c38e 100644
--- a/packages/ai/src/generate-text/generate-text.ts
+++ b/packages/ai/src/generate-text/generate-text.ts
@@ -115,9 +115,10 @@ type GenerateTextIncludeSettings = {
  */
 export type GenerateTextOnStartCallback<
   TOOLS extends ToolSet,
+  CONTEXT extends ExpandedContext<TOOLS>,
   OUTPUT extends Output = Output,
 > = (
-  event: OnStartEvent<TOOLS, OUTPUT, GenerateTextIncludeSettings>,
+  event: OnStartEvent<TOOLS, CONTEXT, OUTPUT, GenerateTextIncludeSettings>,
 ) => PromiseLike<void> | void;
 
 /**
@@ -370,7 +371,11 @@ export async function generateText<
      * Callback that is called when the generateText operation begins,
      * before any LLM calls are made.
      */
-    experimental_onStart?: GenerateTextOnStartCallback<NoInfer<TOOLS>, OUTPUT>;
+    experimental_onStart?: GenerateTextOnStartCallback<
+      NoInfer<TOOLS>,
+      CONTEXT,
+      OUTPUT
+    >;
 
     /**
      * Callback that is called when a step (LLM call) begins,
@@ -520,7 +525,7 @@ export async function generateText<
       onStart,
       globalTelemetry.onStart as
         | undefined
-        | GenerateTextOnStartCallback<TOOLS, OUTPUT>,
+        | GenerateTextOnStartCallback<TOOLS, CONTEXT, OUTPUT>,
     ],
   });
 
diff --git a/packages/ai/src/generate-text/run-tools-transformation.ts b/packages/ai/src/generate-text/run-tools-transformation.ts
index 05237710d4bb..2b0f2158cef8 100644
--- a/packages/ai/src/generate-text/run-tools-transformation.ts
+++ b/packages/ai/src/generate-text/run-tools-transformation.ts
@@ -24,7 +24,7 @@ import { TypedToolCall } from './tool-call';
 import { ToolCallRepairFunction } from './tool-call-repair-function';
 import { TypedToolError } from './tool-error';
 import { TypedToolResult } from './tool-result';
-import { ToolSet } from './tool-set';
+import { ExpandedContext, ToolSet } from './tool-set';
 
 export type SingleRequestTextStreamPart<TOOLS extends ToolSet> =
   // Text blocks:
@@ -108,7 +108,10 @@ export type SingleRequestTextStreamPart<TOOLS extends ToolSet> =
   | { type: 'error'; error: unknown }
   | { type: 'raw'; rawValue: unknown };
 
-export function runToolsTransformation<TOOLS extends ToolSet>({
+export function runToolsTransformation<
+  TOOLS extends ToolSet,
+  CONTEXT extends ExpandedContext<TOOLS>,
+>({
   tools,
   generatorStream,
   tracer,
@@ -132,7 +135,7 @@ export function runToolsTransformation<TOOLS extends ToolSet>({
   messages: ModelMessage[];
   abortSignal: AbortSignal | undefined;
   repairToolCall: ToolCallRepairFunction<TOOLS> | undefined;
-  experimental_context: unknown;
+  experimental_context: CONTEXT;
   generateId: IdGenerator;
   stepNumber?: number;
   model?: { provider: string; modelId: string };
diff --git a/packages/ai/src/telemetry/telemetry-integration.ts b/packages/ai/src/telemetry/telemetry-integration.ts
index fe4e710581f7..2ce023d7a46b 100644
--- a/packages/ai/src/telemetry/telemetry-integration.ts
+++ b/packages/ai/src/telemetry/telemetry-integration.ts
@@ -16,7 +16,7 @@ import { Listener } from '../util/notify';
  * Methods can be sync or return a PromiseLike.
  */
 export interface TelemetryIntegration {
-  onStart?: Listener<OnStartEvent<ToolSet, Output>>;
+  onStart?: Listener<OnStartEvent<ToolSet, Context, Output>>;
   onStepStart?: Listener<OnStepStartEvent<ToolSet, Context, Output>>;
   onToolCallStart?: Listener<OnToolCallStartEvent<ToolSet>>;
   onToolCallFinish?: Listener<OnToolCallFinishEvent<ToolSet>>;

From 4195a90e4b93ef7a26618beacea2249d832dd4f1 Mon Sep 17 00:00:00 2001
From: Lars Grammel <lars.grammel@gmail.com>
Date: Mon, 9 Mar 2026 14:38:19 +0100
Subject: [PATCH 20/41] x

---
 .../ai/src/generate-text/generate-text.ts     | 24 ++++++++++---------
 packages/ai/src/generate-text/prepare-step.ts | 16 ++++++-------
 2 files changed, 21 insertions(+), 19 deletions(-)

diff --git a/packages/ai/src/generate-text/generate-text.ts b/packages/ai/src/generate-text/generate-text.ts
index b9656f15c38e..f46c8c8eaac0 100644
--- a/packages/ai/src/generate-text/generate-text.ts
+++ b/packages/ai/src/generate-text/generate-text.ts
@@ -4,7 +4,6 @@ import {
   LanguageModelV4ToolCall,
 } from '@ai-sdk/provider';
 import {
-  Context,
   createIdGenerator,
   getErrorMessage,
   IdGenerator,
@@ -89,7 +88,7 @@ import { ToolCallRepairFunction } from './tool-call-repair-function';
 import { TypedToolError } from './tool-error';
 import { ToolOutput } from './tool-output';
 import { TypedToolResult } from './tool-result';
-import { ExpandedContext, InferToolSetContext, ToolSet } from './tool-set';
+import { ExpandedContext, ToolSet } from './tool-set';
 
 const originalGenerateId = createIdGenerator({
   prefix: 'aitxt',
@@ -355,12 +354,12 @@ export async function generateText<
     /**
      * @deprecated Use `prepareStep` instead.
      */
-    experimental_prepareStep?: PrepareStepFunction<CONTEXT, NoInfer<TOOLS>>;
+    experimental_prepareStep?: PrepareStepFunction<NoInfer<TOOLS>, CONTEXT>;
 
     /**
      * Optional function that you can use to provide different settings for a step.
      */
-    prepareStep?: PrepareStepFunction<CONTEXT, NoInfer<TOOLS>>;
+    prepareStep?: PrepareStepFunction<NoInfer<TOOLS>, CONTEXT>;
 
     /**
      * A function that attempts to repair a tool call that failed to parse.
@@ -373,8 +372,8 @@ export async function generateText<
      */
     experimental_onStart?: GenerateTextOnStartCallback<
       NoInfer<TOOLS>,
-      CONTEXT,
-      OUTPUT
+      NoInfer<CONTEXT>,
+      NoInfer<OUTPUT>
     >;
 
     /**
@@ -382,9 +381,9 @@ export async function generateText<
      * before the provider is called.
      */
     experimental_onStepStart?: GenerateTextOnStepStartCallback<
-      TOOLS,
-      CONTEXT,
-      OUTPUT
+      NoInfer<TOOLS>,
+      NoInfer<CONTEXT>,
+      NoInfer<OUTPUT>
     >;
 
     /**
@@ -404,12 +403,15 @@ export async function generateText<
     /**
      * Callback that is called when each step (LLM call) is finished, including intermediate steps.
      */
-    onStepFinish?: GenerateTextOnStepFinishCallback<TOOLS, CONTEXT>;
+    onStepFinish?: GenerateTextOnStepFinishCallback<
+      NoInfer<TOOLS>,
+      NoInfer<CONTEXT>
+    >;
 
     /**
      * Callback that is called when all steps are finished and the response is complete.
      */
-    onFinish?: GenerateTextOnFinishCallback<TOOLS, CONTEXT>;
+    onFinish?: GenerateTextOnFinishCallback<NoInfer<TOOLS>, NoInfer<CONTEXT>>;
 
     /**
      * Context that is passed into tool execution.
diff --git a/packages/ai/src/generate-text/prepare-step.ts b/packages/ai/src/generate-text/prepare-step.ts
index 3dcab84207c0..cf06e5e3c5e7 100644
--- a/packages/ai/src/generate-text/prepare-step.ts
+++ b/packages/ai/src/generate-text/prepare-step.ts
@@ -6,7 +6,7 @@ import {
 } from '@ai-sdk/provider-utils';
 import { LanguageModel, ToolChoice } from '../types/language-model';
 import { StepResult } from './step-result';
-import { ToolSet } from './tool-set';
+import { ExpandedContext, ToolSet } from './tool-set';
 
 /**
  * Function that you can use to provide different settings for a step.
@@ -22,13 +22,13 @@ import { ToolSet } from './tool-set';
  * If you return undefined (or for undefined settings), the settings from the outer level will be used.
  */
 export type PrepareStepFunction<
-  CONTEXT extends Context,
-  TOOLS extends ToolSet<CONTEXT> = ToolSet<CONTEXT>,
+  TOOLS extends ToolSet,
+  CONTEXT extends ExpandedContext<TOOLS>,
 > = (options: {
   /**
    * The steps that have been executed so far.
    */
-  steps: Array<StepResult<NoInfer<TOOLS>>>;
+  steps: Array<StepResult<TOOLS, CONTEXT>>;
 
   /**
    * The number of the step that is being executed.
@@ -50,16 +50,16 @@ export type PrepareStepFunction<
    */
   experimental_context: CONTEXT;
 }) =>
-  | PromiseLike<PrepareStepResult<CONTEXT, TOOLS>>
-  | PrepareStepResult<CONTEXT, TOOLS>;
+  | PromiseLike<PrepareStepResult<TOOLS, CONTEXT>>
+  | PrepareStepResult<TOOLS, CONTEXT>;
 
 /**
  * The result type returned by a {@link PrepareStepFunction},
  * allowing per-step overrides of model, tools, or messages.
  */
 export type PrepareStepResult<
-  CONTEXT extends Context,
-  TOOLS extends ToolSet<CONTEXT> = ToolSet<CONTEXT>,
+  TOOLS extends ToolSet,
+  CONTEXT extends ExpandedContext<TOOLS>,
 > =
   | {
       /**

From dfdbdc6aec47f0b2d3f6ee949756fada9f8c2ab0 Mon Sep 17 00:00:00 2001
From: Lars Grammel <lars.grammel@gmail.com>
Date: Mon, 9 Mar 2026 14:39:41 +0100
Subject: [PATCH 21/41] a

---
 packages/ai/src/generate-text/generate-text-result.ts | 5 +++--
 packages/ai/src/generate-text/generate-text.ts        | 4 ++--
 2 files changed, 5 insertions(+), 4 deletions(-)

diff --git a/packages/ai/src/generate-text/generate-text-result.ts b/packages/ai/src/generate-text/generate-text-result.ts
index 6f37f8d69eaf..81f0757400cd 100644
--- a/packages/ai/src/generate-text/generate-text-result.ts
+++ b/packages/ai/src/generate-text/generate-text-result.ts
@@ -16,7 +16,7 @@ import {
   StaticToolResult,
   TypedToolResult,
 } from './tool-result';
-import { ToolSet } from './tool-set';
+import { ExpandedContext, ToolSet } from './tool-set';
 
 /**
  * The result of a `generateText` call.
@@ -24,6 +24,7 @@ import { ToolSet } from './tool-set';
  */
 export interface GenerateTextResult<
   TOOLS extends ToolSet,
+  CONTEXT extends ExpandedContext<TOOLS>,
   OUTPUT extends Output,
 > {
   /**
@@ -151,7 +152,7 @@ export interface GenerateTextResult<
    * You can use this to get information about intermediate steps,
    * such as the tool calls or the response headers.
    */
-  readonly steps: Array<StepResult<TOOLS>>;
+  readonly steps: Array<StepResult<TOOLS, CONTEXT>>;
 
   /**
    * The generated structured output. It uses the `output` specification.
diff --git a/packages/ai/src/generate-text/generate-text.ts b/packages/ai/src/generate-text/generate-text.ts
index f46c8c8eaac0..124d422060c3 100644
--- a/packages/ai/src/generate-text/generate-text.ts
+++ b/packages/ai/src/generate-text/generate-text.ts
@@ -450,7 +450,7 @@ export async function generateText<
     _internal?: {
       generateId?: IdGenerator;
     };
-  }): Promise<GenerateTextResult<TOOLS, OUTPUT>> {
+  }): Promise<GenerateTextResult<TOOLS, CONTEXT, OUTPUT>> {
   const model = resolveLanguageModel(modelArg);
   const createGlobalTelemetry = getGlobalTelemetryIntegration<any, OUTPUT>();
   const stopConditions = asArray(stopWhen);
@@ -1332,7 +1332,7 @@ class DefaultGenerateTextResult<
   TOOLS extends ToolSet,
   CONTEXT extends ExpandedContext<TOOLS>,
   OUTPUT extends Output,
-> implements GenerateTextResult<TOOLS, OUTPUT>
+> implements GenerateTextResult<TOOLS, CONTEXT, OUTPUT>
 {
   readonly steps: StepResult<TOOLS, CONTEXT>[];
   readonly totalUsage: LanguageModelUsage;

From 342136df118efe8b7b1c5d58cdedaf2f4d701199 Mon Sep 17 00:00:00 2001
From: Lars Grammel <lars.grammel@gmail.com>
Date: Mon, 9 Mar 2026 14:53:40 +0100
Subject: [PATCH 22/41] a

---
 .../openai/tool-call-with-context.ts           | 10 +++++++++-
 packages/ai/src/generate-text/tool-set.ts      | 18 ++++++++++--------
 2 files changed, 19 insertions(+), 9 deletions(-)

diff --git a/examples/ai-functions/src/generate-text/openai/tool-call-with-context.ts b/examples/ai-functions/src/generate-text/openai/tool-call-with-context.ts
index cae53ba9f0e1..4436ffce9098 100644
--- a/examples/ai-functions/src/generate-text/openai/tool-call-with-context.ts
+++ b/examples/ai-functions/src/generate-text/openai/tool-call-with-context.ts
@@ -47,7 +47,15 @@ run(async () => {
         },
       }),
     },
-    experimental_context: { weatherApiKey: '123', calculatorApiKey: '456' },
+    experimental_context: {
+      weatherApiKey: '123',
+      calculatorApiKey: '456',
+      somethingElse: 'context',
+    },
+    prepareStep: async ({ experimental_context: context }) => {
+      console.log(context);
+      return {};
+    },
     prompt: 'What is the weather in San Francisco?',
   });
 
diff --git a/packages/ai/src/generate-text/tool-set.ts b/packages/ai/src/generate-text/tool-set.ts
index e6b039a9ddb5..104c96b30162 100644
--- a/packages/ai/src/generate-text/tool-set.ts
+++ b/packages/ai/src/generate-text/tool-set.ts
@@ -4,13 +4,13 @@ import { UnionToIntersection } from '../util/union-to-intersection';
 export type ToolSet = Record<
   string,
   (
-    | Tool<never, never, Context>
-    | Tool<any, any, Context>
-    | Tool<any, never, Context>
-    | Tool<never, any, Context>
+    | Tool<never, never, any>
+    | Tool<any, any, any>
+    | Tool<any, never, any>
+    | Tool<never, any, any>
   ) &
     Pick<
-      Tool<any, any, Context>,
+      Tool<any, any, any>,
       | 'execute'
       | 'onInputAvailable'
       | 'onInputStart'
@@ -21,9 +21,11 @@ export type ToolSet = Record<
 
 export type InferToolSetContext<TOOLS extends ToolSet> = UnionToIntersection<
   {
-    [K in keyof TOOLS]: InferToolContext<TOOLS[K]>;
+    [K in keyof TOOLS]: InferToolContext<NoInfer<TOOLS[K]>>;
   }[keyof TOOLS]
 >;
 
-export type ExpandedContext<TOOLS extends ToolSet> =
-  InferToolSetContext<TOOLS> & Context;
+export type ExpandedContext<TOOLS extends ToolSet> = InferToolSetContext<
+  NoInfer<TOOLS>
+> &
+  Context;

From 9921f52a89e794714ca6d7969edd2ba5cc56e47a Mon Sep 17 00:00:00 2001
From: Lars Grammel <lars.grammel@gmail.com>
Date: Mon, 9 Mar 2026 15:52:03 +0100
Subject: [PATCH 23/41] s

---
 packages/ai/src/generate-text/stream-text.ts | 158 ++++++++++++-------
 1 file changed, 103 insertions(+), 55 deletions(-)

diff --git a/packages/ai/src/generate-text/stream-text.ts b/packages/ai/src/generate-text/stream-text.ts
index 939d30df27c8..cd8138844c77 100644
--- a/packages/ai/src/generate-text/stream-text.ts
+++ b/packages/ai/src/generate-text/stream-text.ts
@@ -122,7 +122,7 @@ import { TypedToolCall } from './tool-call';
 import { ToolCallRepairFunction } from './tool-call-repair-function';
 import { ToolOutput } from './tool-output';
 import { StaticToolOutputDenied } from './tool-output-denied';
-import { ToolSet } from './tool-set';
+import { ExpandedContext, ToolSet } from './tool-set';
 
 const originalGenerateId = createIdGenerator({
   prefix: 'aitxt',
@@ -154,9 +154,10 @@ export type StreamTextOnErrorCallback = (event: {
  *
  * @param stepResult - The result of the step.
  */
-export type StreamTextOnStepFinishCallback<TOOLS extends ToolSet> = (
-  event: OnStepFinishEvent<TOOLS>,
-) => PromiseLike<void> | void;
+export type StreamTextOnStepFinishCallback<
+  TOOLS extends ToolSet,
+  CONTEXT extends ExpandedContext<TOOLS>,
+> = (event: OnStepFinishEvent<TOOLS, CONTEXT>) => PromiseLike<void> | void;
 
 /**
  * Callback that is set using the `onChunk` option.
@@ -185,20 +186,24 @@ export type StreamTextOnChunkCallback<TOOLS extends ToolSet> = (event: {
  *
  * @param event - The event that is passed to the callback.
  */
-export type StreamTextOnFinishCallback<TOOLS extends ToolSet> = (
-  event: OnFinishEvent<TOOLS>,
-) => PromiseLike<void> | void;
+export type StreamTextOnFinishCallback<
+  TOOLS extends ToolSet,
+  CONTEXT extends ExpandedContext<TOOLS>,
+> = (event: OnFinishEvent<TOOLS, CONTEXT>) => PromiseLike<void> | void;
 
 /**
  * Callback that is set using the `onAbort` option.
  *
  * @param event - The event that is passed to the callback.
  */
-export type StreamTextOnAbortCallback<TOOLS extends ToolSet> = (event: {
+export type StreamTextOnAbortCallback<
+  TOOLS extends ToolSet,
+  CONTEXT extends ExpandedContext<TOOLS>,
+> = (event: {
   /**
    * Details for all previously finished steps.
    */
-  readonly steps: StepResult<TOOLS>[];
+  readonly steps: StepResult<TOOLS, CONTEXT>[];
 }) => PromiseLike<void> | void;
 
 /**
@@ -217,9 +222,10 @@ type StreamTextIncludeSettings = { requestBody?: boolean };
  */
 export type StreamTextOnStartCallback<
   TOOLS extends ToolSet = ToolSet,
+  CONTEXT extends ExpandedContext<TOOLS> = ExpandedContext<TOOLS>,
   OUTPUT extends Output = Output,
 > = (
-  event: OnStartEvent<TOOLS, OUTPUT, StreamTextIncludeSettings>,
+  event: OnStartEvent<TOOLS, CONTEXT, OUTPUT, StreamTextIncludeSettings>,
 ) => PromiseLike<void> | void;
 
 /**
@@ -233,9 +239,10 @@ export type StreamTextOnStartCallback<
  */
 export type StreamTextOnStepStartCallback<
   TOOLS extends ToolSet = ToolSet,
+  CONTEXT extends ExpandedContext<TOOLS> = ExpandedContext<TOOLS>,
   OUTPUT extends Output = Output,
 > = (
-  event: OnStepStartEvent<TOOLS, OUTPUT, StreamTextIncludeSettings>,
+  event: OnStepStartEvent<TOOLS, CONTEXT, OUTPUT, StreamTextIncludeSettings>,
 ) => PromiseLike<void> | void;
 
 export type StreamTextOnToolCallStartCallback<TOOLS extends ToolSet = ToolSet> =
@@ -293,6 +300,7 @@ export type StreamTextOnToolCallFinishCallback<
  */
 export function streamText<
   TOOLS extends ToolSet,
+  CONTEXT extends ExpandedContext<TOOLS> = ExpandedContext<TOOLS>,
   OUTPUT extends Output = Output<string, string, never>,
 >({
   model,
@@ -328,7 +336,7 @@ export function streamText<
   experimental_onStepStart: onStepStart,
   experimental_onToolCallStart: onToolCallStart,
   experimental_onToolCallFinish: onToolCallFinish,
-  experimental_context,
+  experimental_context = {} as CONTEXT,
   experimental_include: include,
   _internal: { now = originalNow, generateId = originalGenerateId } = {},
   ...settings
@@ -356,8 +364,8 @@ export function streamText<
      * @default stepCountIs(1)
      */
     stopWhen?:
-      | StopCondition<NoInfer<TOOLS>>
-      | Array<StopCondition<NoInfer<TOOLS>>>;
+      | StopCondition<NoInfer<TOOLS>, CONTEXT>
+      | Array<StopCondition<NoInfer<TOOLS>, CONTEXT>>;
 
     /**
      * Optional telemetry configuration (experimental).
@@ -405,7 +413,7 @@ export function streamText<
      * @returns An object that contains the settings for the step.
      * If you return undefined (or for undefined settings), the settings from the outer level will be used.
      */
-    prepareStep?: PrepareStepFunction<NoInfer<TOOLS>>;
+    prepareStep?: PrepareStepFunction<NoInfer<TOOLS>, CONTEXT>;
 
     /**
      * A function that attempts to repair a tool call that failed to parse.
@@ -455,20 +463,27 @@ export function streamText<
      *
      * The usage is the combined usage of all steps.
      */
-    onFinish?: StreamTextOnFinishCallback<TOOLS>;
+    onFinish?: StreamTextOnFinishCallback<NoInfer<TOOLS>, NoInfer<CONTEXT>>;
 
-    onAbort?: StreamTextOnAbortCallback<TOOLS>;
+    onAbort?: StreamTextOnAbortCallback<NoInfer<TOOLS>, NoInfer<CONTEXT>>;
 
     /**
      * Callback that is called when each step (LLM call) is finished, including intermediate steps.
      */
-    onStepFinish?: StreamTextOnStepFinishCallback<TOOLS>;
+    onStepFinish?: StreamTextOnStepFinishCallback<
+      NoInfer<TOOLS>,
+      NoInfer<CONTEXT>
+    >;
 
     /**
      * Callback that is called when the streamText operation begins,
      * before any LLM calls are made.
      */
-    experimental_onStart?: StreamTextOnStartCallback<NoInfer<TOOLS>, OUTPUT>;
+    experimental_onStart?: StreamTextOnStartCallback<
+      NoInfer<TOOLS>,
+      NoInfer<CONTEXT>,
+      NoInfer<OUTPUT>
+    >;
 
     /**
      * Callback that is called when a step (LLM call) begins,
@@ -476,7 +491,8 @@ export function streamText<
      */
     experimental_onStepStart?: StreamTextOnStepStartCallback<
       NoInfer<TOOLS>,
-      OUTPUT
+      NoInfer<CONTEXT>,
+      NoInfer<OUTPUT>
     >;
 
     /**
@@ -500,7 +516,7 @@ export function streamText<
      *
      * @default undefined
      */
-    experimental_context?: unknown;
+    experimental_context?: CONTEXT;
 
     /**
      * Settings for controlling what data is included in step results.
@@ -533,7 +549,7 @@ export function streamText<
     stepTimeoutMs != null ? new AbortController() : undefined;
   const chunkAbortController =
     chunkTimeoutMs != null ? new AbortController() : undefined;
-  return new DefaultStreamTextResult<TOOLS, OUTPUT>({
+  return new DefaultStreamTextResult<TOOLS, CONTEXT, OUTPUT>({
     model: resolveLanguageModel(model),
     telemetry,
     headers,
@@ -684,8 +700,11 @@ function createOutputTransformStream<
   });
 }
 
-class DefaultStreamTextResult<TOOLS extends ToolSet, OUTPUT extends Output>
-  implements StreamTextResult<TOOLS, OUTPUT>
+class DefaultStreamTextResult<
+  TOOLS extends ToolSet,
+  CONTEXT extends ExpandedContext<TOOLS>,
+  OUTPUT extends Output,
+> implements StreamTextResult<TOOLS, OUTPUT>
 {
   private readonly _totalUsage = new DelayedPromise<
     Awaited<StreamTextResult<TOOLS, OUTPUT>['usage']>
@@ -776,31 +795,51 @@ class DefaultStreamTextResult<TOOLS extends ToolSet, OUTPUT extends Output>
     transforms: Array<StreamTextTransform<TOOLS>>;
     activeTools: Array<keyof TOOLS> | undefined;
     repairToolCall: ToolCallRepairFunction<TOOLS> | undefined;
-    stopConditions: Array<StopCondition<NoInfer<TOOLS>>>;
+    stopConditions: Array<StopCondition<NoInfer<TOOLS>, NoInfer<CONTEXT>>>;
     output: OUTPUT | undefined;
     providerOptions: ProviderOptions | undefined;
-    prepareStep: PrepareStepFunction<NoInfer<TOOLS>> | undefined;
+    prepareStep:
+      | PrepareStepFunction<NoInfer<TOOLS>, NoInfer<CONTEXT>>
+      | undefined;
     includeRawChunks: boolean;
     now: () => number;
     generateId: () => string;
     timeout: TimeoutConfiguration | undefined;
     stopWhen:
-      | StopCondition<NoInfer<TOOLS>>
-      | Array<StopCondition<NoInfer<TOOLS>>>
+      | StopCondition<NoInfer<TOOLS>, NoInfer<CONTEXT>>
+      | Array<StopCondition<NoInfer<TOOLS>, NoInfer<CONTEXT>>>
       | undefined;
     originalAbortSignal: AbortSignal | undefined;
-    experimental_context: unknown;
+    experimental_context: CONTEXT;
     download: DownloadFunction | undefined;
     include: { requestBody?: boolean } | undefined;
 
     // callbacks:
     onChunk: undefined | StreamTextOnChunkCallback<TOOLS>;
     onError: StreamTextOnErrorCallback;
-    onFinish: undefined | StreamTextOnFinishCallback<TOOLS>;
-    onAbort: undefined | StreamTextOnAbortCallback<TOOLS>;
-    onStepFinish: undefined | StreamTextOnStepFinishCallback<TOOLS>;
-    onStart: undefined | StreamTextOnStartCallback<TOOLS, OUTPUT>;
-    onStepStart: undefined | StreamTextOnStepStartCallback<TOOLS, OUTPUT>;
+    onFinish:
+      | undefined
+      | StreamTextOnFinishCallback<NoInfer<TOOLS>, NoInfer<CONTEXT>>;
+    onAbort:
+      | undefined
+      | StreamTextOnAbortCallback<NoInfer<TOOLS>, NoInfer<CONTEXT>>;
+    onStepFinish:
+      | undefined
+      | StreamTextOnStepFinishCallback<NoInfer<TOOLS>, NoInfer<CONTEXT>>;
+    onStart:
+      | undefined
+      | StreamTextOnStartCallback<
+          NoInfer<TOOLS>,
+          NoInfer<CONTEXT>,
+          NoInfer<OUTPUT>
+        >;
+    onStepStart:
+      | undefined
+      | StreamTextOnStepStartCallback<
+          NoInfer<TOOLS>,
+          NoInfer<CONTEXT>,
+          NoInfer<OUTPUT>
+        >;
     onToolCallStart: undefined | StreamTextOnToolCallStartCallback<TOOLS>;
     onToolCallFinish: undefined | StreamTextOnToolCallFinishCallback<TOOLS>;
   }) {
@@ -826,7 +865,7 @@ class DefaultStreamTextResult<TOOLS extends ToolSet, OUTPUT extends Output>
     let recordedTotalUsage: LanguageModelUsage | undefined = undefined;
     let recordedRequest: LanguageModelRequestMetadata = {};
     let recordedWarnings: Array<CallWarning> = [];
-    const recordedSteps: StepResult<TOOLS>[] = [];
+    const recordedSteps: StepResult<TOOLS, CONTEXT>[] = [];
 
     // Track provider-executed tool calls that support deferred results
     // (e.g., code_execution in programmatic tool calling scenarios).
@@ -1024,23 +1063,24 @@ class DefaultStreamTextResult<TOOLS extends ToolSet, OUTPUT extends Output>
           });
 
           // Add step information (after response messages are updated):
-          const currentStepResult: StepResult<TOOLS> = new DefaultStepResult({
-            stepNumber: recordedSteps.length,
-            model: modelInfo,
-            ...callbackTelemetryProps,
-            experimental_context,
-            content: recordedContent,
-            finishReason: part.finishReason,
-            rawFinishReason: part.rawFinishReason,
-            usage: part.usage,
-            warnings: recordedWarnings,
-            request: recordedRequest,
-            response: {
-              ...part.response,
-              messages: [...recordedResponseMessages, ...stepMessages],
-            },
-            providerMetadata: part.providerMetadata,
-          });
+          const currentStepResult: StepResult<TOOLS, CONTEXT> =
+            new DefaultStepResult({
+              stepNumber: recordedSteps.length,
+              model: modelInfo,
+              ...callbackTelemetryProps,
+              experimental_context,
+              content: recordedContent,
+              finishReason: part.finishReason,
+              rawFinishReason: part.rawFinishReason,
+              usage: part.usage,
+              warnings: recordedWarnings,
+              request: recordedRequest,
+              response: {
+                ...part.response,
+                messages: [...recordedResponseMessages, ...stepMessages],
+              },
+              providerMetadata: part.providerMetadata,
+            });
 
           await notify({
             event: currentStepResult,
@@ -1135,7 +1175,7 @@ class DefaultStreamTextResult<TOOLS extends ToolSet, OUTPUT extends Output>
               onFinish,
               globalTelemetry.onFinish as
                 | undefined
-                | StreamTextOnFinishCallback<TOOLS>,
+                | StreamTextOnFinishCallback<NoInfer<TOOLS>, NoInfer<CONTEXT>>,
             ],
           });
 
@@ -1340,7 +1380,11 @@ class DefaultStreamTextResult<TOOLS extends ToolSet, OUTPUT extends Output>
             onStart,
             globalTelemetry.onStart as
               | undefined
-              | StreamTextOnStartCallback<TOOLS, OUTPUT>,
+              | StreamTextOnStartCallback<
+                  NoInfer<TOOLS>,
+                  NoInfer<CONTEXT>,
+                  NoInfer<OUTPUT>
+                >,
           ],
         });
 
@@ -1622,7 +1666,11 @@ class DefaultStreamTextResult<TOOLS extends ToolSet, OUTPUT extends Output>
                 onStepStart,
                 globalTelemetry.onStepStart as
                   | undefined
-                  | StreamTextOnStepStartCallback<TOOLS, OUTPUT>,
+                  | StreamTextOnStepStartCallback<
+                      NoInfer<TOOLS>,
+                      NoInfer<CONTEXT>,
+                      NoInfer<OUTPUT>
+                    >,
               ],
             });
 

From d0e70cb380732bf18b24b81dd42216448fa397a5 Mon Sep 17 00:00:00 2001
From: Lars Grammel <lars.grammel@gmail.com>
Date: Mon, 9 Mar 2026 16:00:15 +0100
Subject: [PATCH 24/41] st

---
 packages/ai/src/agent/agent.ts                | 17 ++++---
 .../agent/create-agent-ui-stream-response.ts  |  7 +--
 .../ai/src/agent/create-agent-ui-stream.ts    |  5 +-
 .../agent/pipe-agent-ui-stream-to-response.ts |  5 +-
 .../ai/src/agent/tool-loop-agent-settings.ts  | 48 ++++++++++-------
 packages/ai/src/agent/tool-loop-agent.ts      | 51 ++++++++++++++-----
 .../src/generate-text/stream-text-result.ts   |  5 +-
 packages/ai/src/generate-text/stream-text.ts  | 12 ++---
 packages/ai/src/ui/direct-chat-transport.ts   | 22 +++++---
 9 files changed, 111 insertions(+), 61 deletions(-)

diff --git a/packages/ai/src/agent/agent.ts b/packages/ai/src/agent/agent.ts
index 2dbb9ba1c198..c8ce19c941a4 100644
--- a/packages/ai/src/agent/agent.ts
+++ b/packages/ai/src/agent/agent.ts
@@ -3,7 +3,7 @@ import { GenerateTextResult } from '../generate-text/generate-text-result';
 import { Output } from '../generate-text/output';
 import { StreamTextTransform } from '../generate-text/stream-text';
 import { StreamTextResult } from '../generate-text/stream-text-result';
-import { ToolSet } from '../generate-text/tool-set';
+import { ExpandedContext, ToolSet } from '../generate-text/tool-set';
 import { TimeoutConfiguration } from '../prompt/call-settings';
 import type {
   ToolLoopAgentOnFinishCallback,
@@ -17,9 +17,11 @@ import type {
 /**
  * Parameters for calling an agent.
  */
-export type AgentCallParameters<CALL_OPTIONS, TOOLS extends ToolSet = {}> = ([
+export type AgentCallParameters<
   CALL_OPTIONS,
-] extends [never]
+  TOOLS extends ToolSet = {},
+  CONTEXT extends ExpandedContext<TOOLS> = ExpandedContext<TOOLS>,
+> = ([CALL_OPTIONS] extends [never]
   ? { options?: never }
   : { options: CALL_OPTIONS }) &
   (
@@ -67,12 +69,12 @@ export type AgentCallParameters<CALL_OPTIONS, TOOLS extends ToolSet = {}> = ([
     /**
      * Callback that is called when the agent operation begins, before any LLM calls.
      */
-    experimental_onStart?: ToolLoopAgentOnStartCallback<TOOLS>;
+    experimental_onStart?: ToolLoopAgentOnStartCallback<TOOLS, CONTEXT>;
 
     /**
      * Callback that is called when a step (LLM call) begins, before the provider is called.
      */
-    experimental_onStepStart?: ToolLoopAgentOnStepStartCallback<TOOLS>;
+    experimental_onStepStart?: ToolLoopAgentOnStepStartCallback<TOOLS, CONTEXT>;
 
     /**
      * Callback that is called before each tool execution begins.
@@ -122,6 +124,7 @@ export type AgentStreamParameters<
 export interface Agent<
   CALL_OPTIONS = never,
   TOOLS extends ToolSet = {},
+  CONTEXT extends ExpandedContext<TOOLS> = ExpandedContext<TOOLS>,
   OUTPUT extends Output = never,
 > {
   /**
@@ -145,12 +148,12 @@ export interface Agent<
    */
   generate(
     options: AgentCallParameters<CALL_OPTIONS, TOOLS>,
-  ): PromiseLike<GenerateTextResult<TOOLS, OUTPUT>>;
+  ): PromiseLike<GenerateTextResult<TOOLS, CONTEXT, OUTPUT>>;
 
   /**
    * Streams an output from the agent (streaming).
    */
   stream(
     options: AgentStreamParameters<CALL_OPTIONS, TOOLS>,
-  ): PromiseLike<StreamTextResult<TOOLS, OUTPUT>>;
+  ): PromiseLike<StreamTextResult<TOOLS, CONTEXT, OUTPUT>>;
 }
diff --git a/packages/ai/src/agent/create-agent-ui-stream-response.ts b/packages/ai/src/agent/create-agent-ui-stream-response.ts
index a8f6646895c1..d47e3db798cd 100644
--- a/packages/ai/src/agent/create-agent-ui-stream-response.ts
+++ b/packages/ai/src/agent/create-agent-ui-stream-response.ts
@@ -1,6 +1,6 @@
 import { StreamTextTransform, UIMessageStreamOptions } from '../generate-text';
 import { Output } from '../generate-text/output';
-import { ToolSet } from '../generate-text/tool-set';
+import { ExpandedContext, ToolSet } from '../generate-text/tool-set';
 import { TimeoutConfiguration } from '../prompt/call-settings';
 import { createUIMessageStreamResponse } from '../ui-message-stream';
 import { UIMessageStreamResponseInit } from '../ui-message-stream/ui-message-stream-response-init';
@@ -29,7 +29,8 @@ import type { ToolLoopAgentOnStepFinishCallback } from './tool-loop-agent-settin
 export async function createAgentUIStreamResponse<
   CALL_OPTIONS = never,
   TOOLS extends ToolSet = {},
-  OUTPUT extends Output = never,
+  CONTEXT extends ExpandedContext<TOOLS> = ExpandedContext<TOOLS>,
+  OUTPUT extends Output<TOOLS, CONTEXT> = never,
   MESSAGE_METADATA = unknown,
 >({
   headers,
@@ -38,7 +39,7 @@ export async function createAgentUIStreamResponse<
   consumeSseStream,
   ...options
 }: {
-  agent: Agent<CALL_OPTIONS, TOOLS, OUTPUT>;
+  agent: Agent<CALL_OPTIONS, TOOLS, CONTEXT, OUTPUT>;
   uiMessages: unknown[];
   abortSignal?: AbortSignal;
   timeout?: TimeoutConfiguration;
diff --git a/packages/ai/src/agent/create-agent-ui-stream.ts b/packages/ai/src/agent/create-agent-ui-stream.ts
index cc8da02822d5..5460b21a94f0 100644
--- a/packages/ai/src/agent/create-agent-ui-stream.ts
+++ b/packages/ai/src/agent/create-agent-ui-stream.ts
@@ -1,6 +1,6 @@
 import { StreamTextTransform, UIMessageStreamOptions } from '../generate-text';
 import { Output } from '../generate-text/output';
-import { ToolSet } from '../generate-text/tool-set';
+import { ExpandedContext, ToolSet } from '../generate-text/tool-set';
 import { TimeoutConfiguration } from '../prompt/call-settings';
 import { InferUIMessageChunk } from '../ui-message-stream';
 import { convertToModelMessages } from '../ui/convert-to-model-messages';
@@ -26,6 +26,7 @@ import type { ToolLoopAgentOnStepFinishCallback } from './tool-loop-agent-settin
 export async function createAgentUIStream<
   CALL_OPTIONS = never,
   TOOLS extends ToolSet = {},
+  CONTEXT extends ExpandedContext<TOOLS> = ExpandedContext<TOOLS>,
   OUTPUT extends Output = never,
   MESSAGE_METADATA = unknown,
 >({
@@ -38,7 +39,7 @@ export async function createAgentUIStream<
   onStepFinish,
   ...uiMessageStreamOptions
 }: {
-  agent: Agent<CALL_OPTIONS, TOOLS, OUTPUT>;
+  agent: Agent<CALL_OPTIONS, TOOLS, CONTEXT, OUTPUT>;
   uiMessages: unknown[];
   abortSignal?: AbortSignal;
   timeout?: TimeoutConfiguration;
diff --git a/packages/ai/src/agent/pipe-agent-ui-stream-to-response.ts b/packages/ai/src/agent/pipe-agent-ui-stream-to-response.ts
index 048921f4bd78..8ed3f80132dd 100644
--- a/packages/ai/src/agent/pipe-agent-ui-stream-to-response.ts
+++ b/packages/ai/src/agent/pipe-agent-ui-stream-to-response.ts
@@ -1,7 +1,7 @@
 import { ServerResponse } from 'node:http';
 import { StreamTextTransform, UIMessageStreamOptions } from '../generate-text';
 import { Output } from '../generate-text/output';
-import { ToolSet } from '../generate-text/tool-set';
+import { ExpandedContext, ToolSet } from '../generate-text/tool-set';
 import { TimeoutConfiguration } from '../prompt/call-settings';
 import { pipeUIMessageStreamToResponse } from '../ui-message-stream';
 import { UIMessageStreamResponseInit } from '../ui-message-stream/ui-message-stream-response-init';
@@ -29,6 +29,7 @@ import type { ToolLoopAgentOnStepFinishCallback } from './tool-loop-agent-settin
 export async function pipeAgentUIStreamToResponse<
   CALL_OPTIONS = never,
   TOOLS extends ToolSet = {},
+  CONTEXT extends ExpandedContext<TOOLS> = ExpandedContext<TOOLS>,
   OUTPUT extends Output = never,
   MESSAGE_METADATA = unknown,
 >({
@@ -40,7 +41,7 @@ export async function pipeAgentUIStreamToResponse<
   ...options
 }: {
   response: ServerResponse;
-  agent: Agent<CALL_OPTIONS, TOOLS, OUTPUT>;
+  agent: Agent<CALL_OPTIONS, TOOLS, CONTEXT, OUTPUT>;
   uiMessages: unknown[];
   abortSignal?: AbortSignal;
   timeout?: TimeoutConfiguration;
diff --git a/packages/ai/src/agent/tool-loop-agent-settings.ts b/packages/ai/src/agent/tool-loop-agent-settings.ts
index 8c0e7dd50643..025e81f80ca3 100644
--- a/packages/ai/src/agent/tool-loop-agent-settings.ts
+++ b/packages/ai/src/agent/tool-loop-agent-settings.ts
@@ -16,7 +16,7 @@ import { Output } from '../generate-text/output';
 import { PrepareStepFunction } from '../generate-text/prepare-step';
 import { StopCondition } from '../generate-text/stop-condition';
 import { ToolCallRepairFunction } from '../generate-text/tool-call-repair-function';
-import { ToolSet } from '../generate-text/tool-set';
+import { ExpandedContext, ToolSet } from '../generate-text/tool-set';
 import { CallSettings } from '../prompt/call-settings';
 import { Prompt } from '../prompt/prompt';
 import { TelemetrySettings } from '../telemetry/telemetry-settings';
@@ -26,13 +26,17 @@ import { AgentCallParameters } from './agent';
 
 export type ToolLoopAgentOnStartCallback<
   TOOLS extends ToolSet = ToolSet,
+  CONTEXT extends ExpandedContext<TOOLS> = ExpandedContext<TOOLS>,
   OUTPUT extends Output = Output,
-> = (event: OnStartEvent<TOOLS, OUTPUT>) => PromiseLike<void> | void;
+> = (event: OnStartEvent<TOOLS, CONTEXT, OUTPUT>) => PromiseLike<void> | void;
 
 export type ToolLoopAgentOnStepStartCallback<
   TOOLS extends ToolSet = ToolSet,
+  CONTEXT extends ExpandedContext<TOOLS> = ExpandedContext<TOOLS>,
   OUTPUT extends Output = Output,
-> = (event: OnStepStartEvent<TOOLS, OUTPUT>) => PromiseLike<void> | void;
+> = (
+  event: OnStepStartEvent<TOOLS, CONTEXT, OUTPUT>,
+) => PromiseLike<void> | void;
 
 export type ToolLoopAgentOnToolCallStartCallback<
   TOOLS extends ToolSet = ToolSet,
@@ -42,13 +46,15 @@ export type ToolLoopAgentOnToolCallFinishCallback<
   TOOLS extends ToolSet = ToolSet,
 > = (event: OnToolCallFinishEvent<TOOLS>) => PromiseLike<void> | void;
 
-export type ToolLoopAgentOnStepFinishCallback<TOOLS extends ToolSet = {}> = (
-  stepResult: OnStepFinishEvent<TOOLS>,
-) => Promise<void> | void;
+export type ToolLoopAgentOnStepFinishCallback<
+  TOOLS extends ToolSet = ToolSet,
+  CONTEXT extends ExpandedContext<TOOLS> = ExpandedContext<TOOLS>,
+> = (stepResult: OnStepFinishEvent<TOOLS, CONTEXT>) => Promise<void> | void;
 
-export type ToolLoopAgentOnFinishCallback<TOOLS extends ToolSet = {}> = (
-  event: OnFinishEvent<TOOLS>,
-) => PromiseLike<void> | void;
+export type ToolLoopAgentOnFinishCallback<
+  TOOLS extends ToolSet = ToolSet,
+  CONTEXT extends ExpandedContext<TOOLS> = ExpandedContext<TOOLS>,
+> = (event: OnFinishEvent<TOOLS, CONTEXT>) => PromiseLike<void> | void;
 
 /**
  * Configuration options for an agent.
@@ -56,6 +62,7 @@ export type ToolLoopAgentOnFinishCallback<TOOLS extends ToolSet = {}> = (
 export type ToolLoopAgentSettings<
   CALL_OPTIONS = never,
   TOOLS extends ToolSet = {},
+  CONTEXT extends ExpandedContext<TOOLS> = ExpandedContext<TOOLS>,
   OUTPUT extends Output = never,
 > = Omit<CallSettings, 'abortSignal'> & {
   /**
@@ -92,8 +99,8 @@ export type ToolLoopAgentSettings<
    * @default stepCountIs(20)
    */
   stopWhen?:
-    | StopCondition<NoInfer<TOOLS>>
-    | Array<StopCondition<NoInfer<TOOLS>>>;
+    | StopCondition<NoInfer<TOOLS>, CONTEXT>
+    | Array<StopCondition<NoInfer<TOOLS>, CONTEXT>>;
 
   /**
    * Optional telemetry configuration (experimental).
@@ -114,7 +121,7 @@ export type ToolLoopAgentSettings<
   /**
    * Optional function that you can use to provide different settings for a step.
    */
-  prepareStep?: PrepareStepFunction<NoInfer<TOOLS>>;
+  prepareStep?: PrepareStepFunction<NoInfer<TOOLS>, CONTEXT>;
 
   /**
    * A function that attempts to repair a tool call that failed to parse.
@@ -124,14 +131,19 @@ export type ToolLoopAgentSettings<
   /**
    * Callback that is called when the agent operation begins, before any LLM calls.
    */
-  experimental_onStart?: ToolLoopAgentOnStartCallback<NoInfer<TOOLS>, OUTPUT>;
+  experimental_onStart?: ToolLoopAgentOnStartCallback<
+    NoInfer<TOOLS>,
+    CONTEXT,
+    NoInfer<OUTPUT>
+  >;
 
   /**
    * Callback that is called when a step (LLM call) begins, before the provider is called.
    */
   experimental_onStepStart?: ToolLoopAgentOnStepStartCallback<
     NoInfer<TOOLS>,
-    OUTPUT
+    NoInfer<CONTEXT>,
+    NoInfer<OUTPUT>
   >;
 
   /**
@@ -169,10 +181,8 @@ export type ToolLoopAgentSettings<
    * Context that is passed into tool calls.
    *
    * Experimental (can break in patch releases).
-   *
-   * @default undefined
    */
-  experimental_context?: unknown;
+  experimental_context?: CONTEXT;
 
   /**
    * Custom download function to use for URLs.
@@ -197,7 +207,7 @@ export type ToolLoopAgentSettings<
       'onStepFinish'
     > &
       Pick<
-        ToolLoopAgentSettings<CALL_OPTIONS, TOOLS, OUTPUT>,
+        ToolLoopAgentSettings<CALL_OPTIONS, TOOLS, CONTEXT, NoInfer<OUTPUT>>,
         | 'model'
         | 'tools'
         | 'maxOutputTokens'
@@ -219,7 +229,7 @@ export type ToolLoopAgentSettings<
       >,
   ) => MaybePromiseLike<
     Pick<
-      ToolLoopAgentSettings<CALL_OPTIONS, TOOLS, OUTPUT>,
+      ToolLoopAgentSettings<CALL_OPTIONS, TOOLS, CONTEXT, NoInfer<OUTPUT>>,
       | 'model'
       | 'tools'
       | 'maxOutputTokens'
diff --git a/packages/ai/src/agent/tool-loop-agent.ts b/packages/ai/src/agent/tool-loop-agent.ts
index 70206687a7b6..acc3bc97333f 100644
--- a/packages/ai/src/agent/tool-loop-agent.ts
+++ b/packages/ai/src/agent/tool-loop-agent.ts
@@ -4,10 +4,14 @@ import { Output } from '../generate-text/output';
 import { stepCountIs } from '../generate-text/stop-condition';
 import { streamText } from '../generate-text/stream-text';
 import { StreamTextResult } from '../generate-text/stream-text-result';
-import { ToolSet } from '../generate-text/tool-set';
+import { ExpandedContext, ToolSet } from '../generate-text/tool-set';
 import { Prompt } from '../prompt';
 import { Agent, AgentCallParameters, AgentStreamParameters } from './agent';
-import { ToolLoopAgentSettings } from './tool-loop-agent-settings';
+import {
+  ToolLoopAgentOnStartCallback,
+  ToolLoopAgentOnStepStartCallback,
+  ToolLoopAgentSettings,
+} from './tool-loop-agent-settings';
 
 /**
  * A tool loop agent is an agent that runs tools in a loop. In each step,
@@ -23,14 +27,22 @@ import { ToolLoopAgentSettings } from './tool-loop-agent-settings';
 export class ToolLoopAgent<
   CALL_OPTIONS = never,
   TOOLS extends ToolSet = {},
+  CONTEXT extends ExpandedContext<TOOLS> = ExpandedContext<TOOLS>,
   OUTPUT extends Output = never,
-> implements Agent<CALL_OPTIONS, TOOLS, OUTPUT>
+> implements Agent<CALL_OPTIONS, TOOLS, CONTEXT, OUTPUT>
 {
   readonly version = 'agent-v1';
 
-  private readonly settings: ToolLoopAgentSettings<CALL_OPTIONS, TOOLS, OUTPUT>;
+  private readonly settings: ToolLoopAgentSettings<
+    CALL_OPTIONS,
+    TOOLS,
+    CONTEXT,
+    OUTPUT
+  >;
 
-  constructor(settings: ToolLoopAgentSettings<CALL_OPTIONS, TOOLS, OUTPUT>) {
+  constructor(
+    settings: ToolLoopAgentSettings<CALL_OPTIONS, TOOLS, CONTEXT, OUTPUT>,
+  ) {
     this.settings = settings;
   }
 
@@ -54,7 +66,7 @@ export class ToolLoopAgent<
     options?: CALL_OPTIONS;
   }): Promise<
     Omit<
-      ToolLoopAgentSettings<CALL_OPTIONS, TOOLS, OUTPUT>,
+      ToolLoopAgentSettings<CALL_OPTIONS, TOOLS, CONTEXT, OUTPUT>,
       | 'prepareCall'
       | 'instructions'
       | 'experimental_onStart'
@@ -86,7 +98,12 @@ export class ToolLoopAgent<
       (await this.settings.prepareCall?.(
         baseCallArgs as Parameters<
           NonNullable<
-            ToolLoopAgentSettings<CALL_OPTIONS, TOOLS, OUTPUT>['prepareCall']
+            ToolLoopAgentSettings<
+              CALL_OPTIONS,
+              TOOLS,
+              CONTEXT,
+              OUTPUT
+            >['prepareCall']
           >
         >[0],
       )) ?? baseCallArgs;
@@ -128,7 +145,7 @@ export class ToolLoopAgent<
     onFinish,
     ...options
   }: AgentCallParameters<CALL_OPTIONS, TOOLS>): Promise<
-    GenerateTextResult<TOOLS, OUTPUT>
+    GenerateTextResult<TOOLS, CONTEXT, OUTPUT>
   > {
     return generateText({
       ...(await this.prepareCall(options)),
@@ -136,11 +153,15 @@ export class ToolLoopAgent<
       timeout,
       experimental_onStart: this.mergeCallbacks(
         this.settings.experimental_onStart,
-        experimental_onStart,
+        experimental_onStart as
+          | ToolLoopAgentOnStartCallback<TOOLS, CONTEXT, OUTPUT>
+          | undefined,
       ),
       experimental_onStepStart: this.mergeCallbacks(
         this.settings.experimental_onStepStart,
-        experimental_onStepStart,
+        experimental_onStepStart as
+          | ToolLoopAgentOnStepStartCallback<TOOLS, CONTEXT, OUTPUT>
+          | undefined,
       ),
       experimental_onToolCallStart: this.mergeCallbacks(
         this.settings.experimental_onToolCallStart,
@@ -173,7 +194,7 @@ export class ToolLoopAgent<
     onFinish,
     ...options
   }: AgentStreamParameters<CALL_OPTIONS, TOOLS>): Promise<
-    StreamTextResult<TOOLS, OUTPUT>
+    StreamTextResult<TOOLS, CONTEXT, OUTPUT>
   > {
     return streamText({
       ...(await this.prepareCall(options)),
@@ -182,11 +203,15 @@ export class ToolLoopAgent<
       experimental_transform,
       experimental_onStart: this.mergeCallbacks(
         this.settings.experimental_onStart,
-        experimental_onStart,
+        experimental_onStart as
+          | ToolLoopAgentOnStartCallback<TOOLS, CONTEXT, OUTPUT>
+          | undefined,
       ),
       experimental_onStepStart: this.mergeCallbacks(
         this.settings.experimental_onStepStart,
-        experimental_onStepStart,
+        experimental_onStepStart as
+          | ToolLoopAgentOnStepStartCallback<TOOLS, CONTEXT, OUTPUT>
+          | undefined,
       ),
       experimental_onToolCallStart: this.mergeCallbacks(
         this.settings.experimental_onToolCallStart,
diff --git a/packages/ai/src/generate-text/stream-text-result.ts b/packages/ai/src/generate-text/stream-text-result.ts
index bc2840be90f1..07fc05f9d7c5 100644
--- a/packages/ai/src/generate-text/stream-text-result.ts
+++ b/packages/ai/src/generate-text/stream-text-result.ts
@@ -35,7 +35,7 @@ import {
   StaticToolResult,
   TypedToolResult,
 } from './tool-result';
-import { ToolSet } from './tool-set';
+import { ExpandedContext, ToolSet } from './tool-set';
 
 export type UIMessageStreamOptions<UI_MESSAGE extends UIMessage> = {
   /**
@@ -108,6 +108,7 @@ export type ConsumeStreamOptions = {
  */
 export interface StreamTextResult<
   TOOLS extends ToolSet,
+  CONTEXT extends ExpandedContext<TOOLS>,
   OUTPUT extends Output,
 > {
   /**
@@ -237,7 +238,7 @@ export interface StreamTextResult<
    *
    * Automatically consumes the stream.
    */
-  readonly steps: PromiseLike<Array<StepResult<TOOLS>>>;
+  readonly steps: PromiseLike<Array<StepResult<TOOLS, CONTEXT>>>;
 
   /**
    * Additional request information from the last step.
diff --git a/packages/ai/src/generate-text/stream-text.ts b/packages/ai/src/generate-text/stream-text.ts
index cd8138844c77..2fdc18603b04 100644
--- a/packages/ai/src/generate-text/stream-text.ts
+++ b/packages/ai/src/generate-text/stream-text.ts
@@ -541,7 +541,7 @@ export function streamText<
       now?: () => number;
       generateId?: IdGenerator;
     };
-  }): StreamTextResult<TOOLS, OUTPUT> {
+  }): StreamTextResult<TOOLS, CONTEXT, OUTPUT> {
   const totalTimeoutMs = getTotalTimeoutMs(timeout);
   const stepTimeoutMs = getStepTimeoutMs(timeout);
   const chunkTimeoutMs = getChunkTimeoutMs(timeout);
@@ -704,19 +704,19 @@ class DefaultStreamTextResult<
   TOOLS extends ToolSet,
   CONTEXT extends ExpandedContext<TOOLS>,
   OUTPUT extends Output,
-> implements StreamTextResult<TOOLS, OUTPUT>
+> implements StreamTextResult<TOOLS, CONTEXT, OUTPUT>
 {
   private readonly _totalUsage = new DelayedPromise<
-    Awaited<StreamTextResult<TOOLS, OUTPUT>['usage']>
+    Awaited<StreamTextResult<TOOLS, CONTEXT, OUTPUT>['usage']>
   >();
   private readonly _finishReason = new DelayedPromise<
-    Awaited<StreamTextResult<TOOLS, OUTPUT>['finishReason']>
+    Awaited<StreamTextResult<TOOLS, CONTEXT, OUTPUT>['finishReason']>
   >();
   private readonly _rawFinishReason = new DelayedPromise<
-    Awaited<StreamTextResult<TOOLS, OUTPUT>['rawFinishReason']>
+    Awaited<StreamTextResult<TOOLS, CONTEXT, OUTPUT>['rawFinishReason']>
   >();
   private readonly _steps = new DelayedPromise<
-    Awaited<StreamTextResult<TOOLS, OUTPUT>['steps']>
+    Awaited<StreamTextResult<TOOLS, CONTEXT, OUTPUT>['steps']>
   >();
 
   private readonly addStream: (
diff --git a/packages/ai/src/ui/direct-chat-transport.ts b/packages/ai/src/ui/direct-chat-transport.ts
index 6611b4ba7d01..60b20fe01cad 100644
--- a/packages/ai/src/ui/direct-chat-transport.ts
+++ b/packages/ai/src/ui/direct-chat-transport.ts
@@ -1,6 +1,6 @@
 import { Output } from '../generate-text/output';
 import { UIMessageStreamOptions } from '../generate-text/stream-text-result';
-import { ToolSet } from '../generate-text/tool-set';
+import { ExpandedContext, ToolSet } from '../generate-text/tool-set';
 import { UIMessageChunk } from '../ui-message-stream/ui-message-chunks';
 import { Agent } from '../agent/agent';
 import { ChatTransport } from './chat-transport';
@@ -14,13 +14,14 @@ import { validateUIMessages } from './validate-ui-messages';
 export type DirectChatTransportOptions<
   CALL_OPTIONS,
   TOOLS extends ToolSet,
-  OUTPUT extends Output,
+  CONTEXT extends ExpandedContext<TOOLS>,
+  OUTPUT extends Output<TOOLS, CONTEXT>,
   UI_MESSAGE extends UIMessage<unknown, never, InferUITools<TOOLS>>,
 > = {
   /**
    * The agent to use for generating responses.
    */
-  agent: Agent<CALL_OPTIONS, TOOLS, OUTPUT>;
+  agent: Agent<CALL_OPTIONS, TOOLS, CONTEXT, OUTPUT>;
 
   /**
    * Options to pass to the agent when calling it.
@@ -49,7 +50,8 @@ export type DirectChatTransportOptions<
 export class DirectChatTransport<
   CALL_OPTIONS = never,
   TOOLS extends ToolSet = {},
-  OUTPUT extends Output = never,
+  CONTEXT extends ExpandedContext<TOOLS> = ExpandedContext<TOOLS>,
+  OUTPUT extends Output<TOOLS, CONTEXT> = never,
   UI_MESSAGE extends UIMessage<unknown, never, InferUITools<TOOLS>> = UIMessage<
     unknown,
     never,
@@ -57,7 +59,7 @@ export class DirectChatTransport<
   >,
 > implements ChatTransport<UI_MESSAGE>
 {
-  private readonly agent: Agent<CALL_OPTIONS, TOOLS, OUTPUT>;
+  private readonly agent: Agent<CALL_OPTIONS, TOOLS, CONTEXT, OUTPUT>;
   private readonly agentOptions: CALL_OPTIONS | undefined;
   private readonly uiMessageStreamOptions: Omit<
     UIMessageStreamOptions<UI_MESSAGE>,
@@ -68,7 +70,13 @@ export class DirectChatTransport<
     agent,
     options,
     ...uiMessageStreamOptions
-  }: DirectChatTransportOptions<CALL_OPTIONS, TOOLS, OUTPUT, UI_MESSAGE>) {
+  }: DirectChatTransportOptions<
+    CALL_OPTIONS,
+    TOOLS,
+    CONTEXT,
+    OUTPUT,
+    UI_MESSAGE
+  >) {
     this.agent = agent;
     this.agentOptions = options;
     this.uiMessageStreamOptions = uiMessageStreamOptions;
@@ -98,7 +106,7 @@ export class DirectChatTransport<
       ...(this.agentOptions !== undefined
         ? { options: this.agentOptions }
         : {}),
-    } as Parameters<Agent<CALL_OPTIONS, TOOLS, OUTPUT>['stream']>[0]);
+    } as Parameters<Agent<CALL_OPTIONS, TOOLS, CONTEXT, OUTPUT>['stream']>[0]);
 
     // Return the UI message stream
     return result.toUIMessageStream(this.uiMessageStreamOptions);

From eebb671eca5a86328700646e18c0da381f762777 Mon Sep 17 00:00:00 2001
From: Lars Grammel <lars.grammel@gmail.com>
Date: Mon, 9 Mar 2026 16:03:56 +0100
Subject: [PATCH 25/41] re

---
 .../ai/src/generate-text/callback-events.ts   | 24 +++++++++----------
 1 file changed, 12 insertions(+), 12 deletions(-)

diff --git a/packages/ai/src/generate-text/callback-events.ts b/packages/ai/src/generate-text/callback-events.ts
index e3f45081fc8a..add4a7b4ba38 100644
--- a/packages/ai/src/generate-text/callback-events.ts
+++ b/packages/ai/src/generate-text/callback-events.ts
@@ -29,8 +29,8 @@ export interface CallbackModelInfo {
  * Called when the generation operation begins, before any LLM calls.
  */
 export interface OnStartEvent<
-  TOOLS extends ToolSet,
-  CONTEXT extends ExpandedContext<TOOLS>,
+  TOOLS extends ToolSet = ToolSet,
+  CONTEXT extends ExpandedContext<TOOLS> = ExpandedContext<TOOLS>,
   OUTPUT extends Output = Output,
   INCLUDE = { requestBody?: boolean; responseBody?: boolean },
 > {
@@ -95,8 +95,8 @@ export interface OnStartEvent<
    * When the condition is an array, any of the conditions can be met to stop.
    */
   readonly stopWhen:
-    | StopCondition<TOOLS, CONTEXT>
-    | Array<StopCondition<TOOLS, CONTEXT>>
+    | StopCondition<NoInfer<TOOLS>, CONTEXT>
+    | Array<StopCondition<NoInfer<TOOLS>, CONTEXT>>
     | undefined;
 
   /** The output specification for structured outputs, if configured. */
@@ -130,8 +130,8 @@ export interface OnStartEvent<
  * Each step represents a single LLM invocation.
  */
 export interface OnStepStartEvent<
-  TOOLS extends ToolSet,
-  CONTEXT extends ExpandedContext<TOOLS>,
+  TOOLS extends ToolSet = ToolSet,
+  CONTEXT extends ExpandedContext<TOOLS> = ExpandedContext<TOOLS>,
   OUTPUT extends Output = Output,
   INCLUDE = { requestBody?: boolean; responseBody?: boolean },
 > {
@@ -218,7 +218,7 @@ export interface OnStepStartEvent<
  *
  * Called when a tool execution begins, before the tool's `execute` function is invoked.
  */
-export interface OnToolCallStartEvent<TOOLS extends ToolSet> {
+export interface OnToolCallStartEvent<TOOLS extends ToolSet = ToolSet> {
   /** Zero-based index of the current step where this tool call occurs. */
   readonly stepNumber: number | undefined;
 
@@ -250,7 +250,7 @@ export interface OnToolCallStartEvent<TOOLS extends ToolSet> {
  * Called when a tool execution completes, either successfully or with an error.
  * Uses a discriminated union on the `success` field.
  */
-export type OnToolCallFinishEvent<TOOLS extends ToolSet> = {
+export type OnToolCallFinishEvent<TOOLS extends ToolSet = ToolSet> = {
   /** Zero-based index of the current step where this tool call occurred. */
   readonly stepNumber: number | undefined;
 
@@ -301,8 +301,8 @@ export type OnToolCallFinishEvent<TOOLS extends ToolSet> = {
  * This is simply the StepResult for that step.
  */
 export type OnStepFinishEvent<
-  TOOLS extends ToolSet,
-  CONTEXT extends ExpandedContext<TOOLS>,
+  TOOLS extends ToolSet = ToolSet,
+  CONTEXT extends ExpandedContext<TOOLS> = ExpandedContext<TOOLS>,
 > = StepResult<TOOLS, CONTEXT>;
 
 /**
@@ -312,8 +312,8 @@ export type OnStepFinishEvent<
  * Includes the final step's result along with aggregated data from all steps.
  */
 export type OnFinishEvent<
-  TOOLS extends ToolSet,
-  CONTEXT extends ExpandedContext<TOOLS>,
+  TOOLS extends ToolSet = ToolSet,
+  CONTEXT extends ExpandedContext<TOOLS> = ExpandedContext<TOOLS>,
 > = StepResult<TOOLS, CONTEXT> & {
   /** Array containing results from all steps in the generation. */
   readonly steps: StepResult<TOOLS, CONTEXT>[];

From 61a61c7f5b5755268ae3993ef09dce5f02752480 Mon Sep 17 00:00:00 2001
From: Lars Grammel <lars.grammel@gmail.com>
Date: Tue, 10 Mar 2026 16:33:08 +0100
Subject: [PATCH 26/41] ex

---
 .../openai/tool-call-with-context.ts          | 26 +++++++++++--------
 1 file changed, 15 insertions(+), 11 deletions(-)

diff --git a/examples/ai-functions/src/generate-text/openai/tool-call-with-context.ts b/examples/ai-functions/src/generate-text/openai/tool-call-with-context.ts
index 4436ffce9098..2ac1cb6dbb06 100644
--- a/examples/ai-functions/src/generate-text/openai/tool-call-with-context.ts
+++ b/examples/ai-functions/src/generate-text/openai/tool-call-with-context.ts
@@ -5,7 +5,7 @@ import { run } from '../../lib/run';
 
 run(async () => {
   const result = await generateText({
-    model: openai('gpt-4o'),
+    model: openai('gpt-5-mini'),
     tools: {
       weather: tool({
         description: 'Get the weather in a location',
@@ -15,10 +15,11 @@ run(async () => {
         contextSchema: z.object({
           weatherApiKey: z.string().describe('The API key for the weather API'),
         }),
-        execute: async ({ location }, { experimental_context: context }) => {
-          console.log(context);
-
-          context satisfies { weatherApiKey: string };
+        execute: async (
+          { location },
+          { experimental_context: { weatherApiKey } },
+        ) => {
+          console.log('weather tool api key:', weatherApiKey);
 
           return {
             location,
@@ -38,8 +39,11 @@ run(async () => {
             .string()
             .describe('The API key for the calculator API'),
         }),
-        execute: async ({ expression }, { experimental_context: context }) => {
-          console.log(context);
+        execute: async (
+          { expression },
+          { experimental_context: { calculatorApiKey } },
+        ) => {
+          console.log('calculator tool api key:', calculatorApiKey);
           return {
             expression,
             result: eval(expression),
@@ -48,12 +52,12 @@ run(async () => {
       }),
     },
     experimental_context: {
-      weatherApiKey: '123',
-      calculatorApiKey: '456',
-      somethingElse: 'context',
+      weatherApiKey: 'weather-123',
+      calculatorApiKey: 'calculator-456',
+      somethingElse: 'other-context',
     },
     prepareStep: async ({ experimental_context: context }) => {
-      console.log(context);
+      console.log('prepareStep context:', context);
       return {};
     },
     prompt: 'What is the weather in San Francisco?',

From ec894c2a457d6c31f18e8e679d5e01deff9adb2b Mon Sep 17 00:00:00 2001
From: Lars Grammel <lars.grammel@gmail.com>
Date: Wed, 1 Apr 2026 12:32:39 +0200
Subject: [PATCH 27/41] types

---
 .../ai/src/agent/tool-loop-agent.test-d.ts    | 14 ++--
 ...reate-execute-tools-transformation.test.ts | 24 +++----
 .../create-execute-tools-transformation.ts    |  9 ++-
 .../generate-text/execute-tool-call.test.ts   | 60 ++++++++--------
 .../src/generate-text/generate-text.test.ts   | 72 ++++++++++---------
 .../ai/src/generate-text/stream-text.test.ts  | 70 +++++++++---------
 packages/ai/src/generate-text/tool-call.ts    |  4 +-
 .../telemetry/open-telemetry-integration.ts   |  5 +-
 .../src/provider-tool-factory.ts              | 12 ++--
 9 files changed, 141 insertions(+), 129 deletions(-)

diff --git a/packages/ai/src/agent/tool-loop-agent.test-d.ts b/packages/ai/src/agent/tool-loop-agent.test-d.ts
index 86be0b7c0109..c2169b1e02aa 100644
--- a/packages/ai/src/agent/tool-loop-agent.test-d.ts
+++ b/packages/ai/src/agent/tool-loop-agent.test-d.ts
@@ -11,11 +11,13 @@ import type { ToolLoopAgentOnFinishCallback } from './tool-loop-agent-settings';
 describe('ToolLoopAgent', () => {
   describe('onFinish callback type compatibility', () => {
     it('should allow StreamTextOnFinishCallback where ToolLoopAgentOnFinishCallback is expected', () => {
-      const streamTextCallback: StreamTextOnFinishCallback<{}> =
-        async event => {
-          const context: unknown = event.experimental_context;
-          context;
-        };
+      const streamTextCallback: StreamTextOnFinishCallback<
+        {},
+        {}
+      > = async event => {
+        const context: unknown = event.experimental_context;
+        context;
+      };
 
       expectTypeOf(streamTextCallback).toMatchTypeOf<
         ToolLoopAgentOnFinishCallback<{}>
@@ -29,7 +31,7 @@ describe('ToolLoopAgent', () => {
       };
 
       expectTypeOf(agentCallback).toMatchTypeOf<
-        StreamTextOnFinishCallback<{}>
+        StreamTextOnFinishCallback<{}, {}>
       >();
     });
   });
diff --git a/packages/ai/src/generate-text/create-execute-tools-transformation.test.ts b/packages/ai/src/generate-text/create-execute-tools-transformation.test.ts
index b68a63d76ae6..69d56e8a194d 100644
--- a/packages/ai/src/generate-text/create-execute-tools-transformation.test.ts
+++ b/packages/ai/src/generate-text/create-execute-tools-transformation.test.ts
@@ -58,7 +58,7 @@ describe('createExecuteToolsTransformation', () => {
         messages: [],
         timeout: undefined,
         abortSignal: undefined,
-        experimental_context: undefined,
+        experimental_context: {},
       }),
     );
 
@@ -136,7 +136,7 @@ describe('createExecuteToolsTransformation', () => {
       messages: [],
       abortSignal: undefined,
       timeout: undefined,
-      experimental_context: undefined,
+      experimental_context: {},
     });
 
     expect(
@@ -231,7 +231,7 @@ describe('createExecuteToolsTransformation', () => {
         messages: [],
         abortSignal: undefined,
         timeout: undefined,
-        experimental_context: undefined,
+        experimental_context: {},
       }),
     );
 
@@ -274,7 +274,7 @@ describe('createExecuteToolsTransformation', () => {
           messages: [],
           timeout: undefined,
           abortSignal: undefined,
-          experimental_context: undefined,
+          experimental_context: {},
           onToolCallStart: async () => {
             callOrder.push('onToolCallStart');
           },
@@ -324,7 +324,7 @@ describe('createExecuteToolsTransformation', () => {
           messages: [],
           timeout: undefined,
           abortSignal: undefined,
-          experimental_context: undefined,
+          experimental_context: {},
           stepNumber: 2,
           provider: 'test-provider',
           modelId: 'test-model',
@@ -397,7 +397,7 @@ describe('createExecuteToolsTransformation', () => {
           messages: [],
           timeout: undefined,
           abortSignal: undefined,
-          experimental_context: undefined,
+          experimental_context: {},
           onToolCallFinish: async event => {
             finishEvents.push(event);
           },
@@ -454,7 +454,7 @@ describe('createExecuteToolsTransformation', () => {
           messages: [],
           timeout: undefined,
           abortSignal: undefined,
-          experimental_context: undefined,
+          experimental_context: {},
           onToolCallFinish: async event => {
             finishEvents.push(event);
           },
@@ -504,7 +504,7 @@ describe('createExecuteToolsTransformation', () => {
           messages: [],
           timeout: undefined,
           abortSignal: undefined,
-          experimental_context: undefined,
+          experimental_context: {},
           onToolCallStart: async event => {
             startEvents.push(event);
           },
@@ -557,7 +557,7 @@ describe('createExecuteToolsTransformation', () => {
           messages: [],
           timeout: undefined,
           abortSignal: undefined,
-          experimental_context: undefined,
+          experimental_context: {},
           onToolCallStart: async event => {
             startEvents.push(event.toolCall.toolCallId);
           },
@@ -613,7 +613,7 @@ describe('createExecuteToolsTransformation', () => {
           messages: [],
           timeout: undefined,
           abortSignal: undefined,
-          experimental_context: undefined,
+          experimental_context: {},
           onToolCallStart: async event => {
             startEvents.push(event);
           },
@@ -667,7 +667,7 @@ describe('createExecuteToolsTransformation', () => {
           messages: [],
           abortSignal: undefined,
           timeout: undefined,
-          experimental_context: undefined,
+          experimental_context: {},
         }),
       );
 
@@ -753,7 +753,7 @@ describe('createExecuteToolsTransformation', () => {
           callId: 'test-telemetry-call-id',
           messages: [],
           abortSignal: undefined,
-          experimental_context: undefined,
+          experimental_context: {},
         }),
       );
 
diff --git a/packages/ai/src/generate-text/create-execute-tools-transformation.ts b/packages/ai/src/generate-text/create-execute-tools-transformation.ts
index f2271ecfad78..736f5886b932 100644
--- a/packages/ai/src/generate-text/create-execute-tools-transformation.ts
+++ b/packages/ai/src/generate-text/create-execute-tools-transformation.ts
@@ -9,10 +9,13 @@ import {
   StreamTextOnToolCallStartCallback,
 } from './stream-text';
 import { TypedToolCall } from './tool-call';
-import { ToolSet } from './tool-set';
+import { ExpandedContext, ToolSet } from './tool-set';
 import { ModelCallStreamPart } from './stream-model-call';
 
-export function createExecuteToolsTransformation<TOOLS extends ToolSet>({
+export function createExecuteToolsTransformation<
+  TOOLS extends ToolSet,
+  CONTEXT extends ExpandedContext<TOOLS>,
+>({
   tools,
   telemetry,
   callId,
@@ -34,7 +37,7 @@ export function createExecuteToolsTransformation<TOOLS extends ToolSet>({
   messages: ModelMessage[];
   abortSignal: AbortSignal | undefined;
   timeout?: TimeoutConfiguration<TOOLS>;
-  experimental_context: unknown;
+  experimental_context: CONTEXT;
   generateId: IdGenerator;
   stepNumber?: number;
   provider?: string;
diff --git a/packages/ai/src/generate-text/execute-tool-call.test.ts b/packages/ai/src/generate-text/execute-tool-call.test.ts
index f19ce9febd65..87c29f8dcd74 100644
--- a/packages/ai/src/generate-text/execute-tool-call.test.ts
+++ b/packages/ai/src/generate-text/execute-tool-call.test.ts
@@ -48,7 +48,7 @@ describe('executeToolCall', () => {
         callId: 'test-telemetry-call-id',
         messages: [],
         abortSignal: undefined,
-        experimental_context: undefined,
+        experimental_context: {},
       });
 
       expect(result).toBeUndefined();
@@ -69,7 +69,7 @@ describe('executeToolCall', () => {
         callId: 'test-telemetry-call-id',
         messages: [],
         abortSignal: undefined,
-        experimental_context: undefined,
+        experimental_context: {},
       });
 
       expect(result).toEqual({
@@ -97,7 +97,7 @@ describe('executeToolCall', () => {
         callId: 'test-telemetry-call-id',
         messages: [],
         abortSignal: undefined,
-        experimental_context: undefined,
+        experimental_context: {},
       });
 
       expect(result).toMatchObject({
@@ -125,7 +125,7 @@ describe('executeToolCall', () => {
         callId: 'test-telemetry-call-id',
         messages: [],
         abortSignal: undefined,
-        experimental_context: undefined,
+        experimental_context: {},
       });
 
       expect(result).toEqual({
@@ -155,7 +155,7 @@ describe('executeToolCall', () => {
         callId: 'test-telemetry-call-id',
         messages: [],
         abortSignal: undefined,
-        experimental_context: undefined,
+        experimental_context: {},
       });
 
       expect(result).toMatchObject({
@@ -229,7 +229,7 @@ describe('executeToolCall', () => {
         callId: 'test-telemetry-call-id',
         messages: [],
         abortSignal: undefined,
-        experimental_context: undefined,
+        experimental_context: {},
         onToolCallStart: async () => {
           throw new Error('callback error');
         },
@@ -357,7 +357,7 @@ describe('executeToolCall', () => {
         callId: 'test-telemetry-call-id',
         messages: [],
         abortSignal: undefined,
-        experimental_context: undefined,
+        experimental_context: {},
         onToolCallFinish: async () => {
           throw new Error('callback error');
         },
@@ -386,7 +386,7 @@ describe('executeToolCall', () => {
         callId: 'test-telemetry-call-id',
         messages: [],
         abortSignal: undefined,
-        experimental_context: undefined,
+        experimental_context: {},
         onToolCallFinish: async () => {
           throw new Error('callback error');
         },
@@ -419,7 +419,7 @@ describe('executeToolCall', () => {
         callId: 'test-telemetry-call-id',
         messages: [],
         abortSignal: undefined,
-        experimental_context: undefined,
+        experimental_context: {},
         onToolCallFinish: async event => {
           finishEvents.push(event);
         },
@@ -449,7 +449,7 @@ describe('executeToolCall', () => {
         callId: 'test-telemetry-call-id',
         messages: [],
         abortSignal: undefined,
-        experimental_context: undefined,
+        experimental_context: {},
         onToolCallFinish: async event => {
           finishEvents.push(event);
         },
@@ -479,7 +479,7 @@ describe('executeToolCall', () => {
         callId: 'test-telemetry-call-id',
         messages: [],
         abortSignal: undefined,
-        experimental_context: undefined,
+        experimental_context: {},
         onPreliminaryToolResult: result => {
           preliminaryResults.push(result);
         },
@@ -516,7 +516,7 @@ describe('executeToolCall', () => {
         callId: 'test-telemetry-call-id',
         messages: [],
         abortSignal: undefined,
-        experimental_context: undefined,
+        experimental_context: {},
         onPreliminaryToolResult: result => {
           preliminaryResults.push(result);
         },
@@ -611,7 +611,7 @@ describe('executeToolCall', () => {
         callId: 'test-telemetry-call-id',
         messages: [],
         abortSignal: undefined,
-        experimental_context: undefined,
+        experimental_context: {},
         onToolCallFinish: async event => {
           finishEvents.push(event);
         },
@@ -645,7 +645,7 @@ describe('executeToolCall', () => {
         callId: 'test-telemetry-call-id',
         messages: [],
         abortSignal: undefined,
-        experimental_context: undefined,
+        experimental_context: {},
         executeToolInTelemetryContext,
       });
 
@@ -669,7 +669,7 @@ describe('executeToolCall', () => {
         callId: 'test-telemetry-call-id',
         messages: [],
         abortSignal: undefined,
-        experimental_context: undefined,
+        experimental_context: {},
       });
 
       expect(result).toMatchObject({
@@ -694,7 +694,7 @@ describe('executeToolCall', () => {
         messages: [],
         abortSignal: undefined,
         timeout: { toolMs: 5000 },
-        experimental_context: undefined,
+        experimental_context: {},
       });
 
       expect(result).toMatchObject({
@@ -722,7 +722,7 @@ describe('executeToolCall', () => {
         messages: [],
         abortSignal: undefined,
         timeout: { toolMs: 5000 },
-        experimental_context: undefined,
+        experimental_context: {},
       });
 
       expect(receivedSignal).toBeDefined();
@@ -747,7 +747,7 @@ describe('executeToolCall', () => {
         callId: 'test-telemetry-call-id',
         messages: [],
         abortSignal: undefined,
-        experimental_context: undefined,
+        experimental_context: {},
       });
 
       expect(receivedSignal).toBeUndefined();
@@ -773,7 +773,7 @@ describe('executeToolCall', () => {
         messages: [],
         abortSignal: controller.signal,
         timeout: { toolMs: 5000 },
-        experimental_context: undefined,
+        experimental_context: {},
       });
 
       expect(receivedSignal).toBeDefined();
@@ -800,7 +800,7 @@ describe('executeToolCall', () => {
         messages: [],
         abortSignal: undefined,
         timeout: { toolMs: 10000, tools: { testToolMs: 2000 } },
-        experimental_context: undefined,
+        experimental_context: {},
       });
 
       expect(receivedSignal).toBeDefined();
@@ -826,7 +826,7 @@ describe('executeToolCall', () => {
         messages: [],
         abortSignal: undefined,
         timeout: { toolMs: 5000, tools: { otherToolMs: 2000 } },
-        experimental_context: undefined,
+        experimental_context: {},
       });
 
       expect(receivedSignal).toBeDefined();
@@ -852,7 +852,7 @@ describe('executeToolCall', () => {
         messages: [],
         abortSignal: undefined,
         timeout: { tools: { otherToolMs: 2000 } },
-        experimental_context: undefined,
+        experimental_context: {},
       });
 
       expect(receivedSignal).toBeUndefined();
@@ -875,7 +875,7 @@ describe('executeToolCall', () => {
         callId: 'test-telemetry-call-id',
         messages: [],
         abortSignal: undefined,
-        experimental_context: undefined,
+        experimental_context: {},
       });
 
       expect(result).toMatchObject({
@@ -901,7 +901,7 @@ describe('executeToolCall', () => {
         callId: 'test-telemetry-call-id',
         messages: [],
         abortSignal: undefined,
-        experimental_context: undefined,
+        experimental_context: {},
       });
 
       expect(result).toMatchObject({
@@ -920,7 +920,7 @@ describe('executeToolCall', () => {
         callId: 'test-telemetry-call-id',
         messages: [],
         abortSignal: undefined,
-        experimental_context: undefined,
+        experimental_context: {},
       });
 
       expect(result).toBeUndefined();
@@ -941,7 +941,7 @@ describe('executeToolCall', () => {
         callId: 'test-telemetry-call-id',
         messages: [],
         abortSignal: undefined,
-        experimental_context: undefined,
+        experimental_context: {},
       });
 
       expect(result).toBeUndefined();
@@ -964,7 +964,7 @@ describe('executeToolCall', () => {
         callId: 'test-telemetry-call-id',
         messages: [],
         abortSignal: undefined,
-        experimental_context: undefined,
+        experimental_context: {},
         onToolCallStart: [
           async () => {
             calls.push('first');
@@ -993,7 +993,7 @@ describe('executeToolCall', () => {
         callId: 'test-telemetry-call-id',
         messages: [],
         abortSignal: undefined,
-        experimental_context: undefined,
+        experimental_context: {},
         onToolCallFinish: [
           async () => {
             calls.push('first');
@@ -1022,7 +1022,7 @@ describe('executeToolCall', () => {
         callId: 'test-telemetry-call-id',
         messages: [],
         abortSignal: undefined,
-        experimental_context: undefined,
+        experimental_context: {},
         onToolCallStart: [
           undefined,
           async () => {
@@ -1057,7 +1057,7 @@ describe('executeToolCall', () => {
         callId: 'test-telemetry-call-id',
         messages: [],
         abortSignal: undefined,
-        experimental_context: undefined,
+        experimental_context: {},
         onToolCallStart: [
           async () => {
             throw new Error('listener error');
diff --git a/packages/ai/src/generate-text/generate-text.test.ts b/packages/ai/src/generate-text/generate-text.test.ts
index f2e10805eb71..26afabfd70a9 100644
--- a/packages/ai/src/generate-text/generate-text.test.ts
+++ b/packages/ai/src/generate-text/generate-text.test.ts
@@ -812,7 +812,9 @@ describe('generateText', () => {
 
   describe('options.experimental_onStart', () => {
     it('should send correct information with text prompt', async () => {
-      let startEvent!: Parameters<GenerateTextOnStartCallback>[0];
+      let startEvent!: Parameters<
+        GenerateTextOnStartCallback<any, any, any>
+      >[0];
 
       await generateText({
         model: new MockLanguageModelV4({
@@ -839,7 +841,9 @@ describe('generateText', () => {
     });
 
     it('should pass experimental_context', async () => {
-      let startEvent!: Parameters<GenerateTextOnStartCallback>[0];
+      let startEvent!: Parameters<
+        GenerateTextOnStartCallback<any, any, any>
+      >[0];
 
       await generateText({
         model: new MockLanguageModelV4({
@@ -862,7 +866,9 @@ describe('generateText', () => {
     });
 
     it('should send correct information with system and messages', async () => {
-      let startEvent!: Parameters<GenerateTextOnStartCallback>[0];
+      let startEvent!: Parameters<
+        GenerateTextOnStartCallback<any, any, any>
+      >[0];
 
       await generateText({
         model: new MockLanguageModelV4({
@@ -1365,7 +1371,7 @@ describe('generateText', () => {
   describe('options.onStepFinish stepNumber', () => {
     it('should pass stepNumber 0 for a single step', async () => {
       let stepFinishEvent!: Parameters<
-        GenerateTextOnStepFinishCallback<any>
+        GenerateTextOnStepFinishCallback<any, any>
       >[0];
 
       await generateText({
@@ -2178,7 +2184,7 @@ describe('generateText', () => {
 
   describe('options.onFinish', () => {
     it('should send correct information', async () => {
-      let result!: Parameters<GenerateTextOnFinishCallback<any>>[0];
+      let result!: Parameters<GenerateTextOnFinishCallback<any, any>>[0];
 
       await generateText({
         model: new MockLanguageModelV4({
@@ -2513,9 +2519,9 @@ describe('generateText', () => {
   });
 
   describe('options.stopWhen', () => {
-    let result: GenerateTextResult<any, any>;
-    let onFinishResult: Parameters<GenerateTextOnFinishCallback<any>>[0];
-    let onStepFinishResults: StepResult<any>[];
+    let result: GenerateTextResult<any, any, any>;
+    let onFinishResult: Parameters<GenerateTextOnFinishCallback<any, any>>[0];
+    let onStepFinishResults: StepResult<any, any>[];
 
     beforeEach(() => {
       result = undefined as any;
@@ -2711,13 +2717,13 @@ describe('generateText', () => {
     });
 
     describe('2 steps: initial, tool-result with prepareStep', () => {
-      let result: GenerateTextResult<any, any>;
-      let onStepFinishResults: StepResult<any>[];
+      let result: GenerateTextResult<any, any, any>;
+      let onStepFinishResults: StepResult<any, any>[];
       let doGenerateCalls: Array<LanguageModelV4CallOptions>;
       let prepareStepCalls: Array<{
         modelId: string;
         stepNumber: number;
-        steps: Array<StepResult<any>>;
+        steps: Array<StepResult<any, any>>;
         messages: Array<ModelMessage>;
         experimental_context: unknown;
       }>;
@@ -3481,10 +3487,10 @@ describe('generateText', () => {
     });
 
     describe('2 stop conditions', () => {
-      let result: GenerateTextResult<any, any>;
+      let result: GenerateTextResult<any, any, any>;
       let stopConditionCalls: Array<{
         number: number;
-        steps: StepResult<any>[];
+        steps: StepResult<any, any>[];
       }>;
 
       beforeEach(async () => {
@@ -4476,7 +4482,7 @@ describe('generateText', () => {
 
   describe('provider-executed tools', () => {
     describe('two provider-executed tool calls and results', () => {
-      let result: GenerateTextResult<any, any>;
+      let result: GenerateTextResult<any, any, any>;
 
       beforeEach(async () => {
         result = await generateText({
@@ -5069,7 +5075,7 @@ describe('generateText', () => {
   });
 
   describe('tool execution errors', () => {
-    let result: GenerateTextResult<any, any>;
+    let result: GenerateTextResult<any, any, any>;
 
     beforeEach(async () => {
       result = await generateText({
@@ -5260,14 +5266,14 @@ describe('generateText', () => {
 
   describe('programmatic tool calling', () => {
     describe('5 steps: code_execution triggers client tool across multiple turns (dice game fixture)', () => {
-      let result: GenerateTextResult<any, any>;
-      let onFinishResult: Parameters<GenerateTextOnFinishCallback<any>>[0];
-      let onStepFinishResults: StepResult<any>[];
+      let result: GenerateTextResult<any, any, any>;
+      let onFinishResult: Parameters<GenerateTextOnFinishCallback<any, any>>[0];
+      let onStepFinishResults: StepResult<any, any>[];
       let doGenerateCalls: Array<LanguageModelV4CallOptions>;
       let prepareStepCalls: Array<{
         modelId: string;
         stepNumber: number;
-        steps: Array<StepResult<any>>;
+        steps: Array<StepResult<any, any>>;
         messages: Array<ModelMessage>;
       }>;
       let rollDieExecutions: Array<{ player: string }>;
@@ -6598,7 +6604,7 @@ describe('generateText', () => {
 
   describe('invalid tool calls', () => {
     describe('single invalid tool call', () => {
-      let result: GenerateTextResult<any, any>;
+      let result: GenerateTextResult<any, any, any>;
 
       beforeEach(async () => {
         result = await generateText({
@@ -6741,7 +6747,7 @@ describe('generateText', () => {
 
   describe('tools with preliminary results', () => {
     describe('single tool with preliminary results', () => {
-      let result: GenerateTextResult<any, any>;
+      let result: GenerateTextResult<any, any, any>;
 
       beforeEach(async () => {
         result = await generateText({
@@ -7062,7 +7068,7 @@ describe('generateText', () => {
 
   describe('tool execution approval', () => {
     describe('when a single tool needs approval', () => {
-      let result: GenerateTextResult<any, any>;
+      let result: GenerateTextResult<any, any, any>;
 
       beforeEach(async () => {
         result = await generateText({
@@ -7167,7 +7173,7 @@ describe('generateText', () => {
     });
 
     describe('when a single tool has a needsApproval function', () => {
-      let result: GenerateTextResult<any, any>;
+      let result: GenerateTextResult<any, any, any>;
       let needsApprovalCalls: Array<{ input: any; options: any }> = [];
 
       beforeEach(async () => {
@@ -7366,9 +7372,9 @@ describe('generateText', () => {
     });
 
     describe('when a call from a single tool that needs approval is approved', () => {
-      let result: GenerateTextResult<any, any>;
+      let result: GenerateTextResult<any, any, any>;
       let prompts: LanguageModelV4Prompt[];
-      let executeFunction: ToolExecuteFunction<any, any>;
+      let executeFunction: ToolExecuteFunction<any, any, any>;
 
       beforeEach(async () => {
         prompts = [];
@@ -7532,7 +7538,7 @@ describe('generateText', () => {
     });
 
     describe('when a call from a single tool that needs approval is approved and the tool throws', () => {
-      let result: GenerateTextResult<any, any>;
+      let result: GenerateTextResult<any, any, any>;
       let prompts: LanguageModelV4Prompt[];
 
       beforeEach(async () => {
@@ -7646,9 +7652,9 @@ describe('generateText', () => {
     });
 
     describe('when a call from a single tool that needs approval is denied', () => {
-      let result: GenerateTextResult<any, any>;
+      let result: GenerateTextResult<any, any, any>;
       let prompts: LanguageModelV4Prompt[];
-      let executeFunction: ToolExecuteFunction<any, any>;
+      let executeFunction: ToolExecuteFunction<any, any, any>;
 
       beforeEach(async () => {
         prompts = [];
@@ -7805,9 +7811,9 @@ describe('generateText', () => {
     });
 
     describe('when two calls from a single tool that needs approval are approved', () => {
-      let result: GenerateTextResult<any, any>;
+      let result: GenerateTextResult<any, any, any>;
       let prompts: LanguageModelV4Prompt[];
-      let executeFunction: ToolExecuteFunction<any, any>;
+      let executeFunction: ToolExecuteFunction<any, any, any>;
 
       beforeEach(async () => {
         prompts = [];
@@ -8010,7 +8016,7 @@ describe('generateText', () => {
 
     describe('provider-executed tool (MCP) approval', () => {
       describe('when a provider-executed tool emits tool-approval-request', () => {
-        let result: GenerateTextResult<any, any>;
+        let result: GenerateTextResult<any, any, any>;
 
         beforeEach(async () => {
           result = await generateText({
@@ -8121,7 +8127,7 @@ describe('generateText', () => {
       });
 
       describe('when a provider-executed tool approval is approved', () => {
-        let result: GenerateTextResult<any, any>;
+        let result: GenerateTextResult<any, any, any>;
         let prompts: LanguageModelV4Prompt[];
 
         beforeEach(async () => {
@@ -8307,7 +8313,7 @@ describe('generateText', () => {
       });
 
       describe('when a provider-executed tool approval is denied', () => {
-        let result: GenerateTextResult<any, any>;
+        let result: GenerateTextResult<any, any, any>;
         let prompts: LanguageModelV4Prompt[];
 
         beforeEach(async () => {
diff --git a/packages/ai/src/generate-text/stream-text.test.ts b/packages/ai/src/generate-text/stream-text.test.ts
index ce7a7ed67e87..a3ed17513aba 100644
--- a/packages/ai/src/generate-text/stream-text.test.ts
+++ b/packages/ai/src/generate-text/stream-text.test.ts
@@ -7041,7 +7041,7 @@ describe('streamText', () => {
 
   describe('options.onFinish', () => {
     it('should send correct information', async () => {
-      let result!: Parameters<StreamTextOnFinishCallback<any>>[0];
+      let result!: Parameters<StreamTextOnFinishCallback<any, any>>[0];
 
       const resultObject = streamText({
         model: createTestModel({
@@ -8023,9 +8023,9 @@ describe('streamText', () => {
   });
 
   describe('options.stopWhen', () => {
-    let result: StreamTextResult<any, any>;
-    let onFinishResult: Parameters<StreamTextOnFinishCallback<any>>[0];
-    let onStepFinishResults: StepResult<any>[];
+    let result: StreamTextResult<any, any, any>;
+    let onFinishResult: Parameters<StreamTextOnFinishCallback<any, any>>[0];
+    let onStepFinishResults: StepResult<any, any>[];
     let tracer: MockTracer;
     let stepInputs: Array<any>;
 
@@ -9344,7 +9344,7 @@ describe('streamText', () => {
       let prepareStepCalls: Array<{
         modelId: string;
         stepNumber: number;
-        steps: Array<StepResult<any>>;
+        steps: Array<StepResult<any, any>>;
         messages: Array<ModelMessage>;
         experimental_context: unknown;
       }>;
@@ -11392,7 +11392,7 @@ describe('streamText', () => {
     describe('2 stop conditions', () => {
       let stopConditionCalls: Array<{
         number: number;
-        steps: StepResult<any>[];
+        steps: StepResult<any, any>[];
       }>;
 
       beforeEach(async () => {
@@ -11820,7 +11820,7 @@ describe('streamText', () => {
 
   describe('provider-executed tools', () => {
     describe('single provider-executed tool call and result', () => {
-      let result: StreamTextResult<any, any>;
+      let result: StreamTextResult<any, any, any>;
 
       beforeEach(async () => {
         result = streamText({
@@ -12455,7 +12455,7 @@ describe('streamText', () => {
 
   describe('dynamic tools', () => {
     describe('single dynamic tool call and result', () => {
-      let result: StreamTextResult<any, any>;
+      let result: StreamTextResult<any, any, any>;
 
       beforeEach(async () => {
         result = streamText({
@@ -13713,7 +13713,7 @@ describe('streamText', () => {
   });
 
   describe('tool execution errors', () => {
-    let result: StreamTextResult<any, any>;
+    let result: StreamTextResult<any, any, any>;
 
     beforeEach(async () => {
       result = streamText({
@@ -16045,7 +16045,7 @@ describe('streamText', () => {
     });
 
     describe('array output', () => {
-      let result: StreamTextResult<any, any> | undefined;
+      let result: StreamTextResult<any, any, any> | undefined;
 
       let onFinishResult:
         | Parameters<Required<Parameters<typeof streamText>[0]>['onFinish']>[0]
@@ -16796,7 +16796,7 @@ describe('streamText', () => {
 
   describe('mixed multi content streaming with interleaving parts', () => {
     describe('mixed text and reasoning blocks', () => {
-      let result: StreamTextResult<any, any>;
+      let result: StreamTextResult<any, any, any>;
 
       beforeEach(async () => {
         result = streamText({
@@ -17120,9 +17120,9 @@ describe('streamText', () => {
 
   describe('abort signal', () => {
     describe('basic abort', () => {
-      let result: StreamTextResult<ToolSet, never>;
+      let result: StreamTextResult<ToolSet, any, never>;
       let onErrorCalls: Array<{ error: unknown }> = [];
-      let onAbortCalls: Array<{ steps: StepResult<ToolSet>[] }> = [];
+      let onAbortCalls: Array<{ steps: StepResult<ToolSet, any>[] }> = [];
 
       beforeEach(() => {
         onErrorCalls = [];
@@ -17301,9 +17301,9 @@ describe('streamText', () => {
     });
 
     describe('abort in 2nd step', () => {
-      let result: StreamTextResult<any, never>;
+      let result: StreamTextResult<any, any, never>;
       let onErrorCalls: Array<{ error: unknown }> = [];
-      let onAbortCalls: Array<{ steps: StepResult<any>[] }> = [];
+      let onAbortCalls: Array<{ steps: StepResult<any, any>[] }> = [];
 
       beforeEach(() => {
         onErrorCalls = [];
@@ -17619,9 +17619,9 @@ describe('streamText', () => {
     });
 
     describe('abort during tool execution', () => {
-      let result: StreamTextResult<any, never>;
+      let result: StreamTextResult<any, any, never>;
       let onErrorCalls: Array<{ error: unknown }> = [];
-      let onAbortCalls: Array<{ steps: StepResult<any>[] }> = [];
+      let onAbortCalls: Array<{ steps: StepResult<any, any>[] }> = [];
 
       beforeEach(() => {
         onErrorCalls = [];
@@ -17906,7 +17906,7 @@ describe('streamText', () => {
 
   describe('invalid tool calls', () => {
     describe('single invalid tool call', () => {
-      let result: StreamTextResult<any, any>;
+      let result: StreamTextResult<any, any, any>;
 
       beforeEach(async () => {
         result = streamText({
@@ -18203,7 +18203,7 @@ describe('streamText', () => {
 
   describe('tools with preliminary results', () => {
     describe('single tool with preliminary results', () => {
-      let result: StreamTextResult<any, any>;
+      let result: StreamTextResult<any, any, any>;
 
       beforeEach(async () => {
         result = streamText({
@@ -18574,7 +18574,7 @@ describe('streamText', () => {
 
   describe('provider-executed dynamic tools', () => {
     describe('single provider-executed dynamic tool with input streaming', () => {
-      let result: StreamTextResult<any, any>;
+      let result: StreamTextResult<any, any, any>;
 
       beforeEach(async () => {
         result = streamText({
@@ -18882,14 +18882,14 @@ describe('streamText', () => {
 
   describe('programmatic tool calling', () => {
     describe('5 steps: code_execution triggers client tool across multiple turns (dice game fixture)', () => {
-      let result: StreamTextResult<any, any>;
-      let onFinishResult: Parameters<StreamTextOnFinishCallback<any>>[0];
-      let onStepFinishResults: StepResult<any>[];
+      let result: StreamTextResult<any, any, any>;
+      let onFinishResult: Parameters<StreamTextOnFinishCallback<any, any>>[0];
+      let onStepFinishResults: StepResult<any, any>[];
       let doStreamCalls: Array<LanguageModelV4CallOptions>;
       let prepareStepCalls: Array<{
         modelId: string;
         stepNumber: number;
-        steps: Array<StepResult<any>>;
+        steps: Array<StepResult<any, any>>;
         messages: Array<ModelMessage>;
       }>;
       let rollDieExecutions: Array<{ player: string }>;
@@ -20788,7 +20788,7 @@ describe('streamText', () => {
 
   describe('tool execution approval', () => {
     describe('when a single tool needs approval', () => {
-      let result: StreamTextResult<any, any>;
+      let result: StreamTextResult<any, any, any>;
 
       beforeEach(async () => {
         result = streamText({
@@ -21012,7 +21012,7 @@ describe('streamText', () => {
     });
 
     describe('when a single tool has a needsApproval function', () => {
-      let result: StreamTextResult<any, any>;
+      let result: StreamTextResult<any, any, any>;
       let needsApprovalCalls: Array<{ input: any; options: any }> = [];
 
       beforeEach(async () => {
@@ -21340,9 +21340,9 @@ describe('streamText', () => {
     });
 
     describe('when a call from a single tool that needs approval is approved', () => {
-      let result: StreamTextResult<any, any>;
+      let result: StreamTextResult<any, any, any>;
       let prompts: LanguageModelV4Prompt[];
-      let executeFunction: ToolExecuteFunction<any, any>;
+      let executeFunction: ToolExecuteFunction<any, any, any>;
 
       beforeEach(async () => {
         prompts = [];
@@ -21644,7 +21644,7 @@ describe('streamText', () => {
     });
 
     describe('when a call from a single tool that needs approval is approved and the tool throws', () => {
-      let result: StreamTextResult<any, any>;
+      let result: StreamTextResult<any, any, any>;
       let prompts: LanguageModelV4Prompt[];
 
       beforeEach(async () => {
@@ -21767,7 +21767,7 @@ describe('streamText', () => {
     });
 
     describe('when a call from a single tool with preliminary results that needs approval is approved', () => {
-      let result: StreamTextResult<any, any>;
+      let result: StreamTextResult<any, any, any>;
       let prompts: LanguageModelV4Prompt[];
 
       beforeEach(async () => {
@@ -22097,9 +22097,9 @@ describe('streamText', () => {
     });
 
     describe('when a call from a single tool that needs approval is denied', () => {
-      let result: StreamTextResult<any, any>;
+      let result: StreamTextResult<any, any, any>;
       let prompts: LanguageModelV4Prompt[];
-      let executeFunction: ToolExecuteFunction<any, any>;
+      let executeFunction: ToolExecuteFunction<any, any, any>;
 
       beforeEach(async () => {
         prompts = [];
@@ -22389,7 +22389,7 @@ describe('streamText', () => {
 
     describe('provider-executed tool (MCP) approval', () => {
       describe('when a provider-executed tool emits tool-approval-request', () => {
-        let result: StreamTextResult<any, any>;
+        let result: StreamTextResult<any, any, any>;
 
         beforeEach(async () => {
           result = streamText({
@@ -22624,7 +22624,7 @@ describe('streamText', () => {
       });
 
       describe('when a provider-executed tool approval is approved', () => {
-        let result: StreamTextResult<any, any>;
+        let result: StreamTextResult<any, any, any>;
         let prompts: LanguageModelV4Prompt[];
 
         beforeEach(async () => {
@@ -22813,7 +22813,7 @@ describe('streamText', () => {
       });
 
       describe('when a provider-executed tool approval is denied', () => {
-        let result: StreamTextResult<any, any>;
+        let result: StreamTextResult<any, any, any>;
         let prompts: LanguageModelV4Prompt[];
 
         beforeEach(async () => {
diff --git a/packages/ai/src/generate-text/tool-call.ts b/packages/ai/src/generate-text/tool-call.ts
index 59698ce0b550..5c5fa589b32a 100644
--- a/packages/ai/src/generate-text/tool-call.ts
+++ b/packages/ai/src/generate-text/tool-call.ts
@@ -1,4 +1,4 @@
-import { Tool } from '@ai-sdk/provider-utils';
+import { InferToolInput } from '@ai-sdk/provider-utils';
 import { ProviderMetadata } from '../types';
 import { ValueOf } from '../util/value-of';
 import { ToolSet } from './tool-set';
@@ -13,7 +13,7 @@ type BaseToolCall = {
 export type StaticToolCall<TOOLS extends ToolSet = ToolSet> = ValueOf<{
   [NAME in keyof TOOLS]: BaseToolCall & {
     toolName: NAME & string;
-    input: TOOLS[NAME] extends Tool<any, infer PARAMETERS> ? PARAMETERS : never;
+    input: InferToolInput<TOOLS[NAME]>;
     dynamic?: false | undefined;
     invalid?: false | undefined;
     error?: never;
diff --git a/packages/ai/src/telemetry/open-telemetry-integration.ts b/packages/ai/src/telemetry/open-telemetry-integration.ts
index 3b4fd00ea77e..89ad5f5e029a 100644
--- a/packages/ai/src/telemetry/open-telemetry-integration.ts
+++ b/packages/ai/src/telemetry/open-telemetry-integration.ts
@@ -31,7 +31,7 @@ import type {
   OnToolCallStartEvent,
 } from '../generate-text/core-events';
 import type { Output } from '../generate-text/output';
-import type { ToolSet } from '../generate-text/tool-set';
+import type { ExpandedContext, ToolSet } from '../generate-text/tool-set';
 import type {
   ObjectOnStartEvent,
   ObjectOnFinishEvent,
@@ -115,8 +115,9 @@ function selectAttributes(
 
 interface OtelStepStartEvent<
   TOOLS extends ToolSet = ToolSet,
+  CONTEXT extends ExpandedContext<TOOLS> = ExpandedContext<TOOLS>,
   OUTPUT extends Output = Output,
-> extends OnStepStartEvent<TOOLS, OUTPUT> {
+> extends OnStepStartEvent<TOOLS, CONTEXT, OUTPUT> {
   readonly promptMessages?: LanguageModelV4Prompt;
   readonly stepTools?: ReadonlyArray<Record<string, unknown>>;
   readonly stepToolChoice?: unknown;
diff --git a/packages/provider-utils/src/provider-tool-factory.ts b/packages/provider-utils/src/provider-tool-factory.ts
index 9766b664a530..77e63ec5bfd1 100644
--- a/packages/provider-utils/src/provider-tool-factory.ts
+++ b/packages/provider-utils/src/provider-tool-factory.ts
@@ -3,9 +3,9 @@ import { FlexibleSchema } from './schema';
 import { Context } from './types/context';
 
 export type ProviderToolFactory<
-  CONTEXT extends Context,
   INPUT,
   ARGS extends object,
+  CONTEXT extends Context = {},
 > = <OUTPUT>(
   options: ARGS & {
     execute?: ToolExecuteFunction<INPUT, OUTPUT, CONTEXT>;
@@ -18,16 +18,16 @@ export type ProviderToolFactory<
 ) => Tool<INPUT, OUTPUT, CONTEXT>;
 
 export function createProviderToolFactory<
-  CONTEXT extends Context,
   INPUT,
   ARGS extends object,
+  CONTEXT extends Context = {},
 >({
   id,
   inputSchema,
 }: {
   id: `${string}.${string}`;
   inputSchema: FlexibleSchema<INPUT>;
-}): ProviderToolFactory<CONTEXT, INPUT, ARGS> {
+}): ProviderToolFactory<INPUT, ARGS, CONTEXT> {
   return <OUTPUT>({
     execute,
     outputSchema,
@@ -62,10 +62,10 @@ export function createProviderToolFactory<
 }
 
 export type ProviderToolFactoryWithOutputSchema<
-  CONTEXT extends Context,
   INPUT,
   OUTPUT,
   ARGS extends object,
+  CONTEXT extends Context = {},
 > = (
   options: ARGS & {
     execute?: ToolExecuteFunction<INPUT, OUTPUT, CONTEXT>;
@@ -78,10 +78,10 @@ export type ProviderToolFactoryWithOutputSchema<
 ) => Tool<INPUT, OUTPUT, CONTEXT>;
 
 export function createProviderToolFactoryWithOutputSchema<
-  CONTEXT extends Context,
   INPUT,
   OUTPUT,
   ARGS extends object,
+  CONTEXT extends Context = {},
 >({
   id,
   inputSchema,
@@ -102,7 +102,7 @@ export function createProviderToolFactoryWithOutputSchema<
    * @default false
    */
   supportsDeferredResults?: boolean;
-}): ProviderToolFactoryWithOutputSchema<CONTEXT, INPUT, OUTPUT, ARGS> {
+}): ProviderToolFactoryWithOutputSchema<INPUT, OUTPUT, ARGS, CONTEXT> {
   return ({
     execute,
     needsApproval,

From a966d4594805960af321362f332f1fec17b1e36a Mon Sep 17 00:00:00 2001
From: Lars Grammel <lars.grammel@gmail.com>
Date: Wed, 1 Apr 2026 12:35:33 +0200
Subject: [PATCH 28/41] tests

---
 .../__snapshots__/generate-text.test.ts.snap  | 24 ++---
 .../__snapshots__/stream-text.test.ts.snap    |  4 +-
 .../src/generate-text/generate-text.test.ts   | 18 ++--
 .../ai/src/generate-text/stream-text.test.ts  | 93 ++++++++++---------
 4 files changed, 71 insertions(+), 68 deletions(-)

diff --git a/packages/ai/src/generate-text/__snapshots__/generate-text.test.ts.snap b/packages/ai/src/generate-text/__snapshots__/generate-text.test.ts.snap
index a67498cf5668..30103ffbd5d0 100644
--- a/packages/ai/src/generate-text/__snapshots__/generate-text.test.ts.snap
+++ b/packages/ai/src/generate-text/__snapshots__/generate-text.test.ts.snap
@@ -5,7 +5,7 @@ exports[`generateText > options.experimental_onStart > should send correct infor
   "abortSignal": undefined,
   "activeTools": undefined,
   "callId": "test-telemetry-call-id",
-  "experimental_context": undefined,
+  "experimental_context": {},
   "frequencyPenalty": undefined,
   "functionId": "test-function",
   "headers": {
@@ -46,7 +46,7 @@ exports[`generateText > options.experimental_onToolCallStart > should be called
 {
   "abortSignal": undefined,
   "callId": "test-telemetry-call-id",
-  "experimental_context": undefined,
+  "experimental_context": {},
   "functionId": undefined,
   "messages": [
     {
@@ -83,7 +83,7 @@ exports[`generateText > options.stopWhen > 2 steps: initial, tool-result > callb
   ],
   "dynamicToolCalls": [],
   "dynamicToolResults": [],
-  "experimental_context": undefined,
+  "experimental_context": {},
   "files": [],
   "finishReason": "stop",
   "functionId": undefined,
@@ -177,7 +177,7 @@ exports[`generateText > options.stopWhen > 2 steps: initial, tool-result > callb
           "type": "tool-result",
         },
       ],
-      "experimental_context": undefined,
+      "experimental_context": {},
       "finishReason": "tool-calls",
       "functionId": undefined,
       "metadata": undefined,
@@ -254,7 +254,7 @@ exports[`generateText > options.stopWhen > 2 steps: initial, tool-result > callb
           "type": "text",
         },
       ],
-      "experimental_context": undefined,
+      "experimental_context": {},
       "finishReason": "stop",
       "functionId": undefined,
       "metadata": undefined,
@@ -403,7 +403,7 @@ exports[`generateText > options.stopWhen > 2 steps: initial, tool-result > callb
         "type": "tool-result",
       },
     ],
-    "experimental_context": undefined,
+    "experimental_context": {},
     "finishReason": "tool-calls",
     "functionId": undefined,
     "metadata": undefined,
@@ -480,7 +480,7 @@ exports[`generateText > options.stopWhen > 2 steps: initial, tool-result > callb
         "type": "text",
       },
     ],
-    "experimental_context": undefined,
+    "experimental_context": {},
     "finishReason": "stop",
     "functionId": undefined,
     "metadata": undefined,
@@ -635,7 +635,7 @@ exports[`generateText > options.stopWhen > 2 steps: initial, tool-result > resul
         "type": "tool-result",
       },
     ],
-    "experimental_context": undefined,
+    "experimental_context": {},
     "finishReason": "tool-calls",
     "functionId": undefined,
     "metadata": undefined,
@@ -712,7 +712,7 @@ exports[`generateText > options.stopWhen > 2 steps: initial, tool-result > resul
         "type": "text",
       },
     ],
-    "experimental_context": undefined,
+    "experimental_context": {},
     "finishReason": "stop",
     "functionId": undefined,
     "metadata": undefined,
@@ -1673,7 +1673,7 @@ exports[`generateText > result.steps > should add the reasoning from the model r
         "type": "text",
       },
     ],
-    "experimental_context": undefined,
+    "experimental_context": {},
     "finishReason": "stop",
     "functionId": undefined,
     "metadata": undefined,
@@ -1774,7 +1774,7 @@ exports[`generateText > result.steps > should contain files 1`] = `
         "type": "file",
       },
     ],
-    "experimental_context": undefined,
+    "experimental_context": {},
     "finishReason": "stop",
     "functionId": undefined,
     "metadata": undefined,
@@ -1873,7 +1873,7 @@ exports[`generateText > result.steps > should contain sources 1`] = `
         "url": "https://example.com/2",
       },
     ],
-    "experimental_context": undefined,
+    "experimental_context": {},
     "finishReason": "stop",
     "functionId": undefined,
     "metadata": undefined,
diff --git a/packages/ai/src/generate-text/__snapshots__/stream-text.test.ts.snap b/packages/ai/src/generate-text/__snapshots__/stream-text.test.ts.snap
index c32c3bf645a1..b5603beb4166 100644
--- a/packages/ai/src/generate-text/__snapshots__/stream-text.test.ts.snap
+++ b/packages/ai/src/generate-text/__snapshots__/stream-text.test.ts.snap
@@ -7,7 +7,7 @@ exports[`streamText > options.experimental_onStart > should send correct informa
   "abortSignal": undefined,
   "activeTools": undefined,
   "callId": "test-telemetry-call-id",
-  "experimental_context": undefined,
+  "experimental_context": {},
   "frequencyPenalty": undefined,
   "functionId": "test-function",
   "headers": undefined,
@@ -46,7 +46,7 @@ exports[`streamText > options.experimental_onToolCallStart > should be called wi
 {
   "abortSignal": undefined,
   "callId": "test-telemetry-call-id",
-  "experimental_context": undefined,
+  "experimental_context": {},
   "functionId": undefined,
   "messages": [
     {
diff --git a/packages/ai/src/generate-text/generate-text.test.ts b/packages/ai/src/generate-text/generate-text.test.ts
index 26afabfd70a9..042631852e0d 100644
--- a/packages/ai/src/generate-text/generate-text.test.ts
+++ b/packages/ai/src/generate-text/generate-text.test.ts
@@ -2260,7 +2260,7 @@ describe('generateText', () => {
           ],
           "dynamicToolCalls": [],
           "dynamicToolResults": [],
-          "experimental_context": undefined,
+          "experimental_context": {},
           "files": [],
           "finishReason": "stop",
           "functionId": undefined,
@@ -2376,7 +2376,7 @@ describe('generateText', () => {
                   "type": "tool-result",
                 },
               ],
-              "experimental_context": undefined,
+              "experimental_context": {},
               "finishReason": "stop",
               "functionId": undefined,
               "metadata": undefined,
@@ -3605,7 +3605,7 @@ describe('generateText', () => {
                       "type": "tool-result",
                     },
                   ],
-                  "experimental_context": undefined,
+                  "experimental_context": {},
                   "finishReason": "tool-calls",
                   "functionId": undefined,
                   "metadata": undefined,
@@ -3704,7 +3704,7 @@ describe('generateText', () => {
                       "type": "tool-result",
                     },
                   ],
-                  "experimental_context": undefined,
+                  "experimental_context": {},
                   "finishReason": "tool-calls",
                   "functionId": undefined,
                   "metadata": undefined,
@@ -3951,6 +3951,7 @@ describe('generateText', () => {
           abortSignal: abortController.signal,
           toolCallId: 'call-1',
           messages: expect.any(Array),
+          experimental_context: {},
         },
       );
     });
@@ -4054,6 +4055,7 @@ describe('generateText', () => {
           abortSignal: expect.any(AbortSignal),
           toolCallId: 'call-1',
           messages: expect.any(Array),
+          experimental_context: {},
         },
       );
     });
@@ -4352,7 +4354,7 @@ describe('generateText', () => {
           {
             "options": {
               "abortSignal": undefined,
-              "experimental_context": undefined,
+              "experimental_context": {},
               "input": {
                 "value": "value",
               },
@@ -6868,7 +6870,7 @@ describe('generateText', () => {
                   "type": "tool-result",
                 },
               ],
-              "experimental_context": undefined,
+              "experimental_context": {},
               "finishReason": "tool-calls",
               "functionId": undefined,
               "metadata": undefined,
@@ -7341,7 +7343,7 @@ describe('generateText', () => {
                 "value": "value-needs-approval",
               },
               "options": {
-                "experimental_context": undefined,
+                "experimental_context": {},
                 "messages": [
                   {
                     "content": "test-input",
@@ -7356,7 +7358,7 @@ describe('generateText', () => {
                 "value": "value-no-approval",
               },
               "options": {
-                "experimental_context": undefined,
+                "experimental_context": {},
                 "messages": [
                   {
                     "content": "test-input",
diff --git a/packages/ai/src/generate-text/stream-text.test.ts b/packages/ai/src/generate-text/stream-text.test.ts
index a3ed17513aba..344e8cd28303 100644
--- a/packages/ai/src/generate-text/stream-text.test.ts
+++ b/packages/ai/src/generate-text/stream-text.test.ts
@@ -4968,7 +4968,7 @@ describe('streamText', () => {
                 "type": "text",
               },
             ],
-            "experimental_context": undefined,
+            "experimental_context": {},
             "finishReason": "stop",
             "functionId": undefined,
             "metadata": undefined,
@@ -5111,7 +5111,7 @@ describe('streamText', () => {
                 "url": "https://example.com/2",
               },
             ],
-            "experimental_context": undefined,
+            "experimental_context": {},
             "finishReason": "stop",
             "functionId": undefined,
             "metadata": undefined,
@@ -5199,7 +5199,7 @@ describe('streamText', () => {
                 "type": "file",
               },
             ],
-            "experimental_context": undefined,
+            "experimental_context": {},
             "finishReason": "stop",
             "functionId": undefined,
             "metadata": undefined,
@@ -5309,7 +5309,7 @@ describe('streamText', () => {
                 "type": "text",
               },
             ],
-            "experimental_context": undefined,
+            "experimental_context": {},
             "finishReason": "stop",
             "functionId": undefined,
             "metadata": undefined,
@@ -7121,7 +7121,7 @@ describe('streamText', () => {
           ],
           "dynamicToolCalls": [],
           "dynamicToolResults": [],
-          "experimental_context": undefined,
+          "experimental_context": {},
           "files": [],
           "finishReason": "stop",
           "functionId": undefined,
@@ -7241,7 +7241,7 @@ describe('streamText', () => {
                   "type": "tool-result",
                 },
               ],
-              "experimental_context": undefined,
+              "experimental_context": {},
               "finishReason": "stop",
               "functionId": undefined,
               "metadata": undefined,
@@ -7436,7 +7436,7 @@ describe('streamText', () => {
           ],
           "dynamicToolCalls": [],
           "dynamicToolResults": [],
-          "experimental_context": undefined,
+          "experimental_context": {},
           "files": [],
           "finishReason": "stop",
           "functionId": undefined,
@@ -7531,7 +7531,7 @@ describe('streamText', () => {
                   "url": "https://example.com/2",
                 },
               ],
-              "experimental_context": undefined,
+              "experimental_context": {},
               "finishReason": "stop",
               "functionId": undefined,
               "metadata": undefined,
@@ -7698,7 +7698,7 @@ describe('streamText', () => {
           ],
           "dynamicToolCalls": [],
           "dynamicToolResults": [],
-          "experimental_context": undefined,
+          "experimental_context": {},
           "files": [
             DefaultGeneratedFileWithType {
               "base64Data": "Hello World",
@@ -7787,7 +7787,7 @@ describe('streamText', () => {
                   "type": "file",
                 },
               ],
-              "experimental_context": undefined,
+              "experimental_context": {},
               "finishReason": "stop",
               "functionId": undefined,
               "metadata": undefined,
@@ -8428,7 +8428,7 @@ describe('streamText', () => {
               ],
               "dynamicToolCalls": [],
               "dynamicToolResults": [],
-              "experimental_context": undefined,
+              "experimental_context": {},
               "files": [],
               "finishReason": "stop",
               "functionId": undefined,
@@ -8531,7 +8531,7 @@ describe('streamText', () => {
                       "type": "tool-result",
                     },
                   ],
-                  "experimental_context": undefined,
+                  "experimental_context": {},
                   "finishReason": "tool-calls",
                   "functionId": undefined,
                   "metadata": undefined,
@@ -8615,7 +8615,7 @@ describe('streamText', () => {
                       "type": "text",
                     },
                   ],
-                  "experimental_context": undefined,
+                  "experimental_context": {},
                   "finishReason": "stop",
                   "functionId": undefined,
                   "metadata": undefined,
@@ -8775,7 +8775,7 @@ describe('streamText', () => {
                     "type": "tool-result",
                   },
                 ],
-                "experimental_context": undefined,
+                "experimental_context": {},
                 "finishReason": "tool-calls",
                 "functionId": undefined,
                 "metadata": undefined,
@@ -8859,7 +8859,7 @@ describe('streamText', () => {
                     "type": "text",
                   },
                 ],
-                "experimental_context": undefined,
+                "experimental_context": {},
                 "finishReason": "stop",
                 "functionId": undefined,
                 "metadata": undefined,
@@ -9037,7 +9037,7 @@ describe('streamText', () => {
                     "type": "tool-result",
                   },
                 ],
-                "experimental_context": undefined,
+                "experimental_context": {},
                 "finishReason": "tool-calls",
                 "functionId": undefined,
                 "metadata": undefined,
@@ -9121,7 +9121,7 @@ describe('streamText', () => {
                     "type": "text",
                   },
                 ],
-                "experimental_context": undefined,
+                "experimental_context": {},
                 "finishReason": "stop",
                 "functionId": undefined,
                 "metadata": undefined,
@@ -10332,7 +10332,7 @@ describe('streamText', () => {
               ],
               "dynamicToolCalls": [],
               "dynamicToolResults": [],
-              "experimental_context": undefined,
+              "experimental_context": {},
               "files": [],
               "finishReason": "stop",
               "functionId": undefined,
@@ -10435,7 +10435,7 @@ describe('streamText', () => {
                       "type": "tool-result",
                     },
                   ],
-                  "experimental_context": undefined,
+                  "experimental_context": {},
                   "finishReason": "tool-calls",
                   "functionId": undefined,
                   "metadata": undefined,
@@ -10519,7 +10519,7 @@ describe('streamText', () => {
                       "type": "text",
                     },
                   ],
-                  "experimental_context": undefined,
+                  "experimental_context": {},
                   "finishReason": "stop",
                   "functionId": undefined,
                   "metadata": undefined,
@@ -10679,7 +10679,7 @@ describe('streamText', () => {
                     "type": "tool-result",
                   },
                 ],
-                "experimental_context": undefined,
+                "experimental_context": {},
                 "finishReason": "tool-calls",
                 "functionId": undefined,
                 "metadata": undefined,
@@ -10763,7 +10763,7 @@ describe('streamText', () => {
                     "type": "text",
                   },
                 ],
-                "experimental_context": undefined,
+                "experimental_context": {},
                 "finishReason": "stop",
                 "functionId": undefined,
                 "metadata": undefined,
@@ -10937,7 +10937,7 @@ describe('streamText', () => {
                     "type": "tool-result",
                   },
                 ],
-                "experimental_context": undefined,
+                "experimental_context": {},
                 "finishReason": "tool-calls",
                 "functionId": undefined,
                 "metadata": undefined,
@@ -11021,7 +11021,7 @@ describe('streamText', () => {
                     "type": "text",
                   },
                 ],
-                "experimental_context": undefined,
+                "experimental_context": {},
                 "finishReason": "stop",
                 "functionId": undefined,
                 "metadata": undefined,
@@ -11518,7 +11518,7 @@ describe('streamText', () => {
                       "type": "tool-result",
                     },
                   ],
-                  "experimental_context": undefined,
+                  "experimental_context": {},
                   "finishReason": "tool-calls",
                   "functionId": undefined,
                   "metadata": undefined,
@@ -11628,7 +11628,7 @@ describe('streamText', () => {
                       "type": "tool-result",
                     },
                   ],
-                  "experimental_context": undefined,
+                  "experimental_context": {},
                   "finishReason": "tool-calls",
                   "functionId": undefined,
                   "metadata": undefined,
@@ -12824,6 +12824,7 @@ describe('streamText', () => {
         {
           abortSignal: abortController.signal,
           toolCallId: 'call-1',
+          experimental_context: {},
           messages: expect.any(Array),
         },
       );
@@ -13451,7 +13452,7 @@ describe('streamText', () => {
           {
             "options": {
               "abortSignal": undefined,
-              "experimental_context": undefined,
+              "experimental_context": {},
               "messages": [
                 {
                   "content": "test-input",
@@ -13465,7 +13466,7 @@ describe('streamText', () => {
           {
             "options": {
               "abortSignal": undefined,
-              "experimental_context": undefined,
+              "experimental_context": {},
               "inputTextDelta": "{"",
               "messages": [
                 {
@@ -13480,7 +13481,7 @@ describe('streamText', () => {
           {
             "options": {
               "abortSignal": undefined,
-              "experimental_context": undefined,
+              "experimental_context": {},
               "inputTextDelta": "value",
               "messages": [
                 {
@@ -13495,7 +13496,7 @@ describe('streamText', () => {
           {
             "options": {
               "abortSignal": undefined,
-              "experimental_context": undefined,
+              "experimental_context": {},
               "inputTextDelta": "":"",
               "messages": [
                 {
@@ -13510,7 +13511,7 @@ describe('streamText', () => {
           {
             "options": {
               "abortSignal": undefined,
-              "experimental_context": undefined,
+              "experimental_context": {},
               "inputTextDelta": "Spark",
               "messages": [
                 {
@@ -13525,7 +13526,7 @@ describe('streamText', () => {
           {
             "options": {
               "abortSignal": undefined,
-              "experimental_context": undefined,
+              "experimental_context": {},
               "inputTextDelta": "le",
               "messages": [
                 {
@@ -13540,7 +13541,7 @@ describe('streamText', () => {
           {
             "options": {
               "abortSignal": undefined,
-              "experimental_context": undefined,
+              "experimental_context": {},
               "inputTextDelta": " Day",
               "messages": [
                 {
@@ -13555,7 +13556,7 @@ describe('streamText', () => {
           {
             "options": {
               "abortSignal": undefined,
-              "experimental_context": undefined,
+              "experimental_context": {},
               "inputTextDelta": ""}",
               "messages": [
                 {
@@ -13570,7 +13571,7 @@ describe('streamText', () => {
           {
             "options": {
               "abortSignal": undefined,
-              "experimental_context": undefined,
+              "experimental_context": {},
               "input": {
                 "value": "Sparkle Day",
               },
@@ -13865,7 +13866,7 @@ describe('streamText', () => {
                 "type": "tool-error",
               },
             ],
-            "experimental_context": undefined,
+            "experimental_context": {},
             "finishReason": "stop",
             "functionId": undefined,
             "metadata": undefined,
@@ -14370,7 +14371,7 @@ describe('streamText', () => {
                   "type": "tool-result",
                 },
               ],
-              "experimental_context": undefined,
+              "experimental_context": {},
               "finishReason": "stop",
               "functionId": undefined,
               "metadata": undefined,
@@ -14604,7 +14605,7 @@ describe('streamText', () => {
             ],
             "dynamicToolCalls": [],
             "dynamicToolResults": [],
-            "experimental_context": undefined,
+            "experimental_context": {},
             "files": [],
             "finishReason": "stop",
             "functionId": undefined,
@@ -14724,7 +14725,7 @@ describe('streamText', () => {
                     "type": "tool-result",
                   },
                 ],
-                "experimental_context": undefined,
+                "experimental_context": {},
                 "finishReason": "stop",
                 "functionId": undefined,
                 "metadata": undefined,
@@ -14954,7 +14955,7 @@ describe('streamText', () => {
                 "type": "tool-result",
               },
             ],
-            "experimental_context": undefined,
+            "experimental_context": {},
             "finishReason": "stop",
             "functionId": undefined,
             "metadata": undefined,
@@ -15461,7 +15462,7 @@ describe('streamText', () => {
                 "type": "text",
               },
             ],
-            "experimental_context": undefined,
+            "experimental_context": {},
             "finishReason": "stop",
             "functionId": undefined,
             "metadata": undefined,
@@ -15906,7 +15907,7 @@ describe('streamText', () => {
             ],
             "dynamicToolCalls": [],
             "dynamicToolResults": [],
-            "experimental_context": undefined,
+            "experimental_context": {},
             "files": [],
             "finishReason": "stop",
             "functionId": undefined,
@@ -15952,7 +15953,7 @@ describe('streamText', () => {
                     "type": "text",
                   },
                 ],
-                "experimental_context": undefined,
+                "experimental_context": {},
                 "finishReason": "stop",
                 "functionId": undefined,
                 "metadata": undefined,
@@ -17048,7 +17049,7 @@ describe('streamText', () => {
                   "type": "reasoning",
                 },
               ],
-              "experimental_context": undefined,
+              "experimental_context": {},
               "finishReason": "stop",
               "functionId": undefined,
               "metadata": undefined,
@@ -17439,7 +17440,7 @@ describe('streamText', () => {
                       "type": "tool-result",
                     },
                   ],
-                  "experimental_context": undefined,
+                  "experimental_context": {},
                   "finishReason": "tool-calls",
                   "functionId": undefined,
                   "metadata": undefined,
@@ -18494,7 +18495,7 @@ describe('streamText', () => {
                   "type": "tool-result",
                 },
               ],
-              "experimental_context": undefined,
+              "experimental_context": {},
               "finishReason": "stop",
               "functionId": undefined,
               "metadata": undefined,

From 974866e63ee7e83f002e1354dd1bf11ac59dd3a4 Mon Sep 17 00:00:00 2001
From: Lars Grammel <lars.grammel@gmail.com>
Date: Wed, 1 Apr 2026 12:42:06 +0200
Subject: [PATCH 29/41] types

---
 .../app/api/chat/human-in-the-loop/utils.ts   |  3 +-
 .../ai-functions/src/lib/print-full-stream.ts |  2 +-
 .../ai-functions/src/lib/save-raw-chunks.ts   |  2 +-
 examples/mcp/src/image-content/client.ts      |  5 +-
 examples/mcp/src/output-schema/client.ts      |  6 +-
 packages/mcp/src/tool/mcp-client.test.ts      | 80 ++++++++++++++-----
 packages/mcp/src/tool/mcp-client.ts           |  4 +-
 .../openai/src/tool/local-shell.test-d.ts     |  3 +-
 packages/openai/src/tool/web-search.test-d.ts |  2 +-
 9 files changed, 74 insertions(+), 33 deletions(-)

diff --git a/examples/ai-e2e-next/app/api/chat/human-in-the-loop/utils.ts b/examples/ai-e2e-next/app/api/chat/human-in-the-loop/utils.ts
index 561ecf49a683..752f4d99216c 100644
--- a/examples/ai-e2e-next/app/api/chat/human-in-the-loop/utils.ts
+++ b/examples/ai-e2e-next/app/api/chat/human-in-the-loop/utils.ts
@@ -51,7 +51,7 @@ export async function processToolCalls<
   executeFunctions: {
     [K in keyof Tools & keyof ExecutableTools]?: (
       args: ExecutableTools[K] extends Tool<infer P> ? P : never,
-      context: ToolExecutionOptions,
+      context: ToolExecutionOptions<{}>,
     ) => Promise<any>;
   },
 ): Promise<HumanInTheLoopUIMessage[]> {
@@ -86,6 +86,7 @@ export async function processToolCalls<
           result = await toolInstance(part.input, {
             messages: await convertToModelMessages(messages),
             toolCallId: part.toolCallId,
+            experimental_context: {},
           });
         } else {
           result = 'Error: No execute function found on tool';
diff --git a/examples/ai-functions/src/lib/print-full-stream.ts b/examples/ai-functions/src/lib/print-full-stream.ts
index dc63f0b30c55..a8994ba4e493 100644
--- a/examples/ai-functions/src/lib/print-full-stream.ts
+++ b/examples/ai-functions/src/lib/print-full-stream.ts
@@ -3,7 +3,7 @@ import { StreamTextResult } from 'ai';
 export async function printFullStream({
   result,
 }: {
-  result: StreamTextResult<any, any>;
+  result: StreamTextResult<any, any, any>;
 }) {
   for await (const chunk of result.fullStream) {
     switch (chunk.type) {
diff --git a/examples/ai-functions/src/lib/save-raw-chunks.ts b/examples/ai-functions/src/lib/save-raw-chunks.ts
index 2bb0a0f8662f..57661394945f 100644
--- a/examples/ai-functions/src/lib/save-raw-chunks.ts
+++ b/examples/ai-functions/src/lib/save-raw-chunks.ts
@@ -5,7 +5,7 @@ export async function saveRawChunks({
   result,
   filename,
 }: {
-  result: StreamTextResult<any, any>;
+  result: StreamTextResult<any, any, any>;
   filename: string;
 }) {
   const rawChunks: unknown[] = [];
diff --git a/examples/mcp/src/image-content/client.ts b/examples/mcp/src/image-content/client.ts
index c7955957c56e..e36da502074e 100644
--- a/examples/mcp/src/image-content/client.ts
+++ b/examples/mcp/src/image-content/client.ts
@@ -20,7 +20,10 @@ async function main() {
     const tool = tools['get-image'];
 
     console.log('Calling get-image tool...\n');
-    const result = await tool.execute!({}, { messages: [], toolCallId: '1' });
+    const result = await tool.execute!(
+      {},
+      { messages: [], toolCallId: '1', experimental_context: {} },
+    );
 
     console.log('Raw execute() result (MCP format):');
     console.log(JSON.stringify(result, null, 2));
diff --git a/examples/mcp/src/output-schema/client.ts b/examples/mcp/src/output-schema/client.ts
index f7ec49b6a7da..9832ae124a6a 100644
--- a/examples/mcp/src/output-schema/client.ts
+++ b/examples/mcp/src/output-schema/client.ts
@@ -59,7 +59,7 @@ async function main() {
   const weatherTool = tools['get-weather'];
   const weatherResult = await weatherTool.execute(
     { location: 'New York' },
-    { messages: [], toolCallId: 'weather-1' },
+    { messages: [], toolCallId: 'weather-1', experimental_context: {} },
   );
 
   const weather = weatherResult as {
@@ -78,7 +78,7 @@ async function main() {
   const usersTool = tools['list-users'];
   const usersResult = await usersTool.execute(
     {},
-    { messages: [], toolCallId: 'users-1' },
+    { messages: [], toolCallId: 'users-1', experimental_context: {} },
   );
 
   const users = usersResult as {
@@ -96,7 +96,7 @@ async function main() {
   const echoTool = tools['echo'];
   const echoResult = await echoTool.execute(
     { message: 'Hello, MCP!' },
-    { messages: [], toolCallId: 'echo-1' },
+    { messages: [], toolCallId: 'echo-1', experimental_context: {} },
   );
 
   console.log('Raw result:', JSON.stringify(echoResult, null, 2));
diff --git a/packages/mcp/src/tool/mcp-client.test.ts b/packages/mcp/src/tool/mcp-client.test.ts
index 629d785b00cc..7747b0f72377 100644
--- a/packages/mcp/src/tool/mcp-client.test.ts
+++ b/packages/mcp/src/tool/mcp-client.test.ts
@@ -76,6 +76,7 @@ describe('MCPClient', () => {
         {
           messages: [],
           toolCallId: '1',
+          experimental_context: {},
         },
       ),
     ).toMatchInlineSnapshot(`
@@ -135,7 +136,7 @@ describe('MCPClient', () => {
     // Verify the execute function works
     const result = await tool.execute(
       { foo: 'bar' },
-      { messages: [], toolCallId: '1' },
+      { messages: [], toolCallId: '1', experimental_context: {} },
     );
     expect(result).toMatchObject({
       content: [{ type: 'text', text: 'Mock tool call result' }],
@@ -192,8 +193,12 @@ describe('MCPClient', () => {
     const tools = await client.tools();
     const tool = tools['get-image'];
 
-    expect(await tool.execute!({}, { messages: [], toolCallId: '1' }))
-      .toMatchInlineSnapshot(`
+    expect(
+      await tool.execute!(
+        {},
+        { messages: [], toolCallId: '1', experimental_context: {} },
+      ),
+    ).toMatchInlineSnapshot(`
       {
         "content": [
           {
@@ -268,8 +273,12 @@ describe('MCPClient', () => {
     const tools = await client.tools();
     const tool = tools['get-text'];
 
-    expect(await tool.execute!({}, { messages: [], toolCallId: '1' }))
-      .toMatchInlineSnapshot(`
+    expect(
+      await tool.execute!(
+        {},
+        { messages: [], toolCallId: '1', experimental_context: {} },
+      ),
+    ).toMatchInlineSnapshot(`
       {
         "content": [
           {
@@ -334,8 +343,12 @@ describe('MCPClient', () => {
     const tools = await client.tools();
     const tool = tools['get-mixed'];
 
-    expect(await tool.execute!({}, { messages: [], toolCallId: '1' }))
-      .toMatchInlineSnapshot(`
+    expect(
+      await tool.execute!(
+        {},
+        { messages: [], toolCallId: '1', experimental_context: {} },
+      ),
+    ).toMatchInlineSnapshot(`
         {
           "content": [
             {
@@ -411,8 +424,12 @@ describe('MCPClient', () => {
     const tools = await client.tools();
     const tool = tools['get-unknown'];
 
-    expect(await tool.execute!({}, { messages: [], toolCallId: '1' }))
-      .toMatchInlineSnapshot(`
+    expect(
+      await tool.execute!(
+        {},
+        { messages: [], toolCallId: '1', experimental_context: {} },
+      ),
+    ).toMatchInlineSnapshot(`
         {
           "content": [
             {
@@ -477,8 +494,12 @@ describe('MCPClient', () => {
     const tools = await client.tools();
     const tool = tools['get-raw'];
 
-    expect(await tool.execute!({}, { messages: [], toolCallId: '1' }))
-      .toMatchInlineSnapshot(`
+    expect(
+      await tool.execute!(
+        {},
+        { messages: [], toolCallId: '1', experimental_context: {} },
+      ),
+    ).toMatchInlineSnapshot(`
         {
           "isError": false,
           "toolResult": undefined,
@@ -716,6 +737,7 @@ describe('MCPClient', () => {
       {
         messages: [],
         toolCallId: '1',
+        experimental_context: {},
       },
     );
 
@@ -758,7 +780,10 @@ describe('MCPClient', () => {
     });
     const toolCall = tools['mock-tool'].execute;
     await expect(
-      toolCall({ bar: 'bar' }, { messages: [], toolCallId: '1' }),
+      toolCall(
+        { bar: 'bar' },
+        { messages: [], toolCallId: '1', experimental_context: {} },
+      ),
     ).rejects.toThrow(MCPClientError);
   });
 
@@ -782,7 +807,10 @@ describe('MCPClient', () => {
     const toolCall = tools['mock-tool'].execute;
 
     try {
-      await toolCall({ bar: 'bar' }, { messages: [], toolCallId: '1' });
+      await toolCall(
+        { bar: 'bar' },
+        { messages: [], toolCallId: '1', experimental_context: {} },
+      );
       throw new Error('Expected error to be thrown');
     } catch (error) {
       expect(MCPClientError.isInstance(error)).toBe(true);
@@ -940,6 +968,7 @@ describe('MCPClient', () => {
           messages: [],
           toolCallId: '1',
           abortSignal: abortController.signal,
+          experimental_context: {},
         },
       ),
     ).rejects.toSatisfy(
@@ -1054,6 +1083,7 @@ describe('MCPClient', () => {
       {
         messages: [],
         toolCallId: '1',
+        experimental_context: {},
       },
     );
 
@@ -1091,7 +1121,10 @@ describe('MCPClient', () => {
       },
     });
 
-    const result = await tool.execute({}, { messages: [], toolCallId: '1' });
+    const result = await tool.execute(
+      {},
+      { messages: [], toolCallId: '1', experimental_context: {} },
+    );
     expect(result).toMatchInlineSnapshot(`
       {
         "content": [
@@ -1209,7 +1242,7 @@ describe('MCPClient', () => {
 
       const result = await tool.execute(
         { location: 'New York' },
-        { messages: [], toolCallId: '1' },
+        { messages: [], toolCallId: '1', experimental_context: {} },
       );
 
       expectTypeOf<Exclude<typeof result, AsyncIterable<any>>>().toEqualTypeOf<{
@@ -1265,7 +1298,7 @@ describe('MCPClient', () => {
 
       const result = await tools['json-tool'].execute(
         {},
-        { messages: [], toolCallId: '1' },
+        { messages: [], toolCallId: '1', experimental_context: {} },
       );
 
       expect(result).toEqual({
@@ -1319,7 +1352,7 @@ describe('MCPClient', () => {
 
       const result = await tool.execute(
         { input: 'test' },
-        { messages: [], toolCallId: '1' },
+        { messages: [], toolCallId: '1', experimental_context: {} },
       );
 
       expectTypeOf<
@@ -1380,7 +1413,10 @@ describe('MCPClient', () => {
       });
 
       await expect(
-        tools['bad-output-tool'].execute({}, { messages: [], toolCallId: '1' }),
+        tools['bad-output-tool'].execute(
+          {},
+          { messages: [], toolCallId: '1', experimental_context: {} },
+        ),
       ).rejects.toThrow(MCPClientError);
     });
 
@@ -1426,7 +1462,7 @@ describe('MCPClient', () => {
       await expect(
         tools['invalid-json-tool'].execute(
           {},
-          { messages: [], toolCallId: '1' },
+          { messages: [], toolCallId: '1', experimental_context: {} },
         ),
       ).rejects.toThrow(MCPClientError);
     });
@@ -1473,7 +1509,7 @@ describe('MCPClient', () => {
       await expect(
         tools['mismatched-json-tool'].execute(
           {},
-          { messages: [], toolCallId: '1' },
+          { messages: [], toolCallId: '1', experimental_context: {} },
         ),
       ).rejects.toThrow(MCPClientError);
     });
@@ -1487,7 +1523,7 @@ describe('MCPClient', () => {
 
       const result = await tools['mock-tool'].execute(
         { foo: 'bar' },
-        { messages: [], toolCallId: '1' },
+        { messages: [], toolCallId: '1', experimental_context: {} },
       );
 
       // With automatic discovery, result is CallToolResult
@@ -1568,7 +1604,7 @@ describe('MCPClient', () => {
 
       const result = await tools['complex-tool'].execute(
         {},
-        { messages: [], toolCallId: '1' },
+        { messages: [], toolCallId: '1', experimental_context: {} },
       );
 
       expect(result).toEqual({
diff --git a/packages/mcp/src/tool/mcp-client.ts b/packages/mcp/src/tool/mcp-client.ts
index debcaf497b0c..978fb44ed9a7 100644
--- a/packages/mcp/src/tool/mcp-client.ts
+++ b/packages/mcp/src/tool/mcp-client.ts
@@ -420,7 +420,7 @@ class DefaultMCPClient implements MCPClient {
   }: {
     name: string;
     args: Record<string, unknown>;
-    options?: ToolExecutionOptions;
+    options?: ToolExecutionOptions<{}>;
   }): Promise<CallToolResult> {
     try {
       return this.request({
@@ -579,7 +579,7 @@ class DefaultMCPClient implements MCPClient {
 
       const execute = async (
         args: any,
-        options: ToolExecutionOptions,
+        options: ToolExecutionOptions<{}>,
       ): Promise<unknown> => {
         options?.abortSignal?.throwIfAborted();
         const result = await self.callTool({ name, args, options });
diff --git a/packages/openai/src/tool/local-shell.test-d.ts b/packages/openai/src/tool/local-shell.test-d.ts
index 45b198d660a4..8dc6ef798de7 100644
--- a/packages/openai/src/tool/local-shell.test-d.ts
+++ b/packages/openai/src/tool/local-shell.test-d.ts
@@ -13,7 +13,8 @@ describe('local-shell tool type', () => {
     expectTypeOf(localShellTool).toEqualTypeOf<
       Tool<
         InferSchema<typeof localShellInputSchema>,
-        InferSchema<typeof localShellOutputSchema>
+        InferSchema<typeof localShellOutputSchema>,
+        {}
       >
     >();
   });
diff --git a/packages/openai/src/tool/web-search.test-d.ts b/packages/openai/src/tool/web-search.test-d.ts
index fc78329bf87f..8c0b8e86dc62 100644
--- a/packages/openai/src/tool/web-search.test-d.ts
+++ b/packages/openai/src/tool/web-search.test-d.ts
@@ -7,7 +7,7 @@ describe('web-search tool type', () => {
     const webSearchTool = webSearch();
 
     expectTypeOf(webSearchTool).toEqualTypeOf<
-      Tool<{}, InferSchema<typeof webSearchOutputSchema>>
+      Tool<{}, InferSchema<typeof webSearchOutputSchema>, {}>
     >();
   });
 });

From c6cc59f3c901b3f871522703ea03b0b697a65ce3 Mon Sep 17 00:00:00 2001
From: Lars Grammel <lars.grammel@gmail.com>
Date: Wed, 1 Apr 2026 12:49:11 +0200
Subject: [PATCH 30/41] revert

---
 packages/ai/src/generate-text/tool-call.ts   | 4 ++--
 packages/ai/src/generate-text/tool-result.ts | 4 ++--
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/packages/ai/src/generate-text/tool-call.ts b/packages/ai/src/generate-text/tool-call.ts
index 5c5fa589b32a..1d72085b9091 100644
--- a/packages/ai/src/generate-text/tool-call.ts
+++ b/packages/ai/src/generate-text/tool-call.ts
@@ -10,7 +10,7 @@ type BaseToolCall = {
   providerMetadata?: ProviderMetadata;
 };
 
-export type StaticToolCall<TOOLS extends ToolSet = ToolSet> = ValueOf<{
+export type StaticToolCall<TOOLS extends ToolSet> = ValueOf<{
   [NAME in keyof TOOLS]: BaseToolCall & {
     toolName: NAME & string;
     input: InferToolInput<TOOLS[NAME]>;
@@ -42,6 +42,6 @@ export type DynamicToolCall = BaseToolCall & {
   error?: unknown;
 };
 
-export type TypedToolCall<TOOLS extends ToolSet = ToolSet> =
+export type TypedToolCall<TOOLS extends ToolSet> =
   | StaticToolCall<TOOLS>
   | DynamicToolCall;
diff --git a/packages/ai/src/generate-text/tool-result.ts b/packages/ai/src/generate-text/tool-result.ts
index 422f3dafb283..406deef0231e 100644
--- a/packages/ai/src/generate-text/tool-result.ts
+++ b/packages/ai/src/generate-text/tool-result.ts
@@ -3,7 +3,7 @@ import { ProviderMetadata } from '../types';
 import { ValueOf } from '../../src/util/value-of';
 import { ToolSet } from './tool-set';
 
-export type StaticToolResult<TOOLS extends ToolSet = ToolSet> = ValueOf<{
+export type StaticToolResult<TOOLS extends ToolSet> = ValueOf<{
   [NAME in keyof TOOLS]: {
     type: 'tool-result';
     toolCallId: string;
@@ -31,6 +31,6 @@ export type DynamicToolResult = {
   title?: string;
 };
 
-export type TypedToolResult<TOOLS extends ToolSet = ToolSet> =
+export type TypedToolResult<TOOLS extends ToolSet> =
   | StaticToolResult<TOOLS>
   | DynamicToolResult;

From 69fc8178d3db56bdb07fac55064d2c1fbc6efb5f Mon Sep 17 00:00:00 2001
From: Lars Grammel <lars.grammel@gmail.com>
Date: Wed, 1 Apr 2026 13:12:36 +0200
Subject: [PATCH 31/41] extract infer

---
 packages/provider-utils/src/types/context.ts  |  3 ++
 packages/provider-utils/src/types/index.ts    |  6 +--
 .../src/types/infer-tool-context.test-d.ts    | 38 ++++++++++++++++++
 .../src/types/infer-tool-context.ts           |  7 ++++
 .../src/types/infer-tool-input.test-d.ts      | 21 ++++++++++
 .../src/types/infer-tool-input.ts             |  7 ++++
 .../src/types/infer-tool-output.test-d.ts     | 40 +++++++++++++++++++
 .../src/types/infer-tool-output.ts            |  7 ++++
 packages/provider-utils/src/types/tool.ts     | 18 ---------
 9 files changed, 126 insertions(+), 21 deletions(-)
 create mode 100644 packages/provider-utils/src/types/infer-tool-context.test-d.ts
 create mode 100644 packages/provider-utils/src/types/infer-tool-context.ts
 create mode 100644 packages/provider-utils/src/types/infer-tool-input.test-d.ts
 create mode 100644 packages/provider-utils/src/types/infer-tool-input.ts
 create mode 100644 packages/provider-utils/src/types/infer-tool-output.test-d.ts
 create mode 100644 packages/provider-utils/src/types/infer-tool-output.ts

diff --git a/packages/provider-utils/src/types/context.ts b/packages/provider-utils/src/types/context.ts
index 9a1d29a311fd..d595a8e7d3a3 100644
--- a/packages/provider-utils/src/types/context.ts
+++ b/packages/provider-utils/src/types/context.ts
@@ -1 +1,4 @@
+/**
+ * A context object that is passed into tool execution.
+ */
 export type Context = Record<string, unknown>;
diff --git a/packages/provider-utils/src/types/index.ts b/packages/provider-utils/src/types/index.ts
index ddda36df5ff7..e896a6e14521 100644
--- a/packages/provider-utils/src/types/index.ts
+++ b/packages/provider-utils/src/types/index.ts
@@ -16,15 +16,15 @@ export type {
 export type { Context } from './context';
 export type { DataContent } from './data-content';
 export { executeTool } from './execute-tool';
+export type { InferToolContext } from './infer-tool-context';
+export type { InferToolInput } from './infer-tool-input';
+export type { InferToolOutput } from './infer-tool-output';
 export type { ModelMessage } from './model-message';
 export type { ProviderOptions } from './provider-options';
 export type { SystemModelMessage } from './system-model-message';
 export {
   dynamicTool,
   tool,
-  type InferToolContext,
-  type InferToolInput,
-  type InferToolOutput,
   type Tool,
   type ToolExecuteFunction,
   type ToolExecutionOptions,
diff --git a/packages/provider-utils/src/types/infer-tool-context.test-d.ts b/packages/provider-utils/src/types/infer-tool-context.test-d.ts
new file mode 100644
index 000000000000..e32503d7328e
--- /dev/null
+++ b/packages/provider-utils/src/types/infer-tool-context.test-d.ts
@@ -0,0 +1,38 @@
+import { describe, expectTypeOf, it } from 'vitest';
+import { z } from 'zod/v4';
+import type { InferToolContext } from './infer-tool-context';
+import { tool } from './tool';
+import { Context } from './context';
+
+describe('InferToolContext', () => {
+  it('infers the context type from a tool with contextSchema', () => {
+    const weatherTool = tool({
+      inputSchema: z.object({
+        city: z.string(),
+      }),
+      contextSchema: z.object({
+        userId: z.string(),
+        role: z.string(),
+      }),
+      execute: async () => ({ temperature: 72 }),
+    });
+
+    expectTypeOf<InferToolContext<typeof weatherTool>>().toEqualTypeOf<{
+      userId: string;
+      role: string;
+    }>();
+  });
+
+  it('infers the generic context type from a tool without contextSchema', () => {
+    const weatherTool = tool({
+      inputSchema: z.object({
+        city: z.string(),
+      }),
+      execute: async () => ({ temperature: 72 }),
+    });
+
+    expectTypeOf<
+      InferToolContext<typeof weatherTool>
+    >().toEqualTypeOf<Context>();
+  });
+});
diff --git a/packages/provider-utils/src/types/infer-tool-context.ts b/packages/provider-utils/src/types/infer-tool-context.ts
new file mode 100644
index 000000000000..294074e3f096
--- /dev/null
+++ b/packages/provider-utils/src/types/infer-tool-context.ts
@@ -0,0 +1,7 @@
+import type { Tool } from './tool';
+
+/**
+ * Infer the context type of a tool.
+ */
+export type InferToolContext<TOOL extends Tool<any, any, any>> =
+  TOOL extends Tool<any, any, infer CONTEXT> ? CONTEXT : never;
diff --git a/packages/provider-utils/src/types/infer-tool-input.test-d.ts b/packages/provider-utils/src/types/infer-tool-input.test-d.ts
new file mode 100644
index 000000000000..0045b7ac8779
--- /dev/null
+++ b/packages/provider-utils/src/types/infer-tool-input.test-d.ts
@@ -0,0 +1,21 @@
+import { describe, expectTypeOf, it } from 'vitest';
+import { z } from 'zod/v4';
+import type { InferToolInput } from './infer-tool-input';
+import { tool } from './tool';
+
+describe('InferToolInput', () => {
+  it('infers the input type from a tool with inputSchema', () => {
+    const weatherTool = tool({
+      inputSchema: z.object({
+        city: z.string(),
+        countryCode: z.string().length(2),
+      }),
+      execute: async () => ({ temperature: 72 }),
+    });
+
+    expectTypeOf<InferToolInput<typeof weatherTool>>().toEqualTypeOf<{
+      city: string;
+      countryCode: string;
+    }>();
+  });
+});
diff --git a/packages/provider-utils/src/types/infer-tool-input.ts b/packages/provider-utils/src/types/infer-tool-input.ts
new file mode 100644
index 000000000000..8ae9c8678f8f
--- /dev/null
+++ b/packages/provider-utils/src/types/infer-tool-input.ts
@@ -0,0 +1,7 @@
+import type { Tool } from './tool';
+
+/**
+ * Infer the input type of a tool.
+ */
+export type InferToolInput<TOOL extends Tool<any, any, any>> =
+  TOOL extends Tool<infer INPUT, any, any> ? INPUT : never;
diff --git a/packages/provider-utils/src/types/infer-tool-output.test-d.ts b/packages/provider-utils/src/types/infer-tool-output.test-d.ts
new file mode 100644
index 000000000000..10c2a08c1ecd
--- /dev/null
+++ b/packages/provider-utils/src/types/infer-tool-output.test-d.ts
@@ -0,0 +1,40 @@
+import { describe, expectTypeOf, it } from 'vitest';
+import { z } from 'zod/v4';
+import type { InferToolOutput } from './infer-tool-output';
+import { tool } from './tool';
+
+describe('InferToolOutput', () => {
+  it('infers the output type from a tool with execute function', () => {
+    const weatherTool = tool({
+      inputSchema: z.object({
+        city: z.string(),
+      }),
+      execute: async () => ({
+        temperature: 72,
+        conditions: 'sunny' as const,
+      }),
+    });
+
+    expectTypeOf<InferToolOutput<typeof weatherTool>>().toEqualTypeOf<{
+      temperature: number;
+      conditions: 'sunny';
+    }>();
+  });
+
+  it('infers the output type from a tool with outputSchema', () => {
+    const weatherTool = tool({
+      inputSchema: z.object({
+        city: z.string(),
+      }),
+      outputSchema: z.object({
+        temperature: z.number(),
+        conditions: z.literal('sunny'),
+      }),
+    });
+
+    expectTypeOf<InferToolOutput<typeof weatherTool>>().toEqualTypeOf<{
+      temperature: number;
+      conditions: 'sunny';
+    }>();
+  });
+});
diff --git a/packages/provider-utils/src/types/infer-tool-output.ts b/packages/provider-utils/src/types/infer-tool-output.ts
new file mode 100644
index 000000000000..6401e1e9407d
--- /dev/null
+++ b/packages/provider-utils/src/types/infer-tool-output.ts
@@ -0,0 +1,7 @@
+import type { Tool } from './tool';
+
+/**
+ * Infer the output type of a tool.
+ */
+export type InferToolOutput<TOOL extends Tool<any, any, any>> =
+  TOOL extends Tool<any, infer OUTPUT, any> ? OUTPUT : never;
diff --git a/packages/provider-utils/src/types/tool.ts b/packages/provider-utils/src/types/tool.ts
index 663aa0ecd098..ebbe371ffe43 100644
--- a/packages/provider-utils/src/types/tool.ts
+++ b/packages/provider-utils/src/types/tool.ts
@@ -272,24 +272,6 @@ export type Tool<
       }
   );
 
-/**
- * Infer the input type of a tool.
- */
-export type InferToolInput<TOOL extends Tool> =
-  TOOL extends Tool<infer INPUT, any, any> ? INPUT : never;
-
-/**
- * Infer the output type of a tool.
- */
-export type InferToolOutput<TOOL extends Tool> =
-  TOOL extends Tool<any, infer OUTPUT, any> ? OUTPUT : never;
-
-/**
- * Infer the context type of a tool.
- */
-export type InferToolContext<TOOL extends Tool> =
-  TOOL extends Tool<any, any, infer CONTEXT> ? CONTEXT : never;
-
 /**
  * Helper function for inferring the execute args of a tool.
  */

From 9d2db4861f14fc5ec0fc3b7cc5def2a0e4cc5d29 Mon Sep 17 00:00:00 2001
From: Lars Grammel <lars.grammel@gmail.com>
Date: Wed, 1 Apr 2026 14:57:32 +0200
Subject: [PATCH 32/41] type test

---
 .../src/generate-text/generate-text.test-d.ts | 44 +++++++++++++++++++
 1 file changed, 44 insertions(+)

diff --git a/packages/ai/src/generate-text/generate-text.test-d.ts b/packages/ai/src/generate-text/generate-text.test-d.ts
index fd7ccd47fabd..d1fb9110da4e 100644
--- a/packages/ai/src/generate-text/generate-text.test-d.ts
+++ b/packages/ai/src/generate-text/generate-text.test-d.ts
@@ -1,4 +1,5 @@
 import { JSONValue } from '@ai-sdk/provider';
+import { tool } from '@ai-sdk/provider-utils';
 import { describe, expectTypeOf, it } from 'vitest';
 import { z } from 'zod';
 import { generateText, Output } from '../generate-text';
@@ -65,4 +66,47 @@ describe('generateText types', () => {
       expectTypeOf<typeof result.output>().toEqualTypeOf<JSONValue>();
     });
   });
+
+  describe('experimental_context', () => {
+    it('should infer typed experimental_context with one tool context and prepareStep', async () => {
+      generateText({
+        model: new MockLanguageModelV4(),
+        prompt: 'Hello, world!',
+        tools: {
+          weather: tool({
+            inputSchema: z.object({
+              city: z.string(),
+            }),
+            contextSchema: z.object({
+              userId: z.string(),
+            }),
+            execute: async (_input, { experimental_context }) => {
+              expectTypeOf(experimental_context).toMatchTypeOf<{
+                userId: string;
+              }>();
+
+              return 'sunny';
+            },
+          }),
+        },
+        experimental_context: {
+          userId: 'test-user',
+          role: 'admin',
+        },
+        prepareStep: ({ experimental_context }) => {
+          expectTypeOf(experimental_context).toMatchTypeOf<{
+            userId: string;
+            role: string;
+          }>();
+
+          return {
+            experimental_context: {
+              userId: experimental_context.userId,
+              role: experimental_context.role,
+            },
+          };
+        },
+      });
+    });
+  });
 });

From fef307ae7bb346e464f1ed8d91d5f8e0e1cae27f Mon Sep 17 00:00:00 2001
From: Lars Grammel <lars.grammel@gmail.com>
Date: Wed, 1 Apr 2026 15:02:19 +0200
Subject: [PATCH 33/41] cs

---
 .changeset/bright-glasses-happen.md | 6 ++++++
 1 file changed, 6 insertions(+)
 create mode 100644 .changeset/bright-glasses-happen.md

diff --git a/.changeset/bright-glasses-happen.md b/.changeset/bright-glasses-happen.md
new file mode 100644
index 000000000000..41930e4f56e5
--- /dev/null
+++ b/.changeset/bright-glasses-happen.md
@@ -0,0 +1,6 @@
+---
+"@ai-sdk/provider-utils": major
+"ai": major
+---
+
+feat(ai): change type of experimental_context from unknown to generic

From 3a4dc0aab3b4625f020044cdc83858164883f462 Mon Sep 17 00:00:00 2001
From: Lars Grammel <lars.grammel@gmail.com>
Date: Wed, 1 Apr 2026 15:05:05 +0200
Subject: [PATCH 34/41] stream example

---
 .../openai/tool-call-with-context.ts          | 68 +++++++++++++++++++
 1 file changed, 68 insertions(+)
 create mode 100644 examples/ai-functions/src/stream-text/openai/tool-call-with-context.ts

diff --git a/examples/ai-functions/src/stream-text/openai/tool-call-with-context.ts b/examples/ai-functions/src/stream-text/openai/tool-call-with-context.ts
new file mode 100644
index 000000000000..41fe8dca9508
--- /dev/null
+++ b/examples/ai-functions/src/stream-text/openai/tool-call-with-context.ts
@@ -0,0 +1,68 @@
+import { openai } from '@ai-sdk/openai';
+import { streamText, tool } from 'ai';
+import { z } from 'zod';
+import { run } from '../../lib/run';
+import { printFullStream } from '../../lib/print-full-stream';
+
+run(async () => {
+  const result = streamText({
+    model: openai('gpt-5-mini'),
+    tools: {
+      weather: tool({
+        description: 'Get the weather in a location',
+        inputSchema: z.object({
+          location: z.string().describe('The location to get the weather for'),
+        }),
+        contextSchema: z.object({
+          weatherApiKey: z.string().describe('The API key for the weather API'),
+        }),
+        execute: async (
+          { location },
+          { experimental_context: { weatherApiKey } },
+        ) => {
+          console.log('weather tool api key:', weatherApiKey);
+
+          return {
+            location,
+            temperature: 72 + Math.floor(Math.random() * 21) - 10,
+          };
+        },
+      }),
+      calculator: tool({
+        description: 'Calculate mathematical expressions',
+        inputSchema: z.object({
+          expression: z
+            .string()
+            .describe('The mathematical expression to calculate'),
+        }),
+        contextSchema: z.object({
+          calculatorApiKey: z
+            .string()
+            .describe('The API key for the calculator API'),
+        }),
+        execute: async (
+          { expression },
+          { experimental_context: { calculatorApiKey } },
+        ) => {
+          console.log('calculator tool api key:', calculatorApiKey);
+          return {
+            expression,
+            result: eval(expression),
+          };
+        },
+      }),
+    },
+    experimental_context: {
+      weatherApiKey: 'weather-123',
+      calculatorApiKey: 'calculator-456',
+      somethingElse: 'other-context',
+    },
+    prepareStep: async ({ experimental_context: context }) => {
+      console.log('prepareStep context:', context);
+      return {};
+    },
+    prompt: 'What is the weather in San Francisco?',
+  });
+
+  await printFullStream({ result });
+});

From 91b1a2f972d1134f81d28e10aa5e3702fc2ecdbf Mon Sep 17 00:00:00 2001
From: Lars Grammel <lars.grammel@gmail.com>
Date: Wed, 1 Apr 2026 15:15:02 +0200
Subject: [PATCH 35/41] Add tests for UnionToIntersection utility type

---
 .../src/util/union-to-intersection.test-d.ts  | 29 +++++++++++++++++++
 1 file changed, 29 insertions(+)
 create mode 100644 packages/ai/src/util/union-to-intersection.test-d.ts

diff --git a/packages/ai/src/util/union-to-intersection.test-d.ts b/packages/ai/src/util/union-to-intersection.test-d.ts
new file mode 100644
index 000000000000..f2fde5a0fd41
--- /dev/null
+++ b/packages/ai/src/util/union-to-intersection.test-d.ts
@@ -0,0 +1,29 @@
+import { describe, expectTypeOf, it } from 'vitest';
+import type { UnionToIntersection } from './union-to-intersection';
+
+describe('UnionToIntersection', () => {
+  it('returns never when given no input', () => {
+    type Result = UnionToIntersection<never>;
+
+    expectTypeOf<Result>().toEqualTypeOf<never>();
+  });
+
+  it('returns the same type for a single input', () => {
+    type Result = UnionToIntersection<{ city: string }>;
+
+    expectTypeOf<Result>().toEqualTypeOf<{
+      city: string;
+    }>();
+  });
+
+  it('converts a union of object types into an intersection', () => {
+    type Result = UnionToIntersection<
+      { city: string } | { countryCode: string }
+    >;
+
+    expectTypeOf<Result>().toEqualTypeOf<{
+      city: string;
+      countryCode: string;
+    }>();
+  });
+});

From 6e56bd9b5b8eb9b5afc95dc526bfe1fdb9430f6e Mon Sep 17 00:00:00 2001
From: Lars Grammel <lars.grammel@gmail.com>
Date: Wed, 1 Apr 2026 15:30:08 +0200
Subject: [PATCH 36/41] Add InferToolSetContext type and corresponding tests

---
 .../infer-tool-set-context.test-d.ts          | 65 +++++++++++++++++++
 .../generate-text/infer-tool-set-context.ts   | 12 ++++
 packages/ai/src/generate-text/tool-set.ts     | 10 +--
 3 files changed, 79 insertions(+), 8 deletions(-)
 create mode 100644 packages/ai/src/generate-text/infer-tool-set-context.test-d.ts
 create mode 100644 packages/ai/src/generate-text/infer-tool-set-context.ts

diff --git a/packages/ai/src/generate-text/infer-tool-set-context.test-d.ts b/packages/ai/src/generate-text/infer-tool-set-context.test-d.ts
new file mode 100644
index 000000000000..58d92f6dba55
--- /dev/null
+++ b/packages/ai/src/generate-text/infer-tool-set-context.test-d.ts
@@ -0,0 +1,65 @@
+import { Context, tool } from '@ai-sdk/provider-utils';
+import { describe, expectTypeOf, it } from 'vitest';
+import { z } from 'zod/v4';
+import type { InferToolSetContext } from './infer-tool-set-context';
+
+describe('InferToolSetContext', () => {
+  it('infers the intersection of context types across a tool set', () => {
+    const tools = {
+      weather: tool({
+        inputSchema: z.object({
+          city: z.string(),
+        }),
+        contextSchema: z.object({
+          userId: z.string(),
+        }),
+      }),
+      forecast: tool({
+        inputSchema: z.object({
+          days: z.number(),
+        }),
+        contextSchema: z.object({
+          role: z.string(),
+        }),
+      }),
+    };
+
+    expectTypeOf<InferToolSetContext<typeof tools>>().toMatchObjectType<{
+      userId: string;
+      role: string;
+    }>();
+  });
+
+  it('infers a single tool context type from a tool set', () => {
+    const tools = {
+      weather: tool({
+        inputSchema: z.object({
+          city: z.string(),
+        }),
+        contextSchema: z.object({
+          userId: z.string(),
+          role: z.string(),
+        }),
+      }),
+    };
+
+    expectTypeOf<InferToolSetContext<typeof tools>>().toMatchObjectType<{
+      userId: string;
+      role: string;
+    }>();
+  });
+
+  it('falls back to the generic context type for tools without contextSchema', () => {
+    const tools = {
+      weather: tool({
+        inputSchema: z.object({
+          city: z.string(),
+        }),
+      }),
+    };
+
+    expectTypeOf<
+      InferToolSetContext<typeof tools>
+    >().toMatchObjectType<Context>();
+  });
+});
diff --git a/packages/ai/src/generate-text/infer-tool-set-context.ts b/packages/ai/src/generate-text/infer-tool-set-context.ts
new file mode 100644
index 000000000000..8c549e2d8afc
--- /dev/null
+++ b/packages/ai/src/generate-text/infer-tool-set-context.ts
@@ -0,0 +1,12 @@
+import { InferToolContext } from '@ai-sdk/provider-utils';
+import { UnionToIntersection } from '../util/union-to-intersection';
+import type { ToolSet } from './tool-set';
+
+/**
+ * Infer the context type of a tool set.
+ */
+export type InferToolSetContext<TOOLS extends ToolSet> = UnionToIntersection<
+  {
+    [K in keyof TOOLS]: InferToolContext<NoInfer<TOOLS[K]>>;
+  }[keyof TOOLS]
+>;
diff --git a/packages/ai/src/generate-text/tool-set.ts b/packages/ai/src/generate-text/tool-set.ts
index 104c96b30162..fa2fe650ad7d 100644
--- a/packages/ai/src/generate-text/tool-set.ts
+++ b/packages/ai/src/generate-text/tool-set.ts
@@ -1,5 +1,5 @@
-import { Context, InferToolContext, Tool } from '@ai-sdk/provider-utils';
-import { UnionToIntersection } from '../util/union-to-intersection';
+import { Context, Tool } from '@ai-sdk/provider-utils';
+import type { InferToolSetContext } from './infer-tool-set-context';
 
 export type ToolSet = Record<
   string,
@@ -19,12 +19,6 @@ export type ToolSet = Record<
     >
 >;
 
-export type InferToolSetContext<TOOLS extends ToolSet> = UnionToIntersection<
-  {
-    [K in keyof TOOLS]: InferToolContext<NoInfer<TOOLS[K]>>;
-  }[keyof TOOLS]
->;
-
 export type ExpandedContext<TOOLS extends ToolSet> = InferToolSetContext<
   NoInfer<TOOLS>
 > &

From 66a0c37ddc081652b0327eeb725925b8f960f095 Mon Sep 17 00:00:00 2001
From: Lars Grammel <lars.grammel@gmail.com>
Date: Wed, 1 Apr 2026 15:36:12 +0200
Subject: [PATCH 37/41] 1

---
 packages/ai/src/generate-text/generate-text.ts | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/packages/ai/src/generate-text/generate-text.ts b/packages/ai/src/generate-text/generate-text.ts
index 126462bf2abf..80d47d93223c 100644
--- a/packages/ai/src/generate-text/generate-text.ts
+++ b/packages/ai/src/generate-text/generate-text.ts
@@ -1195,7 +1195,7 @@ class DefaultGenerateTextResult<
   private readonly _output: InferCompleteOutput<OUTPUT> | undefined;
 
   constructor(options: {
-    steps: StepResult<TOOLS, CONTEXT>[];
+    steps: GenerateTextResult<TOOLS, CONTEXT, OUTPUT>['steps'];
     output: InferCompleteOutput<OUTPUT> | undefined;
     totalUsage: LanguageModelUsage;
   }) {

From 251ac1b8648c0e355d1b2f15049ab9ed27fa2b77 Mon Sep 17 00:00:00 2001
From: Lars Grammel <lars.grammel@gmail.com>
Date: Wed, 1 Apr 2026 15:43:17 +0200
Subject: [PATCH 38/41] Refactor context types: Replace ExpandedContext with
 GenerationContext across multiple files for improved type handling in agent
 and generation processes.

---
 packages/ai/src/agent/agent.ts                  |  7 ++++---
 .../agent/create-agent-ui-stream-response.ts    |  5 +++--
 packages/ai/src/agent/create-agent-ui-stream.ts |  5 +++--
 .../agent/pipe-agent-ui-stream-to-response.ts   |  5 +++--
 .../ai/src/agent/tool-loop-agent-settings.ts    | 13 +++++++------
 packages/ai/src/agent/tool-loop-agent.ts        |  5 +++--
 packages/ai/src/generate-text/core-events.ts    | 11 ++++++-----
 .../create-execute-tools-transformation.ts      |  5 +++--
 .../ai/src/generate-text/execute-tool-call.ts   |  5 +++--
 .../src/generate-text/generate-text-result.ts   |  5 +++--
 packages/ai/src/generate-text/generate-text.ts  | 17 +++++++++--------
 .../ai/src/generate-text/generation-context.ts  | 14 ++++++++++++++
 packages/ai/src/generate-text/prepare-step.ts   |  7 ++++---
 packages/ai/src/generate-text/step-result.ts    |  7 ++++---
 packages/ai/src/generate-text/stop-condition.ts |  7 ++++---
 .../ai/src/generate-text/stream-text-result.ts  |  5 +++--
 packages/ai/src/generate-text/stream-text.ts    | 17 +++++++++--------
 packages/ai/src/generate-text/tool-set.ts       |  8 +-------
 .../src/telemetry/open-telemetry-integration.ts |  5 +++--
 packages/ai/src/ui/direct-chat-transport.ts     |  7 ++++---
 20 files changed, 93 insertions(+), 67 deletions(-)
 create mode 100644 packages/ai/src/generate-text/generation-context.ts

diff --git a/packages/ai/src/agent/agent.ts b/packages/ai/src/agent/agent.ts
index ceced8c26a83..714496e2a259 100644
--- a/packages/ai/src/agent/agent.ts
+++ b/packages/ai/src/agent/agent.ts
@@ -3,7 +3,8 @@ import { GenerateTextResult } from '../generate-text/generate-text-result';
 import { Output } from '../generate-text/output';
 import { StreamTextTransform } from '../generate-text/stream-text';
 import { StreamTextResult } from '../generate-text/stream-text-result';
-import { ExpandedContext, ToolSet } from '../generate-text/tool-set';
+import type { GenerationContext } from '../generate-text/generation-context';
+import type { ToolSet } from '../generate-text/tool-set';
 import { TimeoutConfiguration } from '../prompt/call-settings';
 import type {
   ToolLoopAgentOnFinishCallback,
@@ -20,7 +21,7 @@ import type {
 export type AgentCallParameters<
   CALL_OPTIONS,
   TOOLS extends ToolSet = {},
-  CONTEXT extends ExpandedContext<TOOLS> = ExpandedContext<TOOLS>,
+  CONTEXT extends GenerationContext<TOOLS> = GenerationContext<TOOLS>,
 > = ([CALL_OPTIONS] extends [never]
   ? { options?: never }
   : { options: CALL_OPTIONS }) &
@@ -124,7 +125,7 @@ export type AgentStreamParameters<
 export interface Agent<
   CALL_OPTIONS = never,
   TOOLS extends ToolSet = {},
-  CONTEXT extends ExpandedContext<TOOLS> = ExpandedContext<TOOLS>,
+  CONTEXT extends GenerationContext<TOOLS> = GenerationContext<TOOLS>,
   OUTPUT extends Output = never,
 > {
   /**
diff --git a/packages/ai/src/agent/create-agent-ui-stream-response.ts b/packages/ai/src/agent/create-agent-ui-stream-response.ts
index 6ee6581c78f6..957bb104111f 100644
--- a/packages/ai/src/agent/create-agent-ui-stream-response.ts
+++ b/packages/ai/src/agent/create-agent-ui-stream-response.ts
@@ -1,6 +1,7 @@
 import { StreamTextTransform, UIMessageStreamOptions } from '../generate-text';
 import { Output } from '../generate-text/output';
-import { ExpandedContext, ToolSet } from '../generate-text/tool-set';
+import type { GenerationContext } from '../generate-text/generation-context';
+import type { ToolSet } from '../generate-text/tool-set';
 import { TimeoutConfiguration } from '../prompt/call-settings';
 import { createUIMessageStreamResponse } from '../ui-message-stream';
 import { UIMessageStreamResponseInit } from '../ui-message-stream/ui-message-stream-response-init';
@@ -29,7 +30,7 @@ import type { ToolLoopAgentOnStepFinishCallback } from './tool-loop-agent-settin
 export async function createAgentUIStreamResponse<
   CALL_OPTIONS = never,
   TOOLS extends ToolSet = {},
-  CONTEXT extends ExpandedContext<TOOLS> = ExpandedContext<TOOLS>,
+  CONTEXT extends GenerationContext<TOOLS> = GenerationContext<TOOLS>,
   OUTPUT extends Output<TOOLS, CONTEXT> = never,
   MESSAGE_METADATA = unknown,
 >({
diff --git a/packages/ai/src/agent/create-agent-ui-stream.ts b/packages/ai/src/agent/create-agent-ui-stream.ts
index 39e9d4d39f1d..7ceb950a99ec 100644
--- a/packages/ai/src/agent/create-agent-ui-stream.ts
+++ b/packages/ai/src/agent/create-agent-ui-stream.ts
@@ -1,6 +1,7 @@
 import { StreamTextTransform, UIMessageStreamOptions } from '../generate-text';
 import { Output } from '../generate-text/output';
-import { ExpandedContext, ToolSet } from '../generate-text/tool-set';
+import type { GenerationContext } from '../generate-text/generation-context';
+import type { ToolSet } from '../generate-text/tool-set';
 import { TimeoutConfiguration } from '../prompt/call-settings';
 import { InferUIMessageChunk } from '../ui-message-stream';
 import { convertToModelMessages } from '../ui/convert-to-model-messages';
@@ -26,7 +27,7 @@ import type { ToolLoopAgentOnStepFinishCallback } from './tool-loop-agent-settin
 export async function createAgentUIStream<
   CALL_OPTIONS = never,
   TOOLS extends ToolSet = {},
-  CONTEXT extends ExpandedContext<TOOLS> = ExpandedContext<TOOLS>,
+  CONTEXT extends GenerationContext<TOOLS> = GenerationContext<TOOLS>,
   OUTPUT extends Output = never,
   MESSAGE_METADATA = unknown,
 >({
diff --git a/packages/ai/src/agent/pipe-agent-ui-stream-to-response.ts b/packages/ai/src/agent/pipe-agent-ui-stream-to-response.ts
index 655b6e0fe1b0..942c0871fbdc 100644
--- a/packages/ai/src/agent/pipe-agent-ui-stream-to-response.ts
+++ b/packages/ai/src/agent/pipe-agent-ui-stream-to-response.ts
@@ -1,7 +1,8 @@
 import { ServerResponse } from 'node:http';
 import { StreamTextTransform, UIMessageStreamOptions } from '../generate-text';
 import { Output } from '../generate-text/output';
-import { ExpandedContext, ToolSet } from '../generate-text/tool-set';
+import type { GenerationContext } from '../generate-text/generation-context';
+import type { ToolSet } from '../generate-text/tool-set';
 import { TimeoutConfiguration } from '../prompt/call-settings';
 import { pipeUIMessageStreamToResponse } from '../ui-message-stream';
 import { UIMessageStreamResponseInit } from '../ui-message-stream/ui-message-stream-response-init';
@@ -29,7 +30,7 @@ import type { ToolLoopAgentOnStepFinishCallback } from './tool-loop-agent-settin
 export async function pipeAgentUIStreamToResponse<
   CALL_OPTIONS = never,
   TOOLS extends ToolSet = {},
-  CONTEXT extends ExpandedContext<TOOLS> = ExpandedContext<TOOLS>,
+  CONTEXT extends GenerationContext<TOOLS> = GenerationContext<TOOLS>,
   OUTPUT extends Output = never,
   MESSAGE_METADATA = unknown,
 >({
diff --git a/packages/ai/src/agent/tool-loop-agent-settings.ts b/packages/ai/src/agent/tool-loop-agent-settings.ts
index dcba514ba66b..5f6cf1cab0db 100644
--- a/packages/ai/src/agent/tool-loop-agent-settings.ts
+++ b/packages/ai/src/agent/tool-loop-agent-settings.ts
@@ -16,7 +16,8 @@ import { Output } from '../generate-text/output';
 import { PrepareStepFunction } from '../generate-text/prepare-step';
 import { StopCondition } from '../generate-text/stop-condition';
 import { ToolCallRepairFunction } from '../generate-text/tool-call-repair-function';
-import { ExpandedContext, ToolSet } from '../generate-text/tool-set';
+import type { GenerationContext } from '../generate-text/generation-context';
+import type { ToolSet } from '../generate-text/tool-set';
 import { CallSettings, TimeoutConfiguration } from '../prompt/call-settings';
 import { Prompt } from '../prompt/prompt';
 import { TelemetrySettings } from '../telemetry/telemetry-settings';
@@ -26,13 +27,13 @@ import { AgentCallParameters } from './agent';
 
 export type ToolLoopAgentOnStartCallback<
   TOOLS extends ToolSet = ToolSet,
-  CONTEXT extends ExpandedContext<TOOLS> = ExpandedContext<TOOLS>,
+  CONTEXT extends GenerationContext<TOOLS> = GenerationContext<TOOLS>,
   OUTPUT extends Output = Output,
 > = (event: OnStartEvent<TOOLS, CONTEXT, OUTPUT>) => PromiseLike<void> | void;
 
 export type ToolLoopAgentOnStepStartCallback<
   TOOLS extends ToolSet = ToolSet,
-  CONTEXT extends ExpandedContext<TOOLS> = ExpandedContext<TOOLS>,
+  CONTEXT extends GenerationContext<TOOLS> = GenerationContext<TOOLS>,
   OUTPUT extends Output = Output,
 > = (
   event: OnStepStartEvent<TOOLS, CONTEXT, OUTPUT>,
@@ -48,12 +49,12 @@ export type ToolLoopAgentOnToolCallFinishCallback<
 
 export type ToolLoopAgentOnStepFinishCallback<
   TOOLS extends ToolSet = ToolSet,
-  CONTEXT extends ExpandedContext<TOOLS> = ExpandedContext<TOOLS>,
+  CONTEXT extends GenerationContext<TOOLS> = GenerationContext<TOOLS>,
 > = (stepResult: OnStepFinishEvent<TOOLS, CONTEXT>) => Promise<void> | void;
 
 export type ToolLoopAgentOnFinishCallback<
   TOOLS extends ToolSet = ToolSet,
-  CONTEXT extends ExpandedContext<TOOLS> = ExpandedContext<TOOLS>,
+  CONTEXT extends GenerationContext<TOOLS> = GenerationContext<TOOLS>,
 > = (event: OnFinishEvent<TOOLS, CONTEXT>) => PromiseLike<void> | void;
 
 /**
@@ -62,7 +63,7 @@ export type ToolLoopAgentOnFinishCallback<
 export type ToolLoopAgentSettings<
   CALL_OPTIONS = never,
   TOOLS extends ToolSet = {},
-  CONTEXT extends ExpandedContext<TOOLS> = ExpandedContext<TOOLS>,
+  CONTEXT extends GenerationContext<TOOLS> = GenerationContext<TOOLS>,
   OUTPUT extends Output = never,
 > = Omit<CallSettings, 'abortSignal'> & {
   /**
diff --git a/packages/ai/src/agent/tool-loop-agent.ts b/packages/ai/src/agent/tool-loop-agent.ts
index 92ff75573f73..1c6b33569f1e 100644
--- a/packages/ai/src/agent/tool-loop-agent.ts
+++ b/packages/ai/src/agent/tool-loop-agent.ts
@@ -4,7 +4,8 @@ import { Output } from '../generate-text/output';
 import { isStepCount } from '../generate-text/stop-condition';
 import { streamText } from '../generate-text/stream-text';
 import { StreamTextResult } from '../generate-text/stream-text-result';
-import { ExpandedContext, ToolSet } from '../generate-text/tool-set';
+import type { GenerationContext } from '../generate-text/generation-context';
+import type { ToolSet } from '../generate-text/tool-set';
 import { Prompt } from '../prompt';
 import { Agent, AgentCallParameters, AgentStreamParameters } from './agent';
 import {
@@ -27,7 +28,7 @@ import {
 export class ToolLoopAgent<
   CALL_OPTIONS = never,
   TOOLS extends ToolSet = {},
-  CONTEXT extends ExpandedContext<TOOLS> = ExpandedContext<TOOLS>,
+  CONTEXT extends GenerationContext<TOOLS> = GenerationContext<TOOLS>,
   OUTPUT extends Output = never,
 > implements Agent<CALL_OPTIONS, TOOLS, CONTEXT, OUTPUT> {
   readonly version = 'agent-v1';
diff --git a/packages/ai/src/generate-text/core-events.ts b/packages/ai/src/generate-text/core-events.ts
index b4545c1d5029..cd30610e6beb 100644
--- a/packages/ai/src/generate-text/core-events.ts
+++ b/packages/ai/src/generate-text/core-events.ts
@@ -12,7 +12,8 @@ import type { StepResult } from './step-result';
 import type { StopCondition } from './stop-condition';
 import { TextStreamPart } from './stream-text-result';
 import type { TypedToolCall } from './tool-call';
-import type { ExpandedContext, ToolSet } from './tool-set';
+import type { GenerationContext } from './generation-context';
+import type { ToolSet } from './tool-set';
 
 /**
  * Common model information used across callback events.
@@ -31,7 +32,7 @@ export interface CallbackModelInfo {
  */
 export interface OnStartEvent<
   TOOLS extends ToolSet = ToolSet,
-  CONTEXT extends ExpandedContext<TOOLS> = ExpandedContext<TOOLS>,
+  CONTEXT extends GenerationContext<TOOLS> = GenerationContext<TOOLS>,
   OUTPUT extends Output = Output,
   INCLUDE = { requestBody?: boolean; responseBody?: boolean },
 > {
@@ -150,7 +151,7 @@ export interface OnStartEvent<
  */
 export interface OnStepStartEvent<
   TOOLS extends ToolSet = ToolSet,
-  CONTEXT extends ExpandedContext<TOOLS> = ExpandedContext<TOOLS>,
+  CONTEXT extends GenerationContext<TOOLS> = GenerationContext<TOOLS>,
   OUTPUT extends Output = Output,
   INCLUDE = { requestBody?: boolean; responseBody?: boolean },
 > {
@@ -370,7 +371,7 @@ export interface OnChunkEvent<TOOLS extends ToolSet = ToolSet> {
  */
 export type OnStepFinishEvent<
   TOOLS extends ToolSet = ToolSet,
-  CONTEXT extends ExpandedContext<TOOLS> = ExpandedContext<TOOLS>,
+  CONTEXT extends GenerationContext<TOOLS> = GenerationContext<TOOLS>,
 > = StepResult<TOOLS, CONTEXT>;
 
 /**
@@ -381,7 +382,7 @@ export type OnStepFinishEvent<
  */
 export type OnFinishEvent<
   TOOLS extends ToolSet = ToolSet,
-  CONTEXT extends ExpandedContext<TOOLS> = ExpandedContext<TOOLS>,
+  CONTEXT extends GenerationContext<TOOLS> = GenerationContext<TOOLS>,
 > = StepResult<TOOLS, CONTEXT> & {
   /** Array containing results from all steps in the generation. */
   readonly steps: StepResult<TOOLS, CONTEXT>[];
diff --git a/packages/ai/src/generate-text/create-execute-tools-transformation.ts b/packages/ai/src/generate-text/create-execute-tools-transformation.ts
index 736f5886b932..a1e3e967b43e 100644
--- a/packages/ai/src/generate-text/create-execute-tools-transformation.ts
+++ b/packages/ai/src/generate-text/create-execute-tools-transformation.ts
@@ -9,12 +9,13 @@ import {
   StreamTextOnToolCallStartCallback,
 } from './stream-text';
 import { TypedToolCall } from './tool-call';
-import { ExpandedContext, ToolSet } from './tool-set';
+import type { GenerationContext } from './generation-context';
+import type { ToolSet } from './tool-set';
 import { ModelCallStreamPart } from './stream-model-call';
 
 export function createExecuteToolsTransformation<
   TOOLS extends ToolSet,
-  CONTEXT extends ExpandedContext<TOOLS>,
+  CONTEXT extends GenerationContext<TOOLS>,
 >({
   tools,
   telemetry,
diff --git a/packages/ai/src/generate-text/execute-tool-call.ts b/packages/ai/src/generate-text/execute-tool-call.ts
index d251b6bc517a..409b75c43478 100644
--- a/packages/ai/src/generate-text/execute-tool-call.ts
+++ b/packages/ai/src/generate-text/execute-tool-call.ts
@@ -14,7 +14,8 @@ import { TypedToolCall } from './tool-call';
 import { TypedToolError } from './tool-error';
 import { ToolOutput } from './tool-output';
 import { TypedToolResult } from './tool-result';
-import { ExpandedContext, ToolSet } from './tool-set';
+import type { GenerationContext } from './generation-context';
+import type { ToolSet } from './tool-set';
 
 /**
  * Executes a single tool call and manages its lifecycle callbacks.
@@ -29,7 +30,7 @@ import { ExpandedContext, ToolSet } from './tool-set';
  */
 export async function executeToolCall<
   TOOLS extends ToolSet,
-  CONTEXT extends ExpandedContext<TOOLS>,
+  CONTEXT extends GenerationContext<TOOLS>,
 >({
   toolCall,
   tools,
diff --git a/packages/ai/src/generate-text/generate-text-result.ts b/packages/ai/src/generate-text/generate-text-result.ts
index 5722e5e8e31e..a80bb66df7ab 100644
--- a/packages/ai/src/generate-text/generate-text-result.ts
+++ b/packages/ai/src/generate-text/generate-text-result.ts
@@ -16,7 +16,8 @@ import {
   StaticToolResult,
   TypedToolResult,
 } from './tool-result';
-import { ExpandedContext, ToolSet } from './tool-set';
+import type { GenerationContext } from './generation-context';
+import type { ToolSet } from './tool-set';
 
 /**
  * The result of a `generateText` call.
@@ -24,7 +25,7 @@ import { ExpandedContext, ToolSet } from './tool-set';
  */
 export interface GenerateTextResult<
   TOOLS extends ToolSet,
-  CONTEXT extends ExpandedContext<TOOLS>,
+  CONTEXT extends GenerationContext<TOOLS>,
   OUTPUT extends Output,
 > {
   /**
diff --git a/packages/ai/src/generate-text/generate-text.ts b/packages/ai/src/generate-text/generate-text.ts
index 80d47d93223c..317b2b7ff977 100644
--- a/packages/ai/src/generate-text/generate-text.ts
+++ b/packages/ai/src/generate-text/generate-text.ts
@@ -64,6 +64,7 @@ import { executeToolCall } from './execute-tool-call';
 import { filterActiveTools } from './filter-active-tool';
 import { GenerateTextResult } from './generate-text-result';
 import { DefaultGeneratedFile } from './generated-file';
+import type { GenerationContext } from './generation-context';
 import { isApprovalNeeded } from './is-approval-needed';
 import { Output, text } from './output';
 import { InferCompleteOutput } from './output-utils';
@@ -84,7 +85,7 @@ import { ToolCallRepairFunction } from './tool-call-repair-function';
 import { TypedToolError } from './tool-error';
 import { ToolOutput } from './tool-output';
 import { TypedToolResult } from './tool-result';
-import { ExpandedContext, ToolSet } from './tool-set';
+import type { ToolSet } from './tool-set';
 
 const originalGenerateId = createIdGenerator({
   prefix: 'aitxt',
@@ -115,7 +116,7 @@ type GenerateTextIncludeSettings = {
  */
 export type GenerateTextOnStartCallback<
   TOOLS extends ToolSet,
-  CONTEXT extends ExpandedContext<TOOLS>,
+  CONTEXT extends GenerationContext<TOOLS>,
   OUTPUT extends Output = Output,
 > = (
   event: OnStartEvent<TOOLS, CONTEXT, OUTPUT, GenerateTextIncludeSettings>,
@@ -132,7 +133,7 @@ export type GenerateTextOnStartCallback<
  */
 export type GenerateTextOnStepStartCallback<
   TOOLS extends ToolSet,
-  CONTEXT extends ExpandedContext<TOOLS>,
+  CONTEXT extends GenerationContext<TOOLS>,
   OUTPUT extends Output = Output,
 > = (
   event: OnStepStartEvent<TOOLS, CONTEXT, OUTPUT, GenerateTextIncludeSettings>,
@@ -176,7 +177,7 @@ export type GenerateTextOnToolCallFinishCallback<TOOLS extends ToolSet> = (
  */
 export type GenerateTextOnStepFinishCallback<
   TOOLS extends ToolSet,
-  CONTEXT extends ExpandedContext<TOOLS>,
+  CONTEXT extends GenerationContext<TOOLS>,
 > = (event: OnStepFinishEvent<TOOLS, CONTEXT>) => Promise<void> | void;
 
 /**
@@ -190,7 +191,7 @@ export type GenerateTextOnStepFinishCallback<
  */
 export type GenerateTextOnFinishCallback<
   TOOLS extends ToolSet,
-  CONTEXT extends ExpandedContext<TOOLS>,
+  CONTEXT extends GenerationContext<TOOLS>,
 > = (event: OnFinishEvent<TOOLS, CONTEXT>) => PromiseLike<void> | void;
 
 /**
@@ -249,7 +250,7 @@ export type GenerateTextOnFinishCallback<
  */
 export async function generateText<
   TOOLS extends ToolSet,
-  CONTEXT extends ExpandedContext<TOOLS>,
+  CONTEXT extends GenerationContext<TOOLS>,
   OUTPUT extends Output = Output<string, string>,
 >({
   model: modelArg,
@@ -1127,7 +1128,7 @@ export async function generateText<
 
 async function executeTools<
   TOOLS extends ToolSet,
-  CONTEXT extends ExpandedContext<TOOLS>,
+  CONTEXT extends GenerationContext<TOOLS>,
 >({
   toolCalls,
   tools,
@@ -1187,7 +1188,7 @@ async function executeTools<
 
 class DefaultGenerateTextResult<
   TOOLS extends ToolSet,
-  CONTEXT extends ExpandedContext<TOOLS>,
+  CONTEXT extends GenerationContext<TOOLS>,
   OUTPUT extends Output,
 > implements GenerateTextResult<TOOLS, CONTEXT, OUTPUT> {
   readonly steps: GenerateTextResult<TOOLS, CONTEXT, OUTPUT>['steps'];
diff --git a/packages/ai/src/generate-text/generation-context.ts b/packages/ai/src/generate-text/generation-context.ts
new file mode 100644
index 000000000000..510b2050f133
--- /dev/null
+++ b/packages/ai/src/generate-text/generation-context.ts
@@ -0,0 +1,14 @@
+import type { Context } from '@ai-sdk/provider-utils';
+import type { InferToolSetContext } from './infer-tool-set-context';
+import type { ToolSet } from './tool-set';
+
+/**
+ * The context type for a generation call.
+ *
+ * It expands the tool set context with the generic context type for
+ * e.g. prepareStep or telemetry.
+ */
+export type GenerationContext<TOOLS extends ToolSet> = InferToolSetContext<
+  NoInfer<TOOLS>
+> &
+  Context;
diff --git a/packages/ai/src/generate-text/prepare-step.ts b/packages/ai/src/generate-text/prepare-step.ts
index cf06e5e3c5e7..890f934de36c 100644
--- a/packages/ai/src/generate-text/prepare-step.ts
+++ b/packages/ai/src/generate-text/prepare-step.ts
@@ -6,7 +6,8 @@ import {
 } from '@ai-sdk/provider-utils';
 import { LanguageModel, ToolChoice } from '../types/language-model';
 import { StepResult } from './step-result';
-import { ExpandedContext, ToolSet } from './tool-set';
+import type { GenerationContext } from './generation-context';
+import type { ToolSet } from './tool-set';
 
 /**
  * Function that you can use to provide different settings for a step.
@@ -23,7 +24,7 @@ import { ExpandedContext, ToolSet } from './tool-set';
  */
 export type PrepareStepFunction<
   TOOLS extends ToolSet,
-  CONTEXT extends ExpandedContext<TOOLS>,
+  CONTEXT extends GenerationContext<TOOLS>,
 > = (options: {
   /**
    * The steps that have been executed so far.
@@ -59,7 +60,7 @@ export type PrepareStepFunction<
  */
 export type PrepareStepResult<
   TOOLS extends ToolSet,
-  CONTEXT extends ExpandedContext<TOOLS>,
+  CONTEXT extends GenerationContext<TOOLS>,
 > =
   | {
       /**
diff --git a/packages/ai/src/generate-text/step-result.ts b/packages/ai/src/generate-text/step-result.ts
index 399aa1cb3770..dddfe648f61f 100644
--- a/packages/ai/src/generate-text/step-result.ts
+++ b/packages/ai/src/generate-text/step-result.ts
@@ -23,14 +23,15 @@ import {
   StaticToolResult,
   TypedToolResult,
 } from './tool-result';
-import { ExpandedContext, ToolSet } from './tool-set';
+import type { GenerationContext } from './generation-context';
+import type { ToolSet } from './tool-set';
 
 /**
  * The result of a single step in the generation process.
  */
 export type StepResult<
   TOOLS extends ToolSet,
-  CONTEXT extends ExpandedContext<TOOLS>,
+  CONTEXT extends GenerationContext<TOOLS>,
 > = {
   /**
    * Unique identifier for the generation call this step belongs to.
@@ -181,7 +182,7 @@ export type StepResult<
 
 export class DefaultStepResult<
   TOOLS extends ToolSet,
-  CONTEXT extends ExpandedContext<TOOLS>,
+  CONTEXT extends GenerationContext<TOOLS>,
 > implements StepResult<TOOLS, CONTEXT> {
   readonly callId: StepResult<TOOLS, CONTEXT>['callId'];
   readonly stepNumber: StepResult<TOOLS, CONTEXT>['stepNumber'];
diff --git a/packages/ai/src/generate-text/stop-condition.ts b/packages/ai/src/generate-text/stop-condition.ts
index 7a849fceb802..6527b7f73218 100644
--- a/packages/ai/src/generate-text/stop-condition.ts
+++ b/packages/ai/src/generate-text/stop-condition.ts
@@ -1,9 +1,10 @@
 import { StepResult } from './step-result';
-import { ExpandedContext, ToolSet } from './tool-set';
+import type { GenerationContext } from './generation-context';
+import type { ToolSet } from './tool-set';
 
 export type StopCondition<
   TOOLS extends ToolSet,
-  CONTEXT extends ExpandedContext<TOOLS>,
+  CONTEXT extends GenerationContext<TOOLS>,
 > = (options: {
   steps: Array<StepResult<TOOLS, CONTEXT>>;
 }) => PromiseLike<boolean> | boolean;
@@ -25,7 +26,7 @@ export function hasToolCall(toolName: string): StopCondition<any, any> {
 
 export async function isStopConditionMet<
   TOOLS extends ToolSet,
-  CONTEXT extends ExpandedContext<TOOLS>,
+  CONTEXT extends GenerationContext<TOOLS>,
 >({
   stopConditions,
   steps,
diff --git a/packages/ai/src/generate-text/stream-text-result.ts b/packages/ai/src/generate-text/stream-text-result.ts
index a75da8fc0bbd..cbef38b52777 100644
--- a/packages/ai/src/generate-text/stream-text-result.ts
+++ b/packages/ai/src/generate-text/stream-text-result.ts
@@ -35,7 +35,8 @@ import {
   StaticToolResult,
   TypedToolResult,
 } from './tool-result';
-import { ExpandedContext, ToolSet } from './tool-set';
+import type { GenerationContext } from './generation-context';
+import type { ToolSet } from './tool-set';
 
 export type UIMessageStreamOptions<UI_MESSAGE extends UIMessage> = {
   /**
@@ -108,7 +109,7 @@ export type ConsumeStreamOptions = {
  */
 export interface StreamTextResult<
   TOOLS extends ToolSet,
-  CONTEXT extends ExpandedContext<TOOLS>,
+  CONTEXT extends GenerationContext<TOOLS>,
   OUTPUT extends Output,
 > {
   /**
diff --git a/packages/ai/src/generate-text/stream-text.ts b/packages/ai/src/generate-text/stream-text.ts
index 8c77fe47547f..455bbefd6b9f 100644
--- a/packages/ai/src/generate-text/stream-text.ts
+++ b/packages/ai/src/generate-text/stream-text.ts
@@ -113,7 +113,8 @@ import { TypedToolCall } from './tool-call';
 import { ToolCallRepairFunction } from './tool-call-repair-function';
 import { ToolOutput } from './tool-output';
 import { StaticToolOutputDenied } from './tool-output-denied';
-import { ExpandedContext, ToolSet } from './tool-set';
+import type { GenerationContext } from './generation-context';
+import type { ToolSet } from './tool-set';
 
 const originalGenerateId = createIdGenerator({
   prefix: 'aitxt',
@@ -152,7 +153,7 @@ export type StreamTextOnErrorCallback = (event: {
  */
 export type StreamTextOnStepFinishCallback<
   TOOLS extends ToolSet,
-  CONTEXT extends ExpandedContext<TOOLS>,
+  CONTEXT extends GenerationContext<TOOLS>,
 > = (event: OnStepFinishEvent<TOOLS, CONTEXT>) => PromiseLike<void> | void;
 
 /**
@@ -185,7 +186,7 @@ export type StreamTextOnChunkCallback<TOOLS extends ToolSet> = (event: {
  */
 export type StreamTextOnFinishCallback<
   TOOLS extends ToolSet,
-  CONTEXT extends ExpandedContext<TOOLS>,
+  CONTEXT extends GenerationContext<TOOLS>,
 > = (event: OnFinishEvent<TOOLS, CONTEXT>) => PromiseLike<void> | void;
 
 /**
@@ -195,7 +196,7 @@ export type StreamTextOnFinishCallback<
  */
 export type StreamTextOnAbortCallback<
   TOOLS extends ToolSet,
-  CONTEXT extends ExpandedContext<TOOLS>,
+  CONTEXT extends GenerationContext<TOOLS>,
 > = (event: {
   /**
    * Details for all previously finished steps.
@@ -219,7 +220,7 @@ type StreamTextIncludeSettings = { requestBody?: boolean };
  */
 export type StreamTextOnStartCallback<
   TOOLS extends ToolSet = ToolSet,
-  CONTEXT extends ExpandedContext<TOOLS> = ExpandedContext<TOOLS>,
+  CONTEXT extends GenerationContext<TOOLS> = GenerationContext<TOOLS>,
   OUTPUT extends Output = Output,
 > = (
   event: OnStartEvent<TOOLS, CONTEXT, OUTPUT, StreamTextIncludeSettings>,
@@ -236,7 +237,7 @@ export type StreamTextOnStartCallback<
  */
 export type StreamTextOnStepStartCallback<
   TOOLS extends ToolSet = ToolSet,
-  CONTEXT extends ExpandedContext<TOOLS> = ExpandedContext<TOOLS>,
+  CONTEXT extends GenerationContext<TOOLS> = GenerationContext<TOOLS>,
   OUTPUT extends Output = Output,
 > = (
   event: OnStepStartEvent<TOOLS, CONTEXT, OUTPUT, StreamTextIncludeSettings>,
@@ -297,7 +298,7 @@ export type StreamTextOnToolCallFinishCallback<
  */
 export function streamText<
   TOOLS extends ToolSet,
-  CONTEXT extends ExpandedContext<TOOLS> = ExpandedContext<TOOLS>,
+  CONTEXT extends GenerationContext<TOOLS> = GenerationContext<TOOLS>,
   OUTPUT extends Output = Output<string, string, never>,
 >({
   model,
@@ -713,7 +714,7 @@ function createOutputTransformStream<
 
 class DefaultStreamTextResult<
   TOOLS extends ToolSet,
-  CONTEXT extends ExpandedContext<TOOLS>,
+  CONTEXT extends GenerationContext<TOOLS>,
   OUTPUT extends Output,
 > implements StreamTextResult<TOOLS, CONTEXT, OUTPUT> {
   private readonly _totalUsage = new DelayedPromise<
diff --git a/packages/ai/src/generate-text/tool-set.ts b/packages/ai/src/generate-text/tool-set.ts
index fa2fe650ad7d..57f40d946f9c 100644
--- a/packages/ai/src/generate-text/tool-set.ts
+++ b/packages/ai/src/generate-text/tool-set.ts
@@ -1,5 +1,4 @@
-import { Context, Tool } from '@ai-sdk/provider-utils';
-import type { InferToolSetContext } from './infer-tool-set-context';
+import type { Tool } from '@ai-sdk/provider-utils';
 
 export type ToolSet = Record<
   string,
@@ -18,8 +17,3 @@ export type ToolSet = Record<
       | 'needsApproval'
     >
 >;
-
-export type ExpandedContext<TOOLS extends ToolSet> = InferToolSetContext<
-  NoInfer<TOOLS>
-> &
-  Context;
diff --git a/packages/ai/src/telemetry/open-telemetry-integration.ts b/packages/ai/src/telemetry/open-telemetry-integration.ts
index 89ad5f5e029a..f9a156937511 100644
--- a/packages/ai/src/telemetry/open-telemetry-integration.ts
+++ b/packages/ai/src/telemetry/open-telemetry-integration.ts
@@ -31,7 +31,8 @@ import type {
   OnToolCallStartEvent,
 } from '../generate-text/core-events';
 import type { Output } from '../generate-text/output';
-import type { ExpandedContext, ToolSet } from '../generate-text/tool-set';
+import type { GenerationContext } from '../generate-text/generation-context';
+import type { ToolSet } from '../generate-text/tool-set';
 import type {
   ObjectOnStartEvent,
   ObjectOnFinishEvent,
@@ -115,7 +116,7 @@ function selectAttributes(
 
 interface OtelStepStartEvent<
   TOOLS extends ToolSet = ToolSet,
-  CONTEXT extends ExpandedContext<TOOLS> = ExpandedContext<TOOLS>,
+  CONTEXT extends GenerationContext<TOOLS> = GenerationContext<TOOLS>,
   OUTPUT extends Output = Output,
 > extends OnStepStartEvent<TOOLS, CONTEXT, OUTPUT> {
   readonly promptMessages?: LanguageModelV4Prompt;
diff --git a/packages/ai/src/ui/direct-chat-transport.ts b/packages/ai/src/ui/direct-chat-transport.ts
index d1d5cd6c75ff..52a999944f8a 100644
--- a/packages/ai/src/ui/direct-chat-transport.ts
+++ b/packages/ai/src/ui/direct-chat-transport.ts
@@ -1,7 +1,8 @@
 import { Agent } from '../agent/agent';
 import { Output } from '../generate-text/output';
 import { UIMessageStreamOptions } from '../generate-text/stream-text-result';
-import { ExpandedContext, ToolSet } from '../generate-text/tool-set';
+import type { GenerationContext } from '../generate-text/generation-context';
+import type { ToolSet } from '../generate-text/tool-set';
 import { UIMessageChunk } from '../ui-message-stream/ui-message-chunks';
 import { ChatTransport } from './chat-transport';
 import { convertToModelMessages } from './convert-to-model-messages';
@@ -14,7 +15,7 @@ import { validateUIMessages } from './validate-ui-messages';
 export type DirectChatTransportOptions<
   CALL_OPTIONS,
   TOOLS extends ToolSet,
-  CONTEXT extends ExpandedContext<TOOLS>,
+  CONTEXT extends GenerationContext<TOOLS>,
   OUTPUT extends Output<TOOLS, CONTEXT>,
   UI_MESSAGE extends UIMessage<unknown, never, InferUITools<TOOLS>>,
 > = {
@@ -50,7 +51,7 @@ export type DirectChatTransportOptions<
 export class DirectChatTransport<
   CALL_OPTIONS = never,
   TOOLS extends ToolSet = {},
-  CONTEXT extends ExpandedContext<TOOLS> = ExpandedContext<TOOLS>,
+  CONTEXT extends GenerationContext<TOOLS> = GenerationContext<TOOLS>,
   OUTPUT extends Output<TOOLS, CONTEXT> = never,
   UI_MESSAGE extends UIMessage<unknown, never, InferUITools<TOOLS>> = UIMessage<
     unknown,

From 780835f11245c33d43d19961d6ca062bfde0e96d Mon Sep 17 00:00:00 2001
From: Lars Grammel <lars.grammel@gmail.com>
Date: Wed, 1 Apr 2026 15:46:49 +0200
Subject: [PATCH 39/41] Add tests for GenerationContext type to validate
 context merging and fallback behavior

---
 .../generation-context.test-d.ts              | 53 +++++++++++++++++++
 .../src/generate-text/generation-context.ts   |  3 +-
 2 files changed, 55 insertions(+), 1 deletion(-)
 create mode 100644 packages/ai/src/generate-text/generation-context.test-d.ts

diff --git a/packages/ai/src/generate-text/generation-context.test-d.ts b/packages/ai/src/generate-text/generation-context.test-d.ts
new file mode 100644
index 000000000000..942581e206a6
--- /dev/null
+++ b/packages/ai/src/generate-text/generation-context.test-d.ts
@@ -0,0 +1,53 @@
+import { Context, tool } from '@ai-sdk/provider-utils';
+import { describe, expectTypeOf, it } from 'vitest';
+import { z } from 'zod/v4';
+import type { GenerationContext } from './generation-context';
+
+describe('GenerationContext', () => {
+  it('combines inferred tool context with the generic context type', () => {
+    const tools = {
+      weather: tool({
+        inputSchema: z.object({
+          city: z.string(),
+        }),
+        contextSchema: z.object({
+          userId: z.string(),
+        }),
+      }),
+      forecast: tool({
+        inputSchema: z.object({
+          days: z.number(),
+        }),
+        contextSchema: z.object({
+          role: z.string(),
+        }),
+      }),
+    };
+
+    expectTypeOf<GenerationContext<typeof tools>>().toEqualTypeOf<
+      {
+        userId: string;
+      } & {
+        role: string;
+      } & Context
+    >();
+
+    expectTypeOf<
+      GenerationContext<typeof tools>
+    >().toMatchObjectType<Context>();
+  });
+
+  it('falls back to the generic context type when tools have no contextSchema', () => {
+    const tools = {
+      weather: tool({
+        inputSchema: z.object({
+          city: z.string(),
+        }),
+      }),
+    };
+
+    expectTypeOf<
+      GenerationContext<typeof tools>
+    >().toMatchObjectType<Context>();
+  });
+});
diff --git a/packages/ai/src/generate-text/generation-context.ts b/packages/ai/src/generate-text/generation-context.ts
index 510b2050f133..7b9a0dce877a 100644
--- a/packages/ai/src/generate-text/generation-context.ts
+++ b/packages/ai/src/generate-text/generation-context.ts
@@ -6,7 +6,8 @@ import type { ToolSet } from './tool-set';
  * The context type for a generation call.
  *
  * It expands the tool set context with the generic context type for
- * e.g. prepareStep or telemetry.
+ * e.g. prepareStep or telemetry,
+ * while keeping the inferred tool set context for autocompletion.
  */
 export type GenerationContext<TOOLS extends ToolSet> = InferToolSetContext<
   NoInfer<TOOLS>

From bc2fbf52343a77aa62e29af1b1841c0ecbe4a5a9 Mon Sep 17 00:00:00 2001
From: Lars Grammel <lars.grammel@gmail.com>
Date: Wed, 1 Apr 2026 15:58:34 +0200
Subject: [PATCH 40/41] test

---
 .../src/generate-text/generate-text.test-d.ts |  4 +-
 .../src/generate-text/stream-text.test-d.ts   | 44 +++++++++++++++++++
 2 files changed, 46 insertions(+), 2 deletions(-)

diff --git a/packages/ai/src/generate-text/generate-text.test-d.ts b/packages/ai/src/generate-text/generate-text.test-d.ts
index d1fb9110da4e..3bb3dc08a1b0 100644
--- a/packages/ai/src/generate-text/generate-text.test-d.ts
+++ b/packages/ai/src/generate-text/generate-text.test-d.ts
@@ -81,7 +81,7 @@ describe('generateText types', () => {
               userId: z.string(),
             }),
             execute: async (_input, { experimental_context }) => {
-              expectTypeOf(experimental_context).toMatchTypeOf<{
+              expectTypeOf(experimental_context).toMatchObjectType<{
                 userId: string;
               }>();
 
@@ -94,7 +94,7 @@ describe('generateText types', () => {
           role: 'admin',
         },
         prepareStep: ({ experimental_context }) => {
-          expectTypeOf(experimental_context).toMatchTypeOf<{
+          expectTypeOf(experimental_context).toMatchObjectType<{
             userId: string;
             role: string;
           }>();
diff --git a/packages/ai/src/generate-text/stream-text.test-d.ts b/packages/ai/src/generate-text/stream-text.test-d.ts
index 3a9afb5228af..824d3a8b0909 100644
--- a/packages/ai/src/generate-text/stream-text.test-d.ts
+++ b/packages/ai/src/generate-text/stream-text.test-d.ts
@@ -1,4 +1,5 @@
 import { JSONValue } from '@ai-sdk/provider';
+import { tool } from '@ai-sdk/provider-utils';
 import { describe, expectTypeOf, it } from 'vitest';
 import { z } from 'zod';
 import { Output, streamText } from '../generate-text';
@@ -197,4 +198,47 @@ describe('streamText types', () => {
       >();
     });
   });
+
+  describe('experimental_context', () => {
+    it('should infer typed experimental_context with one tool context and prepareStep', async () => {
+      streamText({
+        model: new MockLanguageModelV4(),
+        prompt: 'Hello, world!',
+        tools: {
+          weather: tool({
+            inputSchema: z.object({
+              city: z.string(),
+            }),
+            contextSchema: z.object({
+              userId: z.string(),
+            }),
+            execute: async (_input, { experimental_context }) => {
+              expectTypeOf(experimental_context).toMatchObjectType<{
+                userId: string;
+              }>();
+
+              return 'sunny';
+            },
+          }),
+        },
+        experimental_context: {
+          userId: 'test-user',
+          role: 'admin',
+        },
+        prepareStep: ({ experimental_context }) => {
+          expectTypeOf(experimental_context).toMatchObjectType<{
+            userId: string;
+            role: string;
+          }>();
+
+          return {
+            experimental_context: {
+              userId: experimental_context.userId,
+              role: experimental_context.role,
+            },
+          };
+        },
+      });
+    });
+  });
 });

From bf081149df2cbb8dd1b33aa4adc4c4a72bc55d21 Mon Sep 17 00:00:00 2001
From: Lars Grammel <lars.grammel@gmail.com>
Date: Wed, 1 Apr 2026 17:35:29 +0200
Subject: [PATCH 41/41] tf

---
 .../ai/src/util/union-to-intersection.test-d.ts     | 13 ++++++++-----
 1 file changed, 8 insertions(+), 5 deletions(-)

diff --git a/packages/ai/src/util/union-to-intersection.test-d.ts b/packages/ai/src/util/union-to-intersection.test-d.ts
index f2fde5a0fd41..3a94464b5c58 100644
--- a/packages/ai/src/util/union-to-intersection.test-d.ts
+++ b/packages/ai/src/util/union-to-intersection.test-d.ts
@@ -5,7 +5,7 @@ describe('UnionToIntersection', () => {
   it('returns never when given no input', () => {
     type Result = UnionToIntersection<never>;
 
-    expectTypeOf<Result>().toEqualTypeOf<never>();
+    expectTypeOf<Result>().toEqualTypeOf<unknown>();
   });
 
   it('returns the same type for a single input', () => {
@@ -21,9 +21,12 @@ describe('UnionToIntersection', () => {
       { city: string } | { countryCode: string }
     >;
 
-    expectTypeOf<Result>().toEqualTypeOf<{
-      city: string;
-      countryCode: string;
-    }>();
+    expectTypeOf<Result>().toEqualTypeOf<
+      {
+        city: string;
+      } & {
+        countryCode: string;
+      }
+    >();
   });
 });