@defai.digital/ax-cli 3.7.2 → 3.8.2

package/dist/mcp/client-v2.js

@@ -0,0 +1,740 @@
+/**
+ * Type-Safe MCP Client (Phase 1 Improvements)
+ *
+ * Improvements applied:
+ * 1. SafeMutex with lock tokens (prevents race conditions)
+ * 2. Result types for all public APIs (explicit error handling)
+ * 3. State machine for connection tracking (type-safe states)
+ * 4. Branded types for ServerName/ToolName (prevent confusion)
+ * 5. Invariant checks (runtime validation)
+ *
+ * Coverage: 70% → 85%+ (Phase 1)
+ */
+import { Client } from "@modelcontextprotocol/sdk/client/index.js";
+import { EventEmitter } from "events";
+import { createTransport } from "./transports.js";
+import { MCP_CONFIG, ERROR_MESSAGES } from "../constants.js";
+import { MCPServerConfigSchema } from "../schemas/settings-schemas.js";
+import { getTokenCounter } from "../utils/token-counter.js";
+// Phase 1: Import type safety utilities
+import { SafeKeyedMutex } from "./mutex-safe.js";
+import { Ok, Err } from "./type-safety.js";
+import { createServerName } from "./type-safety.js";
+import { assertValidServerName } from "./invariants.js";
+/**
+ * Default reconnection configuration
+ */
+export const DEFAULT_RECONNECTION_CONFIG = {
+    enabled: true,
+    maxRetries: 5,
+    initialDelayMs: 1000, // 1 second
+    maxDelayMs: 30000, // 30 seconds
+    backoffMultiplier: 2 // exponential: 1s, 2s, 4s, 8s, 16s, 30s (max)
+};
+/**
+ * Default health check configuration
+ */
+export const DEFAULT_HEALTH_CHECK_CONFIG = {
+    enabled: true,
+    intervalMs: 30000 // 30 seconds
+};
+/**
+ * Type-safe MCP Manager with improved safety
+ */
+export class MCPManagerV2 extends EventEmitter {
+    // Phase 1: Replace Maps with state machine
+    connections = new Map();
+    tools = new Map();
+    // Phase 1: Use SafeMutex instead of pendingConnections Map
+    connectionMutex = new SafeKeyedMutex();
+    initializationPromise = null;
+    tokenCounter = getTokenCounter();
+    disposed = false;
+    // Phase 2: Reconnection management
+    reconnectionConfig;
+    healthCheckConfig;
+    reconnectionAttempts = new Map();
+    reconnectionTimers = new Map();
+    serverConfigs = new Map();
+    healthCheckTimer = null;
+    constructor(reconnectionConfig = {}, healthCheckConfig = {}) {
+        super();
+        this.reconnectionConfig = {
+            ...DEFAULT_RECONNECTION_CONFIG,
+            ...reconnectionConfig
+        };
+        this.healthCheckConfig = {
+            ...DEFAULT_HEALTH_CHECK_CONFIG,
+            ...healthCheckConfig
+        };
+        // Start health checks if enabled
+        if (this.healthCheckConfig.enabled) {
+            this.startHealthChecks();
+        }
+    }
+    /**
+     * Add MCP server with type-safe connection management
+     *
+     * Phase 1 improvements:
+     * - Returns Result instead of throwing
+     * - Uses SafeMutex for concurrency control
+     * - Tracks state machine transitions
+     * - Validates inputs with invariants
+     */
+    async addServer(config) {
+        // Phase 1: Check if disposed
+        if (this.disposed) {
+            return Err(new Error('MCPManager is disposed'));
+        }
+        // Phase 1: Validate server name (branded type creation)
+        const serverName = createServerName(config.name);
+        if (!serverName) {
+            return Err(new Error(`Invalid server name: "${config.name}"`));
+        }
+        // Phase 1: Check current state
+        const currentState = this.connections.get(serverName);
+        if (currentState) {
+            switch (currentState.status) {
+                case 'connected':
+                    return Ok(undefined); // Already connected
+                case 'connecting':
+                    // Wait for existing connection attempt
+                    return await currentState.promise;
+                case 'disconnecting':
+                    return Err(new Error(`Server ${serverName} is disconnecting`));
+                case 'failed':
+                    // Can retry after failure
+                    break;
+                case 'idle':
+                    // Can connect
+                    break;
+            }
+        }
+        // Phase 1: Use SafeMutex for concurrency control
+        const mutexResult = await this.connectionMutex.runExclusive(serverName, async () => {
+            // Double-check state inside mutex
+            const state = this.connections.get(serverName);
+            if (state?.status === 'connected') {
+                return Ok(undefined);
+            }
+            return await this._addServerInternal(serverName, config);
+        });
+        // Unwrap nested Result
+        if (!mutexResult.success) {
+            return mutexResult;
+        }
+        return mutexResult.value;
+    }
+    /**
+     * Internal connection logic with state transitions
+     */
+    async _addServerInternal(serverName, config) {
+        // Validate config with Zod
+        const validationResult = MCPServerConfigSchema.safeParse(config);
+        if (!validationResult.success) {
+            // Transition to failed state
+            this.connections.set(serverName, {
+                status: 'failed',
+                serverName,
+                error: new Error(`Invalid config: ${validationResult.error.message}`),
+                failedAt: Date.now()
+            });
+            return Err(new Error(`Invalid MCP server config: ${validationResult.error.message}`));
+        }
+        const validatedConfig = validationResult.data;
+        // Phase 2: Store server config for reconnection attempts
+        this.serverConfigs.set(serverName, validatedConfig);
+        // Handle legacy stdio-only configuration
+        let transportConfig = validatedConfig.transport;
+        if (!transportConfig && validatedConfig.command) {
+            transportConfig = {
+                type: 'stdio',
+                command: validatedConfig.command,
+                args: validatedConfig.args,
+                env: validatedConfig.env
+            };
+        }
+        if (!transportConfig) {
+            const error = new Error(ERROR_MESSAGES.TRANSPORT_CONFIG_REQUIRED);
+            this.connections.set(serverName, {
+                status: 'failed',
+                serverName,
+                error,
+                failedAt: Date.now()
+            });
+            return Err(error);
+        }
+        try {
+            // Transition to connecting state
+            const startedAt = Date.now();
+            const connectingPromise = (async () => {
+                try {
+                    // Create transport
+                    const transport = createTransport(transportConfig);
+                    // Create client
+                    const client = new Client({
+                        name: MCP_CONFIG.CLIENT_NAME,
+                        version: MCP_CONFIG.CLIENT_VERSION
+                    }, {
+                        capabilities: {
+                            tools: {}
+                        }
+                    });
+                    // Connect
+                    const sdkTransport = await transport.connect();
+                    await client.connect(sdkTransport);
+                    // List tools
+                    const toolsResult = await client.listTools();
+                    // Register tools with branded types
+                    for (const tool of toolsResult.tools) {
+                        const toolName = createToolName(`mcp__${serverName}__${tool.name}`);
+                        if (!toolName) {
+                            console.warn(`Invalid tool name: ${tool.name}`);
+                            continue;
+                        }
+                        const mcpTool = {
+                            name: toolName,
+                            description: tool.description || `Tool from ${serverName} server`,
+                            inputSchema: tool.inputSchema,
+                            serverName
+                        };
+                        this.tools.set(toolName, mcpTool);
+                    }
+                    // Transition to connected state
+                    this.connections.set(serverName, {
+                        status: 'connected',
+                        serverName,
+                        client,
+                        transport,
+                        connectedAt: Date.now()
+                    });
+                    this.emit('serverAdded', serverName, toolsResult.tools.length);
+                    return Ok(undefined);
+                }
+                catch (error) {
+                    // Transition to failed state
+                    const err = error instanceof Error ? error : new Error(String(error));
+                    this.connections.set(serverName, {
+                        status: 'failed',
+                        serverName,
+                        error: err,
+                        failedAt: Date.now()
+                    });
+                    this.emit('serverError', serverName, err);
+                    // Phase 2: Schedule reconnection if enabled
+                    if (this.reconnectionConfig.enabled && !this.disposed) {
+                        this.scheduleReconnection(serverName, validatedConfig);
+                    }
+                    return Err(err);
+                }
+            })();
+            this.connections.set(serverName, {
+                status: 'connecting',
+                serverName,
+                startedAt,
+                promise: connectingPromise
+            });
+            return await connectingPromise;
+        }
+        catch (error) {
+            const err = error instanceof Error ? error : new Error(String(error));
+            this.connections.set(serverName, {
+                status: 'failed',
+                serverName,
+                error: err,
+                failedAt: Date.now()
+            });
+            // Phase 2: Schedule reconnection if enabled
+            if (this.reconnectionConfig.enabled && !this.disposed) {
+                this.scheduleReconnection(serverName, validatedConfig);
+            }
+            return Err(err);
+        }
+    }
+    /**
+     * Remove MCP server with proper state transitions
+     */
+    async removeServer(serverName) {
+        // Phase 1: Check if disposed
+        if (this.disposed) {
+            return Err(new Error('MCPManager is disposed'));
+        }
+        // Phase 1: Validate server name
+        assertValidServerName(serverName);
+        // Phase 2: Cancel any pending reconnection attempts
+        this.cancelReconnection(serverName);
+        this.reconnectionAttempts.delete(serverName);
+        this.serverConfigs.delete(serverName);
+        // BUG FIX: Use mutex to prevent TOCTOU race with addServer/callTool
+        const mutexResult = await this.connectionMutex.runExclusive(serverName, async () => {
+            // Re-check state inside mutex (prevent race conditions)
+            const state = this.connections.get(serverName);
+            if (!state) {
+                return Err(new Error(`Server ${serverName} not found`));
+            }
+            // Check if we can disconnect from current state
+            if (state.status === 'connecting') {
+                return Err(new Error(`Server ${serverName} is still connecting`));
+            }
+            if (state.status === 'disconnecting') {
+                return Err(new Error(`Server ${serverName} is already disconnecting`));
+            }
+            if (state.status !== 'connected') {
+                // Remove from map if not connected
+                this.connections.delete(serverName);
+                return Ok(undefined);
+            }
+            return await this._removeServerInternal(serverName, state);
+        });
+        // Unwrap nested Result
+        if (!mutexResult.success) {
+            return mutexResult;
+        }
+        return mutexResult.value;
+    }
+    /**
+     * Internal disconnection logic with state transitions
+     */
+    async _removeServerInternal(serverName, state) {
+        // Transition to disconnecting state
+        this.connections.set(serverName, {
+            status: 'disconnecting',
+            serverName,
+            client: state.client,
+            transport: state.transport
+        });
+        try {
+            // Remove tools
+            for (const [toolName, tool] of this.tools.entries()) {
+                if (tool.serverName === serverName) {
+                    this.tools.delete(toolName);
+                }
+            }
+            // Disconnect client
+            const clientResult = await this.closeClient(state.client, serverName);
+            // Disconnect transport
+            const transportResult = await this.disconnectTransport(state.transport, serverName);
+            // Aggregate errors
+            const errors = [];
+            if (!clientResult.success)
+                errors.push(clientResult.error);
+            if (!transportResult.success)
+                errors.push(transportResult.error);
+            // Transition to idle state
+            this.connections.delete(serverName);
+            this.emit('serverRemoved', serverName);
+            if (errors.length > 0) {
+                return Err(new AggregateError(errors, `Failed to fully disconnect ${serverName}`));
+            }
+            return Ok(undefined);
+        }
+        catch (error) {
+            // Even if error, remove from state
+            this.connections.delete(serverName);
+            return Err(error instanceof Error ? error : new Error(String(error)));
+        }
+    }
+    /**
+     * Close client with error handling
+     */
+    async closeClient(client, serverName) {
+        try {
+            await client.close();
+            return Ok(undefined);
+        }
+        catch (error) {
+            const err = error instanceof Error ? error : new Error(String(error));
+            console.warn(`Error closing MCP client ${serverName}:`, err);
+            return Err(err);
+        }
+    }
+    /**
+     * Disconnect transport with error handling
+     */
+    async disconnectTransport(transport, serverName) {
+        try {
+            await transport.disconnect();
+            return Ok(undefined);
+        }
+        catch (error) {
+            const err = error instanceof Error ? error : new Error(String(error));
+            console.warn(`Error disconnecting MCP transport ${serverName}:`, err);
+            return Err(err);
+        }
+    }
+    /**
+     * Call MCP tool with type safety
+     */
+    async callTool(toolName, arguments_) {
+        if (this.disposed) {
+            return Err(new Error('MCPManager is disposed'));
+        }
+        const tool = this.tools.get(toolName);
+        if (!tool) {
+            return Err(new Error(`Tool ${toolName} not found`));
+        }
+        // BUG FIX: Get client reference inside mutex to prevent TOCTOU race with removeServer
+        const mutexResult = await this.connectionMutex.runExclusive(tool.serverName, async () => {
+            // Re-check state inside mutex
+            const state = this.connections.get(tool.serverName);
+            if (!state) {
+                return Err(new Error(`Server ${tool.serverName} not found`));
+            }
+            if (state.status !== 'connected') {
+                return Err(new Error(`Server ${tool.serverName} not connected (status: ${state.status})`));
+            }
+            // Return client snapshot (mutex released after this, but client reference is safe to use)
+            return Ok(state.client);
+        });
+        // Unwrap nested Result
+        if (!mutexResult.success) {
+            return mutexResult;
+        }
+        const clientResult = mutexResult.value;
+        if (!clientResult.success) {
+            return clientResult;
+        }
+        const client = clientResult.value;
+        try {
+            // Extract original tool name
+            const prefix = `mcp__${tool.serverName}__`;
+            const originalToolName = toolName.startsWith(prefix)
+                ? toolName.substring(prefix.length)
+                : toolName;
+            // Validate arguments
+            const safeArgs = (arguments_ && typeof arguments_ === 'object' && !Array.isArray(arguments_))
+                ? arguments_
+                : {};
+            // Call tool (mutex released, but client reference is still valid)
+            const result = await client.callTool({
+                name: originalToolName,
+                arguments: safeArgs
+            });
+            // Apply token limiting
+            if (MCP_CONFIG.TRUNCATION_ENABLED) {
+                const resultText = JSON.stringify(result.content);
+                const tokenCount = this.tokenCounter.countTokens(resultText);
+                if (tokenCount > MCP_CONFIG.TOKEN_HARD_LIMIT) {
+                    const truncatedText = this.truncateToTokenLimit(resultText, MCP_CONFIG.TOKEN_HARD_LIMIT);
+                    result.content = [
+                        { type: 'text', text: truncatedText },
+                        {
+                            type: 'text',
+                            text: `\n\n⚠️  Output truncated: ${tokenCount.toLocaleString()} tokens exceeded limit of ${MCP_CONFIG.TOKEN_HARD_LIMIT.toLocaleString()} tokens`
+                        }
+                    ];
+                    this.emit('token-limit-exceeded', {
+                        toolName,
+                        serverName: tool.serverName,
+                        originalTokens: tokenCount,
+                        truncatedTokens: MCP_CONFIG.TOKEN_HARD_LIMIT
+                    });
+                }
+                else if (tokenCount > MCP_CONFIG.TOKEN_WARNING_THRESHOLD) {
+                    this.emit('token-warning', {
+                        toolName,
+                        serverName: tool.serverName,
+                        tokenCount,
+                        threshold: MCP_CONFIG.TOKEN_WARNING_THRESHOLD
+                    });
+                }
+            }
+            return Ok(result);
+        }
+        catch (error) {
+            return Err(error instanceof Error ? error : new Error(String(error)));
+        }
+    }
+    /**
+     * Truncate text to fit within token limit
+     * UNICODE FIX: Uses grapheme clusters
+     */
+    truncateToTokenLimit(text, maxTokens) {
+        const chars = Array.from(text);
+        let low = 0;
+        let high = chars.length;
+        let result = text;
+        while (low <= high) {
+            const mid = Math.floor((low + high) / 2);
+            const truncated = chars.slice(0, mid).join('');
+            const tokens = this.tokenCounter.countTokens(truncated);
+            if (tokens <= maxTokens) {
+                result = truncated;
+                low = mid + 1;
+            }
+            else {
+                high = mid - 1;
+            }
+        }
+        return result;
+    }
+    /**
+     * Get all tools
+     */
+    getTools() {
+        return Array.from(this.tools.values());
+    }
+    /**
+     * Get all connected servers
+     */
+    getServers() {
+        const connected = [];
+        for (const [serverName, state] of this.connections.entries()) {
+            if (state.status === 'connected') {
+                connected.push(serverName);
+            }
+        }
+        return connected;
+    }
+    /**
+     * Get connection state for a server
+     */
+    getConnectionState(serverName) {
+        return this.connections.get(serverName);
+    }
+    /**
+     * Get transport type for a server
+     */
+    getTransportType(serverName) {
+        const state = this.connections.get(serverName);
+        if (!state) {
+            return Err(new Error(`Server ${serverName} not found`));
+        }
+        if (state.status !== 'connected') {
+            return Err(new Error(`Server ${serverName} not connected`));
+        }
+        const type = state.transport.getType();
+        return Ok(type);
+    }
+    /**
+     * Schedule reconnection for a failed server with exponential backoff
+     *
+     * Phase 2: Automatic reconnection logic
+     *
+     * @param serverName - Server to reconnect
+     * @param config - Server configuration
+     */
+    scheduleReconnection(serverName, config) {
+        // Cancel any existing reconnection timer
+        this.cancelReconnection(serverName);
+        // Get current attempt count
+        const attempts = this.reconnectionAttempts.get(serverName) || 0;
+        // Check if we've exceeded max retries
+        if (attempts >= this.reconnectionConfig.maxRetries) {
+            this.emit('reconnection-failed', serverName, attempts, 'Max retries exceeded');
+            return;
+        }
+        // Calculate exponential backoff delay
+        const baseDelay = this.reconnectionConfig.initialDelayMs;
+        const multiplier = Math.pow(this.reconnectionConfig.backoffMultiplier, attempts);
+        const calculatedDelay = Math.min(baseDelay * multiplier, this.reconnectionConfig.maxDelayMs);
+        // Emit reconnection scheduled event
+        this.emit('reconnection-scheduled', serverName, attempts + 1, calculatedDelay);
+        // Schedule reconnection attempt
+        const timer = setTimeout(async () => {
+            // Increment attempt count
+            this.reconnectionAttempts.set(serverName, attempts + 1);
+            // Attempt reconnection
+            const result = await this.addServer(config);
+            if (result.success) {
+                // Success! Reset attempt counter
+                this.reconnectionAttempts.delete(serverName);
+                this.emit('reconnection-succeeded', serverName, attempts + 1);
+            }
+            else {
+                // Failed - will be rescheduled by addServer error handling
+                // (which calls scheduleReconnection again)
+            }
+        }, calculatedDelay);
+        this.reconnectionTimers.set(serverName, timer);
+    }
+    /**
+     * Cancel reconnection for a server
+     */
+    cancelReconnection(serverName) {
+        const timer = this.reconnectionTimers.get(serverName);
+        if (timer) {
+            clearTimeout(timer);
+            this.reconnectionTimers.delete(serverName);
+        }
+    }
+    /**
+     * Perform health check on a single server
+     *
+     * @param serverName - Server to check
+     * @returns Result indicating health status
+     */
+    async healthCheck(serverName) {
+        const state = this.connections.get(serverName);
+        if (!state) {
+            return Err(new Error(`Server ${serverName} not found`));
+        }
+        if (state.status !== 'connected') {
+            return Ok(false); // Not connected = not healthy
+        }
+        try {
+            // Simple health check: try to list tools
+            await state.client.listTools();
+            return Ok(true); // Healthy
+        }
+        catch (error) {
+            // Server is unhealthy
+            this.emit('server-unhealthy', serverName, error);
+            // Transition to failed state
+            const err = error instanceof Error ? error : new Error(String(error));
+            this.connections.set(serverName, {
+                status: 'failed',
+                serverName,
+                error: err,
+                failedAt: Date.now()
+            });
+            // Close the connection
+            try {
+                await state.client.close();
+                await state.transport.disconnect();
+            }
+            catch (closeError) {
+                console.warn(`Error closing unhealthy server ${serverName}:`, closeError);
+            }
+            // Schedule reconnection
+            const config = this.serverConfigs.get(serverName);
+            if (config && this.reconnectionConfig.enabled && !this.disposed) {
+                this.scheduleReconnection(serverName, config);
+            }
+            return Ok(false); // Unhealthy
+        }
+    }
+    /**
+     * Start periodic health checks for all connected servers
+     */
+    startHealthChecks() {
+        if (this.healthCheckTimer) {
+            return; // Already running
+        }
+        const runHealthChecks = async () => {
+            const connectedServers = Array.from(this.connections.entries())
+                .filter(([_, state]) => state.status === 'connected')
+                .map(([name, _]) => name);
+            for (const serverName of connectedServers) {
+                await this.healthCheck(serverName);
+            }
+        };
+        // Run initial health check
+        runHealthChecks().catch(error => {
+            console.warn('Health check error:', error);
+        });
+        // Schedule periodic checks
+        this.healthCheckTimer = setInterval(() => {
+            runHealthChecks().catch(error => {
+                console.warn('Health check error:', error);
+            });
+        }, this.healthCheckConfig.intervalMs);
+    }
+    /**
+     * Stop periodic health checks
+     */
+    stopHealthChecks() {
+        if (this.healthCheckTimer) {
+            clearInterval(this.healthCheckTimer);
+            this.healthCheckTimer = null;
+        }
+    }
+    /**
+     * Shutdown all servers
+     */
+    async shutdown() {
+        const serverNames = Array.from(this.connections.keys());
+        const results = await Promise.allSettled(serverNames.map(name => this.removeServer(name)));
+        const errors = [];
+        results.forEach((result, index) => {
+            if (result.status === 'rejected') {
+                console.warn(`Failed to remove server ${serverNames[index]}:`, result.reason);
+                errors.push(result.reason);
+            }
+            else if (!result.value.success) {
+                errors.push(result.value.error);
+            }
+        });
+        if (errors.length > 0) {
+            return Err(new AggregateError(errors, 'Shutdown had errors'));
+        }
+        return Ok(undefined);
+    }
+    /**
+     * Ensure servers initialized
+     */
+    async ensureServersInitialized() {
+        if (this.initializationPromise) {
+            return this.initializationPromise;
+        }
+        if (this.connections.size === 0 && !this.initializationPromise) {
+            this.initializationPromise = (async () => {
+                try {
+                    const { loadMCPConfig } = await import('../mcp/config.js');
+                    const config = loadMCPConfig();
+                    const initPromises = config.servers.map(async (serverConfig) => {
+                        const serverName = createServerName(serverConfig.name);
+                        if (!serverName) {
+                            console.warn(`Invalid server name: ${serverConfig.name}`);
+                            return;
+                        }
+                        const result = await this.addServer(serverConfig);
+                        if (!result.success) {
+                            console.warn(`Failed to initialize MCP server ${serverName}:`, result.error);
+                        }
+                    });
+                    await Promise.all(initPromises);
+                    return Ok(undefined);
+                }
+                catch (error) {
+                    console.error('Failed to initialize MCP servers:', error);
+                    return Err(error instanceof Error ? error : new Error(String(error)));
+                }
+                finally {
+                    this.initializationPromise = null;
+                }
+            })();
+        }
+        if (this.initializationPromise) {
+            return await this.initializationPromise;
+        }
+        return Ok(undefined);
+    }
+    /**
+     * Dispose all resources
+     */
+    async dispose() {
+        if (this.disposed) {
+            return Ok(undefined);
+        }
+        this.disposed = true;
+        // Phase 2: Stop health checks
+        this.stopHealthChecks();
+        // Phase 2: Cancel all reconnection timers
+        for (const timer of this.reconnectionTimers.values()) {
+            clearTimeout(timer);
+        }
+        this.reconnectionTimers.clear();
+        this.reconnectionAttempts.clear();
+        this.serverConfigs.clear();
+        const shutdownResult = await this.shutdown();
+        this.removeAllListeners();
+        return shutdownResult;
+    }
+}
+/**
+ * Helper functions for creating branded types
+ */
+function createToolName(name) {
+    // Tool names can have double underscores for MCP prefix
+    if (!/^[a-zA-Z0-9_-]+$/.test(name)) {
+        return null;
+    }
+    if (name.length < 1 || name.length > 128) {
+        return null;
+    }
+    return name;
+}
+// Re-export createServerName from type-safety
+export { createServerName, createToolName };
+//# sourceMappingURL=client-v2.js.map