github/codex
4
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-10-22 01:34:54 +00:00
Code Issues Packages Projects Releases Wiki Activity

Deduplicate MCP resource list handling (#36734)

Browse Source 
## What changed

- Share argument normalization, server selection, cursor validation, and pagination setup between `list_mcp_resources` and `list_mcp_resource_templates`.
- Use a generic server-tagged wrapper to build single-server and deterministically sorted all-server payloads for both resources and templates.

## Testing

- Cover shared argument normalization and sorted resource-template payload serialization.

GitOrigin-RevId: 3505e4c1d2b95d8bbbb01f6f930fb914309eb898
This commit is contained in:
jif
2026-08-03 10:41:03 +00:00
committed by copyberry
parent 87e2d41eb3
commit d6407d7359
4 changed files with 106 additions and 104 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/core/src/tools/handlers/mcp_resource.rs
+53 -61
View File
@@ -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<String>,
#[serde(default)]
cursor: Option<String>,
}
#[derive(Debug, Deserialize, Default)]
struct ListResourceTemplatesArgs {
/// Lists all resource templates from all servers if not specified.
#[serde(default)]
server: Option<String>,
#[serde(default)]
cursor: Option<String>,
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<Option<(String, Option<PaginatedRequestParams>)>, 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<T> {
server: String,
#[serde(flatten)]
resource: Resource,
resource: T,
}
impl ResourceWithServer {
fn new(server: String, resource: Resource) -> Self {
impl<T> ResourceWithServer<T> {
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<T>) -> Vec<Self> {
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<String, Vec<T>>) -> Vec<Self> {
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<String>,
resources: Vec<ResourceWithServer>,
resources: Vec<ResourceWithServer<Resource>>,
#[serde(skip_serializing_if = "Option::is_none")]
next_cursor: Option<String>,
}
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<String, Vec<Resource>>) -> Self {
let mut entries: Vec<(String, Vec<Resource>)> = 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<String>,
resource_templates: Vec<ResourceTemplateWithServer>,
resource_templates: Vec<ResourceWithServer<ResourceTemplate>>,
#[serde(skip_serializing_if = "Option::is_none")]
next_cursor: Option<String>,
}
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<String, Vec<ResourceTemplate>>) -> Self {
let mut entries: Vec<(String, Vec<ResourceTemplate>)> =
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,
}
}
codex-rs/core/src/tools/handlers/mcp_resource/list_mcp_resource_templates.rs
+5 -21
View File
@@ -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)
codex-rs/core/src/tools/handlers/mcp_resource/list_mcp_resources.rs
+5 -21
View File
@@ -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)
codex-rs/core/src/tools/handlers/mcp_resource_tests.rs
+43 -1
View File
@@ -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": []});