diff --git a/src/agent/tool-search.test.ts b/src/agent/tool-search.test.ts index e3e2f600d..a65be8e69 100644 --- a/src/agent/tool-search.test.ts +++ b/src/agent/tool-search.test.ts @@ -287,6 +287,81 @@ describe("createToolSearchTool", () => { "No tools matched", ); }); + + test("mid-handshake search waits for a connecting server instead of reporting no match", async () => { + const live: ToolDefinition[] = []; + let resolveConnect!: () => void; + const connected = new Promise((resolve) => { + resolveConnect = resolve; + }); + const tool = createToolSearchTool({ + search: (query) => createToolIndex(() => live).search(query), + lookup: (name) => live.find((def) => def.name === name), + promote: () => undefined, + awaitPendingConnections: async (timeoutMs?: number) => { + await Promise.race([ + connected, + new Promise((resolve) => setTimeout(resolve, timeoutMs ?? 50)), + ]); + return live.length === 0 ? 1 : 0; + }, + }); + const pending = call(tool, { query: "linear tracker" }); + live.push({ + name: "mcp__linear__create_issue", + description: "Create an issue in the tracker", + inputSchema: { type: "object", properties: {}, required: [] }, + }); + resolveConnect(); + const out = await pending; + expect(out).toContain("mcp__linear__create_issue"); + expect(out).not.toContain("No tools matched"); + }); + + test("a hung connection never hangs the search — bounded wait, then a retry signal", async () => { + const tool = createToolSearchTool({ + search: () => [], + lookup: () => undefined, + promote: () => undefined, + awaitPendingConnections: () => + new Promise(() => { + // Never settles: simulates a hung authorization handshake. + }), + }); + const out = await call(tool, { query: "linear" }); + expect(out).toContain("No tools matched"); + expect(out).toMatch(/starting up|still connecting/); + expect(out).toMatch(/retry.*shortly/i); + expect(out).not.toContain("different keywords"); + }); + + test("two pending connectors report the plural connecting copy", async () => { + const tool = createToolSearchTool({ + search: () => [], + lookup: () => undefined, + promote: () => undefined, + awaitPendingConnections: async () => 2, + }); + const out = await call(tool, { query: "linear" }); + expect(out).toContain("2 connectors are still connecting"); + expect(out).toMatch(/retry.*shortly/i); + expect(out).not.toContain("different keywords"); + }); + + test("a genuine miss keeps the keyword advice and omits the retry caveat", async () => { + const tool = createToolSearchTool({ + search: () => [], + lookup: () => undefined, + promote: () => undefined, + awaitPendingConnections: async () => 0, + }); + const out = await call(tool, { query: "nonsense" }); + expect(out).toContain("No tools matched"); + expect(out).toContain("different keywords"); + expect(out).not.toMatch( + /still connecting|still starting up|retry shortly/i, + ); + }); }); describe("advertisedTools", () => { diff --git a/src/agent/tool-search.ts b/src/agent/tool-search.ts index b391adfa9..303017252 100644 --- a/src/agent/tool-search.ts +++ b/src/agent/tool-search.ts @@ -283,8 +283,18 @@ export interface ToolSearchDeps { // invoke them this turn. The next inference also declares them on the wire // for strict providers. promote: (names: string[]) => void; + // Resolves to the remaining in-flight MCP handshake count after waiting up + // to `timeoutMs`. The toolset bounds its own wait; the handler re-races + // below so even a stuck dependency can never hang the call. Omitted callers + // (tests, ad-hoc indexes) have no pending handshakes to wait for. + awaitPendingConnections?: (timeoutMs?: number) => Promise; } +// Brief bound a tool_search miss waits for in-flight MCP handshakes before +// answering. A hung authorization must never hang the call, so both the +// toolset wait and the handler race below are capped by this. +export const TOOL_SEARCH_PENDING_WAIT_MS = 1_000; + const ToolSearchArgs = type({ query: "string" }); // Render one discovered tool as name, description, and pretty-printed input @@ -305,6 +315,28 @@ function indent(text: string, pad: string): string { .join("\n"); } +// Race the dependency's pending-count wait against the same bound, so a +// stuck dependency (hung OAuth that never settles) cannot hang the call. +// Resolves undefined when this race itself times out. +async function racePendingCount( + awaitPending: (timeoutMs?: number) => Promise, +): Promise { + let timer: ReturnType | undefined; + try { + return await Promise.race([ + awaitPending(TOOL_SEARCH_PENDING_WAIT_MS), + new Promise((resolve) => { + timer = setTimeout( + () => resolve(undefined), + TOOL_SEARCH_PENDING_WAIT_MS, + ); + }), + ]); + } finally { + if (timer !== undefined) clearTimeout(timer); + } +} + export function createToolSearchTool(deps: ToolSearchDeps): AgentTool { return stringTool({ definition: toolSearchDefinition, @@ -316,7 +348,25 @@ export function createToolSearchTool(deps: ToolSearchDeps): AgentTool { const query = parsed.query.trim(); if (query.length === 0) return "Error: tool_search requires a non-empty query."; - const names = deps.search(query); + let names = deps.search(query); + if (names.length === 0 && deps.awaitPendingConnections !== undefined) { + // Miss while connectors start up: wait briefly, then re-search so + // late-mounting tools land. The race bounds even a stuck dependency + // (hung OAuth) — undefined means the wait itself timed out. + const stillPending = await racePendingCount( + deps.awaitPendingConnections, + ); + names = deps.search(query); + if (names.length === 0 && (stillPending ?? 1) > 0) { + const detail = + stillPending === undefined + ? "a connector may still be starting up" + : stillPending === 1 + ? "1 connector is still connecting" + : `${stillPending} connectors are still connecting`; + return `No tools matched "${query}" yet — ${detail}. Retry this search shortly.`; + } + } if (names.length === 0) { return `No tools matched "${query}". Try different keywords describing the capability.`; } diff --git a/src/agent/tools.ts b/src/agent/tools.ts index 621de824a..65f6deb82 100644 --- a/src/agent/tools.ts +++ b/src/agent/tools.ts @@ -103,7 +103,11 @@ import { } from "../tools/web-search.js"; import { createUseSkillTool } from "./use-skill.js"; import { createSkillSearchTool } from "./skill-search.js"; -import { createToolIndex, createToolSearchTool } from "./tool-search.js"; +import { + createToolIndex, + createToolSearchTool, + TOOL_SEARCH_PENDING_WAIT_MS, +} from "./tool-search.js"; import { createSearchAgentsTool } from "./agent-search.js"; import { createReadAgentTraceTool } from "../subagent/trace-tool.js"; import { @@ -319,6 +323,10 @@ export interface AgentToolset { // second add of an active name; failed rows retry through connectMCPServer // without a second persist. Still true while disable is in progress. hasMCPServer: (name: string) => boolean; + // Bounded wait for in-flight MCP handshakes; resolves to the remaining + // count. Capped by `timeoutMs` so a hung authorization never hangs the + // caller — the tool_search bound passes briefly by default. + awaitPendingMcpConnections: (timeoutMs?: number) => Promise; // Catalog unshadow can change local → global/none without rebuilding the // toolset; connectOne reads this on every late connect. setMcpServersSource: (source: "local" | "global" | "none") => void; @@ -719,6 +727,12 @@ export async function createAgentToolset( lookup: (name) => runnerHolder.current?.currentDefinitions().find((d) => d.name === name), promote: (names) => promoter.promote(names), + // Misses wait briefly for in-flight MCP handshakes (bounded, so hung + // OAuth cannot hang the call) and re-search before answering. Reads the + // connection map live — declared below, populated by the time any + // search runs. + awaitPendingConnections: (timeoutMs = TOOL_SEARCH_PENDING_WAIT_MS) => + awaitPendingMcpConnections(timeoutMs), }), ); @@ -744,6 +758,26 @@ export async function createAgentToolset( const connectedClients = new Map(); const inFlightConnections = new Map>(); const inFlightEpochs = new Map(); + // Bounded wait for in-flight handshakes; resolves to the remaining count. + // Capped by `timeoutMs` so a hung authorization never hangs the caller. + const awaitPendingMcpConnections = async ( + timeoutMs = TOOL_SEARCH_PENDING_WAIT_MS, + ): Promise => { + if (inFlightConnections.size === 0) return 0; + const pending = [...inFlightConnections.values()]; + let timer: ReturnType | undefined; + try { + await Promise.race([ + Promise.allSettled(pending), + new Promise((resolve) => { + timer = setTimeout(resolve, timeoutMs); + }), + ]); + } finally { + if (timer !== undefined) clearTimeout(timer); + } + return inFlightConnections.size; + }; const disabledNames = new Set(); const serverAborts = new Map(); const serverEpochs = new Map(); @@ -1218,6 +1252,7 @@ export async function createAgentToolset( disconnectMCPServer: publicDisconnectMCPServer, hasMCPServer: (name) => connectedClients.has(name) || inFlightConnections.has(name), + awaitPendingMcpConnections, setMcpServersSource: (source) => { mcpServersSource = source; },