diff --git a/codex-rs/core/src/api_bridge.rs b/codex-rs/core/src/api_bridge.rs index b9f802ae6..79fd67d65 100644 --- a/codex-rs/core/src/api_bridge.rs +++ b/codex-rs/core/src/api_bridge.rs @@ -33,12 +33,20 @@ pub(crate) fn map_api_error(err: ApiError) -> CodexErr { headers, body, } => { - if status == http::StatusCode::INTERNAL_SERVER_ERROR { + let body_text = body.unwrap_or_default(); + + if status == http::StatusCode::BAD_REQUEST { + if body_text + .contains("The image data you provided does not represent a valid image") + { + CodexErr::InvalidImageRequest() + } else { + CodexErr::InvalidRequest(body_text) + } + } else if status == http::StatusCode::INTERNAL_SERVER_ERROR { CodexErr::InternalServerError } else if status == http::StatusCode::TOO_MANY_REQUESTS { - if let Some(body) = body - && let Ok(err) = serde_json::from_str::(&body) - { + if let Ok(err) = serde_json::from_str::(&body_text) { if err.error.error_type.as_deref() == Some("usage_limit_reached") { let rate_limits = headers.as_ref().and_then(parse_rate_limit); let resets_at = err @@ -62,7 +70,7 @@ pub(crate) fn map_api_error(err: ApiError) -> CodexErr { } else { CodexErr::UnexpectedStatus(UnexpectedResponseError { status, - body: body.unwrap_or_default(), + body: body_text, request_id: extract_request_id(headers.as_ref()), }) } diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index 8b1dfafd2..f76cd7de7 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -2039,6 +2039,13 @@ pub(crate) async fn run_task( // Aborted turn is reported via a different event. break; } + Err(CodexErr::InvalidImageRequest()) => { + let mut state = sess.state.lock().await; + error_or_panic( + "Invalid image detected, replacing it in the last turn to prevent poisoning", + ); + state.history.replace_last_turn_images("Invalid image"); + } Err(e) => { info!("Turn error: {e:#}"); let event = EventMsg::Error(e.to_error_event(None)); @@ -2146,6 +2153,8 @@ async fn run_turn( } Err(CodexErr::UsageNotIncluded) => return Err(CodexErr::UsageNotIncluded), Err(e @ CodexErr::QuotaExceeded) => return Err(e), + Err(e @ CodexErr::InvalidImageRequest()) => return Err(e), + Err(e @ CodexErr::InvalidRequest(_)) => return Err(e), Err(e @ CodexErr::RefreshTokenFailed(_)) => return Err(e), Err(e) => { // Use the configured provider-specific stream retry budget. diff --git a/codex-rs/core/src/context_manager/history.rs b/codex-rs/core/src/context_manager/history.rs index 7277f0b0c..b9a9c58f6 100644 --- a/codex-rs/core/src/context_manager/history.rs +++ b/codex-rs/core/src/context_manager/history.rs @@ -5,6 +5,8 @@ use crate::truncate::approx_token_count; use crate::truncate::approx_tokens_from_byte_count; use crate::truncate::truncate_function_output_items_with_policy; use crate::truncate::truncate_text; +use codex_protocol::models::ContentItem; +use codex_protocol::models::FunctionCallOutputContentItem; use codex_protocol::models::FunctionCallOutputPayload; use codex_protocol::models::ResponseItem; use codex_protocol::protocol::TokenUsage; @@ -118,6 +120,37 @@ impl ContextManager { self.items = items; } + pub(crate) fn replace_last_turn_images(&mut self, placeholder: &str) { + let Some(last_item) = self.items.last_mut() else { + return; + }; + + match last_item { + ResponseItem::Message { role, content, .. } if role == "user" => { + for item in content.iter_mut() { + if matches!(item, ContentItem::InputImage { .. }) { + *item = ContentItem::InputText { + text: placeholder.to_string(), + }; + } + } + } + ResponseItem::FunctionCallOutput { output, .. } => { + let Some(content_items) = output.content_items.as_mut() else { + return; + }; + for item in content_items.iter_mut() { + if matches!(item, FunctionCallOutputContentItem::InputImage { .. }) { + *item = FunctionCallOutputContentItem::InputText { + text: placeholder.to_string(), + }; + } + } + } + _ => {} + } + } + pub(crate) fn update_token_info( &mut self, usage: &TokenUsage, diff --git a/codex-rs/core/src/error.rs b/codex-rs/core/src/error.rs index 9130b40e1..a25261d64 100644 --- a/codex-rs/core/src/error.rs +++ b/codex-rs/core/src/error.rs @@ -103,6 +103,14 @@ pub enum CodexErr { #[error("{0}")] UnexpectedStatus(UnexpectedResponseError), + /// Invalid request. + #[error("{0}")] + InvalidRequest(String), + + /// Invalid image. + #[error("Image poisoning")] + InvalidImageRequest(), + #[error("{0}")] UsageLimitReached(UsageLimitReachedError), diff --git a/codex-rs/core/src/util.rs b/codex-rs/core/src/util.rs index 698e02fed..5304a89ac 100644 --- a/codex-rs/core/src/util.rs +++ b/codex-rs/core/src/util.rs @@ -16,11 +16,11 @@ pub(crate) fn backoff(attempt: u64) -> Duration { Duration::from_millis((base as f64 * jitter) as u64) } -pub(crate) fn error_or_panic(message: String) { +pub(crate) fn error_or_panic(message: impl std::string::ToString) { if cfg!(debug_assertions) || env!("CARGO_PKG_VERSION").contains("alpha") { - panic!("{message}"); + panic!("{}", message.to_string()); } else { - error!("{message}"); + error!("{}", message.to_string()); } } diff --git a/codex-rs/core/tests/common/responses.rs b/codex-rs/core/tests/common/responses.rs index 932ddb51a..a8209b513 100644 --- a/codex-rs/core/tests/common/responses.rs +++ b/codex-rs/core/tests/common/responses.rs @@ -518,6 +518,32 @@ pub fn sse_response(body: String) -> ResponseTemplate { .set_body_raw(body, "text/event-stream") } +pub async fn mount_response_once(server: &MockServer, response: ResponseTemplate) -> ResponseMock { + let (mock, response_mock) = base_mock(); + mock.respond_with(response) + .up_to_n_times(1) + .mount(server) + .await; + response_mock +} + +pub async fn mount_response_once_match( + server: &MockServer, + matcher: M, + response: ResponseTemplate, +) -> ResponseMock +where + M: wiremock::Match + Send + Sync + 'static, +{ + let (mock, response_mock) = base_mock(); + mock.and(matcher) + .respond_with(response) + .up_to_n_times(1) + .mount(server) + .await; + response_mock +} + fn base_mock() -> (MockBuilder, ResponseMock) { let response_mock = ResponseMock::new(); let mock = Mock::given(method("POST")) diff --git a/codex-rs/core/tests/suite/view_image.rs b/codex-rs/core/tests/suite/view_image.rs index 394c5c2f1..6c0f6dcc8 100644 --- a/codex-rs/core/tests/suite/view_image.rs +++ b/codex-rs/core/tests/suite/view_image.rs @@ -474,3 +474,82 @@ async fn view_image_tool_errors_when_file_missing() -> anyhow::Result<()> { Ok(()) } + +#[cfg(not(debug_assertions))] +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn replaces_invalid_local_image_after_bad_request() -> anyhow::Result<()> { + skip_if_no_network!(Ok(())); + + let server = start_mock_server().await; + + const INVALID_IMAGE_ERROR: &str = + "The image data you provided does not represent a valid image"; + + let invalid_image_mock = responses::mount_response_once_match( + &server, + body_string_contains("\"input_image\""), + ResponseTemplate::new(400) + .insert_header("content-type", "text/plain") + .set_body_string(INVALID_IMAGE_ERROR), + ) + .await; + + let success_response = sse(vec![ + ev_response_created("resp-2"), + ev_assistant_message("msg-1", "done"), + ev_completed("resp-2"), + ]); + + let completion_mock = responses::mount_sse_once(&server, success_response).await; + + let TestCodex { + codex, + cwd, + session_configured, + .. + } = test_codex().build(&server).await?; + + let rel_path = "assets/poisoned.png"; + let abs_path = cwd.path().join(rel_path); + if let Some(parent) = abs_path.parent() { + std::fs::create_dir_all(parent)?; + } + let image = ImageBuffer::from_pixel(1024, 512, Rgba([10u8, 20, 30, 255])); + image.save(&abs_path)?; + + let session_model = session_configured.model.clone(); + + codex + .submit(Op::UserTurn { + items: vec![UserInput::LocalImage { + path: abs_path.clone(), + }], + final_output_json_schema: None, + cwd: cwd.path().to_path_buf(), + approval_policy: AskForApproval::Never, + sandbox_policy: SandboxPolicy::DangerFullAccess, + model: session_model, + effort: None, + summary: ReasoningSummary::Auto, + }) + .await?; + + wait_for_event(&codex, |event| matches!(event, EventMsg::TaskComplete(_))).await; + + let first_body = invalid_image_mock.single_request().body_json(); + assert!( + find_image_message(&first_body).is_some(), + "initial request should include the uploaded image" + ); + + let second_request = completion_mock.single_request(); + let second_body = second_request.body_json(); + assert!( + find_image_message(&second_body).is_none(), + "second request should replace the invalid image" + ); + let user_texts = second_request.message_input_texts("user"); + assert!(user_texts.iter().any(|text| text == "Invalid image")); + + Ok(()) +}