Harden remote image materialization

This commit is contained in:
Rohit Arunachalam
2026-06-09 14:40:17 -07:00
parent 09e690afa1
commit 0bb8d84568
+104 -90
View File
@@ -1,55 +1,28 @@
use std::borrow::Cow;
use std::path::Path;
use std::sync::LazyLock;
use std::time::Duration;
use anyhow::Context;
use anyhow::Result;
use anyhow::bail;
use codex_login::default_client::build_reqwest_client;
use codex_protocol::models::ContentItem;
use codex_protocol::models::FunctionCallOutputContentItem;
use codex_protocol::models::ResponseItem;
use codex_utils_image::ImageProcessingError;
use codex_utils_image::PromptImageMode;
use codex_utils_image::load_for_prompt_bytes;
use reqwest::Client;
use thiserror::Error;
use reqwest::Response;
use tracing::warn;
use url::Url;
const IMAGE_DOWNLOAD_TIMEOUT: Duration = Duration::from_secs(60);
const MAX_IMAGE_DOWNLOAD_BYTES: u64 = 50 * 1024 * 1024;
const IMAGE_MATERIALIZATION_ERROR_PLACEHOLDER: &str =
"image content omitted because it could not be downloaded or processed";
static IMAGE_URL_CLIENT: LazyLock<Client> = LazyLock::new(Client::new);
#[derive(Debug, Error)]
pub(crate) enum ImageUrlMaterializationError {
#[error("failed to download image: {0}")]
Download(#[from] reqwest::Error),
#[error("failed to process downloaded image: {0}")]
Processing(#[from] ImageProcessingError),
}
/// Downloads an HTTP(S) image and returns it as an inline data URL.
///
/// Non-HTTP(S) values, including existing data URLs, are returned unchanged.
pub(crate) async fn materialize_http_image_url<'a>(
client: &Client,
image_url: &'a str,
) -> Result<Cow<'a, str>, ImageUrlMaterializationError> {
let Ok(url) = Url::parse(image_url) else {
return Ok(Cow::Borrowed(image_url));
};
if !matches!(url.scheme(), "http" | "https") {
return Ok(Cow::Borrowed(image_url));
}
let response = client.get(url.as_str()).send().await?.error_for_status()?;
let bytes = response.bytes().await?;
let image = load_for_prompt_bytes(
Path::new(url.path()),
bytes.to_vec(),
PromptImageMode::Original,
)?;
Ok(Cow::Owned(image.into_data_url()))
}
static IMAGE_URL_CLIENT: LazyLock<Client> = LazyLock::new(build_reqwest_client);
/// Materializes HTTP(S) image URLs in a newly recorded batch.
///
@@ -65,11 +38,40 @@ pub(crate) async fn materialize_conversation_item_images<'a>(
let mut items = items.to_vec();
for item in &mut items {
match item {
ResponseItem::Message { content, .. } => materialize_content_items(content).await,
ResponseItem::Message { content, .. } => {
for content_item in content {
let result = match content_item {
ContentItem::InputImage { image_url, .. } => {
materialize_image_url(image_url).await
}
ContentItem::InputText { .. } | ContentItem::OutputText { .. } => continue,
};
if let Err(err) = result {
warn!(error = %err, "failed to materialize remote message image");
*content_item = ContentItem::InputText {
text: IMAGE_MATERIALIZATION_ERROR_PLACEHOLDER.to_string(),
};
}
}
}
ResponseItem::FunctionCallOutput { output, .. }
| ResponseItem::CustomToolCallOutput { output, .. } => {
if let Some(content_items) = output.content_items_mut() {
materialize_function_call_output_items(content_items).await;
for content_item in content_items {
let result = match content_item {
FunctionCallOutputContentItem::InputImage { image_url, .. } => {
materialize_image_url(image_url).await
}
FunctionCallOutputContentItem::InputText { .. }
| FunctionCallOutputContentItem::EncryptedContent { .. } => continue,
};
if let Err(err) = result {
warn!(error = %err, "failed to materialize remote tool output image");
*content_item = FunctionCallOutputContentItem::InputText {
text: IMAGE_MATERIALIZATION_ERROR_PLACEHOLDER.to_string(),
};
}
}
}
}
ResponseItem::Reasoning { .. }
@@ -94,13 +96,21 @@ pub(crate) async fn materialize_conversation_item_images<'a>(
fn response_item_has_http_image(item: &ResponseItem) -> bool {
match item {
ResponseItem::Message { content, .. } => content.iter().any(|item| match item {
ContentItem::InputImage { image_url, .. } => is_http_url(image_url),
ContentItem::InputImage { image_url, .. } => parse_http_url(image_url).is_some(),
ContentItem::InputText { .. } | ContentItem::OutputText { .. } => false,
}),
ResponseItem::FunctionCallOutput { output, .. }
| ResponseItem::CustomToolCallOutput { output, .. } => output
.content_items()
.is_some_and(|items| items.iter().any(function_call_output_item_has_http_image)),
| ResponseItem::CustomToolCallOutput { output, .. } => {
output.content_items().is_some_and(|items| {
items.iter().any(|item| match item {
FunctionCallOutputContentItem::InputImage { image_url, .. } => {
parse_http_url(image_url).is_some()
}
FunctionCallOutputContentItem::InputText { .. }
| FunctionCallOutputContentItem::EncryptedContent { .. } => false,
})
})
}
ResponseItem::Reasoning { .. }
| ResponseItem::AgentMessage { .. }
| ResponseItem::LocalShellCall { .. }
@@ -117,56 +127,60 @@ fn response_item_has_http_image(item: &ResponseItem) -> bool {
}
}
fn function_call_output_item_has_http_image(item: &FunctionCallOutputContentItem) -> bool {
match item {
FunctionCallOutputContentItem::InputImage { image_url, .. } => is_http_url(image_url),
FunctionCallOutputContentItem::InputText { .. }
| FunctionCallOutputContentItem::EncryptedContent { .. } => false,
fn parse_http_url(value: &str) -> Option<Url> {
// Avoid parsing potentially large inline data URLs on the common no-op path.
let (scheme, _) = value.split_once(':')?;
if !scheme.eq_ignore_ascii_case("http") && !scheme.eq_ignore_ascii_case("https") {
return None;
}
Url::parse(value).ok()
}
fn is_http_url(value: &str) -> bool {
Url::parse(value).is_ok_and(|url| matches!(url.scheme(), "http" | "https"))
}
async fn materialize_content_items(items: &mut [ContentItem]) {
for item in items {
let result = match item {
ContentItem::InputImage { image_url, .. } => materialize_image_url(image_url).await,
ContentItem::InputText { .. } | ContentItem::OutputText { .. } => continue,
};
if let Err(err) = result {
warn!(error = %err, "failed to materialize remote message image");
*item = ContentItem::InputText {
text: IMAGE_MATERIALIZATION_ERROR_PLACEHOLDER.to_string(),
};
}
}
}
async fn materialize_function_call_output_items(items: &mut [FunctionCallOutputContentItem]) {
for item in items {
let result = match item {
FunctionCallOutputContentItem::InputImage { image_url, .. } => {
materialize_image_url(image_url).await
}
FunctionCallOutputContentItem::InputText { .. }
| FunctionCallOutputContentItem::EncryptedContent { .. } => continue,
};
if let Err(err) = result {
warn!(error = %err, "failed to materialize remote tool output image");
*item = FunctionCallOutputContentItem::InputText {
text: IMAGE_MATERIALIZATION_ERROR_PLACEHOLDER.to_string(),
};
}
}
}
async fn materialize_image_url(image_url: &mut String) -> Result<(), ImageUrlMaterializationError> {
if let Cow::Owned(materialized_url) =
materialize_http_image_url(&IMAGE_URL_CLIENT, image_url).await?
{
*image_url = materialized_url;
}
async fn materialize_image_url(image_url: &mut String) -> Result<()> {
let Some(url) = parse_http_url(image_url) else {
return Ok(());
};
let response = IMAGE_URL_CLIENT
.get(url)
.timeout(IMAGE_DOWNLOAD_TIMEOUT)
.send()
.await
.map_err(reqwest::Error::without_url)
.context("failed to download image")?
.error_for_status()
.map_err(reqwest::Error::without_url)
.context("failed to download image")?;
let bytes = read_response_body_with_limit(response).await?;
let image = load_for_prompt_bytes(
Path::new("<remote image>"),
bytes,
PromptImageMode::Original,
)
.context("failed to process downloaded image")?;
*image_url = image.into_data_url();
Ok(())
}
async fn read_response_body_with_limit(mut response: Response) -> Result<Vec<u8>> {
if response
.content_length()
.is_some_and(|length| length > MAX_IMAGE_DOWNLOAD_BYTES)
{
bail!("downloaded image exceeded the maximum size of {MAX_IMAGE_DOWNLOAD_BYTES} bytes");
}
let mut body = Vec::new();
while let Some(chunk) = response
.chunk()
.await
.map_err(reqwest::Error::without_url)
.context("failed to download image")?
{
let next_length = body.len() as u64 + chunk.len() as u64;
if next_length > MAX_IMAGE_DOWNLOAD_BYTES {
bail!("downloaded image exceeded the maximum size of {MAX_IMAGE_DOWNLOAD_BYTES} bytes");
}
body.extend_from_slice(&chunk);
}
Ok(body)
}