mirror of
https://github.com/openai/codex.git
synced 2026-10-24 18:18:41 +00:00
Harden remote image materialization
This commit is contained in:
@@ -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)
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user