fix(agent): settle runs after invalid image preparation

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,

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),
 }

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,16 +341,15 @@ 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",
                                     )));
                                 };
                                 output = next;
                             }
                         })
                         .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,26 +778,27 @@ async fn probe_source(
     source: &SourceBytes,
     initial: &[u8],
     budget: Arc<dyn SourceBudget>,
-) -> Result<codec::ImageProbe, String> {
+) -> Result<codec::ImageProbe, RunError> {
     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,
     };
     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<u8, String> {
+    async fn byte(&mut self) -> Result<u8, RunError> {
         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);

plugins/agent/native/src/owner/driver/tests/model_recovery/preset_acceptance.rs

@@ -154,6 +154,112 @@ fn text_messages(request: &Value) -> String {
         .join("\n")
 }
 
+#[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()?;
+    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<dyn Error + Send + Sync>>(())
+    }))
+    .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,

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 {