diff --git a/crates/goose-mcp/src/computercontroller/mod.rs b/crates/goose-mcp/src/computercontroller/mod.rs index d6d281adc407..934ca24d3891 100644 --- a/crates/goose-mcp/src/computercontroller/mod.rs +++ b/crates/goose-mcp/src/computercontroller/mod.rs @@ -12,7 +12,9 @@ use tokio::{process::Command, sync::mpsc}; use std::os::unix::fs::PermissionsExt; use mcp_core::{ - handler::{PromptError, ResourceError, ToolError}, + handler::{ + require_str_parameter, require_u64_parameter, PromptError, ResourceError, ToolError, + }, protocol::ServerCapabilities, }; use mcp_server::router::CapabilitiesBuilder; @@ -595,11 +597,7 @@ impl ComputerControllerRouter { } async fn web_scrape(&self, params: Value) -> Result, ToolError> { - let url = params - .get("url") - .and_then(|v| v.as_str()) - .ok_or_else(|| ToolError::InvalidParameters("Missing 'url' parameter".into()))?; - + let url = require_str_parameter(¶ms, "url")?; let save_as = params .get("save_as") .and_then(|v| v.as_str()) @@ -916,20 +914,9 @@ impl ComputerControllerRouter { ))]) } "update_cell" => { - let row = params.get("row").and_then(|v| v.as_u64()).ok_or_else(|| { - ToolError::InvalidParameters("Missing 'row' parameter".into()) - })?; - - let col = params.get("col").and_then(|v| v.as_u64()).ok_or_else(|| { - ToolError::InvalidParameters("Missing 'col' parameter".into()) - })?; - - let value = params - .get("value") - .and_then(|v| v.as_str()) - .ok_or_else(|| { - ToolError::InvalidParameters("Missing 'value' parameter".into()) - })?; + let row = require_u64_parameter(¶ms, "row")?; + let col = require_u64_parameter(¶ms, "col")?; + let value = require_str_parameter(¶ms, "value")?; let worksheet_name = params .get("worksheet") diff --git a/crates/goose-mcp/src/developer/mod.rs b/crates/goose-mcp/src/developer/mod.rs index 8accfe0e8298..9a9ad98e49f8 100644 --- a/crates/goose-mcp/src/developer/mod.rs +++ b/crates/goose-mcp/src/developer/mod.rs @@ -23,7 +23,7 @@ use url::Url; use include_dir::{include_dir, Dir}; use mcp_core::{ - handler::{PromptError, ResourceError, ToolError}, + handler::{require_str_parameter, PromptError, ResourceError, ToolError}, protocol::ServerCapabilities, }; @@ -918,12 +918,7 @@ impl DeveloperRouter { self.text_editor_view(&path, view_range).await } "write" => { - let file_text = params - .get("file_text") - .and_then(|v| v.as_str()) - .ok_or_else(|| { - ToolError::InvalidParameters("Missing 'file_text' parameter".into()) - })?; + let file_text = require_str_parameter(¶ms, "file_text")?; self.text_editor_write(&path, file_text).await } diff --git a/crates/goose/src/agents/extension_manager.rs b/crates/goose/src/agents/extension_manager.rs index 72328d2c86ed..b3769fef4572 100644 --- a/crates/goose/src/agents/extension_manager.rs +++ b/crates/goose/src/agents/extension_manager.rs @@ -3,6 +3,7 @@ use axum::http::{HeaderMap, HeaderName}; use chrono::{DateTime, TimeZone, Utc}; use futures::stream::{FuturesUnordered, StreamExt}; use futures::{future, FutureExt}; +use mcp_core::handler::require_str_parameter; use mcp_core::{ToolCall, ToolError}; use rmcp::service::ClientInitializeError; use rmcp::transport::streamable_http_client::StreamableHttpClientTransportConfig; @@ -551,11 +552,7 @@ impl ExtensionManager { // Function that gets executed for read_resource tool pub async fn read_resource(&self, params: Value) -> Result, ToolError> { - let uri = params - .get("uri") - .and_then(|v| v.as_str()) - .ok_or_else(|| ToolError::InvalidParameters("Missing 'uri' parameter".to_string()))?; - + let uri = require_str_parameter(¶ms, "uri")?; let extension_name = params.get("extension_name").and_then(|v| v.as_str()); // If extension name is provided, we can just look it up diff --git a/crates/mcp-core/src/handler.rs b/crates/mcp-core/src/handler.rs index 4724cb97edfd..8d97a3fd1f5d 100644 --- a/crates/mcp-core/src/handler.rs +++ b/crates/mcp-core/src/handler.rs @@ -1,6 +1,6 @@ use serde::{Deserialize, Serialize}; #[allow(unused_imports)] // this is used in schema below -use serde_json::json; +use serde_json::{json, Value}; use thiserror::Error; #[non_exhaustive] @@ -35,3 +35,32 @@ pub enum PromptError { #[error("Prompt not found: {0}")] NotFound(String), } + +/// Helper function to require a string, returning a ToolError +pub fn require_str_parameter<'a>( + v: &'a serde_json::Value, + name: &str, +) -> Result<&'a str, ToolError> { + let v = v + .get(name) + .ok_or_else(|| ToolError::InvalidParameters(format!("The parameter {name} is required")))?; + match v.as_str() { + Some(r) => Ok(r), + None => Err(ToolError::InvalidParameters(format!( + "The parameter {name} must be a string" + ))), + } +} + +/// Helper function to require a u64, returning a ToolError +pub fn require_u64_parameter(v: &serde_json::Value, name: &str) -> Result { + let v = v + .get(name) + .ok_or_else(|| ToolError::InvalidParameters(format!("The parameter {name} is required")))?; + match v.as_u64() { + Some(r) => Ok(r), + None => Err(ToolError::InvalidParameters(format!( + "The parameter {name} must be a number" + ))), + } +}