integration tests

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::<ThreadStartResponse>(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::<TurnStartResponse>(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::<ThreadStartResponse>(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::<TurnStartResponse>(turn_resp)?;
+
+    Ok(())
+}
+
 async fn wait_for_image_generation_completed(
-    mcp: &mut McpProcess,
+    mcp: &mut TestAppServer,
 ) -> Result<ItemCompletedNotification> {
     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"

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;