@skydiveai/pi-extensions 0.1.0-beta.1965 → 0.1.0-beta.1968

package/dist/index.mjs

@@ -1069,6 +1069,126 @@ async function connectClient(id, config, opts = {}) {
 	};
 }
 //#endregion
+//#region src/extensions/mcp/limits.ts
+/**
+* Tool-budget guardrails for MCP registration.
+*
+* Why this exists: the agent loop sends the *full* tool list — every tool's
+* name, description, and JSON-schema `parameters` — to the model on every
+* prompt. MCP servers add tools without bound: a handful of chatty servers (or
+* one server that exposes 100+ tools, or a few tools with enormous schemas) can
+* push the registered set past what the model can accept, and the request is
+* rejected before the turn even runs. Because `reconcile` re-registers from
+* `mcp.config.json` on *every* turn, an over-limit config bricks the harness on
+* a loop — the agent can't get a turn to run in order to edit the config back
+* down. Worse, the person can't tell *why*: tools just stop working.
+*
+* What actually overflows the request is *tokens*, not tool count — a few tools
+* with deeply-nested schemas and long descriptions cost more than a hundred
+* trivial ones. And how many tokens are safe depends on the *model*: a 200K
+* context window can afford far more tool surface than a 32K one. So the primary
+* limiter is a **token budget derived from the active model's context window**,
+* with a fixed tool-count cap as a coarse secondary guard (and the fallback
+* when the model — hence its window — isn't known at reconcile time).
+*
+* The fix is to make registration bounded and fail-soft. We register in
+* deterministic config order and stop before we blow the budget, recording how
+* much we dropped so the agent is *told* it hit the limit and which servers
+* were truncated. A config that would have bricked the harness now degrades to
+* "a bounded set of tools plus a loud warning", which the agent can act on by
+* pruning servers.
+*
+* Everything is env-overridable so the ceilings can be tuned per deployment
+* without a release, but ships with conservative defaults. A cap value of 0 (or
+* a non-finite / negative override) disables that cap — an explicit escape
+* hatch, not the default.
+*/
+const env = process.env;
+/**
+* Fraction of the model's context window we're willing to spend on MCP tool
+* schemas. Tool definitions are sent on every prompt, so they permanently eat
+* into the window available for the conversation — a small slice is plenty for
+* a healthy tool set while leaving the vast majority for actual work. 0.15 of a
+* 200K window is ~30K tokens of tool schema, comfortably more than any sane
+* MCP setup; of a 32K window it's ~4.8K, which correctly forces truncation
+* before a small model chokes.
+*/
+const DEFAULT_MCP_TOOL_TOKEN_BUDGET_FRACTION = .15;
+/**
+* Floor for the token budget when the model's context window is unknown at
+* reconcile time (e.g. the model hasn't been resolved yet). Generous enough not
+* to truncate an ordinary tool set, low enough to still catch a runaway.
+*/
+const DEFAULT_MCP_TOOL_TOKEN_BUDGET_FLOOR = 16e3;
+/**
+* Read a positive-integer cap from an env var, falling back to `fallback`.
+* A `0` override (or any non-finite / negative value) means "no cap" and is
+* returned as `Infinity`, so callers can compare against it directly.
+*/
+function readCap(raw, fallback) {
+	if (raw === void 0 || raw.trim() === "") return fallback;
+	const parsed = Number(raw);
+	if (!Number.isFinite(parsed) || parsed < 0) return Number.POSITIVE_INFINITY;
+	if (parsed === 0) return Number.POSITIVE_INFINITY;
+	return Math.floor(parsed);
+}
+/** Read a fraction in (0, 1] from an env var, falling back to `fallback`. */
+function readFraction(raw, fallback) {
+	if (raw === void 0 || raw.trim() === "") return fallback;
+	const parsed = Number(raw);
+	if (!Number.isFinite(parsed) || parsed <= 0 || parsed > 1) return fallback;
+	return parsed;
+}
+/** Read a non-negative integer from an env var, falling back to `fallback`. */
+function readNonNegativeInt(raw, fallback) {
+	if (raw === void 0 || raw.trim() === "") return fallback;
+	const parsed = Number(raw);
+	if (!Number.isFinite(parsed) || parsed < 0) return fallback;
+	return Math.floor(parsed);
+}
+/**
+* Resolve the active caps from the environment. Read once per reconcile so a
+* deployment can retune without a restart, cheap enough not to cache.
+*/
+function resolveMcpToolLimits() {
+	return {
+		maxTotalTools: readCap(env["SKYDIVE_MCP_MAX_TOTAL_TOOLS"], 128),
+		maxToolsPerServer: readCap(env["SKYDIVE_MCP_MAX_TOOLS_PER_SERVER"], 50),
+		tokenBudgetFraction: readFraction(env["SKYDIVE_MCP_TOOL_TOKEN_BUDGET_FRACTION"], DEFAULT_MCP_TOOL_TOKEN_BUDGET_FRACTION),
+		tokenBudgetFloor: readNonNegativeInt(env["SKYDIVE_MCP_TOOL_TOKEN_BUDGET_FLOOR"], DEFAULT_MCP_TOOL_TOKEN_BUDGET_FLOOR)
+	};
+}
+/**
+* The token budget for MCP tool schemas, given the active model's context
+* window (or undefined when the model isn't known yet).
+*
+* With a known window we spend `tokenBudgetFraction` of it, but never less than
+* the floor — a tiny window shouldn't collapse the budget to near-zero and
+* strand every tool. With no window we fall back to the floor outright.
+*/
+function resolveTokenBudget(limits, contextWindow) {
+	if (contextWindow === void 0 || !Number.isFinite(contextWindow)) return limits.tokenBudgetFloor;
+	return Math.max(limits.tokenBudgetFloor, Math.floor(contextWindow * limits.tokenBudgetFraction));
+}
+/**
+* Estimate the tokens an MCP tool's *definition* costs in the request. We
+* serialize what's actually sent to the model — the tool name, its description,
+* and its JSON-schema parameters — and apply pi's own ~chars/4 heuristic
+* (`estimateTokens` in the coding agent uses the same convention; there is no
+* per-provider tokenizer to lean on, and the provider's real usage is only
+* known after the response). Rough by design, but it tracks the true cost far
+* better than a flat per-tool count: a fat schema is charged for its fatness.
+*/
+function estimateToolTokens(tool) {
+	let chars = tool.name.length + (tool.description?.length ?? 0);
+	if (tool.inputSchema !== void 0 && tool.inputSchema !== null) try {
+		chars += JSON.stringify(tool.inputSchema).length;
+	} catch {
+		chars += 256;
+	}
+	return Math.ceil(chars / 4);
+}
+//#endregion
 //#region src/extensions/mcp/mcp-config.ts
 /**
 * mcp.config.json schema + loader. Split out from the extension so it has
@@ -1275,19 +1395,26 @@ var McpExtension = class {
 		}));
 		this.registeredMcpToolNames.add(name);
 	}
-	async reconcile({ pi, configPath, connectTimeoutMs }) {
+	async reconcile({ pi, configPath, connectTimeoutMs, contextWindow }) {
 		let config;
 		try {
 			config = await loadMcpConfig(configPath);
 		} catch (err) {
 			throw new Error(`Failed to load MCP config: ${err instanceof Error ? err.message : String(err)}`);
 		}
+		const limits = resolveMcpToolLimits();
+		const tokenBudget = resolveTokenBudget(limits, contextWindow);
 		const summary = {
 			added: [],
 			removed: [],
 			refreshed: [],
 			errors: [],
 			totalTools: 0,
+			droppedTools: 0,
+			limits,
+			tokenBudget,
+			tokensUsed: 0,
+			toolCounts: {},
 			servers: {}
 		};
 		const desiredIds = new Set(Object.keys(config.servers));
@@ -1311,14 +1438,46 @@ var McpExtension = class {
 			if (outcome.error) summary.errors.push(outcome.error);
 			if (outcome.change === "added") summary.added.push(id);
 			else if (outcome.change === "refreshed") summary.refreshed.push(id);
-			if (outcome.tools) for (const tool of outcome.tools.list) {
-				this.registerMcpTool({
-					pi,
-					serverId: id,
-					client: outcome.tools.client,
-					tool
-				});
-				summary.totalTools++;
+			if (outcome.tools) {
+				const advertised = outcome.tools.list.length;
+				const perServerRoom = Math.min(advertised, limits.maxToolsPerServer);
+				let registered = 0;
+				let serverTokens = 0;
+				let droppedReason = null;
+				for (const tool of outcome.tools.list) {
+					if (summary.totalTools >= limits.maxTotalTools) {
+						droppedReason = "total";
+						break;
+					}
+					if (registered >= perServerRoom) {
+						droppedReason = "per_server";
+						break;
+					}
+					const cost = estimateToolTokens(tool);
+					if (summary.tokensUsed + cost > tokenBudget && summary.totalTools > 0) {
+						droppedReason = "tokens";
+						break;
+					}
+					this.registerMcpTool({
+						pi,
+						serverId: id,
+						client: outcome.tools.client,
+						tool
+					});
+					registered++;
+					serverTokens += cost;
+					summary.totalTools++;
+					summary.tokensUsed += cost;
+				}
+				const dropped = advertised - registered;
+				if (dropped > 0) summary.droppedTools += dropped;
+				else droppedReason = null;
+				summary.toolCounts[id] = {
+					advertised,
+					registered,
+					tokens: serverTokens,
+					droppedReason
+				};
 			}
 		}
 		return summary;
@@ -1563,10 +1722,11 @@ var McpExtension = class {
 			}
 		};
 	}
-	async reconcileAndRecordMtime({ pi, configPath, reason }) {
+	async reconcileAndRecordMtime({ pi, configPath, reason, contextWindow }) {
 		const summary = await this.reconcile({
 			pi,
-			configPath
+			configPath,
+			contextWindow
 		});
 		this.lastConfigMtimeMs = await readConfigMtimeMs(configPath);
 		if (reason !== "session_start" && summaryHasChanges(summary)) this.pendingMcpUpdate = summary;
@@ -1574,6 +1734,9 @@ var McpExtension = class {
 			event: "mcp_reconcile",
 			reason,
 			total_tools: summary.totalTools,
+			tokens_used: summary.tokensUsed,
+			token_budget: summary.tokenBudget,
+			dropped_tools: summary.droppedTools,
 			added: summary.added,
 			removed: summary.removed,
 			refreshed: summary.refreshed,
@@ -1590,7 +1753,8 @@ var McpExtension = class {
 					await this.reconcileAndRecordMtime({
 						pi,
 						configPath,
-						reason: "session_start"
+						reason: "session_start",
+						contextWindow: ctx.model?.contextWindow
 					});
 				} catch (err) {
 					log$12.error({
@@ -1616,7 +1780,8 @@ var McpExtension = class {
 					await this.reconcileAndRecordMtime({
 						pi,
 						configPath,
-						reason: "auto_reload"
+						reason: "auto_reload",
+						contextWindow: ctx.model?.contextWindow
 					});
 				} catch (err) {
 					log$12.error({
@@ -1636,7 +1801,8 @@ var McpExtension = class {
 						const summary = await this.reconcileAndRecordMtime({
 							pi,
 							configPath,
-							reason: "tool"
+							reason: "tool",
+							contextWindow: ctx.model?.contextWindow
 						});
 						return {
 							content: [{
@@ -1695,9 +1861,29 @@ function appendStderrBlock(lines, stderr) {
 	for (const line of stderr.trimEnd().split("\n")) lines.push(`    ${line}`);
 	lines.push("    ---");
 }
+/**
+* Human-readable lines describing any budget-driven truncation, shared by
+* `summaryText` (the reload_mcp tool output) and `formatMcpUpdateMessage` (the
+* synthetic continuation). Empty when nothing was dropped.
+*/
+function truncationLines(summary) {
+	if (summary.droppedTools <= 0) return [];
+	const lines = [];
+	const { maxTotalTools, maxToolsPerServer } = summary.limits;
+	const totalCapHint = Number.isFinite(maxTotalTools) ? `${maxTotalTools}` : "unlimited";
+	lines.push(`  WARNING: ${summary.droppedTools} MCP tool(s) were NOT registered because a tool budget was hit (token budget: ~${summary.tokenBudget} tokens, used ~${summary.tokensUsed}; total-tool cap: ${totalCapHint}, per-server cap: ${Number.isFinite(maxToolsPerServer) ? maxToolsPerServer : "unlimited"}).`);
+	lines.push("  Tool definitions are sent to the model on every prompt; too many (or too-large) tool schemas push the request over the limit, so they are budgeted against the model context window. Prune servers from mcp.config.json to bring the tool surface down.");
+	for (const [id, count] of Object.entries(summary.toolCounts)) {
+		if (count.droppedReason === null) continue;
+		const why = count.droppedReason === "tokens" ? "token budget exhausted" : count.droppedReason === "total" ? "total-tool cap reached" : "per-server cap";
+		lines.push(`    - ${id}: registered ${count.registered}/${count.advertised} tools (~${count.tokens} tokens, ${why}).`);
+	}
+	return lines;
+}
 function summaryText(summary) {
 	const lines = [];
 	lines.push(`MCP reconcile complete: ${summary.totalTools} tool(s) live.`);
+	lines.push(...truncationLines(summary));
 	if (summary.added.length > 0) lines.push(`  Added: ${summary.added.join(", ")}`);
 	if (summary.refreshed.length > 0) lines.push(`  Refreshed: ${summary.refreshed.join(", ")}`);
 	if (summary.removed.length > 0) lines.push(`  Removed: ${summary.removed.join(", ")}`);
@@ -1723,7 +1909,7 @@ function summaryText(summary) {
 	return lines.join("\n");
 }
 function summaryHasChanges(summary) {
-	return summary.added.length > 0 || summary.removed.length > 0 || summary.refreshed.length > 0 || summary.errors.length > 0;
+	return summary.added.length > 0 || summary.removed.length > 0 || summary.refreshed.length > 0 || summary.errors.length > 0 || summary.droppedTools > 0;
 }
 /**
 * Format a queued tool-update as a synthetic system-style message for
@@ -1736,6 +1922,10 @@ function formatMcpUpdateMessage(summary) {
 	if (summary.added.length > 0) lines.push(`Newly available servers: ${summary.added.join(", ")}`);
 	if (summary.refreshed.length > 0) lines.push(`Refreshed servers: ${summary.refreshed.join(", ")}`);
 	if (summary.removed.length > 0) lines.push(`Removed servers (and their tools): ${summary.removed.join(", ")}`);
+	if (summary.droppedTools > 0) {
+		lines.push("");
+		lines.push(...truncationLines(summary));
+	}
 	const pending = pendingAuthEntries(summary);
 	if (pending.length > 0) {
 		lines.push("");

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@skydiveai/pi-extensions",
-  "version": "0.1.0-beta.1965",
+  "version": "0.1.0-beta.1968",
   "homepage": "https://skydive.com",
   "license": "MIT",
   "author": "Create, Inc.",