diff --git a/codex-rs/ext/image-generation/src/extension.rs b/codex-rs/ext/image-generation/src/extension.rs index 8195c08e61..5525c10f5c 100644 --- a/codex-rs/ext/image-generation/src/extension.rs +++ b/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, @@ -29,6 +31,7 @@ type SaveRootResolver = dyn Fn(&Config) -> Option + Send + Sync #[derive(Clone)] struct ImageGenerationExtensionConfig { available: bool, + uses_provider_auth: bool, provider: ModelProviderInfo, save_root: Option, } @@ -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 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::() 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::::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::>(); + + assert_eq!( + tool_names, + vec![ToolName::namespaced( + IMAGE_GEN_NAMESPACE, + IMAGEGEN_TOOL_NAME + )] + ); + } +}