diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index ec5bfebf7b92..8b62cb5210d7 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -3106,6 +3106,7 @@ dependencies = [ "codex-protocol", "codex-tools", "codex-utils-absolute-path", + "codex-utils-image", "http 1.4.0", "pretty_assertions", "schemars 0.8.22", 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 b6593bc820f1..c73b5593c6cb 100644 --- a/codex-rs/app-server/tests/suite/v2/imagegen_extension.rs +++ b/codex-rs/app-server/tests/suite/v2/imagegen_extension.rs @@ -29,6 +29,11 @@ use wiremock::matchers::method; use wiremock::matchers::path; const RESULT: &str = "cG5n"; +const TINY_PNG_BYTES: &[u8] = &[ + 137, 80, 78, 71, 13, 10, 26, 10, 0, 0, 0, 13, 73, 72, 68, 82, 0, 0, 0, 1, 0, 0, 0, 1, 8, 6, 0, + 0, 0, 31, 21, 196, 137, 0, 0, 0, 11, 73, 68, 65, 84, 120, 156, 99, 96, 0, 2, 0, 0, 5, 0, 1, + 122, 94, 171, 63, 0, 0, 0, 0, 73, 69, 78, 68, 174, 66, 96, 130, +]; #[derive(Clone, Copy)] enum ImagegenTestMode { @@ -59,7 +64,6 @@ async fn standalone_image_generation_returns_saved_path_hint_to_model() -> Resul "image_gen", "imagegen", &json!({ - "action": "generate", "prompt": "paint a blue whale", }) .to_string(), @@ -142,6 +146,67 @@ async fn standalone_image_generation_returns_saved_path_hint_to_model() -> Resul Ok(()) } +#[tokio::test] +async fn standalone_image_edit_uses_attached_model_visible_image() -> Result<()> { + let edit_request = run_image_edit_test(|codex_home| { + let image_path = codex_home.join("attached.png"); + std::fs::write(&image_path, TINY_PNG_BYTES)?; + Ok(( + json!({ + "prompt": "add a red hat", + "referenced_image_paths": [image_path.display().to_string()], + }), + vec![ + V2UserInput::Text { + text: "Edit the attached image".to_string(), + text_elements: Vec::new(), + }, + V2UserInput::LocalImage { + path: image_path, + detail: None, + }, + ], + )) + }) + .await?; + assert_eq!(edit_request["prompt"], "add a red hat"); + assert!( + edit_request["images"][0]["image_url"] + .as_str() + .is_some_and(|image_url| image_url.starts_with("data:image/png;base64,")) + ); + + Ok(()) +} + +#[tokio::test] +async fn standalone_image_edit_uses_recent_pathless_image() -> Result<()> { + let image_url = "https://example.com/reference.png"; + let edit_request = run_image_edit_test(|_| { + Ok(( + json!({ + "prompt": "add a red hat", + "num_last_images_to_include": 1, + }), + vec![ + V2UserInput::Text { + text: "Edit the attached image".to_string(), + text_elements: Vec::new(), + }, + V2UserInput::Image { + url: image_url.to_string(), + detail: None, + }, + ], + )) + }) + .await?; + assert_eq!(edit_request["prompt"], "add a red hat"); + assert_eq!(edit_request["images"][0]["image_url"], image_url); + + Ok(()) +} + #[tokio::test] async fn standalone_image_generation_is_exposed_in_code_mode_only() -> Result<()> { let server = responses::start_mock_server().await; @@ -202,7 +267,6 @@ async fn standalone_image_generation_is_callable_from_code_mode_only() -> Result "exec", r#" const result = await tools.image_gen__imagegen({ - action: "generate", prompt: "paint a blue whale", }); generatedImage(result); @@ -263,6 +327,81 @@ generatedImage(result); } async fn start_image_generation_turn(mcp: &mut TestAppServer) -> Result<()> { + start_turn( + mcp, + vec![V2UserInput::Text { + text: "Generate an image".to_string(), + text_elements: Vec::new(), + }], + ) + .await +} + +async fn run_image_edit_test( + input: impl FnOnce(&Path) -> Result<(serde_json::Value, Vec)>, +) -> Result { + let call_id = "image-edit-1"; + let server = responses::start_mock_server().await; + mount_image_edit_response(&server).await; + + let codex_home = TempDir::new()?; + let (arguments, input) = input(codex_home.path())?; + 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", + &arguments.to_string(), + ), + responses::ev_completed("resp-1"), + ]), + responses::sse(vec![ + responses::ev_assistant_message("msg-1", "Done"), + responses::ev_completed("resp-2"), + ]), + ], + ) + .await; + + create_config_toml(codex_home.path(), &server.uri(), ImagegenTestMode::Direct)?; + 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_turn(&mut mcp, input).await?; + 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!(response_mock.requests().len(), 2); + let requests = server + .received_requests() + .await + .context("failed to fetch received requests")?; + Ok(requests + .iter() + .find(|request| request.url.path() == "/api/codex/images/edits") + .context("image edit request should be sent")? + .body_json::()?) +} + +async fn start_turn(mcp: &mut TestAppServer, input: Vec) -> Result<()> { let thread_req = mcp .send_thread_start_request(ThreadStartParams::default()) .await?; @@ -277,10 +416,7 @@ async fn start_image_generation_turn(mcp: &mut TestAppServer) -> Result<()> { .send_turn_start_request(TurnStartParams { thread_id: thread.id, client_user_message_id: None, - input: vec![V2UserInput::Text { - text: "Generate an image".to_string(), - text_elements: Vec::new(), - }], + input, ..Default::default() }) .await?; @@ -324,6 +460,18 @@ async fn mount_image_response(server: &MockServer) { .await; } +async fn mount_image_edit_response(server: &MockServer) { + Mock::given(method("POST")) + .and(path("/api/codex/images/edits")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "created": 1, + "data": [{"b64_json": RESULT}], + }))) + .expect(1) + .mount(server) + .await; +} + fn create_config_toml( codex_home: &Path, server_uri: &str, diff --git a/codex-rs/ext/image-generation/Cargo.toml b/codex-rs/ext/image-generation/Cargo.toml index c6b87e1d4fd0..cb5ddda0a816 100644 --- a/codex-rs/ext/image-generation/Cargo.toml +++ b/codex-rs/ext/image-generation/Cargo.toml @@ -23,6 +23,7 @@ codex-model-provider-info = { workspace = true } codex-protocol = { workspace = true } codex-tools = { workspace = true } codex-utils-absolute-path = { workspace = true } +codex-utils-image = { workspace = true } http = { workspace = true } schemars = { workspace = true } serde = { workspace = true, features = ["derive"] } diff --git a/codex-rs/ext/image-generation/imagegen_description.md b/codex-rs/ext/image-generation/imagegen_description.md index 1e5a46390f7b..53e662013534 100644 --- a/codex-rs/ext/image-generation/imagegen_description.md +++ b/codex-rs/ext/image-generation/imagegen_description.md @@ -5,8 +5,13 @@ The `image_gen.imagegen` tool enables image generation from descriptions and edi Guidelines: - In code mode, pass the result to `generatedImage(result)`. -- Set `action` to `generate` when the user asks for a brand new image. -- Set `action` to `edit` when the user asks to modify an existing image from the conversation history. -- Directly generate the image without reconfirmation or clarification. +- Omit both `referenced_image_paths` and `num_last_images_to_include` when generating a brand new image. +- For edits, use `referenced_image_paths` when every target image has a local file path. +- If you have not seen a local image yet, use `view_image` to inspect it before editing. +- Use `num_last_images_to_include` only when at least one target image has no local file path. +- Set `num_last_images_to_include` to the smallest number of recent conversation images that includes every target image, up to 5. +- Never provide both `referenced_image_paths` and `num_last_images_to_include`. +- If neither mechanism can include every target image, ask the user to attach the missing images again. +- Directly generate the image without reconfirmation or clarification 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 followup 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/tests.rs b/codex-rs/ext/image-generation/src/tests.rs index 6a974e2d512c..3f729726562a 100644 --- a/codex-rs/ext/image-generation/src/tests.rs +++ b/codex-rs/ext/image-generation/src/tests.rs @@ -19,10 +19,9 @@ use pretty_assertions::assert_eq; use super::GeneratedImageOutput; use super::ImageRequest; -use super::ImagegenAction; use super::ImagegenArgs; use super::imagegen_tool_spec; -use super::request_for_action; +use super::request_for_args; use crate::IMAGE_GEN_NAMESPACE; use crate::IMAGEGEN_TOOL_NAME; @@ -39,10 +38,17 @@ fn uses_reserved_image_gen_namespace() { } #[test] -fn generate_uses_fixed_request_defaults() { +fn omitted_references_generate_with_fixed_defaults() { assert_eq!( - request_for_action(&args(ImagegenAction::Generate, "paint a moonlit lake"), &[]) - .expect("generation request should build"), + request_for_args( + &ImagegenArgs { + prompt: "paint a moonlit lake".to_string(), + referenced_image_paths: None, + num_last_images_to_include: None, + }, + &[] + ) + .expect("generation request should build"), ImageRequest::Generate(ImageGenerationRequest { prompt: "paint a moonlit lake".to_string(), background: Some(ImageBackground::Auto), @@ -54,6 +60,144 @@ fn generate_uses_fixed_request_defaults() { ); } +#[test] +fn recent_image_fallback_selects_newest_images_in_chronological_order() { + let history = vec![ + ResponseItem::Message { + id: None, + role: "user".to_string(), + content: vec![ + input_image("user-1"), + input_image("user-2"), + ContentItem::InputText { + text: "edit these".to_string(), + }, + ], + phase: None, + }, + ResponseItem::FunctionCall { + id: None, + name: "mcp_image".to_string(), + namespace: None, + arguments: "{}".to_string(), + call_id: "mcp-call".to_string(), + }, + ResponseItem::FunctionCallOutput { + call_id: "mcp-call".to_string(), + output: image_output("mcp"), + }, + ResponseItem::CustomToolCall { + id: None, + status: Some("completed".to_string()), + call_id: "code-mode-call".to_string(), + name: "exec".to_string(), + input: String::new(), + }, + ResponseItem::CustomToolCallOutput { + call_id: "code-mode-call".to_string(), + name: Some("exec".to_string()), + output: image_output("code-mode"), + }, + ResponseItem::ImageGenerationCall { + id: "generated-call".to_string(), + status: "completed".to_string(), + revised_prompt: None, + result: "generated".to_string(), + }, + ResponseItem::FunctionCallOutput { + call_id: "orphan-call".to_string(), + output: image_output("orphan"), + }, + ]; + + assert_eq!( + request_for_args( + &ImagegenArgs { + prompt: "change the lighting".to_string(), + referenced_image_paths: None, + num_last_images_to_include: Some(4), + }, + &history, + ) + .expect("history-backed edit request should build"), + ImageRequest::Edit(expected_edit_request( + "change the lighting", + &["user-2", "mcp", "code-mode", "generated"], + )) + ); +} + +#[test] +fn conflicting_image_selectors_return_tool_error() { + let error = request_for_args( + &ImagegenArgs { + prompt: "change the lighting".to_string(), + referenced_image_paths: Some(vec![ + "/tmp/image.png" + .try_into() + .expect("test path should be absolute"), + ]), + num_last_images_to_include: Some(1), + }, + &[], + ) + .expect_err("conflicting selectors should fail"); + + assert_eq!( + error.to_string(), + "provide only one of `referenced_image_paths` or `num_last_images_to_include`" + ); +} + +#[test] +fn too_many_referenced_image_paths_return_tool_error() { + let error = request_for_args( + &ImagegenArgs { + prompt: "change the lighting".to_string(), + referenced_image_paths: Some( + (0..6) + .map(|index| { + format!("/tmp/image-{index}.png") + .try_into() + .expect("test path should be absolute") + }) + .collect(), + ), + num_last_images_to_include: None, + }, + &[], + ) + .expect_err("too many paths should fail before reading files"); + + assert_eq!( + error.to_string(), + "`referenced_image_paths` must contain at most 5 paths" + ); +} + +#[test] +fn recent_image_fallback_requires_requested_count() { + let error = request_for_args( + &ImagegenArgs { + prompt: "change the lighting".to_string(), + referenced_image_paths: None, + num_last_images_to_include: Some(2), + }, + &[ResponseItem::Message { + id: None, + role: "user".to_string(), + content: vec![input_image("only-image")], + phase: None, + }], + ) + .expect_err("history-backed edit should require the requested image count"); + + assert_eq!( + error.to_string(), + "requested the last 2 conversation images, but only 1 were available" + ); +} + #[test] fn generated_output_returns_image_input_and_output_hint() { let output_hint = @@ -128,195 +272,26 @@ fn generated_output_omits_oversized_output_hint() { ); } -#[test] -fn edit_matches_context_selector_for_generated_images_after_latest_user_anchor() { - let history = vec![ - generated_item("g1"), - generated_item("g2"), - generated_item("g3"), - ResponseItem::Message { - id: None, - role: "user".to_string(), - content: vec![ - ContentItem::InputImage { - image_url: "data:image/png;base64,u1".to_string(), - detail: None, - }, - ContentItem::InputImage { - image_url: "data:image/png;base64,u2".to_string(), - detail: None, - }, - ], - phase: None, - }, - generated_item("g4"), - generated_item("g5"), - generated_item("g6"), - generated_item("g7"), - ]; - - assert_eq!( - edit_request("change the lighting", &history), - expected_edit_request( - "change the lighting", - &[ - "data:image/png;base64,u1", - "data:image/png;base64,u2", - "data:image/png;base64,g5", - "data:image/png;base64,g6", - "data:image/png;base64,g7", - ] - ) - ); -} - -#[test] -fn edit_preserves_a_generated_image_when_user_anchor_fills_the_limit() { - let history = vec![ - ResponseItem::Message { - id: None, - role: "user".to_string(), - content: ["a", "b", "c", "d", "e"] - .into_iter() - .map(|image| ContentItem::InputImage { - image_url: format!("data:image/png;base64,{image}"), - detail: None, - }) - .collect(), - phase: None, - }, - generated_item("generated"), - ]; - - assert_eq!( - edit_request("edit the last generated image", &history), - expected_edit_request( - "edit the last generated image", - &[ - "data:image/png;base64,b", - "data:image/png;base64,c", - "data:image/png;base64,d", - "data:image/png;base64,e", - "data:image/png;base64,generated", - ] - ) - ); -} - -#[test] -fn edit_uses_latest_user_upload_before_a_text_only_follow_up() { - let history = vec![ - ResponseItem::Message { - id: None, - role: "user".to_string(), - content: vec![ContentItem::InputImage { - image_url: "data:image/png;base64,user".to_string(), - detail: None, - }], - phase: None, - }, - ResponseItem::Message { - id: None, - role: "user".to_string(), - content: vec![ContentItem::InputText { - text: "edit this image".to_string(), - }], - phase: None, - }, - ]; - - assert_eq!( - edit_request("change the lighting", &history), - expected_edit_request("change the lighting", &["data:image/png;base64,user"]) - ); -} - -#[test] -fn edit_reuses_images_from_prior_standalone_imagegen_calls() { - let history = vec![ - ResponseItem::FunctionCall { - id: None, - name: IMAGEGEN_TOOL_NAME.to_string(), - namespace: Some(IMAGE_GEN_NAMESPACE.to_string()), - arguments: "{}".to_string(), - call_id: "imagegen-1".to_string(), - }, - generated_function_output("imagegen-1", "standalone"), - ]; - - assert_eq!( - edit_request("change the lighting", &history), - expected_edit_request("change the lighting", &["data:image/png;base64,standalone"]) - ); -} - -#[test] -fn edit_keeps_newest_standalone_generated_images_when_over_limit() { - let history = (1..=6) - .flat_map(|index| { - let call_id = format!("imagegen-{index}"); - vec![ - ResponseItem::FunctionCall { - id: None, - name: IMAGEGEN_TOOL_NAME.to_string(), - namespace: Some(IMAGE_GEN_NAMESPACE.to_string()), - arguments: "{}".to_string(), - call_id: call_id.clone(), - }, - generated_function_output(&call_id, &index.to_string()), - ] - }) - .collect::>(); - - assert_eq!( - edit_request("change the lighting", &history), - expected_edit_request( - "change the lighting", - &[ - "data:image/png;base64,2", - "data:image/png;base64,3", - "data:image/png;base64,4", - "data:image/png;base64,5", - "data:image/png;base64,6", - ] - ) - ); -} - -#[test] -fn edit_without_image_history_returns_tool_error() { - let error = request_for_action(&args(ImagegenAction::Edit, "change the lighting"), &[]) - .expect_err("edit should require image context"); - - assert_eq!( - error.to_string(), - "image edit requested without any usable image in conversation history" - ); -} - -fn args(action: ImagegenAction, prompt: &str) -> ImagegenArgs { - ImagegenArgs { - prompt: prompt.to_string(), - action, +fn input_image(image: &str) -> ContentItem { + ContentItem::InputImage { + image_url: format!("data:image/png;base64,{image}"), + detail: None, } } -fn edit_request(prompt: &str, history: &[ResponseItem]) -> ImageEditRequest { - let ImageRequest::Edit(request) = - request_for_action(&args(ImagegenAction::Edit, prompt), history) - .expect("edit request should build") - else { - panic!("expected edit request"); - }; - request +fn image_output(image: &str) -> FunctionCallOutputPayload { + FunctionCallOutputPayload::from_content_items(vec![FunctionCallOutputContentItem::InputImage { + image_url: format!("data:image/png;base64,{image}"), + detail: None, + }]) } fn expected_edit_request(prompt: &str, images: &[&str]) -> ImageEditRequest { ImageEditRequest { images: images .iter() - .map(|image_url| ImageUrl { - image_url: (*image_url).to_string(), + .map(|image| ImageUrl { + image_url: format!("data:image/png;base64,{image}"), }) .collect(), prompt: prompt.to_string(), @@ -328,33 +303,6 @@ fn expected_edit_request(prompt: &str, images: &[&str]) -> ImageEditRequest { } } -fn generated_item(result: &str) -> ResponseItem { - ResponseItem::ImageGenerationCall { - id: format!("id-{result}"), - status: "completed".to_string(), - revised_prompt: None, - result: result.to_string(), - } -} - -fn generated_function_output(call_id: &str, result: &str) -> ResponseItem { - ResponseItem::FunctionCallOutput { - call_id: call_id.to_string(), - output: FunctionCallOutputPayload { - body: FunctionCallOutputBody::ContentItems(vec![ - FunctionCallOutputContentItem::InputImage { - image_url: format!("data:image/png;base64,{result}"), - detail: Some(DEFAULT_IMAGE_DETAIL), - }, - FunctionCallOutputContentItem::InputText { - text: "generated image save hint".to_string(), - }, - ]), - success: Some(true), - }, - } -} - fn function_payload() -> ToolPayload { ToolPayload::Function { arguments: "{}".to_string(), diff --git a/codex-rs/ext/image-generation/src/tool.rs b/codex-rs/ext/image-generation/src/tool.rs index 41c06c8007c2..69b99bdbc42e 100644 --- a/codex-rs/ext/image-generation/src/tool.rs +++ b/codex-rs/ext/image-generation/src/tool.rs @@ -1,3 +1,5 @@ +use std::collections::HashSet; + use codex_api::ImageBackground; use codex_api::ImageEditRequest; use codex_api::ImageGenerationRequest; @@ -28,6 +30,8 @@ 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 schemars::JsonSchema; use schemars::r#gen::SchemaSettings; use serde::Deserialize; @@ -68,14 +72,10 @@ impl ImageGenerationTool { #[serde(deny_unknown_fields)] struct ImagegenArgs { prompt: String, - action: ImagegenAction, -} - -#[derive(Debug, Deserialize, JsonSchema)] -#[serde(rename_all = "lowercase")] -enum ImagegenAction { - Generate, - Edit, + #[schemars(length(max = 5))] + referenced_image_paths: Option>, + #[schemars(range(min = 1, max = 5))] + num_last_images_to_include: Option, } #[async_trait::async_trait] @@ -85,7 +85,7 @@ impl ToolExecutor for ImageGenerationTool { ToolName::namespaced(IMAGE_GEN_NAMESPACE, IMAGEGEN_TOOL_NAME) } - /// Advertises the model contract: a rewritten prompt and semantic action. + /// Advertises the model contract: a rewritten prompt and optional edit references. fn spec(&self) -> ToolSpec { imagegen_tool_spec() } @@ -98,7 +98,7 @@ impl ToolExecutor for ImageGenerationTool { /// Executes the selected image operation and returns the completed image result. async fn handle(&self, call: ToolCall) -> Result, FunctionCallError> { let args = parse_args(&call)?; - let request = request_for_action(&args, call.conversation_history.items())?; + let request = request_for_args(&args, call.conversation_history.items())?; call.turn_item_emitter .emit_started(ExtensionTurnItem::ImageGeneration(ImageGenerationItem { id: call.call_id.clone(), @@ -149,106 +149,125 @@ enum ImageRequest { Edit(ImageEditRequest), } -/// Maps the model-selected action to the fixed image API request parameters. -fn request_for_action( +/// Builds a generation or edit request from the mutually exclusive image selectors. +fn request_for_args( args: &ImagegenArgs, history: &[ResponseItem], ) -> Result { - 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, + 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) => paths.iter().map(image_url).collect::, _>>()?, + (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()), + })) } -/// Selects edit context using the hosted imagegen anchor and truncation behavior. -fn edit_images(history: &[ResponseItem]) -> Vec { - 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; +/// Selects the newest requested images while preserving their conversation order. +fn recent_images(history: &[ResponseItem], count: usize) -> Vec { + 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::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 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::>(); - (!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..] { + } + + let mut images = Vec::with_capacity(count); + 'history: for item in history.iter().rev() { + let mut image_urls = Vec::new(); match item { - ResponseItem::ImageGenerationCall { result, .. } if !result.is_empty() => { - generated_images.push(ImageUrl { - image_url: format!("data:image/png;base64,{result}"), - }); + 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 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 - ) - }) => + if function_call_ids.contains(call_id.as_str()) => { - 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, - }, - )); + image_urls.extend(output_image_urls(output)); } - ResponseItem::Message { .. } + 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::Reasoning { .. } | ResponseItem::AgentMessage { .. } - | ResponseItem::Reasoning { .. } | ResponseItem::LocalShellCall { .. } | ResponseItem::FunctionCall { .. } | ResponseItem::ToolSearchCall { .. } - | ResponseItem::FunctionCallOutput { .. } | ResponseItem::CustomToolCall { .. } + | ResponseItem::FunctionCallOutput { .. } | ResponseItem::CustomToolCallOutput { .. } | ResponseItem::ToolSearchOutput { .. } | ResponseItem::WebSearchCall { .. } @@ -258,26 +277,49 @@ fn edit_images(history: &[ResponseItem]) -> Vec { | ResponseItem::ContextCompaction { .. } | ResponseItem::Other => {} } + for image_url in image_urls { + images.push(ImageUrl { image_url }); + if images.len() == count { + break 'history; + } + } } - truncate_images(user_images, generated_images) + images.reverse(); + images } -/// Truncates edit inputs while preserving the newest generated image when possible. -fn truncate_images( - mut user_images: Vec, - mut generated_images: Vec, -) -> Vec { - 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); +/// Extracts image URLs from a tool output in newest-first order. +fn output_image_urls(output: &FunctionCallOutputPayload) -> impl Iterator + '_ { + 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, + }) +} - user_images.extend(generated_images); - user_images +fn image_url(path: &AbsolutePathBuf) -> Result { + let bytes = std::fs::read(path).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.