Update start_mcp.py

start_mcp.py

@@ -64,9 +64,11 @@ def run_sse_server(host: str, port: int):
                 
                 while True:
                     try:
-                        msg = await asyncio.wait_for(queue.get(), timeout=30)
+                        # Shorter timeout for more frequent keepalives
+                        msg = await asyncio.wait_for(queue.get(), timeout=15)
                         yield f"event: message\ndata: {json.dumps(msg)}\n\n"
                     except asyncio.TimeoutError:
+                        # Send keepalive more frequently
                         yield ": keepalive\n\n"
             except asyncio.CancelledError:
                 pass
@@ -98,71 +100,82 @@ def run_sse_server(host: str, port: int):
         if not mcp_available:
             return JSONResponse({"error": "MCP service not available"}, status_code=500)
         
-        try:
-            body = await request.json()
-            method = body.get("method", "")
-            params = body.get("params", {})
-            msg_id = body.get("id")
-            
-            logger.info(f"Method: {method}, ID: {msg_id}")
-            
-            # Process MCP methods
-            if method == "initialize":
-                result = {
-                    "protocolVersion": "2024-11-05",
-                    "serverInfo": {"name": "MaTableGPT-MCP", "version": "1.0.0"},
-                    "capabilities": {"tools": {}}
-                }
-            elif method == "notifications/initialized":
-                return JSONResponse({"ok": True})
-            elif method == "tools/list":
-                tools = []
-                for name, tool in mcp._tool_manager._tools.items():
-                    # Tool object has name, description, parameters attributes
-                    tools.append({
-                        "name": tool.name,
-                        "description": tool.description or name,
-                        "inputSchema": tool.parameters if hasattr(tool, 'parameters') else {"type": "object", "properties": {}}
-                    })
-                result = {"tools": tools}
-                logger.info(f"Listed {len(tools)} tools")
-            elif method == "tools/call":
-                tool_name = params.get("name")
-                tool_args = params.get("arguments", {})
-                
-                if tool_name not in mcp._tool_manager._tools:
-                    raise Exception(f"Unknown tool: {tool_name}")
-                
-                logger.info(f"Calling tool: {tool_name} with args: {list(tool_args.keys())}")
-                
-                # Get the Tool object and call its fn (the original function)
-                tool = mcp._tool_manager._tools[tool_name]
-                
-                # Call the underlying function directly
-                if tool.is_async:
-                    tool_result = await tool.fn(**tool_args)
+        body = await request.json()
+        method = body.get("method", "")
+        params = body.get("params", {})
+        msg_id = body.get("id")
+        
+        logger.info(f"Method: {method}, ID: {msg_id}")
+        
+        async def process_request():
+            """Process the request in background."""
+            try:
+                # Process MCP methods
+                if method == "initialize":
+                    result = {
+                        "protocolVersion": "2024-11-05",
+                        "serverInfo": {"name": "MaTableGPT-MCP", "version": "1.0.0"},
+                        "capabilities": {"tools": {}}
+                    }
+                elif method == "tools/list":
+                    tools = []
+                    for name, tool in mcp._tool_manager._tools.items():
+                        tools.append({
+                            "name": tool.name,
+                            "description": tool.description or name,
+                            "inputSchema": tool.parameters if hasattr(tool, 'parameters') else {"type": "object", "properties": {}}
+                        })
+                    result = {"tools": tools}
+                    logger.info(f"Listed {len(tools)} tools")
+                elif method == "tools/call":
+                    tool_name = params.get("name")
+                    tool_args = params.get("arguments", {})
+                    
+                    if tool_name not in mcp._tool_manager._tools:
+                        raise Exception(f"Unknown tool: {tool_name}")
+                    
+                    logger.info(f"Calling tool: {tool_name}")
+                    
+                    tool = mcp._tool_manager._tools[tool_name]
+                    
+                    # Run in executor if sync to avoid blocking
+                    if tool.is_async:
+                        tool_result = await tool.fn(**tool_args)
+                    else:
+                        loop = asyncio.get_event_loop()
+                        tool_result = await loop.run_in_executor(None, lambda: tool.fn(**tool_args))
+                    
+                    result = {"content": [{"type": "text", "text": json.dumps(tool_result)}]}
+                    logger.info(f"Tool {tool_name} completed")
                 else:
-                    tool_result = tool.fn(**tool_args)
+                    raise Exception(f"Unknown method: {method}")
                 
-                result = {"content": [{"type": "text", "text": json.dumps(tool_result)}]}
-                logger.info(f"Tool {tool_name} completed successfully")
-            else:
-                raise Exception(f"Unknown method: {method}")
-            
-            response = {"jsonrpc": "2.0", "id": msg_id, "result": result}
+                response = {"jsonrpc": "2.0", "id": msg_id, "result": result}
+                
+            except Exception as e:
+                logger.error(f"Error: {e}")
+                import traceback
+                traceback.print_exc()
+                response = {
+                    "jsonrpc": "2.0",
+                    "id": msg_id,
+                    "error": {"code": -32000, "message": str(e)}
+                }
             
-        except Exception as e:
-            logger.error(f"Error processing request: {e}")
-            import traceback
-            traceback.print_exc()
-            response = {
-                "jsonrpc": "2.0",
-                "id": body.get("id") if 'body' in dir() else None,
-                "error": {"code": -32000, "message": str(e)}
-            }
+            # Send response via SSE
+            if session_id in connections:
+                await connections[session_id].put(response)
+                logger.info(f"Response sent for {method}")
+            else:
+                logger.error(f"Session {session_id} disconnected before response")
+        
+        # Handle notifications immediately
+        if method == "notifications/initialized":
+            return JSONResponse({"ok": True})
+        
+        # Start background task for other requests
+        asyncio.create_task(process_request())
         
-        # Send via SSE
-        await connections[session_id].put(response)
         return JSONResponse({"ok": True})
     
     async def health(request: Request):