Allow imagegen extension with provider auth

codex-rs/ext/image-generation/src/extension.rs

@@ -18,6 +18,8 @@ use codex_utils_absolute_path::AbsolutePathBuf;
 use crate::backend::CodexImagesBackend;
 use crate::tool::ImageGenerationTool;
 
+const ACTOR_AUTHORIZATION_HEADER: &str = "x-openai-actor-authorization";
+
 #[derive(Clone)]
 struct ImageGenerationExtension {
     auth_manager: Arc<AuthManager>,
@@ -29,6 +31,7 @@ type SaveRootResolver = dyn Fn(&Config) -> Option<AbsolutePathBuf> + Send + Sync
 #[derive(Clone)]
 struct ImageGenerationExtensionConfig {
     available: bool,
+    uses_provider_auth: bool,
     provider: ModelProviderInfo,
     save_root: Option<AbsolutePathBuf>,
 }
@@ -36,15 +39,26 @@ struct ImageGenerationExtensionConfig {
 impl ImageGenerationExtensionConfig {
     /// Resolves whether standalone image generation should be available for a thread.
     fn from_config(config: &Config, resolve_save_root: &SaveRootResolver) -> Self {
+        let uses_provider_auth = provider_uses_actor_authorization(&config.model_provider);
         Self {
             // Core selects this executor per turn using the feature flag or model metadata.
-            available: config.model_provider.is_openai(),
+            available: config.model_provider.is_openai() || uses_provider_auth,
+            uses_provider_auth,
             provider: config.model_provider.clone(),
             save_root: resolve_save_root(config),
         }
     }
 }
 
+fn provider_uses_actor_authorization(provider: &ModelProviderInfo) -> bool {
+    !provider.requires_openai_auth
+        && provider.http_headers.as_ref().is_some_and(|headers| {
+            headers.iter().any(|(name, value)| {
+                name.eq_ignore_ascii_case(ACTOR_AUTHORIZATION_HEADER) && !value.trim().is_empty()
+            })
+        })
+}
+
 impl ThreadLifecycleContributor<Config> for ImageGenerationExtension {
     /// Seeds image-generation availability when a thread begins.
     fn on_thread_start<'a>(
@@ -88,7 +102,9 @@ impl ToolContributor for ImageGenerationExtension {
         let Some(config) = thread_store.get::<ImageGenerationExtensionConfig>() else {
             return Vec::new();
         };
-        if !config.available || !self.auth_manager.current_auth_uses_codex_backend() {
+        if !config.available
+            || (!config.uses_provider_auth && !self.auth_manager.current_auth_uses_codex_backend())
+        {
             return Vec::new();
         }
 
@@ -117,3 +133,80 @@ pub fn install(
     registry.config_contributor(extension.clone());
     registry.tool_contributor(extension);
 }
+
+#[cfg(test)]
+mod tests {
+    use std::collections::HashMap;
+
+    use codex_extension_api::ExtensionData;
+    use codex_extension_api::ToolName;
+    use codex_login::CodexAuth;
+
+    use super::*;
+    use crate::IMAGE_GEN_NAMESPACE;
+    use crate::IMAGEGEN_TOOL_NAME;
+
+    #[test]
+    fn actor_authorization_requires_a_nonempty_header_on_a_provider_auth_path() {
+        let provider = |requires_openai_auth, value: Option<&str>| {
+            let mut provider = ModelProviderInfo::default();
+            provider.requires_openai_auth = requires_openai_auth;
+            provider.http_headers = value.map(|value| {
+                HashMap::from([(ACTOR_AUTHORIZATION_HEADER.to_string(), value.to_string())])
+            });
+            provider
+        };
+
+        assert!(provider_uses_actor_authorization(&provider(
+            false,
+            Some("actor-biscuit")
+        )));
+        assert!(!provider_uses_actor_authorization(&provider(
+            false,
+            Some("  ")
+        )));
+        assert!(!provider_uses_actor_authorization(&provider(
+            true,
+            Some("actor-biscuit")
+        )));
+    }
+
+    #[test]
+    fn installed_extension_contributes_imagegen_with_provider_auth() {
+        let mut builder = ExtensionRegistryBuilder::<Config>::new();
+        install(
+            &mut builder,
+            AuthManager::from_auth_for_testing(CodexAuth::from_api_key("dummy")),
+            |_| None,
+        );
+        let registry = builder.build();
+        let session_store = ExtensionData::new("session");
+        let thread_store = ExtensionData::new("11111111-1111-4111-8111-111111111111");
+        let mut provider = ModelProviderInfo::default();
+        provider.http_headers = Some(HashMap::from([(
+            ACTOR_AUTHORIZATION_HEADER.to_string(),
+            "actor-biscuit".to_string(),
+        )]));
+        thread_store.insert(ImageGenerationExtensionConfig {
+            available: true,
+            uses_provider_auth: true,
+            provider,
+            save_root: None,
+        });
+
+        let tool_names = registry
+            .tool_contributors()
+            .iter()
+            .flat_map(|contributor| contributor.tools(&session_store, &thread_store))
+            .map(|tool| tool.tool_name())
+            .collect::<Vec<_>>();
+
+        assert_eq!(
+            tool_names,
+            vec![ToolName::namespaced(
+                IMAGE_GEN_NAMESPACE,
+                IMAGEGEN_TOOL_NAME
+            )]
+        );
+    }
+}