From 7988f73ad98033457581998d44cf04672c56629c Mon Sep 17 00:00:00 2001 From: rick12345 Date: Fri, 13 Nov 2026 03:35:28 +0800 Subject: [PATCH 1/2] fix(agent): settle runs after invalid image preparation Generated-by: OpenAI Codex --- plugins/agent/engine/src/driver.rs | 2 +- plugins/agent/engine/src/lib.rs | 2 + .../agent/engine/src/prepared_prompt/media.rs | 65 +++++++------ .../tests/model_recovery/preset_acceptance.rs | 97 +++++++++++++++++++ plugins/agent/native/src/owner/queue.rs | 1 + 5 files changed, 137 insertions(+), 30 deletions(-) diff --git a/plugins/agent/engine/src/driver.rs b/plugins/agent/engine/src/driver.rs index be6053cd81..6dc840d9fc 100644 --- a/plugins/agent/engine/src/driver.rs +++ b/plugins/agent/engine/src/driver.rs @@ -1275,7 +1275,7 @@ impl Engine { .extend(compaction.and_then(|commit| commit.candidate)); return Ok(()); } - Err(error @ RunError::CompactionMismatch) => { + Err(error @ (RunError::CompactionMismatch | RunError::InvalidImage(_))) => { self.publish( Some(claim), &view, diff --git a/plugins/agent/engine/src/lib.rs b/plugins/agent/engine/src/lib.rs index b04956f8b5..d990257924 100644 --- a/plugins/agent/engine/src/lib.rs +++ b/plugins/agent/engine/src/lib.rs @@ -102,6 +102,8 @@ pub enum RunError { Json(#[source] serde_json::Error), #[error("model input preparation failed: {0}")] Preparation(#[from] std::io::Error), + #[error("invalid image in model input: {0}")] + InvalidImage(String), #[error("model input cleanup was not confirmed: {0}")] PreparationCleanup(#[source] std::io::Error), } diff --git a/plugins/agent/engine/src/prepared_prompt/media.rs b/plugins/agent/engine/src/prepared_prompt/media.rs index 9302948c5b..bc1d811122 100644 --- a/plugins/agent/engine/src/prepared_prompt/media.rs +++ b/plugins/agent/engine/src/prepared_prompt/media.rs @@ -317,7 +317,10 @@ impl MediaView { // Keep the typed codec result inside the arena result so the // worker cannot flatten cancellation or writer errors. let _lease = lease; - let image = codec::decode(&encoded, &plan).map_err(io::Error::other)?; + let image = match codec::decode(&encoded, &plan) { + Ok(image) => image, + Err(error) => return Ok(Err(RunError::InvalidImage(error))), + }; loop { let start = writer.position(); if let Err(error) = image.write_image( @@ -328,7 +331,7 @@ impl MediaView { cancellation: &cancelled, }, ) { - return Ok(Err(error)); + return Ok(Err(RunError::preparation(error))); } let range = start..writer.position(); if limits @@ -338,7 +341,7 @@ impl MediaView { return Ok(Ok((range, output))); } let Some(next) = smaller_geometry(original, output, &limits) else { - return Ok(Err(io::Error::other( + return Ok(Err(failure( "whole image cannot satisfy the selected route's image byte limit and minimum dimensions", ))); }; @@ -346,8 +349,7 @@ impl MediaView { } }) .await - .map_err(RunError::Preparation)? - .map_err(RunError::preparation)?; + .map_err(RunError::Preparation)??; item.pending .0 .set(MaterializedImage { @@ -579,8 +581,7 @@ pub(super) async fn prepare( &read.buffer.as_ref()[..read.bytes_read], budget.clone(), ) - .await - .map_err(failure)?; + .await?; drop(read); let output = geometry(probe.geometry, &limits)?; let length = media @@ -777,18 +778,17 @@ async fn probe_source( source: &SourceBytes, initial: &[u8], budget: Arc, -) -> Result { +) -> Result { match codec::probe(initial) { Ok(probe) => return Ok(probe), - Err(error) if !initial.starts_with(&[0xff, 0xd8]) => return Err(error), + Err(error) if !initial.starts_with(&[0xff, 0xd8]) => { + return Err(RunError::InvalidImage(error)); + } Err(_) => {} } let mut input = HeaderStream { cursor: source.open(budget.clone()), - buffer: budget - .allocate(32 * 1024) - .await - .map_err(|error| error.to_string())?, + buffer: budget.allocate(32 * 1024).await.map_err(failure)?, position: 0, used: 0, eof: false, @@ -796,7 +796,9 @@ async fn probe_source( input.skip(2).await?; loop { if input.byte().await? != 0xff { - return Err("JPEG header marker is missing".into()); + return Err(RunError::InvalidImage( + "JPEG header marker is missing".into(), + )); } let marker = loop { let marker = input.byte().await?; @@ -805,20 +807,24 @@ async fn probe_source( } }; if marker == 0xda || marker == 0xd9 { - return Err("JPEG dimensions are missing before pixel data".into()); + return Err(RunError::InvalidImage( + "JPEG dimensions are missing before pixel data".into(), + )); } if marker == 0x01 || (0xd0..=0xd8).contains(&marker) { continue; } let length = u16::from_be_bytes([input.byte().await?, input.byte().await?]); if length < 2 { - return Err("JPEG header segment has an invalid length".into()); + return Err(RunError::InvalidImage( + "JPEG header segment has an invalid length".into(), + )); } if (0xc0..=0xcf).contains(&marker) && !matches!(marker, 0xc4 | 0xc8 | 0xcc) { let mut header = budget .allocate(usize::from(length) + 4) .await - .map_err(|error| error.to_string())?; + .map_err(failure)?; header.as_mut()[..6].copy_from_slice(&[ 0xff, 0xd8, @@ -830,7 +836,8 @@ async fn probe_source( input .copy(&mut header.as_mut()[6..usize::from(length) + 4]) .await?; - return codec::probe(&header.as_ref()[..usize::from(length) + 4]); + return codec::probe(&header.as_ref()[..usize::from(length) + 4]) + .map_err(RunError::InvalidImage); } input.skip(usize::from(length) - 2).await?; } @@ -843,37 +850,37 @@ struct HeaderStream { eof: bool, } impl HeaderStream { - async fn fill(&mut self) -> Result<(), String> { + async fn fill(&mut self) -> Result<(), RunError> { if self.position < self.used { return Ok(()); } if self.eof { - return Err("image header ended before dimensions".into()); + return Err(RunError::InvalidImage( + "image header ended before dimensions".into(), + )); } // A cursor returns its supplied work buffer. A temporary empty buffer // only transfers ownership across the awaited bounded read. let buffer = std::mem::replace(&mut self.buffer, ChargedBuffer::new(Vec::new(), ())); - let read = self - .cursor - .read_chunk(buffer) - .await - .map_err(|error| error.to_string())?; + let read = self.cursor.read_chunk(buffer).await.map_err(failure)?; self.buffer = read.buffer; self.position = 0; self.used = read.bytes_read; self.eof = read.eof; if self.used == 0 { - return Err("image header made no progress".into()); + return Err(RunError::InvalidImage( + "image header made no progress".into(), + )); } Ok(()) } - async fn byte(&mut self) -> Result { + async fn byte(&mut self) -> Result { self.fill().await?; let value = self.buffer.as_ref()[self.position]; self.position += 1; Ok(value) } - async fn skip(&mut self, mut bytes: usize) -> Result<(), String> { + async fn skip(&mut self, mut bytes: usize) -> Result<(), RunError> { while bytes != 0 { self.fill().await?; let count = bytes.min(self.used - self.position); @@ -882,7 +889,7 @@ impl HeaderStream { } Ok(()) } - async fn copy(&mut self, mut output: &mut [u8]) -> Result<(), String> { + async fn copy(&mut self, mut output: &mut [u8]) -> Result<(), RunError> { while !output.is_empty() { self.fill().await?; let count = output.len().min(self.used - self.position); diff --git a/plugins/agent/native/src/owner/driver/tests/model_recovery/preset_acceptance.rs b/plugins/agent/native/src/owner/driver/tests/model_recovery/preset_acceptance.rs index ad6fbc6fd2..74a442abe6 100644 --- a/plugins/agent/native/src/owner/driver/tests/model_recovery/preset_acceptance.rs +++ b/plugins/agent/native/src/owner/driver/tests/model_recovery/preset_acceptance.rs @@ -154,6 +154,103 @@ fn text_messages(request: &Value) -> String { .join("\n") } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn truncated_cell_image_terminates_the_run() -> TestResult<()> { + use base64::Engine as _; + + let mut rig = Rig::new()?; + let mut provider = Provider::start().await?; + let outcome = AssertUnwindSafe(tokio::time::timeout(Duration::from_secs(45), async { + let base = rig.directory.path().to_path_buf(); + let boot = bootstrap(&base, provider.url.clone(), true).await?; + attach(&mut rig, &boot).await?; + let client = rig.client.as_ref().unwrap(); + let model = model_options(client, &provider.url, true, json!({"vision":true})).await?; + + let mut state = 1_u32; + let pixels = image::RgbImage::from_fn(256, 128, |_, _| { + state ^= state << 13; + state ^= state >> 17; + state ^= state << 5; + let [r, g, b, _] = state.to_le_bytes(); + image::Rgb([r, g, b]) + }); + let mut png = std::io::Cursor::new(Vec::new()); + pixels.write_to(&mut png, image::ImageFormat::Png)?; + let encoded = base64::engine::general_purpose::STANDARD.encode(png.into_inner()); + let tail = &encoded[encoded.len() - 65_536..]; + assert!(!tail.starts_with("iVBOR")); + + for (data, expected) in [(tail, "failed"), (encoded.as_str(), "completed")] { + let session = create(client, &base, &model, None).await?; + let id = session["id"].as_str().unwrap(); + let sent = call( + client, + "agent", + "send", + json!({ + "sessionId": id, "message": "Inspect the image.", "delivery": "queue" + }), + ) + .await?; + assert_eq!(sent["kind"], "queued", "{sent}"); + let (_, reply) = provider.next().await?; + reply + .send(tool_response( + "exec", + json!({ + "code": format!("image('data:image/png;base64,{data}');") + }), + )) + .map_err(|_| "image tool response lost")?; + + if expected == "completed" { + let (request, reply) = provider.next().await?; + assert!( + request["messages"] + .to_string() + .contains("data:image/png;base64,") + ); + reply + .send(completion( + json!({"content":"Valid image received."}), + "stop", + )) + .map_err(|_| "final model response lost")?; + } + let target = json!({"kind":"message","messageId":sent["messageId"]}); + loop { + let work = call( + client, + "agent", + "work", + json!({"sessionId":id,"target":target}), + ) + .await?; + let kind = work["execution"]["state"]["kind"].as_str(); + if matches!(kind, Some("completed" | "failed" | "cancelled")) { + assert_eq!(kind, Some(expected), "{work}"); + break; + } + tokio::task::yield_now().await; + } + assert!( + provider.requests.try_recv().is_err(), + "unexpected model request" + ); + } + Ok::<_, Box>(()) + })) + .catch_unwind() + .await; + let cleanup = rig.finish().await; + assert!(cleanup.is_empty(), "{cleanup:?}"); + match outcome { + Ok(result) => result?, + Err(panic) => std::panic::resume_unwind(panic), + } +} + async fn tool_round( rig: &Rig, provider: &mut Provider, diff --git a/plugins/agent/native/src/owner/queue.rs b/plugins/agent/native/src/owner/queue.rs index 1dd5898d3d..a97173cc74 100644 --- a/plugins/agent/native/src/owner/queue.rs +++ b/plugins/agent/native/src/owner/queue.rs @@ -615,6 +615,7 @@ impl Owner { RunError::Internal(_) => "internal", RunError::Json(_) => "json", RunError::Preparation(_) => "preparation", + RunError::InvalidImage(_) => "invalid_image", RunError::PreparationCleanup(_) => "preparation_cleanup", }; let tool_error_variant = match error { From f3be91d5246cf459654c265b7e421c7f96c0c096 Mon Sep 17 00:00:00 2001 From: rick12345 Date: Fri, 13 Nov 2026 04:14:40 +0800 Subject: [PATCH 2/2] test(agent): size the runtime stack for the image regression fixture Generated-by: OpenAI Codex --- .../tests/model_recovery/preset_acceptance.rs | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/plugins/agent/native/src/owner/driver/tests/model_recovery/preset_acceptance.rs b/plugins/agent/native/src/owner/driver/tests/model_recovery/preset_acceptance.rs index 74a442abe6..da49619c66 100644 --- a/plugins/agent/native/src/owner/driver/tests/model_recovery/preset_acceptance.rs +++ b/plugins/agent/native/src/owner/driver/tests/model_recovery/preset_acceptance.rs @@ -154,8 +154,17 @@ fn text_messages(request: &Value) -> String { .join("\n") } -#[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn truncated_cell_image_terminates_the_run() -> TestResult<()> { +#[test] +fn truncated_cell_image_terminates_the_run() -> TestResult<()> { + tokio::runtime::Builder::new_multi_thread() + .worker_threads(2) + .thread_stack_size(16 * 1024 * 1024) + .enable_all() + .build()? + .block_on(async { tokio::spawn(check_truncated_cell_image()).await? }) +} + +async fn check_truncated_cell_image() -> TestResult<()> { use base64::Engine as _; let mut rig = Rig::new()?;