codex_image_generation_extension/
extension.rs

1use std::sync::Arc;
2
3use codex_core::config::Config;
4use codex_extension_api::ConfigContributor;
5use codex_extension_api::ExtensionData;
6use codex_extension_api::ExtensionFuture;
7use codex_extension_api::ExtensionRegistryBuilder;
8use codex_extension_api::ThreadLifecycleContributor;
9use codex_extension_api::ThreadOriginator;
10use codex_extension_api::ThreadStartInput;
11use codex_extension_api::ToolCall;
12use codex_extension_api::ToolContributor;
13use codex_extension_api::ToolExecutor;
14use codex_login::AuthManager;
15use codex_model_provider::create_model_provider;
16use codex_model_provider_info::ModelProviderInfo;
17use codex_utils_absolute_path::AbsolutePathBuf;
18
19use crate::backend::CodexImagesBackend;
20use crate::tool::ImageGenerationTool;
21
22#[derive(Clone)]
23struct ImageGenerationExtension {
24    auth_manager: Arc<AuthManager>,
25    resolve_save_root: Arc<SaveRootResolver>,
26}
27
28type SaveRootResolver = dyn Fn(&Config) -> Option<AbsolutePathBuf> + Send + Sync;
29
30#[derive(Clone)]
31struct ImageGenerationExtensionConfig {
32    available: bool,
33    provider: ModelProviderInfo,
34    save_root: Option<AbsolutePathBuf>,
35}
36
37impl ImageGenerationExtensionConfig {
38    /// Resolves the image provider and save root for a thread.
39    fn from_config(config: &Config, resolve_save_root: &SaveRootResolver) -> Self {
40        Self {
41            available: config.model_provider.is_openai()
42                || config.model_provider.requires_openai_auth
43                || config.model_provider.uses_openai_actor_authorization(),
44            provider: config.model_provider.clone(),
45            save_root: resolve_save_root(config),
46        }
47    }
48}
49
50impl ThreadLifecycleContributor<Config> for ImageGenerationExtension {
51    /// Seeds image-generation configuration when a thread begins.
52    fn on_thread_start<'a>(
53        &'a self,
54        input: ThreadStartInput<'a, Config>,
55    ) -> ExtensionFuture<'a, ()> {
56        Box::pin(async move {
57            input
58                .thread_store
59                .insert(ImageGenerationExtensionConfig::from_config(
60                    input.config,
61                    self.resolve_save_root.as_ref(),
62                ));
63        })
64    }
65}
66
67impl ConfigContributor<Config> for ImageGenerationExtension {
68    /// Refreshes image-generation configuration after thread configuration changes.
69    fn on_config_changed(
70        &self,
71        _session_store: &ExtensionData,
72        thread_store: &ExtensionData,
73        _previous_config: &Config,
74        new_config: &Config,
75    ) {
76        thread_store.insert(ImageGenerationExtensionConfig::from_config(
77            new_config,
78            self.resolve_save_root.as_ref(),
79        ));
80    }
81}
82
83impl ToolContributor for ImageGenerationExtension {
84    /// Creates the image-generation tool exposed by this installed extension.
85    fn tools(
86        &self,
87        _session_store: &ExtensionData,
88        thread_store: &ExtensionData,
89        _step_store: &ExtensionData,
90    ) -> Vec<Arc<dyn ToolExecutor<ToolCall>>> {
91        let Some(config) = thread_store.get::<ImageGenerationExtensionConfig>() else {
92            return Vec::new();
93        };
94        if !config.available {
95            return Vec::new();
96        }
97
98        vec![Arc::new(ImageGenerationTool::new(
99            CodexImagesBackend::new(
100                create_model_provider(config.provider.clone(), Some(self.auth_manager.clone())),
101                thread_store
102                    .get::<ThreadOriginator>()
103                    .map(|originator| originator.0.clone()),
104            ),
105            config.save_root.clone(),
106            thread_store.level_id().to_string(),
107        ))]
108    }
109}
110
111/// Installs the standalone image-generation extension contributors.
112pub fn install(
113    registry: &mut ExtensionRegistryBuilder<Config>,
114    auth_manager: Arc<AuthManager>,
115    resolve_save_root: impl Fn(&Config) -> Option<AbsolutePathBuf> + Send + Sync + 'static,
116) {
117    let extension = Arc::new(ImageGenerationExtension {
118        auth_manager,
119        resolve_save_root: Arc::new(resolve_save_root),
120    });
121    registry.thread_lifecycle_contributor(extension.clone());
122    registry.config_contributor(extension.clone());
123    registry.tool_contributor(extension);
124}
codex_image_generation_extension/extension.rs

codex_image_generation_extension/
extension.rs
