From 1bfc3d9773e7d61ed2d15be991924194e5b0f5b9 Mon Sep 17 00:00:00 2001 From: starr-openai Date: Thu, 7 May 2026 18:29:03 -0700 Subject: [PATCH] Route view_image through selected environments Route view_image through selected environments so image reads use the selected turn environment and cwd, with schema exposure limited to multi-environment toolsets.\n\nCo-authored-by: Codex --- .../core/src/tools/handlers/view_image.rs | 38 ++-- .../src/tools/handlers/view_image_spec.rs | 10 + codex-rs/core/src/tools/spec_plan.rs | 3 + codex-rs/core/src/tools/spec_plan_tests.rs | 43 ++++ codex-rs/core/tests/suite/view_image.rs | 186 ++++++++++++++++++ 5 files changed, 266 insertions(+), 14 deletions(-) diff --git a/codex-rs/core/src/tools/handlers/view_image.rs b/codex-rs/core/src/tools/handlers/view_image.rs index d52a84415..6e4a25c0a 100644 --- a/codex-rs/core/src/tools/handlers/view_image.rs +++ b/codex-rs/core/src/tools/handlers/view_image.rs @@ -17,6 +17,7 @@ use crate::tools::context::ToolInvocation; use crate::tools::context::ToolOutput; use crate::tools::context::ToolPayload; use crate::tools::handlers::parse_arguments; +use crate::tools::handlers::resolve_tool_environment; use crate::tools::handlers::view_image_spec::ViewImageToolOptions; use crate::tools::handlers::view_image_spec::create_view_image_tool; use crate::tools::registry::ToolHandler; @@ -33,6 +34,7 @@ impl Default for ViewImageHandler { Self { options: ViewImageToolOptions { can_request_original_image_detail: false, + include_environment_id: false, }, } } @@ -50,6 +52,8 @@ const VIEW_IMAGE_UNSUPPORTED_MESSAGE: &str = #[derive(Deserialize)] struct ViewImageArgs { path: String, + #[serde(default)] + environment_id: Option, detail: Option, } @@ -106,12 +110,16 @@ impl ToolHandler for ViewImageHandler { } }; - let args: ViewImageArgs = parse_arguments(&arguments)?; + let ViewImageArgs { + path, + environment_id, + detail, + } = parse_arguments(&arguments)?; // `view_image` accepts only its documented detail values: omit // `detail` for the default path or set it to `original`. // Other string values remain invalid rather than being silently // reinterpreted. - let detail = match args.detail.as_deref() { + let detail = match detail.as_deref() { None => None, Some("original") => Some(ViewImageDetail::Original), Some(detail) => { @@ -121,20 +129,24 @@ impl ToolHandler for ViewImageHandler { } }; - let abs_path = turn.resolve_path(Some(args.path)); - let Some(environment) = turn.environments.primary() else { + let Some(turn_environment) = + resolve_tool_environment(turn.as_ref(), environment_id.as_deref())? + else { return Err(FunctionCallError::RespondToModel( "view_image is unavailable in this session".to_string(), )); }; - let sandbox = environment - .environment - .is_remote() - .then(|| turn.file_system_sandbox_context(/*additional_permissions*/ None)); + let cwd = turn_environment.cwd.clone(); + let abs_path = cwd.join(path); + let sandbox = turn_environment.environment.is_remote().then(|| { + let mut sandbox = + turn.file_system_sandbox_context(/*additional_permissions*/ None); + sandbox.cwd = Some(cwd.clone()); + sandbox + }); + let fs = turn_environment.environment.get_filesystem(); - let metadata = environment - .environment - .get_filesystem() + let metadata = fs .get_metadata(&abs_path, sandbox.as_ref()) .await .map_err(|error| { @@ -150,9 +162,7 @@ impl ToolHandler for ViewImageHandler { abs_path.display() ))); } - let file_bytes = environment - .environment - .get_filesystem() + let file_bytes = fs .read_file(&abs_path, sandbox.as_ref()) .await .map_err(|error| { diff --git a/codex-rs/core/src/tools/handlers/view_image_spec.rs b/codex-rs/core/src/tools/handlers/view_image_spec.rs index 28953cc9b..7d1422a03 100644 --- a/codex-rs/core/src/tools/handlers/view_image_spec.rs +++ b/codex-rs/core/src/tools/handlers/view_image_spec.rs @@ -9,6 +9,7 @@ use std::collections::BTreeMap; #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub struct ViewImageToolOptions { pub can_request_original_image_detail: bool, + pub include_environment_id: bool, } pub fn create_view_image_tool(options: ViewImageToolOptions) -> ToolSpec { @@ -24,6 +25,15 @@ pub fn create_view_image_tool(options: ViewImageToolOptions) -> ToolSpec { )), ); } + if options.include_environment_id { + properties.insert( + "environment_id".to_string(), + JsonSchema::string(Some( + "Optional selected environment id to target. Omit this to use the primary environment." + .to_string(), + )), + ); + } ToolSpec::Function(ResponsesApiTool { name: VIEW_IMAGE_TOOL_NAME.to_string(), diff --git a/codex-rs/core/src/tools/spec_plan.rs b/codex-rs/core/src/tools/spec_plan.rs index 644509961..ddad2cbf4 100644 --- a/codex-rs/core/src/tools/spec_plan.rs +++ b/codex-rs/core/src/tools/spec_plan.rs @@ -289,8 +289,11 @@ pub fn build_tool_registry_builder( } if config.environment_mode.has_environment() { + let include_environment_id = + matches!(config.environment_mode, ToolEnvironmentMode::Multiple); builder.register_handler(Arc::new(ViewImageHandler::new(ViewImageToolOptions { can_request_original_image_detail: config.can_request_original_image_detail, + include_environment_id, }))); } diff --git a/codex-rs/core/src/tools/spec_plan_tests.rs b/codex-rs/core/src/tools/spec_plan_tests.rs index a77c5c2e1..3fa8797c5 100644 --- a/codex-rs/core/src/tools/spec_plan_tests.rs +++ b/codex-rs/core/src/tools/spec_plan_tests.rs @@ -129,6 +129,7 @@ fn test_full_toolset_specs_for_gpt5_codex_unified_exec_web_search() { create_image_generation_tool("png"), create_view_image_tool(ViewImageToolOptions { can_request_original_image_detail: config.can_request_original_image_detail, + include_environment_id: false, }), ] { expected.insert(spec.name().to_string(), spec); @@ -628,6 +629,48 @@ fn disabled_environment_omits_environment_backed_tools() { assert_lacks_tool_name(&tools, VIEW_IMAGE_TOOL_NAME); } +#[test] +fn view_image_spec_includes_environment_id_only_for_multiple_selected_environments() { + let model_info = model_info(); + let available_models = Vec::new(); + let tools_config = ToolsConfig::new(&ToolsConfigParams { + model_info: &model_info, + available_models: &available_models, + features: &Features::with_defaults(), + image_generation_tool_auth_allowed: true, + web_search_mode: Some(WebSearchMode::Cached), + session_source: SessionSource::Cli, + permission_profile: &PermissionProfile::Disabled, + windows_sandbox_level: WindowsSandboxLevel::Disabled, + }); + + let (single_environment_tools, _) = build_specs( + &tools_config, + /*mcp_tools*/ None, + /*deferred_mcp_tools*/ None, + &[], + ); + assert_process_tool_environment_id( + &single_environment_tools, + VIEW_IMAGE_TOOL_NAME, + /*expected_present*/ false, + ); + + let multi_environment_config = + tools_config.with_environment_mode(ToolEnvironmentMode::Multiple); + let (multi_environment_tools, _) = build_specs( + &multi_environment_config, + /*mcp_tools*/ None, + /*deferred_mcp_tools*/ None, + &[], + ); + assert_process_tool_environment_id( + &multi_environment_tools, + VIEW_IMAGE_TOOL_NAME, + /*expected_present*/ true, + ); +} + #[test] fn test_build_specs_agent_job_worker_tools_enabled() { let model_info = model_info(); diff --git a/codex-rs/core/tests/suite/view_image.rs b/codex-rs/core/tests/suite/view_image.rs index 19c2e6673..cf06f9479 100644 --- a/codex-rs/core/tests/suite/view_image.rs +++ b/codex-rs/core/tests/suite/view_image.rs @@ -4,6 +4,9 @@ use anyhow::Context; use base64::Engine; use base64::engine::general_purpose::STANDARD as BASE64_STANDARD; use codex_exec_server::CreateDirectoryOptions; +use codex_exec_server::LOCAL_ENVIRONMENT_ID; +use codex_exec_server::REMOTE_ENVIRONMENT_ID; +use codex_exec_server::RemoveOptions; use codex_login::CodexAuth; use codex_protocol::config_types::ReasoningSummary; use codex_protocol::models::PermissionProfile; @@ -18,13 +21,18 @@ use codex_protocol::openai_models::TruncationPolicyConfig; use codex_protocol::protocol::AskForApproval; use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::Op; +use codex_protocol::protocol::TurnEnvironmentSelection; use codex_protocol::user_input::UserInput; +use core_test_support::PathBufExt; +use core_test_support::PathExt; +use core_test_support::get_remote_test_env; use core_test_support::responses; use core_test_support::responses::ev_assistant_message; use core_test_support::responses::ev_completed; use core_test_support::responses::ev_function_call; use core_test_support::responses::ev_response_created; use core_test_support::responses::mount_models_once; +use core_test_support::responses::mount_sse_sequence; use core_test_support::responses::sse; use core_test_support::responses::start_mock_server; use core_test_support::skip_if_no_network; @@ -39,8 +47,13 @@ use image::Rgba; use image::load_from_memory; use pretty_assertions::assert_eq; use serde_json::Value; +use serde_json::json; +use std::fs; use std::io::Cursor; use std::path::PathBuf; +use std::time::SystemTime; +use std::time::UNIX_EPOCH; +use tempfile::TempDir; use tokio::time::Duration; use wiremock::BodyPrintLimit; use wiremock::MockServer; @@ -390,6 +403,179 @@ async fn view_image_tool_attaches_local_image() -> anyhow::Result<()> { Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn view_image_routes_to_selected_local_environment() -> anyhow::Result<()> { + skip_if_no_network!(Ok(())); + + let server = start_mock_server().await; + let mut builder = test_codex(); + let test = builder.build(&server).await?; + write_workspace_file( + &test, + "local.png", + png_bytes(/*width*/ 1, /*height*/ 1, [0, 255, 0, 255])?, + ) + .await?; + let call_id = "call-view-image-local-env"; + let response_mock = mount_sse_sequence( + &server, + vec![ + sse(vec![ + ev_response_created("resp-1"), + ev_function_call( + call_id, + "view_image", + &json!({ + "path": "local.png", + "environment_id": LOCAL_ENVIRONMENT_ID, + }) + .to_string(), + ), + ev_completed("resp-1"), + ]), + sse(vec![ + ev_response_created("resp-2"), + ev_assistant_message("msg-1", "done"), + ev_completed("resp-2"), + ]), + ], + ) + .await; + + test.submit_turn_with_environments( + "route local view image", + Some(vec![TurnEnvironmentSelection { + environment_id: LOCAL_ENVIRONMENT_ID.to_string(), + cwd: test.config.cwd.clone(), + }]), + ) + .await?; + + let output = response_mock + .last_request() + .context("missing request containing local view_image output")? + .function_call_output(call_id); + let output_items = output + .get("output") + .and_then(Value::as_array) + .context("view_image output should be content items")?; + assert_eq!(output_items.len(), 1); + let image_url = output_items[0] + .get("image_url") + .and_then(Value::as_str) + .context("view_image output should include image_url")?; + assert!( + image_url.starts_with("data:image/png;base64,"), + "unexpected image_url: {image_url}", + ); + + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn view_image_routes_to_selected_remote_environment() -> anyhow::Result<()> { + skip_if_no_network!(Ok(())); + let Some(_remote_env) = get_remote_test_env() else { + return Ok(()); + }; + + let server = start_mock_server().await; + let mut builder = test_codex(); + let test = builder.build_remote_aware(&server).await?; + let local_cwd = TempDir::new()?; + fs::write(local_cwd.path().join("remote.png"), b"not a remote image")?; + let local_selection = TurnEnvironmentSelection { + environment_id: LOCAL_ENVIRONMENT_ID.to_string(), + cwd: local_cwd.path().abs(), + }; + let remote_cwd = PathBuf::from(format!( + "/tmp/codex-view-image-routing-{}", + SystemTime::now().duration_since(UNIX_EPOCH)?.as_millis() + )) + .abs(); + let image_path = remote_cwd.join("remote.png"); + test.fs() + .create_directory( + &remote_cwd, + CreateDirectoryOptions { recursive: true }, + /*sandbox*/ None, + ) + .await?; + let png = BASE64_STANDARD.decode( + "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mP8/x8AAwMCAO+/p9sAAAAASUVORK5CYII=", + )?; + test.fs() + .write_file(&image_path, png, /*sandbox*/ None) + .await?; + let remote_selection = TurnEnvironmentSelection { + environment_id: REMOTE_ENVIRONMENT_ID.to_string(), + cwd: remote_cwd.clone(), + }; + let call_id = "call-view-image-multi-env"; + let response_mock = mount_sse_sequence( + &server, + vec![ + sse(vec![ + ev_response_created("resp-1"), + ev_function_call( + call_id, + "view_image", + &json!({ + "path": "remote.png", + "environment_id": REMOTE_ENVIRONMENT_ID, + }) + .to_string(), + ), + ev_completed("resp-1"), + ]), + sse(vec![ + ev_response_created("resp-2"), + ev_assistant_message("msg-1", "done"), + ev_completed("resp-2"), + ]), + ], + ) + .await; + + test.submit_turn_with_environments( + "route view image", + Some(vec![local_selection, remote_selection]), + ) + .await?; + + let output = response_mock + .last_request() + .context("missing request containing view_image output")? + .function_call_output(call_id) + .clone(); + let output_items = output + .get("output") + .and_then(Value::as_array) + .context("view_image output should be content items")?; + assert_eq!(output_items.len(), 1); + let image_url = output_items[0] + .get("image_url") + .and_then(Value::as_str) + .context("view_image output should include image_url")?; + assert!( + image_url.starts_with("data:image/png;base64,"), + "unexpected image_url: {image_url}", + ); + + test.fs() + .remove( + &remote_cwd, + RemoveOptions { + recursive: true, + force: true, + }, + /*sandbox*/ None, + ) + .await?; + + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn view_image_tool_can_preserve_original_resolution_when_requested_on_gpt5_3_codex() -> anyhow::Result<()> {