Codex Handbook
core/src/tools/handlers/mcp_resource.rs 335 lines
use std::collections::HashMap;use std::sync::Arc;use std::time::Duration;use codex_protocol::items::McpToolCallError;use codex_protocol::items::McpToolCallItem;use codex_protocol::items::McpToolCallStatus;use codex_protocol::items::TurnItem;use codex_protocol::mcp::CallToolResult;use codex_protocol::protocol::TruncationPolicy;use codex_utils_output_truncation::truncate_text;use rmcp::model::ListResourceTemplatesResult;use rmcp::model::ListResourcesResult;use rmcp::model::ReadResourceResult;use rmcp::model::Resource;use rmcp::model::ResourceTemplate;use serde::Deserialize;use serde::Serialize;use serde::de::DeserializeOwned;use serde_json::Value;use crate::function_tool::FunctionCallError;use crate::session::session::Session;use crate::session::turn_context::TurnContext;use crate::tools::context::FunctionToolOutput;use codex_protocol::protocol::McpInvocation;mod list_mcp_resource_templates;mod list_mcp_resources;mod read_mcp_resource;pub use list_mcp_resource_templates::ListMcpResourceTemplatesHandler;pub use list_mcp_resources::ListMcpResourcesHandler;pub use read_mcp_resource::ReadMcpResourceHandler;#[derive(Debug, Deserialize, Default)]struct ListResourcesArgs {    /// Lists all resources from all servers if not specified.    #[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>,}#[derive(Debug, Deserialize)]struct ReadResourceArgs {    server: String,    uri: String,}#[derive(Debug, Serialize)]struct ResourceWithServer {    server: String,    #[serde(flatten)]    resource: Resource,}impl ResourceWithServer {    fn new(server: String, resource: Resource) -> Self {        Self { server, resource }    }}#[derive(Debug, Serialize)]struct ResourceTemplateWithServer {    server: String,    #[serde(flatten)]    template: ResourceTemplate,}impl ResourceTemplateWithServer {    fn new(server: String, template: ResourceTemplate) -> Self {        Self { server, template }    }}#[derive(Debug, Serialize)]#[serde(rename_all = "camelCase")]struct ListResourcesPayload {    #[serde(skip_serializing_if = "Option::is_none")]    server: Option<String>,    resources: Vec<ResourceWithServer>,    #[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 {            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,            next_cursor: None,        }    }}#[derive(Debug, Serialize)]#[serde(rename_all = "camelCase")]struct ListResourceTemplatesPayload {    #[serde(skip_serializing_if = "Option::is_none")]    server: Option<String>,    resource_templates: Vec<ResourceTemplateWithServer>,    #[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 {            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,            next_cursor: None,        }    }}#[derive(Debug, Serialize)]struct ReadResourcePayload {    server: String,    uri: String,    #[serde(flatten)]    result: ReadResourceResult,}fn call_tool_result_from_content(content: &str, success: Option<bool>) -> CallToolResult {    CallToolResult {        content: vec![serde_json::json!({"type": "text", "text": content})],        structured_content: None,        is_error: success.map(|value| !value),        meta: None,    }}async fn emit_tool_call_begin(    session: &Arc<Session>,    turn: &TurnContext,    call_id: &str,    invocation: McpInvocation,) {    let McpInvocation {        server,        tool,        arguments,    } = invocation;    let item = TurnItem::McpToolCall(McpToolCallItem {        id: call_id.to_string(),        server,        tool,        arguments: arguments.unwrap_or(Value::Null),        mcp_app_resource_uri: None,        plugin_id: None,        status: McpToolCallStatus::InProgress,        result: None,        error: None,        duration: None,    });    session.emit_turn_item_started(turn, &item).await;}async fn emit_tool_call_end(    session: &Arc<Session>,    turn: &TurnContext,    call_id: &str,    invocation: McpInvocation,    duration: Duration,    result: Result<CallToolResult, String>,) {    let (status, result, error) = match result {        Ok(result) if result.is_error.unwrap_or(false) => {            (McpToolCallStatus::Failed, Some(result), None)        }        Ok(result) => (McpToolCallStatus::Completed, Some(result), None),        Err(message) => (            McpToolCallStatus::Failed,            None,            Some(McpToolCallError { message }),        ),    };    let McpInvocation {        server,        tool,        arguments,    } = invocation;    let item = TurnItem::McpToolCall(McpToolCallItem {        id: call_id.to_string(),        server,        tool,        arguments: arguments.unwrap_or(Value::Null),        mcp_app_resource_uri: None,        plugin_id: None,        status,        result,        error,        duration: Some(duration),    });    session.emit_turn_item_completed(turn, item).await;}fn normalize_optional_string(input: Option<String>) -> Option<String> {    input.and_then(|value| {        let trimmed = value.trim().to_string();        if trimmed.is_empty() {            None        } else {            Some(trimmed)        }    })}fn normalize_required_string(field: &str, value: String) -> Result<String, FunctionCallError> {    match normalize_optional_string(Some(value)) {        Some(normalized) => Ok(normalized),        None => Err(FunctionCallError::RespondToModel(format!(            "{field} must be provided"        ))),    }}fn serialize_function_output<T>(    payload: T,    truncation_policy: TruncationPolicy,) -> Result<FunctionToolOutput, FunctionCallError>where    T: Serialize,{    let content = serde_json::to_string(&payload).map_err(|err| {        FunctionCallError::RespondToModel(format!(            "failed to serialize MCP resource response: {err}"        ))    })?;    // Match regular MCP tool outputs by bounding the copy persisted to the    // rollout and injected into model context.    let content = truncate_text(&content, truncation_policy * 1.2);    Ok(FunctionToolOutput::from_text(content, Some(true)))}fn parse_arguments(raw_args: &str) -> Result<Option<Value>, FunctionCallError> {    if raw_args.trim().is_empty() {        Ok(None)    } else {        let value: Value = serde_json::from_str(raw_args).map_err(|err| {            FunctionCallError::RespondToModel(format!("failed to parse function arguments: {err}"))        })?;        if value.is_null() {            Ok(None)        } else {            Ok(Some(value))        }    }}fn parse_args<T>(arguments: Option<Value>) -> Result<T, FunctionCallError>where    T: DeserializeOwned,{    match arguments {        Some(value) => serde_json::from_value(value).map_err(|err| {            FunctionCallError::RespondToModel(format!("failed to parse function arguments: {err}"))        }),        None => Err(FunctionCallError::RespondToModel(            "failed to parse function arguments: expected value".to_string(),        )),    }}fn parse_args_with_default<T>(arguments: Option<Value>) -> Result<T, FunctionCallError>where    T: DeserializeOwned + Default,{    match arguments {        Some(value) => parse_args(Some(value)),        None => Ok(T::default()),    }}#[cfg(test)]#[path = "mcp_resource_tests.rs"]mod tests;