diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 99e5d1525..71ed77e9e 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -3177,8 +3177,10 @@ dependencies = [ name = "codex-image-generation-extension" version = "0.0.0" dependencies = [ + "base64 0.22.1", "codex-api", "codex-core", + "codex-exec-server", "codex-extension-api", "codex-login", "codex-model-provider", @@ -3194,6 +3196,7 @@ dependencies = [ "serde", "serde_json", "tokio", + "tracing", ] [[package]] diff --git a/codex-rs/app-server/src/extensions.rs b/codex-rs/app-server/src/extensions.rs index d804e139b..ed040aa8d 100644 --- a/codex-rs/app-server/src/extensions.rs +++ b/codex-rs/app-server/src/extensions.rs @@ -76,7 +76,9 @@ where codex_mcp_extension::install(&mut builder); codex_mcp_extension::install_executor_plugins(&mut builder, environment_manager); codex_web_search_extension::install(&mut builder, auth_manager.clone()); - codex_image_generation_extension::install(&mut builder, auth_manager); + codex_image_generation_extension::install(&mut builder, auth_manager, |config: &Config| { + Some(config.codex_home.clone()) + }); let skill_providers = codex_skills_extension::SkillProviders::new() .with_executor_provider(executor_skill_provider) .with_orchestrator_provider(Arc::new( diff --git a/codex-rs/app-server/tests/suite/v2/imagegen_extension.rs b/codex-rs/app-server/tests/suite/v2/imagegen_extension.rs index 23b97dc9c..8c1643e14 100644 --- a/codex-rs/app-server/tests/suite/v2/imagegen_extension.rs +++ b/codex-rs/app-server/tests/suite/v2/imagegen_extension.rs @@ -134,7 +134,7 @@ async fn standalone_image_generation_returns_saved_path_hint_to_model() -> Resul .context("image output should include model-visible path hint")?; assert!( output_hint.contains(&saved_path.display().to_string()), - "output hint should identify the path core saved" + "output hint should identify the path the extension saved" ); assert!( !requests[1] diff --git a/codex-rs/core/src/stream_events_utils.rs b/codex-rs/core/src/stream_events_utils.rs index 71a9ca7d2..b07c180b1 100644 --- a/codex-rs/core/src/stream_events_utils.rs +++ b/codex-rs/core/src/stream_events_utils.rs @@ -321,7 +321,7 @@ pub(crate) struct HandleOutputCtx { pub cancellation_token: CancellationToken, } -async fn apply_turn_item_contributors( +pub(crate) async fn apply_turn_item_contributors( sess: &Session, turn_store: &ExtensionData, item: &mut TurnItem, diff --git a/codex-rs/core/src/tools/handlers/extension_tools.rs b/codex-rs/core/src/tools/handlers/extension_tools.rs index a89697d90..1aa7754a1 100644 --- a/codex-rs/core/src/tools/handlers/extension_tools.rs +++ b/codex-rs/core/src/tools/handlers/extension_tools.rs @@ -16,6 +16,7 @@ use crate::sandboxing::SandboxPermissions; use crate::session::session::Session; use crate::session::turn_context::TurnContext; use crate::stream_events_utils::TurnItemContributorPolicy; +use crate::stream_events_utils::apply_turn_item_contributors; use crate::stream_events_utils::finalize_turn_item; use crate::tools::context::ToolInvocation; use crate::tools::context::ToolPayload; @@ -71,10 +72,7 @@ struct CoreTurnItemEmitter { fn extension_turn_item(item: ExtensionTurnItem) -> TurnItem { match item { ExtensionTurnItem::WebSearch(item) => TurnItem::WebSearch(item), - ExtensionTurnItem::ImageGeneration(mut item) => { - item.saved_path = None; - TurnItem::ImageGeneration(item) - } + ExtensionTurnItem::ImageGeneration(item) => TurnItem::ImageGeneration(item), } } @@ -95,15 +93,31 @@ impl TurnItemEmitter for CoreTurnItemEmitter { let (Some(session), Some(turn)) = (self.session.upgrade(), self.turn.upgrade()) else { return; }; - let mut item = extension_turn_item(item); - finalize_turn_item( - session.as_ref(), - turn.as_ref(), - TurnItemContributorPolicy::Run(turn.extension_data.as_ref()), - &mut item, - turn.collaboration_mode.mode == codex_protocol::config_types::ModeKind::Plan, - ) - .await; + let item = match item { + ExtensionTurnItem::ImageGeneration(item) => { + let mut item = TurnItem::ImageGeneration(item); + apply_turn_item_contributors( + session.as_ref(), + turn.extension_data.as_ref(), + &mut item, + ) + .await; + item + } + ExtensionTurnItem::WebSearch(item) => { + let mut item = TurnItem::WebSearch(item); + finalize_turn_item( + session.as_ref(), + turn.as_ref(), + TurnItemContributorPolicy::Run(turn.extension_data.as_ref()), + &mut item, + turn.collaboration_mode.mode + == codex_protocol::config_types::ModeKind::Plan, + ) + .await; + item + } + }; session.emit_turn_item_completed(turn.as_ref(), item).await; }) } @@ -540,10 +554,11 @@ mod tests { } #[tokio::test] - async fn image_generation_publication_is_finalized_by_core() { - let handler = ExtensionToolAdapter::new(Arc::new(ImageGenerationExtensionExecutor)); + async fn image_generation_publication_preserves_extension_saved_path() { let (session, turn, rx) = crate::session::tests::make_session_and_context_with_rx().await; - let expected_path = crate::stream_events_utils::image_generation_artifact_path( + let handler = ExtensionToolAdapter::new(Arc::new(ImageGenerationExtensionExecutor)); + let expected_path = test_path_buf("/tmp/extension-claimed.png").abs(); + let default_path = crate::stream_events_utils::image_generation_artifact_path( &turn.config.codex_home, &session.thread_id.to_string(), "call-image", @@ -606,9 +621,6 @@ mod tests { saved_path: Some(expected_path.clone()), } ); - assert_eq!( - std::fs::read(&expected_path).expect("generated artifact should be saved"), - b"png" - ); + assert!(!default_path.exists()); } } diff --git a/codex-rs/core/tests/suite/extension_sandbox.rs b/codex-rs/core/tests/suite/extension_sandbox.rs index aa835fa9a..6d998a469 100644 --- a/codex-rs/core/tests/suite/extension_sandbox.rs +++ b/codex-rs/core/tests/suite/extension_sandbox.rs @@ -26,6 +26,7 @@ use codex_protocol::request_permissions::PermissionGrantScope; use codex_protocol::request_permissions::RequestPermissionProfile; use codex_protocol::request_permissions::RequestPermissionsResponse; use codex_protocol::user_input::UserInput; +use codex_utils_absolute_path::AbsolutePathBuf; use core_test_support::responses; use core_test_support::skip_if_no_network; use core_test_support::skip_if_sandbox; @@ -48,10 +49,13 @@ const TINY_PNG_BYTES: &[u8] = &[ const TINY_PNG_BASE64: &str = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR4nGP4z8DwHwAFAAH/iZk9HQAAAABJRU5ErkJggg=="; const TINY_PNG_DATA_URL: &str = "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR4nGP4z8DwHwAFAAH/iZk9HQAAAABJRU5ErkJggg=="; -fn image_generation_extensions(auth: &CodexAuth) -> Arc> { +fn image_generation_extensions( + auth: &CodexAuth, + resolve_save_root: impl Fn(&Config) -> Option + Send + Sync + 'static, +) -> Arc> { let auth_manager = codex_core::test_support::auth_manager_from_auth(auth.clone()); let mut extension_builder = ExtensionRegistryBuilder::::new(); - install_image_generation_extension(&mut extension_builder, auth_manager); + install_image_generation_extension(&mut extension_builder, auth_manager, resolve_save_root); Arc::new(extension_builder.build()) } @@ -61,7 +65,7 @@ async fn extension_tool_receives_turn_environment_sandbox() -> Result<()> { let server = responses::start_mock_server().await; let auth = CodexAuth::create_dummy_chatgpt_auth_for_testing(); - let extensions = image_generation_extensions(&auth); + let extensions = image_generation_extensions(&auth, |config| Some(config.codex_home.clone())); let mut builder = test_codex() .with_auth(auth) .with_extensions(extensions) @@ -141,7 +145,7 @@ async fn extension_tool_receives_turn_environment_sandbox() -> Result<()> { } #[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn extension_tool_uses_granted_turn_permissions() -> Result<()> { +async fn extension_tool_uses_granted_turn_permissions_without_local_persistence() -> Result<()> { skip_if_no_network!(Ok(())); skip_if_sandbox!(Ok(())); @@ -157,7 +161,7 @@ async fn extension_tool_uses_granted_turn_permissions() -> Result<()> { .await; let auth = CodexAuth::create_dummy_chatgpt_auth_for_testing(); - let extensions = image_generation_extensions(&auth); + let extensions = image_generation_extensions(&auth, |_config| None); let base_permission_profile = PermissionProfile::workspace_write_with( &[], NetworkSandboxPolicy::Restricted, @@ -299,9 +303,14 @@ async fn extension_tool_uses_granted_turn_permissions() -> Result<()> { .last_request() .context("missing request containing extension output")?; let output = request.function_call_output(image_call_id); - let image = &output["output"][0]; - assert_eq!(image["type"], "input_image"); - assert_eq!(image["image_url"], TINY_PNG_DATA_URL); + assert_eq!( + output["output"], + json!([{ + "type": "input_image", + "image_url": TINY_PNG_DATA_URL, + }]) + ); + assert!(!test.config.codex_home.join("generated_images").exists()); Ok(()) } diff --git a/codex-rs/core/tests/suite/responses_lite.rs b/codex-rs/core/tests/suite/responses_lite.rs index 4d2a38f6f..807434411 100644 --- a/codex-rs/core/tests/suite/responses_lite.rs +++ b/codex-rs/core/tests/suite/responses_lite.rs @@ -28,7 +28,9 @@ fn responses_extensions(auth: &CodexAuth) -> Arc> { let auth_manager = codex_core::test_support::auth_manager_from_auth(auth.clone()); let mut extension_builder = ExtensionRegistryBuilder::::new(); install_web_search_extension(&mut extension_builder, Arc::clone(&auth_manager)); - install_image_generation_extension(&mut extension_builder, auth_manager); + install_image_generation_extension(&mut extension_builder, auth_manager, |config| { + Some(config.codex_home.clone()) + }); Arc::new(extension_builder.build()) } diff --git a/codex-rs/ext/image-generation/Cargo.toml b/codex-rs/ext/image-generation/Cargo.toml index 393709a82..c5139ac12 100644 --- a/codex-rs/ext/image-generation/Cargo.toml +++ b/codex-rs/ext/image-generation/Cargo.toml @@ -13,8 +13,10 @@ doctest = false workspace = true [dependencies] +base64 = { workspace = true } codex-api = { workspace = true } codex-core = { workspace = true } +codex-exec-server = { workspace = true } codex-extension-api = { workspace = true } codex-login = { workspace = true } codex-model-provider = { workspace = true } @@ -28,6 +30,7 @@ http = { workspace = true } schemars = { workspace = true } serde = { workspace = true, features = ["derive"] } serde_json = { workspace = true } +tracing = { workspace = true } [dev-dependencies] pretty_assertions = { workspace = true } diff --git a/codex-rs/ext/image-generation/src/extension.rs b/codex-rs/ext/image-generation/src/extension.rs index 6a0016c03..8195c08e6 100644 --- a/codex-rs/ext/image-generation/src/extension.rs +++ b/codex-rs/ext/image-generation/src/extension.rs @@ -21,23 +21,26 @@ use crate::tool::ImageGenerationTool; #[derive(Clone)] struct ImageGenerationExtension { auth_manager: Arc, + resolve_save_root: Arc, } +type SaveRootResolver = dyn Fn(&Config) -> Option + Send + Sync; + #[derive(Clone)] struct ImageGenerationExtensionConfig { available: bool, provider: ModelProviderInfo, - codex_home: AbsolutePathBuf, + save_root: Option, } -impl From<&Config> for ImageGenerationExtensionConfig { +impl ImageGenerationExtensionConfig { /// Resolves whether standalone image generation should be available for a thread. - fn from(config: &Config) -> Self { + fn from_config(config: &Config, resolve_save_root: &SaveRootResolver) -> 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(), + save_root: resolve_save_root(config), } } } @@ -51,7 +54,10 @@ impl ThreadLifecycleContributor for ImageGenerationExtension { Box::pin(async move { input .thread_store - .insert(ImageGenerationExtensionConfig::from(input.config)); + .insert(ImageGenerationExtensionConfig::from_config( + input.config, + self.resolve_save_root.as_ref(), + )); }) } } @@ -65,7 +71,10 @@ impl ConfigContributor for ImageGenerationExtension { _previous_config: &Config, new_config: &Config, ) { - thread_store.insert(ImageGenerationExtensionConfig::from(new_config)); + thread_store.insert(ImageGenerationExtensionConfig::from_config( + new_config, + self.resolve_save_root.as_ref(), + )); } } @@ -88,15 +97,22 @@ impl ToolContributor for ImageGenerationExtension { config.provider.clone(), Some(self.auth_manager.clone()), )), - config.codex_home.clone(), + config.save_root.clone(), thread_store.level_id().to_string(), ))] } } /// Installs the standalone image-generation extension contributors. -pub fn install(registry: &mut ExtensionRegistryBuilder, auth_manager: Arc) { - let extension = Arc::new(ImageGenerationExtension { auth_manager }); +pub fn install( + registry: &mut ExtensionRegistryBuilder, + auth_manager: Arc, + resolve_save_root: impl Fn(&Config) -> Option + Send + Sync + 'static, +) { + let extension = Arc::new(ImageGenerationExtension { + auth_manager, + resolve_save_root: Arc::new(resolve_save_root), + }); registry.thread_lifecycle_contributor(extension.clone()); registry.config_contributor(extension.clone()); registry.tool_contributor(extension); diff --git a/codex-rs/ext/image-generation/src/tool.rs b/codex-rs/ext/image-generation/src/tool.rs index d4451a892..3bd4dc534 100644 --- a/codex-rs/ext/image-generation/src/tool.rs +++ b/codex-rs/ext/image-generation/src/tool.rs @@ -1,5 +1,8 @@ use std::collections::HashSet; +use std::io; +use base64::Engine; +use base64::engine::general_purpose::STANDARD as BASE64_STANDARD; use codex_api::ImageBackground; use codex_api::ImageEditRequest; use codex_api::ImageGenerationRequest; @@ -7,6 +10,9 @@ use codex_api::ImageQuality; use codex_api::ImageUrl; use codex_core::context::extension_image_generation_output_hint; use codex_core::image_generation_artifact_path; +use codex_exec_server::CreateDirectoryOptions; +use codex_exec_server::ExecutorFileSystem; +use codex_exec_server::LOCAL_FS; use codex_extension_api::ExtensionTurnItem; use codex_extension_api::FunctionCallError; use codex_extension_api::ToolCall; @@ -51,7 +57,7 @@ const IMAGEGEN_DESCRIPTION: &str = include_str!("../imagegen_description.md"); #[derive(Clone)] pub(crate) struct ImageGenerationTool { backend: CodexImagesBackend, - codex_home: AbsolutePathBuf, + save_root: Option, thread_id: String, } @@ -59,12 +65,12 @@ impl ImageGenerationTool { /// Creates an image-generation tool backed by an image API executor. pub(crate) fn new( backend: CodexImagesBackend, - codex_home: AbsolutePathBuf, + save_root: Option, thread_id: String, ) -> Self { Self { backend, - codex_home, + save_root, thread_id, } } @@ -145,22 +151,44 @@ impl ImageGenerationTool { return Err(FunctionCallError::RespondToModel(message)); } }; + let saved_path = match self.save_root.as_ref() { + Some(save_root) => match save_image_generation_result( + LOCAL_FS.as_ref(), + save_root, + &self.thread_id, + &call.call_id, + &result, + ) + .await + { + Ok(path) => Some(path), + Err(error) => { + let output_path = + image_generation_artifact_path(save_root, &self.thread_id, &call.call_id); + let output_dir = output_path.parent().unwrap_or_else(|| save_root.clone()); + tracing::warn!( + call_id = %call.call_id, + output_dir = %output_dir.display(), + "failed to save generated image: {error}" + ); + None + } + }, + None => None, + }; call.turn_item_emitter .emit_completed(ExtensionTurnItem::ImageGeneration(ImageGenerationItem { id: call.call_id.clone(), status: "completed".to_string(), revised_prompt: Some(args.prompt), result: result.clone(), - saved_path: None, + saved_path: saved_path.clone(), })) .await; - let output_path = - image_generation_artifact_path(&self.codex_home, &self.thread_id, &call.call_id); - let output_dir = output_path - .parent() - .unwrap_or_else(|| self.codex_home.clone()); - let output_hint = - extension_image_generation_output_hint(output_dir.display(), output_path.display()); + let output_hint = saved_path.as_ref().and_then(|output_path| { + let output_dir = output_path.parent()?; + extension_image_generation_output_hint(output_dir.display(), output_path.display()) + }); Ok(Box::new(GeneratedImageOutput { result, output_hint, @@ -168,6 +196,30 @@ impl ImageGenerationTool { } } +async fn save_image_generation_result( + fs: &dyn ExecutorFileSystem, + save_root: &AbsolutePathBuf, + session_id: &str, + call_id: &str, + result: &str, +) -> io::Result { + let bytes = BASE64_STANDARD + .decode(result.trim().as_bytes()) + .map_err(|error| io::Error::new(io::ErrorKind::InvalidData, error))?; + let path = image_generation_artifact_path(save_root, session_id, call_id); + if let Some(parent) = path.parent() { + fs.create_directory( + &PathUri::from_abs_path(&parent), + CreateDirectoryOptions { recursive: true }, + /*sandbox*/ None, + ) + .await?; + } + fs.write_file(&PathUri::from_abs_path(&path), bytes, /*sandbox*/ None) + .await?; + Ok(path) +} + #[derive(Debug, PartialEq)] enum ImageRequest { Generate(ImageGenerationRequest), diff --git a/codex-rs/tools/src/tool_call.rs b/codex-rs/tools/src/tool_call.rs index fdbfffad8..0bcee5622 100644 --- a/codex-rs/tools/src/tool_call.rs +++ b/codex-rs/tools/src/tool_call.rs @@ -43,12 +43,12 @@ pub enum ExtensionTurnItem { /// Host-provided capability for extension tools to emit visible turn items. /// /// Implementations route lifecycle events through the host's normal item event -/// pipeline, including any persistence and client delivery owned by the host. +/// pipeline and client delivery. pub trait TurnItemEmitter: Send + Sync { /// Emits the beginning of one visible turn item. fn emit_started<'a>(&'a self, item: ExtensionTurnItem) -> TurnItemEmissionFuture<'a>; - /// Emits one visible turn item after host-owned finalization. + /// Emits one completed visible turn item. fn emit_completed<'a>(&'a self, item: ExtensionTurnItem) -> TurnItemEmissionFuture<'a>; }