agent/codex
2
0
Fork 0
mirror of https://github.com/pchuan98/codex.git synced 2026-07-01 00:31:56 +08:00
Code Issues Packages Projects Releases Wiki Activity

Let image generation extension hosts control output persistence (#29711)

Browse Source 
## Why

Some extension hosts need generated images returned without writing them
to the local filesystem or giving the model a local path.

## What changed

**tl;dr**: we now conduct all extension operations in the image gen
extension

- Let hosts provide an optional image save root when installing the
extension.
- Save images and return path hints only when a save root is configured.
- Return image data without saving or adding a path hint when no save
root is configured.
- Preserve the extension-provided `saved_path` instead of persisting
extension images again in core.
- Leave built-in image generation unchanged.

## Validation

- `just test -p codex-image-generation-extension`
- `just test -p codex-app-server
standalone_image_generation_returns_saved_path_hint_to_model`
- `just test -p codex-core
extension_tool_uses_granted_turn_permissions_without_local_persistence`
- `just test -p codex-core tools::handlers::extension_tools::tests`
- tested on CODEX CLI on both save_root: CODEX_HOME and None 
- tested on CODEX APP on both as well
This commit is contained in:
Won Park
2026-06-23 18:51:49 -07:00
committed by GitHub
Unverified
parent 283bc4cf01
commit 61f5a84930
11 changed files with 153 additions and 54 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/Cargo.lock
Generated
+3
View File
@@ -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]]
codex-rs/app-server/src/extensions.rs
+3 -1
View File
@@ -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(
codex-rs/app-server/tests/suite/v2/imagegen_extension.rs
+1 -1
View File
@@ -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]
codex-rs/core/src/stream_events_utils.rs
+1 -1
View File
@@ -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,
codex-rs/core/src/tools/handlers/extension_tools.rs
+32 -20
View File
@@ -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());
}
}
codex-rs/core/tests/suite/extension_sandbox.rs
+17 -8
View File
@@ -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<ExtensionRegistry<Config>> {
fn image_generation_extensions(
auth: &CodexAuth,
resolve_save_root: impl Fn(&Config) -> Option<AbsolutePathBuf> + Send + Sync + 'static,
) -> Arc<ExtensionRegistry<Config>> {
let auth_manager = codex_core::test_support::auth_manager_from_auth(auth.clone());
let mut extension_builder = ExtensionRegistryBuilder::<Config>::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(())
}
codex-rs/core/tests/suite/responses_lite.rs
+3 -1
View File
@@ -28,7 +28,9 @@ fn responses_extensions(auth: &CodexAuth) -> Arc<ExtensionRegistry<Config>> {
let auth_manager = codex_core::test_support::auth_manager_from_auth(auth.clone());
let mut extension_builder = ExtensionRegistryBuilder::<Config>::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())
}
codex-rs/ext/image-generation/Cargo.toml
+3
View File
@@ -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 }
codex-rs/ext/image-generation/src/extension.rs
+25 -9
View File
@@ -21,23 +21,26 @@ use crate::tool::ImageGenerationTool;
#[derive(Clone)]
struct ImageGenerationExtension {
auth_manager: Arc<AuthManager>,
resolve_save_root: Arc<SaveRootResolver>,
}
type SaveRootResolver = dyn Fn(&Config) -> Option<AbsolutePathBuf> + Send + Sync;
#[derive(Clone)]
struct ImageGenerationExtensionConfig {
available: bool,
provider: ModelProviderInfo,
codex_home: AbsolutePathBuf,
save_root: Option<AbsolutePathBuf>,
}
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<Config> 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<Config> 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<Config>, auth_manager: Arc<AuthManager>) {
let extension = Arc::new(ImageGenerationExtension { auth_manager });
pub fn install(
registry: &mut ExtensionRegistryBuilder<Config>,
auth_manager: Arc<AuthManager>,
resolve_save_root: impl Fn(&Config) -> Option<AbsolutePathBuf> + 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);
codex-rs/ext/image-generation/src/tool.rs
+63 -11
View File
@@ -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<AbsolutePathBuf>,
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<AbsolutePathBuf>,
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<AbsolutePathBuf> {
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),
codex-rs/tools/src/tool_call.rs
+2 -2
View File
@@ -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>;
}