Simplify MCP arch #2

agent.py

@@ -18,7 +18,6 @@ from pathlib import Path
 try:
     from mcp.server import Server
     from mcp.types import Tool, TextContent, ImageContent, EmbeddedResource
-    from mcp.server.models import InitializationOptions
 except ImportError:
     print("Error: MCP SDK not installed. Install with: pip install mcp", file=sys.stderr)
     sys.exit(1)
@@ -282,47 +281,15 @@ async def main():
     
     try:
         async with stdio_server() as streams:
-            # Create initialization options
-            # The Server class will automatically provide its capabilities based on
-            # the registered @app.list_tools() and @app.call_tool() handlers
-            try:
-                # Try to get capabilities from the server if the method exists
-                if hasattr(app, 'get_capabilities'):
-                    try:
-                        # Try with NotificationOptions if available
-                        from mcp.server.lowlevel.server import NotificationOptions
-                        server_capabilities = app.get_capabilities(
-                            notification_options=NotificationOptions(),
-                            experimental_capabilities={}
-                        )
-                    except (ImportError, AttributeError, TypeError):
-                        # Fallback: try without NotificationOptions
-                        try:
-                            server_capabilities = app.get_capabilities()
-                        except Exception:
-                            # If get_capabilities doesn't work, create minimal capabilities
-                            server_capabilities = {}
-                else:
-                    # Server will provide capabilities automatically, use empty dict
-                    server_capabilities = {}
-            except Exception:
-                # Server will handle capabilities automatically
-                server_capabilities = {}
-            
-            # Create initialization options
-            init_options = InitializationOptions(
-                server_name="gemini-mcp-server",
-                server_version="1.0.0",
-                capabilities=server_capabilities
-            )
-            
-            # Run the server with initialization options
+            # The Server class automatically handles initialization and provides capabilities
+            # based on the registered @app.list_tools() and @app.call_tool() handlers
+            # No need to manually create InitializationOptions - the server handles this
             logger.info("MCP server ready")
             try:
+                # Run the server - it will automatically handle the initialization handshake
                 await app.run(
                     read_stream=streams[0],
-                    write_stream=streams[1],
-                    initialization_options=init_options
+                    write_stream=streams[1]
                 )
             except Exception as run_error:
                 logger.error(f"Error in app.run(): {run_error}")

app.py

@@ -264,9 +264,6 @@ async def get_mcp_session():
         stdio_ctx = stdio_client(server_params)
         read, write = await stdio_ctx.__aenter__()
         
-        # Wait for the server process to start and be ready
-        await asyncio.sleep(1.5)  # Wait for server startup
-        
         # Create ClientSession from the streams
         session = ClientSession(read, write)
         
@@ -278,12 +275,11 @@ async def get_mcp_session():
         # According to MCP protocol spec, the client MUST wait for the initialized notification
         # before sending any other requests (like list_tools)
         try:
+            # The __aenter__() method properly handles the full initialization sequence
+            # including waiting for the server's initialized notification
+            # This is a blocking call that completes only after the server sends initialized
             await session.__aenter__()
             logger.info("✅ MCP session initialized")
-            
-            # After __aenter__() completes, wait a bit longer to ensure server's internal state is ready
-            # The server may have sent the initialized notification but needs time to set up handlers
-            await asyncio.sleep(1.0)  # Give server time to finalize internal state
         except Exception as e:
             error_msg = str(e)
             error_type = type(e).__name__
@@ -322,11 +318,24 @@ async def call_agent(user_prompt: str, system_prompt: str = None, files: list =
             return ""
         
         # List tools - session should be ready after proper initialization
+        # Add a small delay to ensure server has fully processed initialization
+        await asyncio.sleep(0.1)
         try:
             tools = await session.list_tools()
         except Exception as e:
-            logger.error(f"❌ Failed to list MCP tools: {e}")
-            return ""
+            error_msg = str(e)
+            # Check if it's an initialization error
+            if "initialization" in error_msg.lower() or "before initialization" in error_msg.lower():
+                logger.warning(f"⚠️ Server not ready yet, waiting a bit more...: {error_msg}")
+                await asyncio.sleep(0.5)
+                try:
+                    tools = await session.list_tools()
+                except Exception as retry_error:
+                    logger.error(f"❌ Failed to list MCP tools after retry: {retry_error}")
+                    return ""
+            else:
+                logger.error(f"❌ Failed to list MCP tools: {error_msg}")
+                return ""
         
         if not tools or not hasattr(tools, 'tools'):
             logger.error("Invalid tools response from MCP server")
@@ -694,11 +703,24 @@ async def search_web_mcp_tool(query: str, max_results: int = 5) -> list:
             return []
         
         # List tools - session should be ready after proper initialization
+        # Add a small delay to ensure server has fully processed initialization
+        await asyncio.sleep(0.1)
         try:
             tools = await session.list_tools()
         except Exception as e:
-            logger.error(f"Failed to list MCP tools: {e}")
-            return []
+            error_msg = str(e)
+            # Check if it's an initialization error
+            if "initialization" in error_msg.lower() or "before initialization" in error_msg.lower():
+                logger.warning(f"⚠️ Server not ready yet, waiting a bit more...: {error_msg}")
+                await asyncio.sleep(0.5)
+                try:
+                    tools = await session.list_tools()
+                except Exception as retry_error:
+                    logger.error(f"Failed to list MCP tools after retry: {retry_error}")
+                    return []
+            else:
+                logger.error(f"Failed to list MCP tools: {error_msg}")
+                return []
         
         if not tools or not hasattr(tools, 'tools'):
             return []