@oh-my-pi/pi-coding-agent 6.2.0 → 6.7.0

package/src/core/mcp/manager.ts

@@ -7,11 +7,48 @@
 
 import type { TSchema } from "@sinclair/typebox";
 import type { CustomTool } from "../custom-tools/types";
+import { logger } from "../logger";
 import { connectToServer, disconnectServer, listTools } from "./client";
 import { loadAllMCPConfigs, validateServerConfig } from "./config";
 import type { MCPToolDetails } from "./tool-bridge";
-import { createMCPTools } from "./tool-bridge";
-import type { MCPServerConfig, MCPServerConnection } from "./types";
+import { DeferredMCPTool, MCPTool } from "./tool-bridge";
+import type { MCPToolCache } from "./tool-cache";
+import type { MCPServerConfig, MCPServerConnection, MCPToolDefinition } from "./types";
+
+type SourceMeta = import("../../capability/types").SourceMeta;
+
+type ToolLoadResult = {
+	connection: MCPServerConnection;
+	serverTools: MCPToolDefinition[];
+};
+
+type TrackedPromise<T> = {
+	promise: Promise<T>;
+	status: "pending" | "fulfilled" | "rejected";
+	value?: T;
+	reason?: unknown;
+};
+
+const STARTUP_TIMEOUT_MS = 250;
+
+function trackPromise<T>(promise: Promise<T>): TrackedPromise<T> {
+	const tracked: TrackedPromise<T> = { promise, status: "pending" };
+	promise.then(
+		(value) => {
+			tracked.status = "fulfilled";
+			tracked.value = value;
+		},
+		(reason) => {
+			tracked.status = "rejected";
+			tracked.reason = reason;
+		},
+	);
+	return tracked;
+}
+
+function delay(ms: number): Promise<void> {
+	return new Promise((resolve) => setTimeout(resolve, ms));
+}
 
 /** Result of loading MCP tools */
 export interface MCPLoadResult {
@@ -43,8 +80,14 @@ export interface MCPDiscoverOptions {
 export class MCPManager {
 	private connections = new Map<string, MCPServerConnection>();
 	private tools: CustomTool<TSchema, MCPToolDetails>[] = [];
+	private pendingConnections = new Map<string, Promise<MCPServerConnection>>();
+	private pendingToolLoads = new Map<string, Promise<ToolLoadResult>>();
+	private sources = new Map<string, SourceMeta>();
 
-	constructor(private cwd: string) {}
+	constructor(
+		private cwd: string,
+		private toolCache: MCPToolCache | null = null,
+	) {}
 
 	/**
 	 * Discover and connect to all MCP servers from .mcp.json files.
@@ -66,24 +109,41 @@ export class MCPManager {
 	 */
 	async connectServers(
 		configs: Record<string, MCPServerConfig>,
-		sources: Record<string, import("../../capability/types").SourceMeta>,
+		sources: Record<string, SourceMeta>,
 		onConnecting?: (serverNames: string[]) => void,
 	): Promise<MCPLoadResult> {
+		type ConnectionTask = {
+			name: string;
+			config: MCPServerConfig;
+			tracked: TrackedPromise<ToolLoadResult>;
+			toolsPromise: Promise<ToolLoadResult>;
+		};
+
 		const errors = new Map<string, string>();
-		const connectedServers: string[] = [];
+		const connectedServers = new Set<string>();
 		const allTools: CustomTool<TSchema, MCPToolDetails>[] = [];
+		const reportedErrors = new Set<string>();
+		let allowBackgroundLogging = false;
 
 		// Prepare connection tasks
-		const connectionTasks: Array<{
-			name: string;
-			config: MCPServerConfig;
-			validationErrors: string[];
-		}> = [];
+		const connectionTasks: ConnectionTask[] = [];
 
 		for (const [name, config] of Object.entries(configs)) {
+			if (sources[name]) {
+				this.sources.set(name, sources[name]);
+				const existing = this.connections.get(name);
+				if (existing) {
+					existing._source = sources[name];
+				}
+			}
+
 			// Skip if already connected
 			if (this.connections.has(name)) {
-				connectedServers.push(name);
+				connectedServers.add(name);
+				continue;
+			}
+
+			if (this.pendingConnections.has(name) || this.pendingToolLoads.has(name)) {
 				continue;
 			}
 
@@ -91,59 +151,128 @@ export class MCPManager {
 			const validationErrors = validateServerConfig(name, config);
 			if (validationErrors.length > 0) {
 				errors.set(name, validationErrors.join("; "));
+				reportedErrors.add(name);
 				continue;
 			}
 
-			connectionTasks.push({ name, config, validationErrors });
+			const connectionPromise = connectToServer(name, config).then(
+				(connection) => {
+					if (sources[name]) {
+						connection._source = sources[name];
+					}
+					if (this.pendingConnections.get(name) === connectionPromise) {
+						this.pendingConnections.delete(name);
+						this.connections.set(name, connection);
+					}
+					return connection;
+				},
+				(error) => {
+					if (this.pendingConnections.get(name) === connectionPromise) {
+						this.pendingConnections.delete(name);
+					}
+					throw error;
+				},
+			);
+			this.pendingConnections.set(name, connectionPromise);
+
+			const toolsPromise = connectionPromise.then(async (connection) => {
+				const serverTools = await listTools(connection);
+				return { connection, serverTools };
+			});
+			this.pendingToolLoads.set(name, toolsPromise);
+
+			const tracked = trackPromise(toolsPromise);
+			connectionTasks.push({ name, config, tracked, toolsPromise });
+
+			void toolsPromise
+				.then(({ connection, serverTools }) => {
+					if (this.pendingToolLoads.get(name) !== toolsPromise) return;
+					this.pendingToolLoads.delete(name);
+					const customTools = MCPTool.fromTools(connection, serverTools);
+					this.replaceServerTools(name, customTools);
+					void this.toolCache?.set(name, config, serverTools);
+				})
+				.catch((error) => {
+					if (this.pendingToolLoads.get(name) !== toolsPromise) return;
+					this.pendingToolLoads.delete(name);
+					if (!allowBackgroundLogging || reportedErrors.has(name)) return;
+					const message = error instanceof Error ? error.message : String(error);
+					logger.error("MCP tool load failed", { path: `mcp:${name}`, error: message });
+				});
 		}
 
 		// Notify about servers we're connecting to
 		if (connectionTasks.length > 0 && onConnecting) {
-			onConnecting(connectionTasks.map((t) => t.name));
+			onConnecting(connectionTasks.map((task) => task.name));
 		}
 
-		// Connect to all servers in parallel
-		const results = await Promise.allSettled(
-			connectionTasks.map(async ({ name, config }) => {
-				const connection = await connectToServer(name, config);
-				// Attach source metadata to connection
-				if (sources[name]) {
-					connection._source = sources[name];
+		if (connectionTasks.length > 0) {
+			await Promise.race([
+				Promise.allSettled(connectionTasks.map((task) => task.tracked.promise)),
+				delay(STARTUP_TIMEOUT_MS),
+			]);
+
+			const cachedTools = new Map<string, MCPToolDefinition[]>();
+			const pendingTasks = connectionTasks.filter((task) => task.tracked.status === "pending");
+
+			if (pendingTasks.length > 0) {
+				if (this.toolCache) {
+					await Promise.all(
+						pendingTasks.map(async (task) => {
+							const cached = await this.toolCache?.get(task.name, task.config);
+							if (cached) {
+								cachedTools.set(task.name, cached);
+							}
+						}),
+					);
+				}
+
+				const pendingWithoutCache = pendingTasks.filter((task) => !cachedTools.has(task.name));
+				if (pendingWithoutCache.length > 0) {
+					await Promise.allSettled(pendingWithoutCache.map((task) => task.tracked.promise));
+				}
+			}
+
+			for (const task of connectionTasks) {
+				const { name } = task;
+				if (task.tracked.status === "fulfilled") {
+					const value = task.tracked.value;
+					if (!value) continue;
+					const { connection, serverTools } = value;
+					connectedServers.add(name);
+					allTools.push(...MCPTool.fromTools(connection, serverTools));
+				} else if (task.tracked.status === "rejected") {
+					const message =
+						task.tracked.reason instanceof Error ? task.tracked.reason.message : String(task.tracked.reason);
+					errors.set(name, message);
+					reportedErrors.add(name);
+				} else {
+					const cached = cachedTools.get(name);
+					if (cached) {
+						const source = this.sources.get(name);
+						allTools.push(...DeferredMCPTool.fromTools(name, cached, () => this.waitForConnection(name), source));
+					}
 				}
-				const serverTools = await listTools(connection);
-				return { name, connection, serverTools };
-			}),
-		);
-
-		// Process results
-		for (let i = 0; i < results.length; i++) {
-			const result = results[i];
-			const { name } = connectionTasks[i];
-
-			if (result.status === "fulfilled") {
-				const { connection, serverTools } = result.value;
-				this.connections.set(name, connection);
-				connectedServers.push(name);
-
-				const customTools = createMCPTools(connection, serverTools);
-				allTools.push(...customTools);
-			} else {
-				const message = result.reason instanceof Error ? result.reason.message : String(result.reason);
-				errors.set(name, message);
 			}
 		}
 
 		// Update cached tools
 		this.tools = allTools;
+		allowBackgroundLogging = true;
 
 		return {
 			tools: allTools,
 			errors,
-			connectedServers,
+			connectedServers: Array.from(connectedServers),
 			exaApiKeys: [], // Will be populated by discoverAndConnect
 		};
 	}
 
+	private replaceServerTools(name: string, tools: CustomTool<TSchema, MCPToolDetails>[]): void {
+		this.tools = this.tools.filter((t) => !t.name.startsWith(`mcp_${name}_`));
+		this.tools.push(...tools);
+	}
+
 	/**
 	 * Get all loaded tools.
 	 */
@@ -158,6 +287,24 @@ export class MCPManager {
 		return this.connections.get(name);
 	}
 
+	/**
+	 * Get the source metadata for a server.
+	 */
+	getSource(name: string): SourceMeta | undefined {
+		return this.sources.get(name) ?? this.connections.get(name)?._source;
+	}
+
+	/**
+	 * Wait for a connection to complete (or fail).
+	 */
+	async waitForConnection(name: string): Promise<MCPServerConnection> {
+		const connection = this.connections.get(name);
+		if (connection) return connection;
+		const pending = this.pendingConnections.get(name);
+		if (pending) return pending;
+		throw new Error(`MCP server not connected: ${name}`);
+	}
+
 	/**
 	 * Get all connected server names.
 	 */
@@ -169,11 +316,15 @@ export class MCPManager {
 	 * Disconnect from a specific server.
 	 */
 	async disconnectServer(name: string): Promise<void> {
-		const connection = this.connections.get(name);
-		if (!connection) return;
+		this.pendingConnections.delete(name);
+		this.pendingToolLoads.delete(name);
+		this.sources.delete(name);
 
-		await disconnectServer(connection);
-		this.connections.delete(name);
+		const connection = this.connections.get(name);
+		if (connection) {
+			await disconnectServer(connection);
+			this.connections.delete(name);
+		}
 
 		// Remove tools from this server
 		this.tools = this.tools.filter((t) => !t.name.startsWith(`mcp_${name}_`));
@@ -186,6 +337,9 @@ export class MCPManager {
 		const promises = Array.from(this.connections.values()).map((conn) => disconnectServer(conn));
 		await Promise.allSettled(promises);
 
+		this.pendingConnections.clear();
+		this.pendingToolLoads.clear();
+		this.sources.clear();
 		this.connections.clear();
 		this.tools = [];
 	}
@@ -202,11 +356,11 @@ export class MCPManager {
 
 		// Reload tools
 		const serverTools = await listTools(connection);
-		const customTools = createMCPTools(connection, serverTools);
+		const customTools = MCPTool.fromTools(connection, serverTools);
+		void this.toolCache?.set(name, connection.config, serverTools);
 
 		// Replace tools from this server
-		this.tools = this.tools.filter((t) => !t.name.startsWith(`mcp_${name}_`));
-		this.tools.push(...customTools);
+		this.replaceServerTools(name, customTools);
 	}
 
 	/**

package/src/core/mcp/tool-bridge.ts

@@ -4,8 +4,10 @@
  * Converts MCP tool definitions to CustomTool format for the agent.
  */
 
+import type { AgentToolUpdateCallback } from "@oh-my-pi/pi-agent-core";
 import type { TSchema } from "@sinclair/typebox";
-import type { CustomTool, CustomToolResult } from "../custom-tools/types";
+import type { SourceMeta } from "../../capability/types";
+import type { CustomTool, CustomToolContext, CustomToolResult } from "../custom-tools/types";
 import { callTool } from "./client";
 import type { MCPContent, MCPServerConnection, MCPToolDefinition } from "./types";
 
@@ -88,69 +90,155 @@ export function parseMCPToolName(name: string): { serverName: string; toolName:
 }
 
 /**
- * Convert an MCP tool definition to a CustomTool.
+ * CustomTool wrapping an MCP tool with an active connection.
  */
-export function createMCPTool(
-	connection: MCPServerConnection,
-	tool: MCPToolDefinition,
-): CustomTool<TSchema, MCPToolDetails> {
-	const name = createMCPToolName(connection.name, tool.name);
-	const schema = convertSchema(tool.inputSchema);
-
-	return {
-		name,
-		label: `${connection.name}/${tool.name}`,
-		description: tool.description ?? `MCP tool from ${connection.name}`,
-		parameters: schema,
-
-		async execute(_toolCallId, params, _onUpdate, _ctx, _signal): Promise<CustomToolResult<MCPToolDetails>> {
-			try {
-				const result = await callTool(connection, tool.name, params as Record<string, unknown>);
-
-				const text = formatMCPContent(result.content);
-				const details: MCPToolDetails = {
-					serverName: connection.name,
-					mcpToolName: tool.name,
-					isError: result.isError,
-					rawContent: result.content,
-					provider: connection._source?.provider,
-					providerName: connection._source?.providerName,
-				};
+export class MCPTool implements CustomTool<TSchema, MCPToolDetails> {
+	public readonly name: string;
+	public readonly label: string;
+	public readonly description: string;
+	public readonly parameters: TSchema;
+
+	/** Create MCPTool instances for all tools from an MCP server connection */
+	static fromTools(connection: MCPServerConnection, tools: MCPToolDefinition[]): MCPTool[] {
+		return tools.map((tool) => new MCPTool(connection, tool));
+	}
 
-				if (result.isError) {
-					return {
-						content: [{ type: "text", text: `Error: ${text}` }],
-						details,
-					};
-				}
+	constructor(
+		private readonly connection: MCPServerConnection,
+		private readonly tool: MCPToolDefinition,
+	) {
+		this.name = createMCPToolName(connection.name, tool.name);
+		this.label = `${connection.name}/${tool.name}`;
+		this.description = tool.description ?? `MCP tool from ${connection.name}`;
+		this.parameters = convertSchema(tool.inputSchema);
+	}
 
+	async execute(
+		_toolCallId: string,
+		params: unknown,
+		_onUpdate: AgentToolUpdateCallback<MCPToolDetails> | undefined,
+		_ctx: CustomToolContext,
+		_signal?: AbortSignal,
+	): Promise<CustomToolResult<MCPToolDetails>> {
+		try {
+			const result = await callTool(this.connection, this.tool.name, params as Record<string, unknown>);
+
+			const text = formatMCPContent(result.content);
+			const details: MCPToolDetails = {
+				serverName: this.connection.name,
+				mcpToolName: this.tool.name,
+				isError: result.isError,
+				rawContent: result.content,
+				provider: this.connection._source?.provider,
+				providerName: this.connection._source?.providerName,
+			};
+
+			if (result.isError) {
 				return {
-					content: [{ type: "text", text }],
+					content: [{ type: "text", text: `Error: ${text}` }],
 					details,
 				};
-			} catch (error) {
-				const message = error instanceof Error ? error.message : String(error);
-				return {
-					content: [{ type: "text", text: `MCP error: ${message}` }],
-					details: {
-						serverName: connection.name,
-						mcpToolName: tool.name,
-						isError: true,
-						provider: connection._source?.provider,
-						providerName: connection._source?.providerName,
-					},
-				};
 			}
-		},
-	};
+
+			return {
+				content: [{ type: "text", text }],
+				details,
+			};
+		} catch (error) {
+			const message = error instanceof Error ? error.message : String(error);
+			return {
+				content: [{ type: "text", text: `MCP error: ${message}` }],
+				details: {
+					serverName: this.connection.name,
+					mcpToolName: this.tool.name,
+					isError: true,
+					provider: this.connection._source?.provider,
+					providerName: this.connection._source?.providerName,
+				},
+			};
+		}
+	}
 }
 
 /**
- * Convert all tools from an MCP server to CustomTools.
+ * CustomTool wrapping an MCP tool with deferred connection resolution.
  */
-export function createMCPTools(
-	connection: MCPServerConnection,
-	tools: MCPToolDefinition[],
-): CustomTool<TSchema, MCPToolDetails>[] {
-	return tools.map((tool) => createMCPTool(connection, tool));
+export class DeferredMCPTool implements CustomTool<TSchema, MCPToolDetails> {
+	public readonly name: string;
+	public readonly label: string;
+	public readonly description: string;
+	public readonly parameters: TSchema;
+	private readonly fallbackProvider: string | undefined;
+	private readonly fallbackProviderName: string | undefined;
+
+	/** Create DeferredMCPTool instances for all tools from an MCP server */
+	static fromTools(
+		serverName: string,
+		tools: MCPToolDefinition[],
+		getConnection: () => Promise<MCPServerConnection>,
+		source?: SourceMeta,
+	): DeferredMCPTool[] {
+		return tools.map((tool) => new DeferredMCPTool(serverName, tool, getConnection, source));
+	}
+
+	constructor(
+		private readonly serverName: string,
+		private readonly tool: MCPToolDefinition,
+		private readonly getConnection: () => Promise<MCPServerConnection>,
+		source?: SourceMeta,
+	) {
+		this.name = createMCPToolName(serverName, tool.name);
+		this.label = `${serverName}/${tool.name}`;
+		this.description = tool.description ?? `MCP tool from ${serverName}`;
+		this.parameters = convertSchema(tool.inputSchema);
+		this.fallbackProvider = source?.provider;
+		this.fallbackProviderName = source?.providerName;
+	}
+
+	async execute(
+		_toolCallId: string,
+		params: unknown,
+		_onUpdate: AgentToolUpdateCallback<MCPToolDetails> | undefined,
+		_ctx: CustomToolContext,
+		_signal?: AbortSignal,
+	): Promise<CustomToolResult<MCPToolDetails>> {
+		try {
+			const connection = await this.getConnection();
+			const result = await callTool(connection, this.tool.name, params as Record<string, unknown>);
+
+			const text = formatMCPContent(result.content);
+			const details: MCPToolDetails = {
+				serverName: this.serverName,
+				mcpToolName: this.tool.name,
+				isError: result.isError,
+				rawContent: result.content,
+				provider: connection._source?.provider ?? this.fallbackProvider,
+				providerName: connection._source?.providerName ?? this.fallbackProviderName,
+			};
+
+			if (result.isError) {
+				return {
+					content: [{ type: "text", text: `Error: ${text}` }],
+					details,
+				};
+			}
+
+			return {
+				content: [{ type: "text", text }],
+				details,
+			};
+		} catch (error) {
+			const message = error instanceof Error ? error.message : String(error);
+			return {
+				content: [{ type: "text", text: `MCP error: ${message}` }],
+				details: {
+					serverName: this.serverName,
+					mcpToolName: this.tool.name,
+					isError: true,
+					provider: this.fallbackProvider,
+					providerName: this.fallbackProviderName,
+				},
+			};
+		}
+	}
 }

package/src/core/mcp/tool-cache.ts

@@ -0,0 +1,122 @@
+/**
+ * MCP tool cache.
+ *
+ * Stores tool definitions per server in agent.db for fast startup.
+ */
+
+import type { AgentStorage } from "../agent-storage";
+import { logger } from "../logger";
+import type { MCPServerConfig, MCPToolDefinition } from "./types";
+
+const CACHE_VERSION = 1;
+const CACHE_PREFIX = "mcp_tools:";
+const CACHE_TTL_MS = 30 * 24 * 60 * 60 * 1000;
+
+type MCPToolCachePayload = {
+	version: number;
+	configHash: string;
+	tools: MCPToolDefinition[];
+};
+
+function isRecord(value: unknown): value is Record<string, unknown> {
+	return !!value && typeof value === "object" && !Array.isArray(value);
+}
+
+function stableClone(value: unknown): unknown {
+	if (Array.isArray(value)) {
+		return value.map((item) => stableClone(item));
+	}
+	if (isRecord(value)) {
+		const sorted: Record<string, unknown> = {};
+		for (const key of Object.keys(value).sort()) {
+			sorted[key] = stableClone(value[key]);
+		}
+		return sorted;
+	}
+	return value;
+}
+
+function stableStringify(value: unknown): string {
+	return JSON.stringify(stableClone(value));
+}
+
+function toHex(buffer: ArrayBuffer): string {
+	const bytes = new Uint8Array(buffer);
+	let output = "";
+	for (const byte of bytes) {
+		output += byte.toString(16).padStart(2, "0");
+	}
+	return output;
+}
+
+async function hashConfig(config: MCPServerConfig): Promise<string> {
+	const stable = stableStringify(config);
+	const digest = await crypto.subtle.digest("SHA-256", new TextEncoder().encode(stable));
+	return toHex(digest);
+}
+
+function cacheKey(serverName: string): string {
+	return `${CACHE_PREFIX}${serverName}`;
+}
+
+export class MCPToolCache {
+	constructor(private storage: AgentStorage) {}
+
+	async get(serverName: string, config: MCPServerConfig): Promise<MCPToolDefinition[] | null> {
+		const key = cacheKey(serverName);
+		const raw = this.storage.getCache(key);
+		if (!raw) return null;
+
+		let parsed: unknown;
+		try {
+			parsed = JSON.parse(raw);
+		} catch (error) {
+			logger.warn("MCP tool cache parse failed", { serverName, error: String(error) });
+			return null;
+		}
+
+		if (!isRecord(parsed)) return null;
+		if (parsed.version !== CACHE_VERSION) return null;
+		if (typeof parsed.configHash !== "string") return null;
+		if (!Array.isArray(parsed.tools)) return null;
+
+		let currentHash: string;
+		try {
+			currentHash = await hashConfig(config);
+		} catch (error) {
+			logger.warn("MCP tool cache hash failed", { serverName, error: String(error) });
+			return null;
+		}
+
+		if (parsed.configHash !== currentHash) return null;
+
+		return parsed.tools as MCPToolDefinition[];
+	}
+
+	async set(serverName: string, config: MCPServerConfig, tools: MCPToolDefinition[]): Promise<void> {
+		let configHash: string;
+		try {
+			configHash = await hashConfig(config);
+		} catch (error) {
+			logger.warn("MCP tool cache hash failed", { serverName, error: String(error) });
+			return;
+		}
+
+		const payload: MCPToolCachePayload = {
+			version: CACHE_VERSION,
+			configHash,
+			tools,
+		};
+
+		let serialized: string;
+		try {
+			serialized = JSON.stringify(payload);
+		} catch (error) {
+			logger.warn("MCP tool cache serialize failed", { serverName, error: String(error) });
+			return;
+		}
+
+		const expiresAtSec = Math.floor((Date.now() + CACHE_TTL_MS) / 1000);
+		this.storage.setCache(cacheKey(serverName), serialized, expiresAtSec);
+	}
+}