diff --git a/.changeset/next-turn-tool-choice.md b/.changeset/next-turn-tool-choice.md new file mode 100644 index 00000000..c388f206 --- /dev/null +++ b/.changeset/next-turn-tool-choice.md @@ -0,0 +1,40 @@ +--- +"@openrouter/agent": minor +--- + +Add `toolChoice` to `nextTurnParams`, so a tool can change which tools the model may call on the following turn without touching the `tools` array. + +This is what a tool-search tool needs: declare every tool up front, keep the not-yet-needed ones out of reach behind an `allowed_tools` choice, and widen that choice as the model discovers what it wants. Because `tools` is byte-identical across turns, the provider's prompt-cache prefix survives — which is the whole reason to withhold tools rather than send them all. + +```ts +import { callModel, OpenRouter, tool } from '@openrouter/agent'; +import { z } from 'zod/v4'; + +const allowed = (names: string[]) => ({ + type: 'allowed_tools' as const, + mode: 'auto' as const, + tools: names.map((name) => ({ type: 'function', name })), +}); + +const toolSearch = tool({ + name: 'tool_search', + inputSchema: z.object({ pattern: z.string() }), + execute: ({ pattern }) => findMatchingToolNames(pattern), + nextTurnParams: { + // Append, never rebuild: dropping a name revokes a tool the model may + // already have used, and reordering churns the request for nothing. + toolChoice: ({ pattern }, context) => + allowed([...namesIn(context.toolChoice), ...findMatchingToolNames(pattern)]), + }, +}); + +const client = new OpenRouter({ apiKey: process.env['OPENROUTER_API_KEY'] }); + +const result = callModel(client, { + model: 'openai/gpt-4o-mini', + input: 'What is the weather in Tokyo?', + tools: [toolSearch, getWeather, sendEmail, listRepos], + // Only the search tool is reachable until it finds something. + toolChoice: allowed(['tool_search']), +}); +``` diff --git a/packages/agent/src/lib/model-result.ts b/packages/agent/src/lib/model-result.ts index 17bd4928..c23ea3ac 100644 --- a/packages/agent/src/lib/model-result.ts +++ b/packages/agent/src/lib/model-result.ts @@ -4386,9 +4386,25 @@ export class ModelResult< this.resolvedRequest, ); - if (Object.keys(computedParams).length > 0) { - this.resolvedRequest = applyNextTurnParamsToRequest(this.resolvedRequest, computedParams); + if (Object.keys(computedParams).length === 0) { + return; } + + const nextRequest = applyNextTurnParamsToRequest(this.resolvedRequest, computedParams); + + /* + * A tool-computed `toolChoice` becomes the new caller-level policy, not a + * one-turn override. Merging it onto the request alone is not enough: + * `makeFollowupRequest` re-derives the wire choice from + * `configuredToolChoice` via `applyForcedToolChoicePolicy`, which would + * discard the tool's value before dispatch. Re-running the resolved-policy + * bookkeeping re-stamps the configured choice and its forced-choice + * consumption key together, so relaxation stays consistent for later turns. + */ + this.resolvedRequest = + 'toolChoice' in computedParams + ? this.applyResolvedForcedToolChoicePolicy(nextRequest) + : nextRequest; } /** diff --git a/packages/agent/src/lib/next-turn-params.test.ts b/packages/agent/src/lib/next-turn-params.test.ts index b6a2a8be..80a57b04 100644 --- a/packages/agent/src/lib/next-turn-params.test.ts +++ b/packages/agent/src/lib/next-turn-params.test.ts @@ -112,3 +112,89 @@ describe('applyNextTurnParamsToRequest', () => { expect(result.instructions).toBe(''); }); }); + +/* + * `toolChoice` is how a tool-search tool widens the model's reach mid-run: the + * `tools` array stays byte-identical (preserving the provider's prompt-cache + * prefix) while `{ type: 'allowed_tools', tools: [...] }` grows. + */ +describe('applyNextTurnParamsToRequest with allowed_tools', () => { + const allowed = (names: string[]): models.ResponsesRequest['toolChoice'] => ({ + type: 'allowed_tools', + mode: 'auto', + tools: names.map((name) => ({ + type: 'function', + name, + })), + }); + + it('replaces toolChoice with a widened allowed_tools set', () => { + const request = createBaseRequest({ + toolChoice: allowed([ + 'tool_search', + ]), + }); + + const result = applyNextTurnParamsToRequest(request, { + toolChoice: allowed([ + 'tool_search', + 'get_weather', + ]), + }); + + expect(result.toolChoice).toEqual( + allowed([ + 'tool_search', + 'get_weather', + ]), + ); + }); + + it('leaves the tools array untouched so the prompt-cache prefix survives', () => { + const tools = [ + { + type: 'function' as const, + name: 'tool_search', + parameters: {}, + }, + { + type: 'function' as const, + name: 'get_weather', + parameters: {}, + }, + ]; + const request = createBaseRequest({ + tools, + toolChoice: allowed([ + 'tool_search', + ]), + }); + + const result = applyNextTurnParamsToRequest(request, { + toolChoice: allowed([ + 'tool_search', + 'get_weather', + ]), + }); + + expect(result.tools).toBe(tools); + }); + + it('leaves toolChoice alone when no tool computed one', () => { + const request = createBaseRequest({ + toolChoice: allowed([ + 'tool_search', + ]), + }); + + const result = applyNextTurnParamsToRequest(request, { + temperature: 0.2, + }); + + expect(result.toolChoice).toEqual( + allowed([ + 'tool_search', + ]), + ); + }); +}); diff --git a/packages/agent/src/lib/next-turn-params.ts b/packages/agent/src/lib/next-turn-params.ts index c9e813d2..7ba0807b 100644 --- a/packages/agent/src/lib/next-turn-params.ts +++ b/packages/agent/src/lib/next-turn-params.ts @@ -21,6 +21,7 @@ export function buildNextTurnParamsContext( ): NextTurnParamsContext { return { input: request.input ?? [], + toolChoice: request.toolChoice, model: request.model ?? '', models: request.models ?? [], temperature: request.temperature ?? null, @@ -115,7 +116,7 @@ async function processNextTurnParamsForCall( if (process.env['NODE_ENV'] !== 'production') { console.warn( `Invalid nextTurnParams key "${paramKey}" in tool "${toolName}". ` + - 'Valid keys: input, model, models, temperature, maxOutputTokens, topP, topK, instructions', + 'Valid keys: input, toolChoice, model, models, temperature, maxOutputTokens, topP, topK, instructions', ); } continue; @@ -137,6 +138,7 @@ async function processNextTurnParamsForCall( function isValidNextTurnParamKey(key: string): key is keyof NextTurnParamsContext { const validKeys: ReadonlySet = new Set([ 'input', + 'toolChoice', 'model', 'models', 'temperature', diff --git a/packages/agent/src/lib/tool-types.ts b/packages/agent/src/lib/tool-types.ts index e89d7490..ee80fb6f 100644 --- a/packages/agent/src/lib/tool-types.ts +++ b/packages/agent/src/lib/tool-types.ts @@ -287,6 +287,17 @@ export const SHARED_CONTEXT_KEY = 'shared' as const; export type NextTurnParamsContext = { /** Current input (messages) */ input: models.InputsUnion; + /** + * Current tool choice. + * + * Returning a new value changes which tools the model may call next turn + * without touching the `tools` array — the hook a tool-search tool uses to + * widen `{ type: 'allowed_tools', tools: [...] }` once it has found what it + * was looking for. Leaving `tools` alone is the point: rewriting it would + * invalidate the provider's prompt-cache prefix, which is usually the reason + * the caller is withholding tools in the first place. + */ + toolChoice: models.ResponsesRequest['toolChoice']; /** Current model selection */ model: string; /** Current models array */ diff --git a/packages/agent/tests/unit/next-turn-tool-choice-loop.test.ts b/packages/agent/tests/unit/next-turn-tool-choice-loop.test.ts new file mode 100644 index 00000000..f96a1d15 --- /dev/null +++ b/packages/agent/tests/unit/next-turn-tool-choice-loop.test.ts @@ -0,0 +1,249 @@ +/** + * `nextTurnParams.toolChoice` driven through the real callModel loop, with + * betaResponsesSend mocked at the module level. + * + * Asserted against the DISPATCHED request rather than the helper in isolation: + * `makeFollowupRequest` re-derives the wire tool choice from the caller- + * configured value after `applyNextTurnParams` runs, so a hook that merges + * correctly into `resolvedRequest` can still be discarded before dispatch. + * Only a loop-level assertion catches that. + */ +import type * as models from '@openrouter/sdk/models'; +import { afterEach, describe, expect, it, vi } from 'vitest'; +import { z } from 'zod/v4'; + +const mockBetaResponsesSend = vi.hoisted(() => vi.fn()); + +vi.mock('@openrouter/sdk/funcs/betaResponsesSend', () => ({ + betaResponsesSend: mockBetaResponsesSend, +})); + +import type { OpenRouterCore } from '@openrouter/sdk/core'; +import { callModel } from '../../src/inner-loop/call-model.js'; +import { tool } from '../../src/lib/tool.js'; + +afterEach(() => { + mockBetaResponsesSend.mockReset(); + vi.restoreAllMocks(); +}); + +const client = {} as unknown as OpenRouterCore; + +function allowed(names: string[]): models.ResponsesRequest['toolChoice'] { + return { + type: 'allowed_tools', + mode: 'auto', + tools: names.map((name) => ({ + type: 'function', + name, + })), + } as unknown as models.ResponsesRequest['toolChoice']; +} + +function toolCallResponse(): models.OpenResponsesResult { + return { + id: 'resp_tool', + output: [ + { + type: 'function_call', + id: 'out_1', + callId: 'call_1', + name: 'tool_search', + arguments: '{"pattern":"weather"}', + status: 'completed', + }, + ], + } as unknown as models.OpenResponsesResult; +} + +function textResponse(): models.OpenResponsesResult { + return { + id: 'resp_text', + output: [ + { + type: 'message', + id: 'msg_1', + role: 'assistant', + content: [ + { + type: 'output_text', + text: 'done', + }, + ], + status: 'completed', + }, + ], + } as unknown as models.OpenResponsesResult; +} + +/** The tool choice on each request actually handed to the SDK. */ +function dispatchedToolChoices(): unknown[] { + return mockBetaResponsesSend.mock.calls.map( + (call) => + ( + call[1] as { + responsesRequest: models.ResponsesRequest; + } + ).responsesRequest.toolChoice, + ); +} + +function makeSearchTool(widened: string[]) { + return tool({ + name: 'tool_search', + inputSchema: z.object({ + pattern: z.string(), + }), + execute: () => ({ + found: widened, + }), + nextTurnParams: { + toolChoice: () => allowed(widened), + }, + }); +} + +const getWeather = tool({ + name: 'get_weather', + inputSchema: z.object({}), + execute: () => ({ + ok: true, + }), +}); + +describe('nextTurnParams.toolChoice through the callModel loop', () => { + it('reaches the dispatched follow-up request', async () => { + mockBetaResponsesSend + .mockResolvedValueOnce({ + ok: true, + value: toolCallResponse(), + }) + .mockResolvedValue({ + ok: true, + value: textResponse(), + }); + + await callModel(client, { + model: 'test-model', + input: 'hi', + tools: [ + makeSearchTool([ + 'tool_search', + 'get_weather', + ]), + getWeather, + ], + toolChoice: allowed([ + 'tool_search', + ]), + }).getText(); + + const choices = dispatchedToolChoices(); + expect(choices.length).toBeGreaterThanOrEqual(2); + + // Turn 1 uses the caller's narrow set. + expect(choices[0]).toEqual( + allowed([ + 'tool_search', + ]), + ); + + // Turn 2 must carry the tool's widened set, not the caller's original. + expect(choices[1]).toEqual( + allowed([ + 'tool_search', + 'get_weather', + ]), + ); + }); + + it('leaves the tools array untouched while the choice widens', async () => { + mockBetaResponsesSend + .mockResolvedValueOnce({ + ok: true, + value: toolCallResponse(), + }) + .mockResolvedValue({ + ok: true, + value: textResponse(), + }); + + await callModel(client, { + model: 'test-model', + input: 'hi', + tools: [ + makeSearchTool([ + 'tool_search', + 'get_weather', + ]), + getWeather, + ], + toolChoice: allowed([ + 'tool_search', + ]), + }).getText(); + + const toolNames = mockBetaResponsesSend.mock.calls.map((call) => + ( + ( + call[1] as { + responsesRequest: models.ResponsesRequest; + } + ).responsesRequest.tools ?? [] + ).map( + (t) => + ( + t as { + name?: string; + } + ).name, + ), + ); + + /* Identical across turns: widening must not change the request prefix, or + * the provider's prompt cache is lost — the reason to withhold tools at all. */ + expect(toolNames[1]).toEqual(toolNames[0]); + }); + + it('does not disturb toolChoice when no tool computes one', async () => { + const plain = tool({ + name: 'tool_search', + inputSchema: z.object({ + pattern: z.string(), + }), + execute: () => ({ + ok: true, + }), + }); + + mockBetaResponsesSend + .mockResolvedValueOnce({ + ok: true, + value: toolCallResponse(), + }) + .mockResolvedValue({ + ok: true, + value: textResponse(), + }); + + await callModel(client, { + model: 'test-model', + input: 'hi', + tools: [ + plain, + getWeather, + ], + toolChoice: allowed([ + 'tool_search', + ]), + }).getText(); + + for (const choice of dispatchedToolChoices()) { + expect(choice).toEqual( + allowed([ + 'tool_search', + ]), + ); + } + }); +});