codex/codex-rs/ext/image-generation/src/tool.rs

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;
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;
use codex_extension_api::ToolEnvironment;
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::items::ImageGenerationItem;
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 codex_utils_absolute_path::AbsolutePathBuf;
use codex_utils_image::PromptImageMode;
use codex_utils_image::load_for_prompt_bytes;
use codex_utils_path_uri::PathUri;
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");

#[derive(Clone)]
pub(crate) struct ImageGenerationTool {
    backend: CodexImagesBackend,
    save_root: Option<AbsolutePathBuf>,
    thread_id: String,
}

impl ImageGenerationTool {
    /// Creates an image-generation tool backed by an image API executor.
    pub(crate) fn new(
        backend: CodexImagesBackend,
        save_root: Option<AbsolutePathBuf>,
        thread_id: String,
    ) -> Self {
        Self {
            backend,
            save_root,
            thread_id,
        }
    }
}

#[derive(Debug, Deserialize, JsonSchema)]
#[serde(deny_unknown_fields)]
struct ImagegenArgs {
    prompt: String,
    #[schemars(length(max = 5))]
    referenced_image_paths: Option<Vec<AbsolutePathBuf>>,
    #[schemars(range(min = 1, max = 5))]
    num_last_images_to_include: Option<usize>,
}

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 optional edit references.
    fn spec(&self) -> ToolSpec {
        imagegen_tool_spec()
    }

    /// Exposes image generation directly and through the nested code-mode tool surface.
    fn exposure(&self) -> ToolExposure {
        ToolExposure::Direct
    }

    /// Executes the selected image operation and returns the completed image result.
    fn handle(&self, call: ToolCall) -> codex_extension_api::ToolExecutorFuture<'_> {
        Box::pin(self.handle_call(call))
    }
}

impl ImageGenerationTool {
    async fn handle_call(&self, call: ToolCall) -> Result<Box<dyn ToolOutput>, FunctionCallError> {
        let args = parse_args(&call)?;
        let request =
            request_for_call_args(&args, call.conversation_history.items(), &call.environments)
                .await?;
        call.turn_item_emitter
            .emit_started(ExtensionTurnItem::ImageGeneration(ImageGenerationItem {
                id: call.call_id.clone(),
                status: "in_progress".to_string(),
                revised_prompt: None,
                result: String::new(),
                saved_path: None,
            }))
            .await;
        let result = match request {
            ImageRequest::Generate(request) => self.backend.generate(request).await,
            ImageRequest::Edit(request) => self.backend.edit(request).await,
        }
        .map_err(|err| format!("image generation failed: {err}"))
        .and_then(|response| {
            response
                .data
                .into_iter()
                .next()
                .map(|data| data.b64_json)
                .ok_or_else(|| "image generation returned no image data".to_string())
        });
        let result = match result {
            Ok(result) => result,
            Err(message) => {
                call.turn_item_emitter
                    .emit_completed(ExtensionTurnItem::ImageGeneration(ImageGenerationItem {
                        id: call.call_id.clone(),
                        status: "failed".to_string(),
                        revised_prompt: Some(args.prompt.clone()),
                        result: String::new(),
                        saved_path: None,
                    }))
                    .await;
                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: saved_path.clone(),
            }))
            .await;
        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,
        }))
    }
}

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),
    Edit(ImageEditRequest),
}

async fn request_for_call_args(
    args: &ImagegenArgs,
    history: &[ResponseItem],
    environments: &[ToolEnvironment],
) -> Result<ImageRequest, FunctionCallError> {
    let paths = args.referenced_image_paths.as_deref().unwrap_or_default();
    if paths.len() > MAX_EDIT_IMAGES {
        return Err(FunctionCallError::RespondToModel(format!(
            "`referenced_image_paths` must contain at most {MAX_EDIT_IMAGES} paths"
        )));
    }
    let images = match (paths.is_empty(), args.num_last_images_to_include) {
        (true, None) => {
            return 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()),
            }));
        }
        (false, None) => {
            let Some(environment) = environments.first() else {
                return Err(FunctionCallError::RespondToModel(
                    "referenced image paths are unavailable in this session".to_string(),
                ));
            };
            let mut images = Vec::with_capacity(paths.len());
            for path in paths {
                images.push(image_url(path, environment).await?);
            }
            images
        }
        (true, Some(count)) => {
            if !(1..=MAX_EDIT_IMAGES).contains(&count) {
                return Err(FunctionCallError::RespondToModel(format!(
                    "`num_last_images_to_include` must be between 1 and {MAX_EDIT_IMAGES}"
                )));
            }
            // Pathless images have no stable reference, so this bounded window may include newer
            // unrelated images. This remains best-effort until the harness provides stable refs.
            let images = recent_images(history, count);
            if images.len() != count {
                return Err(FunctionCallError::RespondToModel(format!(
                    "requested the last {count} conversation images, but only {} were available",
                    images.len()
                )));
            }
            images
        }
        (false, Some(_)) => {
            return Err(FunctionCallError::RespondToModel(
                "provide only one of `referenced_image_paths` or \
                 `num_last_images_to_include`"
                    .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()),
    }))
}

fn recent_images(history: &[ResponseItem], count: usize) -> Vec<ImageUrl> {
    let mut function_call_ids = HashSet::new();
    let mut custom_tool_call_ids = HashSet::new();
    for item in history {
        match item {
            ResponseItem::FunctionCall { call_id, .. } => {
                function_call_ids.insert(call_id.as_str());
            }
            ResponseItem::CustomToolCall { call_id, .. } => {
                custom_tool_call_ids.insert(call_id.as_str());
            }
            ResponseItem::AdditionalTools { .. }
            | ResponseItem::Message { .. }
            | ResponseItem::AgentMessage { .. }
            | ResponseItem::Reasoning { .. }
            | ResponseItem::LocalShellCall { .. }
            | ResponseItem::ToolSearchCall { .. }
            | ResponseItem::FunctionCallOutput { .. }
            | ResponseItem::CustomToolCallOutput { .. }
            | ResponseItem::ToolSearchOutput { .. }
            | ResponseItem::WebSearchCall { .. }
            | ResponseItem::ImageGenerationCall { .. }
            | ResponseItem::Compaction { .. }
            | ResponseItem::CompactionTrigger { .. }
            | ResponseItem::ContextCompaction { .. }
            | ResponseItem::Other => {}
        }
    }

    let mut images = Vec::with_capacity(count);
    'history: for item in history.iter().rev() {
        let mut image_urls = Vec::new();
        match item {
            ResponseItem::Message { content, .. } => {
                image_urls.extend(content.iter().rev().filter_map(|item| match item {
                    ContentItem::InputImage { image_url, .. } => Some(image_url.clone()),
                    ContentItem::InputText { .. } | ContentItem::OutputText { .. } => None,
                }));
            }
            ResponseItem::FunctionCallOutput {
                call_id, output, ..
            } if function_call_ids.contains(call_id.as_str()) => {
                image_urls.extend(output_image_urls(output));
            }
            ResponseItem::CustomToolCallOutput {
                call_id, output, ..
            } if custom_tool_call_ids.contains(call_id.as_str()) => {
                image_urls.extend(output_image_urls(output));
            }
            ResponseItem::ImageGenerationCall { result, .. } if !result.is_empty() => {
                image_urls.push(format!("data:image/png;base64,{result}"));
            }
            ResponseItem::AdditionalTools { .. }
            | ResponseItem::Reasoning { .. }
            | ResponseItem::AgentMessage { .. }
            | ResponseItem::LocalShellCall { .. }
            | ResponseItem::FunctionCall { .. }
            | ResponseItem::ToolSearchCall { .. }
            | ResponseItem::CustomToolCall { .. }
            | ResponseItem::FunctionCallOutput { .. }
            | ResponseItem::CustomToolCallOutput { .. }
            | ResponseItem::ToolSearchOutput { .. }
            | ResponseItem::WebSearchCall { .. }
            | ResponseItem::ImageGenerationCall { .. }
            | ResponseItem::Compaction { .. }
            | ResponseItem::CompactionTrigger { .. }
            | ResponseItem::ContextCompaction { .. }
            | ResponseItem::Other => {}
        }
        for image_url in image_urls {
            images.push(ImageUrl { image_url });
            if images.len() == count {
                break 'history;
            }
        }
    }
    images.reverse();
    images
}

/// Extracts image URLs from a tool output in newest-first order.
fn output_image_urls(output: &FunctionCallOutputPayload) -> impl Iterator<Item = String> + '_ {
    output
        .content_items()
        .into_iter()
        .flatten()
        .rev()
        .filter_map(|item| match item {
            FunctionCallOutputContentItem::InputImage { image_url, .. } => Some(image_url.clone()),
            FunctionCallOutputContentItem::InputText { .. }
            | FunctionCallOutputContentItem::EncryptedContent { .. } => None,
        })
}

async fn image_url(
    path: &AbsolutePathBuf,
    environment: &ToolEnvironment,
) -> Result<ImageUrl, FunctionCallError> {
    let path_uri = PathUri::from_abs_path(path);
    let sandbox = environment.file_system_sandbox_context.clone();
    let bytes = environment
        .file_system
        .read_file(&path_uri, Some(&sandbox))
        .await
        .map_err(|error| {
            FunctionCallError::RespondToModel(format!(
                "unable to read referenced image at `{}`: {error}",
                path.display()
            ))
        })?;
    let image = load_for_prompt_bytes(path.as_path(), bytes, PromptImageMode::Original).map_err(
        |error| {
            FunctionCallError::RespondToModel(format!(
                "unable to process referenced image at `{}`: {error}",
                path.display()
            ))
        },
    )?;
    Ok(ImageUrl {
        image_url: image.into_data_url(),
    })
}

/// 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()))
}

/// 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 the object consumed by the code-mode `generatedImage()` helper.
    fn code_mode_result(&self, _payload: &ToolPayload) -> Value {
        let mut result = Map::from_iter([(
            "image_url".to_string(),
            Value::String(format!("data:image/png;base64,{}", self.result)),
        )]);
        if let Some(output_hint) = &self.output_hint {
            result.insert(
                "output_hint".to_string(),
                Value::String(output_hint.clone()),
            );
        }
        Value::Object(result)
    }

    /// Returns generated bytes and persisted-artifact context for model follow-up.
    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;