diff --git a/codex-rs/codex-api/src/files.rs b/codex-rs/codex-api/src/files.rs index 3bb376d618..3b046a881a 100644 --- a/codex-rs/codex-api/src/files.rs +++ b/codex-rs/codex-api/src/files.rs @@ -10,6 +10,7 @@ use serde::Deserialize; use tokio::fs::File; use tokio::time::Instant; use tokio_util::io::ReaderStream; +use url::Url; pub const OPENAI_FILE_URI_PREFIX: &str = "sediment://"; pub const OPENAI_FILE_UPLOAD_LIMIT_BYTES: u64 = 512 * 1024 * 1024; @@ -81,6 +82,12 @@ pub enum OpenAiFileError { #[source] source: serde_json::Error, }, + #[error("failed to resolve OpenAI file URL `{url}`: {source}")] + InvalidUrl { + url: String, + #[source] + source: url::ParseError, + }, #[error("OpenAI file upload for `{file_id}` is not ready yet")] UploadNotReady { file_id: String }, #[error("OpenAI file upload for `{file_id}` failed: {message}")] @@ -107,6 +114,45 @@ pub fn openai_file_uri(file_id: &str) -> String { format!("{OPENAI_FILE_URI_PREFIX}{file_id}") } +pub async fn download_openai_file( + base_url: &str, + auth: &impl AuthProvider, + download_url: &str, +) -> Result, OpenAiFileError> { + let resolved_url = resolve_openai_file_download_url(base_url, download_url)?; + let request_builder = if should_attach_auth_to_openai_file_url(&resolved_url, base_url) { + authorized_request(auth, reqwest::Method::GET, resolved_url.as_str()) + } else { + build_reqwest_client() + .request(reqwest::Method::GET, resolved_url.as_str()) + .timeout(OPENAI_FILE_REQUEST_TIMEOUT) + }; + let response = request_builder + .send() + .await + .map_err(|source| OpenAiFileError::Request { + url: resolved_url.to_string(), + source, + })?; + let status = response.status(); + if !status.is_success() { + let body = response.text().await.unwrap_or_default(); + return Err(OpenAiFileError::UnexpectedStatus { + url: resolved_url.to_string(), + status, + body, + }); + } + let bytes = response + .bytes() + .await + .map_err(|source| OpenAiFileError::Request { + url: resolved_url.to_string(), + source, + })?; + Ok(bytes.to_vec()) +} + pub async fn upload_local_file( base_url: &str, auth: &impl AuthProvider, @@ -284,6 +330,46 @@ fn authorized_request( .headers(headers) } +fn resolve_openai_file_download_url( + base_url: &str, + download_url: &str, +) -> Result { + match Url::parse(download_url) { + Ok(url) => Ok(url), + Err(url::ParseError::RelativeUrlWithoutBase) => { + let normalized_base_url = if base_url.ends_with('/') { + base_url.to_string() + } else { + format!("{base_url}/") + }; + let base = + Url::parse(&normalized_base_url).map_err(|source| OpenAiFileError::InvalidUrl { + url: normalized_base_url.clone(), + source, + })?; + base.join(download_url) + .map_err(|source| OpenAiFileError::InvalidUrl { + url: download_url.to_string(), + source, + }) + } + Err(source) => Err(OpenAiFileError::InvalidUrl { + url: download_url.to_string(), + source, + }), + } +} + +fn should_attach_auth_to_openai_file_url(download_url: &Url, base_url: &str) -> bool { + let Ok(base_url) = Url::parse(base_url) else { + return false; + }; + match (download_url.host_str(), base_url.host_str()) { + (Some(download_host), Some(base_host)) => download_host.eq_ignore_ascii_case(base_host), + _ => false, + } +} + fn build_reqwest_client() -> reqwest::Client { build_reqwest_client_with_custom_ca(reqwest::Client::builder()).unwrap_or_else(|error| { tracing::warn!(error = %error, "failed to build OpenAI file upload client"); @@ -317,6 +403,32 @@ mod tests { format!("{}/backend-api", server.uri()) } + #[tokio::test] + async fn download_openai_file_resolves_relative_url_and_attaches_auth() { + let server = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/files/download/file_123")) + .and(header("authorization", "Bearer token")) + .and(header("chatgpt-account-id", "account_id")) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "text/plain") + .set_body_bytes(b"hello".to_vec()), + ) + .mount(&server) + .await; + + let downloaded = download_openai_file( + &format!("{}/backend-api/codex", server.uri()), + &chatgpt_auth(), + "/files/download/file_123", + ) + .await + .expect("download succeeds"); + + assert_eq!(downloaded, b"hello".to_vec()); + } + #[tokio::test] async fn upload_local_file_returns_canonical_uri() { let server = MockServer::start().await; diff --git a/codex-rs/codex-api/src/lib.rs b/codex-rs/codex-api/src/lib.rs index 9e501a8ca6..b4e0b32290 100644 --- a/codex-rs/codex-api/src/lib.rs +++ b/codex-rs/codex-api/src/lib.rs @@ -55,6 +55,7 @@ pub use crate::endpoint::ResponsesWebsocketConnection; pub use crate::endpoint::session_update_session_json; pub use crate::error::ApiError; pub use crate::files::OpenAiFileUploadOptions; +pub use crate::files::download_openai_file; pub use crate::files::upload_local_file; pub use crate::provider::Provider; pub use crate::provider::RetryConfig; diff --git a/codex-rs/core/src/codex_apps_file_download.rs b/codex-rs/core/src/codex_apps_file_download.rs new file mode 100644 index 0000000000..6d5a87343e --- /dev/null +++ b/codex-rs/core/src/codex_apps_file_download.rs @@ -0,0 +1,413 @@ +use crate::codex::Session; +use crate::codex::TurnContext; +use crate::codex_apps_mcp_tools::should_materialize_codex_apps_file_download; +use codex_api::CoreAuthProvider; +use codex_api::download_openai_file; +use codex_login::CodexAuth; +use codex_protocol::mcp::CallToolResult; +use serde::Deserialize; +use serde::Serialize; +use serde_json::Map as JsonMap; +use serde_json::Value as JsonValue; +use tracing::warn; + +const CODEX_APPS_FILE_DOWNLOAD_ARTIFACTS_DIR: &str = ".tmp/codex_apps_downloads"; + +#[derive(Debug, Deserialize, Serialize)] +struct CodexAppsFileDownloadPayload { + file_id: String, + #[serde(default)] + file_name: Option, + file_uri: CodexAppsFileUri, +} + +#[derive(Debug, Deserialize, Serialize)] +struct CodexAppsFileUri { + download_url: String, + #[serde(default)] + file_name: Option, +} + +pub(crate) async fn maybe_materialize_codex_apps_file_download_result( + sess: &Session, + turn_context: &TurnContext, + server: &str, + codex_apps_meta: Option<&JsonMap>, + result: CallToolResult, +) -> CallToolResult { + let auth = sess.services.auth_manager.auth().await; + maybe_materialize_codex_apps_file_download_result_with_auth( + turn_context, + &sess.conversation_id.to_string(), + auth.as_ref(), + server, + codex_apps_meta, + result, + ) + .await +} + +async fn maybe_materialize_codex_apps_file_download_result_with_auth( + turn_context: &TurnContext, + session_id: &str, + auth: Option<&CodexAuth>, + server: &str, + codex_apps_meta: Option<&JsonMap>, + mut result: CallToolResult, +) -> CallToolResult { + if !should_materialize_codex_apps_file_download(server, codex_apps_meta) + || result.is_error == Some(true) + { + return result; + } + + let Some(payload) = extract_codex_apps_file_download_payload(&result) else { + return result; + }; + if result.structured_content.is_none() + && let Ok(structured_content) = serde_json::to_value(&payload) + { + result.structured_content = Some(structured_content); + } + + let Some(auth) = auth else { + warn!( + "skipping codex_apps file download materialization because ChatGPT auth is unavailable" + ); + return result; + }; + let token_data = match auth.get_token_data() { + Ok(token_data) => token_data, + Err(error) => { + warn!(error = %error, "failed to read ChatGPT auth for codex_apps file download materialization"); + return result; + } + }; + let download_auth = CoreAuthProvider { + token: Some(token_data.access_token), + account_id: token_data + .id_token + .chatgpt_account_id + .clone() + .or(token_data.account_id), + }; + let downloaded = match download_openai_file( + turn_context.config.chatgpt_base_url.trim_end_matches('/'), + &download_auth, + &payload.file_uri.download_url, + ) + .await + { + Ok(downloaded) => downloaded, + Err(error) => { + warn!( + error = %error, + file_id = payload.file_id, + "failed to materialize codex_apps file download via app-server", + ); + return result; + } + }; + + let artifact_path = codex_apps_file_download_artifact_path( + &turn_context.config.codex_home, + session_id, + &payload.file_id, + payload + .file_name + .as_deref() + .or(payload.file_uri.file_name.as_deref()) + .unwrap_or("downloaded_file"), + ); + if let Some(parent) = artifact_path.parent() + && let Err(error) = tokio::fs::create_dir_all(parent.as_path()).await + { + warn!( + error = %error, + path = %parent.display(), + "failed to create codex_apps file download artifact directory", + ); + return result; + } + if let Err(error) = tokio::fs::write(artifact_path.as_path(), &downloaded).await { + warn!( + error = %error, + path = %artifact_path.display(), + "failed to write codex_apps file download artifact", + ); + return result; + } + + let local_path = artifact_path.to_string_lossy().to_string(); + if let Some(JsonValue::Object(map)) = result.structured_content.as_mut() { + map.insert( + "local_path".to_string(), + JsonValue::String(local_path.clone()), + ); + } + result.content.push(serde_json::json!({ + "type": "text", + "text": format!("Downloaded file to local path: {local_path}"), + })); + result +} + +fn extract_codex_apps_file_download_payload( + result: &CallToolResult, +) -> Option { + if let Some(structured_content) = result.structured_content.clone() + && let Ok(payload) = + serde_json::from_value::(structured_content) + { + return Some(payload); + } + + result + .content + .iter() + .filter_map(|item| item.as_object()) + .find_map(|item| { + let text = item.get("text")?.as_str()?; + serde_json::from_str::(text).ok() + }) +} + +fn codex_apps_file_download_artifact_path( + codex_home: &codex_utils_absolute_path::AbsolutePathBuf, + session_id: &str, + file_id: &str, + file_name: &str, +) -> codex_utils_absolute_path::AbsolutePathBuf { + codex_home + .join(CODEX_APPS_FILE_DOWNLOAD_ARTIFACTS_DIR) + .join(sanitize_path_component(session_id, "session")) + .join(sanitize_path_component(file_id, "file")) + .join(sanitize_file_name(file_name)) +} + +fn sanitize_path_component(value: &str, fallback: &str) -> String { + let sanitized: String = value + .chars() + .map(|ch| { + if ch.is_ascii_alphanumeric() || ch == '-' || ch == '_' { + ch + } else { + '_' + } + }) + .collect(); + if sanitized.is_empty() { + fallback.to_string() + } else { + sanitized + } +} + +fn sanitize_file_name(value: &str) -> String { + let sanitized: String = value + .chars() + .map(|ch| { + if ch.is_ascii_alphanumeric() || ch == '-' || ch == '_' || ch == '.' { + ch + } else { + '_' + } + }) + .collect(); + if sanitized.is_empty() { + "downloaded_file".to_string() + } else { + sanitized + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::codex::make_session_and_context; + use codex_login::CodexAuth; + use pretty_assertions::assert_eq; + use std::sync::Arc; + use wiremock::Mock; + use wiremock::MockServer; + use wiremock::ResponseTemplate; + use wiremock::matchers::header; + use wiremock::matchers::method; + use wiremock::matchers::path; + + fn download_materialization_meta() -> JsonMap { + serde_json::json!({ + "provider": "builtin", + "materialize_file_download": true, + }) + .as_object() + .cloned() + .expect("_codex_apps metadata object") + } + + #[tokio::test] + async fn codex_apps_file_download_materialization_ignores_results_without_metadata_flag() { + let (_, turn_context) = make_session_and_context().await; + let original = CallToolResult { + content: vec![serde_json::json!({"type": "text", "text": "hello"})], + structured_content: Some(serde_json::json!({"x": 1})), + is_error: Some(false), + meta: None, + }; + + let result = maybe_materialize_codex_apps_file_download_result_with_auth( + &turn_context, + "session-1", + Some(&CodexAuth::create_dummy_chatgpt_auth_for_testing()), + "custom_server", + /*codex_apps_meta*/ None, + original.clone(), + ) + .await; + + assert_eq!(result, original); + } + + #[tokio::test] + async fn codex_apps_file_download_materialization_adds_local_path_for_marked_tools() { + let server = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/api/codex/files/file_123/content")) + .and(header("authorization", "Bearer Access Token")) + .and(header("chatgpt-account-id", "account_id")) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "text/plain") + .set_body_bytes(b"downloaded contents".to_vec()), + ) + .mount(&server) + .await; + + let (_, mut turn_context) = make_session_and_context().await; + let mut config = (*turn_context.config).clone(); + config.chatgpt_base_url = format!("{}/backend-api/codex", server.uri()); + turn_context.config = Arc::new(config); + let original = CallToolResult { + content: vec![serde_json::json!({ + "type": "text", + "text": "{\"file_id\":\"file_123\"}", + })], + structured_content: Some(serde_json::json!({ + "file_id": "file_123", + "file_name": "testing-file.txt", + "file_uri": { + "download_url": "/api/codex/files/file_123/content", + "file_id": "file_123", + "file_name": "testing-file.txt", + "mime_type": "text/plain", + } + })), + is_error: Some(false), + meta: None, + }; + + let result = maybe_materialize_codex_apps_file_download_result_with_auth( + &turn_context, + "session-1", + Some(&CodexAuth::create_dummy_chatgpt_auth_for_testing()), + codex_mcp::CODEX_APPS_MCP_SERVER_NAME, + Some(&download_materialization_meta()), + original, + ) + .await; + + let local_path = result + .structured_content + .as_ref() + .and_then(|value| value.get("local_path")) + .and_then(JsonValue::as_str) + .expect("local_path in structured content"); + assert!(local_path.contains("codex_apps_downloads")); + let saved = tokio::fs::read(local_path) + .await + .expect("saved local file should exist"); + assert_eq!(saved, b"downloaded contents".to_vec()); + assert!(result.content.iter().any(|block| { + block.get("type").and_then(JsonValue::as_str) == Some("text") + && block + .get("text") + .and_then(JsonValue::as_str) + .is_some_and(|text| text.contains("Downloaded file to local path:")) + })); + } + + #[tokio::test] + async fn codex_apps_file_download_materialization_uses_json_text_when_structured_content_is_missing() + { + let server = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/api/codex/files/file_123/content")) + .and(header("authorization", "Bearer Access Token")) + .and(header("chatgpt-account-id", "account_id")) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "text/plain") + .set_body_bytes(b"downloaded contents".to_vec()), + ) + .mount(&server) + .await; + + let (_, mut turn_context) = make_session_and_context().await; + let mut config = (*turn_context.config).clone(); + config.chatgpt_base_url = format!("{}/backend-api/codex", server.uri()); + turn_context.config = Arc::new(config); + let original = CallToolResult { + content: vec![serde_json::json!({ + "type": "text", + "text": serde_json::json!({ + "file_id": "file_123", + "file_name": "testing-file.txt", + "file_uri": { + "download_url": "/api/codex/files/file_123/content", + "file_name": "testing-file.txt", + } + }) + .to_string(), + })], + structured_content: None, + is_error: Some(false), + meta: None, + }; + + let result = maybe_materialize_codex_apps_file_download_result_with_auth( + &turn_context, + "session-1", + Some(&CodexAuth::create_dummy_chatgpt_auth_for_testing()), + codex_mcp::CODEX_APPS_MCP_SERVER_NAME, + Some(&download_materialization_meta()), + original, + ) + .await; + + let local_path = result + .content + .iter() + .find_map(|item| { + item.get("text") + .and_then(|text| text.as_str()) + .and_then(|text| text.strip_prefix("Downloaded file to local path: ")) + }) + .expect("expected local path announcement"); + assert_eq!( + result.structured_content, + Some(serde_json::json!({ + "file_id": "file_123", + "file_name": "testing-file.txt", + "file_uri": { + "download_url": "/api/codex/files/file_123/content", + "file_name": "testing-file.txt", + }, + "local_path": local_path, + })) + ); + assert_eq!( + tokio::fs::read(local_path).await.expect("downloaded file"), + b"downloaded contents" + ); + } +} diff --git a/codex-rs/core/src/codex_apps_mcp_tools.rs b/codex-rs/core/src/codex_apps_mcp_tools.rs index 055109e186..e2c4778887 100644 --- a/codex-rs/core/src/codex_apps_mcp_tools.rs +++ b/codex-rs/core/src/codex_apps_mcp_tools.rs @@ -8,13 +8,26 @@ pub(crate) const CODEX_APPS_META_KEY: &str = "_codex_apps"; const CODEX_APPS_PROVIDER_BUILTIN: &str = "builtin"; const CODEX_APPS_META_PROVIDER_KEY: &str = "provider"; const CODEX_APPS_META_DIRECT_EXPOSE_KEY: &str = "direct_expose"; +const CODEX_APPS_META_MATERIALIZE_FILE_DOWNLOAD_KEY: &str = "materialize_file_download"; pub(crate) fn is_direct_exposed_codex_apps_builtin_tool_info(tool: &McpToolInfo) -> bool { - if tool.server_name != CODEX_APPS_MCP_SERVER_NAME || tool.connector_id.is_some() { + is_direct_exposed_codex_apps_builtin( + &tool.server_name, + tool.connector_id.as_deref(), + codex_apps_meta_from_tool_info(tool), + ) +} + +pub(crate) fn is_direct_exposed_codex_apps_builtin( + server_name: &str, + connector_id: Option<&str>, + codex_apps_meta: Option<&Map>, +) -> bool { + if server_name != CODEX_APPS_MCP_SERVER_NAME || connector_id.is_some() { return false; } - let Some(codex_apps_meta) = codex_apps_meta_from_tool_info(tool) else { + let Some(codex_apps_meta) = codex_apps_meta else { return false; }; @@ -28,6 +41,28 @@ pub(crate) fn is_direct_exposed_codex_apps_builtin_tool_info(tool: &McpToolInfo) == Some(true) } +pub(crate) fn should_materialize_codex_apps_file_download( + server_name: &str, + codex_apps_meta: Option<&Map>, +) -> bool { + if server_name != CODEX_APPS_MCP_SERVER_NAME { + return false; + } + + let Some(codex_apps_meta) = codex_apps_meta else { + return false; + }; + + codex_apps_meta + .get(CODEX_APPS_META_PROVIDER_KEY) + .and_then(Value::as_str) + == Some(CODEX_APPS_PROVIDER_BUILTIN) + && codex_apps_meta + .get(CODEX_APPS_META_MATERIALIZE_FILE_DOWNLOAD_KEY) + .and_then(Value::as_bool) + == Some(true) +} + fn codex_apps_meta_from_tool_info(tool: &McpToolInfo) -> Option<&Map> { tool.tool .meta diff --git a/codex-rs/core/src/lib.rs b/codex-rs/core/src/lib.rs index 6b96b6d830..7995cadde9 100644 --- a/codex-rs/core/src/lib.rs +++ b/codex-rs/core/src/lib.rs @@ -21,6 +21,7 @@ mod compact_remote; pub use codex_thread::CodexThread; pub use codex_thread::ThreadConfigSnapshot; mod agent; +mod codex_apps_file_download; mod codex_apps_mcp_tools; mod codex_delegate; mod command_canonicalization; diff --git a/codex-rs/core/src/mcp_tool_call.rs b/codex-rs/core/src/mcp_tool_call.rs index 9a982aece0..89c890e017 100644 --- a/codex-rs/core/src/mcp_tool_call.rs +++ b/codex-rs/core/src/mcp_tool_call.rs @@ -14,6 +14,7 @@ use crate::arc_monitor::ArcMonitorOutcome; use crate::arc_monitor::monitor_action; use crate::codex::Session; use crate::codex::TurnContext; +use crate::codex_apps_file_download::maybe_materialize_codex_apps_file_download_result; use crate::codex_apps_mcp_tools::CODEX_APPS_META_KEY; use crate::config::Config; use crate::config::edit::ConfigEdit; @@ -478,6 +479,14 @@ async fn execute_mcp_tool_call( .call_tool(server, tool_name, rewritten_arguments, request_meta) .await .map_err(|e| format!("tool call error: {e:?}"))?; + let result = maybe_materialize_codex_apps_file_download_result( + sess, + turn_context, + server, + metadata.and_then(|metadata| metadata.codex_apps_meta.as_ref()), + result, + ) + .await; sanitize_mcp_tool_result_for_model( turn_context .model_info