From 68e23658f6213b43c504bc27c811caffc9a6b766 Mon Sep 17 00:00:00 2001 From: Andrew Barnes Date: Thu, 5 Mar 2026 12:33:43 -0500 Subject: [PATCH 1/2] fix: replace unwrap() with proper error handling in MCP server Replace panic-prone unwrap() calls with safe alternatives: - Use unwrap_or for clap arg with default value - Use if-let pattern for JSON-RPC id extraction - Use match for tools_cache access - Return validation error instead of unwrap on parts.last() - Break server loop on stdout pipe errors instead of silently continuing to process requests that can never be delivered Also adds unit test coverage for build_mcp_cli, walk_resources, and handle_request. --- .changeset/fix-mcp-server-unwrap.md | 12 ++ src/mcp_server.rs | 173 ++++++++++++++++++++++++++-- 2 files changed, 175 insertions(+), 10 deletions(-) create mode 100644 .changeset/fix-mcp-server-unwrap.md diff --git a/.changeset/fix-mcp-server-unwrap.md b/.changeset/fix-mcp-server-unwrap.md new file mode 100644 index 00000000..fe8b091d --- /dev/null +++ b/.changeset/fix-mcp-server-unwrap.md @@ -0,0 +1,12 @@ +--- +"@anthropic/gws": patch +--- + +Replace unwrap() calls with proper error handling in MCP server + +- Replace `.unwrap()` on `get_one("services")` with `.unwrap_or("")` +- Replace `req.get("id").unwrap()` with `if let Some(id)` pattern +- Replace `tools_cache.as_ref().unwrap()` with match expression +- Replace `parts.last().unwrap()` with `.ok_or_else()` returning a proper validation error +- Handle broken stdout pipe by breaking the server loop instead of silently continuing +- Add unit tests for `build_mcp_cli`, `walk_resources`, and `handle_request` diff --git a/src/mcp_server.rs b/src/mcp_server.rs index 331e70d8..3ff56656 100644 --- a/src/mcp_server.rs +++ b/src/mcp_server.rs @@ -84,7 +84,10 @@ pub async fn start(args: &[String]) -> Result<(), GwsError> { tool_mode, }; - let svc_str = matches.get_one::("services").unwrap(); + let svc_str = matches + .get_one::("services") + .map(|s| s.as_str()) + .unwrap_or(""); if !svc_str.is_empty() { if svc_str == "all" { config.services = services::SERVICES @@ -121,14 +124,12 @@ pub async fn start(args: &[String]) -> Result<(), GwsError> { match serde_json::from_str::(&line) { Ok(req) => { - let is_notification = req.get("id").is_none(); let method = req.get("method").and_then(|m| m.as_str()).unwrap_or(""); let params = req.get("params").cloned().unwrap_or_else(|| json!({})); let result = handle_request(method, ¶ms, &config, &mut tools_cache).await; - if !is_notification { - let id = req.get("id").unwrap(); + if let Some(id) = req.get("id") { let response = match result { Ok(res) => json!({ "jsonrpc": "2.0", @@ -153,8 +154,12 @@ pub async fn start(args: &[String]) -> Result<(), GwsError> { } }; out.push('\n'); - let _ = stdout.write_all(out.as_bytes()).await; - let _ = stdout.flush().await; + if stdout.write_all(out.as_bytes()).await.is_err() + || stdout.flush().await.is_err() + { + eprintln!("[gws mcp] Stdout pipe broken, shutting down."); + break; + } } } Err(_) => { @@ -174,8 +179,12 @@ pub async fn start(args: &[String]) -> Result<(), GwsError> { } }; out.push('\n'); - let _ = stdout.write_all(out.as_bytes()).await; - let _ = stdout.flush().await; + if stdout.write_all(out.as_bytes()).await.is_err() + || stdout.flush().await.is_err() + { + eprintln!("[gws mcp] Stdout pipe broken, shutting down."); + break; + } } } } @@ -208,8 +217,12 @@ async fn handle_request( if tools_cache.is_none() { *tools_cache = Some(build_tools_list(config).await?); } + let tools = match tools_cache.as_ref() { + Some(t) => json!(t), + None => json!([]), + }; Ok(json!({ - "tools": tools_cache.as_ref().unwrap() + "tools": tools })) } "tools/call" => { @@ -733,7 +746,9 @@ async fn handle_tools_call(params: &Value, config: &ServerConfig) -> Result= 2); // list + get + let names: Vec<&str> = tools.iter().filter_map(|t| t["name"].as_str()).collect(); + assert!(names.contains(&"drive_files_list")); + assert!(names.contains(&"drive_files_get")); + } + + #[test] + fn test_walk_resources_nested() { + let doc = mock_nested_doc(); + let mut tools = Vec::new(); + walk_resources("gmail", &doc.resources, &mut tools); + let names: Vec<&str> = tools.iter().filter_map(|t| t["name"].as_str()).collect(); + assert!(names.contains(&"gmail_users_messages_list")); + assert!(names.contains(&"gmail_users_messages_get")); + assert!(names.contains(&"gmail_users_getProfile")); + } + + #[test] + fn test_walk_resources_empty_description_fallback() { + let mut methods = HashMap::new(); + methods.insert( + "delete".to_string(), + RestMethod { + description: None, + ..Default::default() + }, + ); + let mut resources = HashMap::new(); + resources.insert( + "items".to_string(), + RestResource { + methods, + resources: HashMap::new(), + }, + ); + + let mut tools = Vec::new(); + walk_resources("tasks", &resources, &mut tools); + assert_eq!(tools.len(), 1); + let desc = tools[0]["description"].as_str().unwrap(); + assert!(desc.contains("tasks_items_delete")); + } + + #[test] + fn test_walk_resources_empty() { + let resources = HashMap::new(); + let mut tools = Vec::new(); + walk_resources("empty", &resources, &mut tools); + assert!(tools.is_empty()); + } + + // -- handle_request tests -- + + #[tokio::test] + async fn test_handle_request_initialize() { + let config = ServerConfig { + services: vec![], + workflows: false, + _helpers: false, + tool_mode: ToolMode::Full, + }; + let mut cache = None; + let result = handle_request("initialize", &json!({}), &config, &mut cache) + .await + .unwrap(); + + assert_eq!(result["protocolVersion"], "2024-11-05"); + assert_eq!(result["serverInfo"]["name"], "gws-mcp"); + assert!(result["capabilities"]["tools"].is_object()); + } + + #[tokio::test] + async fn test_handle_request_unsupported_method() { + let config = ServerConfig { + services: vec![], + workflows: false, + _helpers: false, + tool_mode: ToolMode::Full, + }; + let mut cache = None; + let result = handle_request("foo/bar", &json!({}), &config, &mut cache).await; + assert!(result.is_err()); + assert!(result.unwrap_err().to_string().contains("Method not supported")); + } + + #[tokio::test] + async fn test_handle_request_notifications_initialized() { + let config = ServerConfig { + services: vec![], + workflows: false, + _helpers: false, + tool_mode: ToolMode::Full, + }; + let mut cache = None; + let result = handle_request("notifications/initialized", &json!({}), &config, &mut cache) + .await + .unwrap(); + assert_eq!(result, json!({})); + } + + // -- build_mcp_cli tests -- + + #[test] + fn test_build_mcp_cli_defaults() { + let cmd = build_mcp_cli(); + let matches = cmd.get_matches_from(vec!["mcp"]); + let svc = matches + .get_one::("services") + .map(|s| s.as_str()) + .unwrap_or(""); + assert_eq!(svc, ""); + assert!(!matches.get_flag("workflows")); + assert!(!matches.get_flag("helpers")); + } + + #[test] + fn test_build_mcp_cli_with_services() { + let cmd = build_mcp_cli(); + let matches = cmd.get_matches_from(vec!["mcp", "-s", "drive,gmail"]); + let svc = matches.get_one::("services").unwrap(); + assert_eq!(svc, "drive,gmail"); + } + + #[test] + fn test_build_mcp_cli_with_flags() { + let cmd = build_mcp_cli(); + let matches = cmd.get_matches_from(vec!["mcp", "--workflows", "--helpers"]); + assert!(matches.get_flag("workflows")); + assert!(matches.get_flag("helpers")); + } } From 4721bedfe91c96148153ca2fd15159d92b441e6d Mon Sep 17 00:00:00 2001 From: Andrew Barnes Date: Thu, 5 Mar 2026 19:31:28 -0500 Subject: [PATCH 2/2] fix: use expect/unreachable for provably safe unwrap replacements Co-Authored-By: Claude Opus 4.6 --- src/mcp_server.rs | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/src/mcp_server.rs b/src/mcp_server.rs index 3ff56656..b0edbbfd 100644 --- a/src/mcp_server.rs +++ b/src/mcp_server.rs @@ -219,7 +219,7 @@ async fn handle_request( } let tools = match tools_cache.as_ref() { Some(t) => json!(t), - None => json!([]), + None => unreachable!("tools_cache is guaranteed to be Some by the preceding check"), }; Ok(json!({ "tools": tools @@ -746,9 +746,9 @@ async fn handle_tools_call(params: &Value, config: &ServerConfig) -> Result