From c2d4ea45d5de92e0f821c89659adda3c1c4cdf2d Mon Sep 17 00:00:00 2001
From: luvs01 <luvs01@hanmail.net>
Date: Thu, 20 Aug 2026 22:31:23 +0900
Subject: [PATCH] fix(tools): index tool-choice candidates

---
 src/types/tools.ts                    | 39 ++++++++++++++++++++++-----
 tests/tool-choice-performance.test.ts | 28 +++++++++++++++++++
 2 files changed, 61 insertions(+), 6 deletions(-)
 create mode 100644 tests/tool-choice-performance.test.ts

diff --git a/src/types/tools.ts b/src/types/tools.ts
index 9e3dc37fd0..553fdb9260 100644
--- a/src/types/tools.ts
+++ b/src/types/tools.ts
@@ -43,6 +43,38 @@ function sameToolIdentity(
   return left.namespace === right.namespace && left.name === right.name;
 }
 
+type ToolIdentity = Pick<OcxTool, "namespace" | "name">;
+
+const toolChoiceCandidateIndexes = new WeakMap<
+  readonly ToolIdentity[],
+  ReadonlyMap<string, readonly ToolIdentity[]>
+>();
+
+function toolChoiceCandidateIndex(
+  tools: readonly ToolIdentity[],
+): ReadonlyMap<string, readonly ToolIdentity[]> {
+  const cached = toolChoiceCandidateIndexes.get(tools);
+  if (cached) return cached;
+
+  const index = new Map<string, ToolIdentity[]>();
+  const identities = new Map<string, Set<string>>();
+  for (const tool of tools) {
+    const identity = JSON.stringify([tool.namespace ?? null, tool.name]);
+    for (const selector of [...toolChoiceAliases(tool), tool.name]) {
+      const candidates = index.get(selector);
+      if (!candidates) {
+        index.set(selector, [tool]);
+        identities.set(selector, new Set([identity]));
+      } else if (!identities.get(selector)!.has(identity)) {
+        candidates.push(tool);
+        identities.get(selector)!.add(identity);
+      }
+    }
+  }
+  toolChoiceCandidateIndexes.set(tools, index);
+  return index;
+}
+
 /**
  * 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 +85,7 @@ export function toolChoiceCandidates(
   name: string,
 ): Pick<OcxTool, "namespace" | "name">[] {
   if (!tools) return [];
-  const candidates: Pick<OcxTool, "namespace" | "name">[] = [];
-  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 [...(toolChoiceCandidateIndex(tools).get(name) ?? [])];
 }
 
 /**
diff --git a/tests/tool-choice-performance.test.ts b/tests/tool-choice-performance.test.ts
new file mode 100644
index 0000000000..95c086609c
--- /dev/null
+++ b/tests/tool-choice-performance.test.ts
@@ -0,0 +1,28 @@
+import { expect, test } from "bun:test";
+import type { OcxTool } from "../src/types";
+import { toolChoiceToolPredicate } from "../src/types";
+
+test("allowed_tools resolves an ambiguous bare name with one catalog pass", () => {
+  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 allowed = toolChoiceToolPredicate(
+    { allowedTools: ["shared_name"], mode: "auto" },
+    tools,
+  );
+
+  expect(tools.filter(allowed)).toEqual([]);
+  expect(catalogReads).toBeLessThanOrEqual(size * 2);
+});