Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
75 changes: 75 additions & 0 deletions src/agent/tool-search.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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<void>((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<number>(() => {
// 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", () => {
Expand Down
52 changes: 51 additions & 1 deletion src/agent/tool-search.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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<number>;
}

// 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
Expand All @@ -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<number>,
): Promise<number | undefined> {
let timer: ReturnType<typeof setTimeout> | undefined;
try {
return await Promise.race([
awaitPending(TOOL_SEARCH_PENDING_WAIT_MS),
new Promise<undefined>((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,
Expand All @@ -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.`;
}
Expand Down
37 changes: 36 additions & 1 deletion src/agent/tools.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down Expand Up @@ -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<number>;
// Catalog unshadow can change local → global/none without rebuilding the
// toolset; connectOne reads this on every late connect.
setMcpServersSource: (source: "local" | "global" | "none") => void;
Expand Down Expand Up @@ -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),
}),
);

Expand All @@ -744,6 +758,26 @@ export async function createAgentToolset(
const connectedClients = new Map<string, MCPClient>();
const inFlightConnections = new Map<string, Promise<void>>();
const inFlightEpochs = new Map<string, number>();
// 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<number> => {
if (inFlightConnections.size === 0) return 0;
const pending = [...inFlightConnections.values()];
let timer: ReturnType<typeof setTimeout> | undefined;
try {
await Promise.race([
Promise.allSettled(pending),
new Promise<void>((resolve) => {
timer = setTimeout(resolve, timeoutMs);
}),
]);
} finally {
if (timer !== undefined) clearTimeout(timer);
}
return inFlightConnections.size;
};
const disabledNames = new Set<string>();
const serverAborts = new Map<string, AbortController>();
const serverEpochs = new Map<string, number>();
Expand Down Expand Up @@ -1218,6 +1252,7 @@ export async function createAgentToolset(
disconnectMCPServer: publicDisconnectMCPServer,
hasMCPServer: (name) =>
connectedClients.has(name) || inFlightConnections.has(name),
awaitPendingMcpConnections,
setMcpServersSource: (source) => {
mcpServersSource = source;
},
Expand Down
Loading