Codex Handbook
ext/image-generation/src/extension.rs 103 lines
use std::sync::Arc;use codex_core::config::Config;use codex_extension_api::ConfigContributor;use codex_extension_api::ExtensionData;use codex_extension_api::ExtensionFuture;use codex_extension_api::ExtensionRegistryBuilder;use codex_extension_api::ThreadLifecycleContributor;use codex_extension_api::ThreadStartInput;use codex_extension_api::ToolCall;use codex_extension_api::ToolContributor;use codex_extension_api::ToolExecutor;use codex_login::AuthManager;use codex_model_provider::create_model_provider;use codex_model_provider_info::ModelProviderInfo;use codex_utils_absolute_path::AbsolutePathBuf;use crate::backend::CodexImagesBackend;use crate::tool::ImageGenerationTool;#[derive(Clone)]struct ImageGenerationExtension {    auth_manager: Arc<AuthManager>,}#[derive(Clone)]struct ImageGenerationExtensionConfig {    available: bool,    provider: ModelProviderInfo,    codex_home: AbsolutePathBuf,}impl From<&Config> for ImageGenerationExtensionConfig {    /// Resolves whether standalone image generation should be available for a thread.    fn from(config: &Config) -> Self {        Self {            // Core selects this executor per turn using the feature flag or model metadata.            available: config.model_provider.is_openai(),            provider: config.model_provider.clone(),            codex_home: config.codex_home.clone(),        }    }}impl ThreadLifecycleContributor<Config> for ImageGenerationExtension {    /// Seeds image-generation availability when a thread begins.    fn on_thread_start<'a>(        &'a self,        input: ThreadStartInput<'a, Config>,    ) -> ExtensionFuture<'a, ()> {        Box::pin(async move {            input                .thread_store                .insert(ImageGenerationExtensionConfig::from(input.config));        })    }}impl ConfigContributor<Config> for ImageGenerationExtension {    /// Refreshes image-generation availability after thread configuration changes.    fn on_config_changed(        &self,        _session_store: &ExtensionData,        thread_store: &ExtensionData,        _previous_config: &Config,        new_config: &Config,    ) {        thread_store.insert(ImageGenerationExtensionConfig::from(new_config));    }}impl ToolContributor for ImageGenerationExtension {    /// Creates the image-generation tool exposed by this installed extension.    fn tools(        &self,        _session_store: &ExtensionData,        thread_store: &ExtensionData,    ) -> Vec<Arc<dyn ToolExecutor<ToolCall>>> {        let Some(config) = thread_store.get::<ImageGenerationExtensionConfig>() else {            return Vec::new();        };        if !config.available || !self.auth_manager.current_auth_uses_codex_backend() {            return Vec::new();        }        vec![Arc::new(ImageGenerationTool::new(            CodexImagesBackend::new(create_model_provider(                config.provider.clone(),                Some(self.auth_manager.clone()),            )),            config.codex_home.clone(),            thread_store.level_id().to_string(),        ))]    }}/// Installs the standalone image-generation extension contributors.pub fn install(registry: &mut ExtensionRegistryBuilder<Config>, auth_manager: Arc<AuthManager>) {    let extension = Arc::new(ImageGenerationExtension { auth_manager });    registry.thread_lifecycle_contributor(extension.clone());    registry.config_contributor(extension.clone());    registry.tool_contributor(extension);}