diff --git a/packages/aitm/docs/mcp.md b/packages/aitm/docs/mcp.md index 6d9ff38..8eb5d7d 100644 --- a/packages/aitm/docs/mcp.md +++ b/packages/aitm/docs/mcp.md @@ -63,6 +63,25 @@ Per-role allowlists scope which servers reach which subagent. Useful for sandbox Unlisted servers default to all roles. Optional — most users skip this. +## Deferred tool loading + +Mounting a large MCP registry (a GitHub, browser, or database server can expose dozens of tools) would serialize every tool's full JSON schema into every request, on every step — burning context window and per-step cost on schemas the subagent never calls. `mcpDeferToolsOver` bounds that: once a role's MCP tools (beyond the fixed local slots like `readFile`/`bash`) exceed the threshold, the surplus is **deferred** — presented name-only in the system prompt, with schemas absent from requests until fetched. + +```jsonc +{ + "mcpServers": { /* ... */ }, + "mcpDeferToolsOver": 20 // default 20; 0 = always defer surplus tools +} +``` + +At or below the threshold, surplus tools mount directly (full schema) — identical to plain mounting. Above it: + +- The system prompt carries a name-only index (`: ` per tool) plus a fetch-before-call contract. +- A `tool_search` tool is mounted. `select:mcp__server__tool,…` fetches those exact tools by name; any other query is keyword-ranked over names + descriptions (capped by `max_results`, default 5). Each match's full schema is returned and the tool becomes callable on subsequent steps. +- Calling a deferred tool before fetching its schema returns a typed "fetch it via `tool_search` first" result — never a provider-level error. + +Activation is scoped to one subagent invocation; a fresh invocation starts fully deferred again. Applies to the **Worker** and **Reviewer** surfaces (the Planner keeps its read-only trio; the CI-fix session keeps its fixed record). Config-only, resolved project > global. See `src/mcp/tool-search.ts` and `src/loop/run-loop-adapter.ts`. + ## Lifecycle `McpClientManager` (`src/mcp/mcp-client.ts`) owns the lifecycle: diff --git a/packages/aitm/src/config/config-loader.test.ts b/packages/aitm/src/config/config-loader.test.ts index bebe025..7d40b2e 100644 --- a/packages/aitm/src/config/config-loader.test.ts +++ b/packages/aitm/src/config/config-loader.test.ts @@ -58,6 +58,7 @@ test('resolve: uses built-in defaults when only env key is set', async () => { assert.equal(resolved.logLevel, 'info'); assert.equal(resolved.concurrency, 1); assert.equal(resolved.allowForcePush, true); + assert.equal(resolved.mcpDeferToolsOver, 20); assert.deepEqual(resolved.models, DEFAULT_MODELS); } finally { await home.cleanup(); @@ -65,6 +66,23 @@ test('resolve: uses built-in defaults when only env key is set', async () => { } }); +test('resolve: mcpDeferToolsOver is project over global, and a configured 0 is honored (issue #119)', async () => { + const home = await tempDir('aitm-home-'); + const cwd = await tempDir('aitm-cwd-'); + try { + await writeGlobalConfig(home.path, { mcpDeferToolsOver: 50 }); + const loader = new ConfigLoader(cwd.path, home.path, { OPENROUTER_API_KEY: 'sk-env' }); + assert.equal((await loader.resolve({})).mcpDeferToolsOver, 50, 'global applies'); + + // 0 means "always defer" — pick() must treat it as set, not fall through to the default. + await writeProjectConfig(cwd.path, { mcpDeferToolsOver: 0 }); + assert.equal((await loader.resolve({})).mcpDeferToolsOver, 0, 'project 0 wins over global'); + } finally { + await home.cleanup(); + await cwd.cleanup(); + } +}); + test('resolve: formatCommand is read from project, then global (issue #48)', async () => { const home = await tempDir('aitm-home-'); const cwd = await tempDir('aitm-cwd-'); diff --git a/packages/aitm/src/config/config-loader.ts b/packages/aitm/src/config/config-loader.ts index bb00e18..45e5b7a 100644 --- a/packages/aitm/src/config/config-loader.ts +++ b/packages/aitm/src/config/config-loader.ts @@ -9,6 +9,7 @@ import { ZodError, z } from 'zod'; import { DEFAULT_MODELS } from '../credentials/defaults.ts'; import { atomicWrite } from '../fs/atomic-write.ts'; import { DEFAULT_MAX_CI_FIX_ATTEMPTS } from '../loop/constants.ts'; +import { DEFAULT_MCP_DEFER_TOOLS_OVER } from '../mcp/mcp-client.ts'; import { type McpServers, McpServersSchema } from '../mcp/schema.ts'; import { DEFAULT_LLM_STEP_TIMEOUT_MS } from '../subagents/factory.ts'; import { @@ -55,6 +56,7 @@ const KNOWN_KEYS = new Set([ 'reasoningEffort', 'mcpServers', 'mcpRoleAllowlist', + 'mcpDeferToolsOver', 'hooks', ]); @@ -85,6 +87,7 @@ const DEFAULTS = { logLevel: 'info' as const, concurrency: 1, allowForcePush: true, + mcpDeferToolsOver: DEFAULT_MCP_DEFER_TOOLS_OVER, }; type WarnFn = (msg: string) => void; @@ -249,6 +252,14 @@ export class ConfigLoader { ...((project?.mcpRoleAllowlist ?? global?.mcpRoleAllowlist) ? { mcpRoleAllowlist: project?.mcpRoleAllowlist ?? global?.mcpRoleAllowlist } : {}), + // Defer-tools threshold — aitm config only, project over global (issue #119). pick() treats a + // configured 0 as set (!== undefined), so "always defer" survives the default. + mcpDeferToolsOver: pick( + undefined, + project?.mcpDeferToolsOver, + global?.mcpDeferToolsOver, + DEFAULTS.mcpDeferToolsOver, + ), // Tool-registry hooks (issue #121). Hooks run shell commands with the operator's privileges, so // they are honored ONLY from the user-owned global config (~/.aitm.json) — NEVER from the // per-repo project config, which an untrusted repo could ship (CR: arbitrary code execution). A diff --git a/packages/aitm/src/config/schema.ts b/packages/aitm/src/config/schema.ts index 8e6e5c1..3d09304 100644 --- a/packages/aitm/src/config/schema.ts +++ b/packages/aitm/src/config/schema.ts @@ -180,6 +180,10 @@ export const ConfigFileSchema = z // patterns. aitm-config-only (project > global); the Claude Code interop sources contribute // mcpServers alone. See src/mcp/schema.ts and src/mcp/mcp-client.ts. mcpRoleAllowlist: McpRoleAllowlistSchema.optional(), + // Defer a role's MCP tools to name-only stubs + `tool_search` once their count exceeds this + // (issue #119), keeping their JSON schemas out of every request. Default 20; 0 = always defer. + // aitm-config-only (project > global). See src/mcp/mcp-client.ts and src/mcp/tool-search.ts. + mcpDeferToolsOver: z.number().int().min(0).optional(), // PreToolUse/PostToolUse shell hooks on the tool registry (issue #121). Hooks run shell commands // with operator privileges, so they are honored ONLY from the user-owned global config // (~/.aitm.json); the same key in a repo-shippable project config is parsed but ignored + warned. @@ -256,6 +260,9 @@ export type ResolvedConfig = { // Per-role MCP allowlist (issue #115), aitm-config-only (project > global). Undefined → every role // gets every connected server. Passed into McpClientManager. mcpRoleAllowlist?: import('../mcp/schema.ts').McpRoleAllowlist | undefined; + // Threshold above which a role's MCP tools are deferred (issue #119). Default + // DEFAULT_MCP_DEFER_TOOLS_OVER; 0 = always defer. Passed into McpClientManager. + mcpDeferToolsOver: number; // Tool-registry hooks (issue #121). Global config only (~/.aitm.json) — project hooks are ignored // as a code-execution trust boundary. Undefined → no hooks; behavior unchanged. Applied over the // resolved tool records in run-loop-adapter via withHooks. diff --git a/packages/aitm/src/loop/run-loop-adapter.test.ts b/packages/aitm/src/loop/run-loop-adapter.test.ts index 535182c..7dee381 100644 --- a/packages/aitm/src/loop/run-loop-adapter.test.ts +++ b/packages/aitm/src/loop/run-loop-adapter.test.ts @@ -12,8 +12,10 @@ import { test } from 'node:test'; import { AUTONOMY_CONTRACT_TEXT, COMMUNICATION_CONTRACT_TEXT, + SUBMIT_TOOL_NAME, SYSTEM_REMINDER_CONTRACT, } from '@developerz.ai/ai-claude-compat'; +import type { ToolSet } from 'ai'; import { tool } from 'ai'; import { MockLanguageModelV3 } from 'ai/test'; import { z } from 'zod'; @@ -21,6 +23,8 @@ import type { RunLoopInput } from '../cli/commands.ts'; import { Credentials } from '../credentials/credentials.ts'; import { GitHubClient } from '../github/github-client.ts'; import type { PullRequest, ReviewThread } from '../github/schema.ts'; +import { McpClientManager } from '../mcp/mcp-client.ts'; +import { TOOL_SEARCH_TOOL_NAME } from '../mcp/tool-search.ts'; import type { Plan } from '../plan/schema.ts'; import type { PrGroup, RunState } from '../state/schema.ts'; import { StateStore } from '../state/state-store.ts'; @@ -28,6 +32,7 @@ import type { ReviewerResult } from '../subagents/reviewer.ts'; import type { WorkerDelivery, WorkerResult } from '../subagents/worker.ts'; import { type AdapterStatePort, + activeToolNames, applyHooks, branchFor, createRollingContextAccumulator, @@ -38,6 +43,7 @@ import { localEditTools, localReadTools, mcpTool, + mountDeferredTools, type PlanGroupsOutcome, persistRollingContext, planToPrGroups, @@ -713,7 +719,10 @@ test('defaultMakeOrchestrator constructs the Compactor and wires it into the sta }, modelIdForCapability: () => 'openai/gpt-5', }; - const mcp = { toolsForRole: () => ({}) }; + const mcp = { + toolsForRole: () => ({}), + toolSurfaceForRole: () => ({ direct: {}, deferred: {} }), + }; const input = { cwd: '/tmp/adapter-compaction', resolved: { openrouterApiKey: 'sk-or-test', maxSessions: null }, @@ -971,3 +980,107 @@ test('resolvePlannerTools mounts explore only when the caller wires it', () => { // The Planner never gets a memory tool — it reads memory files directly (issue #118). assert.equal('memory' in withExplore, false, 'planner has no memory tool'); }); + +// ---- deferred MCP tool loading (issue #119) ---- + +function mcpFake(desc: string): ToolSet[string] { + return { description: desc, inputSchema: { type: 'object' } } as ToolSet[string]; +} + +test('mountDeferredTools: below threshold (nothing deferred) mounts surplus direct, no tool_search (issue #119)', () => { + const mount = mountDeferredTools({ + direct: { mcp__gh__create_issue: mcpFake('Create an issue.') }, + deferred: {}, + }); + assert.deepEqual(Object.keys(mount.extraTools), ['mcp__gh__create_issue']); + assert.equal(mount.indexBlock, ''); + assert.equal(mount.activated, null); + assert.equal(mount.deferredNames.size, 0); + assert.equal( + TOOL_SEARCH_TOOL_NAME in mount.extraTools, + false, + 'no tool_search when nothing deferred', + ); +}); + +test('mountDeferredTools: above threshold defers surplus behind tool_search + a name-only index (issue #119)', () => { + const mount = mountDeferredTools({ + direct: {}, + deferred: { + mcp__gh__create_issue: mcpFake('Create an issue.'), + mcp__db__query: mcpFake('Query the DB.'), + }, + }); + assert.ok(TOOL_SEARCH_TOOL_NAME in mount.extraTools, 'tool_search mounted'); + assert.ok( + 'mcp__gh__create_issue' in mount.extraTools, + 'deferred tool guard-wrapped into the record', + ); + assert.ok('mcp__db__query' in mount.extraTools); + assert.match(mount.indexBlock, /mcp__gh__create_issue: Create an issue\./); + assert.notEqual(mount.activated, null); + assert.deepEqual([...mount.deferredNames].sort(), ['mcp__db__query', 'mcp__gh__create_issue']); +}); + +test('mountDeferredTools: fixed-slot-named MCP tools are not surplus — excluded from the mount (issue #119)', () => { + const mount = mountDeferredTools({ + direct: {}, + deferred: { mcp__fs__readFile: mcpFake('read'), mcp__gh__x: mcpFake('x') }, + }); + // readFile is a fixed slot (partial-filled elsewhere) → not deferred here; only true surplus is. + assert.deepEqual([...mount.deferredNames], ['mcp__gh__x']); + assert.equal('mcp__fs__readFile' in mount.extraTools, false); +}); + +test('activeToolNames: hides un-activated deferred tools, always keeps submit + non-deferred (issue #119)', () => { + const tools: ToolSet = { + readFile: mcpFake('r'), + mcp__gh__x: mcpFake('x'), + [TOOL_SEARCH_TOOL_NAME]: mcpFake('search'), + }; + const deferredNames = new Set(['mcp__gh__x']); + const before = activeToolNames(tools, deferredNames, new Set()); + assert.equal(before.includes('mcp__gh__x'), false, 'deferred tool inactive until fetched'); + assert.ok(before.includes('readFile') && before.includes(TOOL_SEARCH_TOOL_NAME)); + assert.ok(before.includes(SUBMIT_TOOL_NAME), 'submit always active'); + const after = activeToolNames(tools, deferredNames, new Set(['mcp__gh__x'])); + assert.ok(after.includes('mcp__gh__x'), 'an activated deferred tool becomes active'); +}); + +test('deferred loading end-to-end: an over-threshold MCP server surfaces name-only + tool_search on the Worker (issue #119)', async () => { + const surplus: ToolSet = { + create_issue: mcpFake('Create a GitHub issue.'), + list_prs: mcpFake('List PRs.'), + }; + const mcp = new McpClientManager({ + servers: { gh: { command: 'gh-mcp' } }, + deferToolsOver: 1, // 2 surplus tools > 1 → deferred + createClient: (async () => + ({ tools: async () => surplus, close: async () => {} }) as never) as never, + }); + await mcp.connectAll(); + const mount = mountDeferredTools(mcp.toolSurfaceForRole('worker')); + // resolveWorkerTools fills the fixed slots (local, since the server supplies none); the surplus is + // added by the mount — proving tools beyond the fixed slots now reach the Worker (dropped pre-#119). + const workerTools: ToolSet = { + ...resolveWorkerTools(mcp.toolsForRole('worker'), '/tmp/wt'), + ...mount.extraTools, + }; + assert.ok(TOOL_SEARCH_TOOL_NAME in workerTools, 'tool_search reaches the Worker'); + assert.ok( + 'mcp__gh__create_issue' in workerTools, + 'surplus tools reach the Worker (were dropped before #119)', + ); + assert.ok('readFile' in workerTools, 'fixed slots still present'); + const active = activeToolNames(workerTools, mount.deferredNames, mount.activated ?? new Set()); + assert.equal( + active.includes('mcp__gh__create_issue'), + false, + 'deferred schema absent from active tools until fetched', + ); + assert.ok( + active.includes('readFile') && active.includes(SUBMIT_TOOL_NAME), + 'fixed slots + submit stay active', + ); + await mcp.close(); +}); diff --git a/packages/aitm/src/loop/run-loop-adapter.ts b/packages/aitm/src/loop/run-loop-adapter.ts index 97c7f34..0f7d836 100644 --- a/packages/aitm/src/loop/run-loop-adapter.ts +++ b/packages/aitm/src/loop/run-loop-adapter.ts @@ -29,6 +29,7 @@ import { multiEditTool, type ReminderProvider, readFileTool, + SUBMIT_TOOL_NAME, type SubagentHandle, SYSTEM_REMINDER_CONTRACT, type ToolHooks, @@ -36,14 +37,15 @@ import { withReminders, writeFileTool, } from '@developerz.ai/ai-claude-compat'; -import { type ModelMessage, type Tool, type ToolSet, tool } from 'ai'; +import { type ModelMessage, type Tool, type ToolLoopAgentSettings, type ToolSet, tool } from 'ai'; import { z } from 'zod'; // Type-only import — no runtime cycle with commands.ts, which imports this module's value. import type { RunLoopInput } from '../cli/commands.ts'; import { buildCompactionStep } from '../compaction/compaction-step.ts'; import { Compactor } from '../compaction/compactor.ts'; import type { GitHubClient } from '../github/github-client.ts'; -import { McpClientManager } from '../mcp/mcp-client.ts'; +import { McpClientManager, type ToolSurface } from '../mcp/mcp-client.ts'; +import { guardDeferred, TOOL_SEARCH_TOOL_NAME, toolSearch } from '../mcp/tool-search.ts'; import { roleUsageSink } from '../observability/usage-tracker.ts'; import { OpenRouterClient } from '../openrouter/client.ts'; import { ModelLimitsRegistry } from '../openrouter/model-limits.ts'; @@ -715,17 +717,28 @@ export function defaultMakeOrchestrator(ctx: OrchestratorBridgeCtx): WorkLoopOrc // conversations share one agent, so this records them into one transcript rather than one-per-thread. const runReviewerThreads = async ({ pr, threads, worktree }: ReviewerInvocation) => { const github = githubThreadTool(input.github); + // Surplus MCP tools beyond the fixed slots reach the Reviewer too, deferred above the threshold + // (issue #119). The Reviewer's local `github` glue is a fixed slot, never deferred. + const reviewerSurface = mcp.toolSurfaceForRole('reviewer'); + const reviewerMount = mountDeferredTools(reviewerSurface); const tools = applyHooks( - resolveReviewerTools( - mcp.toolsForRole('reviewer'), - worktree.path, - github, - input.resolved.bashRules, - fetchHtmlAvailable, - ), + { + ...resolveReviewerTools( + mcp.toolsForRole('reviewer'), + worktree.path, + github, + input.resolved.bashRules, + fetchHtmlAvailable, + ), + ...reviewerMount.extraTools, + } as ReviewerTools, input, worktree.path, ); + const reviewerCompaction = buildCompactionStep({ + compactor, + modelId: input.credentials.modelIdFor('reviewer'), + }); const recorder = await beginTranscript(state.transcripts?.(), { group: worktree.groupId, stage: 'addressing-reviews', @@ -733,17 +746,25 @@ export function defaultMakeOrchestrator(ctx: OrchestratorBridgeCtx): WorkLoopOrc const agent = createReviewerAgent({ model: input.credentials.modelFor('reviewer'), tools, - systemPrompt: reminderAgentSystemPrompt({ - style, - roleGuidance: REVIEWER_SYSTEM_PREFIX, - cwd: worktree.path, - maxSteps: REVIEWER_MAX_STEPS, - modelId: input.credentials.modelIdFor('reviewer'), - }), - prepareStep: buildCompactionStep({ - compactor, - modelId: input.credentials.modelIdFor('reviewer'), - }), + systemPrompt: appendIndexBlock( + reminderAgentSystemPrompt({ + style, + roleGuidance: REVIEWER_SYSTEM_PREFIX, + cwd: worktree.path, + maxSteps: REVIEWER_MAX_STEPS, + modelId: input.credentials.modelIdFor('reviewer'), + }), + reviewerMount.indexBlock, + ), + prepareStep: + reviewerMount.activated === null + ? reviewerCompaction + : withActiveTools( + reviewerCompaction, + tools, + reviewerMount.deferredNames, + reviewerMount.activated, + ), timeout: stepTimeout, ...(reviewerUsage ? { onUsage: reviewerUsage } : {}), ...(recorder @@ -766,18 +787,29 @@ export function defaultMakeOrchestrator(ctx: OrchestratorBridgeCtx): WorkLoopOrc // Prefer MCP-supplied tools; partial-fill any the server omits from the local set so a // bare `aitm start` (no mcpServers configured) can still edit, commit and open a PR. // memory (#118) is mounted on the manifest Worker so it can record durable repo facts. + // Surplus MCP tools beyond the fixed slots are mounted too (issue #119): directly below the + // defer threshold, else name-only + `tool_search`. + const workerSurface = mcp.toolSurfaceForRole('worker'); + const workerMount = mountDeferredTools(workerSurface); const tools = applyHooks( - resolveWorkerTools( - mcp.toolsForRole('worker'), - worktree.path, - input.resolved.bashRules, - fetchHtmlAvailable, - buildExploreFor(input, worktree.path), - memoryToolFor(state), - ), + { + ...resolveWorkerTools( + mcp.toolsForRole('worker'), + worktree.path, + input.resolved.bashRules, + fetchHtmlAvailable, + buildExploreFor(input, worktree.path), + memoryToolFor(state), + ), + ...workerMount.extraTools, + } as WithExplore & WithMemory, input, worktree.path, ); + const workerCompaction = buildCompactionStep({ + compactor, + modelId: input.credentials.modelIdFor('worker'), + }); const memoryIndex = await memoryIndexFor(state); // Transcript (issue #108): resume from an interrupted 'working' transcript for this group if one // exists — looked up BEFORE we begin the new one so it can't self-resume — then record this run. @@ -789,18 +821,26 @@ export function defaultMakeOrchestrator(ctx: OrchestratorBridgeCtx): WorkLoopOrc const agent = createWorkerAgent({ model: input.credentials.modelFor('worker'), tools, - systemPrompt: reminderAgentSystemPrompt({ - style, - roleGuidance: WORKER_SYSTEM_PREFIX, - cwd: worktree.path, - maxSteps: WORKER_MAX_STEPS, - modelId: input.credentials.modelIdFor('worker'), - memoryIndex, - }), - prepareStep: buildCompactionStep({ - compactor, - modelId: input.credentials.modelIdFor('worker'), - }), + systemPrompt: appendIndexBlock( + reminderAgentSystemPrompt({ + style, + roleGuidance: WORKER_SYSTEM_PREFIX, + cwd: worktree.path, + maxSteps: WORKER_MAX_STEPS, + modelId: input.credentials.modelIdFor('worker'), + memoryIndex, + }), + workerMount.indexBlock, + ), + prepareStep: + workerMount.activated === null + ? workerCompaction + : withActiveTools( + workerCompaction, + tools, + workerMount.deferredNames, + workerMount.activated, + ), timeout: stepTimeout, ...(providerOptions !== undefined ? { providerOptions } : {}), ...(workerUsage ? { onUsage: workerUsage } : {}), @@ -957,6 +997,110 @@ function mcpBaseName(key: string): string | undefined { return parts.slice(2).join('__'); } +// Fixed-slot canonical names that resolveWorkerTools/resolvePlannerTools partial-fill. An MCP tool +// whose base name is one of these is consumed into that slot (never dropped), so it is not part of +// the "surplus" deferred loading operates on (issue #119). +const FIXED_SLOT_NAMES = new Set([ + 'readFile', + 'writeFile', + 'editFile', + 'multiEdit', + 'grep', + 'glob', + 'bash', + 'multiBash', + 'webFetch', + 'datetime', + 'fetchHtml', +]); + +// The role's surplus MCP tools — everything beyond the fixed slots. Deferred loading operates on +// these; today they are dropped by the fixed-record resolvers (issue #119). +function surplusMcpTools(set: ToolSet): ToolSet { + const out: ToolSet = {}; + for (const [key, entry] of Object.entries(set)) { + if (!FIXED_SLOT_NAMES.has(mcpBaseName(key) ?? key)) out[key] = entry; + } + return out; +} + +type PrepareStep = NonNullable< + ToolLoopAgentSettings['prepareStep'] +>; + +// Deferred-loading mount for a role (issue #119). Splits the surface's surplus into directly-mounted +// (full schema) vs. deferred (name-only + `tool_search` + guard). Below the threshold nothing is +// deferred: the direct surplus still mounts, but there is no tool_search, no index block, and +// `activated` is null so the caller keeps its plain compaction prepareStep — the surface is +// byte-identical to today for that role. +type DeferredMount = { + extraTools: ToolSet; + indexBlock: string; + deferredNames: ReadonlySet; + activated: ReadonlySet | null; +}; + +export function mountDeferredTools(surface: ToolSurface): DeferredMount { + const directSurplus = surplusMcpTools(surface.direct); + const deferredSurplus = surplusMcpTools(surface.deferred); + const deferredKeys = Object.keys(deferredSurplus); + if (deferredKeys.length === 0) { + return { + extraTools: directSurplus, + indexBlock: '', + deferredNames: new Set(), + activated: null, + }; + } + const search = toolSearch(deferredSurplus); + const guarded: ToolSet = {}; + for (const [name, entry] of Object.entries(deferredSurplus)) { + guarded[name] = guardDeferred(name, entry, search.activated); + } + return { + extraTools: { ...directSurplus, ...guarded, [TOOL_SEARCH_TOOL_NAME]: search.tool }, + indexBlock: search.indexBlock, + deferredNames: new Set(deferredKeys), + activated: search.activated, + }; +} + +// The active-tool subset for a step: every mounted tool (plus `submit`) except deferred tools not +// yet activated. Grows as activations accumulate within one invocation. +export function activeToolNames( + tools: ToolSet, + deferredNames: ReadonlySet, + activated: ReadonlySet, +): string[] { + return [...Object.keys(tools), SUBMIT_TOOL_NAME].filter( + (name) => !deferredNames.has(name) || activated.has(name), + ); +} + +// Compose activeTools onto the #102 compaction prepareStep: one function returns `activeTools` every +// step and the compaction `messages` override when it triggers (issue #119 §"Activation plumbing"). +// activeTools is recomputed per step so newly activated deferred tools become callable. +function withActiveTools( + base: PrepareStep, + tools: ToolSet, + deferredNames: ReadonlySet, + activated: ReadonlySet, +): PrepareStep { + return async (options) => { + const result = await base(options); + return { + ...(result ?? {}), + activeTools: activeToolNames(tools, deferredNames, activated) as Array, + }; + }; +} + +// Append the deferred-tool index block to a role's system prompt, or return it unchanged when the +// role has nothing deferred (issue #119). +function appendIndexBlock(prompt: string, indexBlock: string): string { + return indexBlock ? `${prompt}\n\n${indexBlock}` : prompt; +} + export function resolveWorkerTools( set: ToolSet, cwd: string, diff --git a/packages/aitm/src/mcp/mcp-client.test.ts b/packages/aitm/src/mcp/mcp-client.test.ts index 3be386d..98c8461 100644 --- a/packages/aitm/src/mcp/mcp-client.test.ts +++ b/packages/aitm/src/mcp/mcp-client.test.ts @@ -409,3 +409,79 @@ test('http transport propagates headers through client config', async () => { assert.equal(transport.url, 'https://example.com/mcp'); assert.equal(transport.headers?.Authorization, 'Bearer xyz'); }); + +// ---- toolSurfaceForRole: deferred-loading split (issue #119) ---- + +test('toolSurfaceForRole: at or below the threshold, every tool mounts direct — no deferral (issue #119)', async () => { + const { createClient } = recordingFactory({ srv: { a: fakeTool(), b: fakeTool() } }); + const m = new McpClientManager({ + servers: { srv: { command: 'x' } }, + deferToolsOver: 2, + createClient, + }); + await m.connectAll(); + const surface = m.toolSurfaceForRole('worker'); + assert.deepEqual(Object.keys(surface.direct).sort(), ['mcp__srv__a', 'mcp__srv__b']); + assert.deepEqual(surface.deferred, {}); + await m.close(); +}); + +test('toolSurfaceForRole: above the threshold, the whole surplus defers (issue #119)', async () => { + const { createClient } = recordingFactory({ + srv: { a: fakeTool(), b: fakeTool(), c: fakeTool() }, + }); + const m = new McpClientManager({ + servers: { srv: { command: 'x' } }, + deferToolsOver: 2, + createClient, + }); + await m.connectAll(); + const surface = m.toolSurfaceForRole('worker'); + assert.deepEqual(surface.direct, {}); + assert.deepEqual(Object.keys(surface.deferred).sort(), [ + 'mcp__srv__a', + 'mcp__srv__b', + 'mcp__srv__c', + ]); + await m.close(); +}); + +test('toolSurfaceForRole: threshold 0 always defers any surplus (issue #119)', async () => { + const { createClient } = recordingFactory({ srv: { only: fakeTool() } }); + const m = new McpClientManager({ + servers: { srv: { command: 'x' } }, + deferToolsOver: 0, + createClient, + }); + await m.connectAll(); + const surface = m.toolSurfaceForRole('worker'); + assert.deepEqual(surface.direct, {}); + assert.deepEqual(Object.keys(surface.deferred), ['mcp__srv__only']); + await m.close(); +}); + +test('toolSurfaceForRole: no threshold set uses the default (20) — a small registry stays direct (issue #119)', async () => { + const { createClient } = recordingFactory({ srv: { a: fakeTool() } }); + const m = new McpClientManager({ servers: { srv: { command: 'x' } }, createClient }); + await m.connectAll(); + const surface = m.toolSurfaceForRole('worker'); + assert.deepEqual(Object.keys(surface.direct), ['mcp__srv__a']); + assert.deepEqual(surface.deferred, {}); + await m.close(); +}); + +test('toolSurfaceForRole: the split honors the per-role allowlist (issue #119)', async () => { + const { createClient } = recordingFactory({ a: { x: fakeTool() }, b: { y: fakeTool() } }); + const m = new McpClientManager({ + servers: { a: { command: 'x' }, b: { command: 'y' } }, + roleAllowlist: { worker: ['a'] }, + deferToolsOver: 0, + createClient, + }); + await m.connectAll(); + const surface = m.toolSurfaceForRole('worker'); + // Only server `a` is mounted for the worker, so only its tool appears — then defers at threshold 0. + assert.deepEqual(Object.keys(surface.deferred), ['mcp__a__x']); + assert.deepEqual(surface.direct, {}); + await m.close(); +}); diff --git a/packages/aitm/src/mcp/mcp-client.ts b/packages/aitm/src/mcp/mcp-client.ts index ee55eb1..0ffcb9f 100644 --- a/packages/aitm/src/mcp/mcp-client.ts +++ b/packages/aitm/src/mcp/mcp-client.ts @@ -20,12 +20,24 @@ export type TransportKind = 'stdio' | 'http' | 'sse'; export type CreateMcpClient = (config: MCPClientConfig) => Promise; +// Above this many role-visible MCP tools, the surplus is deferred (name-only stubs + a `tool_search` +// tool) instead of mounted directly, so their JSON schemas stay out of every request (issue #119). +// `0` = always defer. Overridable via the `mcpDeferToolsOver` config key. +export const DEFAULT_MCP_DEFER_TOOLS_OVER = 20; + +// The role's MCP tools split by how they enter the agent: `direct` are mounted with full schemas; +// `deferred` are surfaced name-only and fetched on demand through `tool_search` (issue #119). +export type ToolSurface = { direct: ToolSet; deferred: ToolSet }; + export type McpClientInit = { servers: McpServers; // Optional per-role allowlist (issue #115). Per role, either whole servers by name // (`{ worker: ['filesystem'] }`) or per-server tool patterns with `*` wildcards // (`{ planner: { filesystem: ['read_*', 'list_*'] } }`). Unlisted roles get every connected server. roleAllowlist?: McpRoleAllowlist; + // Defer a role's MCP tools once their count exceeds this (issue #119). Default + // DEFAULT_MCP_DEFER_TOOLS_OVER; 0 = always defer. See toolSurfaceForRole. + deferToolsOver?: number; // Injection seam for tests — defaults to the AI SDK factory. createClient?: CreateMcpClient; logger?: LoggerLike; @@ -92,6 +104,17 @@ export class McpClientManager { return merged; } + // Split the role's MCP tools into directly-mounted vs. deferred (issue #119). At or below the + // threshold the whole set is `direct` — byte-identical to plain mounting, no stubs, no tool_search. + // Above it the whole surplus is `deferred`, surfaced name-only and fetched via tool_search. The + // fixed local slots (readFile, bash, …) live in the adapter, not here, so every MCP tool counts. + toolSurfaceForRole(role: Role): ToolSurface { + const tools = this.toolsForRole(role); + const threshold = this.init.deferToolsOver ?? DEFAULT_MCP_DEFER_TOOLS_OVER; + if (Object.keys(tools).length <= threshold) return { direct: tools, deferred: {} }; + return { direct: {}, deferred: tools }; + } + async close(): Promise { const toClose = this.servers; this.servers = []; diff --git a/packages/aitm/src/mcp/tool-search.test.ts b/packages/aitm/src/mcp/tool-search.test.ts new file mode 100644 index 0000000..6c2013b --- /dev/null +++ b/packages/aitm/src/mcp/tool-search.test.ts @@ -0,0 +1,156 @@ +import assert from 'node:assert/strict'; +import { test } from 'node:test'; +import type { ToolSet } from 'ai'; +import { + DEFERRED_TOOLS_CONTRACT, + deferredToolIndexBlock, + guardDeferred, + TOOL_SEARCH_TOOL_NAME, + type ToolSearchInput, + type ToolSearchOutput, + toolSearch, +} from './tool-search.ts'; + +function fakeTool(description: string, extra: Partial = {}): ToolSet[string] { + return { + description, + inputSchema: { type: 'object', properties: { path: { type: 'string' } } }, + ...extra, + } as ToolSet[string]; +} + +const deferred: ToolSet = { + mcp__gh__create_issue: fakeTool('Create a GitHub issue in a repo. Accepts title and body.'), + mcp__gh__list_prs: fakeTool('List pull requests for a repository.'), + mcp__db__query: fakeTool('Run a read-only SQL query against the database.'), +}; + +function search(surface: ReturnType, input: ToolSearchInput): ToolSearchOutput { + const exec = surface.tool.execute; + if (typeof exec !== 'function') throw new Error('tool_search has no execute'); + return ( + exec as (i: ToolSearchInput, o: { toolCallId: string; messages: never[] }) => ToolSearchOutput + )(input, { toolCallId: 'c', messages: [] }); +} + +function renderSearch( + surface: ReturnType, + input: ToolSearchInput, + output: ToolSearchOutput, +): string { + const fn = surface.tool.toModelOutput; + if (typeof fn !== 'function') throw new Error('tool_search has no toModelOutput'); + return ( + fn as (o: { toolCallId: string; input: ToolSearchInput; output: ToolSearchOutput }) => { + type: string; + value: string; + } + )({ toolCallId: 'c', input, output }).value; +} + +test('tool_search select: fetches exact names, activates them, and reports unknowns (issue #119)', () => { + const surface = toolSearch(deferred); + assert.equal(surface.activated.size, 0, 'starts fully deferred'); + const out = search(surface, { query: 'select:mcp__gh__create_issue,mcp__nope__x' }); + assert.deepEqual( + out.matches.map((m) => m.name), + ['mcp__gh__create_issue'], + ); + assert.deepEqual(out.notFound, ['mcp__nope__x']); + assert.ok(surface.activated.has('mcp__gh__create_issue'), 'the match is activated'); + assert.equal(surface.activated.has('mcp__db__query'), false); + assert.deepEqual(out.matches[0]?.parameters, { + type: 'object', + properties: { path: { type: 'string' } }, + }); +}); + +test('tool_search keyword query ranks over names+descriptions and caps at max_results (issue #119)', () => { + let surface = toolSearch(deferred); + const all = search(surface, { query: 'repo' }); + assert.equal(all.matches.length, 2, 'both GitHub tools mention a repo'); + assert.equal( + all.matches.every((m) => m.name.startsWith('mcp__gh__')), + true, + ); + + surface = toolSearch(deferred); + const capped = search(surface, { query: 'repo', max_results: 1 }); + assert.equal(capped.matches.length, 1, 'capped to max_results'); + assert.equal(surface.activated.size, 1); +}); + +test('tool_search keyword query with no hits returns empty matches, activates nothing (issue #119)', () => { + const surface = toolSearch(deferred); + const out = search(surface, { query: 'zzz-nonexistent' }); + assert.deepEqual(out.matches, []); + assert.equal(surface.activated.size, 0); +}); + +test('tool_search toModelOutput renders name, description, and JSON schema; notes activation (issue #119)', () => { + const surface = toolSearch(deferred); + const out = search(surface, { query: 'select:mcp__db__query' }); + const text = renderSearch(surface, { query: 'select:mcp__db__query' }, out); + assert.match(text, /Activated 1 tool/); + assert.match(text, /mcp__db__query/); + assert.match(text, /read-only SQL/); + assert.match(text, /Parameters \(JSON Schema\)/); +}); + +test('deferredToolIndexBlock: contract preamble + one first-sentence line per tool (issue #119)', () => { + const block = deferredToolIndexBlock(deferred); + assert.ok(block.startsWith(DEFERRED_TOOLS_CONTRACT)); + assert.match(block, //); + assert.match(block, /mcp__gh__create_issue: Create a GitHub issue in a repo\./); + assert.equal(block.includes('Accepts title and body'), false, 'only the first sentence'); +}); + +test('deferredToolIndexBlock: empty when nothing is deferred (issue #119)', () => { + assert.equal(deferredToolIndexBlock({}), ''); +}); + +test('guardDeferred: a call before activation returns the fetch-first result, never throws (issue #119)', async () => { + const activated = new Set(); + let ran = 0; + const base = fakeTool('desc', { + execute: async () => { + ran += 1; + return { ok: true }; + }, + }); + const guarded = guardDeferred('mcp__db__query', base, activated); + const exec = guarded.execute as (i: unknown, o: unknown) => Promise; + const out = await exec({}, { toolCallId: 'c', messages: [] }); + assert.equal(ran, 0, 'the real executor did not run'); + const fn = guarded.toModelOutput as (o: unknown) => { type: string; value: string }; + const text = fn({ toolCallId: 'c', input: {}, output: out }).value; + assert.match(text, /deferred tool/); + assert.match(text, /tool_search/); + assert.match(text, /select:mcp__db__query/); +}); + +test('guardDeferred: after activation the real executor runs and its result renders normally (issue #119)', async () => { + const activated = new Set(['mcp__db__query']); + let ran = 0; + const base = fakeTool('desc', { + execute: async () => { + ran += 1; + return { rows: 3 }; + }, + toModelOutput: ({ output }: { output: { rows: number } }) => ({ + type: 'text', + value: `rows=${output.rows}`, + }), + }); + const guarded = guardDeferred('mcp__db__query', base, activated); + const exec = guarded.execute as (i: unknown, o: unknown) => Promise; + const out = await exec({}, { toolCallId: 'c', messages: [] }); + assert.equal(ran, 1); + assert.deepEqual(out, { rows: 3 }); + const fn = guarded.toModelOutput as (o: unknown) => { type: string; value: string }; + assert.equal(fn({ toolCallId: 'c', input: {}, output: out }).value, 'rows=3'); +}); + +test('TOOL_SEARCH_TOOL_NAME is the Claude Code convention name', () => { + assert.equal(TOOL_SEARCH_TOOL_NAME, 'tool_search'); +}); diff --git a/packages/aitm/src/mcp/tool-search.ts b/packages/aitm/src/mcp/tool-search.ts new file mode 100644 index 0000000..2b1f8cc --- /dev/null +++ b/packages/aitm/src/mcp/tool-search.ts @@ -0,0 +1,199 @@ +// Deferred MCP tool loading (issue #119). When a role's MCP surplus exceeds the threshold +// (McpClientManager.toolSurfaceForRole), those tools are mounted name-only: their JSON schemas stay +// out of every request until the model fetches them with the `tool_search` tool. This module owns +// that mechanism — the tool, the per-invocation activation set, the name-only prompt index, and the +// guard that turns a call to an un-activated tool into a typed "fetch first" result instead of a +// provider-level validation error. +// +// Activation is scoped to one subagent invocation: build a fresh surface per agent construction, and +// a new invocation starts fully deferred again. The adapter (run-loop-adapter.ts) composes the +// activation set into the single #102 `prepareStep`, adding activated names to `activeTools`. + +import type { Tool, ToolCallOptions, ToolSet } from 'ai'; +import { tool } from 'ai'; +import { z } from 'zod'; + +export const TOOL_SEARCH_TOOL_NAME = 'tool_search'; +const SELECT_PREFIX = 'select:'; +const DEFAULT_MAX_RESULTS = 5; + +type AnyTool = ToolSet[string]; + +const toolSearchInputSchema = z.object({ + query: z + .string() + .min(1) + .describe( + 'Either `select:name1,name2` to fetch those exact tools by namespaced name, or free-text keywords ranked over tool names + descriptions.', + ), + max_results: z + .number() + .int() + .positive() + .optional() + .describe('Cap on keyword matches (default 5). Ignored for `select:` queries.'), +}); +export type ToolSearchInput = z.infer; + +export type ToolSearchMatch = { name: string; description: string; parameters: unknown }; +export type ToolSearchOutput = { matches: ToolSearchMatch[]; notFound: string[] }; + +export type ToolSearchSurface = { + // The `tool_search` tool. Mount under TOOL_SEARCH_TOOL_NAME exactly when `deferred` is non-empty. + tool: Tool; + // Names activated so far this invocation. The adapter reads this in prepareStep to grow activeTools. + activated: ReadonlySet; + // The name-only system-prompt index for the deferred tools. + indexBlock: string; +}; + +// Build the tool_search surface over a set of deferred (namespaced) tools. The returned `activated` +// set is mutated by the tool's execute as matches are fetched. +export function toolSearch(deferred: ToolSet): ToolSearchSurface { + const activated = new Set(); + const entries = Object.entries(deferred); + const byName = new Map(entries); + + const searchTool = tool({ + description: + 'Fetch the parameter schemas of deferred tools (listed name-only in your prompt) so you can call them. `select:name1,name2` fetches exact names; any other query is keyword-ranked over names and descriptions. Each returned tool becomes callable on your next step.', + inputSchema: toolSearchInputSchema, + execute: (input: ToolSearchInput): ToolSearchOutput => { + const result = input.query.startsWith(SELECT_PREFIX) + ? selectByName(input.query.slice(SELECT_PREFIX.length), byName) + : keywordSearch(input.query, entries, input.max_results ?? DEFAULT_MAX_RESULTS); + for (const match of result.matches) activated.add(match.name); + return result; + }, + toModelOutput: ({ output }) => ({ type: 'text', value: renderToolSearch(output) }), + }); + + return { tool: searchTool, activated, indexBlock: deferredToolIndexBlock(deferred) }; +} + +function selectByName(list: string, byName: Map): ToolSearchOutput { + const names = list + .split(',') + .map((s) => s.trim()) + .filter((s) => s !== ''); + const matches: ToolSearchMatch[] = []; + const notFound: string[] = []; + for (const name of names) { + const found = byName.get(name); + if (found) matches.push(toMatch(name, found)); + else notFound.push(name); + } + return { matches, notFound }; +} + +function keywordSearch(query: string, entries: [string, AnyTool][], max: number): ToolSearchOutput { + const terms = query + .toLowerCase() + .split(/\s+/) + .filter((t) => t !== ''); + const scored = entries + .map(([name, t]) => { + const haystack = `${name} ${toolDescription(t)}`.toLowerCase(); + const score = terms.reduce((n, term) => (haystack.includes(term) ? n + 1 : n), 0); + return { name, tool: t, score }; + }) + .filter((e) => e.score > 0) + .sort((a, b) => b.score - a.score || a.name.localeCompare(b.name)) + .slice(0, max); + return { matches: scored.map((e) => toMatch(e.name, e.tool)), notFound: [] }; +} + +function toMatch(name: string, t: AnyTool): ToolSearchMatch { + return { name, description: toolDescription(t), parameters: toolParameters(t) }; +} + +function toolDescription(t: AnyTool): string { + return typeof t.description === 'string' ? t.description : ''; +} + +// MCP tools carry a JSON-schema `inputSchema`; render it verbatim so the model can construct a call. +function toolParameters(t: AnyTool): unknown { + return t.inputSchema ?? {}; +} + +function renderToolSearch(output: ToolSearchOutput): string { + const parts: string[] = []; + if (output.matches.length === 0) { + parts.push('No tools matched.'); + } else { + parts.push(`Activated ${output.matches.length} tool(s) — callable on your next step:`); + for (const m of output.matches) { + parts.push( + `\n${m.name}\n${m.description}\nParameters (JSON Schema): ${JSON.stringify(m.parameters)}`, + ); + } + } + if (output.notFound.length > 0) { + parts.push(`\nNot found (no such deferred tool): ${output.notFound.join(', ')}`); + } + return parts.join('\n'); +} + +// The always-visible name-only tier: contract preamble + one `: ` line per +// deferred tool. Empty string when nothing is deferred (the caller then mounts no index, no tool). +export const DEFERRED_TOOLS_CONTRACT = + 'The tools below are available but their parameter schemas are NOT loaded. Before calling one, fetch its schema with the `tool_search` tool (`select:` for an exact name, or keywords to search), then call it on a later step. Calling a deferred tool before fetching its schema returns an error, not a result.'; + +export function deferredToolIndexBlock(deferred: ToolSet): string { + const entries = Object.entries(deferred); + if (entries.length === 0) return ''; + const lines = entries.map(([name, t]) => `${name}: ${firstSentence(toolDescription(t))}`); + return [DEFERRED_TOOLS_CONTRACT, '', ...lines, ''].join('\n'); +} + +function firstSentence(text: string): string { + const trimmed = text.replace(/\s+/g, ' ').trim(); + const match = /^(.*?[.!?])(\s|$)/.exec(trimmed); + return match?.[1] ?? trimmed; +} + +// --- guard: a deferred tool called before activation returns a typed result, never throws --- + +const DEFERRED_GUARD_MARKER = '__aitmDeferredGuard'; +type GuardResult = { [DEFERRED_GUARD_MARKER]: true; error: string }; + +function guardResult(name: string): GuardResult { + return { + [DEFERRED_GUARD_MARKER]: true, + error: `${name} is a deferred tool whose schema is not loaded. Call \`${TOOL_SEARCH_TOOL_NAME}\` with \`select:${name}\` first, then call ${name} on a later step.`, + }; +} + +function isGuardResult(v: unknown): v is GuardResult { + return typeof v === 'object' && v !== null && DEFERRED_GUARD_MARKER in v; +} + +// Wrap a deferred tool so a call before its name is activated short-circuits to the fetch-first +// result (with the same guidance rendered as the model-visible text), and a call after activation +// runs the real executor. Preserves the base tool's own toModelOutput for genuine results. +export function guardDeferred( + name: string, + base: AnyTool, + activated: ReadonlySet, +): AnyTool { + const baseExecute = base.execute; + const baseToModelOutput = base.toModelOutput; + return { + ...base, + execute: async (input: unknown, options: ToolCallOptions) => { + if (!activated.has(name)) return guardResult(name); + if (!baseExecute) return guardResult(name); + return baseExecute(input as never, options); + }, + toModelOutput: (options: { toolCallId: string; input: unknown; output: unknown }) => { + if (isGuardResult(options.output)) return { type: 'text', value: options.output.error }; + if (baseToModelOutput) { + return (baseToModelOutput as (o: typeof options) => { type: 'text'; value: string })( + options, + ); + } + const out = options.output; + return { type: 'text', value: typeof out === 'string' ? out : JSON.stringify(out) }; + }, + } as AnyTool; +}