diff --git a/codex-rs/core/src/session/inject.rs b/codex-rs/core/src/session/inject.rs index 8264d23886..d66602b205 100644 --- a/codex-rs/core/src/session/inject.rs +++ b/codex-rs/core/src/session/inject.rs @@ -1,6 +1,7 @@ use super::TurnInput as PendingTurnInput; use super::session::Session; use super::turn_context::TurnContext; +use codex_analytics::ImagePreparationMetadata; use codex_features::Feature; use codex_history::CodexHarnessMetadata; use codex_history::ResponseItemEnvelope; @@ -124,6 +125,24 @@ impl Session { return; } + let (annotated_items, image_preparations) = self + .prepare_annotated_conversation_items_for_history(turn_context, model_info, items) + .await; + self.record_prepared_conversation_items( + turn_context, + model_info, + annotated_items, + image_preparations, + ) + .await; + } + + pub(super) async fn prepare_annotated_conversation_items_for_history( + &self, + turn_context: &TurnContext, + model_info: &ModelInfo, + items: Vec, + ) -> (Vec, Vec) { let mut annotated_items = Vec::with_capacity(items.len()); let mut image_preparations = Vec::new(); for envelope in items { @@ -144,13 +163,7 @@ impl Session { } })); } - self.record_prepared_conversation_items( - turn_context, - model_info, - annotated_items, - image_preparations, - ) - .await; + (annotated_items, image_preparations) } /// Injects items into active work, or records them without starting a turn. diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 3b22567947..7a19a9078d 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -122,6 +122,7 @@ use codex_protocol::models::BaseInstructions; use codex_protocol::models::BaseInstructionsProvenance; use codex_protocol::models::ContentItem; use codex_protocol::models::ContentItemKind; +use codex_protocol::models::ImageReference; use codex_protocol::models::InternalChatMessageMetadataPassthrough; use codex_protocol::models::PermissionProfile; use codex_protocol::models::SandboxEnforcement; @@ -3548,9 +3549,18 @@ impl Session { } pub(crate) fn response_item_from_user_input(&self, input: Vec) -> ResponseItem { + self.response_item_from_user_input_with_image_positions(input, &mut HashMap::new()) + } + + fn response_item_from_user_input_with_image_positions( + &self, + input: Vec, + user_image_content_indices: &mut HashMap, + ) -> ResponseItem { let mut item = ResponseItem::from(ResponseInputItem::from_user_input( input, LocalImagePreparation::Defer, + user_image_content_indices, )); if let ResponseItem::Message { content, @@ -4843,20 +4853,37 @@ impl Session { // Persist the user message to history, but emit the turn item from `UserInput` so // UI-only `text_elements` are preserved. `ResponseItem::Message` does not carry // those spans, and `record_response_item_and_emit_turn_item` would drop them. - let response_item = self.response_item_from_user_input(input.to_vec()); - self.record_annotated_conversation_items( + let mut user_image_content_indices = HashMap::new(); + let response_item = self.response_item_from_user_input_with_image_positions( + input.to_vec(), + &mut user_image_content_indices, + ); + let (prepared_items, image_preparations) = self + .prepare_annotated_conversation_items_for_history( + turn_context, + model_info, + vec![ResponseItemEnvelope { + item: response_item, + metadata: acceptance_order.map(|order| CodexHarnessMetadata { + user_input_order: Some(order), + ..Default::default() + }), + }], + ) + .await; + let mut user_message_item = UserMessageItem::new(input); + apply_prepared_image_file_ids( + &mut user_message_item, + &prepared_items, + &user_image_content_indices, + ); + self.record_prepared_conversation_items( turn_context, model_info, - vec![ResponseItemEnvelope { - item: response_item, - metadata: acceptance_order.map(|order| CodexHarnessMetadata { - user_input_order: Some(order), - ..Default::default() - }), - }], + prepared_items, + image_preparations, ) .await; - let mut user_message_item = UserMessageItem::new(input); user_message_item.client_id = client_id; let turn_item = TurnItem::UserMessage(user_message_item); self.emit_turn_item_started(turn_context, &turn_item).await; @@ -4943,6 +4970,37 @@ impl Session { } } +fn apply_prepared_image_file_ids( + user_message_item: &mut UserMessageItem, + prepared_items: &[ResponseItemEnvelope], + user_image_content_indices: &HashMap, +) { + // Preparation replaces image slots in place; resize notices are separate messages. + // Read the references from the original message while retaining UI-only input spans. + if let Some(ResponseItemEnvelope { + item: ResponseItem::Message { content, .. }, + .. + }) = prepared_items.first() + { + for (&input_index, &content_index) in user_image_content_indices { + if let Some(ContentItem::InputImage { + image: ImageReference::File { file_id }, + .. + }) = content.get(content_index) + && let UserInput::Image { detail, .. } | UserInput::LocalImage { detail, .. } = + &user_message_item.content[input_index] + { + user_message_item.content[input_index] = UserInput::Image { + image: ImageReference::File { + file_id: file_id.clone(), + }, + detail: *detail, + }; + } + } + } +} + pub(crate) fn emit_subagent_session_started( analytics_events_client: &AnalyticsEventsClient, client_metadata: AppServerClientMetadata, diff --git a/codex-rs/core/tests/suite/image_rollout.rs b/codex-rs/core/tests/suite/image_rollout.rs index fb749dc813..806f04bc81 100644 --- a/codex-rs/core/tests/suite/image_rollout.rs +++ b/codex-rs/core/tests/suite/image_rollout.rs @@ -15,6 +15,7 @@ use codex_history::RolloutItem; use codex_protocol::config_types::CollaborationMode; use codex_protocol::config_types::ModeKind; use codex_protocol::config_types::Settings; +use codex_protocol::items::TurnItem; use codex_protocol::models::ContentItem; use codex_protocol::models::DEFAULT_IMAGE_DETAIL; use codex_protocol::models::ImageDetail; @@ -25,7 +26,10 @@ use codex_protocol::models::ResponseItem; use codex_protocol::protocol::AskForApproval; use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::Op; +use codex_protocol::protocol::ThreadHistoryMode; use codex_protocol::protocol::ThreadSettingsOverrides; +use codex_protocol::user_input::ByteRange; +use codex_protocol::user_input::TextElement; use codex_protocol::user_input::UserInput; use codex_utils_image::data_url_from_bytes; use core_test_support::TempDirExt; @@ -43,6 +47,7 @@ use core_test_support::test_codex::local_selections; use core_test_support::test_codex::test_codex; use core_test_support::test_codex::turn_permission_fields; use core_test_support::wait_for_event; +use core_test_support::wait_for_event_match; use image::GenericImageView; use image::ImageBuffer; use image::Rgba; @@ -406,6 +411,166 @@ async fn file_image_passes_through_request_and_rollout() -> anyhow::Result<()> { Ok(()) } +/// Uploaded images keep their original input positions and text spans in live and durable display +/// history, even when earlier images fail preparation or expand into multiple model content items. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn uploaded_images_preserve_user_message_display_history() -> anyhow::Result<()> { + skip_if_no_network!(Ok(())); + + let server = start_mock_server().await; + let test = test_codex() + .with_image_store(Arc::new(RecordingFileAttachmentStore::default())) + .with_history_mode(ThreadHistoryMode::Paginated) + .build_with_auto_env(&server) + .await?; + let image_path = test.cwd.path().join("display-image.png"); + write_test_png(&image_path, [12, 34, 56, 255])?; + let image_url = data_url_from_bytes("image/png", &fs::read(&image_path)?); + let input = vec![ + UserInput::Text { + text: "inspect these images".to_string(), + text_elements: vec![TextElement::new( + ByteRange { start: 0, end: 7 }, + Some("".to_string()), + )], + }, + UserInput::LocalImage { + path: test.cwd.path().join("missing-image.png"), + detail: None, + }, + UserInput::Image { + image: ImageReference::Inline { + image_url: "data:image/png;base64,bm90IGFuIGltYWdl".to_string(), + }, + detail: None, + }, + UserInput::LocalImage { + path: image_path, + detail: Some(ImageDetail::High), + }, + UserInput::Image { + image: ImageReference::Inline { + image_url: image_url.clone(), + }, + detail: None, + }, + UserInput::Text { + text: "and this copy".to_string(), + text_elements: Vec::new(), + }, + UserInput::Image { + image: ImageReference::File { + file_id: "file_existing".to_string(), + }, + detail: Some(ImageDetail::Original), + }, + UserInput::Image { + image: ImageReference::Inline { image_url }, + detail: None, + }, + ]; + let mut expected_input = input.clone(); + for index in [3, 4, 7] { + let detail = match &input[index] { + UserInput::Image { detail, .. } | UserInput::LocalImage { detail, .. } => *detail, + _ => unreachable!("these inputs are images"), + }; + expected_input[index] = UserInput::Image { + image: ImageReference::File { + file_id: UPLOADED_FILE_ID.to_string(), + }, + detail, + }; + } + let response_mock = responses::mount_sse_once( + &server, + sse(vec![ + ev_response_created("resp-display"), + ev_completed("resp-display"), + ]), + ) + .await; + test.codex + .start_or_steer_turn(TurnInputRequest::user_input(input)) + .await?; + + let started = wait_for_event_match(&test.codex, |event| match event { + EventMsg::ItemStarted(event) => match &event.item { + TurnItem::UserMessage(item) => Some(item.clone()), + _ => None, + }, + _ => None, + }) + .await; + let completed = wait_for_event_match(&test.codex, |event| match event { + EventMsg::ItemCompleted(event) => match &event.item { + TurnItem::UserMessage(item) => Some(item.clone()), + _ => None, + }, + _ => None, + }) + .await; + assert_eq!(started.content, expected_input); + assert_eq!(completed.id, started.id); + assert_eq!(completed.client_id, started.client_id); + assert_eq!(completed.content, started.content); + wait_for_event(&test.codex, |event| { + matches!(event, EventMsg::TurnComplete(_)) + }) + .await; + + let request = response_mock.single_request(); + let mut model_file_ids = Vec::new(); + for item in request.input() { + if item.get("role").and_then(Value::as_str) != Some("user") { + continue; + } + if let Some(content) = item.get("content").and_then(Value::as_array) { + for content_item in content { + if content_item.get("type").and_then(Value::as_str) == Some("input_image") { + model_file_ids.push( + content_item + .get("file_id") + .and_then(Value::as_str) + .context("prepared image file ID")? + .to_string(), + ); + } + } + } + } + assert_eq!( + model_file_ids, + vec![ + UPLOADED_FILE_ID, + UPLOADED_FILE_ID, + "file_existing", + UPLOADED_FILE_ID, + ] + ); + + test.codex.shutdown_and_wait().await?; + let rollout_path = test.codex.rollout_path().context("rollout path")?; + let rollout_text = read_rollout_text(&rollout_path).await?; + let mut persisted_user_items = Vec::new(); + for line in rollout_text.lines() { + if let RolloutItem::EventMsg(EventMsg::ItemCompleted(event)) = + codex_rollout::parse_rollout_line(line)?.item + && let TurnItem::UserMessage(item) = event.item + { + persisted_user_items.push(item); + } + } + let [persisted] = persisted_user_items.as_slice() else { + panic!("expected exactly one persisted user message"); + }; + assert_eq!(persisted.id, completed.id); + assert_eq!(persisted.client_id, completed.client_id); + assert_eq!(persisted.content, completed.content); + + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn resumed_history_only_emits_resize_notices_for_new_images() -> anyhow::Result<()> { skip_if_no_network!(Ok(())); diff --git a/codex-rs/protocol/src/models.rs b/codex-rs/protocol/src/models.rs index ee7fecbf06..d51ded354c 100644 --- a/codex-rs/protocol/src/models.rs +++ b/codex-rs/protocol/src/models.rs @@ -1993,78 +1993,91 @@ pub enum ReasoningItemContent { impl From> for ResponseInputItem { fn from(items: Vec) -> Self { - Self::from_user_input(items, LocalImagePreparation::Process) + Self::from_user_input(items, LocalImagePreparation::Process, &mut HashMap::new()) } } impl ResponseInputItem { + /// Records original user-input indices to image slots in the generated message content. + /// Inputs that do not produce an image, including failed local reads, have no entry. pub fn from_user_input( items: Vec, local_image_preparation: LocalImagePreparation, + user_image_content_indices: &mut HashMap, ) -> Self { + user_image_content_indices.clear(); let mut image_index = 0; let mut audio_index = 0; - Self::Message { - role: "user".to_string(), - content: items - .into_iter() - .flat_map(|c| match c { - UserInput::Text { text, .. } => vec![ContentItem::InputText { text }], - UserInput::Image { image, detail, .. } => { - image_index += 1; - let detail = detail.unwrap_or(DEFAULT_IMAGE_DETAIL); - vec![ContentItem::InputImage { - image, - detail: Some(detail), - }] - } - UserInput::LocalImage { path, detail, .. } => { - image_index += 1; - let detail = detail.unwrap_or(DEFAULT_IMAGE_DETAIL); - match std::fs::read(&path) { - Ok(file_bytes) => match local_image_preparation { - LocalImagePreparation::Process => { - local_image_content_items_with_label_number( - &path, - file_bytes, - Some(image_index), - detail, - ) - } - LocalImagePreparation::Defer => local_image_content_items( + let mut content = Vec::new(); + for (input_index, input) in items.into_iter().enumerate() { + let generated = match input { + UserInput::Text { text, .. } => vec![ContentItem::InputText { text }], + UserInput::Image { image, detail, .. } => { + image_index += 1; + let detail = detail.unwrap_or(DEFAULT_IMAGE_DETAIL); + vec![ContentItem::InputImage { + image, + detail: Some(detail), + }] + } + UserInput::LocalImage { path, detail, .. } => { + image_index += 1; + let detail = detail.unwrap_or(DEFAULT_IMAGE_DETAIL); + match std::fs::read(&path) { + Ok(file_bytes) => match local_image_preparation { + LocalImagePreparation::Process => { + local_image_content_items_with_label_number( &path, - data_url_from_bytes("application/octet-stream", &file_bytes), + file_bytes, Some(image_index), detail, - ), - }, - Err(err) => vec![local_media_error_placeholder( - &path, - err, - LocalMediaKind::Image, - )], - } - } - UserInput::Audio { audio_url } => { - audio_index += 1; - vec![ContentItem::InputAudio { audio_url }] - } - UserInput::LocalAudio { path } => { - audio_index += 1; - match std::fs::read(&path) { - Ok(file_bytes) => { - local_audio_content_items(&path, &file_bytes, audio_index) + ) } - Err(err) => vec![local_media_error_placeholder( + LocalImagePreparation::Defer => local_image_content_items( &path, - err, - LocalMediaKind::Audio, - )], - } + data_url_from_bytes("application/octet-stream", &file_bytes), + Some(image_index), + detail, + ), + }, + Err(err) => vec![local_media_error_placeholder( + &path, + err, + LocalMediaKind::Image, + )], } - UserInput::Skill { .. } | UserInput::Mention { .. } => Vec::new(), // Tool bodies are injected later in core - }) - .collect::>(), + } + UserInput::Audio { audio_url } => { + audio_index += 1; + vec![ContentItem::InputAudio { audio_url }] + } + UserInput::LocalAudio { path } => { + audio_index += 1; + match std::fs::read(&path) { + Ok(file_bytes) => { + local_audio_content_items(&path, &file_bytes, audio_index) + } + Err(err) => vec![local_media_error_placeholder( + &path, + err, + LocalMediaKind::Audio, + )], + } + } + UserInput::Skill { .. } | UserInput::Mention { .. } => Vec::new(), // Tool bodies are injected later in core + }; + // Local images add framing text, while mentions and skills add no content. + // Capture the image's actual slot before preparation replaces its reference. + for (offset, item) in generated.iter().enumerate() { + if matches!(item, ContentItem::InputImage { .. }) { + user_image_content_indices.insert(input_index, content.len() + offset); + } + } + content.extend(generated); + } + Self::Message { + role: "user".to_string(), + content, phase: None, } } @@ -4259,6 +4272,61 @@ mod tests { Ok(()) } + /// Image associations use original input and expanded content positions, so omitted inputs, + /// local-image framing, failed reads, and duplicate URLs cannot shift a later substitution. + #[test] + fn user_input_image_indices_survive_content_expansion() -> Result<()> { + let dir = tempdir()?; + let local_path = dir.path().join("local.png"); + std::fs::write(&local_path, TINY_PNG_BYTES)?; + let inline_image = UserInput::Image { + image: ImageReference::Inline { + image_url: data_url_from_bytes("image/png", TINY_PNG_BYTES), + }, + detail: None, + }; + let inputs = vec![ + UserInput::Text { + text: "compare images".to_string(), + text_elements: Vec::new(), + }, + UserInput::Mention { + name: "app".to_string(), + path: "app://example".to_string(), + }, + UserInput::Skill { + name: "skill".to_string(), + path: dir.path().join("SKILL.md"), + }, + UserInput::LocalImage { + path: dir.path().join("missing.png"), + detail: None, + }, + UserInput::LocalImage { + path: local_path, + detail: None, + }, + inline_image.clone(), + inline_image, + ]; + for preparation in [LocalImagePreparation::Defer, LocalImagePreparation::Process] { + let mut indices = HashMap::new(); + let item = + ResponseInputItem::from_user_input(inputs.clone(), preparation, &mut indices); + assert_eq!(indices, HashMap::from([(4, 3), (5, 5), (6, 6)])); + let ResponseInputItem::Message { content, .. } = item else { + panic!("expected a user message"); + }; + for content_index in indices.into_values() { + assert!(matches!( + content[content_index], + ContentItem::InputImage { .. } + )); + } + } + Ok(()) + } + #[test] fn mixed_remote_and_local_images_share_label_sequence() -> Result<()> { let image_url = "data:image/png;base64,abc".to_string();