Skip to content

Commit 80658ad

Browse files
committed
Wait briefly for connecting servers before answering an empty tool search
A search that misses while a connector handshake is still in flight now waits briefly for the handshake, then reports the miss as possibly incomplete instead of definitive.
1 parent 0bf3447 commit 80658ad

3 files changed

Lines changed: 128 additions & 2 deletions

File tree

‎src/agent/tool-search.test.ts‎

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -317,6 +317,38 @@ describe("createToolSearchTool", () => {
317317
expect(out).toContain("mcp__linear__create_issue");
318318
expect(out).not.toContain("No tools matched");
319319
});
320+
321+
test("a hung connection never hangs the search — bounded wait, then a retry signal", async () => {
322+
const tool = createToolSearchTool({
323+
search: () => [],
324+
lookup: () => undefined,
325+
promote: () => undefined,
326+
awaitPendingConnections: () =>
327+
new Promise<number>(() => {
328+
// Never settles: simulates a hung authorization handshake.
329+
}),
330+
});
331+
const out = await call(tool, { query: "linear" });
332+
expect(out).toContain("No tools matched");
333+
expect(out).toMatch(/starting up|still connecting/);
334+
expect(out).toMatch(/retry.*shortly/i);
335+
expect(out).not.toContain("different keywords");
336+
});
337+
338+
test("a genuine miss keeps the keyword advice and omits the retry caveat", async () => {
339+
const tool = createToolSearchTool({
340+
search: () => [],
341+
lookup: () => undefined,
342+
promote: () => undefined,
343+
awaitPendingConnections: async () => 0,
344+
});
345+
const out = await call(tool, { query: "nonsense" });
346+
expect(out).toContain("No tools matched");
347+
expect(out).toContain("different keywords");
348+
expect(out).not.toMatch(
349+
/still connecting|still starting up|retry shortly/i,
350+
);
351+
});
320352
});
321353

322354
describe("advertisedTools", () => {

‎src/agent/tool-search.ts‎

Lines changed: 51 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -283,8 +283,18 @@ export interface ToolSearchDeps {
283283
// invoke them this turn. The next inference also declares them on the wire
284284
// for strict providers.
285285
promote: (names: string[]) => void;
286+
// Resolves to the remaining in-flight MCP handshake count after waiting up
287+
// to `timeoutMs`. The toolset bounds its own wait; the handler re-races
288+
// below so even a stuck dependency can never hang the call. Omitted callers
289+
// (tests, ad-hoc indexes) have no pending handshakes to wait for.
290+
awaitPendingConnections?: (timeoutMs?: number) => Promise<number>;
286291
}
287292

293+
// Brief bound a tool_search miss waits for in-flight MCP handshakes before
294+
// answering. A hung authorization must never hang the call, so both the
295+
// toolset wait and the handler race below are capped by this.
296+
export const TOOL_SEARCH_PENDING_WAIT_MS = 1_000;
297+
288298
const ToolSearchArgs = type({ query: "string" });
289299

290300
// Render one discovered tool as name, description, and pretty-printed input
@@ -305,6 +315,28 @@ function indent(text: string, pad: string): string {
305315
.join("\n");
306316
}
307317

318+
// Race the dependency's pending-count wait against the same bound, so a
319+
// stuck dependency (hung OAuth that never settles) cannot hang the call.
320+
// Resolves undefined when this race itself times out.
321+
async function racePendingCount(
322+
awaitPending: (timeoutMs?: number) => Promise<number>,
323+
): Promise<number | undefined> {
324+
let timer: ReturnType<typeof setTimeout> | undefined;
325+
try {
326+
return await Promise.race([
327+
awaitPending(TOOL_SEARCH_PENDING_WAIT_MS),
328+
new Promise<undefined>((resolve) => {
329+
timer = setTimeout(
330+
() => resolve(undefined),
331+
TOOL_SEARCH_PENDING_WAIT_MS,
332+
);
333+
}),
334+
]);
335+
} finally {
336+
if (timer !== undefined) clearTimeout(timer);
337+
}
338+
}
339+
308340
export function createToolSearchTool(deps: ToolSearchDeps): AgentTool {
309341
return stringTool({
310342
definition: toolSearchDefinition,
@@ -316,7 +348,25 @@ export function createToolSearchTool(deps: ToolSearchDeps): AgentTool {
316348
const query = parsed.query.trim();
317349
if (query.length === 0)
318350
return "Error: tool_search requires a non-empty query.";
319-
const names = deps.search(query);
351+
let names = deps.search(query);
352+
if (names.length === 0 && deps.awaitPendingConnections !== undefined) {
353+
// Miss while connectors start up: wait briefly, then re-search so
354+
// late-mounting tools land. The race bounds even a stuck dependency
355+
// (hung OAuth) — undefined means the wait itself timed out.
356+
const stillPending = await racePendingCount(
357+
deps.awaitPendingConnections,
358+
);
359+
names = deps.search(query);
360+
if (names.length === 0 && (stillPending ?? 1) > 0) {
361+
const detail =
362+
stillPending === undefined
363+
? "a connector may still be starting up"
364+
: stillPending === 1
365+
? "1 connector is still connecting"
366+
: `${stillPending} connectors are still connecting`;
367+
return `No tools matched "${query}" yet — ${detail}. Retry this search shortly.`;
368+
}
369+
}
320370
if (names.length === 0) {
321371
return `No tools matched "${query}". Try different keywords describing the capability.`;
322372
}

‎src/agent/tools.ts‎

Lines changed: 45 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -103,7 +103,11 @@ import {
103103
} from "../tools/web-search.js";
104104
import { createUseSkillTool } from "./use-skill.js";
105105
import { createSkillSearchTool } from "./skill-search.js";
106-
import { createToolIndex, createToolSearchTool } from "./tool-search.js";
106+
import {
107+
createToolIndex,
108+
createToolSearchTool,
109+
TOOL_SEARCH_PENDING_WAIT_MS,
110+
} from "./tool-search.js";
107111
import { createSearchAgentsTool } from "./agent-search.js";
108112
import { createReadAgentTraceTool } from "../subagent/trace-tool.js";
109113
import {
@@ -319,6 +323,14 @@ export interface AgentToolset {
319323
// second add of an active name; failed rows retry through connectMCPServer
320324
// without a second persist. Still true while disable is in progress.
321325
hasMCPServer: (name: string) => boolean;
326+
// Live MCP handshake depth: non-zero while any server is still connecting.
327+
// tool_search consults this on a miss — waiting briefly and reporting a
328+
// retry signal instead of a definitive negative mid-handshake.
329+
pendingMcpConnectionCount: () => number;
330+
// Bounded wait for in-flight MCP handshakes; resolves to the remaining
331+
// count. Capped by `timeoutMs` so a hung authorization never hangs the
332+
// caller — the tool_search bound passes briefly by default.
333+
awaitPendingMcpConnections: (timeoutMs?: number) => Promise<number>;
322334
// Catalog unshadow can change local → global/none without rebuilding the
323335
// toolset; connectOne reads this on every late connect.
324336
setMcpServersSource: (source: "local" | "global" | "none") => void;
@@ -719,6 +731,12 @@ export async function createAgentToolset(
719731
lookup: (name) =>
720732
runnerHolder.current?.currentDefinitions().find((d) => d.name === name),
721733
promote: (names) => promoter.promote(names),
734+
// Misses wait briefly for in-flight MCP handshakes (bounded, so hung
735+
// OAuth cannot hang the call) and re-search before answering. Reads the
736+
// connection map live — declared below, populated by the time any
737+
// search runs.
738+
awaitPendingConnections: (timeoutMs = TOOL_SEARCH_PENDING_WAIT_MS) =>
739+
awaitPendingMcpConnections(timeoutMs),
722740
}),
723741
);
724742

@@ -744,6 +762,30 @@ export async function createAgentToolset(
744762
const connectedClients = new Map<string, MCPClient>();
745763
const inFlightConnections = new Map<string, Promise<void>>();
746764
const inFlightEpochs = new Map<string, number>();
765+
// Live handshake depth for tool_search: a miss waits briefly for these
766+
// before answering, so a mid-handshake search finds late-mounting tools
767+
// instead of reporting a definitive negative.
768+
const pendingMcpConnectionCount = (): number => inFlightConnections.size;
769+
// Bounded wait for in-flight handshakes; resolves to the remaining count.
770+
// Capped by `timeoutMs` so a hung authorization never hangs the caller.
771+
const awaitPendingMcpConnections = async (
772+
timeoutMs = TOOL_SEARCH_PENDING_WAIT_MS,
773+
): Promise<number> => {
774+
if (inFlightConnections.size === 0) return 0;
775+
const pending = [...inFlightConnections.values()];
776+
let timer: ReturnType<typeof setTimeout> | undefined;
777+
try {
778+
await Promise.race([
779+
Promise.allSettled(pending),
780+
new Promise<void>((resolve) => {
781+
timer = setTimeout(resolve, timeoutMs);
782+
}),
783+
]);
784+
} finally {
785+
if (timer !== undefined) clearTimeout(timer);
786+
}
787+
return inFlightConnections.size;
788+
};
747789
const disabledNames = new Set<string>();
748790
const serverAborts = new Map<string, AbortController>();
749791
const serverEpochs = new Map<string, number>();
@@ -1218,6 +1260,8 @@ export async function createAgentToolset(
12181260
disconnectMCPServer: publicDisconnectMCPServer,
12191261
hasMCPServer: (name) =>
12201262
connectedClients.has(name) || inFlightConnections.has(name),
1263+
pendingMcpConnectionCount,
1264+
awaitPendingMcpConnections,
12211265
setMcpServersSource: (source) => {
12221266
mcpServersSource = source;
12231267
},

0 commit comments

Comments
 (0)