From 19ce6394af86b5e721f648cb9b68e52c2cb1f1ea Mon Sep 17 00:00:00 2001 From: Won Park Date: Thu, 11 Jun 2026 16:32:52 -0700 Subject: [PATCH 1/4] Route image extension reads through turn environments v2 (#27498) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Why Image generation used `std::fs::read` for referenced image paths, which did not support environment-backed filesystems or their sandbox context. ## What changed - Expose optional turn environments to extension tool calls. - Include each environment’s ID, working directory, filesystem, and sandbox context. - Read referenced images through the selected environment filesystem. - Keep sandbox usage at the extension call site so extensions can choose the appropriate access mode. - Consolidate image request construction into one async function. - Add coverage for successful environment reads and read failures. ## Validation - `cargo check -p codex-image-generation-extension --tests` - `just fmt` - `just bazel-lock-update` - `just bazel-lock-check` `just test -p codex-image-generation-extension` could not complete because the build exhausted available disk space. --- codex-rs/Cargo.lock | 2 + .../tests/suite/v2/imagegen_extension.rs | 7 +- .../src/tools/handlers/extension_tools.rs | 24 ++ .../core/tests/suite/extension_sandbox.rs | 305 ++++++++++++++++++ codex-rs/core/tests/suite/mod.rs | 2 + codex-rs/ext/extension-api/src/lib.rs | 1 + .../ext/goal/tests/goal_extension_backend.rs | 1 + codex-rs/ext/image-generation/Cargo.toml | 1 + codex-rs/ext/image-generation/src/tests.rs | 44 ++- codex-rs/ext/image-generation/src/tool.rs | 44 ++- codex-rs/ext/memories/src/tests.rs | 6 + codex-rs/tools/Cargo.toml | 1 + codex-rs/tools/src/lib.rs | 1 + codex-rs/tools/src/tool_call.rs | 18 ++ 14 files changed, 423 insertions(+), 34 deletions(-) create mode 100644 codex-rs/core/tests/suite/extension_sandbox.rs diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index a8f7f3f1b3..8375f2252f 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -3139,6 +3139,7 @@ dependencies = [ "schemars 0.8.22", "serde", "serde_json", + "tokio", ] [[package]] @@ -3949,6 +3950,7 @@ dependencies = [ "codex-app-server-protocol", "codex-code-mode", "codex-features", + "codex-file-system", "codex-protocol", "codex-utils-absolute-path", "codex-utils-cargo-bin", diff --git a/codex-rs/app-server/tests/suite/v2/imagegen_extension.rs b/codex-rs/app-server/tests/suite/v2/imagegen_extension.rs index c73b5593c6..025c6ec186 100644 --- a/codex-rs/app-server/tests/suite/v2/imagegen_extension.rs +++ b/codex-rs/app-server/tests/suite/v2/imagegen_extension.rs @@ -34,6 +34,7 @@ const TINY_PNG_BYTES: &[u8] = &[ 0, 0, 31, 21, 196, 137, 0, 0, 0, 11, 73, 68, 65, 84, 120, 156, 99, 96, 0, 2, 0, 0, 5, 0, 1, 122, 94, 171, 63, 0, 0, 0, 0, 73, 69, 78, 68, 174, 66, 96, 130, ]; +const TINY_PNG_DATA_URL: &str = "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAAC0lEQVR4nGNgAAIAAAUAAXpeqz8AAAAASUVORK5CYII="; #[derive(Clone, Copy)] enum ImagegenTestMode { @@ -170,11 +171,7 @@ async fn standalone_image_edit_uses_attached_model_visible_image() -> Result<()> }) .await?; assert_eq!(edit_request["prompt"], "add a red hat"); - assert!( - edit_request["images"][0]["image_url"] - .as_str() - .is_some_and(|image_url| image_url.starts_with("data:image/png;base64,")) - ); + assert_eq!(edit_request["images"][0]["image_url"], TINY_PNG_DATA_URL); Ok(()) } diff --git a/codex-rs/core/src/tools/handlers/extension_tools.rs b/codex-rs/core/src/tools/handlers/extension_tools.rs index d429b6c1fc..2da47fc7be 100644 --- a/codex-rs/core/src/tools/handlers/extension_tools.rs +++ b/codex-rs/core/src/tools/handlers/extension_tools.rs @@ -5,18 +5,21 @@ use codex_protocol::items::TurnItem; use codex_tools::ConversationHistory; use codex_tools::ExtensionTurnItem; use codex_tools::ToolCall as ExtensionToolCall; +use codex_tools::ToolEnvironment; use codex_tools::ToolName; use codex_tools::ToolSearchInfo; use codex_tools::ToolSpec; use codex_tools::TurnItemEmissionFuture; use codex_tools::TurnItemEmitter; +use crate::sandboxing::SandboxPermissions; use crate::session::session::Session; use crate::session::turn_context::TurnContext; use crate::stream_events_utils::TurnItemContributorPolicy; use crate::stream_events_utils::finalize_turn_item; use crate::tools::context::ToolInvocation; use crate::tools::context::ToolPayload; +use crate::tools::handlers::apply_granted_turn_permissions; use crate::tools::registry::CoreToolRuntime; use crate::tools::registry::ToolExecutor; @@ -109,6 +112,26 @@ impl TurnItemEmitter for CoreTurnItemEmitter { async fn to_extension_call(invocation: &ToolInvocation) -> ExtensionToolCall { let conversation_history = ConversationHistory::new(invocation.session.clone_history().await.into_raw_items()); + let mut environments = Vec::with_capacity(invocation.turn.environments.turn_environments.len()); + for environment in &invocation.turn.environments.turn_environments { + let additional_permissions = apply_granted_turn_permissions( + invocation.session.as_ref(), + &environment.environment_id, + environment.cwd.as_path(), + SandboxPermissions::UseDefault, + /*additional_permissions*/ None, + ) + .await + .additional_permissions; + environments.push(ToolEnvironment { + environment_id: environment.environment_id.clone(), + cwd: environment.cwd.clone(), + file_system: environment.environment.get_filesystem(), + file_system_sandbox_context: invocation + .turn + .file_system_sandbox_context(additional_permissions, &environment.cwd), + }); + } ExtensionToolCall { turn_id: invocation.turn.sub_id.clone(), call_id: invocation.call_id.clone(), @@ -120,6 +143,7 @@ async fn to_extension_call(invocation: &ToolInvocation) -> ExtensionToolCall { session: Arc::downgrade(&invocation.session), turn: Arc::downgrade(&invocation.turn), }), + environments, payload: invocation.payload.clone(), } } diff --git a/codex-rs/core/tests/suite/extension_sandbox.rs b/codex-rs/core/tests/suite/extension_sandbox.rs new file mode 100644 index 0000000000..e4c01557de --- /dev/null +++ b/codex-rs/core/tests/suite/extension_sandbox.rs @@ -0,0 +1,305 @@ +use std::sync::Arc; + +use anyhow::Context; +use anyhow::Result; +use codex_core::config::Config; +use codex_core::config::Constrained; +use codex_extension_api::ExtensionRegistry; +use codex_extension_api::ExtensionRegistryBuilder; +use codex_features::Feature; +use codex_image_generation_extension::install as install_image_generation_extension; +use codex_login::CodexAuth; +use codex_protocol::config_types::ApprovalsReviewer; +use codex_protocol::config_types::WebSearchMode; +use codex_protocol::models::FileSystemPermissions; +use codex_protocol::models::PermissionProfile; +use codex_protocol::openai_models::InputModality; +use codex_protocol::permissions::FileSystemAccessMode; +use codex_protocol::permissions::FileSystemPath; +use codex_protocol::permissions::FileSystemSandboxEntry; +use codex_protocol::permissions::FileSystemSandboxPolicy; +use codex_protocol::permissions::NetworkSandboxPolicy; +use codex_protocol::protocol::AskForApproval; +use codex_protocol::protocol::EventMsg; +use codex_protocol::protocol::Op; +use codex_protocol::request_permissions::PermissionGrantScope; +use codex_protocol::request_permissions::RequestPermissionProfile; +use codex_protocol::request_permissions::RequestPermissionsResponse; +use codex_protocol::user_input::UserInput; +use core_test_support::responses; +use core_test_support::skip_if_no_network; +use core_test_support::skip_if_sandbox; +use core_test_support::test_codex::local_selections; +use core_test_support::test_codex::test_codex; +use core_test_support::test_codex::turn_permission_fields; +use core_test_support::wait_for_event; +use pretty_assertions::assert_eq; +use serde_json::json; +use wiremock::Mock; +use wiremock::ResponseTemplate; +use wiremock::matchers::method; +use wiremock::matchers::path; + +const TINY_PNG_BYTES: &[u8] = &[ + 137, 80, 78, 71, 13, 10, 26, 10, 0, 0, 0, 13, 73, 72, 68, 82, 0, 0, 0, 1, 0, 0, 0, 1, 8, 6, 0, + 0, 0, 31, 21, 196, 137, 0, 0, 0, 11, 73, 68, 65, 84, 120, 156, 99, 96, 0, 2, 0, 0, 5, 0, 1, + 122, 94, 171, 63, 0, 0, 0, 0, 73, 69, 78, 68, 174, 66, 96, 130, +]; + +fn image_generation_extensions(auth: &CodexAuth) -> Arc> { + let auth_manager = codex_core::test_support::auth_manager_from_auth(auth.clone()); + let mut extension_builder = ExtensionRegistryBuilder::::new(); + install_image_generation_extension(&mut extension_builder, auth_manager); + Arc::new(extension_builder.build()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn extension_tool_receives_turn_environment_sandbox() -> Result<()> { + skip_if_no_network!(Ok(())); + + let server = responses::start_mock_server().await; + let auth = CodexAuth::create_dummy_chatgpt_auth_for_testing(); + let extensions = image_generation_extensions(&auth); + let mut builder = test_codex() + .with_auth(auth) + .with_extensions(extensions) + .with_model_info_override("gpt-5.4", |model_info| { + model_info.use_responses_lite = true; + model_info.input_modalities = vec![InputModality::Text, InputModality::Image]; + }) + .with_config(|config| { + assert!(config.web_search_mode.set(WebSearchMode::Live).is_ok()); + assert!(config.features.enable(Feature::ImageGeneration).is_ok()); + assert!(config.features.disable(Feature::ImageGenExt).is_ok()); + }); + let test = builder.build(&server).await?; + let denied_path = test.config.cwd.join("denied.png"); + std::fs::write(&denied_path, b"not readable")?; + + let call_id = "image-edit-denied"; + let response_mock = responses::mount_sse_sequence( + &server, + vec![ + responses::sse(vec![ + responses::ev_response_created("resp-1"), + responses::ev_function_call_with_namespace( + call_id, + "image_gen", + "imagegen", + &json!({ + "prompt": "edit the image", + "referenced_image_paths": [denied_path.display().to_string()], + }) + .to_string(), + ), + responses::ev_completed("resp-1"), + ]), + responses::sse(vec![ + responses::ev_response_created("resp-2"), + responses::ev_assistant_message("msg-1", "done"), + responses::ev_completed("resp-2"), + ]), + ], + ) + .await; + + let mut file_system_sandbox_policy = FileSystemSandboxPolicy::default(); + file_system_sandbox_policy + .entries + .push(FileSystemSandboxEntry { + path: FileSystemPath::Path { + path: denied_path.clone(), + }, + access: FileSystemAccessMode::Deny, + }); + let permission_profile = PermissionProfile::from_runtime_permissions( + &file_system_sandbox_policy, + NetworkSandboxPolicy::Restricted, + ); + + test.submit_turn_with_permission_profile("edit the denied image", permission_profile) + .await?; + + let request = response_mock + .last_request() + .context("missing request containing extension output")?; + let output = request + .function_call_output_content_and_success(call_id) + .and_then(|(content, _)| content) + .context("extension error text should be present")?; + assert!( + output.starts_with(&format!( + "unable to read referenced image at `{}`:", + denied_path.display() + )), + "unexpected extension error: {output}" + ); + + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn extension_tool_uses_granted_turn_permissions() -> Result<()> { + skip_if_no_network!(Ok(())); + skip_if_sandbox!(Ok(())); + + let server = responses::start_mock_server().await; + Mock::given(method("POST")) + .and(path("/v1/images/edits")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "created": 1, + "data": [{"b64_json": "cG5n"}], + }))) + .expect(1) + .mount(&server) + .await; + + let auth = CodexAuth::create_dummy_chatgpt_auth_for_testing(); + let extensions = image_generation_extensions(&auth); + let base_permission_profile = PermissionProfile::workspace_write_with( + &[], + NetworkSandboxPolicy::Restricted, + /*exclude_tmpdir_env_var*/ true, + /*exclude_slash_tmp*/ true, + ); + let permission_profile_for_config = base_permission_profile.clone(); + let mut builder = test_codex() + .with_auth(auth) + .with_extensions(extensions) + .with_model_info_override("gpt-5.4", |model_info| { + model_info.use_responses_lite = true; + model_info.input_modalities = vec![InputModality::Text, InputModality::Image]; + }) + .with_config(move |config| { + config.permissions.approval_policy = Constrained::allow_any(AskForApproval::OnRequest); + config + .permissions + .set_permission_profile(permission_profile_for_config) + .expect("set permission profile"); + assert!(config.web_search_mode.set(WebSearchMode::Live).is_ok()); + assert!(config.features.enable(Feature::ImageGeneration).is_ok()); + assert!(config.features.disable(Feature::ImageGenExt).is_ok()); + assert!( + config + .features + .enable(Feature::RequestPermissionsTool) + .is_ok() + ); + }); + let test = builder.build(&server).await?; + + let image_dir = tempfile::tempdir()?; + let image_path = image_dir.path().canonicalize()?.join("granted.png"); + std::fs::write(&image_path, TINY_PNG_BYTES)?; + let requested_permissions = RequestPermissionProfile { + file_system: Some(FileSystemPermissions::from_read_write_roots( + Some(vec![image_dir.path().canonicalize()?.try_into()?]), + Some(Vec::new()), + )), + ..RequestPermissionProfile::default() + }; + let permissions_call_id = "permissions-call"; + let image_call_id = "image-edit-granted"; + let response_mock = responses::mount_sse_sequence( + &server, + vec![ + responses::sse(vec![ + responses::ev_response_created("resp-1"), + responses::ev_function_call( + permissions_call_id, + "request_permissions", + &serde_json::to_string(&json!({ + "reason": "Read an image outside the workspace", + "permissions": requested_permissions, + }))?, + ), + responses::ev_completed("resp-1"), + ]), + responses::sse(vec![ + responses::ev_response_created("resp-2"), + responses::ev_function_call_with_namespace( + image_call_id, + "image_gen", + "imagegen", + &json!({ + "prompt": "edit the image", + "referenced_image_paths": [image_path.display().to_string()], + }) + .to_string(), + ), + responses::ev_completed("resp-2"), + ]), + responses::sse(vec![ + responses::ev_response_created("resp-3"), + responses::ev_assistant_message("msg-1", "done"), + responses::ev_completed("resp-3"), + ]), + ], + ) + .await; + + let (sandbox_policy, permission_profile) = + turn_permission_fields(base_permission_profile, test.config.cwd.as_path()); + test.codex + .submit(Op::UserInput { + items: vec![UserInput::Text { + text: "request access and edit the image".to_string(), + text_elements: Vec::new(), + }], + final_output_json_schema: None, + responsesapi_client_metadata: None, + additional_context: Default::default(), + thread_settings: codex_protocol::protocol::ThreadSettingsOverrides { + environments: Some(local_selections(test.config.cwd.clone())), + approval_policy: Some(AskForApproval::OnRequest), + approvals_reviewer: Some(ApprovalsReviewer::User), + sandbox_policy: Some(sandbox_policy), + permission_profile, + collaboration_mode: Some(codex_protocol::config_types::CollaborationMode { + mode: codex_protocol::config_types::ModeKind::Default, + settings: codex_protocol::config_types::Settings { + model: test.session_configured.model.clone(), + reasoning_effort: None, + developer_instructions: None, + }, + }), + ..Default::default() + }, + }) + .await?; + let event = wait_for_event(&test.codex, |event| { + matches!( + event, + EventMsg::RequestPermissions(_) | EventMsg::TurnComplete(_) + ) + }) + .await; + let EventMsg::RequestPermissions(request) = event else { + panic!("expected request_permissions before turn completion"); + }; + assert_eq!(request.call_id, permissions_call_id); + test.codex + .submit(Op::RequestPermissionsResponse { + id: permissions_call_id.to_string(), + response: RequestPermissionsResponse { + permissions: request.permissions, + scope: PermissionGrantScope::Turn, + strict_auto_review: false, + }, + }) + .await?; + wait_for_event(&test.codex, |event| { + matches!(event, EventMsg::TurnComplete(_)) + }) + .await; + + let request = response_mock + .last_request() + .context("missing request containing extension output")?; + let output = request.function_call_output(image_call_id); + let image = &output["output"][0]; + assert_eq!(image["type"], "input_image"); + assert_eq!(image["image_url"], "data:image/png;base64,cG5n"); + + Ok(()) +} diff --git a/codex-rs/core/tests/suite/mod.rs b/codex-rs/core/tests/suite/mod.rs index 5764068bd9..87dc0ffce3 100644 --- a/codex-rs/core/tests/suite/mod.rs +++ b/codex-rs/core/tests/suite/mod.rs @@ -51,6 +51,8 @@ mod compact_resume_fork; mod deprecation_notice; mod exec; mod exec_policy; +#[cfg(not(target_os = "windows"))] +mod extension_sandbox; mod fork_thread; #[cfg(not(target_os = "windows"))] mod guardian_review; diff --git a/codex-rs/ext/extension-api/src/lib.rs b/codex-rs/ext/extension-api/src/lib.rs index 7004af2633..829ef295b9 100644 --- a/codex-rs/ext/extension-api/src/lib.rs +++ b/codex-rs/ext/extension-api/src/lib.rs @@ -20,6 +20,7 @@ pub use codex_tools::JsonToolOutput; pub use codex_tools::NoopTurnItemEmitter; pub use codex_tools::ResponsesApiTool; pub use codex_tools::ToolCall; +pub use codex_tools::ToolEnvironment; pub use codex_tools::ToolExecutor; pub use codex_tools::ToolExecutorFuture; pub use codex_tools::ToolName; diff --git a/codex-rs/ext/goal/tests/goal_extension_backend.rs b/codex-rs/ext/goal/tests/goal_extension_backend.rs index 60c992ec78..13492a6699 100644 --- a/codex-rs/ext/goal/tests/goal_extension_backend.rs +++ b/codex-rs/ext/goal/tests/goal_extension_backend.rs @@ -1341,6 +1341,7 @@ fn tool_call(tool_name: &str, call_id: &str, arguments: serde_json::Value) -> To truncation_policy: TruncationPolicy::Bytes(1024), conversation_history: codex_extension_api::ConversationHistory::default(), turn_item_emitter: Arc::new(NoopTurnItemEmitter), + environments: Vec::new(), payload: ToolPayload::Function { arguments: arguments.to_string(), }, diff --git a/codex-rs/ext/image-generation/Cargo.toml b/codex-rs/ext/image-generation/Cargo.toml index dfa85ee82e..ef01e2d2be 100644 --- a/codex-rs/ext/image-generation/Cargo.toml +++ b/codex-rs/ext/image-generation/Cargo.toml @@ -30,3 +30,4 @@ serde_json = { workspace = true } [dev-dependencies] pretty_assertions = { workspace = true } +tokio = { workspace = true, features = ["macros", "rt"] } diff --git a/codex-rs/ext/image-generation/src/tests.rs b/codex-rs/ext/image-generation/src/tests.rs index 3f72972656..f0f483a2aa 100644 --- a/codex-rs/ext/image-generation/src/tests.rs +++ b/codex-rs/ext/image-generation/src/tests.rs @@ -21,7 +21,7 @@ use super::GeneratedImageOutput; use super::ImageRequest; use super::ImagegenArgs; use super::imagegen_tool_spec; -use super::request_for_args; +use super::request_for_call_args; use crate::IMAGE_GEN_NAMESPACE; use crate::IMAGEGEN_TOOL_NAME; @@ -37,17 +37,19 @@ fn uses_reserved_image_gen_namespace() { assert_eq!(function.name, IMAGEGEN_TOOL_NAME); } -#[test] -fn omitted_references_generate_with_fixed_defaults() { +#[tokio::test] +async fn omitted_references_generate_with_fixed_defaults() { assert_eq!( - request_for_args( + request_for_call_args( &ImagegenArgs { prompt: "paint a moonlit lake".to_string(), referenced_image_paths: None, num_last_images_to_include: None, }, - &[] + &[], + &[], ) + .await .expect("generation request should build"), ImageRequest::Generate(ImageGenerationRequest { prompt: "paint a moonlit lake".to_string(), @@ -60,8 +62,8 @@ fn omitted_references_generate_with_fixed_defaults() { ); } -#[test] -fn recent_image_fallback_selects_newest_images_in_chronological_order() { +#[tokio::test] +async fn recent_image_fallback_selects_newest_images_in_chronological_order() { let history = vec![ ResponseItem::Message { id: None, @@ -111,14 +113,16 @@ fn recent_image_fallback_selects_newest_images_in_chronological_order() { ]; assert_eq!( - request_for_args( + request_for_call_args( &ImagegenArgs { prompt: "change the lighting".to_string(), referenced_image_paths: None, num_last_images_to_include: Some(4), }, &history, + &[], ) + .await .expect("history-backed edit request should build"), ImageRequest::Edit(expected_edit_request( "change the lighting", @@ -127,9 +131,9 @@ fn recent_image_fallback_selects_newest_images_in_chronological_order() { ); } -#[test] -fn conflicting_image_selectors_return_tool_error() { - let error = request_for_args( +#[tokio::test] +async fn conflicting_image_selectors_return_tool_error() { + let error = request_for_call_args( &ImagegenArgs { prompt: "change the lighting".to_string(), referenced_image_paths: Some(vec![ @@ -140,7 +144,9 @@ fn conflicting_image_selectors_return_tool_error() { num_last_images_to_include: Some(1), }, &[], + &[], ) + .await .expect_err("conflicting selectors should fail"); assert_eq!( @@ -149,9 +155,9 @@ fn conflicting_image_selectors_return_tool_error() { ); } -#[test] -fn too_many_referenced_image_paths_return_tool_error() { - let error = request_for_args( +#[tokio::test] +async fn too_many_referenced_image_paths_return_tool_error() { + let error = request_for_call_args( &ImagegenArgs { prompt: "change the lighting".to_string(), referenced_image_paths: Some( @@ -166,7 +172,9 @@ fn too_many_referenced_image_paths_return_tool_error() { num_last_images_to_include: None, }, &[], + &[], ) + .await .expect_err("too many paths should fail before reading files"); assert_eq!( @@ -175,9 +183,9 @@ fn too_many_referenced_image_paths_return_tool_error() { ); } -#[test] -fn recent_image_fallback_requires_requested_count() { - let error = request_for_args( +#[tokio::test] +async fn recent_image_fallback_requires_requested_count() { + let error = request_for_call_args( &ImagegenArgs { prompt: "change the lighting".to_string(), referenced_image_paths: None, @@ -189,7 +197,9 @@ fn recent_image_fallback_requires_requested_count() { content: vec![input_image("only-image")], phase: None, }], + &[], ) + .await .expect_err("history-backed edit should require the requested image count"); assert_eq!( diff --git a/codex-rs/ext/image-generation/src/tool.rs b/codex-rs/ext/image-generation/src/tool.rs index 07ed8f4bf6..0a409f3e1b 100644 --- a/codex-rs/ext/image-generation/src/tool.rs +++ b/codex-rs/ext/image-generation/src/tool.rs @@ -10,6 +10,7 @@ use codex_core::image_generation_artifact_path; use codex_extension_api::ExtensionTurnItem; use codex_extension_api::FunctionCallError; use codex_extension_api::ToolCall; +use codex_extension_api::ToolEnvironment; use codex_extension_api::ToolExecutor; use codex_extension_api::ToolName; use codex_extension_api::ToolOutput; @@ -103,7 +104,9 @@ impl ToolExecutor for ImageGenerationTool { impl ImageGenerationTool { async fn handle_call(&self, call: ToolCall) -> Result, FunctionCallError> { let args = parse_args(&call)?; - let request = request_for_args(&args, call.conversation_history.items())?; + let request = + request_for_call_args(&args, call.conversation_history.items(), &call.environments) + .await?; call.turn_item_emitter .emit_started(ExtensionTurnItem::ImageGeneration(ImageGenerationItem { id: call.call_id.clone(), @@ -154,10 +157,10 @@ enum ImageRequest { Edit(ImageEditRequest), } -/// Builds a generation or edit request from the mutually exclusive image selectors. -fn request_for_args( +async fn request_for_call_args( args: &ImagegenArgs, history: &[ResponseItem], + environments: &[ToolEnvironment], ) -> Result { let paths = args.referenced_image_paths.as_deref().unwrap_or_default(); if paths.len() > MAX_EDIT_IMAGES { @@ -176,7 +179,18 @@ fn request_for_args( size: Some("auto".to_string()), })); } - (false, None) => paths.iter().map(image_url).collect::, _>>()?, + (false, None) => { + let Some(environment) = environments.first() else { + return Err(FunctionCallError::RespondToModel( + "referenced image paths are unavailable in this session".to_string(), + )); + }; + let mut images = Vec::with_capacity(paths.len()); + for path in paths { + images.push(image_url(path, environment).await?); + } + images + } (true, Some(count)) => { if !(1..=MAX_EDIT_IMAGES).contains(&count) { return Err(FunctionCallError::RespondToModel(format!( @@ -214,7 +228,6 @@ fn request_for_args( })) } -/// Selects the newest requested images while preserving their conversation order. fn recent_images(history: &[ResponseItem], count: usize) -> Vec { let mut function_call_ids = HashSet::new(); let mut custom_tool_call_ids = HashSet::new(); @@ -307,13 +320,20 @@ fn output_image_urls(output: &FunctionCallOutputPayload) -> impl Iterator Result { - let bytes = std::fs::read(path).map_err(|error| { - FunctionCallError::RespondToModel(format!( - "unable to read referenced image at `{}`: {error}", - path.display() - )) - })?; +async fn image_url( + path: &AbsolutePathBuf, + environment: &ToolEnvironment, +) -> Result { + let bytes = environment + .file_system + .read_file(path, Some(&environment.file_system_sandbox_context)) + .await + .map_err(|error| { + FunctionCallError::RespondToModel(format!( + "unable to read referenced image at `{}`: {error}", + path.display() + )) + })?; let image = load_for_prompt_bytes(path.as_path(), bytes, PromptImageMode::Original).map_err( |error| { FunctionCallError::RespondToModel(format!( diff --git a/codex-rs/ext/memories/src/tests.rs b/codex-rs/ext/memories/src/tests.rs index ce03d74965..b78c20c9f1 100644 --- a/codex-rs/ext/memories/src/tests.rs +++ b/codex-rs/ext/memories/src/tests.rs @@ -215,6 +215,7 @@ async fn add_ad_hoc_note_tool_creates_note_file() { truncation_policy: TruncationPolicy::Bytes(1024), conversation_history: codex_extension_api::ConversationHistory::default(), turn_item_emitter: Arc::new(NoopTurnItemEmitter), + environments: Vec::new(), payload: payload.clone(), }) .await @@ -258,6 +259,7 @@ async fn add_ad_hoc_note_tool_rejects_paths_as_filenames() { truncation_policy: TruncationPolicy::Bytes(1024), conversation_history: codex_extension_api::ConversationHistory::default(), turn_item_emitter: Arc::new(NoopTurnItemEmitter), + environments: Vec::new(), payload, }) .await; @@ -302,6 +304,7 @@ async fn read_tool_reads_memory_file() { truncation_policy: TruncationPolicy::Bytes(1024), conversation_history: codex_extension_api::ConversationHistory::default(), turn_item_emitter: Arc::new(NoopTurnItemEmitter), + environments: Vec::new(), payload: payload.clone(), }) .await @@ -349,6 +352,7 @@ async fn search_tool_accepts_multiple_queries() { truncation_policy: TruncationPolicy::Bytes(1024), conversation_history: codex_extension_api::ConversationHistory::default(), turn_item_emitter: Arc::new(NoopTurnItemEmitter), + environments: Vec::new(), payload: payload.clone(), }) .await @@ -422,6 +426,7 @@ async fn search_tool_accepts_windowed_all_match_mode() { truncation_policy: TruncationPolicy::Bytes(1024), conversation_history: codex_extension_api::ConversationHistory::default(), turn_item_emitter: Arc::new(NoopTurnItemEmitter), + environments: Vec::new(), payload: payload.clone(), }) .await @@ -475,6 +480,7 @@ async fn search_tool_rejects_legacy_single_query() { truncation_policy: TruncationPolicy::Bytes(1024), conversation_history: codex_extension_api::ConversationHistory::default(), turn_item_emitter: Arc::new(NoopTurnItemEmitter), + environments: Vec::new(), payload, }) .await; diff --git a/codex-rs/tools/Cargo.toml b/codex-rs/tools/Cargo.toml index 4b2a3257f3..7d06acb0c0 100644 --- a/codex-rs/tools/Cargo.toml +++ b/codex-rs/tools/Cargo.toml @@ -11,6 +11,7 @@ workspace = true codex-app-server-protocol = { workspace = true } codex-code-mode = { workspace = true } codex-features = { workspace = true } +codex-file-system = { workspace = true } codex-protocol = { workspace = true } codex-utils-absolute-path = { workspace = true } codex-utils-output-truncation = { workspace = true } diff --git a/codex-rs/tools/src/lib.rs b/codex-rs/tools/src/lib.rs index 51c313eb0e..8bcad84981 100644 --- a/codex-rs/tools/src/lib.rs +++ b/codex-rs/tools/src/lib.rs @@ -66,6 +66,7 @@ pub use tool_call::ConversationHistory; pub use tool_call::ExtensionTurnItem; pub use tool_call::NoopTurnItemEmitter; pub use tool_call::ToolCall; +pub use tool_call::ToolEnvironment; pub use tool_call::TurnItemEmissionFuture; pub use tool_call::TurnItemEmitter; pub use tool_config::ShellCommandBackendConfig; diff --git a/codex-rs/tools/src/tool_call.rs b/codex-rs/tools/src/tool_call.rs index e3fc21e864..fdbfffad86 100644 --- a/codex-rs/tools/src/tool_call.rs +++ b/codex-rs/tools/src/tool_call.rs @@ -1,9 +1,12 @@ use crate::FunctionCallError; use crate::ToolName; use crate::ToolPayload; +use codex_file_system::ExecutorFileSystem; +use codex_file_system::FileSystemSandboxContext; use codex_protocol::items::ImageGenerationItem; use codex_protocol::items::WebSearchItem; use codex_protocol::models::ResponseItem; +use codex_utils_absolute_path::AbsolutePathBuf; use codex_utils_output_truncation::TruncationPolicy; use std::future::Future; use std::pin::Pin; @@ -49,6 +52,19 @@ pub trait TurnItemEmitter: Send + Sync { fn emit_completed<'a>(&'a self, item: ExtensionTurnItem) -> TurnItemEmissionFuture<'a>; } +/// Host-owned turn environment summary visible to extension tools. +#[derive(Clone)] +pub struct ToolEnvironment { + /// Stable host environment id used to route executor-scoped capabilities. + pub environment_id: String, + /// Effective working directory for this turn in the environment. + pub cwd: AbsolutePathBuf, + /// Filesystem implementation for this environment. + pub file_system: Arc, + /// Sandbox context to use for filesystem operations. + pub file_system_sandbox_context: FileSystemSandboxContext, +} + /// Turn-item emitter used when a caller does not expose visible item emission. #[derive(Debug, Default, Clone, Copy)] pub struct NoopTurnItemEmitter; @@ -72,6 +88,7 @@ pub struct ToolCall { pub truncation_policy: TruncationPolicy, pub conversation_history: ConversationHistory, pub turn_item_emitter: Arc, + pub environments: Vec, pub payload: ToolPayload, } @@ -85,6 +102,7 @@ impl std::fmt::Debug for ToolCall { .field("truncation_policy", &self.truncation_policy) .field("conversation_history", &self.conversation_history) .field("turn_item_emitter", &"") + .field("environment_count", &self.environments.len()) .field("payload", &self.payload) .finish() } From 92d90369108daefad443183379a55357e8448b1a Mon Sep 17 00:00:00 2001 From: mchen-oai Date: Thu, 11 Jun 2026 16:55:01 -0700 Subject: [PATCH 2/4] Add spans to turn lifecycle gaps (#27623) ## Why Codex app-server latency traces do not granularly cover turn task startup and inter-request handoffs. These spans help attribute time across task execution, startup prewarm, in-flight tool completion, and rollout persistence. ## What changed - Add `session_task.run` spans around task execution and `session_task.flush_rollout` around flushing pending conversation transcript writes to durable storage - Add `regular_task.prepare_run_turn` around regular-turn startup (Send the `TurnStarted` event, reset turn-specific reasoning state, and resolve any startup prewarm) - Add `startup_prewarm.resolve` around waiting for background session prewarming to finish, fail, time out, or be cancelled - Add a function-level trace span around draining in-flight tool calls (Wait for tool calls to complete, record tool result in conversation history, and other bookkeeping) ## Verification Trigger Codex rollout and observe new spans are included --- codex-rs/core/src/session/mod.rs | 1 + codex-rs/core/src/session/turn.rs | 1 + codex-rs/core/src/session_startup_prewarm.rs | 2 ++ codex-rs/core/src/tasks/mod.rs | 2 ++ codex-rs/core/src/tasks/regular.rs | 29 +++++++++++--------- 5 files changed, 22 insertions(+), 13 deletions(-) diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 68c05e2c43..08c691b9c1 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -1106,6 +1106,7 @@ impl Session { } /// Flush rollout writes and return the final durability-barrier result. + #[instrument(name = "session.flush_rollout", level = "trace", skip_all)] pub(crate) async fn flush_rollout(&self) -> std::io::Result<()> { if let Some(live_thread) = self.live_thread() { live_thread.flush().await.map_err(std::io::Error::other) diff --git a/codex-rs/core/src/session/turn.rs b/codex-rs/core/src/session/turn.rs index 7bb34aec87..da1d825310 100644 --- a/codex-rs/core/src/session/turn.rs +++ b/codex-rs/core/src/session/turn.rs @@ -1764,6 +1764,7 @@ async fn handle_assistant_item_done_in_plan_mode( false } +#[instrument(level = "trace", skip_all)] async fn drain_in_flight( in_flight: &mut FuturesOrdered>>, sess: Arc, diff --git a/codex-rs/core/src/session_startup_prewarm.rs b/codex-rs/core/src/session_startup_prewarm.rs index 4fee5e7747..0474eb8704 100644 --- a/codex-rs/core/src/session_startup_prewarm.rs +++ b/codex-rs/core/src/session_startup_prewarm.rs @@ -6,6 +6,7 @@ use tokio::task::JoinHandle; use tokio_util::sync::CancellationToken; use tokio_util::task::AbortOnDropHandle; use tracing::info; +use tracing::instrument; use tracing::warn; use crate::client::ModelClientSession; @@ -53,6 +54,7 @@ impl SessionStartupPrewarmHandle { let _ = self.task.await; } + #[instrument(name = "startup_prewarm.resolve", level = "trace", skip_all)] async fn resolve( self, session_telemetry: &SessionTelemetry, diff --git a/codex-rs/core/src/tasks/mod.rs b/codex-rs/core/src/tasks/mod.rs index 52af0d04db..61f6980156 100644 --- a/codex-rs/core/src/tasks/mod.rs +++ b/codex-rs/core/src/tasks/mod.rs @@ -19,6 +19,7 @@ use tracing::Span; use tracing::field; use tracing::info_span; use tracing::trace; +use tracing::trace_span; use tracing::warn; use crate::codex_thread::BackgroundTerminalInfo; @@ -400,6 +401,7 @@ impl Session { task_input, task_cancellation_token.child_token(), ) + .instrument(trace_span!("session_task.run")) .await; let sess = session_ctx.clone_session(); if let Err(err) = sess.flush_rollout().await { diff --git a/codex-rs/core/src/tasks/regular.rs b/codex-rs/core/src/tasks/regular.rs index c01cae9145..c5cc2da31d 100644 --- a/codex-rs/core/src/tasks/regular.rs +++ b/codex-rs/core/src/tasks/regular.rs @@ -45,19 +45,22 @@ impl SessionTask for RegularTask { let run_turn_span = trace_span!("run_turn"); // Regular turns emit `TurnStarted` inline so first-turn lifecycle does // not wait on startup prewarm resolution. - let event = EventMsg::TurnStarted(TurnStartedEvent { - turn_id: ctx.sub_id.clone(), - trace_id: ctx.trace_id.clone(), - started_at: ctx.turn_timing_state.started_at_unix_secs().await, - model_context_window: ctx.model_context_window(), - collaboration_mode_kind: ctx.collaboration_mode.mode, - }); - sess.send_event(ctx.as_ref(), event).await; - sess.set_server_reasoning_included(/*included*/ false).await; - let prewarmed_client_session = match sess - .consume_startup_prewarm_for_regular_turn(&cancellation_token) - .await - { + let prewarmed_client_session = async { + let event = EventMsg::TurnStarted(TurnStartedEvent { + turn_id: ctx.sub_id.clone(), + trace_id: ctx.trace_id.clone(), + started_at: ctx.turn_timing_state.started_at_unix_secs().await, + model_context_window: ctx.model_context_window(), + collaboration_mode_kind: ctx.collaboration_mode.mode, + }); + sess.send_event(ctx.as_ref(), event).await; + sess.set_server_reasoning_included(/*included*/ false).await; + sess.consume_startup_prewarm_for_regular_turn(&cancellation_token) + .await + } + .instrument(trace_span!("regular_task.prepare_run_turn")) + .await; + let prewarmed_client_session = match prewarmed_client_session { SessionStartupPrewarmResolution::Cancelled => return None, SessionStartupPrewarmResolution::Unavailable { .. } => None, SessionStartupPrewarmResolution::Ready(prewarmed_client_session) => { From 44403dd3ddff55753af49dd9eba1351c28dc256a Mon Sep 17 00:00:00 2001 From: Mitsuhiro Kotake Date: Fri, 12 Jun 2026 09:15:04 +0900 Subject: [PATCH 3/4] tui: clear stale hook row after turn completion (#27619) Fixes #27210. ## Why When the app server reports a visible `HookStarted` event for a `PostToolUse` hook but the turn reaches `TurnCompleted` before a matching hook completion event arrives, the TUI can leave the transient `Running PostToolUse hook` row visible after the agent is done. Interrupted and failed turn cleanup already drops transient live hook rows; the normal completion path did not. ## What Changed - Added `ChatWidget::clear_active_hook_cell()` for dropping transient live hook status without writing it to history. - Call that cleanup from normal task completion, while reusing it for the existing start/finalize cleanup paths. - Added `completed_turn_clears_visible_running_hook` snapshot coverage for the reported `PostToolUse` case. ## Tests - `just test -p codex-tui completed_turn_clears_visible_running_hook` - `just test -p codex-tui` (fails on current `main` in unrelated guardian tests: `update_feature_flags_disabling_guardian_clears_review_policy_and_restores_default` and `update_feature_flags_disabling_guardian_clears_manual_review_policy_without_history`) --- codex-rs/tui/src/chatwidget/hook_lifecycle.rs | 7 +++++ ...eted_turn_clears_visible_running_hook.snap | 8 ++++++ .../src/chatwidget/tests/status_and_layout.rs | 26 +++++++++++++++++++ codex-rs/tui/src/chatwidget/turn_runtime.rs | 9 +++---- 4 files changed, 44 insertions(+), 6 deletions(-) create mode 100644 codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__completed_turn_clears_visible_running_hook.snap diff --git a/codex-rs/tui/src/chatwidget/hook_lifecycle.rs b/codex-rs/tui/src/chatwidget/hook_lifecycle.rs index 6545d839a5..04fa87f7de 100644 --- a/codex-rs/tui/src/chatwidget/hook_lifecycle.rs +++ b/codex-rs/tui/src/chatwidget/hook_lifecycle.rs @@ -6,6 +6,13 @@ use super::*; impl ChatWidget { + /// Drop transient live hook status without flushing it into history. + pub(super) fn clear_active_hook_cell(&mut self) { + if self.active_hook_cell.take().is_some() { + self.bump_active_cell_revision(); + } + } + pub(super) fn on_hook_started(&mut self, run: codex_app_server_protocol::HookRunSummary) { self.record_visible_turn_activity(); self.flush_answer_stream_with_separator(); diff --git a/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__completed_turn_clears_visible_running_hook.snap b/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__completed_turn_clears_visible_running_hook.snap new file mode 100644 index 0000000000..e41954e415 --- /dev/null +++ b/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__completed_turn_clears_visible_running_hook.snap @@ -0,0 +1,8 @@ +--- +source: tui/src/chatwidget/tests/status_and_layout.rs +expression: "format!(\"before completion:\\n{before_completion}after completion:\\n{}\",\nactive_hook_blob(&chat))" +--- +before completion: +• Running PostToolUse hook +after completion: + diff --git a/codex-rs/tui/src/chatwidget/tests/status_and_layout.rs b/codex-rs/tui/src/chatwidget/tests/status_and_layout.rs index b7cd476834..62aede458d 100644 --- a/codex-rs/tui/src/chatwidget/tests/status_and_layout.rs +++ b/codex-rs/tui/src/chatwidget/tests/status_and_layout.rs @@ -2173,6 +2173,32 @@ async fn interrupted_turn_clears_visible_running_hook() { ); } +#[tokio::test] +async fn completed_turn_clears_visible_running_hook() { + let (mut chat, _rx, _op_rx) = make_chatwidget_manual(/*model_override*/ None).await; + + handle_hook_started( + &mut chat, + hook_started_run( + "post-tool-use:0:/tmp/hooks.json", + codex_app_server_protocol::HookEventName::PostToolUse, + /*status_message*/ None, + ), + ); + reveal_running_hooks(&mut chat); + let before_completion = active_hook_blob(&chat); + + handle_turn_completed(&mut chat, "turn-1", /*duration_ms*/ None); + + assert_chatwidget_snapshot!( + "completed_turn_clears_visible_running_hook", + format!( + "before completion:\n{before_completion}after completion:\n{}", + active_hook_blob(&chat) + ) + ); +} + #[tokio::test] async fn status_line_fast_mode_renders_on_and_off() { let (mut chat, _rx, _op_rx) = make_chatwidget_manual(/*model_override*/ None).await; diff --git a/codex-rs/tui/src/chatwidget/turn_runtime.rs b/codex-rs/tui/src/chatwidget/turn_runtime.rs index 24674d710e..3e413c82a4 100644 --- a/codex-rs/tui/src/chatwidget/turn_runtime.rs +++ b/codex-rs/tui/src/chatwidget/turn_runtime.rs @@ -59,9 +59,7 @@ impl ChatWidget { self.quit_shortcut_key = None; self.update_task_running_state(); self.status_state.retry_status_header = None; - if self.active_hook_cell.take().is_some() { - self.bump_active_cell_revision(); - } + self.clear_active_hook_cell(); self.status_state.pending_status_indicator_restore = false; self.bottom_pane .set_interrupt_hint_visible(/*visible*/ true); @@ -162,6 +160,7 @@ impl ChatWidget { // Mark task stopped and request redraw now that all content is in history. self.status_state.pending_status_indicator_restore = false; self.input_queue.user_turn_pending_start = false; + self.clear_active_hook_cell(); self.turn_lifecycle.finish(); self.update_task_running_state(); self.running_commands.clear(); @@ -302,9 +301,7 @@ impl ChatWidget { // Turn-scoped hook rows are transient live state; once the turn is over, // do not leave an orphaned running row behind if no matching completion // event arrived before cancellation. - if self.active_hook_cell.take().is_some() { - self.bump_active_cell_revision(); - } + self.clear_active_hook_cell(); // Reset running state and clear streaming buffers. self.input_queue.user_turn_pending_start = false; self.turn_lifecycle.finish(); From 1829ed11228d9ffad6bf7d6db615324126065b23 Mon Sep 17 00:00:00 2001 From: "Adam Perry @ OpenAI" Date: Thu, 11 Jun 2026 17:15:19 -0700 Subject: [PATCH 4/4] Fix image extension PathUri conversion (#27711) ## Why `main` stopped compiling when #27498 passed an `AbsolutePathBuf` to the `ExecutorFileSystem` API migrated to `PathUri` by #27653. ## What Convert referenced image paths to `PathUri` before filesystem reads, declare the internal path-URI dependency, and refresh `Cargo.lock`. --- codex-rs/Cargo.lock | 1 + codex-rs/ext/image-generation/Cargo.toml | 1 + codex-rs/ext/image-generation/src/tool.rs | 9 ++++++++- 3 files changed, 10 insertions(+), 1 deletion(-) diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 8375f2252f..cb9caa4a64 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -3134,6 +3134,7 @@ dependencies = [ "codex-tools", "codex-utils-absolute-path", "codex-utils-image", + "codex-utils-path-uri", "http 1.4.0", "pretty_assertions", "schemars 0.8.22", diff --git a/codex-rs/ext/image-generation/Cargo.toml b/codex-rs/ext/image-generation/Cargo.toml index ef01e2d2be..393709a82d 100644 --- a/codex-rs/ext/image-generation/Cargo.toml +++ b/codex-rs/ext/image-generation/Cargo.toml @@ -23,6 +23,7 @@ codex-protocol = { workspace = true } codex-tools = { workspace = true } codex-utils-absolute-path = { workspace = true } codex-utils-image = { workspace = true } +codex-utils-path-uri = { workspace = true } http = { workspace = true } schemars = { workspace = true } serde = { workspace = true, features = ["derive"] } diff --git a/codex-rs/ext/image-generation/src/tool.rs b/codex-rs/ext/image-generation/src/tool.rs index 0a409f3e1b..2cbb0ab9f3 100644 --- a/codex-rs/ext/image-generation/src/tool.rs +++ b/codex-rs/ext/image-generation/src/tool.rs @@ -33,6 +33,7 @@ use codex_tools::default_namespace_description; use codex_utils_absolute_path::AbsolutePathBuf; use codex_utils_image::PromptImageMode; use codex_utils_image::load_for_prompt_bytes; +use codex_utils_path_uri::PathUri; use schemars::JsonSchema; use schemars::r#gen::SchemaSettings; use serde::Deserialize; @@ -324,9 +325,15 @@ async fn image_url( path: &AbsolutePathBuf, environment: &ToolEnvironment, ) -> Result { + let path_uri = PathUri::from_abs_path(path).map_err(|error| { + FunctionCallError::RespondToModel(format!( + "unable to read referenced image at `{}`: {error}", + path.display() + )) + })?; let bytes = environment .file_system - .read_file(path, Some(&environment.file_system_sandbox_context)) + .read_file(&path_uri, Some(&environment.file_system_sandbox_context)) .await .map_err(|error| { FunctionCallError::RespondToModel(format!(