diff --git a/codex-rs/app-server/tests/suite/v2/image_generation.rs b/codex-rs/app-server/tests/suite/v2/imagegen_extension.rs similarity index 69% rename from codex-rs/app-server/tests/suite/v2/image_generation.rs rename to codex-rs/app-server/tests/suite/v2/imagegen_extension.rs index 4752f1497e..5113b9fe66 100644 --- a/codex-rs/app-server/tests/suite/v2/image_generation.rs +++ b/codex-rs/app-server/tests/suite/v2/imagegen_extension.rs @@ -4,7 +4,7 @@ use std::time::Duration; use anyhow::Context; use anyhow::Result; use app_test_support::ChatGptAuthFixture; -use app_test_support::McpProcess; +use app_test_support::TestAppServer; use app_test_support::to_response; use app_test_support::write_chatgpt_auth; use codex_app_server_protocol::ItemCompletedNotification; @@ -30,6 +30,12 @@ use wiremock::matchers::path; const RESULT: &str = "cG5n"; +#[derive(Clone, Copy)] +enum ImagegenTestMode { + Direct, + CodeModeOnly, +} + // macOS and Windows Bazel CI can spend tens of seconds starting app-server // subprocesses or processing test RPCs under load. #[cfg(any(target_os = "macos", windows))] @@ -69,43 +75,17 @@ async fn standalone_image_generation_persists_image_and_returns_it_to_model() -> .await; let codex_home = TempDir::new()?; - create_config_toml(codex_home.path(), &server.uri())?; + create_config_toml(codex_home.path(), &server.uri(), ImagegenTestMode::Direct)?; write_chatgpt_auth( codex_home.path(), ChatGptAuthFixture::new("access-chatgpt"), AuthCredentialsStoreMode::File, )?; - let mut mcp = McpProcess::new_with_env(codex_home.path(), &[("OPENAI_API_KEY", None)]).await?; + let mut mcp = + TestAppServer::new_with_env(codex_home.path(), &[("OPENAI_API_KEY", None)]).await?; timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; - - let thread_req = mcp - .send_thread_start_request(ThreadStartParams::default()) - .await?; - let thread_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(thread_req)), - ) - .await??; - let ThreadStartResponse { thread, .. } = to_response::(thread_resp)?; - - let turn_req = mcp - .send_turn_start_request(TurnStartParams { - thread_id: thread.id, - client_user_message_id: None, - input: vec![V2UserInput::Text { - text: "Generate an image".to_string(), - text_elements: Vec::new(), - }], - ..Default::default() - }) - .await?; - let turn_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(turn_req)), - ) - .await??; - let _turn: TurnStartResponse = to_response::(turn_resp)?; + start_image_generation_turn(&mut mcp).await?; let completed = timeout( DEFAULT_READ_TIMEOUT, @@ -156,8 +136,111 @@ async fn standalone_image_generation_persists_image_and_returns_it_to_model() -> Ok(()) } +#[tokio::test] +async fn standalone_image_generation_is_callable_from_code_mode_only() -> Result<()> { + let call_id = "code-mode-image-run-1"; + let server = responses::start_mock_server().await; + mount_image_response(&server).await; + + let response_mock = responses::mount_sse_sequence( + &server, + vec![ + responses::sse(vec![ + responses::ev_response_created("resp-1"), + responses::ev_custom_tool_call( + call_id, + "exec", + r#" +const result = await tools.image_gen__imagegen({ + action: "generate", + prompt: "paint a blue whale", +}); +image(result); +"#, + ), + responses::ev_completed("resp-1"), + ]), + responses::sse(vec![ + responses::ev_assistant_message("msg-1", "Done"), + responses::ev_completed("resp-2"), + ]), + ], + ) + .await; + + let codex_home = TempDir::new()?; + create_config_toml( + codex_home.path(), + &server.uri(), + ImagegenTestMode::CodeModeOnly, + )?; + write_chatgpt_auth( + codex_home.path(), + ChatGptAuthFixture::new("access-chatgpt"), + AuthCredentialsStoreMode::File, + )?; + + let mut mcp = + TestAppServer::new_with_env(codex_home.path(), &[("OPENAI_API_KEY", None)]).await?; + timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; + start_image_generation_turn(&mut mcp).await?; + timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_notification_message("turn/completed"), + ) + .await??; + + let requests = response_mock.requests(); + assert_eq!(requests.len(), 2); + assert!(requests[0].body_contains_text("image_gen__imagegen")); + let output = requests[1].custom_tool_call_output(call_id); + assert_eq!( + output["output"][1], + json!({ + "type": "input_image", + "image_url": format!("data:image/png;base64,{RESULT}"), + "detail": "high", + }) + ); + assert_eq!(output["output"].as_array().map(Vec::len), Some(2)); + + Ok(()) +} + +async fn start_image_generation_turn(mcp: &mut TestAppServer) -> Result<()> { + let thread_req = mcp + .send_thread_start_request(ThreadStartParams::default()) + .await?; + let thread_resp: JSONRPCResponse = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(thread_req)), + ) + .await??; + let ThreadStartResponse { thread, .. } = to_response::(thread_resp)?; + + let turn_req = mcp + .send_turn_start_request(TurnStartParams { + thread_id: thread.id, + client_user_message_id: None, + input: vec![V2UserInput::Text { + text: "Generate an image".to_string(), + text_elements: Vec::new(), + }], + ..Default::default() + }) + .await?; + let turn_resp: JSONRPCResponse = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(turn_req)), + ) + .await??; + let _turn: TurnStartResponse = to_response::(turn_resp)?; + + Ok(()) +} + async fn wait_for_image_generation_completed( - mcp: &mut McpProcess, + mcp: &mut TestAppServer, ) -> Result { loop { let notification = mcp @@ -186,7 +269,15 @@ async fn mount_image_response(server: &MockServer) { .await; } -fn create_config_toml(codex_home: &Path, server_uri: &str) -> std::io::Result<()> { +fn create_config_toml( + codex_home: &Path, + server_uri: &str, + mode: ImagegenTestMode, +) -> std::io::Result<()> { + let code_mode_only = match mode { + ImagegenTestMode::Direct => "", + ImagegenTestMode::CodeModeOnly => "code_mode_only = true", + }; std::fs::write( codex_home.join("config.toml"), format!( @@ -199,6 +290,7 @@ chatgpt_base_url = "{server_uri}" [features] imagegenext = true +{code_mode_only} [model_providers.openai-custom] name = "OpenAI" diff --git a/codex-rs/app-server/tests/suite/v2/mod.rs b/codex-rs/app-server/tests/suite/v2/mod.rs index 950cccebb5..9aef40aade 100644 --- a/codex-rs/app-server/tests/suite/v2/mod.rs +++ b/codex-rs/app-server/tests/suite/v2/mod.rs @@ -17,7 +17,7 @@ mod experimental_feature_list; mod external_agent_config; mod fs; mod hooks_list; -mod image_generation; +mod imagegen_extension; mod initialize; mod marketplace_add; mod marketplace_remove;