Fix unsupported standalone imagegen exposure

codex-rs/app-server/tests/common/lib.rs

@@ -28,6 +28,7 @@ pub use mock_model_server::create_mock_responses_server_repeating_assistant;
 pub use mock_model_server::create_mock_responses_server_sequence;
 pub use mock_model_server::create_mock_responses_server_sequence_unchecked;
 pub use models_cache::write_mock_provider_models_cache;
+pub use models_cache::write_mock_provider_models_cache_with_auth;
 pub use models_cache::write_models_cache;
 pub use models_cache::write_models_cache_with_models;
 pub use responses::create_apply_patch_sse_response;

codex-rs/app-server/tests/common/models_cache.rs

@@ -1,6 +1,8 @@
 use chrono::DateTime;
 use chrono::Utc;
+use codex_config::types::AuthCredentialsStoreMode;
 use codex_core::test_support::all_model_presets;
+use codex_login::AuthManager;
 use codex_model_provider::model_cache_key_for_provider;
 use codex_model_provider_info::ModelProviderInfo;
 use codex_model_provider_info::OPENAI_PROVIDER_ID;
@@ -135,6 +137,29 @@ pub fn write_mock_provider_models_cache(codex_home: &Path) -> std::io::Result<()
     write_models_cache_for_provider(codex_home, &provider_cache_key)
 }
 
+pub async fn write_mock_provider_models_cache_with_auth(codex_home: &Path) -> std::io::Result<()> {
+    let provider_info = mock_provider_info_from_config(codex_home);
+    let auth_manager = AuthManager::new(
+        codex_home.to_path_buf(),
+        /*enable_codex_api_key_env*/ false,
+        AuthCredentialsStoreMode::File,
+        /*chatgpt_base_url*/ None,
+    )
+    .await;
+    let provider_cache_key =
+        model_cache_key_for_provider("mock_provider", &provider_info, Some(&auth_manager));
+    let mut models = bundled_models_response()
+        .map_err(|err| io::Error::new(io::ErrorKind::InvalidData, err))?
+        .models;
+    if let Some(mut lite_model) = models.iter().find(|model| model.slug == "gpt-5.4").cloned() {
+        lite_model.slug = "imagegen-lite-model".to_string();
+        lite_model.display_name = "Imagegen Lite Model".to_string();
+        lite_model.use_responses_lite = true;
+        models.push(lite_model);
+    }
+    write_models_cache_with_models_for_provider(codex_home, models, &provider_cache_key)
+}
+
 fn mock_provider_cache_key(codex_home: &Path) -> String {
     let provider_info = mock_provider_info_from_config(codex_home);
     model_cache_key_for_provider("mock_provider", &provider_info, /*auth_manager*/ None)

codex-rs/app-server/tests/suite/v2/imagegen_extension.rs

@@ -7,6 +7,7 @@ use app_test_support::ChatGptAuthFixture;
 use app_test_support::TestAppServer;
 use app_test_support::to_response;
 use app_test_support::write_chatgpt_auth;
+use app_test_support::write_mock_provider_models_cache_with_auth;
 use codex_app_server_protocol::ItemCompletedNotification;
 use codex_app_server_protocol::JSONRPCResponse;
 use codex_app_server_protocol::RequestId;
@@ -17,6 +18,7 @@ use codex_app_server_protocol::TurnStartParams;
 use codex_app_server_protocol::TurnStartResponse;
 use codex_app_server_protocol::UserInput as V2UserInput;
 use codex_config::types::AuthCredentialsStoreMode;
+use codex_models_manager::bundled_models_response;
 use core_test_support::responses;
 use pretty_assertions::assert_eq;
 use serde_json::json;
@@ -47,6 +49,7 @@ const DEFAULT_READ_TIMEOUT: Duration = Duration::from_secs(10);
 async fn standalone_image_generation_returns_saved_path_hint_to_model() -> Result<()> {
     let call_id = "image-run-1";
     let server = responses::start_mock_server().await;
+    mount_imagegen_models(&server).await?;
     mount_image_response(&server).await;
 
     let response_mock = responses::mount_sse_sequence(
@@ -81,9 +84,17 @@ async fn standalone_image_generation_returns_saved_path_hint_to_model() -> Resul
         ChatGptAuthFixture::new("access-chatgpt"),
         AuthCredentialsStoreMode::File,
     )?;
+    write_mock_provider_models_cache_with_auth(codex_home.path()).await?;
 
-    let mut mcp =
-        TestAppServer::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),
+            ("CODEWITH_AUTH_PROFILE", None),
+            ("CODEX_AUTH_PROFILE", None),
+        ],
+    )
+    .await?;
     timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??;
     start_image_generation_turn(&mut mcp).await?;
 
@@ -145,6 +156,7 @@ async fn standalone_image_generation_returns_saved_path_hint_to_model() -> Resul
 #[tokio::test]
 async fn standalone_image_generation_is_exposed_in_code_mode_only() -> Result<()> {
     let server = responses::start_mock_server().await;
+    mount_imagegen_models(&server).await?;
     let response_mock = responses::mount_sse_once(
         &server,
         responses::sse(vec![
@@ -165,9 +177,17 @@ async fn standalone_image_generation_is_exposed_in_code_mode_only() -> Result<()
         ChatGptAuthFixture::new("access-chatgpt"),
         AuthCredentialsStoreMode::File,
     )?;
+    write_mock_provider_models_cache_with_auth(codex_home.path()).await?;
 
-    let mut mcp =
-        TestAppServer::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),
+            ("CODEWITH_AUTH_PROFILE", None),
+            ("CODEX_AUTH_PROFILE", None),
+        ],
+    )
+    .await?;
     timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??;
     start_image_generation_turn(&mut mcp).await?;
     timeout(
@@ -190,6 +210,7 @@ async fn standalone_image_generation_is_exposed_in_code_mode_only() -> Result<()
 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_imagegen_models(&server).await?;
     mount_image_response(&server).await;
 
     let response_mock = responses::mount_sse_sequence(
@@ -229,9 +250,17 @@ generatedImage(result);
         ChatGptAuthFixture::new("access-chatgpt"),
         AuthCredentialsStoreMode::File,
     )?;
+    write_mock_provider_models_cache_with_auth(codex_home.path()).await?;
 
-    let mut mcp =
-        TestAppServer::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),
+            ("CODEWITH_AUTH_PROFILE", None),
+            ("CODEX_AUTH_PROFILE", None),
+        ],
+    )
+    .await?;
     timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??;
     start_image_generation_turn(&mut mcp).await?;
     timeout(
@@ -324,30 +353,48 @@ async fn mount_image_response(server: &MockServer) {
         .await;
 }
 
+async fn mount_imagegen_models(server: &MockServer) -> Result<()> {
+    let mut models = bundled_models_response()?;
+    let Some(mut lite_model) = models
+        .models
+        .iter()
+        .find(|model| model.slug == "gpt-5.4")
+        .cloned()
+    else {
+        anyhow::bail!("expected bundled gpt-5.4 model metadata");
+    };
+    lite_model.slug = "imagegen-lite-model".to_string();
+    lite_model.display_name = "Imagegen Lite Model".to_string();
+    lite_model.use_responses_lite = true;
+    models.models.push(lite_model);
+    responses::mount_models_once(server, models).await;
+    Ok(())
+}
+
 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",
+    let (model, code_mode_only) = match mode {
+        ImagegenTestMode::Direct => ("imagegen-lite-model", ""),
+        ImagegenTestMode::CodeModeOnly => ("gpt-5.4", "code_mode_only = true"),
     };
     std::fs::write(
         codex_home.join("config.toml"),
         format!(
             r#"
-model = "mock-model"
+model = "{model}"
 approval_policy = "never"
 sandbox_mode = "read-only"
-model_provider = "openai-custom"
+model_provider = "mock_provider"
 chatgpt_base_url = "{server_uri}"
 
 [features]
 imagegenext = true
 {code_mode_only}
 
-[model_providers.openai-custom]
+[model_providers.mock_provider]
 name = "OpenAI"
 base_url = "{server_uri}/api/codex"
 wire_api = "responses"

codex-rs/core/src/tools/spec_plan.rs

@@ -208,8 +208,7 @@ fn build_model_visible_specs_and_registry(
             continue;
         }
         let exposure = runtime.exposure();
-        if exposure.is_direct() && !is_hidden_by_code_mode_only(turn_context, &tool_name, exposure)
-        {
+        if exposure.is_direct() && !is_hidden_by_code_mode(turn_context, &tool_name, exposure) {
             let spec = runtime.spec();
             specs.push(spec_for_model_request(
                 turn_context,
@@ -364,27 +363,51 @@ fn image_generation_runtime_enabled(turn_context: &TurnContext) -> bool {
 }
 
 fn standalone_image_generation_model_visible(turn_context: &TurnContext) -> bool {
-    if !image_generation_runtime_enabled(turn_context) || !namespace_tools_enabled(turn_context) {
-        return false;
-    }
-
-    if turn_context.model_info.use_responses_lite {
-        return true;
-    }
-
-    turn_context.features.get().enabled(Feature::ImageGenExt)
+    standalone_image_generation_runtime_available(turn_context)
+        && turn_context.model_info.use_responses_lite
 }
 
 fn standalone_image_generation_available(
     turn_context: &TurnContext,
     extension_tools: &[Arc<dyn ToolExecutor<ExtensionToolCall>>],
 ) -> bool {
-    standalone_image_generation_model_visible(turn_context)
+    if image_generation_excluded_from_code_mode(turn_context) {
+        return false;
+    }
+
+    standalone_image_generation_extension_registerable(turn_context)
         && extension_tools.iter().any(|executor| {
             executor.tool_name() == ToolName::namespaced(IMAGE_GEN_NAMESPACE, IMAGEGEN_TOOL_NAME)
         })
 }
 
+fn standalone_image_generation_runtime_available(turn_context: &TurnContext) -> bool {
+    image_generation_runtime_enabled(turn_context)
+        && namespace_tools_enabled(turn_context)
+        && turn_context.features.get().enabled(Feature::ImageGenExt)
+}
+
+fn standalone_image_generation_extension_registerable(turn_context: &TurnContext) -> bool {
+    standalone_image_generation_model_visible(turn_context)
+        || (standalone_image_generation_runtime_available(turn_context)
+            && matches!(
+                turn_context.tool_mode,
+                ToolMode::CodeMode | ToolMode::CodeModeOnly
+            ))
+}
+
+fn image_generation_excluded_from_code_mode(turn_context: &TurnContext) -> bool {
+    matches!(
+        turn_context.tool_mode,
+        ToolMode::CodeMode | ToolMode::CodeModeOnly
+    ) && turn_context
+        .config
+        .code_mode
+        .excluded_tool_namespaces
+        .iter()
+        .any(|namespace| namespace == IMAGE_GEN_NAMESPACE)
+}
+
 fn wait_agent_timeout_options(turn_context: &TurnContext) -> WaitAgentTimeoutOptions {
     if multi_agent_v2_enabled(turn_context) {
         return WaitAgentTimeoutOptions {
@@ -423,16 +446,25 @@ fn agent_type_description(
     }
 }
 
-fn is_hidden_by_code_mode_only(
+fn is_hidden_by_code_mode(
     turn_context: &TurnContext,
     tool_name: &ToolName,
     exposure: ToolExposure,
 ) -> bool {
-    turn_context.tool_mode == ToolMode::CodeModeOnly
-        && exposure != ToolExposure::DirectModelOnly
-        && codex_code_mode::is_code_mode_nested_tool(&codex_tools::code_mode_name_for_tool_name(
-            tool_name,
-        ))
+    if exposure == ToolExposure::DirectModelOnly {
+        return false;
+    }
+
+    match turn_context.tool_mode {
+        ToolMode::CodeMode => {
+            tool_name == &ToolName::namespaced(IMAGE_GEN_NAMESPACE, IMAGEGEN_TOOL_NAME)
+                && standalone_image_generation_extension_registerable(turn_context)
+        }
+        ToolMode::CodeModeOnly => codex_code_mode::is_code_mode_nested_tool(
+            &codex_tools::code_mode_name_for_tool_name(tool_name),
+        ),
+        ToolMode::Direct => false,
+    }
 }
 
 fn is_excluded_from_code_mode(turn_context: &TurnContext, tool_name: &ToolName) -> bool {
@@ -936,7 +968,8 @@ fn append_extension_tool_executors(
             continue;
         }
         if tool_name == ToolName::namespaced(IMAGE_GEN_NAMESPACE, IMAGEGEN_TOOL_NAME)
-            && !standalone_image_generation_model_visible(turn_context)
+            && (!standalone_image_generation_extension_registerable(turn_context)
+                || image_generation_excluded_from_code_mode(turn_context))
         {
             continue;
         }