From e418f91a7af34d2dcfbb35e6a2c8bf9c5ed4c47e Mon Sep 17 00:00:00 2001 From: luvs01 Date: Thu, 20 Aug 2026 22:31:23 +0900 Subject: [PATCH] fix(tools): index tool-choice candidates --- src/adapters/anthropic.ts | 9 +-- src/adapters/command-code.ts | 5 +- src/adapters/google.ts | 9 +-- src/responses/parser.ts | 5 +- src/types.ts | 2 +- src/types/tools.ts | 98 ++++++++++++++++++++++++--- tests/tool-choice-performance.test.ts | 84 +++++++++++++++++++++++ tests/types-barrel-identity.test.ts | 1 + 8 files changed, 184 insertions(+), 29 deletions(-) create mode 100644 tests/tool-choice-performance.test.ts diff --git a/src/adapters/anthropic.ts b/src/adapters/anthropic.ts index 626a45f286..7d49e7ca70 100644 --- a/src/adapters/anthropic.ts +++ b/src/adapters/anthropic.ts @@ -14,7 +14,7 @@ import type { OcxToolResultMessage, OcxUsage, } from "../types"; -import { isAllowedToolChoice, namespacedToolName, resolveToolChoiceWireName, toolAllowedByChoice } from "../types"; +import { isAllowedToolChoice, namespacedToolName, resolveToolChoiceWireName, toolChoiceToolPredicate } from "../types"; import { ANTHROPIC_OAUTH_BETA, CLAUDE_CODE_SYSTEM_INSTRUCTION, applyClaudeToolPrefix, stripClaudeToolPrefix } from "../oauth/anthropic"; import { parseDataUrl } from "./image"; import { enforceAnthropicImageLimits } from "./anthropic-image-guard"; @@ -796,11 +796,8 @@ function messagesToAnthropicFormat( function toolsToAnthropicFormat(parsed: OcxParsedRequest, toolNames: { toWire: (name: string) => string }): unknown[] | undefined { if (!parsed.context.tools || parsed.context.tools.length === 0) return undefined; - const allowed = isAllowedToolChoice(parsed.options.toolChoice) - ? new Set(parsed.options.toolChoice.allowedTools) - : undefined; - const tools = allowed - ? parsed.context.tools.filter(t => toolAllowedByChoice(t, allowed, parsed.context.tools)) + const tools = isAllowedToolChoice(parsed.options.toolChoice) + ? parsed.context.tools.filter(toolChoiceToolPredicate(parsed.options.toolChoice, parsed.context.tools)) : parsed.context.tools; if (tools.length === 0) return undefined; const converted = tools.map(t => ({ diff --git a/src/adapters/command-code.ts b/src/adapters/command-code.ts index 8f2edbdb55..1cd7975d83 100644 --- a/src/adapters/command-code.ts +++ b/src/adapters/command-code.ts @@ -3,7 +3,7 @@ import { execFile as execFileCallback } from "node:child_process"; import { promisify } from "node:util"; import { opendir } from "node:fs/promises"; import type { AdapterEvent, OcxContentPart, OcxMessage, OcxParsedRequest, OcxProviderConfig, OcxTool, OcxUsage } from "../types"; -import { isAllowedToolChoice, namespacedToolName, resolveToolChoiceWireName, toolAllowedByChoice } from "../types"; +import { isAllowedToolChoice, namespacedToolName, resolveToolChoiceWireName, toolChoiceToolPredicate } from "../types"; import type { AdapterFetchContext, AdapterRequest, ProviderAdapter } from "./base"; import type { TranslatorBudget } from "../lib/translator-budget"; import { readBoundedResponseBody } from "../lib/bounded-body"; @@ -156,8 +156,7 @@ function visibleTools(parsed: OcxParsedRequest): OcxTool[] { if (choice === "none") return []; const tools = parsed.context.tools ?? []; if (isAllowedToolChoice(choice)) { - const allowed = new Set(choice.allowedTools); - return tools.filter(tool => toolAllowedByChoice(tool, allowed, tools)); + return tools.filter(toolChoiceToolPredicate(choice, tools)); } if (choice && typeof choice !== "string") { const selected = resolveToolChoiceWireName(tools, choice.name); diff --git a/src/adapters/google.ts b/src/adapters/google.ts index 746f9490a4..233d0ed5d1 100644 --- a/src/adapters/google.ts +++ b/src/adapters/google.ts @@ -14,7 +14,7 @@ import type { OcxToolResultMessage, OcxUsage, } from "../types"; -import { isAllowedToolChoice, namespacedToolName, resolveToolChoiceWireName, toolAllowedByChoice } from "../types"; +import { isAllowedToolChoice, namespacedToolName, resolveToolChoiceWireName, toolChoiceToolPredicate } from "../types"; import { contentPartsToText, parseDataUrl } from "./image"; import { getVertexAccessToken } from "../lib/gcp-adc"; import { fetchAntigravityWithRetry, fetchVertexWithRetry } from "./google-http"; @@ -328,11 +328,8 @@ function messagesToGeminiFormat( function toolsToGeminiFormat(parsed: OcxParsedRequest): unknown[] | undefined { if (!parsed.context.tools?.length) return undefined; - const allowed = isAllowedToolChoice(parsed.options.toolChoice) - ? new Set(parsed.options.toolChoice.allowedTools) - : undefined; - const tools = allowed - ? parsed.context.tools.filter(t => toolAllowedByChoice(t, allowed, parsed.context.tools)) + const tools = isAllowedToolChoice(parsed.options.toolChoice) + ? parsed.context.tools.filter(toolChoiceToolPredicate(parsed.options.toolChoice, parsed.context.tools)) : parsed.context.tools; if (tools.length === 0) return undefined; return [{ diff --git a/src/responses/parser.ts b/src/responses/parser.ts index de07832a40..28b7c8c2bb 100644 --- a/src/responses/parser.ts +++ b/src/responses/parser.ts @@ -11,7 +11,7 @@ import type { OcxToolCall, OcxReasoningReplayScopeRef, } from "../types"; -import { namespacedToolName, toolChoiceCandidates } from "../types"; +import { createToolChoiceResolver, namespacedToolName } from "../types"; import { responsesRequestSchema } from "./schema"; import { providerMetadataFromResponsesFunctionCall } from "./provider-opaque-metadata"; import { lookupReplayThoughtSignature } from "./thought-signature-replay"; @@ -758,8 +758,9 @@ export function parseRequest( const tc = mapToolChoice(data.tool_choice); if (tc && typeof tc === "object") { const selectors = "allowedTools" in tc ? tc.allowedTools : [tc.name]; + const resolver = createToolChoiceResolver(mergedTools); for (const selector of selectors) { - if (toolChoiceCandidates(mergedTools, selector).length > 1) { + if (resolver.candidateCount(selector) > 1) { throw new Error(`ambiguous tool_choice name: ${selector}`); } } diff --git a/src/types.ts b/src/types.ts index cf045ca891..559e71cbc9 100644 --- a/src/types.ts +++ b/src/types.ts @@ -5,6 +5,7 @@ export type { OcxTool, OcxToolChoice } from "./types/tools"; export { namespacedToolName, toolChoiceAliases, + createToolChoiceResolver, toolChoiceCandidates, toolAllowedByChoice, resolveToolChoiceWireName, @@ -102,4 +103,3 @@ export type { CodexAccountCredentials, CodexAccountCredentialRecord, } from "./types/accounts"; - diff --git a/src/types/tools.ts b/src/types/tools.ts index 9e3dc37fd0..5e4f4547a0 100644 --- a/src/types/tools.ts +++ b/src/types/tools.ts @@ -43,6 +43,74 @@ function sameToolIdentity( return left.namespace === right.namespace && left.name === right.name; } +type ToolIdentity = Readonly>; + +function snapshotToolIdentity(tool: Pick): ToolIdentity { + return Object.freeze({ + name: tool.name, + ...(tool.namespace === undefined ? {} : { namespace: tool.namespace }), + }); +} + +function buildToolChoiceCatalog( + tools: readonly ToolIdentity[], +): { + candidatesByName: ReadonlyMap; + sourceCandidatesByName: ReadonlyMap; + identitiesByTool: WeakMap; +} { + const index = new Map(); + const sourceIndex = new Map(); + const identities = new Map>(); + const identitiesByTool = new WeakMap(); + for (const tool of tools) { + const snapshot = snapshotToolIdentity(tool); + identitiesByTool.set(tool, snapshot); + const identity = JSON.stringify([snapshot.namespace ?? null, snapshot.name]); + for (const selector of [...toolChoiceAliases(snapshot), snapshot.name]) { + const candidates = index.get(selector); + if (!candidates) { + index.set(selector, [snapshot]); + sourceIndex.set(selector, [tool]); + identities.set(selector, new Set([identity])); + } else if (!identities.get(selector)!.has(identity)) { + candidates.push(snapshot); + sourceIndex.get(selector)!.push(tool); + identities.get(selector)!.add(identity); + } + } + } + return { candidatesByName: index, sourceCandidatesByName: sourceIndex, identitiesByTool }; +} + +/** Compile one immutable view of a request's tool catalog for repeated policy checks. */ +export function createToolChoiceResolver(tools: readonly ToolIdentity[] | undefined) { + const compiled = tools ? buildToolChoiceCatalog(tools) : undefined; + const candidatesByName = compiled?.candidatesByName; + const snapshotFor = (tool: ToolIdentity): ToolIdentity | undefined => { + const snapshot = compiled?.identitiesByTool.get(tool); + return snapshot && sameToolIdentity(snapshot, tool) ? snapshot : undefined; + }; + return { + candidates(name: string): ToolIdentity[] { + return (candidatesByName?.get(name) ?? []).map(candidate => ({ ...candidate })); + }, + candidateCount(name: string): number { + return candidatesByName?.get(name)?.length ?? 0; + }, + allows(tool: ToolIdentity, allowedTools: ReadonlySet): boolean { + if (!candidatesByName) return toolChoiceAliases(tool).some(name => allowedTools.has(name)); + const snapshot = snapshotFor(tool); + return snapshot ? toolAllowedByChoiceFromIndex(snapshot, allowedTools, candidatesByName) : false; + }, + selects(tool: ToolIdentity, name: string): boolean { + const snapshot = snapshotFor(tool); + const candidates = candidatesByName?.get(name); + return !!snapshot && candidates?.length === 1 && sameToolIdentity(candidates[0], snapshot); + }, + }; +} + /** * All tools that could be selected by one client-facing name. Bare logical names are included * here because they are a compatibility selector for namespaced tools, while wire and dotted @@ -53,12 +121,7 @@ export function toolChoiceCandidates( name: string, ): Pick[] { if (!tools) return []; - const candidates: Pick[] = []; - for (const tool of tools) { - if (tool.name !== name && !toolChoiceAliases(tool).includes(name)) continue; - if (!candidates.some(candidate => sameToolIdentity(candidate, tool))) candidates.push(tool); - } - return candidates; + return [...(buildToolChoiceCatalog(tools).sourceCandidatesByName.get(name) ?? [])]; } /** @@ -72,10 +135,22 @@ export function toolAllowedByChoice( tools?: readonly Pick[], ): boolean { if (!tools) return toolChoiceAliases(tool).some(name => allowedTools.has(name)); + return toolAllowedByChoiceFromIndex( + snapshotToolIdentity(tool), + allowedTools, + buildToolChoiceCatalog(tools).candidatesByName, + ); +} + +function toolAllowedByChoiceFromIndex( + tool: ToolIdentity, + allowedTools: ReadonlySet, + candidatesByName: ReadonlyMap, +): boolean { for (const name of [...toolChoiceAliases(tool), tool.name]) { if (!allowedTools.has(name)) continue; - const candidates = toolChoiceCandidates(tools, name); - if (candidates.length === 1 && sameToolIdentity(candidates[0], tool)) return true; + const candidates = candidatesByName.get(name); + if (candidates?.length === 1 && sameToolIdentity(candidates[0], tool)) return true; } return false; } @@ -123,9 +198,10 @@ export function toolChoiceToolPredicate( if (choice === "none") return () => false; if (isAllowedToolChoice(choice)) { const allowed = new Set(choice.allowedTools); - return tool => toolAllowedByChoice(tool, allowed, tools); + const resolver = createToolChoiceResolver(tools); + return tool => resolver.allows(tool, allowed); } if (!tools) return tool => toolChoiceAliases(tool).includes(choice.name); - const candidates = toolChoiceCandidates(tools, choice.name); - return tool => candidates.length === 1 && sameToolIdentity(candidates[0], tool); + const resolver = createToolChoiceResolver(tools); + return tool => resolver.selects(tool, choice.name); } diff --git a/tests/tool-choice-performance.test.ts b/tests/tool-choice-performance.test.ts new file mode 100644 index 0000000000..d1d99dee3c --- /dev/null +++ b/tests/tool-choice-performance.test.ts @@ -0,0 +1,84 @@ +import { expect, test } from "bun:test"; +import type { OcxTool } from "../src/types"; +import { toolChoiceCandidates, toolChoiceToolPredicate } from "../src/types"; + +test("allowed_tools resolves an ambiguous bare name without replaying the candidate list", () => { + const size = 256; + const backingTools: OcxTool[] = Array.from({ length: size }, (_, index) => ({ + namespace: `namespace_${index}`, + name: "shared_name", + description: "", + parameters: {}, + })); + let catalogReads = 0; + const tools = new Proxy(backingTools, { + get(target, property, receiver) { + if (typeof property === "string" && /^\d+$/.test(property)) catalogReads += 1; + return Reflect.get(target, property, receiver); + }, + }); + const arrayIterator = Object.getOwnPropertyDescriptor(Array.prototype, Symbol.iterator)!; + let ambiguousCandidateIterations = 0; + Object.defineProperty(Array.prototype, Symbol.iterator, { + ...arrayIterator, + value: function (this: unknown[]) { + const iterator = arrayIterator.value.call(this) as IterableIterator; + const tracksCandidateList = this !== backingTools + && this !== tools + && this.length === size + && this[0] === backingTools[0] + && this[size - 1] === backingTools[size - 1]; + return { + next() { + const result = iterator.next(); + if (tracksCandidateList && !result.done) ambiguousCandidateIterations += 1; + return result; + }, + [Symbol.iterator]() { + return this; + }, + }; + }, + }); + + let filtered: OcxTool[]; + try { + const allowed = toolChoiceToolPredicate( + { allowedTools: ["shared_name"], mode: "auto" }, + tools, + ); + filtered = tools.filter(allowed); + } finally { + Object.defineProperty(Array.prototype, Symbol.iterator, arrayIterator); + } + + expect(filtered).toEqual([]); + expect(catalogReads).toBeLessThanOrEqual(size * 2); + expect(ambiguousCandidateIterations).toBe(0); +}); + +test("public candidate lookups rebuild after a mutable caller changes its catalog", () => { + const tools: OcxTool[] = [{ namespace: "first", name: "shared", description: "", parameters: {} }]; + const firstLookup = toolChoiceCandidates(tools, "shared"); + expect(firstLookup).toHaveLength(1); + expect(firstLookup[0]).toBe(tools[0]); + + tools.push({ namespace: "second", name: "shared", description: "", parameters: {} }); + + const secondLookup = toolChoiceCandidates(tools, "shared"); + expect(secondLookup).toHaveLength(2); + expect(secondLookup[0]).toBe(tools[0]); + expect(secondLookup[1]).toBe(tools[1]); + const allowed = toolChoiceToolPredicate({ allowedTools: ["shared"], mode: "auto" }, tools); + expect(tools.filter(allowed)).toEqual([]); +}); + +test("a compiled resolver fails closed when its catalog objects change", () => { + const tools: OcxTool[] = [{ namespace: "stable", name: "shared", description: "", parameters: {} }]; + const allowed = toolChoiceToolPredicate({ allowedTools: ["shared"], mode: "auto" }, tools); + + tools[0].name = "changed_after_compile"; + tools.push({ namespace: "late", name: "shared", description: "", parameters: {} }); + + expect(tools.filter(allowed)).toEqual([]); +}); diff --git a/tests/types-barrel-identity.test.ts b/tests/types-barrel-identity.test.ts index 7508d45402..975c6650cc 100644 --- a/tests/types-barrel-identity.test.ts +++ b/tests/types-barrel-identity.test.ts @@ -25,6 +25,7 @@ describe("types barrel re-exports the leaves by identity, not by copy", () => { test.each([ "namespacedToolName", "toolChoiceAliases", + "createToolChoiceResolver", "toolChoiceCandidates", "toolAllowedByChoice", "resolveToolChoiceWireName",