diff --git a/codex-rs/core/src/tools/handlers/mcp_resource.rs b/codex-rs/core/src/tools/handlers/mcp_resource.rs index 0cd5857302..deb7491e75 100644 --- a/codex-rs/core/src/tools/handlers/mcp_resource.rs +++ b/codex-rs/core/src/tools/handlers/mcp_resource.rs @@ -15,6 +15,7 @@ use codex_protocol::protocol::TruncationPolicy; use codex_utils_output_truncation::truncate_text; use rmcp::model::ListResourceTemplatesResult; use rmcp::model::ListResourcesResult; +use rmcp::model::PaginatedRequestParams; use rmcp::model::ReadResourceResult; use rmcp::model::Resource; use rmcp::model::ResourceTemplate; @@ -56,22 +57,41 @@ fn ensure_model_can_access_mcp_server( } } -#[derive(Debug, Deserialize, Default)] -struct ListResourcesArgs { - /// Lists all resources from all servers if not specified. +#[derive(Debug, Deserialize, Default, PartialEq, Eq)] +struct ListResourceArgs { #[serde(default)] server: Option, #[serde(default)] cursor: Option, } -#[derive(Debug, Deserialize, Default)] -struct ListResourceTemplatesArgs { - /// Lists all resource templates from all servers if not specified. - #[serde(default)] - server: Option, - #[serde(default)] - cursor: Option, +impl ListResourceArgs { + fn normalized(self) -> Self { + Self { + server: normalize_optional_string(self.server), + cursor: normalize_optional_string(self.cursor), + } + } + + fn target( + &self, + turn: &TurnContext, + ) -> Result)>, FunctionCallError> { + match &self.server { + Some(server) => { + ensure_model_can_access_mcp_server(turn, server)?; + let params = self + .cursor + .clone() + .map(|cursor| PaginatedRequestParams::default().with_cursor(Some(cursor))); + Ok(Some((server.clone(), params))) + } + None if self.cursor.is_some() => Err(FunctionCallError::RespondToModel( + "cursor can only be used when a server is specified".to_string(), + )), + None => Ok(None), + } + } } #[derive(Debug, Deserialize)] @@ -81,28 +101,31 @@ struct ReadResourceArgs { } #[derive(Debug, Serialize)] -struct ResourceWithServer { +struct ResourceWithServer { server: String, #[serde(flatten)] - resource: Resource, + resource: T, } -impl ResourceWithServer { - fn new(server: String, resource: Resource) -> Self { +impl ResourceWithServer { + fn new(server: String, resource: T) -> Self { Self { server, resource } } -} -#[derive(Debug, Serialize)] -struct ResourceTemplateWithServer { - server: String, - #[serde(flatten)] - template: ResourceTemplate, -} + fn from_server(server: &str, resources: Vec) -> Vec { + resources + .into_iter() + .map(|resource| Self::new(server.to_string(), resource)) + .collect() + } -impl ResourceTemplateWithServer { - fn new(server: String, template: ResourceTemplate) -> Self { - Self { server, template } + fn from_all_servers(resources_by_server: HashMap>) -> Vec { + let mut entries: Vec<_> = resources_by_server.into_iter().collect(); + entries.sort_by(|(left, _), (right, _)| left.cmp(right)); + entries + .into_iter() + .flat_map(|(server, resources)| Self::from_server(&server, resources)) + .collect() } } @@ -111,39 +134,24 @@ impl ResourceTemplateWithServer { struct ListResourcesPayload { #[serde(skip_serializing_if = "Option::is_none")] server: Option, - resources: Vec, + resources: Vec>, #[serde(skip_serializing_if = "Option::is_none")] next_cursor: Option, } impl ListResourcesPayload { fn from_single_server(server: String, result: ListResourcesResult) -> Self { - let resources = result - .resources - .into_iter() - .map(|resource| ResourceWithServer::new(server.clone(), resource)) - .collect(); Self { + resources: ResourceWithServer::from_server(&server, result.resources), server: Some(server), - resources, next_cursor: result.next_cursor, } } fn from_all_servers(resources_by_server: HashMap>) -> Self { - let mut entries: Vec<(String, Vec)> = resources_by_server.into_iter().collect(); - entries.sort_by(|a, b| a.0.cmp(&b.0)); - - let mut resources = Vec::new(); - for (server, server_resources) in entries { - for resource in server_resources { - resources.push(ResourceWithServer::new(server.clone(), resource)); - } - } - Self { server: None, - resources, + resources: ResourceWithServer::from_all_servers(resources_by_server), next_cursor: None, } } @@ -154,40 +162,24 @@ impl ListResourcesPayload { struct ListResourceTemplatesPayload { #[serde(skip_serializing_if = "Option::is_none")] server: Option, - resource_templates: Vec, + resource_templates: Vec>, #[serde(skip_serializing_if = "Option::is_none")] next_cursor: Option, } impl ListResourceTemplatesPayload { fn from_single_server(server: String, result: ListResourceTemplatesResult) -> Self { - let resource_templates = result - .resource_templates - .into_iter() - .map(|template| ResourceTemplateWithServer::new(server.clone(), template)) - .collect(); Self { + resource_templates: ResourceWithServer::from_server(&server, result.resource_templates), server: Some(server), - resource_templates, next_cursor: result.next_cursor, } } fn from_all_servers(templates_by_server: HashMap>) -> Self { - let mut entries: Vec<(String, Vec)> = - templates_by_server.into_iter().collect(); - entries.sort_by(|a, b| a.0.cmp(&b.0)); - - let mut resource_templates = Vec::new(); - for (server, server_templates) in entries { - for template in server_templates { - resource_templates.push(ResourceTemplateWithServer::new(server.clone(), template)); - } - } - Self { server: None, - resource_templates, + resource_templates: ResourceWithServer::from_all_servers(templates_by_server), next_cursor: None, } } diff --git a/codex-rs/core/src/tools/handlers/mcp_resource/list_mcp_resource_templates.rs b/codex-rs/core/src/tools/handlers/mcp_resource/list_mcp_resource_templates.rs index 94b0132376..a988038dec 100644 --- a/codex-rs/core/src/tools/handlers/mcp_resource/list_mcp_resource_templates.rs +++ b/codex-rs/core/src/tools/handlers/mcp_resource/list_mcp_resource_templates.rs @@ -8,13 +8,9 @@ use codex_protocol::protocol::McpInvocation; use codex_tools::ToolName; use codex_tools::ToolSpec; -use rmcp::model::PaginatedRequestParams; - -use super::ListResourceTemplatesArgs; +use super::ListResourceArgs; use super::ListResourceTemplatesPayload; -use super::ensure_model_can_access_mcp_server; use super::model_can_access_mcp_server; -use super::normalize_optional_string; use super::parse_args_with_default; use super::parse_arguments; use super::run_resource_operation; @@ -64,23 +60,17 @@ impl ListMcpResourceTemplatesHandler { }; let arguments = parse_arguments(arguments.as_str())?; - let args: ListResourceTemplatesArgs = parse_args_with_default(arguments.clone())?; - let ListResourceTemplatesArgs { server, cursor } = args; - let server = normalize_optional_string(server); - let cursor = normalize_optional_string(cursor); + let args: ListResourceArgs = parse_args_with_default(arguments.clone())?; + let args = args.normalized(); let invocation = McpInvocation { - server: server.clone().unwrap_or_else(|| "codex".to_string()), + server: args.server.clone().unwrap_or_else(|| "codex".to_string()), tool: "list_mcp_resource_templates".to_string(), arguments: arguments.clone(), }; run_resource_operation(&session, turn.as_ref(), &call_id, invocation, async { - if let Some(server_name) = server.clone() { - ensure_model_can_access_mcp_server(turn.as_ref(), &server_name)?; - let params = cursor - .clone() - .map(|value| PaginatedRequestParams::default().with_cursor(Some(value))); + if let Some((server_name, params)) = args.target(turn.as_ref())? { let result = mcp .list_resource_templates(&server_name, params) .await @@ -94,12 +84,6 @@ impl ListMcpResourceTemplatesHandler { result, )) } else { - if cursor.is_some() { - return Err(FunctionCallError::RespondToModel( - "cursor can only be used when a server is specified".to_string(), - )); - } - let templates = mcp .list_all_resource_templates(|server_name| { model_can_access_mcp_server(turn.as_ref(), server_name) diff --git a/codex-rs/core/src/tools/handlers/mcp_resource/list_mcp_resources.rs b/codex-rs/core/src/tools/handlers/mcp_resource/list_mcp_resources.rs index a966fbf4d2..d0d0d0f564 100644 --- a/codex-rs/core/src/tools/handlers/mcp_resource/list_mcp_resources.rs +++ b/codex-rs/core/src/tools/handlers/mcp_resource/list_mcp_resources.rs @@ -8,13 +8,9 @@ use codex_protocol::protocol::McpInvocation; use codex_tools::ToolName; use codex_tools::ToolSpec; -use rmcp::model::PaginatedRequestParams; - -use super::ListResourcesArgs; +use super::ListResourceArgs; use super::ListResourcesPayload; -use super::ensure_model_can_access_mcp_server; use super::model_can_access_mcp_server; -use super::normalize_optional_string; use super::parse_args_with_default; use super::parse_arguments; use super::run_resource_operation; @@ -64,23 +60,17 @@ impl ListMcpResourcesHandler { }; let arguments = parse_arguments(arguments.as_str())?; - let args: ListResourcesArgs = parse_args_with_default(arguments.clone())?; - let ListResourcesArgs { server, cursor } = args; - let server = normalize_optional_string(server); - let cursor = normalize_optional_string(cursor); + let args: ListResourceArgs = parse_args_with_default(arguments.clone())?; + let args = args.normalized(); let invocation = McpInvocation { - server: server.clone().unwrap_or_else(|| "codex".to_string()), + server: args.server.clone().unwrap_or_else(|| "codex".to_string()), tool: "list_mcp_resources".to_string(), arguments: arguments.clone(), }; run_resource_operation(&session, turn.as_ref(), &call_id, invocation, async { - if let Some(server_name) = server.clone() { - ensure_model_can_access_mcp_server(turn.as_ref(), &server_name)?; - let params = cursor - .clone() - .map(|value| PaginatedRequestParams::default().with_cursor(Some(value))); + if let Some((server_name, params)) = args.target(turn.as_ref())? { let result = mcp .list_resources(&server_name, params) .await @@ -92,12 +82,6 @@ impl ListMcpResourcesHandler { result, )) } else { - if cursor.is_some() { - return Err(FunctionCallError::RespondToModel( - "cursor can only be used when a server is specified".to_string(), - )); - } - let resources = mcp .list_all_resources(|server_name| { model_can_access_mcp_server(turn.as_ref(), server_name) diff --git a/codex-rs/core/src/tools/handlers/mcp_resource_tests.rs b/codex-rs/core/src/tools/handlers/mcp_resource_tests.rs index 1155060d7c..7f58175206 100644 --- a/codex-rs/core/src/tools/handlers/mcp_resource_tests.rs +++ b/codex-rs/core/src/tools/handlers/mcp_resource_tests.rs @@ -88,9 +88,26 @@ fn parse_arguments_handles_empty_and_json() { assert_eq!(value["server"], json!("figma")); } +#[test] +fn list_resource_args_normalizes_server_and_cursor() { + let args: ListResourceArgs = serde_json::from_value(json!({ + "server": " hosted ", + "cursor": " next-page " + })) + .expect("parse resource-list arguments"); + + assert_eq!( + args.normalized(), + ListResourceArgs { + server: Some("hosted".to_string()), + cursor: Some("next-page".to_string()), + } + ); +} + #[test] fn template_with_server_serializes_server_field() { - let entry = ResourceTemplateWithServer::new("srv".to_string(), template("memo://{id}", "memo")); + let entry = ResourceWithServer::new("srv".to_string(), template("memo://{id}", "memo")); let value = serde_json::to_value(&entry).expect("serialize template"); assert_eq!( @@ -103,6 +120,31 @@ fn template_with_server_serializes_server_field() { ); } +#[test] +fn list_resource_templates_payload_from_all_servers_is_sorted() { + let mut templates_by_server = HashMap::new(); + templates_by_server.insert( + "beta".to_string(), + vec![template("memo://beta/{id}", "beta")], + ); + templates_by_server.insert( + "alpha".to_string(), + vec![template("memo://alpha/{id}", "alpha")], + ); + + let payload = ListResourceTemplatesPayload::from_all_servers(templates_by_server); + + assert_eq!( + serde_json::to_value(payload).expect("serialize resource templates"), + json!({ + "resourceTemplates": [ + {"server": "alpha", "uriTemplate": "memo://alpha/{id}", "name": "alpha"}, + {"server": "beta", "uriTemplate": "memo://beta/{id}", "name": "beta"} + ] + }) + ); +} + #[test] fn serialize_function_output_preserves_small_payload() { let payload = json!({"server": "hosted", "resources": []});