mirror of
https://github.com/pchuan98/codex.git
synced 2026-07-01 00:31:56 +08:00
Add feature-gated standalone image generation extension (#24723)
## Why Add a standalone image generation path that can be exercised independently of hosted Responses image generation, while retaining the hosted tool as fallback unless the extension is actually available to the model. ## What changed - Added the `codex-image-generation-extension` crate with standalone generate/edit execution, prior-image selection for edits, model-visible image output, and local generated-image persistence. - Installed the extension in app-server behind the disabled-by-default `imagegenext` feature and backend eligibility checks. - Updated core tool planning so eligible `image_gen.imagegen` exposure replaces hosted `image_generation`, while unavailable configurations retain hosted fallback. - Added coverage for extension behavior, edit history reuse, feature gating, auth eligibility, and hosted-tool replacement. - The extension is installed through app-server only in this PR; other execution paths retain hosted image generation because hosted replacement occurs only when the standalone executor is actually registered and model-visible. - The initial extension contract intentionally fixes the image model to `gpt-image-2` and uses automatic image parameters. - Native generated-image history/card parity and rollout persistence cleanup are intentionally deferred follow-up work. ## Validation - `just test -p codex-image-generation-extension` - `just test -p codex-features` - `just test -p codex-core hosted_tools_follow_provider_auth_model_and_config_gates` - `just test -p codex-app-server` - `just fix -p codex-image-generation-extension -p codex-features -p codex-core -p codex-app-server` - `just fmt` - `just bazel-lock-update` - `just bazel-lock-check` --------- Co-authored-by: jif-oai <jif@openai.com>
This commit is contained in:
Won Park
jif-oai
@@ -0,0 +1,395 @@
|
||||
use std::path::Path;
|
||||
use std::path::PathBuf;
|
||||
|
||||
use base64::Engine;
|
||||
use base64::engine::general_purpose::STANDARD as BASE64_STANDARD;
|
||||
use codex_api::ImageBackground;
|
||||
use codex_api::ImageEditRequest;
|
||||
use codex_api::ImageGenerationRequest;
|
||||
use codex_api::ImageQuality;
|
||||
use codex_api::ImageUrl;
|
||||
use codex_extension_api::FunctionCallError;
|
||||
use codex_extension_api::ToolCall;
|
||||
use codex_extension_api::ToolExecutor;
|
||||
use codex_extension_api::ToolName;
|
||||
use codex_extension_api::ToolOutput;
|
||||
use codex_extension_api::ToolPayload;
|
||||
use codex_extension_api::ToolSpec;
|
||||
use codex_extension_api::parse_tool_input_schema;
|
||||
use codex_protocol::models::ContentItem;
|
||||
use codex_protocol::models::DEFAULT_IMAGE_DETAIL;
|
||||
use codex_protocol::models::FunctionCallOutputBody;
|
||||
use codex_protocol::models::FunctionCallOutputContentItem;
|
||||
use codex_protocol::models::FunctionCallOutputPayload;
|
||||
use codex_protocol::models::ResponseInputItem;
|
||||
use codex_protocol::models::ResponseItem;
|
||||
use codex_tools::ResponsesApiNamespace;
|
||||
use codex_tools::ResponsesApiNamespaceTool;
|
||||
use codex_tools::ResponsesApiTool;
|
||||
use codex_tools::ToolExposure;
|
||||
use codex_tools::default_namespace_description;
|
||||
use schemars::JsonSchema;
|
||||
use schemars::r#gen::SchemaSettings;
|
||||
use serde::Deserialize;
|
||||
use serde_json::Map;
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::IMAGE_GEN_NAMESPACE;
|
||||
use crate::IMAGEGEN_TOOL_NAME;
|
||||
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 GENERATED_IMAGE_ARTIFACTS_DIR: &str = "generated_images";
|
||||
|
||||
#[derive(Clone)]
|
||||
pub(crate) struct ImageGenerationTool {
|
||||
backend: CodexImagesBackend,
|
||||
output_dir: PathBuf,
|
||||
}
|
||||
|
||||
impl ImageGenerationTool {
|
||||
/// Creates an image-generation tool backed by an image API executor.
|
||||
pub(crate) fn new(backend: CodexImagesBackend, output_dir: PathBuf) -> Self {
|
||||
Self {
|
||||
backend,
|
||||
output_dir,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize, JsonSchema)]
|
||||
#[serde(deny_unknown_fields)]
|
||||
struct ImagegenArgs {
|
||||
prompt: String,
|
||||
action: ImagegenAction,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize, JsonSchema)]
|
||||
#[serde(rename_all = "lowercase")]
|
||||
enum ImagegenAction {
|
||||
Generate,
|
||||
Edit,
|
||||
}
|
||||
|
||||
#[async_trait::async_trait]
|
||||
impl ToolExecutor<ToolCall> for ImageGenerationTool {
|
||||
/// Keeps the tool in the existing image-generation Responses namespace.
|
||||
fn tool_name(&self) -> ToolName {
|
||||
ToolName::namespaced(IMAGE_GEN_NAMESPACE, IMAGEGEN_TOOL_NAME)
|
||||
}
|
||||
|
||||
/// Advertises the model contract: a rewritten prompt and semantic action.
|
||||
fn spec(&self) -> ToolSpec {
|
||||
imagegen_tool_spec()
|
||||
}
|
||||
|
||||
/// Keeps this model-facing tool out of the nested code-mode tool surface.
|
||||
fn exposure(&self) -> ToolExposure {
|
||||
ToolExposure::DirectModelOnly
|
||||
}
|
||||
|
||||
/// Executes the selected image operation and returns the completed image result.
|
||||
async fn handle(&self, call: ToolCall) -> Result<Box<dyn ToolOutput>, FunctionCallError> {
|
||||
let args = parse_args(&call)?;
|
||||
let request = request_for_action(&args, call.conversation_history.items())?;
|
||||
|
||||
let response = match request {
|
||||
ImageRequest::Generate(request) => self.backend.generate(request).await,
|
||||
ImageRequest::Edit(request) => self.backend.edit(request).await,
|
||||
}
|
||||
.map_err(|err| {
|
||||
FunctionCallError::RespondToModel(format!("image generation failed: {err}"))
|
||||
})?;
|
||||
let Some(result) = response.data.into_iter().next().map(|data| data.b64_json) else {
|
||||
return Err(FunctionCallError::RespondToModel(
|
||||
"image generation returned no image data".to_string(),
|
||||
));
|
||||
};
|
||||
let output_hint =
|
||||
match persist_generated_image(&self.output_dir, &call.call_id, &result).await {
|
||||
Ok(output_hint) => Some(output_hint),
|
||||
Err(err) => {
|
||||
tracing::warn!(
|
||||
call_id = %call.call_id,
|
||||
output_dir = %self.output_dir.display(),
|
||||
"failed to save generated image: {err}"
|
||||
);
|
||||
None
|
||||
}
|
||||
};
|
||||
Ok(Box::new(GeneratedImageOutput {
|
||||
result,
|
||||
output_hint,
|
||||
}))
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, PartialEq)]
|
||||
enum ImageRequest {
|
||||
Generate(ImageGenerationRequest),
|
||||
Edit(ImageEditRequest),
|
||||
}
|
||||
|
||||
/// Maps the model-selected action to the fixed image API request parameters.
|
||||
fn request_for_action(
|
||||
args: &ImagegenArgs,
|
||||
history: &[ResponseItem],
|
||||
) -> Result<ImageRequest, FunctionCallError> {
|
||||
match args.action {
|
||||
ImagegenAction::Generate => Ok(ImageRequest::Generate(ImageGenerationRequest {
|
||||
prompt: args.prompt.clone(),
|
||||
background: Some(ImageBackground::Auto),
|
||||
model: IMAGE_MODEL.to_string(),
|
||||
n: None,
|
||||
quality: Some(ImageQuality::Auto),
|
||||
size: Some("auto".to_string()),
|
||||
})),
|
||||
ImagegenAction::Edit => {
|
||||
let images = edit_images(history);
|
||||
if images.is_empty() {
|
||||
return Err(FunctionCallError::RespondToModel(
|
||||
"image edit requested without any usable image in conversation history"
|
||||
.to_string(),
|
||||
));
|
||||
}
|
||||
Ok(ImageRequest::Edit(ImageEditRequest {
|
||||
images,
|
||||
prompt: args.prompt.clone(),
|
||||
background: Some(ImageBackground::Auto),
|
||||
model: IMAGE_MODEL.to_string(),
|
||||
n: None,
|
||||
quality: Some(ImageQuality::Auto),
|
||||
size: Some("auto".to_string()),
|
||||
}))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Selects edit context using the hosted imagegen anchor and truncation behavior.
|
||||
fn edit_images(history: &[ResponseItem]) -> Vec<ImageUrl> {
|
||||
let latest_uploaded_images = history.iter().enumerate().rev().find_map(|(index, item)| {
|
||||
let ResponseItem::Message { role, content, .. } = item else {
|
||||
return None;
|
||||
};
|
||||
if role != "user" {
|
||||
return None;
|
||||
}
|
||||
let images = content
|
||||
.iter()
|
||||
.filter_map(|item| match item {
|
||||
ContentItem::InputImage { image_url, .. } => Some(ImageUrl {
|
||||
image_url: image_url.clone(),
|
||||
}),
|
||||
ContentItem::InputText { .. } | ContentItem::OutputText { .. } => None,
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
(!images.is_empty()).then_some((index, images))
|
||||
});
|
||||
let (user_images, follow_up_start) = latest_uploaded_images
|
||||
.map_or_else(|| (Vec::new(), 0), |(index, images)| (images, index + 1));
|
||||
let mut generated_images = Vec::new();
|
||||
for item in &history[follow_up_start..] {
|
||||
match item {
|
||||
ResponseItem::ImageGenerationCall { result, .. } if !result.is_empty() => {
|
||||
generated_images.push(ImageUrl {
|
||||
image_url: format!("data:image/png;base64,{result}"),
|
||||
});
|
||||
}
|
||||
ResponseItem::FunctionCallOutput { call_id, output }
|
||||
if history.iter().any(|item| {
|
||||
matches!(
|
||||
item,
|
||||
ResponseItem::FunctionCall {
|
||||
name,
|
||||
namespace: Some(namespace),
|
||||
call_id: function_call_id,
|
||||
..
|
||||
} if function_call_id == call_id
|
||||
&& name == IMAGEGEN_TOOL_NAME
|
||||
&& namespace == IMAGE_GEN_NAMESPACE
|
||||
)
|
||||
}) =>
|
||||
{
|
||||
generated_images.extend(output.content_items().into_iter().flatten().filter_map(
|
||||
|item| match item {
|
||||
FunctionCallOutputContentItem::InputImage { image_url, .. } => {
|
||||
Some(ImageUrl {
|
||||
image_url: image_url.clone(),
|
||||
})
|
||||
}
|
||||
FunctionCallOutputContentItem::InputText { .. }
|
||||
| FunctionCallOutputContentItem::EncryptedContent { .. } => None,
|
||||
},
|
||||
));
|
||||
}
|
||||
ResponseItem::Message { .. }
|
||||
| ResponseItem::Reasoning { .. }
|
||||
| ResponseItem::LocalShellCall { .. }
|
||||
| ResponseItem::FunctionCall { .. }
|
||||
| ResponseItem::ToolSearchCall { .. }
|
||||
| ResponseItem::FunctionCallOutput { .. }
|
||||
| ResponseItem::CustomToolCall { .. }
|
||||
| ResponseItem::CustomToolCallOutput { .. }
|
||||
| ResponseItem::ToolSearchOutput { .. }
|
||||
| ResponseItem::WebSearchCall { .. }
|
||||
| ResponseItem::ImageGenerationCall { .. }
|
||||
| ResponseItem::Compaction { .. }
|
||||
| ResponseItem::CompactionTrigger
|
||||
| ResponseItem::ContextCompaction { .. }
|
||||
| ResponseItem::Other => {}
|
||||
}
|
||||
}
|
||||
truncate_images(user_images, generated_images)
|
||||
}
|
||||
|
||||
/// Truncates edit inputs while preserving the newest generated image when possible.
|
||||
fn truncate_images(
|
||||
mut user_images: Vec<ImageUrl>,
|
||||
mut generated_images: Vec<ImageUrl>,
|
||||
) -> Vec<ImageUrl> {
|
||||
let mut excess = (user_images.len() + generated_images.len()).saturating_sub(MAX_EDIT_IMAGES);
|
||||
let drop_generated = excess.min(generated_images.len().saturating_sub(1));
|
||||
generated_images.drain(..drop_generated);
|
||||
excess -= drop_generated;
|
||||
let drop_user = excess.min(user_images.len());
|
||||
user_images.drain(..drop_user);
|
||||
excess -= drop_user;
|
||||
generated_images.drain(..excess);
|
||||
|
||||
user_images.extend(generated_images);
|
||||
user_images
|
||||
}
|
||||
|
||||
/// Parses the strict model-facing arguments for an image-generation call.
|
||||
fn parse_args(call: &ToolCall) -> Result<ImagegenArgs, FunctionCallError> {
|
||||
serde_json::from_str(call.function_arguments()?)
|
||||
.map_err(|err| FunctionCallError::RespondToModel(err.to_string()))
|
||||
}
|
||||
|
||||
/// Resolves where generated images for one thread are persisted by the extension.
|
||||
pub(crate) fn generated_image_output_dir(codex_home: &Path, thread_id: &str) -> PathBuf {
|
||||
codex_home
|
||||
.join(GENERATED_IMAGE_ARTIFACTS_DIR)
|
||||
.join(sanitize_path_component(thread_id))
|
||||
}
|
||||
|
||||
fn generated_image_output_path(output_dir: &Path, call_id: &str) -> PathBuf {
|
||||
output_dir.join(format!("{}.png", sanitize_path_component(call_id)))
|
||||
}
|
||||
|
||||
fn sanitize_path_component(value: &str) -> String {
|
||||
let sanitized: String = value
|
||||
.chars()
|
||||
.map(|ch| {
|
||||
if ch.is_ascii_alphanumeric() || ch == '-' || ch == '_' {
|
||||
ch
|
||||
} else {
|
||||
'_'
|
||||
}
|
||||
})
|
||||
.collect();
|
||||
if sanitized.is_empty() {
|
||||
"generated_image".to_string()
|
||||
} else {
|
||||
sanitized
|
||||
}
|
||||
}
|
||||
|
||||
async fn persist_generated_image(
|
||||
output_dir: &Path,
|
||||
call_id: &str,
|
||||
result: &str,
|
||||
) -> Result<String, String> {
|
||||
let bytes = BASE64_STANDARD
|
||||
.decode(result.trim().as_bytes())
|
||||
.map_err(|err| format!("invalid image generation payload: {err}"))?;
|
||||
tokio::fs::create_dir_all(output_dir)
|
||||
.await
|
||||
.map_err(|err| err.to_string())?;
|
||||
tokio::fs::write(generated_image_output_path(output_dir, call_id), bytes)
|
||||
.await
|
||||
.map_err(|err| err.to_string())?;
|
||||
|
||||
Ok(format!(
|
||||
"Generated images are saved to {} as {} by default.\n\
|
||||
If you need to use a generated image at another path, copy it and leave the original in place unless the user explicitly asks you to delete it.",
|
||||
output_dir.display(),
|
||||
generated_image_output_path(output_dir, call_id).display(),
|
||||
))
|
||||
}
|
||||
|
||||
/// Builds the namespace function schema exposed to the model.
|
||||
fn imagegen_tool_spec() -> ToolSpec {
|
||||
let mut schema_value = serde_json::to_value(
|
||||
SchemaSettings::draft2019_09()
|
||||
.with(|settings| settings.inline_subschemas = true)
|
||||
.into_generator()
|
||||
.into_root_schema_for::<ImagegenArgs>(),
|
||||
)
|
||||
.unwrap_or_else(|err| panic!("imagegen schema should serialize: {err}"));
|
||||
let Value::Object(ref mut schema) = schema_value else {
|
||||
unreachable!("imagegen root schema must be an object");
|
||||
};
|
||||
let mut input_schema = Map::new();
|
||||
for key in ["properties", "required", "type", "additionalProperties"] {
|
||||
if let Some(value) = schema.remove(key) {
|
||||
input_schema.insert(key.to_string(), value);
|
||||
}
|
||||
}
|
||||
ToolSpec::Namespace(ResponsesApiNamespace {
|
||||
name: IMAGE_GEN_NAMESPACE.to_string(),
|
||||
description: default_namespace_description(IMAGE_GEN_NAMESPACE),
|
||||
tools: vec![ResponsesApiNamespaceTool::Function(ResponsesApiTool {
|
||||
name: IMAGEGEN_TOOL_NAME.to_string(),
|
||||
description: 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}")),
|
||||
output_schema: None,
|
||||
defer_loading: None,
|
||||
})],
|
||||
})
|
||||
}
|
||||
|
||||
struct GeneratedImageOutput {
|
||||
result: String,
|
||||
output_hint: Option<String>,
|
||||
}
|
||||
|
||||
impl ToolOutput for GeneratedImageOutput {
|
||||
/// Avoids copying image bytes into tool-call telemetry.
|
||||
fn log_preview(&self) -> String {
|
||||
"[generated image]".to_string()
|
||||
}
|
||||
|
||||
/// Reports a completed images request as successful tool execution.
|
||||
fn success_for_logging(&self) -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
/// Returns generated bytes and persisted-artifact context for the model's follow-up response.
|
||||
fn to_response_item(&self, call_id: &str, _payload: &ToolPayload) -> ResponseInputItem {
|
||||
let mut content = vec![FunctionCallOutputContentItem::InputImage {
|
||||
image_url: format!("data:image/png;base64,{}", self.result),
|
||||
detail: Some(DEFAULT_IMAGE_DETAIL),
|
||||
}];
|
||||
if let Some(output_hint) = &self.output_hint {
|
||||
content.push(FunctionCallOutputContentItem::InputText {
|
||||
text: output_hint.clone(),
|
||||
});
|
||||
}
|
||||
ResponseInputItem::FunctionCallOutput {
|
||||
call_id: call_id.to_string(),
|
||||
output: FunctionCallOutputPayload {
|
||||
body: FunctionCallOutputBody::ContentItems(content),
|
||||
success: Some(true),
|
||||
},
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "tests.rs"]
|
||||
mod tests;
|
||||
Reference in New Issue
Block a user