From 9ec2811299826e55cb641c2787f083ee9857be46 Mon Sep 17 00:00:00 2001 From: won Date: Tue, 23 Jun 2026 00:05:46 -0700 Subject: [PATCH] draft --- codex-rs/Cargo.lock | 1 + .../tests/suite/v2/imagegen_extension.rs | 93 ++++++++++++++++++- codex-rs/core/config.schema.json | 6 ++ codex-rs/core/src/stream_events_utils.rs | 4 + codex-rs/ext/image-generation/BUILD.bazel | 1 + codex-rs/ext/image-generation/Cargo.toml | 1 + .../imagegen_basic_description.md | 13 +++ .../ext/image-generation/src/extension.rs | 4 + codex-rs/ext/image-generation/src/tests.rs | 24 ++++- codex-rs/ext/image-generation/src/tool.rs | 35 +++++-- codex-rs/features/src/lib.rs | 8 ++ 11 files changed, 177 insertions(+), 13 deletions(-) create mode 100644 codex-rs/ext/image-generation/imagegen_basic_description.md diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 65d581db6f..d1b8399897 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -3164,6 +3164,7 @@ dependencies = [ "codex-api", "codex-core", "codex-extension-api", + "codex-features", "codex-login", "codex-model-provider", "codex-model-provider-info", 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 23b97dc9c5..2b481e879d 100644 --- a/codex-rs/app-server/tests/suite/v2/imagegen_extension.rs +++ b/codex-rs/app-server/tests/suite/v2/imagegen_extension.rs @@ -40,6 +40,7 @@ const TINY_PNG_DATA_URL: &str = "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAA enum ImagegenTestMode { Direct, CodeModeOnly, + Basic, } // macOS and Windows Bazel CI can spend tens of seconds starting app-server @@ -147,6 +148,93 @@ async fn standalone_image_generation_returns_saved_path_hint_to_model() -> Resul Ok(()) } +#[tokio::test] +async fn basic_image_generation_stays_in_conversation() -> Result<()> { + let call_id = "image-run-conversation-only"; + let server = responses::start_mock_server().await; + mount_image_response(&server).await; + let response_mock = responses::mount_sse_sequence( + &server, + vec![ + responses::sse(vec![ + responses::ev_response_created("resp-1"), + responses::ev_function_call_with_namespace( + call_id, + "image_gen", + "imagegen", + &json!({"prompt": "paint a blue whale"}).to_string(), + ), + responses::ev_completed("resp-1"), + ]), + responses::sse(vec![ + responses::ev_assistant_message("msg-1", "Done"), + responses::ev_completed("resp-2"), + ]), + ], + ) + .await; + + let codex_home = TempDir::new()?; + create_config_toml(codex_home.path(), &server.uri(), ImagegenTestMode::Basic)?; + write_chatgpt_auth( + codex_home.path(), + ChatGptAuthFixture::new("access-chatgpt"), + AuthCredentialsStoreMode::File, + )?; + + let mut mcp = + TestAppServer::new_with_env(codex_home.path(), &[("OPENAI_API_KEY", None)]).await?; + timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; + start_image_generation_turn(&mut mcp).await?; + let completed = timeout( + DEFAULT_READ_TIMEOUT, + wait_for_image_generation_completed(&mut mcp), + ) + .await??; + timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_notification_message("turn/completed"), + ) + .await??; + + assert_eq!( + completed.item, + ThreadItem::ImageGeneration { + id: call_id.to_string(), + status: "completed".to_string(), + revised_prompt: Some("paint a blue whale".to_string()), + result: RESULT.to_string(), + saved_path: None, + } + ); + assert!(!codex_home.path().join("generated_images").exists()); + + let requests = response_mock.requests(); + assert_eq!(requests.len(), 2); + let tool = requests[0] + .tool_by_name("image_gen", "imagegen") + .context("basic imagegen tool should be sent to the model")?; + assert!( + tool.pointer("/parameters/properties/referenced_image_paths") + .is_none() + ); + assert!( + tool["description"] + .as_str() + .is_some_and(|description| description.contains("conversation images")) + ); + assert_eq!( + requests[1].function_call_output(call_id)["output"], + json!([{ + "type": "input_image", + "image_url": format!("data:image/png;base64,{RESULT}"), + "detail": "high", + }]) + ); + + Ok(()) +} + #[tokio::test] async fn standalone_image_generation_failure_emits_terminal_item() -> Result<()> { let call_id = "image-run-failed"; @@ -552,9 +640,10 @@ fn create_config_toml( server_uri: &str, mode: ImagegenTestMode, ) -> std::io::Result<()> { - let code_mode_only = match mode { + let mode_feature = match mode { ImagegenTestMode::Direct => "", ImagegenTestMode::CodeModeOnly => "code_mode_only = true", + ImagegenTestMode::Basic => "imagegenbasic = true", }; std::fs::write( codex_home.join("config.toml"), @@ -568,7 +657,7 @@ chatgpt_base_url = "{server_uri}" [features] imagegenext = true -{code_mode_only} +{mode_feature} [model_providers.openai-custom] name = "OpenAI" diff --git a/codex-rs/core/config.schema.json b/codex-rs/core/config.schema.json index a10d01a8a7..2e994564ca 100644 --- a/codex-rs/core/config.schema.json +++ b/codex-rs/core/config.schema.json @@ -533,6 +533,9 @@ "image_generation": { "type": "boolean" }, + "imagegenbasic": { + "type": "boolean" + }, "imagegenext": { "type": "boolean" }, @@ -4837,6 +4840,9 @@ "image_generation": { "type": "boolean" }, + "imagegenbasic": { + "type": "boolean" + }, "imagegenext": { "type": "boolean" }, diff --git a/codex-rs/core/src/stream_events_utils.rs b/codex-rs/core/src/stream_events_utils.rs index 71a9ca7d21..763327cfda 100644 --- a/codex-rs/core/src/stream_events_utils.rs +++ b/codex-rs/core/src/stream_events_utils.rs @@ -4,6 +4,7 @@ use std::sync::Arc; use base64::Engine; use base64::engine::general_purpose::STANDARD as BASE64_STANDARD; use codex_extension_api::ExtensionData; +use codex_features::Feature; use codex_protocol::config_types::ModeKind; use codex_protocol::items::ImageGenerationItem; use codex_protocol::items::TurnItem; @@ -133,6 +134,9 @@ pub(crate) async fn persist_image_generation_item( image_item: &mut ImageGenerationItem, ) -> Option { image_item.saved_path = None; + if turn_context.config.features.enabled(Feature::ImageGenBasic) { + return None; + } let session_id = sess.thread_id.to_string(); match save_image_generation_result( &turn_context.config.codex_home, diff --git a/codex-rs/ext/image-generation/BUILD.bazel b/codex-rs/ext/image-generation/BUILD.bazel index 97698380e6..ac2749fc33 100644 --- a/codex-rs/ext/image-generation/BUILD.bazel +++ b/codex-rs/ext/image-generation/BUILD.bazel @@ -3,6 +3,7 @@ load("//:defs.bzl", "codex_rust_crate") codex_rust_crate( name = "image-generation", compile_data = [ + "imagegen_basic_description.md", "imagegen_description.md", ], crate_name = "codex_image_generation_extension", diff --git a/codex-rs/ext/image-generation/Cargo.toml b/codex-rs/ext/image-generation/Cargo.toml index 393709a82d..e79c7093cf 100644 --- a/codex-rs/ext/image-generation/Cargo.toml +++ b/codex-rs/ext/image-generation/Cargo.toml @@ -16,6 +16,7 @@ workspace = true codex-api = { workspace = true } codex-core = { workspace = true } codex-extension-api = { workspace = true } +codex-features = { workspace = true } codex-login = { workspace = true } codex-model-provider = { workspace = true } codex-model-provider-info = { workspace = true } diff --git a/codex-rs/ext/image-generation/imagegen_basic_description.md b/codex-rs/ext/image-generation/imagegen_basic_description.md new file mode 100644 index 0000000000..3307573df7 --- /dev/null +++ b/codex-rs/ext/image-generation/imagegen_basic_description.md @@ -0,0 +1,13 @@ +The `image_gen.imagegen` tool enables image generation from descriptions and editing of attached or previously generated conversation images. Use it when: + +- The user requests a new image, such as a diagram, portrait, comic, meme, or other visual. +- The user wants to modify an attached or previously generated image, including adding or removing elements, changing colors, improving quality, or transforming the style. + +Guidelines: +- In code mode, pass the result to `generatedImage(result)`. +- Omit `num_last_images_to_include` when generating a brand new image. +- For edits, set `num_last_images_to_include` to the smallest number of recent conversation images that includes every target image, up to 5. +- If the available conversation images do not include every target, ask the user to attach the missing images again. +- Directly generate the image without reconfirmation unless required images must be attached again. +- After each image generation, do not mention anything related to download. Do not summarize the image. Do not ask a follow-up question. Do not say anything after you generate an image. +- Always use this tool for image editing unless the user explicitly requests otherwise. Do not use the `python` tool for image editing unless specifically instructed. diff --git a/codex-rs/ext/image-generation/src/extension.rs b/codex-rs/ext/image-generation/src/extension.rs index 6a0016c037..d4018cd904 100644 --- a/codex-rs/ext/image-generation/src/extension.rs +++ b/codex-rs/ext/image-generation/src/extension.rs @@ -10,6 +10,7 @@ use codex_extension_api::ThreadStartInput; use codex_extension_api::ToolCall; use codex_extension_api::ToolContributor; use codex_extension_api::ToolExecutor; +use codex_features::Feature; use codex_login::AuthManager; use codex_model_provider::create_model_provider; use codex_model_provider_info::ModelProviderInfo; @@ -28,6 +29,7 @@ struct ImageGenerationExtensionConfig { available: bool, provider: ModelProviderInfo, codex_home: AbsolutePathBuf, + basic: bool, } impl From<&Config> for ImageGenerationExtensionConfig { @@ -38,6 +40,7 @@ impl From<&Config> for ImageGenerationExtensionConfig { available: config.model_provider.is_openai(), provider: config.model_provider.clone(), codex_home: config.codex_home.clone(), + basic: config.features.enabled(Feature::ImageGenBasic), } } } @@ -90,6 +93,7 @@ impl ToolContributor for ImageGenerationExtension { )), config.codex_home.clone(), thread_store.level_id().to_string(), + config.basic, ))] } } diff --git a/codex-rs/ext/image-generation/src/tests.rs b/codex-rs/ext/image-generation/src/tests.rs index fe915746f9..df4fe7a94b 100644 --- a/codex-rs/ext/image-generation/src/tests.rs +++ b/codex-rs/ext/image-generation/src/tests.rs @@ -18,6 +18,7 @@ use codex_tools::ResponsesApiNamespaceTool; use pretty_assertions::assert_eq; use super::GeneratedImageOutput; +use super::IMAGEGEN_BASIC_DESCRIPTION; use super::ImageRequest; use super::ImagegenArgs; use super::imagegen_tool_spec; @@ -29,7 +30,7 @@ const RESULT: &str = "cG5n"; #[test] fn uses_reserved_image_gen_namespace() { - let ToolSpec::Namespace(spec) = imagegen_tool_spec() else { + let ToolSpec::Namespace(spec) = imagegen_tool_spec(/*basic*/ false) else { panic!("imagegen should advertise a namespace tool"); }; assert_eq!(spec.name, IMAGE_GEN_NAMESPACE); @@ -37,6 +38,27 @@ fn uses_reserved_image_gen_namespace() { assert_eq!(function.name, IMAGEGEN_TOOL_NAME); } +#[test] +fn basic_model_contract_omits_filesystem_paths() { + let ToolSpec::Namespace(basic) = imagegen_tool_spec(/*basic*/ true) else { + panic!("basic imagegen should advertise a namespace tool"); + }; + let ResponsesApiNamespaceTool::Function(basic) = &basic.tools[0]; + let basic_properties = basic + .parameters + .properties + .as_ref() + .expect("basic imagegen should define properties"); + assert_eq!( + basic_properties.keys().cloned().collect::>(), + vec![ + "num_last_images_to_include".to_string(), + "prompt".to_string(), + ] + ); + assert_eq!(basic.description, IMAGEGEN_BASIC_DESCRIPTION); +} + #[tokio::test] async fn omitted_references_generate_with_fixed_defaults() { assert_eq!( diff --git a/codex-rs/ext/image-generation/src/tool.rs b/codex-rs/ext/image-generation/src/tool.rs index 3e2cfd1792..f0b238dd27 100644 --- a/codex-rs/ext/image-generation/src/tool.rs +++ b/codex-rs/ext/image-generation/src/tool.rs @@ -47,12 +47,14 @@ use crate::backend::CodexImagesBackend; const IMAGE_MODEL: &str = "gpt-image-2"; const MAX_EDIT_IMAGES: usize = 5; const IMAGEGEN_DESCRIPTION: &str = include_str!("../imagegen_description.md"); +const IMAGEGEN_BASIC_DESCRIPTION: &str = include_str!("../imagegen_basic_description.md"); #[derive(Clone)] pub(crate) struct ImageGenerationTool { backend: CodexImagesBackend, codex_home: AbsolutePathBuf, thread_id: String, + basic: bool, } impl ImageGenerationTool { @@ -61,11 +63,13 @@ impl ImageGenerationTool { backend: CodexImagesBackend, codex_home: AbsolutePathBuf, thread_id: String, + basic: bool, ) -> Self { Self { backend, codex_home, thread_id, + basic, } } } @@ -88,7 +92,7 @@ impl ToolExecutor for ImageGenerationTool { /// Advertises the model contract: a rewritten prompt and optional edit references. fn spec(&self) -> ToolSpec { - imagegen_tool_spec() + imagegen_tool_spec(self.basic) } /// Exposes image generation directly and through the nested code-mode tool surface. @@ -154,13 +158,16 @@ impl ImageGenerationTool { saved_path: None, })) .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 = if self.basic { + None + } else { + 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()); + extension_image_generation_output_hint(output_dir.display(), output_path.display()) + }; Ok(Box::new(GeneratedImageOutput { result, output_hint, @@ -373,7 +380,7 @@ fn parse_args(call: &ToolCall) -> Result { } /// Builds the namespace function schema exposed to the model. -fn imagegen_tool_spec() -> ToolSpec { +fn imagegen_tool_spec(basic: bool) -> ToolSpec { let mut schema_value = serde_json::to_value( SchemaSettings::draft2019_09() .with(|settings| settings.inline_subschemas = true) @@ -384,6 +391,9 @@ fn imagegen_tool_spec() -> ToolSpec { let Value::Object(ref mut schema) = schema_value else { unreachable!("imagegen root schema must be an object"); }; + if basic && let Some(Value::Object(properties)) = schema.get_mut("properties") { + properties.remove("referenced_image_paths"); + } let mut input_schema = Map::new(); for key in ["properties", "required", "type", "additionalProperties"] { if let Some(value) = schema.remove(key) { @@ -395,7 +405,12 @@ fn imagegen_tool_spec() -> ToolSpec { description: default_namespace_description(IMAGE_GEN_NAMESPACE), tools: vec![ResponsesApiNamespaceTool::Function(ResponsesApiTool { name: IMAGEGEN_TOOL_NAME.to_string(), - description: IMAGEGEN_DESCRIPTION.to_string(), + description: if basic { + IMAGEGEN_BASIC_DESCRIPTION + } else { + IMAGEGEN_DESCRIPTION + } + .to_string(), strict: false, parameters: parse_tool_input_schema(&Value::Object(input_schema)) .unwrap_or_else(|err| panic!("imagegen input schema should parse: {err}")), diff --git a/codex-rs/features/src/lib.rs b/codex-rs/features/src/lib.rs index 86fdd2e29e..3ab8a6a822 100644 --- a/codex-rs/features/src/lib.rs +++ b/codex-rs/features/src/lib.rs @@ -197,6 +197,8 @@ pub enum Feature { ImageGeneration, /// Replace hosted image generation with the standalone image-generation extension. ImageGenExt, + /// Use the basic image-generation contract without filesystem paths. + ImageGenBasic, /// Removed compatibility flag for always-on centralized image preparation. ResizeAllImages, /// Generate Responses API item IDs for client-created history items. @@ -1167,6 +1169,12 @@ pub const FEATURES: &[FeatureSpec] = &[ stage: Stage::UnderDevelopment, default_enabled: false, }, + FeatureSpec { + id: Feature::ImageGenBasic, + key: "imagegenbasic", + stage: Stage::UnderDevelopment, + default_enabled: false, + }, FeatureSpec { id: Feature::ResizeAllImages, key: "resize_all_images",