[router] Add MCP Tool Handler (#9615)

sgl-router/tests/common/mock_mcp_server.rs

@@ -0,0 +1,237 @@
+// tests/common/mock_mcp_server.rs - Mock MCP server for testing
+
+use axum::{
+    extract::Json, http::StatusCode, response::Json as ResponseJson, routing::post, Router,
+};
+use serde_json::{json, Value};
+use tokio::net::TcpListener;
+
+/// Mock MCP server that returns hardcoded responses for testing
+pub struct MockMCPServer {
+    pub port: u16,
+    pub server_handle: Option<tokio::task::JoinHandle<()>>,
+}
+
+impl MockMCPServer {
+    /// Start a mock MCP server on an available port
+    pub async fn start() -> Result<Self, Box<dyn std::error::Error + Send + Sync>> {
+        // Find an available port
+        let listener = TcpListener::bind("127.0.0.1:0").await?;
+        let port = listener.local_addr()?.port();
+
+        let app = Router::new().route("/mcp", post(handle_mcp_request));
+
+        let server_handle = tokio::spawn(async move {
+            axum::serve(listener, app)
+                .await
+                .expect("Mock MCP server failed to start");
+        });
+
+        // Give the server a moment to start
+        tokio::time::sleep(tokio::time::Duration::from_millis(100)).await;
+
+        Ok(MockMCPServer {
+            port,
+            server_handle: Some(server_handle),
+        })
+    }
+
+    /// Get the full URL for this mock server
+    pub fn url(&self) -> String {
+        format!("http://127.0.0.1:{}/mcp", self.port)
+    }
+
+    /// Stop the mock server
+    pub async fn stop(&mut self) {
+        if let Some(handle) = self.server_handle.take() {
+            handle.abort();
+            // Wait a moment for cleanup
+            tokio::time::sleep(tokio::time::Duration::from_millis(50)).await;
+        }
+    }
+}
+
+impl Drop for MockMCPServer {
+    fn drop(&mut self) {
+        if let Some(handle) = self.server_handle.take() {
+            handle.abort();
+        }
+    }
+}
+
+/// Handle MCP requests and return mock responses
+async fn handle_mcp_request(Json(request): Json<Value>) -> Result<ResponseJson<Value>, StatusCode> {
+    // Parse the JSON-RPC request
+    let method = request.get("method").and_then(|m| m.as_str()).unwrap_or("");
+
+    let id = request
+        .get("id")
+        .and_then(|i| i.as_str())
+        .unwrap_or("unknown");
+
+    let response = match method {
+        "initialize" => {
+            // Mock initialize response
+            json!({
+                "jsonrpc": "2.0",
+                "id": id,
+                "result": {
+                    "serverInfo": {
+                        "name": "Mock MCP Server",
+                        "version": "1.0.0"
+                    },
+                    "instructions": "Mock server for testing"
+                }
+            })
+        }
+        "tools/list" => {
+            // Mock tools list response
+            json!({
+                "jsonrpc": "2.0",
+                "id": id,
+                "result": {
+                    "tools": [
+                        {
+                            "name": "brave_web_search",
+                            "description": "Mock web search tool",
+                            "inputSchema": {
+                                "type": "object",
+                                "properties": {
+                                    "query": {"type": "string"},
+                                    "count": {"type": "integer"}
+                                },
+                                "required": ["query"]
+                            }
+                        },
+                        {
+                            "name": "brave_local_search",
+                            "description": "Mock local search tool",
+                            "inputSchema": {
+                                "type": "object",
+                                "properties": {
+                                    "query": {"type": "string"}
+                                },
+                                "required": ["query"]
+                            }
+                        }
+                    ]
+                }
+            })
+        }
+        "tools/call" => {
+            // Mock tool call response
+            let empty_json = json!({});
+            let params = request.get("params").unwrap_or(&empty_json);
+            let tool_name = params.get("name").and_then(|n| n.as_str()).unwrap_or("");
+            let empty_args = json!({});
+            let arguments = params.get("arguments").unwrap_or(&empty_args);
+
+            match tool_name {
+                "brave_web_search" => {
+                    let query = arguments
+                        .get("query")
+                        .and_then(|q| q.as_str())
+                        .unwrap_or("test");
+                    json!({
+                        "jsonrpc": "2.0",
+                        "id": id,
+                        "result": {
+                            "content": [
+                                {
+                                    "type": "text",
+                                    "text": format!("Mock search results for: {}", query)
+                                }
+                            ],
+                            "isError": false
+                        }
+                    })
+                }
+                "brave_local_search" => {
+                    json!({
+                        "jsonrpc": "2.0",
+                        "id": id,
+                        "result": {
+                            "content": [
+                                {
+                                    "type": "text",
+                                    "text": "Mock local search results"
+                                }
+                            ],
+                            "isError": false
+                        }
+                    })
+                }
+                _ => {
+                    // Unknown tool
+                    json!({
+                        "jsonrpc": "2.0",
+                        "id": id,
+                        "error": {
+                            "code": -1,
+                            "message": format!("Unknown tool: {}", tool_name)
+                        }
+                    })
+                }
+            }
+        }
+        _ => {
+            // Unknown method
+            json!({
+                "jsonrpc": "2.0",
+                "id": id,
+                "error": {
+                    "code": -32601,
+                    "message": format!("Method not found: {}", method)
+                }
+            })
+        }
+    };
+
+    Ok(ResponseJson(response))
+}
+
+#[cfg(test)]
+#[allow(unused_imports)]
+mod tests {
+    use super::MockMCPServer;
+    use serde_json::{json, Value};
+
+    #[tokio::test]
+    async fn test_mock_server_startup() {
+        let mut server = MockMCPServer::start().await.unwrap();
+        assert!(server.port > 0);
+        assert!(server.url().contains(&server.port.to_string()));
+        server.stop().await;
+    }
+
+    #[tokio::test]
+    async fn test_mock_server_responses() {
+        let mut server = MockMCPServer::start().await.unwrap();
+        let client = reqwest::Client::new();
+
+        // Test initialize
+        let init_request = json!({
+            "jsonrpc": "2.0",
+            "id": "1",
+            "method": "initialize",
+            "params": {
+                "protocolVersion": "2024-11-05",
+                "capabilities": {}
+            }
+        });
+
+        let response = client
+            .post(server.url())
+            .json(&init_request)
+            .send()
+            .await
+            .unwrap();
+
+        assert!(response.status().is_success());
+        let json: Value = response.json().await.unwrap();
+        assert_eq!(json["jsonrpc"], "2.0");
+        assert_eq!(json["result"]["serverInfo"]["name"], "Mock MCP Server");
+
+        server.stop().await;
+    }
+}