diff --git a/codex-rs/core/src/tools/file_broker.rs b/codex-rs/core/src/tools/file_broker.rs index 68d581049b..bffe6131e0 100644 --- a/codex-rs/core/src/tools/file_broker.rs +++ b/codex-rs/core/src/tools/file_broker.rs @@ -1,7 +1,10 @@ -#![allow(dead_code)] - +use base64::Engine; +use base64::engine::general_purpose::STANDARD as BASE64_STANDARD; +use codex_mcp::CODEX_APPS_MCP_SERVER_NAME; +use codex_mcp::ToolInfo; use codex_tools::FileRef; use codex_tools::FileScheme; +use std::collections::BTreeMap; use std::fmt; use std::fs; use std::path::Component; @@ -18,6 +21,16 @@ pub(crate) struct CodeModeFileBroker { current_root: PathBuf, } +const MAX_DATA_URI_EXPORT_BYTES: usize = 8 * 1024 * 1024; + +const CONFIGURED_BROKER_ROUTES: &[ConfiguredBrokerRoute] = &[ConfiguredBrokerRoute { + provider: "google_drive", + operation: "upload", + required_server: CODEX_APPS_MCP_SERVER_NAME, + required_namespace: "google_drive", + required_tool: "upload_file", +}]; + impl CodeModeFileBroker { pub(crate) fn new(current_root: impl Into) -> Self { Self { @@ -69,6 +82,36 @@ impl CodeModeFileBroker { }) } + pub(crate) fn export_data_uri( + &self, + source: &FileRef, + mime_type: &str, + ) -> Result { + let bytes = self.read_to_bytes(source)?; + if bytes.len() > MAX_DATA_URI_EXPORT_BYTES { + return Err(FileBrokerError::DataUriTooLarge { + file_ref: source.raw().to_string(), + byte_count: bytes.len() as u64, + max_byte_count: MAX_DATA_URI_EXPORT_BYTES as u64, + }); + } + + let data_uri = format!( + "data:{mime_type};base64,{}", + BASE64_STANDARD.encode(bytes.as_slice()) + ); + Ok(FileBrokerDataUriResult { + source_ref: source.raw().to_string(), + mime_type: mime_type.to_string(), + data_uri, + byte_count: bytes.len() as u64, + }) + } + + pub(crate) fn active_provider_registry(mcp_tools: &[ToolInfo]) -> ActiveProviderRegistry { + ActiveProviderRegistry::from_routes(CONFIGURED_BROKER_ROUTES, mcp_tools) + } + fn resolve_env_path(&self, file_ref: &FileRef) -> Result { if file_ref.scheme() != FileScheme::Env { return Err(FileBrokerError::UnsupportedProvider { @@ -95,6 +138,97 @@ impl CodeModeFileBroker { } } +#[derive(Debug, Clone, Copy, Eq, PartialEq)] +struct ConfiguredBrokerRoute { + provider: &'static str, + operation: &'static str, + required_server: &'static str, + required_namespace: &'static str, + required_tool: &'static str, +} + +#[derive(Debug, Clone, Eq, PartialEq)] +pub(crate) struct ActiveProviderRegistry { + pub(crate) providers: Vec, + pub(crate) unavailable_routes: Vec, +} + +impl ActiveProviderRegistry { + fn from_routes(routes: &[ConfiguredBrokerRoute], mcp_tools: &[ToolInfo]) -> Self { + let mut active_operations = BTreeMap::>::new(); + let mut unavailable_routes = Vec::new(); + + for route in routes { + if route_is_available(route, mcp_tools) { + active_operations + .entry(route.provider.to_string()) + .or_default() + .push(route.operation.to_string()); + } else { + unavailable_routes.push(UnavailableBrokerRoute { + provider: route.provider.to_string(), + operation: route.operation.to_string(), + required_server: route.required_server.to_string(), + required_namespace: route.required_namespace.to_string(), + required_tool: route.required_tool.to_string(), + }); + } + } + + let providers = active_operations + .into_iter() + .map(|(provider, mut operations)| { + operations.sort(); + ActiveBrokerProvider { + provider, + operations, + } + }) + .collect(); + + Self { + providers, + unavailable_routes, + } + } + + pub(crate) fn summary(&self) -> String { + if self.providers.is_empty() { + return "none".to_string(); + } + + self.providers + .iter() + .map(|provider| format!("{}: {}", provider.provider, provider.operations.join(", "))) + .collect::>() + .join("; ") + } +} + +#[derive(Debug, Clone, Eq, PartialEq)] +pub(crate) struct ActiveBrokerProvider { + pub(crate) provider: String, + pub(crate) operations: Vec, +} + +#[derive(Debug, Clone, Eq, PartialEq)] +pub(crate) struct UnavailableBrokerRoute { + pub(crate) provider: String, + pub(crate) operation: String, + pub(crate) required_server: String, + pub(crate) required_namespace: String, + pub(crate) required_tool: String, +} + +fn route_is_available(route: &ConfiguredBrokerRoute, mcp_tools: &[ToolInfo]) -> bool { + mcp_tools.iter().any(|tool| { + tool.server_name == route.required_server + && tool.callable_namespace == route.required_namespace + && (tool.callable_name == route.required_tool + || tool.tool.name.as_ref() == route.required_tool) + }) +} + #[derive(Debug, Eq, PartialEq)] pub(crate) struct FileBrokerWriteResult { pub(crate) file_ref: String, @@ -108,6 +242,14 @@ pub(crate) struct FileBrokerCopyResult { pub(crate) byte_count: u64, } +#[derive(Debug, Eq, PartialEq)] +pub(crate) struct FileBrokerDataUriResult { + pub(crate) source_ref: String, + pub(crate) mime_type: String, + pub(crate) data_uri: String, + pub(crate) byte_count: u64, +} + #[derive(Debug)] pub(crate) enum FileBrokerError { UnsupportedProvider { @@ -123,6 +265,11 @@ pub(crate) enum FileBrokerError { action: &'static str, source: std::io::Error, }, + DataUriTooLarge { + file_ref: String, + byte_count: u64, + max_byte_count: u64, + }, } impl fmt::Display for FileBrokerError { @@ -138,12 +285,29 @@ impl fmt::Display for FileBrokerError { write!(f, "`{file_ref}` must resolve to a relative workspace path") } Self::Io { action, source } => write!(f, "failed to {action} file: {source}"), + Self::DataUriTooLarge { + file_ref, + byte_count, + max_byte_count, + } => write!( + f, + "`{file_ref}` is {byte_count} bytes; max data URI export is {max_byte_count} bytes" + ), } } } impl std::error::Error for FileBrokerError {} +impl FileBrokerError { + pub(crate) fn should_include_active_provider_status(&self) -> bool { + matches!( + self, + Self::UnsupportedProvider { .. } | Self::UnsupportedEnvironment { .. } + ) + } +} + fn clean_relative_path(path: &str) -> Option { let mut clean = PathBuf::new(); for component in Path::new(path).components() { @@ -160,12 +324,40 @@ fn clean_relative_path(path: &str) -> Option { mod tests { use super::*; use pretty_assertions::assert_eq; + use rmcp::model::Tool; + use serde_json::Map; + use std::sync::Arc; use tempfile::TempDir; fn file_ref(raw: &str) -> FileRef { FileRef::parse(raw).expect("file ref should parse") } + fn mcp_tool(server_name: &str, namespace: &str, name: &str) -> ToolInfo { + ToolInfo { + server_name: server_name.to_string(), + supports_parallel_tool_calls: false, + server_origin: None, + callable_name: name.to_string(), + callable_namespace: namespace.to_string(), + namespace_description: None, + tool: Tool { + name: name.to_string().into(), + title: None, + description: None, + input_schema: Arc::new(Map::new()), + output_schema: None, + annotations: None, + execution: None, + icons: None, + meta: None, + }, + connector_id: None, + connector_name: None, + plugin_display_names: Vec::new(), + } + } + #[test] fn writes_and_reads_env_current_refs() { let temp_dir = TempDir::new().expect("temp dir"); @@ -240,4 +432,65 @@ mod tests { Err(FileBrokerError::UnsupportedProvider { .. }) )); } + + #[test] + fn exports_env_current_ref_as_data_uri() { + let temp_dir = TempDir::new().expect("temp dir"); + let broker = CodeModeFileBroker::new(temp_dir.path()); + let source = file_ref("env://current/image.png"); + broker + .write_bytes(&source, b"png") + .expect("write should succeed"); + + let result = broker + .export_data_uri(&source, "image/png") + .expect("export should succeed"); + + assert_eq!( + result, + FileBrokerDataUriResult { + source_ref: "env://current/image.png".to_string(), + mime_type: "image/png".to_string(), + data_uri: "data:image/png;base64,cG5n".to_string(), + byte_count: 3, + } + ); + } + + #[test] + fn active_provider_registry_uses_live_mcp_inventory() { + let registry = CodeModeFileBroker::active_provider_registry(&[mcp_tool( + CODEX_APPS_MCP_SERVER_NAME, + "google_drive", + "upload_file", + )]); + + assert_eq!( + registry.providers, + vec![ActiveBrokerProvider { + provider: "google_drive".to_string(), + operations: vec!["upload".to_string()], + }] + ); + assert!(registry.unavailable_routes.is_empty()); + assert_eq!(registry.summary(), "google_drive: upload"); + } + + #[test] + fn active_provider_registry_reports_missing_route_dependency() { + let registry = CodeModeFileBroker::active_provider_registry(&[]); + + assert!(registry.providers.is_empty()); + assert_eq!( + registry.unavailable_routes, + vec![UnavailableBrokerRoute { + provider: "google_drive".to_string(), + operation: "upload".to_string(), + required_server: CODEX_APPS_MCP_SERVER_NAME.to_string(), + required_namespace: "google_drive".to_string(), + required_tool: "upload_file".to_string(), + }] + ); + assert_eq!(registry.summary(), "none"); + } } diff --git a/codex-rs/core/src/tools/handlers/files.rs b/codex-rs/core/src/tools/handlers/files.rs new file mode 100644 index 0000000000..06d842856b --- /dev/null +++ b/codex-rs/core/src/tools/handlers/files.rs @@ -0,0 +1,170 @@ +use crate::function_tool::FunctionCallError; +use crate::session::session::Session; +use crate::tools::context::ToolInvocation; +use crate::tools::context::ToolPayload; +use crate::tools::file_broker::CodeModeFileBroker; +use crate::tools::file_broker::FileBrokerError; +use crate::tools::handlers::files_spec::FILES_COPY_TOOL_NAME; +use crate::tools::handlers::files_spec::FILES_EXPORT_FOR_TOOL_NAME; +use crate::tools::handlers::files_spec::FILES_MATERIALIZE_TOOL_NAME; +use crate::tools::handlers::files_spec::FILES_NAMESPACE; +use crate::tools::handlers::files_spec::create_files_namespace_tool; +use crate::tools::handlers::parse_arguments; +use crate::tools::registry::ToolExecutor; +use crate::tools::registry::ToolHandler; +use codex_tools::FileRef; +use codex_tools::JsonToolOutput; +use codex_tools::ToolName; +use codex_tools::ToolSpec; +use serde::Deserialize; +use serde_json::json; +use std::sync::Arc; + +pub(crate) struct FilesMaterializeHandler; +pub(crate) struct FilesCopyHandler; +pub(crate) struct FilesExportForToolHandler; + +#[derive(Debug, Deserialize)] +struct SourceTargetArgs { + source_uri: String, + target_uri: String, +} + +#[derive(Debug, Deserialize)] +struct ExportForToolArgs { + file_uri: String, + mime_type: String, +} + +impl ToolExecutor for FilesMaterializeHandler { + type Output = JsonToolOutput; + + fn tool_name(&self) -> ToolName { + ToolName::namespaced(FILES_NAMESPACE, FILES_MATERIALIZE_TOOL_NAME) + } + + fn spec(&self) -> Option { + Some(create_files_namespace_tool()) + } + + async fn handle(&self, invocation: ToolInvocation) -> Result { + let session = invocation.session.clone(); + let args = parse_source_target_args(invocation.payload, FILES_MATERIALIZE_TOOL_NAME)?; + let broker = CodeModeFileBroker::new(invocation.turn.cwd.as_path()); + let result = match broker.copy( + &parse_file_ref(&args.source_uri)?, + &parse_file_ref(&args.target_uri)?, + ) { + Ok(result) => result, + Err(error) => return Err(file_broker_error(error, &session).await), + }; + Ok(JsonToolOutput::new(json!({ + "source_uri": result.source_ref, + "file_uri": result.target_ref, + "byte_count": result.byte_count, + }))) + } +} + +impl ToolHandler for FilesMaterializeHandler {} + +impl ToolExecutor for FilesCopyHandler { + type Output = JsonToolOutput; + + fn tool_name(&self) -> ToolName { + ToolName::namespaced(FILES_NAMESPACE, FILES_COPY_TOOL_NAME) + } + + async fn handle(&self, invocation: ToolInvocation) -> Result { + let session = invocation.session.clone(); + let args = parse_source_target_args(invocation.payload, FILES_COPY_TOOL_NAME)?; + let broker = CodeModeFileBroker::new(invocation.turn.cwd.as_path()); + let result = match broker.copy( + &parse_file_ref(&args.source_uri)?, + &parse_file_ref(&args.target_uri)?, + ) { + Ok(result) => result, + Err(error) => return Err(file_broker_error(error, &session).await), + }; + Ok(JsonToolOutput::new(json!({ + "source_uri": result.source_ref, + "target_uri": result.target_ref, + "byte_count": result.byte_count, + }))) + } +} + +impl ToolHandler for FilesCopyHandler {} + +impl ToolExecutor for FilesExportForToolHandler { + type Output = JsonToolOutput; + + fn tool_name(&self) -> ToolName { + ToolName::namespaced(FILES_NAMESPACE, FILES_EXPORT_FOR_TOOL_NAME) + } + + async fn handle(&self, invocation: ToolInvocation) -> Result { + let session = invocation.session.clone(); + let args = parse_export_for_tool_args(invocation.payload)?; + let broker = CodeModeFileBroker::new(invocation.turn.cwd.as_path()); + let result = match broker.export_data_uri(&parse_file_ref(&args.file_uri)?, &args.mime_type) + { + Ok(result) => result, + Err(error) => return Err(file_broker_error(error, &session).await), + }; + Ok(JsonToolOutput::new(json!({ + "file_uri": result.source_ref, + "mime_type": result.mime_type, + "data_uri": result.data_uri, + "byte_count": result.byte_count, + }))) + } +} + +impl ToolHandler for FilesExportForToolHandler {} + +fn parse_source_target_args( + payload: ToolPayload, + tool_name: &str, +) -> Result { + parse_payload_arguments(payload, tool_name).and_then(|arguments| parse_arguments(&arguments)) +} + +fn parse_export_for_tool_args( + payload: ToolPayload, +) -> Result { + parse_payload_arguments(payload, FILES_EXPORT_FOR_TOOL_NAME) + .and_then(|arguments| parse_arguments(&arguments)) +} + +fn parse_payload_arguments( + payload: ToolPayload, + tool_name: &str, +) -> Result { + match payload { + ToolPayload::Function { arguments } => Ok(arguments), + _ => Err(FunctionCallError::RespondToModel(format!( + "{FILES_NAMESPACE}.{tool_name} received unsupported payload" + ))), + } +} + +fn parse_file_ref(raw: &str) -> Result { + FileRef::parse(raw) + .map_err(|err| FunctionCallError::RespondToModel(format!("invalid file ref: {err}"))) +} + +async fn file_broker_error(error: FileBrokerError, session: &Arc) -> FunctionCallError { + if !error.should_include_active_provider_status() { + return FunctionCallError::RespondToModel(error.to_string()); + } + + let manager = session.services.mcp_connection_manager.read().await; + let mcp_tools = manager.list_all_tools().await; + let registry = CodeModeFileBroker::active_provider_registry(&mcp_tools); + FunctionCallError::RespondToModel(format!( + "{}. Active broker providers: {}", + error, + registry.summary() + )) +} diff --git a/codex-rs/core/src/tools/handlers/files_spec.rs b/codex-rs/core/src/tools/handlers/files_spec.rs new file mode 100644 index 0000000000..829dba57c7 --- /dev/null +++ b/codex-rs/core/src/tools/handlers/files_spec.rs @@ -0,0 +1,97 @@ +use codex_tools::JsonSchema; +use codex_tools::ResponsesApiNamespace; +use codex_tools::ResponsesApiNamespaceTool; +use codex_tools::ResponsesApiTool; +use codex_tools::ToolSpec; +use std::collections::BTreeMap; + +pub(crate) const FILES_NAMESPACE: &str = "files"; +pub(crate) const FILES_MATERIALIZE_TOOL_NAME: &str = "materialize"; +pub(crate) const FILES_COPY_TOOL_NAME: &str = "copy"; +pub(crate) const FILES_EXPORT_FOR_TOOL_NAME: &str = "export_for_tool"; + +pub(crate) fn create_files_namespace_tool() -> ToolSpec { + ToolSpec::Namespace(ResponsesApiNamespace { + name: FILES_NAMESPACE.to_string(), + description: "Move Code Mode file refs between the workspace and provider/tool boundaries." + .to_string(), + tools: vec![ + ResponsesApiNamespaceTool::Function(materialize_tool()), + ResponsesApiNamespaceTool::Function(copy_tool()), + ResponsesApiNamespaceTool::Function(export_for_tool_tool()), + ], + }) +} + +fn materialize_tool() -> ResponsesApiTool { + ResponsesApiTool { + name: FILES_MATERIALIZE_TOOL_NAME.to_string(), + description: "Materialize a source file ref into an environment file ref. The initial POC supports env://current/... refs; provider adapters can be added behind the same contract.".to_string(), + strict: false, + defer_loading: None, + parameters: source_target_schema(), + output_schema: None, + } +} + +fn copy_tool() -> ResponsesApiTool { + ResponsesApiTool { + name: FILES_COPY_TOOL_NAME.to_string(), + description: + "Copy bytes from one file ref to another without exposing provider credentials to the model." + .to_string(), + strict: false, + defer_loading: None, + parameters: source_target_schema(), + output_schema: None, + } +} + +fn export_for_tool_tool() -> ResponsesApiTool { + ResponsesApiTool { + name: FILES_EXPORT_FOR_TOOL_NAME.to_string(), + description: "Export a file ref as a base64 data URI for tools that declare fileParam-style inputs. Use this from Code Mode generated code immediately before invoking the destination tool; do not log the returned data URI.".to_string(), + strict: false, + defer_loading: None, + parameters: JsonSchema::object( + BTreeMap::from([ + ( + "file_uri".to_string(), + JsonSchema::string(Some("Source file ref, for example env://current/out.png.".to_string())), + ), + ( + "mime_type".to_string(), + JsonSchema::string(Some("MIME type to use in the returned data URI, for example image/png.".to_string())), + ), + ]), + Some(vec!["file_uri".to_string(), "mime_type".to_string()]), + Some(false.into()), + ), + output_schema: None, + } +} + +fn source_target_schema() -> JsonSchema { + JsonSchema::object( + BTreeMap::from([ + ( + "source_uri".to_string(), + JsonSchema::string(Some( + "Source file ref, for example env://current/report.pdf.".to_string(), + )), + ), + ( + "target_uri".to_string(), + JsonSchema::string(Some( + "Target file ref, for example env://current/out/report.pdf.".to_string(), + )), + ), + ]), + Some(vec!["source_uri".to_string(), "target_uri".to_string()]), + Some(false.into()), + ) +} + +#[cfg(test)] +#[path = "files_spec_tests.rs"] +mod tests; diff --git a/codex-rs/core/src/tools/handlers/files_spec_tests.rs b/codex-rs/core/src/tools/handlers/files_spec_tests.rs new file mode 100644 index 0000000000..50cfeaa915 --- /dev/null +++ b/codex-rs/core/src/tools/handlers/files_spec_tests.rs @@ -0,0 +1,28 @@ +use super::*; +use codex_tools::ResponsesApiNamespaceTool; +use codex_tools::ToolSpec; +use pretty_assertions::assert_eq; + +#[test] +fn files_namespace_contains_poc_actions() { + let ToolSpec::Namespace(namespace) = create_files_namespace_tool() else { + panic!("files tool should be a namespace"); + }; + + assert_eq!(namespace.name, FILES_NAMESPACE); + let tool_names = namespace + .tools + .iter() + .map(|tool| match tool { + ResponsesApiNamespaceTool::Function(tool) => tool.name.as_str(), + }) + .collect::>(); + assert_eq!( + tool_names, + vec![ + FILES_MATERIALIZE_TOOL_NAME, + FILES_COPY_TOOL_NAME, + FILES_EXPORT_FOR_TOOL_NAME, + ] + ); +} diff --git a/codex-rs/core/src/tools/handlers/mod.rs b/codex-rs/core/src/tools/handlers/mod.rs index 69281acc7d..83496580b7 100644 --- a/codex-rs/core/src/tools/handlers/mod.rs +++ b/codex-rs/core/src/tools/handlers/mod.rs @@ -4,6 +4,8 @@ pub(crate) mod apply_patch; pub(crate) mod apply_patch_spec; mod dynamic; pub(crate) mod extension_tools; +mod files; +pub(crate) mod files_spec; mod goal; pub(crate) mod goal_spec; mod mcp; @@ -51,6 +53,9 @@ pub use apply_patch::ApplyPatchHandler; use codex_protocol::models::AdditionalPermissionProfile; use codex_protocol::protocol::AskForApproval; pub use dynamic::DynamicToolHandler; +pub(crate) use files::FilesCopyHandler; +pub(crate) use files::FilesExportForToolHandler; +pub(crate) use files::FilesMaterializeHandler; pub use goal::CreateGoalHandler; pub use goal::GetGoalHandler; pub use goal::UpdateGoalHandler; diff --git a/codex-rs/core/src/tools/spec_plan.rs b/codex-rs/core/src/tools/spec_plan.rs index c4db408c55..37c21c7e7c 100644 --- a/codex-rs/core/src/tools/spec_plan.rs +++ b/codex-rs/core/src/tools/spec_plan.rs @@ -6,6 +6,9 @@ use crate::tools::handlers::CreateGoalHandler; use crate::tools::handlers::DynamicToolHandler; use crate::tools::handlers::ExecCommandHandler; use crate::tools::handlers::ExecCommandHandlerOptions; +use crate::tools::handlers::FilesCopyHandler; +use crate::tools::handlers::FilesExportForToolHandler; +use crate::tools::handlers::FilesMaterializeHandler; use crate::tools::handlers::GetGoalHandler; use crate::tools::handlers::ListMcpResourceTemplatesHandler; use crate::tools::handlers::ListMcpResourcesHandler; @@ -333,6 +336,16 @@ fn collect_handler_tools( handlers.push(Arc::new(TestSyncHandler)); } + if config + .experimental_supported_tools + .iter() + .any(|tool| tool == "code_mode_files") + { + handlers.push(Arc::new(FilesMaterializeHandler)); + handlers.push(Arc::new(FilesCopyHandler)); + handlers.push(Arc::new(FilesExportForToolHandler)); + } + if config.environment_mode.has_environment() { let include_environment_id = matches!(config.environment_mode, ToolEnvironmentMode::Multiple);