feat: add tool choice with .choice()

Backport of the tool-choice feature to the v1 line (AI SDK v6).

Add a `.choice()` method that resolves the AI SDK `toolChoice` from the same
`{ messages, steps, context }` values as activation. `inferTools()` now returns
`toolChoice` alongside `tools` and `activeTools`, spreadable into
generateText/streamText and prepareStep.

Last-call wins; a resolver returning `undefined` omits `toolChoice` (auto).
The resolved value is passed through as-is, no activeTools validation.

Author: zirkelc

README.md

@@ -12,16 +12,17 @@
 
 </div>
 
-This library provides a type-safe API to manage [`activeTools`](https://ai-sdk.dev/docs/reference/ai-sdk-core/generate-text#active-tools) for [`generateText()`](https://ai-sdk.dev/docs/reference/ai-sdk-core/generate-text) and [`streamText()`](https://ai-sdk.dev/docs/reference/ai-sdk-core/stream-text) in the AI SDK.
+This library provides a type-safe API to manage [`activeTools`](https://ai-sdk.dev/docs/reference/ai-sdk-core/generate-text#active-tools) and [`toolChoice`](https://ai-sdk.dev/docs/ai-sdk-core/tools-and-tool-calling#tool-choice) for [`generateText()`](https://ai-sdk.dev/docs/reference/ai-sdk-core/generate-text) and [`streamText()`](https://ai-sdk.dev/docs/reference/ai-sdk-core/stream-text) in the AI SDK.
 
 ### Why?
 
-The AI SDK provides an `activeTools` parameter to control which tools the model can use at any given time. However, managing tool activation becomes complex when you need to:
+The AI SDK provides an `activeTools` parameter to control which tools the model can use, and a `toolChoice` parameter to force which tool is called. However, managing this becomes complex when you need to:
 
 - **Statically activate/deactivate tools**: Some tools should be inactive by default and only available after being explicitly activated
 - **Dynamically infer tool activation**: Some tools should be activated based on runtime context like the conversation history
+- **Force tool choice**: Sometimes the model should be required to call a specific tool, or any tool, based on context
 
-This library wraps standard AI SDK `tool()` definitions with chainable activation methods and resolves `tools` and `activeTools` for any AI SDK function.
+This library wraps standard AI SDK `tool()` definitions with chainable activation and choice methods and resolves `tools`, `activeTools`, and `toolChoice` for any AI SDK function.
 
 ### Installation
 
@@ -220,6 +221,50 @@ const toolSet = createToolSet({ tools })
   .activateWhen('cancel_order', ({ messages }) => hasUnfulfilledOrders(messages));
 ```
 
+### Tool Choice
+
+The AI SDK can steer tool calls with [`toolChoice`](https://ai-sdk.dev/docs/ai-sdk-core/tools-and-tool-calling#tool-choice), which controls whether and which tool the model must call. `.inferTools()` resolves a `toolChoice` value alongside `tools` and `activeTools`, so a single tool set drives activation and choice together.
+
+Use `.choice()` with a constant for static choices or use a resolver to decide dynamically:
+
+```typescript
+const toolSet = createToolSet({ tools })
+  .deactivate(['cancel_order'])
+  // Force the model to call one of the active tools this turn
+  .choice('required')
+  // Or resolve dynamically, e.g. force unauthenticated users into search
+  .choice(({ context }) => (context?.isAuthenticated ? 'auto' : { type: 'tool', toolName: 'search' }));
+
+const { tools, activeTools, toolChoice } = toolSet.inferTools();
+
+const result = await generateText({ model, tools, activeTools, toolChoice, prompt: 'Show me my orders' });
+```
+
+A `.choice()` entry is either a constant [`ToolChoice`](https://ai-sdk.dev/docs/ai-sdk-core/tools-and-tool-calling#tool-choice) (`'auto'`, `'required'`, `'none'`, or `{ type: 'tool', toolName }` with `toolName` typed to your tools) or a **resolver**. The resolver runs inside `.inferTools()` with your toolset values (`{ messages, steps, context }`) and returns a `ToolChoice`, or `undefined` to leave the choice unconstrained (the AI SDK then defaults to `'auto'`). Since it sets a single value, it follows the same **last-call wins** rule as activation: the last entry decides, and a resolver returning `undefined` does not fall back to an earlier one, so express any fallback inside the resolver itself.
+
+> [!NOTE]
+> `toolChoice` is resolved independently of `activeTools` and passed through as-is. Because the AI SDK filters tools by `activeTools` first, forcing `{ type: 'tool', toolName }` for a tool you have deactivated leaves the model unable to see it and surfaces as a provider error, so keep a forced tool active.
+
+Resolving from `steps` makes it easy to guide a multi-step run. Call `.inferTools()` inside `prepareStep()` to recompute `toolChoice` (and `activeTools`) after each step, for example to force a tool on the first step and then hand control back to the model:
+
+```typescript
+import { generateText, stepCountIs } from 'ai';
+
+const toolSet = createToolSet({ tools }).choice(({ steps }) =>
+  steps?.length === 0 ? { type: 'tool', toolName: 'search' } : 'auto',
+);
+
+const { tools } = toolSet;
+
+const result = await generateText({
+  model,
+  tools,
+  stopWhen: stepCountIs(10),
+  prompt: 'Find my orders and cancel the unfulfilled one',
+  prepareStep: ({ steps }) => toolSet.inferTools({ steps }),
+});
+```
+
 ### Immutable vs Mutable
 
 By default, `createToolSet()` returns an **immutable** tool set, that means every method returns a new instance and the original is never modified. This is ideal when the tool set is created once in the global scope and shared across requests:
@@ -420,32 +465,47 @@ Conditionally deactivate tools. The predicate receives `{ messages, steps, conte
 toolSet.deactivateWhen('search', ({ messages }) => messages && messages.length > 10);
 ```
 
+#### `.choice(entry)`
+
+Set the toolset's `toolChoice` to a constant `ToolChoice` (`'auto'`, `'none'`, `'required'`, or `{ type: 'tool', toolName }`) or a resolver. A resolver receives `{ messages, steps, context }` and returns a `ToolChoice`, or `undefined` to leave it unconstrained (omitted). Last-call wins. The resolved `toolChoice` is passed through as-is; keep a forced tool active so the AI SDK does not filter it out.
+
+```ts
+// Constant
+toolSet.choice('required');
+toolSet.choice({ type: 'tool', toolName: 'search' });
+
+// Resolver — force a tool on the first step, then let the model decide
+toolSet.choice(({ steps }) => (steps?.length === 0 ? { type: 'tool', toolName: 'search' } : 'auto'));
+```
+
 #### `.inferTools(input?)`
 
-Evaluate all predicates and return `{ tools, activeTools }`, directly spreadable into `generateText()` or `streamText()`. The input is optional; all fields are optional. Predicates receive `undefined` for fields not provided.
+Evaluate all predicates and the tool-choice entry and return `{ tools, activeTools, toolChoice }`, directly spreadable into `generateText()` or `streamText()`. The input is optional; all fields are optional. Predicates and resolvers receive `undefined` for fields not provided.
 
 - `input` (optional):
   - `messages` (optional), the current conversation messages
   - `steps` (optional), the completed steps of the current run (the `StepResult` array from `prepareStep`)
-  - `context` (optional), arbitrary values passed to predicates
+  - `context` (optional), arbitrary values passed to predicates and the tool-choice resolver
+
+The result also includes `toolChoice` when set via `.choice()` (otherwise `undefined`).
 
 ```ts
 // Static-only (no predicates)
-const { tools, activeTools } = toolSet.inferTools();
+const { tools, activeTools, toolChoice } = toolSet.inferTools();
 
 // With messages
-const { tools, activeTools } = toolSet.inferTools({ messages });
+const { tools, activeTools, toolChoice } = toolSet.inferTools({ messages });
 
 // With context
-const { tools, activeTools } = toolSet.inferTools({ context: { isAdmin: true } });
+const { tools, activeTools, toolChoice } = toolSet.inferTools({ context: { isAdmin: true } });
 
 // With steps (e.g. inside prepareStep)
-const { tools, activeTools } = toolSet.inferTools({ steps });
+const { tools, activeTools, toolChoice } = toolSet.inferTools({ steps });
 
 // With both
-const { tools, activeTools } = toolSet.inferTools({ messages, context });
+const { tools, activeTools, toolChoice } = toolSet.inferTools({ messages, context });
 
-const result = await generateText({ model, tools, activeTools, messages });
+const result = await generateText({ model, tools, activeTools, toolChoice, messages });
 ```
 
 #### `.clone(options?)`

examples/tool-choice.ts

@@ -0,0 +1,89 @@
+/**
+ * Tool choice — force which tool the model must call, resolved from the same
+ * `{ messages, steps, context }` values that drive activation.
+ *
+ * `.choice()` sets the AI SDK's `toolChoice`. It accepts a constant
+ * (`'auto'` | `'none'` | `'required'` | `{ type: 'tool', toolName }`) or a resolver
+ * that runs inside `inferTools()` and returns one of those (or `undefined` to leave
+ * it unconstrained). `inferTools()` returns `toolChoice` alongside `activeTools`, and
+ * both spread straight into `generateText` and into `prepareStep`'s return.
+ *
+ * Here the resolver forces `search` on the first step (so the run always starts by
+ * gathering data), then hands control back to the model with `'auto'`. The toolChoice
+ * is recomputed per step from the run's `steps`.
+ *
+ * Uses ai-test-kit's MockLanguageModel so it runs without any API keys: the mock
+ * honors the `toolChoice` it receives — it calls `search` while forced, then answers.
+ *
+ * Run: pnpm tsx examples/tool-choice.ts
+ */
+import { generateText, stepCountIs, tool } from 'ai';
+import { Language, MockLanguageModel } from 'ai-test-kit/language';
+import { z } from 'zod';
+import { createToolSet } from '../src/tool-set.js';
+
+const search = tool({
+  description: 'Search the web for information',
+  inputSchema: z.object({ query: z.string() }),
+  execute: async ({ query }) => ({ results: [`result for "${query}"`] }),
+});
+
+const answer = tool({
+  description: 'Give the final answer to the user',
+  inputSchema: z.object({ text: z.string() }),
+  execute: async ({ text }) => ({ answered: text }),
+});
+
+const tools = { search, answer };
+
+/**
+ * Force `search` on the first step (no steps completed yet), then let the model
+ * decide for every step after. Returning `'auto'` is the explicit fallback —
+ * with last-call wins, a resolver returning `undefined` would simply omit
+ * `toolChoice`, which also defaults to `'auto'`.
+ */
+const toolSet = createToolSet({ tools }).choice(({ steps }) =>
+  steps?.length ? 'auto' : { type: 'tool', toolName: 'search' },
+);
+
+/** Mock model (ai-test-kit): obey the forced toolChoice, otherwise produce final text. */
+let searchCallId = 0;
+const model = MockLanguageModel.from({
+  doGenerate: async ({ toolChoice }) => {
+    const forcedTool = toolChoice?.type === 'tool' ? toolChoice.toolName : undefined;
+    if (forcedTool === 'search') {
+      return Language.result([
+        Language.toolCall({
+          toolCallId: `call-${++searchCallId}`,
+          toolName: 'search',
+          input: { query: 'ai sdk tool choice' },
+        }),
+      ]);
+    }
+    return Language.result('Here is my answer based on what I found.');
+  },
+});
+
+const result = await generateText({
+  model,
+  tools,
+  stopWhen: stepCountIs(10),
+  prompt: 'Research the topic and then answer.',
+  prepareStep: ({ stepNumber, steps }) => {
+    /** Recompute activeTools + toolChoice for this step from the run so far. */
+    const { activeTools, toolChoice } = toolSet.inferTools({ steps });
+    console.log(
+      `step ${stepNumber} → toolChoice: ${JSON.stringify(toolChoice)} → activeTools: [${activeTools.join(', ')}]`,
+    );
+    return { activeTools, toolChoice };
+  },
+});
+
+console.log('\nTool-call trace:');
+for (const [i, step] of result.steps.entries()) {
+  for (const call of step.toolCalls) {
+    console.log(`  step ${i} → called ${call.toolName}(${JSON.stringify(call.input)})`);
+  }
+}
+
+console.log('\nFinal text:', result.text);

src/__tests__/tool-set.test-d.ts

@@ -1,4 +1,4 @@
-import type { InferUITool, ModelMessage, StepResult, UIMessage } from 'ai';
+import type { InferUITool, ModelMessage, PrepareStepResult, StepResult, UIMessage } from 'ai';
 import { tool } from 'ai';
 import { describe, expectTypeOf, test } from 'vitest';
 import { z } from 'zod';
@@ -407,6 +407,62 @@ describe('immutable toolset', () => {
     });
   });
 
+  describe('choice', () => {
+    test('should accept the string toolChoice values', () => {
+      const toolSet = createToolSet({ tools: TOOLS });
+      toolSet.choice('auto');
+      toolSet.choice('none');
+      toolSet.choice('required');
+    });
+
+    test('should accept a tool toolChoice with a known tool name', () => {
+      const toolSet = createToolSet({ tools: TOOLS });
+      toolSet.choice({ type: 'tool', toolName: 'plain' });
+      toolSet.choice({ type: 'tool', toolName: 'cancel' });
+    });
+
+    test('should reject a tool toolChoice with an unknown tool name', () => {
+      const toolSet = createToolSet({ tools: TOOLS });
+      // @ts-expect-error — 'unknown' is not in TOOLS
+      toolSet.choice({ type: 'tool', toolName: 'unknown' });
+    });
+
+    test('should type the resolver input as the inferTools input', () => {
+      type MyCtx = { isAdmin: boolean };
+      const toolSet = createToolSet<typeof TOOLS, UIMessage, MyCtx>({ tools: TOOLS });
+      toolSet.choice((input) => {
+        expectTypeOf(input.context).toEqualTypeOf<MyCtx | undefined>();
+        expectTypeOf(input.messages).toEqualTypeOf<Array<UIMessage> | undefined>();
+        return 'required';
+      });
+    });
+
+    test('should allow the resolver to return undefined', () => {
+      const toolSet = createToolSet({ tools: TOOLS });
+      toolSet.choice(({ steps }) => (steps?.length ? 'auto' : undefined));
+    });
+
+    test('should type toolChoice on the inferTools result', () => {
+      const toolSet = createToolSet({ tools: TOOLS });
+      const result = toolSet.inferTools();
+      expectTypeOf(result.toolChoice).toEqualTypeOf<
+        | 'auto'
+        | 'none'
+        | 'required'
+        | { type: 'tool'; toolName: 'plain' | 'calc' | 'cancel' | 'edit' | 'archive' }
+        | undefined
+      >();
+    });
+
+    test('should spread the inferTools result into a prepareStep return', () => {
+      const toolSet = createToolSet({ tools: TOOLS }).choice('required');
+      // The resolved `activeTools` + `toolChoice` are accepted by PrepareStepResult.
+      const prepareStep = ({ steps }: { steps: Array<StepResult<typeof TOOLS>> }): PrepareStepResult<typeof TOOLS> =>
+        toolSet.inferTools({ steps });
+      expectTypeOf(prepareStep).toBeFunction();
+    });
+  });
+
   describe('clone', () => {
     test('should return immutable toolset by default', () => {
       const toolSet = createToolSet({ tools: TOOLS });
@@ -557,6 +613,20 @@ describe('mutable toolset', () => {
     });
   });
 
+  describe('choice', () => {
+    test('should return this', () => {
+      const toolSet = createToolSet({ tools: TOOLS, mutable: true });
+      const result = toolSet.choice('required');
+      expectTypeOf(result).toEqualTypeOf(toolSet);
+    });
+
+    test('should reject a tool toolChoice with an unknown tool name', () => {
+      const toolSet = createToolSet({ tools: TOOLS, mutable: true });
+      // @ts-expect-error — 'unknown' is not in TOOLS
+      toolSet.choice({ type: 'tool', toolName: 'unknown' });
+    });
+  });
+
   describe('clone', () => {
     test('should return immutable toolset by default', () => {
       const toolSet = createToolSet({ tools: TOOLS, mutable: true });