refactor(mcp): enhance type definitions and improve error handling in MCP server and tools

Author: quanru

packages/mcp/src/web-tools.ts

@@ -3,7 +3,11 @@ import {
   MIDSCENE_MCP_USE_PUPPETEER_MODE,
   globalConfigManager,
 } from '@midscene/shared/env';
-import { BaseMidsceneTools, type ToolDefinition } from '@midscene/shared/mcp';
+import {
+  type BaseAgent,
+  BaseMidsceneTools,
+  type ToolDefinition,
+} from '@midscene/shared/mcp';
 import { AgentOverChromeBridge } from '@midscene/web/bridge-mode';
 import { type PuppeteerBrowserAgent, ensureBrowser } from './puppeteer';
 
@@ -13,7 +17,11 @@ export class WebMidsceneTools extends BaseMidsceneTools {
   );
 
   protected createTemporaryDevice() {
-    // Import PuppeteerWebPage class
+    // Import PuppeteerWebPage class using dynamic ESM import
+    // This is intentionally synchronous despite the async nature of createTemporaryDevice
+    // because we need the class constructor immediately for tool initialization
+    // The alternative would be to make all tool initialization async, which is a larger refactor
+    // eslint-disable-next-line @typescript-eslint/no-var-requires
     const { PuppeteerWebPage } = require('@midscene/web');
 
     // Create minimal mock page object that satisfies the interface
@@ -36,11 +44,13 @@ export class WebMidsceneTools extends BaseMidsceneTools {
     return new PuppeteerWebPage(mockPage as any, {});
   }
 
-  protected async ensureAgent(openNewTabWithUrl?: string): Promise<any> {
+  protected async ensureAgent(openNewTabWithUrl?: string): Promise<BaseAgent> {
     // Re-init if URL provided
     if (this.agent && openNewTabWithUrl) {
       try {
-        await this.agent.destroy();
+        if (this.agent.destroy) {
+          await this.agent.destroy();
+        }
       } catch (e) {
         console.debug('Failed to destroy agent during re-init:', e);
       }
@@ -58,9 +68,13 @@ export class WebMidsceneTools extends BaseMidsceneTools {
           'Bridge mode requires a URL. Use web_connect tool to connect to a page first.',
         );
       }
-      this.agent = await this.initAgentByBridgeMode(openNewTabWithUrl);
+      this.agent = (await this.initAgentByBridgeMode(
+        openNewTabWithUrl,
+      )) as unknown as BaseAgent;
     } else {
-      this.agent = await this.initPuppeteerAgent(openNewTabWithUrl);
+      this.agent = (await this.initPuppeteerAgent(
+        openNewTabWithUrl,
+      )) as unknown as BaseAgent;
     }
 
     return this.agent;
@@ -104,9 +118,21 @@ export class WebMidsceneTools extends BaseMidsceneTools {
         schema: {
           url: z.string().url().describe('URL to connect to'),
         },
-        handler: async ({ url }) => {
+        handler: async (args) => {
+          const { url } = args as { url: string };
           const agent = await this.ensureAgent(url);
-          const screenshot = await agent.page.screenshotBase64();
+          const screenshot = await agent.page?.screenshotBase64();
+          if (!screenshot) {
+            return {
+              content: [
+                {
+                  type: 'text',
+                  text: `Connected to: ${url}`,
+                },
+              ],
+            };
+          }
+
           const { parseBase64 } = await import('@midscene/shared/img');
           const { mimeType, body } = parseBase64(screenshot);
 
@@ -122,7 +148,6 @@ export class WebMidsceneTools extends BaseMidsceneTools {
                 mimeType,
               },
             ],
-            isError: false,
           };
         },
         autoDestroy: false,

packages/shared/src/mcp/base-server.ts

@@ -117,9 +117,27 @@ export abstract class BaseMCPServer {
     await this.initializeToolsManager();
 
     const transport = new StdioServerTransport();
-    await this.mcpServer.connect(transport);
 
+    try {
+      await this.mcpServer.connect(transport);
+    } catch (error: unknown) {
+      const message = error instanceof Error ? error.message : String(error);
+      console.error(`Failed to connect MCP stdio transport: ${message}`);
+      throw new Error(`Failed to initialize MCP stdio transport: ${message}`);
+    }
+
+    // Setup cleanup handlers
     process.stdin.on('close', () => this.performCleanup());
+
+    // Setup signal handlers for graceful shutdown
+    const cleanup = () => {
+      console.error(`${this.config.name} shutting down...`);
+      this.performCleanup();
+      process.exit(0);
+    };
+
+    process.once('SIGINT', cleanup);
+    process.once('SIGTERM', cleanup);
   }
 
   /**
@@ -149,6 +167,9 @@ export abstract class BaseMCPServer {
     const sessions = new Map<string, SessionData>();
 
     app.all('/mcp', async (req: Request, res: Response) => {
+      const startTime = Date.now();
+      const requestId = randomUUID().substring(0, 8);
+
       try {
         const rawSessionId = req.headers['mcp-session-id'];
         const sessionId = Array.isArray(rawSessionId)
@@ -159,26 +180,42 @@ export abstract class BaseMCPServer {
         if (!session && req.method === 'POST') {
           // Check session limit to prevent DoS
           if (sessions.size >= MAX_SESSIONS) {
+            console.error(
+              `[${new Date().toISOString()}] [${requestId}] Session limit reached: ${sessions.size}/${MAX_SESSIONS}`,
+            );
             res.status(503).json({
               error: 'Too many active sessions',
               message: 'Server is at maximum capacity. Please try again later.',
             });
             return;
           }
           session = await this.createHttpSession(sessions);
+          console.log(
+            `[${new Date().toISOString()}] [${requestId}] New session created: ${session.transport.sessionId}`,
+          );
         }
 
         if (session) {
           session.lastAccessedAt = new Date();
           await session.transport.handleRequest(req, res, req.body);
+          const duration = Date.now() - startTime;
+          console.log(
+            `[${new Date().toISOString()}] [${requestId}] Request completed in ${duration}ms`,
+          );
         } else {
+          console.error(
+            `[${new Date().toISOString()}] [${requestId}] Invalid session or GET without session`,
+          );
           res
             .status(400)
             .json({ error: 'Invalid session or GET without session' });
         }
       } catch (error: unknown) {
         const message = error instanceof Error ? error.message : String(error);
-        console.error('MCP request error:', message);
+        const duration = Date.now() - startTime;
+        console.error(
+          `[${new Date().toISOString()}] [${requestId}] MCP request error after ${duration}ms: ${message}`,
+        );
         if (!res.headersSent) {
           res.status(500).json({
             error: 'Internal server error',
@@ -200,15 +237,11 @@ export abstract class BaseMCPServer {
       .on('error', (error: NodeJS.ErrnoException) => {
         if (error.code === 'EADDRINUSE') {
           console.error(
-            `ERROR: Port ${options.port} is already in use.\n` +
-              `Please try a different port: --port=<number>\n` +
-              `Example: --mode=http --port=${options.port + 1}`,
+            `ERROR: Port ${options.port} is already in use.\nPlease try a different port: --port=<number>\nExample: --mode=http --port=${options.port + 1}`,
           );
         } else if (error.code === 'EACCES') {
           console.error(
-            `ERROR: Permission denied to bind to port ${options.port}.\n` +
-              `Ports below 1024 require root/admin privileges.\n` +
-              `Please use a port above 1024 or run with elevated privileges.`,
+            `ERROR: Permission denied to bind to port ${options.port}.\nPorts below 1024 require root/admin privileges.\nPlease use a port above 1024 or run with elevated privileges.`,
           );
         } else {
           console.error(
@@ -238,22 +271,28 @@ export abstract class BaseMCPServer {
           createdAt: new Date(),
           lastAccessedAt: new Date(),
         });
-        console.log(`Session ${sid} created`);
+        console.log(
+          `[${new Date().toISOString()}] Session ${sid} initialized (total: ${sessions.size})`,
+        );
       },
     });
 
     transport.onclose = () => {
       if (transport.sessionId) {
         sessions.delete(transport.sessionId);
-        console.log(`Session ${transport.sessionId} closed`);
+        console.log(
+          `[${new Date().toISOString()}] Session ${transport.sessionId} closed (remaining: ${sessions.size})`,
+        );
       }
     };
 
     try {
       await this.mcpServer.connect(transport);
     } catch (error: unknown) {
       const message = error instanceof Error ? error.message : String(error);
-      console.error(`Failed to connect MCP transport: ${message}`);
+      console.error(
+        `[${new Date().toISOString()}] Failed to connect MCP transport: ${message}`,
+      );
       // Clean up the failed transport
       if (transport.sessionId) {
         sessions.delete(transport.sessionId);
@@ -281,12 +320,14 @@ export abstract class BaseMCPServer {
           try {
             session.transport.close();
             sessions.delete(sid);
-            console.log(`Session ${sid} cleaned up due to inactivity`);
+            console.log(
+              `[${new Date().toISOString()}] Session ${sid} cleaned up due to inactivity (remaining: ${sessions.size})`,
+            );
           } catch (error: unknown) {
             const message =
               error instanceof Error ? error.message : String(error);
             console.error(
-              `Failed to close session ${sid} during cleanup: ${message}`,
+              `[${new Date().toISOString()}] Failed to close session ${sid} during cleanup: ${message}`,
             );
             // Still delete from map to prevent retry loops
             sessions.delete(sid);
@@ -313,21 +354,33 @@ export abstract class BaseMCPServer {
         try {
           session.transport.close();
         } catch (error: unknown) {
-          const message = error instanceof Error ? error.message : String(error);
+          const message =
+            error instanceof Error ? error.message : String(error);
           console.error(`Error closing session during shutdown: ${message}`);
         }
       }
       sessions.clear();
 
-      // Close HTTP server
+      // Close HTTP server gracefully
       try {
-        server.close();
+        server.close(() => {
+          // Server closed callback - all connections finished
+          this.performCleanup();
+          process.exit(0);
+        });
+
+        // Set a timeout in case server.close() hangs
+        setTimeout(() => {
+          console.error('Forcefully shutting down after timeout');
+          this.performCleanup();
+          process.exit(1);
+        }, 5000);
       } catch (error: unknown) {
         const message = error instanceof Error ? error.message : String(error);
         console.error(`Error closing HTTP server: ${message}`);
+        this.performCleanup();
+        process.exit(1);
       }
-
-      this.performCleanup();
     };
 
     // Use once() to prevent multiple registrations

packages/shared/src/mcp/base-tools.ts

@@ -4,19 +4,25 @@ import {
   generateCommonTools,
   generateToolsFromActionSpace,
 } from './tool-generator';
-import type { IMidsceneTools, ToolDefinition } from './types';
+import type {
+  ActionSpaceItem,
+  BaseAgent,
+  BaseDevice,
+  IMidsceneTools,
+  ToolDefinition,
+} from './types';
 
 const debug = getDebug('mcp:base-tools');
 
 export abstract class BaseMidsceneTools implements IMidsceneTools {
   protected mcpServer?: McpServer;
-  protected agent?: any;
+  protected agent?: BaseAgent;
   protected toolDefinitions: ToolDefinition[] = [];
 
   /**
    * Must be implemented by subclasses to create platform-specific agent
    */
-  protected abstract ensureAgent(initParam?: string): Promise<any>;
+  protected abstract ensureAgent(initParam?: string): Promise<BaseAgent>;
 
   /**
    * Optional: prepare platform-specific tools (e.g., device connection)
@@ -29,7 +35,7 @@ export abstract class BaseMidsceneTools implements IMidsceneTools {
    * Must be implemented by subclasses to create a temporary device instance
    * This allows getting real actionSpace without connecting to device
    */
-  protected abstract createTemporaryDevice(): any;
+  protected abstract createTemporaryDevice(): BaseDevice;
 
   /**
    * Initialize all tools by querying actionSpace
@@ -46,14 +52,14 @@ export abstract class BaseMidsceneTools implements IMidsceneTools {
     this.toolDefinitions.push(...platformTools);
 
     // 2. Try to get agent and its action space (two-layer fallback)
-    let actionSpace: any[];
+    let actionSpace: ActionSpaceItem[];
     try {
       // Layer 1: Try to use connected agent
       const agent = await this.ensureAgent();
       actionSpace = await agent.getActionSpace();
       debug(
         'Action space from connected agent:',
-        actionSpace.map((a: any) => a.name).join(', '),
+        actionSpace.map((a) => a.name).join(', '),
       );
     } catch (error) {
       // Layer 2: Create temporary device instance to read actionSpace
@@ -62,7 +68,7 @@ export abstract class BaseMidsceneTools implements IMidsceneTools {
       actionSpace = tempDevice.actionSpace();
       debug(
         'Action space from temporary device:',
-        actionSpace.map((a: any) => a.name).join(', '),
+        actionSpace.map((a) => a.name).join(', '),
       );
 
       // Destroy temporary instance using optional chaining
@@ -132,7 +138,9 @@ export abstract class BaseMidsceneTools implements IMidsceneTools {
       } finally {
         if (!process.env.MIDSCENE_MCP_DISABLE_AGENT_AUTO_DESTROY) {
           try {
-            await this.agent?.destroy();
+            if (this.agent?.destroy) {
+              await this.agent.destroy();
+            }
           } catch (e) {
             debug('Failed to destroy agent during cleanup:', e);
           }
@@ -146,6 +154,8 @@ export abstract class BaseMidsceneTools implements IMidsceneTools {
    * Cleanup method
    */
   public async closeBrowser(): Promise<void> {
-    await this.agent?.destroy();
+    if (this.agent?.destroy) {
+      await this.agent.destroy();
+    }
   }
 }