Add opt-in prompt image sanitization (#61341)

## What changed

- Add the disabled-by-default `safe_image_decoder` feature to decode and re-encode prompt images without source metadata, with byte and decoded-pixel limits.
- Apply JPEG orientation and supported ICC color conversion to sRGB, emitting PNG for PNG/GIF inputs and JPEG for JPEG/WebP inputs.
- Use the selected decoder for image preparation, `view_image` validation, and MCP image validation. Include the decoder in image cache keys.
- Preserve original inline images in history metadata so resume and fork can regenerate prepared images without repeated JPEG encoding changing their bytes.

## Testing

Add regression tests for identical JPEG replay after history serialization and for preventing legacy cache entries from bypassing sanitization.

GitOrigin-RevId: 718cd83d0a233e2eb8efd1a64c14d8944e44c350
This commit is contained in:
Vivian Fang
2026-11-13 18:44:31 +00:00
committed by copyberry
parent a4b69a0ae9
commit 7844815bd7
21 changed files with 632 additions and 216 deletions
+1 -1
View File
@@ -1784,7 +1784,7 @@
"pkg-config_0.3.32": "{\"dependencies\":[{\"kind\":\"dev\",\"name\":\"lazy_static\",\"req\":\"^1\"}],\"features\":{}}",
"plain_0.2.3": "{\"dependencies\":[],\"features\":{}}",
"plist_1.9.0": "{\"dependencies\":[{\"name\":\"base64\",\"req\":\"^0.22.0\"},{\"name\":\"indexmap\",\"req\":\"^2.1.0\"},{\"name\":\"quick_xml\",\"package\":\"quick-xml\",\"req\":\"^0.39.2\"},{\"name\":\"serde\",\"optional\":true,\"req\":\"^1.0.2\"},{\"kind\":\"dev\",\"name\":\"serde_derive\",\"req\":\"^1.0.2\"},{\"kind\":\"dev\",\"name\":\"serde_yaml\",\"req\":\"^0.8.21\"},{\"features\":[\"parsing\",\"formatting\"],\"name\":\"time\",\"req\":\"^0.3.47\"}],\"features\":{\"default\":[\"serde\"],\"enable_unstable_features_that_may_break_with_minor_version_bumps\":[]}}",
"png_0.18.0": "{\"dependencies\":[{\"kind\":\"dev\",\"name\":\"approx\",\"req\":\"^0.5.1\"},{\"name\":\"bitflags\",\"req\":\"^2.0\"},{\"kind\":\"dev\",\"name\":\"byteorder\",\"req\":\"^1.5.0\"},{\"features\":[\"derive\"],\"kind\":\"dev\",\"name\":\"clap\",\"req\":\"^4.0\"},{\"name\":\"crc32fast\",\"req\":\"^1.2.0\"},{\"default_features\":false,\"features\":[\"cargo_bench_support\"],\"kind\":\"dev\",\"name\":\"criterion\",\"req\":\"^0.7.0\"},{\"kind\":\"dev\",\"name\":\"criterion\",\"req\":\"^0.7.0\",\"target\":\"cfg(not(target_arch = \\\"wasm32\\\"))\"},{\"name\":\"fdeflate\",\"req\":\"^0.3.3\"},{\"name\":\"flate2\",\"req\":\"^1.0.35\"},{\"kind\":\"dev\",\"name\":\"glob\",\"req\":\"^0.3\"},{\"features\":[\"simd\"],\"name\":\"miniz_oxide\",\"req\":\"^0.8\"},{\"kind\":\"dev\",\"name\":\"rand\",\"req\":\"^0.9.2\"}],\"features\":{\"benchmarks\":[],\"unstable\":[\"crc32fast/nightly\"],\"zlib-rs\":[\"flate2/zlib-rs\"]}}",
"png_0.18.1": "{\"dependencies\":[{\"kind\":\"dev\",\"name\":\"approx\",\"req\":\"^0.5.1\"},{\"name\":\"bitflags\",\"req\":\"^2.0\"},{\"kind\":\"dev\",\"name\":\"byteorder\",\"req\":\"^1.5.0\"},{\"features\":[\"derive\"],\"kind\":\"dev\",\"name\":\"clap\",\"req\":\"^4.0\"},{\"name\":\"crc32fast\",\"req\":\"^1.2.0\"},{\"default_features\":false,\"features\":[\"cargo_bench_support\"],\"kind\":\"dev\",\"name\":\"criterion\",\"req\":\"^0.7.0\"},{\"kind\":\"dev\",\"name\":\"criterion\",\"req\":\"^0.7.0\",\"target\":\"cfg(not(target_arch = \\\"wasm32\\\"))\"},{\"name\":\"fdeflate\",\"req\":\"^0.3.3\"},{\"name\":\"flate2\",\"req\":\"^1.0.35\"},{\"kind\":\"dev\",\"name\":\"glob\",\"req\":\"^0.3\"},{\"features\":[\"simd\"],\"name\":\"miniz_oxide\",\"req\":\"^0.8\"},{\"kind\":\"dev\",\"name\":\"rand\",\"req\":\"^0.9.2\"}],\"features\":{\"benchmarks\":[],\"unstable\":[\"crc32fast/nightly\"],\"zlib-rs\":[\"flate2/zlib-rs\"]}}",
"polling_3.11.0": "{\"dependencies\":[{\"name\":\"cfg-if\",\"req\":\"^1\"},{\"name\":\"concurrent-queue\",\"req\":\"^2.2.0\",\"target\":\"cfg(windows)\"},{\"kind\":\"dev\",\"name\":\"easy-parallel\",\"req\":\"^3.1.0\"},{\"kind\":\"dev\",\"name\":\"fastrand\",\"req\":\"^2.0.0\"},{\"name\":\"hermit-abi\",\"req\":\"^0.5.0\",\"target\":\"cfg(target_os = \\\"hermit\\\")\"},{\"kind\":\"dev\",\"name\":\"libc\",\"req\":\"^0.2\",\"target\":\"cfg(unix)\"},{\"name\":\"pin-project-lite\",\"req\":\"^0.2.9\",\"target\":\"cfg(windows)\"},{\"default_features\":false,\"features\":[\"event\",\"fs\",\"pipe\",\"process\",\"std\",\"time\"],\"name\":\"rustix\",\"req\":\"^1.0.5\",\"target\":\"cfg(any(unix, target_os = \\\"fuchsia\\\", target_os = \\\"vxworks\\\"))\"},{\"kind\":\"dev\",\"name\":\"signal-hook\",\"req\":\"^0.3.17\",\"target\":\"cfg(all(unix, not(target_os=\\\"vita\\\")))\"},{\"kind\":\"dev\",\"name\":\"socket2\",\"req\":\"^0.6.0\"},{\"default_features\":false,\"name\":\"tracing\",\"optional\":true,\"req\":\"^0.1.37\"},{\"features\":[\"Wdk_Foundation\",\"Wdk_Storage_FileSystem\",\"Win32_Foundation\",\"Win32_Networking_WinSock\",\"Win32_Security\",\"Win32_Storage_FileSystem\",\"Win32_System_IO\",\"Win32_System_LibraryLoader\",\"Win32_System_Threading\",\"Win32_System_WindowsProgramming\"],\"name\":\"windows-sys\",\"req\":\"^0.61\",\"target\":\"cfg(windows)\"}],\"features\":{}}",
"poly1305_0.8.0": "{\"dependencies\":[{\"name\":\"cpufeatures\",\"req\":\"^0.2\",\"target\":\"cfg(any(target_arch = \\\"x86_64\\\", target_arch = \\\"x86\\\"))\"},{\"kind\":\"dev\",\"name\":\"hex-literal\",\"req\":\"^0.3\"},{\"name\":\"opaque-debug\",\"req\":\"^0.3\"},{\"default_features\":false,\"name\":\"universal-hash\",\"req\":\"^0.5\"},{\"default_features\":false,\"name\":\"zeroize\",\"optional\":true,\"req\":\"^1\"}],\"features\":{\"std\":[\"universal-hash/std\"]}}",
"polyval_0.6.2": "{\"dependencies\":[{\"name\":\"cfg-if\",\"req\":\"^1\"},{\"name\":\"cpufeatures\",\"req\":\"^0.2\",\"target\":\"cfg(any(target_arch = \\\"aarch64\\\", target_arch = \\\"x86_64\\\", target_arch = \\\"x86\\\"))\"},{\"kind\":\"dev\",\"name\":\"hex-literal\",\"req\":\"^0.3\"},{\"name\":\"opaque-debug\",\"req\":\"^0.3\"},{\"default_features\":false,\"name\":\"universal-hash\",\"req\":\"^0.5\"},{\"default_features\":false,\"name\":\"zeroize\",\"optional\":true,\"req\":\"^1\"}],\"features\":{\"std\":[\"universal-hash/std\"]}}",
+3 -2
View File
@@ -5512,6 +5512,7 @@ dependencies = [
"divan",
"image",
"mime_guess",
"moxcms",
"png",
"thiserror 2.0.18",
"tokio",
@@ -12075,9 +12076,9 @@ dependencies = [
[[package]]
name = "png"
version = "0.18.0"
version = "0.18.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "97baced388464909d42d89643fe4361939af9b7ce7a31ee32a168f832a70f2a0"
checksum = "60769b8b31b2a9f263dae2776c37b1b28ae246943cf719eb6946a1db05128a61"
dependencies = [
"bitflags 2.13.1",
"crc32fast",
+6
View File
@@ -1316,6 +1316,9 @@
"runtime_metrics": {
"type": "boolean"
},
"safe_image_decoder": {
"type": "boolean"
},
"search_tool": {
"type": "boolean"
},
@@ -8450,6 +8453,9 @@
"runtime_metrics": {
"type": "boolean"
},
"safe_image_decoder": {
"type": "boolean"
},
"search_tool": {
"type": "boolean"
},
+1 -1
View File
@@ -148,7 +148,7 @@ pub(crate) async fn prepare_prompt(
})?;
}
restored = session
.prepare_annotated_conversation_items_for_history(step.turn.as_ref(), model, messages)
.prepare_conversation_items_for_history(step.turn.as_ref(), model, messages)
.await
.0;
prompt
@@ -22,7 +22,6 @@ use super::GuardianReviewSessionParams;
use crate::context_manager::ContextManager;
use crate::image_preparation::ImagePreparationMode;
use crate::image_preparation::resize_image;
use crate::image_preparation::unified_image_budget_enabled;
pub(super) async fn build_guardian_prompt_items(
params: &GuardianReviewSessionParams,
@@ -71,11 +70,7 @@ pub(super) async fn build_guardian_prompt_items(
});
let context_window =
context_window.filter(|limit| *limit > 0 && model_info.origin != ModelOrigin::Synthetic);
let image_mode = if unified_image_budget_enabled(&params.spawn_config.features, model_info) {
ImagePreparationMode::UnifiedBudget
} else {
ImagePreparationMode::DetailBased
};
let image_mode = ImagePreparationMode::for_model(&params.spawn_config.features, model_info);
let history_tokens = if context_window.is_some() && collected.sections.has_images() {
let base_instructions = reviewer.session.get_base_instructions().await;
reviewer_history
@@ -121,7 +116,7 @@ impl codex_guardian_context::SyncImagePreparer for ReviewImagePreparer<'_> {
detail: &mut Option<ImageDetail>,
) -> Option<Cow<'image, ImageReference>> {
*detail = match normalize_output_image_detail(self.model_info, *detail) {
_ if self.mode == ImagePreparationMode::UnifiedBudget => Some(ImageDetail::Original),
_ if self.mode.unified_budget => Some(ImageDetail::Original),
Some(ImageDetail::Low) => Some(ImageDetail::High),
detail => detail,
};
+157 -53
View File
@@ -1,4 +1,5 @@
use std::collections::HashMap;
use std::sync::Arc;
use crate::config::ManagedFeatures;
use crate::context::ContextualUserFragment;
@@ -15,6 +16,9 @@ use codex_context_fragments::AnnotatedContent;
use codex_context_fragments::set_annotated_content;
use codex_context_fragments::to_annotated_content;
use codex_features::Feature;
use codex_history::CodexHarnessMetadata;
use codex_history::ImageReplaySource;
use codex_history::ResponseItemEnvelope;
use codex_protocol::models::ContentItem;
use codex_protocol::models::ContentItemKind;
use codex_protocol::models::FunctionCallOutputContentItem;
@@ -23,8 +27,10 @@ use codex_protocol::models::ImageDetail;
use codex_protocol::models::ImageReference;
use codex_protocol::models::ResponseItem;
use codex_protocol::openai_models::ModelInfo;
use codex_utils_cache::sha1_digest;
use codex_utils_image::EncodedImage;
use codex_utils_image::ImageProcessingError;
use codex_utils_image::PromptImageDecoder;
use codex_utils_image::PromptImageMode;
use codex_utils_image::data_url_from_bytes;
use codex_utils_image::load_data_url_for_prompt;
@@ -38,10 +44,37 @@ const UNSUPPORTED_LOW_DETAIL_PLACEHOLDER: &str = "image content omitted because
const REMOTE_IMAGE_URL_PLACEHOLDER: &str =
"image content omitted because remote image URLs are not supported";
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub(crate) enum ImagePreparationMode {
DetailBased,
UnifiedBudget,
#[derive(Clone, Copy, Debug, Default)]
pub(crate) struct ImagePreparationMode {
pub(crate) unified_budget: bool,
pub(crate) decoder: PromptImageDecoder,
replay: bool,
}
impl ImagePreparationMode {
pub(crate) fn for_replay(features: &ManagedFeatures) -> Self {
Self {
unified_budget: false,
decoder: prompt_image_decoder(features),
replay: true,
}
}
pub(crate) fn for_model(features: &ManagedFeatures, model_info: &ModelInfo) -> Self {
Self {
unified_budget: unified_image_budget_enabled(features, model_info),
decoder: prompt_image_decoder(features),
replay: false,
}
}
}
pub(crate) fn prompt_image_decoder(features: &ManagedFeatures) -> PromptImageDecoder {
if features.enabled(Feature::SafeImageDecoder) {
PromptImageDecoder::Sanitizer
} else {
PromptImageDecoder::Legacy
}
}
pub(crate) fn unified_image_budget_enabled(
@@ -74,6 +107,15 @@ struct PreparedImageResize {
prepared_height: u32,
}
struct ImagePreparation<'a> {
mode: ImagePreparationMode,
image_store: &'a dyn AttachmentStore,
analytics: &'a mut Vec<ImagePreparationMetadata>,
history: &'a mut Option<CodexHarnessMetadata>,
}
const REPLAY_IMAGE_DATA_URL_PREFIX: &str = "data:application/octet-stream;base64,";
pub(crate) struct PreparedInlineImage {
encoded: EncodedImage,
effective_detail: ImageDetailSetting,
@@ -115,9 +157,10 @@ impl ImagePreparationError {
match self {
ImagePreparationError::RemoteUrlUnsupported => REMOTE_IMAGE_URL_PLACEHOLDER,
ImagePreparationError::UnsupportedLowDetail => UNSUPPORTED_LOW_DETAIL_PLACEHOLDER,
ImagePreparationError::Processing(ImageProcessingError::ImageTooLarge { .. }) => {
IMAGE_TOO_LARGE_PLACEHOLDER
}
ImagePreparationError::Processing(
ImageProcessingError::ImageTooLarge { .. }
| ImageProcessingError::ImageTooManyPixels { .. },
) => IMAGE_TOO_LARGE_PLACEHOLDER,
ImagePreparationError::Processing(_) => IMAGE_PROCESSING_ERROR_PLACEHOLDER,
}
}
@@ -126,16 +169,23 @@ impl ImagePreparationError {
pub(crate) async fn prepare_response_items(
thread_id: &str,
ephemeral: bool,
items: &mut Vec<ResponseItem>,
items: &mut Vec<ResponseItemEnvelope>,
mode: ImagePreparationMode,
resize_notice_mode: ImageResizeNoticeMode,
image_store: &dyn AttachmentStore,
) -> Vec<ImagePreparationMetadata> {
let mut metadata = Vec::new();
let mut image_preparations = Vec::new();
let mut prepared_items = Vec::with_capacity(items.len());
for mut item in std::mem::take(items) {
let mut annotated_content = to_annotated_content(&mut item);
let resize_notice = match &mut item {
for mut envelope in std::mem::take(items) {
let mut preparation = ImagePreparation {
mode,
image_store,
analytics: &mut image_preparations,
history: &mut envelope.metadata,
};
let item = &mut envelope.item;
let mut annotated_content = to_annotated_content(item);
let resize_notice = match item {
ResponseItem::Message { role, .. } => {
let Some(mut content) = annotated_content.take() else {
continue;
@@ -153,12 +203,10 @@ pub(crate) async fn prepare_response_items(
} else {
ImageResizeNoticeMode::Disabled
},
&mut metadata,
mode,
image_store,
&mut preparation,
)
.await;
let _ = set_annotated_content(&mut item, content);
let _ = set_annotated_content(item, content);
(!resized_images.is_empty()).then(|| {
ImageResizeNotice::new(ImageResizeNoticeSource::UserMessage, resized_images)
})
@@ -172,9 +220,7 @@ pub(crate) async fn prepare_response_items(
output,
call_id.as_deref(),
resize_notice_mode,
&mut metadata,
mode,
image_store,
&mut preparation,
)
.await
}
@@ -187,9 +233,7 @@ pub(crate) async fn prepare_response_items(
output,
Some(call_id.as_str()),
resize_notice_mode,
&mut metadata,
mode,
image_store,
&mut preparation,
)
.await
}
@@ -209,25 +253,57 @@ pub(crate) async fn prepare_response_items(
| ResponseItem::ContextCompaction { .. }
| ResponseItem::Other => None,
};
prepared_items.push(item);
prepared_items.push(envelope);
if let Some(resize_notice) = resize_notice {
prepared_items.push(ContextualUserFragment::into(resize_notice));
prepared_items.push(ResponseItemEnvelope::new(ContextualUserFragment::into(
resize_notice,
)));
}
}
*items = prepared_items;
metadata
image_preparations
}
fn replay_image_url<'a>(
image_url: &str,
detail: Option<ImageDetail>,
sources: &'a [ImageReplaySource],
mode: ImagePreparationMode,
) -> Option<&'a str> {
if sources.is_empty() {
return None;
}
let digest = sha1_digest(image_url.as_bytes());
sources
.iter()
.find(|source| {
if source.detail != detail
|| !source
.original_data_url
.starts_with(REPLAY_IMAGE_DATA_URL_PREFIX)
{
return false;
}
if source.prepared_url_digest.len() == 32 {
// Older builds wrote a 32-byte digest. Match their original by reproducing the recorded URL.
let mut original_detail = detail;
return resize_image(&source.original_data_url, &mut original_detail, mode)
.ok()
.flatten()
.is_some_and(|prepared| prepared.into_data_url() == image_url);
}
source.prepared_url_digest == digest
})
.map(|source| source.original_data_url.as_str())
}
#[allow(clippy::too_many_arguments)]
async fn prepare_tool_output(
thread_id: &str,
ephemeral: bool,
output: &mut FunctionCallOutputPayload,
item_id: Option<&str>,
resize_notice_mode: ImageResizeNoticeMode,
metadata: &mut Vec<ImagePreparationMetadata>,
mode: ImagePreparationMode,
image_store: &dyn AttachmentStore,
preparation: &mut ImagePreparation<'_>,
) -> Option<ImageResizeNotice> {
let content = output.content_items_mut()?;
let resized_images = prepare_tool_output_content(
@@ -239,9 +315,7 @@ async fn prepare_tool_output(
item_id,
},
resize_notice_mode,
metadata,
mode,
image_store,
preparation,
)
.await;
if resized_images.is_empty() {
@@ -258,9 +332,7 @@ async fn prepare_message_content(
items: &mut [AnnotatedContent],
origin: ImageOrigin<'_>,
resize_notice_mode: ImageResizeNoticeMode,
metadata: &mut Vec<ImagePreparationMetadata>,
mode: ImagePreparationMode,
image_store: &dyn AttachmentStore,
preparation: &mut ImagePreparation<'_>,
) -> Vec<ResizedImage> {
// File-backed images are already prepared and pass through while retaining their positions in
// resize notices.
@@ -273,7 +345,7 @@ async fn prepare_message_content(
for item in items {
if let ContentItem::InputImage { image, detail } = item.content_mut() {
image_number += 1;
match prepare_image(image, detail, origin, metadata, mode, image_store).await {
match prepare_image(image, detail, origin, preparation).await {
Ok(Some(resize)) if resize_notice_mode == ImageResizeNoticeMode::Enabled => {
resized_images.push(ResizedImage {
image_number,
@@ -303,9 +375,7 @@ async fn prepare_tool_output_content(
items: &mut [FunctionCallOutputContentItem],
origin: ImageOrigin<'_>,
resize_notice_mode: ImageResizeNoticeMode,
metadata: &mut Vec<ImagePreparationMetadata>,
mode: ImagePreparationMode,
image_store: &dyn AttachmentStore,
preparation: &mut ImagePreparation<'_>,
) -> Vec<ResizedImage> {
let image_count = items
.iter()
@@ -316,7 +386,7 @@ async fn prepare_tool_output_content(
for item in items {
if let FunctionCallOutputContentItem::InputImage { image, detail } = item {
image_number += 1;
match prepare_image(image, detail, origin, metadata, mode, image_store).await {
match prepare_image(image, detail, origin, preparation).await {
Ok(Some(resize)) if resize_notice_mode == ImageResizeNoticeMode::Enabled => {
resized_images.push(ResizedImage {
image_number,
@@ -356,17 +426,25 @@ async fn prepare_image(
image: &mut ImageReference,
detail: &mut Option<ImageDetail>,
origin: ImageOrigin<'_>,
metadata: &mut Vec<ImagePreparationMetadata>,
mode: ImagePreparationMode,
image_store: &dyn AttachmentStore,
preparation: &mut ImagePreparation<'_>,
) -> Result<Option<PreparedImageResize>, ImagePreparationError> {
let ImageReference::Inline { image_url } = image else {
return Ok(None);
};
let Some(prepared) = resize_image(image_url, detail, mode)? else {
let mode = preparation.mode;
// On resume or fork, sanitize the saved original. If there is no match, sanitize
// the inline image from history instead.
let replay_url = if mode.replay && mode.decoder == PromptImageDecoder::Sanitizer {
preparation.history.as_ref().and_then(|history| {
replay_image_url(image_url, *detail, &history.image_replay_sources, mode)
})
} else {
None
};
let Some(prepared) = resize_image(replay_url.unwrap_or(image_url), detail, mode)? else {
return Ok(None);
};
metadata.push(prepared.metadata(origin));
preparation.analytics.push(prepared.metadata(origin));
let resize = prepared.resize;
let request = UploadRequest {
thread_id: origin.thread_id.to_owned(),
@@ -375,7 +453,7 @@ async fn prepare_image(
tags: HashMap::new(),
data: prepared.encoded.bytes.to_vec(),
};
*image = match image_store.upload(request).await {
let prepared_image = match preparation.image_store.upload(request).await {
Ok(UploadResult::Inline { bytes }) => ImageReference::Inline {
image_url: data_url_from_bytes(&prepared.encoded.mime, &bytes),
},
@@ -387,6 +465,31 @@ async fn prepare_image(
}
}
};
// Keep the original so replay can reproduce the image the model saw. Encoding an
// already prepared JPEG again can change its bytes. If only the MIME label changed,
// the saved image already has the original payload.
if !mode.replay
&& mode.decoder == PromptImageDecoder::Sanitizer
&& let ImageReference::Inline {
image_url: prepared_url,
} = &prepared_image
&& let Some((_, source_payload)) = image_url.split_once(',')
&& let Some((_, prepared_payload)) = prepared_url.split_once(',')
&& source_payload != prepared_payload
{
preparation
.history
.get_or_insert_default()
.image_replay_sources
.push(ImageReplaySource {
prepared_url_digest: sha1_digest(prepared_url.as_bytes()).to_vec(),
detail: *detail,
original_data_url: Arc::new(format!(
"{REPLAY_IMAGE_DATA_URL_PREFIX}{source_payload}"
)),
});
}
*image = prepared_image;
Ok(resize)
}
@@ -402,12 +505,13 @@ pub(crate) fn resize_image(
return Ok(None);
}
let (effective_detail, image_mode) = match mode {
ImagePreparationMode::UnifiedBudget => (
let (effective_detail, image_mode) = if mode.unified_budget {
(
ImageDetailSetting::Original,
PromptImageMode::ORIGINAL_DETAIL,
),
ImagePreparationMode::DetailBased => match detail {
)
} else {
match detail {
None | Some(ImageDetail::Auto | ImageDetail::High) => {
(ImageDetailSetting::High, PromptImageMode::HIGH_DETAIL)
}
@@ -416,9 +520,9 @@ pub(crate) fn resize_image(
PromptImageMode::ORIGINAL_DETAIL,
),
Some(ImageDetail::Low) => return Err(ImagePreparationError::UnsupportedLowDetail),
},
}
};
let encoded = load_data_url_for_prompt(image_url, image_mode)?;
let encoded = load_data_url_for_prompt(image_url, image_mode, mode.decoder)?;
let resize = ((encoded.source_width, encoded.source_height) != (encoded.width, encoded.height))
.then_some(PreparedImageResize {
source_width: encoded.source_width,
@@ -426,7 +530,7 @@ pub(crate) fn resize_image(
prepared_width: encoded.width,
prepared_height: encoded.height,
});
if mode == ImagePreparationMode::UnifiedBudget {
if mode.unified_budget {
// Preserve accurate context-window accounting while older transports still require an
// image detail field. Responses Lite removes this compatibility hint before sending.
*detail = Some(ImageDetail::Original);
+82 -6
View File
@@ -9,6 +9,7 @@ use codex_attachment_store::IsResolvableFuture;
use codex_attachment_store::ResolveFuture;
use codex_attachment_store::ResolveRequest;
use codex_attachment_store::UploadFuture;
use codex_history::RolloutItem;
use codex_protocol::models::ContentItemKind;
use codex_protocol::models::FunctionCallOutputBody;
use codex_protocol::models::FunctionCallOutputPayload;
@@ -26,6 +27,28 @@ use serde_json::json;
use super::*;
async fn prepare_response_items(
thread_id: &str,
ephemeral: bool,
items: &mut Vec<ResponseItem>,
mode: ImagePreparationMode,
resize_notice_mode: ImageResizeNoticeMode,
image_store: &dyn AttachmentStore,
) -> Vec<ImagePreparationMetadata> {
let mut envelopes = items.drain(..).map(ResponseItemEnvelope::new).collect();
let metadata = super::prepare_response_items(
thread_id,
ephemeral,
&mut envelopes,
mode,
resize_notice_mode,
image_store,
)
.await;
items.extend(envelopes.into_iter().map(ResponseItemEnvelope::into_item));
metadata
}
pub(crate) struct FailingAttachmentStore;
impl AttachmentStore for FailingAttachmentStore {
@@ -70,6 +93,59 @@ fn decoded_image(image_url: &str) -> (Vec<u8>, DynamicImage) {
(bytes, image)
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn replay_regenerates_the_recorded_jpeg_from_its_original() {
let pixels = image::RgbImage::from_fn(81, 57, |x, y| {
image::Rgb([(19 * x + y) as u8, (11 * y + x) as u8, (x * y) as u8])
});
let mut bytes = Cursor::new(Vec::new());
DynamicImage::ImageRgb8(pixels)
.write_to(&mut bytes, ImageFormat::Jpeg)
.expect("encode original JPEG");
let source = data_url_from_bytes("image/jpeg", &bytes.into_inner());
let message = ResponseItem::Message {
status: None,
encrypted_content: None,
id: None,
role: "user".to_string(),
content: vec![ContentItem::InputImage {
image: ImageReference::Inline { image_url: source },
detail: Some(ImageDetail::High),
}],
phase: None,
internal_chat_message_metadata_passthrough: None,
};
let mut items = vec![ResponseItemEnvelope::new(message)];
let prepare = async |items: &mut Vec<ResponseItemEnvelope>, replay| {
super::prepare_response_items(
"image-preparation-thread",
/*ephemeral*/ false,
items,
ImagePreparationMode {
decoder: PromptImageDecoder::Sanitizer,
replay,
..Default::default()
},
ImageResizeNoticeMode::Disabled,
&InlineAttachmentStore,
)
.await;
};
prepare(&mut items, /*replay*/ false).await;
let recorded = items.pop().expect("prepared item");
assert!(recorded.metadata.is_some());
let json =
serde_json::to_string(&RolloutItem::ResponseItem(recorded.clone())).expect("persist item");
let RolloutItem::ResponseItem(reloaded) =
serde_json::from_str(&json).expect("read persisted item")
else {
panic!("expected persisted response item");
};
let mut replayed = vec![reloaded];
prepare(&mut replayed, /*replay*/ true).await;
assert_eq!(replayed, vec![recorded]);
}
#[tokio::test(flavor = "multi_thread")]
async fn preparation_preserves_small_image_bytes_and_replaces_remote_urls() {
let (data_url, original_bytes) = png_data_url(/*width*/ 64, /*height*/ 32);
@@ -114,7 +190,7 @@ async fn preparation_preserves_small_image_bytes_and_replaces_remote_urls() {
"image-preparation-thread",
/*ephemeral*/ false,
&mut items,
ImagePreparationMode::DetailBased,
ImagePreparationMode::default(),
ImageResizeNoticeMode::Disabled,
&InlineAttachmentStore,
)
@@ -206,7 +282,7 @@ async fn detail_policies_apply_the_expected_budgets() {
"image-preparation-thread",
/*ephemeral*/ false,
&mut items,
ImagePreparationMode::DetailBased,
ImagePreparationMode::default(),
ImageResizeNoticeMode::Disabled,
&InlineAttachmentStore,
)
@@ -261,7 +337,7 @@ async fn preparation_reports_tool_output_item_id() {
"image-preparation-thread",
/*ephemeral*/ false,
&mut items,
ImagePreparationMode::DetailBased,
ImagePreparationMode::default(),
ImageResizeNoticeMode::Disabled,
&InlineAttachmentStore,
)
@@ -301,7 +377,7 @@ async fn upload_failure_keeps_resized_image_inline() {
"image-preparation-thread",
/*ephemeral*/ false,
&mut items,
ImagePreparationMode::DetailBased,
ImagePreparationMode::default(),
ImageResizeNoticeMode::Disabled,
&FailingAttachmentStore,
)
@@ -405,7 +481,7 @@ async fn resize_notices_count_file_backed_images_and_skip_failed_images() {
"image-preparation-thread",
/*ephemeral*/ false,
&mut items,
ImagePreparationMode::DetailBased,
ImagePreparationMode::default(),
ImageResizeNoticeMode::Enabled,
&InlineAttachmentStore,
)
@@ -602,7 +678,7 @@ async fn preparation_replaces_only_failed_tool_images_and_preserves_metadata() {
"image-preparation-thread",
/*ephemeral*/ false,
&mut items,
ImagePreparationMode::DetailBased,
ImagePreparationMode::default(),
ImageResizeNoticeMode::Disabled,
&InlineAttachmentStore,
)
+7 -39
View File
@@ -1,7 +1,6 @@
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::ResponseItemEnvelope;
use codex_protocol::error::CodexErr;
@@ -126,24 +125,20 @@ impl Session {
envelope
}
#[tracing::instrument(
name = "record_conversation_items",
level = "trace",
skip_all,
fields(item_count = items.len())
)]
pub(crate) async fn record_annotated_conversation_items(
&self,
turn_context: &TurnContext,
model_info: &ModelInfo,
items: Vec<ResponseItemEnvelope>,
) {
if items.iter().all(|item| item.metadata.is_none()) {
let items = items
.into_iter()
.map(ResponseItemEnvelope::into_item)
.collect::<Vec<_>>();
self.record_conversation_items(turn_context, model_info, &items)
.await;
return;
}
let (annotated_items, image_preparations) = self
.prepare_annotated_conversation_items_for_history(turn_context, model_info, items)
.prepare_conversation_items_for_history(turn_context, model_info, items)
.await;
self.record_prepared_conversation_items(
turn_context,
@@ -154,33 +149,6 @@ impl Session {
.await;
}
pub(crate) async fn prepare_annotated_conversation_items_for_history(
&self,
turn_context: &TurnContext,
model_info: &ModelInfo,
items: Vec<ResponseItemEnvelope>,
) -> (Vec<ResponseItemEnvelope>, Vec<ImagePreparationMetadata>) {
let mut annotated_items = Vec::with_capacity(items.len());
let mut image_preparations = Vec::new();
for envelope in items {
let (prepared_items, prepared_images) = self
.prepare_conversation_items_for_history(
turn_context,
model_info,
std::slice::from_ref(&envelope.item),
)
.await;
image_preparations.extend(prepared_images);
let mut metadata = envelope.metadata;
annotated_items.extend(prepared_items.into_iter().map(|item| ResponseItemEnvelope {
item,
metadata: metadata.take(),
}));
}
(annotated_items, image_preparations)
}
/// Injects items into active work, or records them without starting a turn.
pub(crate) async fn inject_no_new_turn(
&self,
+31 -53
View File
@@ -49,7 +49,6 @@ use crate::exec_policy::default_policy_path;
use crate::image_preparation::ImagePreparationMode;
use crate::image_preparation::ImageResizeNoticeMode;
use crate::image_preparation::prepare_response_items as prepare_image_response_items;
use crate::image_preparation::unified_image_budget_enabled;
use crate::parse_turn_item;
use crate::realtime_conversation::RealtimeConversationManager;
use crate::realtime_history::RealtimeEventOrder;
@@ -1999,38 +1998,22 @@ impl Session {
.and_then(|context| context.direct_only_tool_namespaces.clone())
.map(|namespaces| (turn.turn_id.clone(), namespaces))
});
// Keep the recorded rollout unchanged. Prepare its reconstructed history before
// installing it, so legacy media is processed once for this resume or fork and
// will be processed again if the rollout is reconstructed in a future session.
// Replay disables image-resize notices, so media preparation remains one-to-one. Keep
// the prior batch behavior and carry history-only metadata in a positional sidecar.
let (mut prepared_history, metadata): (Vec<_>, Vec<_>) = history
.into_iter()
.map(|envelope| (envelope.item, envelope.metadata))
.unzip();
// Replay must not upload or migrate recorded history. The inline store returns prepared
// inline bytes, while existing file references bypass preparation and remain unchanged.
// Prepare inline images in memory on resume or fork. Keep file references as they are.
// The inline store avoids uploads, and replay leaves the saved history untouched.
// Skip resize notices here so replay doesn't add messages to the conversation.
// Bound replay future size now that image preparation can await storage.
let _ = Box::pin(prepare_image_response_items(
&self.thread_id.to_string(),
turn_context.config.ephemeral,
&mut prepared_history,
ImagePreparationMode::DetailBased,
&mut history,
ImagePreparationMode::for_replay(&turn_context.config.features),
ImageResizeNoticeMode::Disabled,
&InlineAttachmentStore,
))
.await;
prepare_audio_response_items(&mut prepared_history);
assert_eq!(
prepared_history.len(),
metadata.len(),
"replay media preparation must remain one-to-one when resize notices are disabled"
);
history = prepared_history
.into_iter()
.zip(metadata)
.map(|(item, metadata)| ResponseItemEnvelope { item, metadata })
.collect();
for envelope in &mut history {
prepare_audio_response_items(std::slice::from_mut(&mut envelope.item));
}
file_audio::resolve(self, file_audio::file_ids(&history).cloned().collect()).await;
// Ordinary replay needs no reviewer refresh. Only an opaque
// checkpoint can make reconstruction depend on reviewer compatibility.
@@ -3946,21 +3929,16 @@ impl Session {
item.set_create_time_if_missing(Self::response_item_create_time());
}
/// Prepares media using the originating model and preserves existing item identity.
/// Prepares media for the model that received these items and keeps their existing IDs.
/// Each saved original stays with the item that contains the prepared image.
pub(crate) async fn prepare_conversation_items_for_history(
&self,
turn_context: &TurnContext,
model_info: &ModelInfo,
items: &[ResponseItem],
) -> (Vec<ResponseItem>, Vec<ImagePreparationMetadata>) {
let mut items = items.to_vec();
let image_preparation_mode =
if unified_image_budget_enabled(&turn_context.config.features, model_info) {
ImagePreparationMode::UnifiedBudget
} else {
ImagePreparationMode::DetailBased
};
let image_resize_notice_mode = if turn_context
mut items: Vec<ResponseItemEnvelope>,
) -> (Vec<ResponseItemEnvelope>, Vec<ImagePreparationMetadata>) {
let image_mode = ImagePreparationMode::for_model(&turn_context.config.features, model_info);
let resize_notices = if turn_context
.config
.features
.enabled(Feature::ImageResizeNotice)
@@ -3974,14 +3952,15 @@ impl Session {
&self.thread_id.to_string(),
turn_context.config.ephemeral,
&mut items,
image_preparation_mode,
image_resize_notice_mode,
image_mode,
resize_notices,
self.services.image_store.as_ref(),
))
.await;
prepare_audio_response_items(&mut items);
// Most response items get their passthrough turn ID at the durable history boundary.
for item in &mut items {
for envelope in &mut items {
let item = &mut envelope.item;
prepare_audio_response_items(std::slice::from_mut(item));
// Most response items get their passthrough turn ID at the durable history boundary.
Self::stamp_response_item_for_history(item, &turn_context.sub_id);
self.assign_missing_response_item_id(item);
}
@@ -4039,22 +4018,20 @@ impl Session {
/// Appends to history, persists the prepared items, then notifies raw-item observers.
/// Execution callers supply their captured model; standalone callers need no tool runtime.
#[tracing::instrument(level = "trace", skip_all, fields(item_count = items.len()))]
pub(crate) async fn record_conversation_items(
&self,
turn_context: &TurnContext,
model_info: &ModelInfo,
items: &[ResponseItem],
) {
let (items, image_preparations) = self
.prepare_conversation_items_for_history(turn_context, model_info, items)
.await;
let items = items.into_iter().map(ResponseItemEnvelope::new).collect();
self.record_prepared_conversation_items(
self.record_annotated_conversation_items(
turn_context,
model_info,
items,
image_preparations,
items
.iter()
.cloned()
.map(ResponseItemEnvelope::new)
.collect(),
)
.await;
}
@@ -4696,14 +4673,15 @@ impl Session {
) {
let communication = communication.into();
let response_item = Self::inter_agent_history_item(&communication, &turn_context.config);
let (items, _) = self
let (mut prepared, _) = self
.prepare_conversation_items_for_history(
turn_context,
model_info,
std::slice::from_ref(&response_item),
vec![ResponseItemEnvelope::new(response_item)],
)
.await;
let mut response_item = ResponseItemEnvelope::new(items[0].clone());
let mut response_item = prepared.remove(0);
let items = [response_item.item.clone()];
// A send confirmed after the pending snapshot must not reach the rollout
// before this boundary; older readers assign deliveries by physical order.
let boundary = self
@@ -5978,7 +5956,7 @@ impl Session {
)]);
}
let (prepared_items, image_preparations) = self
.prepare_annotated_conversation_items_for_history(
.prepare_conversation_items_for_history(
turn_context,
model_info,
vec![ResponseItemEnvelope {
+6 -3
View File
@@ -390,14 +390,17 @@ async fn default_turn_context_assigns_missing_response_item_ids() {
role: "developer".to_string(),
tools: Arc::from([]),
},
];
]
.into_iter()
.map(ResponseItemEnvelope::new)
.collect();
let (items, _) = session
.prepare_conversation_items_for_history(&turn_context, turn_context.model_info(), &items)
.prepare_conversation_items_for_history(&turn_context, turn_context.model_info(), items)
.await;
assert_eq!(items.len(), 3);
for (item, prefix) in items.iter().zip(["msg_", "amsg_", "at_"]) {
let suffix = item.id().unwrap().strip_prefix(prefix).unwrap();
let suffix = item.item.id().unwrap().strip_prefix(prefix).unwrap();
assert_eq!(Uuid::parse_str(suffix).unwrap().get_version_num(), 7);
}
}
+8 -2
View File
@@ -6,6 +6,7 @@ use crate::context::NodeReplReviewEvidence;
use crate::context::NodeReplReviewEvidenceMode;
use crate::context::node_repl_review_evidence_mode;
use crate::function_tool::FunctionCallError;
use crate::image_preparation::prompt_image_decoder;
use crate::mcp_tool_call::handle_mcp_tool_call;
use crate::original_image_detail::can_request_original_image_detail;
use crate::session::session::Session;
@@ -415,6 +416,7 @@ impl CoreToolRuntime for McpHandler {
.and_then(Value::as_bool)
== Some(true)
};
let decoder = prompt_image_decoder(&invocation.turn.config.features);
let mut captured_image_bytes = 0_usize;
let mut items = content
.iter()
@@ -458,8 +460,12 @@ impl CoreToolRuntime for McpHandler {
.and_then(|detail| serde_json::from_value(detail.clone()).ok());
let image_url =
format!("data:{};base64,{payload}", mime_type.to_ascii_lowercase());
load_data_url_for_prompt_uncached(&image_url, PromptImageMode::Original)
.ok()?;
load_data_url_for_prompt_uncached(
&image_url,
PromptImageMode::Original,
decoder,
)
.ok()?;
captured_image_bytes = next_image_bytes;
Some(UserInput::Image {
image: codex_protocol::models::ImageReference::Inline { image_url },
+25 -6
View File
@@ -10,10 +10,14 @@ use codex_protocol::models::ImageDetail;
use codex_protocol::models::ImageReference;
use codex_protocol::models::ResponseInputItem;
use codex_protocol::openai_models::InputModality;
use codex_utils_image::MAX_SANITIZED_IMAGE_BYTES;
use codex_utils_image::PromptImageDecoder;
use codex_utils_image::data_url_from_bytes;
use serde::Deserialize;
use crate::function_tool::FunctionCallError;
use crate::image_preparation::ImagePreparationMode;
use crate::image_preparation::resize_image;
use crate::original_image_detail::can_request_original_image_detail;
use crate::tools::context::ToolEnvironments;
use crate::tools::context::ToolInvocation;
@@ -165,6 +169,16 @@ impl ViewImageHandler {
"image path `{model_visible_path}` is not a file"
)));
}
let mode = ImagePreparationMode::for_model(
&turn.config.features,
&step_context.settings.model_info,
);
let sanitize = mode.decoder == PromptImageDecoder::Sanitizer;
let invalid_image =
|| FunctionCallError::RespondToModel(VIEW_IMAGE_INVALID_MESSAGE.to_string());
if sanitize && metadata.size > MAX_SANITIZED_IMAGE_BYTES as u64 {
return Err(invalid_image());
}
let file_bytes = fs
.read_file(&path_uri, ReadFileOptions::default(), Some(&sandbox))
.await
@@ -173,11 +187,12 @@ impl ViewImageHandler {
"unable to read image at `{model_visible_path}`: {error}"
))
})?;
// Reject non-images before their bytes can reach code mode without changing
// valid image bytes, metadata, or centralized image preparation.
image::load_from_memory(&file_bytes).map_err(|_| {
FunctionCallError::RespondToModel(VIEW_IMAGE_INVALID_MESSAGE.to_string())
})?;
// Validate before passing the image to Code Mode or history.
if !sanitize {
image::load_from_memory(&file_bytes).map_err(|_| invalid_image())?;
} else if file_bytes.len() > MAX_SANITIZED_IMAGE_BYTES {
return Err(invalid_image());
}
let can_request_original_detail =
can_request_original_image_detail(&step_context.settings.model_info);
@@ -189,8 +204,12 @@ impl ViewImageHandler {
DEFAULT_IMAGE_DETAIL
};
// The history insertion path owns image preparation and resizing.
// Return the original so history can resize it and report the change to the model.
let image_url = data_url_from_bytes("application/octet-stream", &file_bytes);
if sanitize {
// Check the image now so view_image can report an error before returning it.
resize_image(&image_url, &mut Some(image_detail), mode).map_err(|_| invalid_image())?;
}
let item = TurnItem::ImageView(ImageViewItem {
id: call_id,
+8
View File
@@ -404,6 +404,8 @@ pub enum Feature {
OmitAppServerNotificationMedia,
/// Tell the model when a prompt image was resized and include its dimensions.
ImageResizeNotice,
/// Normalize prompt images with the reviewed image sanitizer.
SafeImageDecoder,
/// Apply one shared pixel and token budget to every image, regardless of legacy detail hints.
UnifiedImageBudget,
/// Removed compatibility flag for always-on centralized image preparation.
@@ -2046,6 +2048,12 @@ pub const FEATURES: &[FeatureSpec] = &[
stage: Stage::Stable,
default_enabled: true,
},
FeatureSpec {
id: Feature::SafeImageDecoder,
key: "safe_image_decoder",
stage: Stage::UnderDevelopment,
default_enabled: false,
},
FeatureSpec {
id: Feature::ResizeAllImages,
key: "resize_all_images",
+19
View File
@@ -75,6 +75,10 @@ pub struct CodexHarnessMetadata {
#[serde(default, skip_serializing_if = "Option::is_none")]
pub sampled_through_sequence: Option<u64>,
/// Original inline images to sanitize again on resume or fork.
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub image_replay_sources: Vec<ImageReplaySource>,
/// Complete retained records actually delivered by this Guardian message. Host-only.
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub guardian_sources: Vec<RetainedSource>,
@@ -160,6 +164,21 @@ pub struct CodexHarnessMetadata {
pub sender_user_messages: Option<Box<SenderUserMessages>>,
}
#[derive(Serialize, Deserialize, Clone, PartialEq, Eq, JsonSchema)]
pub struct ImageReplaySource {
pub prepared_url_digest: Vec<u8>,
pub detail: Option<codex_protocol::models::ImageDetail>,
pub original_data_url: Arc<String>,
}
impl std::fmt::Debug for ImageReplaySource {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("ImageReplaySource")
.field("original_data_url_len", &self.original_data_url.len())
.finish_non_exhaustive()
}
}
fn deserialize_mcp_attribution_checkpoint<'de, D>(
deserializer: D,
) -> Result<Option<McpAttribution>, D::Error>
+2 -1
View File
@@ -2072,7 +2072,8 @@ pub fn local_image_content_items_with_label_number(
ImageProcessingError::Read { .. }
| ImageProcessingError::Encode { .. }
| ImageProcessingError::InvalidDataUrl { .. }
| ImageProcessingError::ImageTooLarge { .. } => {
| ImageProcessingError::ImageTooLarge { .. }
| ImageProcessingError::ImageTooManyPixels { .. } => {
vec![local_media_error_placeholder(
path,
&err,
+3 -2
View File
@@ -14,8 +14,9 @@ ignored = ["png"]
[dependencies]
base64 = { workspace = true }
image = { workspace = true, features = ["jpeg", "png", "gif", "webp"] }
png = { version = "0.18.0", features = ["zlib-rs"] }
image = { version = "=0.25.9", default-features = false, features = ["jpeg", "png", "gif", "webp"] }
moxcms = { version = "=0.7.11", features = ["options"] }
png = { version = "=0.18.1", features = ["zlib-rs"] }
codex-utils-cache = { workspace = true }
mime_guess = { workspace = true }
thiserror = { workspace = true }
+2
View File
@@ -27,6 +27,8 @@ pub enum ImageProcessingError {
UnsupportedImageFormat { mime: String },
#[error("invalid image data URL: {reason}")]
InvalidDataUrl { reason: String },
#[error("image has too many decoded pixels ({pixels}; max {max})")]
ImageTooManyPixels { pixels: u64, max: u64 },
#[error("image {representation} is too large ({size} bytes; max {max} bytes)")]
ImageTooLarge {
representation: &'static str,
+34 -5
View File
@@ -236,8 +236,12 @@ async fn data_url_processing_preserves_supported_source_bytes() {
.replacen("data:", "DATA:", 1)
.replacen(";base64,", ";BASE64,", 1);
let processed = load_data_url_for_prompt_uncached(&image_url, PromptImageMode::ResizeToFit)
.expect("process data URL image");
let processed = load_data_url_for_prompt_uncached(
&image_url,
PromptImageMode::ResizeToFit,
PromptImageDecoder::Legacy,
)
.expect("process data URL image");
assert_eq!(processed.width, 64);
assert_eq!(processed.height, 32);
@@ -251,8 +255,12 @@ async fn data_url_processing_converts_gif_to_png() {
let gif_bytes = image_bytes(&image, ImageFormat::Gif);
let image_url = data_url_from_bytes("image/gif", &gif_bytes);
let processed = load_data_url_for_prompt(&image_url, PromptImageMode::ResizeToFit)
.expect("process GIF data URL");
let processed = load_data_url_for_prompt(
&image_url,
PromptImageMode::ResizeToFit,
PromptImageDecoder::Legacy,
)
.expect("process GIF data URL");
assert_eq!(processed.mime, "image/png");
assert_eq!(
@@ -270,7 +278,11 @@ fn data_url_processing_rejects_malformed_input() {
"data:image/png;base64,not base64",
] {
assert!(matches!(
load_data_url_for_prompt(image_url, PromptImageMode::ResizeToFit),
load_data_url_for_prompt(
image_url,
PromptImageMode::ResizeToFit,
PromptImageDecoder::Legacy,
),
Err(ImageProcessingError::InvalidDataUrl { .. })
));
}
@@ -382,12 +394,29 @@ async fn reprocesses_updated_file_contents() {
assert_ne!(second.bytes, first.bytes);
}
#[test]
fn legacy_cache_entry_cannot_bypass_sanitization() {
let image = ImageBuffer::from_pixel(7, 11, Rgba([42, 64, 128, 255]));
let bytes = image_bytes_with_metadata(&image, ImageFormat::Jpeg, TEST_RGB_ICC_PROFILE);
let path = Path::new("oriented.jpg");
let mode = PromptImageMode::Original;
let legacy = load_for_prompt_bytes(path, bytes.clone(), mode).expect("load legacy image");
let sanitized =
load_for_prompt_bytes_with_decoder(path, bytes, mode, PromptImageDecoder::Sanitizer)
.expect("sanitize image after warming the legacy cache");
assert_eq!((legacy.width, legacy.height), (7, 11));
assert_eq!((sanitized.width, sanitized.height), (11, 7));
}
#[tokio::test(flavor = "multi_thread")]
async fn bounds_cache_by_encoded_byte_size() {
let cache = ImageCache::new(NonZeroUsize::new(4).expect("non-zero cache capacity"));
let key = |digest_byte| ImageCacheKey {
digest: [digest_byte; 20],
mode: PromptImageMode::Original,
decoder: PromptImageDecoder::Legacy,
};
let image = |size| EncodedImage {
bytes: vec![0; size].into(),
+68 -35
View File
@@ -29,9 +29,13 @@ pub const MAX_DIMENSION: u32 = 2048;
/// This is a high sanity guard against pathological inputs, not a protocol
/// requirement or target upload size.
pub const MAX_PROMPT_IMAGE_INPUT_BYTES: usize = 1024 * 1024 * 1024;
/// Use the same byte limit for input and output so a saved image fits on replay.
/// At the pixel limit, 16-bit RGBA uses 81,920,000 bytes. Round up for PNG overhead.
pub const MAX_SANITIZED_IMAGE_BYTES: usize = 79 * 1024 * 1024;
const MAX_IMAGE_CACHE_BYTES: usize = 64 * 1024 * 1024;
pub mod error;
mod sanitizer;
pub use crate::error::ImageProcessingError;
@@ -64,6 +68,13 @@ pub enum PromptImageMode {
ResizeWithLimits(PromptImageResizeLimits),
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Default)]
pub enum PromptImageDecoder {
#[default]
Legacy,
Sanitizer,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct PromptImageResizeLimits {
pub max_dimension: u32,
@@ -92,6 +103,7 @@ struct ImageMetadata {
struct ImageCacheKey {
digest: [u8; 20],
mode: PromptImageMode,
decoder: PromptImageDecoder,
}
type ImageCache = BlockingLruCache<ImageCacheKey, EncodedImage>;
@@ -104,18 +116,26 @@ pub fn load_for_prompt_bytes(
file_bytes: Vec<u8>,
mode: PromptImageMode,
) -> Result<EncodedImage, ImageProcessingError> {
let path_buf = path.to_path_buf();
load_for_prompt_bytes_with_decoder(path, file_bytes, mode, PromptImageDecoder::Legacy)
}
fn load_for_prompt_bytes_with_decoder(
path: &Path,
file_bytes: Vec<u8>,
mode: PromptImageMode,
decoder: PromptImageDecoder,
) -> Result<EncodedImage, ImageProcessingError> {
let key = ImageCacheKey {
digest: sha1_digest(&file_bytes),
mode,
decoder,
};
if let Some(image) = IMAGE_CACHE.get(&key) {
return Ok(image);
}
let image = load_for_prompt_bytes_uncached(&path_buf, file_bytes, mode)?;
let image = load_for_prompt_bytes_uncached(path, file_bytes, mode, decoder)?;
cache_image(&IMAGE_CACHE, key, image.clone(), MAX_IMAGE_CACHE_BYTES);
Ok(image)
}
@@ -124,7 +144,11 @@ fn load_for_prompt_bytes_uncached(
path: &Path,
file_bytes: Vec<u8>,
mode: PromptImageMode,
decoder: PromptImageDecoder,
) -> Result<EncodedImage, ImageProcessingError> {
if decoder == PromptImageDecoder::Sanitizer {
return sanitizer::sanitize(path, &file_bytes, mode);
}
let path_buf = path.to_path_buf();
(move || {
let guessed_format = image::guess_format(&file_bytes)
@@ -159,26 +183,7 @@ fn load_for_prompt_bytes_uncached(
let (width, height) = dynamic.dimensions();
let target_dimensions = match mode {
PromptImageMode::ResizeToFit if width > MAX_DIMENSION || height > MAX_DIMENSION => {
let resized = dynamic.resize(MAX_DIMENSION, MAX_DIMENSION, FilterType::Triangle);
Some((resized.width(), resized.height(), resized))
}
PromptImageMode::ResizeWithLimits(limits) => {
let (target_width, target_height) =
prompt_image_output_dimensions_for_limits(width, height, limits);
if (target_width, target_height) == (width, height) {
None
} else {
let resized =
dynamic.resize_exact(target_width, target_height, FilterType::Triangle);
Some((target_width, target_height, resized))
}
}
PromptImageMode::ResizeToFit | PromptImageMode::Original => None,
};
let encoded = if let Some((prepared_width, prepared_height, resized)) = target_dimensions {
let encoded = if let Some(resized) = resize_for_prompt(&dynamic, mode) {
let target_format = format
.filter(|format| can_preserve_source_bytes(*format))
.unwrap_or(ImageFormat::Png);
@@ -189,8 +194,8 @@ fn load_for_prompt_bytes_uncached(
mime,
source_width: width,
source_height: height,
width: prepared_width,
height: prepared_height,
width: resized.width(),
height: resized.height(),
}
} else {
if let Some(format) = format.filter(|format| can_preserve_source_bytes(*format)) {
@@ -244,22 +249,25 @@ fn cache_image(cache: &ImageCache, key: ImageCacheKey, image: EncodedImage, byte
pub fn load_data_url_for_prompt(
image_url: &str,
mode: PromptImageMode,
decoder: PromptImageDecoder,
) -> Result<EncodedImage, ImageProcessingError> {
load_data_url_for_prompt_with(image_url, mode, load_for_prompt_bytes)
let bytes = decode_data_url(image_url, decoder)?;
load_for_prompt_bytes_with_decoder(Path::new("<data-url-image>"), bytes, mode, decoder)
}
pub fn load_data_url_for_prompt_uncached(
image_url: &str,
mode: PromptImageMode,
decoder: PromptImageDecoder,
) -> Result<EncodedImage, ImageProcessingError> {
load_data_url_for_prompt_with(image_url, mode, load_for_prompt_bytes_uncached)
let bytes = decode_data_url(image_url, decoder)?;
load_for_prompt_bytes_uncached(Path::new("<data-url-image>"), bytes, mode, decoder)
}
fn load_data_url_for_prompt_with(
fn decode_data_url(
image_url: &str,
mode: PromptImageMode,
load: impl FnOnce(&Path, Vec<u8>, PromptImageMode) -> Result<EncodedImage, ImageProcessingError>,
) -> Result<EncodedImage, ImageProcessingError> {
decoder: PromptImageDecoder,
) -> Result<Vec<u8>, ImageProcessingError> {
let rest = image_url
.get(..DATA_URL_PREFIX.len())
.filter(|prefix| prefix.eq_ignore_ascii_case(DATA_URL_PREFIX))
@@ -281,11 +289,18 @@ fn load_data_url_for_prompt_with(
});
}
if encoded.len() > MAX_PROMPT_IMAGE_INPUT_BYTES {
let (max_bytes, max_base64_bytes) = match decoder {
PromptImageDecoder::Legacy => (MAX_PROMPT_IMAGE_INPUT_BYTES, MAX_PROMPT_IMAGE_INPUT_BYTES),
PromptImageDecoder::Sanitizer => (
MAX_SANITIZED_IMAGE_BYTES,
MAX_SANITIZED_IMAGE_BYTES.div_ceil(3) * 4,
),
};
if encoded.len() > max_base64_bytes {
return Err(ImageProcessingError::ImageTooLarge {
representation: "base64 payload",
size: encoded.len(),
max: MAX_PROMPT_IMAGE_INPUT_BYTES,
max: max_base64_bytes,
});
}
let file_bytes =
@@ -294,15 +309,33 @@ fn load_data_url_for_prompt_with(
.map_err(|source| ImageProcessingError::InvalidDataUrl {
reason: format!("invalid base64 payload: {source}"),
})?;
if file_bytes.len() > MAX_PROMPT_IMAGE_INPUT_BYTES {
if file_bytes.len() > max_bytes {
return Err(ImageProcessingError::ImageTooLarge {
representation: "decoded input",
size: file_bytes.len(),
max: MAX_PROMPT_IMAGE_INPUT_BYTES,
max: max_bytes,
});
}
load(Path::new("<data-url-image>"), file_bytes, mode)
Ok(file_bytes)
}
fn resize_for_prompt(image: &DynamicImage, mode: PromptImageMode) -> Option<DynamicImage> {
let (width, height) = image.dimensions();
match mode {
PromptImageMode::ResizeToFit if width > MAX_DIMENSION || height > MAX_DIMENSION => {
Some(image.resize(MAX_DIMENSION, MAX_DIMENSION, FilterType::Triangle))
}
PromptImageMode::ResizeWithLimits(limits) => {
let (target_width, target_height) =
prompt_image_output_dimensions_for_limits(width, height, limits);
if (target_width, target_height) == (width, height) {
return None;
}
Some(image.resize_exact(target_width, target_height, FilterType::Triangle))
}
PromptImageMode::ResizeToFit | PromptImageMode::Original => None,
}
}
fn prompt_image_output_dimensions_for_limits(
+167
View File
@@ -0,0 +1,167 @@
//! Decode supported prompt images and write fresh PNG or JPEG bytes without the source metadata.
use std::io::Cursor;
use std::path::Path;
use image::DynamicImage;
use image::ImageDecoder;
use image::ImageFormat;
use image::ImageReader;
use image::codecs::jpeg::JpegEncoder;
use image::codecs::png::CompressionType;
use image::codecs::png::FilterType;
use image::codecs::png::PngEncoder;
use moxcms::ColorProfile;
use crate::EncodedImage;
use crate::ImageProcessingError;
use crate::MAX_SANITIZED_IMAGE_BYTES;
use crate::PromptImageMode;
use crate::resize_for_prompt;
// Original detail allows 10,000 patches of 32 × 32 pixels each.
const MAX_DECODED_PIXELS: u64 = 10_240_000;
pub(super) fn sanitize(
path: &Path,
bytes: &[u8],
mode: PromptImageMode,
) -> Result<EncodedImage, ImageProcessingError> {
if bytes.len() > MAX_SANITIZED_IMAGE_BYTES {
return Err(ImageProcessingError::ImageTooLarge {
representation: "input",
size: bytes.len(),
max: MAX_SANITIZED_IMAGE_BYTES,
});
}
let decode_error = |source| ImageProcessingError::decode_error(path, source);
let reader = ImageReader::new(Cursor::new(bytes))
.with_guessed_format()
.map_err(|error| decode_error(error.into()))?;
let (format, output_format) = match reader.format() {
Some(format @ (ImageFormat::Png | ImageFormat::Gif)) => (format, ImageFormat::Png),
Some(format @ (ImageFormat::Jpeg | ImageFormat::WebP)) => (format, ImageFormat::Jpeg),
format => {
return Err(ImageProcessingError::UnsupportedImageFormat {
mime: format
.as_ref()
.map_or("unknown", ImageFormat::to_mime_type)
.to_string(),
});
}
};
let (mut img, orientation, icc_profile) = std::panic::catch_unwind(|| {
let mut decoder = reader.into_decoder().map_err(decode_error)?;
let (width, height) = decoder.dimensions();
let pixels = u64::from(width) * u64::from(height);
if pixels > MAX_DECODED_PIXELS {
return Err(ImageProcessingError::ImageTooManyPixels {
pixels,
max: MAX_DECODED_PIXELS,
});
}
let icc_profile = decoder.icc_profile().ok().flatten();
// Read EXIF orientation only for JPEG.
let orientation = if format == ImageFormat::Jpeg {
decoder.orientation().ok()
} else {
None
};
let img = DynamicImage::from_decoder(decoder).map_err(decode_error)?;
Ok((img, orientation, icc_profile))
})
.map_err(|_| {
decode_error(image::ImageError::Decoding(
image::error::DecodingError::new(
format.into(),
"image decoder rejected malformed data",
),
))
})??;
if let Some(orientation) = orientation {
img.apply_orientation(orientation);
}
let (source_width, source_height) = (img.width(), img.height());
if format == ImageFormat::WebP {
// Drop WebP alpha before resizing because the JPEG output has no alpha channel.
img = img.into_rgb8().into();
}
if let Some(resized) = resize_for_prompt(&img, mode) {
img = resized;
}
if let Some(icc_profile) = icc_profile
&& let Some(new_img) = convert_color_profile(&img, &icc_profile)
{
img = new_img;
}
let mut output = Vec::new();
let encoded = if output_format == ImageFormat::Png {
img.write_with_encoder(PngEncoder::new_with_quality(
&mut output,
CompressionType::Fast,
FilterType::Adaptive,
))
} else {
img.write_with_encoder(JpegEncoder::new_with_quality(
&mut output,
/*quality*/ 90,
))
};
encoded.map_err(|source| ImageProcessingError::Encode {
format: output_format,
source,
})?;
if output.len() > MAX_SANITIZED_IMAGE_BYTES {
return Err(ImageProcessingError::ImageTooLarge {
representation: "sanitized output",
size: output.len(),
max: MAX_SANITIZED_IMAGE_BYTES,
});
}
Ok(EncodedImage {
bytes: output.into(),
mime: output_format.to_mime_type().to_string(),
source_width,
source_height,
width: img.width(),
height: img.height(),
})
}
fn convert_color_profile(img: &DynamicImage, icc_profile: &[u8]) -> Option<DynamicImage> {
// Convert color profiles only for 8-bit RGB and RGBA images.
let layout = match img.color() {
image::ColorType::Rgba8 => moxcms::Layout::Rgba,
image::ColorType::Rgb8 => moxcms::Layout::Rgb,
_ => return None,
};
let parsed_icc_profile = ColorProfile::new_from_slice(icc_profile).ok()?;
let dest_profile = ColorProfile::new_srgb();
let transform_opts = moxcms::TransformOptions {
// Use floating point for more precise color conversion.
prefer_fixed_point: false,
..Default::default()
};
let transform = parsed_icc_profile
.create_transform_8bit(layout, &dest_profile, layout, transform_opts)
.ok()?;
let mut output = vec![0; img.as_bytes().len()];
transform.transform(img.as_bytes(), &mut output).ok()?;
let (width, height) = (img.width(), img.height());
if layout == moxcms::Layout::Rgba {
image::RgbaImage::from_raw(width, height, output).map(DynamicImage::ImageRgba8)
} else {
image::RgbImage::from_raw(width, height, output).map(DynamicImage::ImageRgb8)
}
}