From 74a00f9edc61727cf780093dd281e4a71220695d Mon Sep 17 00:00:00 2001 From: Hare Date: Sun, 4 Oct 2026 23:17:43 +0900 Subject: [PATCH] fix: recover from provider image size rejections --- Cargo.lock | 7 + crates/agen/Cargo.toml | 1 + crates/agen/src/engine.rs | 46 ++- crates/agen/src/image_recovery.rs | 106 ++++++ crates/agen/src/lib.rs | 1 + crates/agen/src/llm_client/error.rs | 29 ++ crates/agen/src/prune.rs | 13 +- crates/agen/tests/image_recovery_test.rs | 350 ++++++++++++++++++ crates/session-store/src/public_index.rs | 1 + crates/session-store/src/public_snapshot.rs | 328 ++++++++++++++-- crates/session-store/src/segment_log.rs | 13 + crates/worker/src/image_recovery_tests.rs | 221 +++++++++++ crates/worker/src/ipc/protocol_session.rs | 1 + crates/worker/src/segment_log_sink.rs | 1 + crates/worker/src/session_history.rs | 11 + crates/worker/src/worker.rs | 22 ++ .../src/lib/workspace/console/model.test.ts | 23 ++ 17 files changed, 1135 insertions(+), 39 deletions(-) create mode 100644 crates/agen/src/image_recovery.rs create mode 100644 crates/agen/tests/image_recovery_test.rs create mode 100644 crates/worker/src/image_recovery_tests.rs diff --git a/Cargo.lock b/Cargo.lock index ce8416376..e7470de00 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -64,6 +64,7 @@ dependencies = [ "dotenv", "eventsource-stream", "futures", + "imagesize", "reqwest", "schemars", "serde", @@ -2188,6 +2189,12 @@ dependencies = [ "winapi-util", ] +[[package]] +name = "imagesize" +version = "0.14.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09e54e57b4c48b40f7aec75635392b12b3421fa26fe8b4332e63138ed278459c" + [[package]] name = "include_dir" version = "0.7.4" diff --git a/crates/agen/Cargo.toml b/crates/agen/Cargo.toml index febf6df18..e01b2ca40 100644 --- a/crates/agen/Cargo.toml +++ b/crates/agen/Cargo.toml @@ -32,6 +32,7 @@ reqwest = { version = "0.13", default-features = false, features = ["stream", "j eventsource-stream = "0.2" zstd = "0.13" base64 = "0.22.1" +imagesize = "0.14" chrono = { version = "0.4", default-features = false, features = ["serde", "clock"], optional = true } agen-macros = { workspace = true } diff --git a/crates/agen/src/engine.rs b/crates/agen/src/engine.rs index decbfbc9b..d35c1778b 100644 --- a/crates/agen/src/engine.rs +++ b/crates/agen/src/engine.rs @@ -626,6 +626,10 @@ pub struct Engine { /// engine history. An error rejects the item and aborts the turn, allowing /// upper layers to make durable storage the commit gate. history_append_cbs: Vec Result<(), String> + Send + Sync>>, + /// Durable host commit gate for image-result corrections. Runs before the + /// correction becomes live or a replacement request can be sent. + image_rejection_handler: + Option, &Item) -> Result<(), String> + Send + Sync>>, /// Request configuration (max_tokens, temperature, etc.) request_config: RequestConfig, /// Cancel notification channel (for interrupting execution) @@ -971,6 +975,16 @@ impl Engine { self.history_append_cbs.push(Box::new(callback)); } + /// Set the host's durable commit gate for correcting a rejected ViewImage + /// result. The original annotation/identity must be preserved. A failed + /// commit leaves both history and the outgoing request untouched. + pub fn set_image_rejection_handler( + &mut self, + handler: impl Fn(&HistoryEntry, &Item) -> Result<(), String> + Send + Sync + 'static, + ) { + self.image_rejection_handler = Some(Box::new(handler)); + } + fn emit_history_append(&self, item: &Item) -> Result<(), EngineError> { for cb in &self.history_append_cbs { cb(item).map_err(EngineError::HistoryAppend)?; @@ -3205,7 +3219,34 @@ impl Engine { "Sending request to LLM" ); - let stream = self.open_stream_with_retry(request, turn, llm_call).await?; + let mut request = request; + let stream = loop { + match self + .open_stream_with_retry(request.clone(), turn, llm_call) + .await + { + Ok(stream) => break stream, + Err(EngineError::Client(error)) if error.is_image_size_rejection() => { + let Some((history_index, request_index)) = + crate::image_recovery::largest_candidate(history, &request) + else { + return Err(EngineError::Client(error)); + }; + let entry = &history.entries()[history_index]; + let replacement = crate::image_recovery::rejected_result(&entry.item); + if let Some(commit) = &self.image_rejection_handler { + commit(entry, &replacement).map_err(EngineError::HistoryAppend)?; + } + history.entries_mut()[history_index].item = replacement.clone(); + // Rebuild the provider request from the corrected projection, + // retaining prior prune/interceptor decisions. No tool runs + // or assistant output have occurred for this failed request. + request.items[request_index] = replacement; + self.emit_warning(crate::image_recovery::REJECTION_MESSAGE); + } + Err(error) => return Err(error), + } + }; self.tool_call_collector.begin_response(); let mut stream = ResponseStreamPump::start(stream); let mut early_tools: Option = None; @@ -3567,6 +3608,7 @@ impl Engine { warning_cbs: Vec::new(), tool_result_cbs: Vec::new(), history_append_cbs: Vec::new(), + image_rejection_handler: None, request_config: RequestConfig::default(), cancel_tx, cancel_rx, @@ -3856,6 +3898,7 @@ impl Engine { warning_cbs: self.warning_cbs, tool_result_cbs: self.tool_result_cbs, history_append_cbs: self.history_append_cbs, + image_rejection_handler: self.image_rejection_handler, request_config: self.request_config, cancel_tx: self.cancel_tx, @@ -4040,6 +4083,7 @@ impl Engine { warning_cbs: self.warning_cbs, tool_result_cbs: self.tool_result_cbs, history_append_cbs: self.history_append_cbs, + image_rejection_handler: self.image_rejection_handler, request_config: self.request_config, cancel_tx: self.cancel_tx, diff --git a/crates/agen/src/image_recovery.rs b/crates/agen/src/image_recovery.rs new file mode 100644 index 000000000..dc38014fa --- /dev/null +++ b/crates/agen/src/image_recovery.rs @@ -0,0 +1,106 @@ +//! Recovery is driven by provider rejection, never a locally chosen image limit. +use crate::{History, Item, Role, ToolResultDisposition, llm_client::Request, tool::Attachment}; + +pub(crate) const REJECTION_MESSAGE: &str = "ViewImage image omitted from model input after the provider rejected image size. This is a recovery candidate, not proof that this particular image caused the rejection. Resize or crop the image and retry ViewImage; do not read the unchanged image again."; + +/// Keep the most recent response's tool results eligible, including early tool +/// results interleaved with that response's later text/calls. Empty boundaries +/// from older versions' failed requests are not evidence of accepted input. +pub(crate) fn accepted_input_end<'a>(items: impl Iterator) -> Option { + let mut response_start = None; + let mut accepted_prefix = None; + for (index, item) in items.enumerate() { + match item { + Item::AssistantResponseBoundary { .. } => response_start = Some(index), + Item::Message { + role: Role::Assistant, + .. + } + | Item::ToolCall { .. } + | Item::Reasoning { .. } => { + accepted_prefix = Some(response_start.unwrap_or(index)); + } + Item::Message { + role: Role::User, .. + } => response_start = None, + _ => {} + } + } + accepted_prefix +} + +/// Select only a ViewImage result actually sent in the rejected request. The +/// tool contract returns exactly one image, so each correction removes one +/// image and preserves the tool call/result pair. Other tools and user uploads +/// are deliberately not rewritten by this recovery policy. +pub(crate) fn largest_candidate( + history: &History, + request: &Request, +) -> Option<(usize, usize)> { + let start = accepted_input_end(history.items()).unwrap_or(0); + let mut largest = None; + for (index, entry) in history.entries().iter().enumerate().skip(start) { + let Item::ToolResult { + call_id, + attachments, + .. + } = &entry.item + else { + continue; + }; + let [Attachment::Image(image)] = attachments.as_slice() else { + continue; + }; + let from_view_image = history.items().any(|item| { + matches!( + item, + Item::ToolCall { call_id: id, name, .. } if id == call_id && name == "ViewImage" + ) + }); + if !from_view_image { + continue; + } + let sent = request.items.iter().position(|item| { + matches!( + item, + Item::ToolResult { call_id: id, attachments: sent, .. } + if id == call_id && sent == attachments + ) + }); + let Some(request_index) = sent else { continue }; + // Header inspection only; no raster allocation or size threshold. + // Unknown dimensions sort by bytes as a deterministic fallback. + let pixels = imagesize::blob_size(image.data()) + .ok() + .map(|size| (size.width as u64).saturating_mul(size.height as u64)) + .unwrap_or(0); + let rank = (pixels, image.data().len(), index); + if largest + .as_ref() + .is_none_or(|(previous, _)| rank > *previous) + { + largest = Some((rank, (index, request_index))); + } + } + largest.map(|(_, candidate)| candidate) +} + +pub(crate) fn rejected_result(item: &Item) -> Item { + let mut item = item.clone(); + if let Item::ToolResult { + summary, + content, + attachments, + disposition, + is_error, + .. + } = &mut item + { + *summary = REJECTION_MESSAGE.to_owned(); + *content = None; + attachments.clear(); + *disposition = ToolResultDisposition::Error; + *is_error = true; + } + item +} diff --git a/crates/agen/src/lib.rs b/crates/agen/src/lib.rs index 869d5c322..fca7069f1 100644 --- a/crates/agen/src/lib.rs +++ b/crates/agen/src/lib.rs @@ -3,6 +3,7 @@ mod engine; mod handler; mod history; +mod image_recovery; mod message; pub(crate) mod callback; diff --git a/crates/agen/src/llm_client/error.rs b/crates/agen/src/llm_client/error.rs index 355940fa1..69334b8a2 100644 --- a/crates/agen/src/llm_client/error.rs +++ b/crates/agen/src/llm_client/error.rs @@ -90,6 +90,35 @@ impl ClientError { } } + /// A provider-declared image size rejection, not a generic invalid input. + /// Some providers (including OpenAI) expose only `invalid_value` here, so + /// narrowly recognize their image-size diagnostic without hardcoded limits. + pub fn is_image_size_rejection(&self) -> bool { + let Self::Api { + status: Some(400 | 413 | 422), + code, + message, + .. + } = self + else { + return false; + }; + if matches!( + code.as_deref(), + Some("image_too_large" | "image_size_exceeded") + ) { + return true; + } + let message = message.to_ascii_lowercase(); + message.contains("image") + && ((message.contains("patches") + && message.contains("exceed") + && message.contains("limit")) + || message.contains("image is too large") + || message.contains("image too large") + || (message.contains("image dimensions") && message.contains("exceed"))) + } + pub fn retry_after(&self) -> Option { match self { ClientError::Api { retry_after, .. } => *retry_after, diff --git a/crates/agen/src/prune.rs b/crates/agen/src/prune.rs index 2ee018904..37af648da 100644 --- a/crates/agen/src/prune.rs +++ b/crates/agen/src/prune.rs @@ -136,14 +136,13 @@ pub fn project(items: &mut [Item], indices: &[usize]) -> usize { /// Remove binary attachments after the first model request that could consume them. /// /// Attachments on tool results following the latest model output are fresh outputs for the -/// immediate follow-up request and remain visible. Once a later assistant/reasoning/tool-call -/// item exists, the model has already had an opportunity to inspect the attachment. Subsequent -/// requests keep the durable tool summary and can call the originating tool again when the image -/// is still needed. Only the request-context clone is modified. +/// immediate follow-up request and remain visible. With early dispatch, later +/// text/calls in the same response have not consumed its tool results: use the +/// confirmed response boundary, not that response's last individual item. +/// Subsequent requests keep the durable tool summary and can call the originating +/// tool again when the image is still needed. Only the request-context clone is modified. pub fn project_consumed_attachments(items: &mut [Item]) -> usize { - let latest_model_output = items.iter().rposition(|item| { - item.is_assistant_message() || item.is_reasoning() || item.is_tool_call() - }); + let latest_model_output = crate::image_recovery::accepted_input_end(items.iter()); let mut count = 0; for (index, item) in items.iter_mut().enumerate() { let is_fresh = latest_model_output.is_some_and(|boundary| index > boundary); diff --git a/crates/agen/tests/image_recovery_test.rs b/crates/agen/tests/image_recovery_test.rs new file mode 100644 index 000000000..2fa11ee06 --- /dev/null +++ b/crates/agen/tests/image_recovery_test.rs @@ -0,0 +1,350 @@ +use agen::llm_client::event::{Event, ResponseStatus, StatusEvent}; +use agen::llm_client::{ClientError, LlmClient, Request, ResponseStream}; +use agen::tool::{Attachment, ImageAttachment}; +use agen::{Engine, EngineRunExit, History, Item, ToolResultDisposition}; +use async_trait::async_trait; +use std::sync::{Arc, Mutex}; + +fn rejection() -> ClientError { + ClientError::Api { status: Some(400), code: Some("invalid_value".into()), message: "The image you provided requires 52150 patches after processing, exceeding the limit of 30000. Please resize the image and try again.".into(), retry_after: None } +} + +fn image_result(id: &str, width: u32, height: u32, padding: usize) -> Item { + // PNG header: dimension inspection never decodes/allocates the pixel plane. + let mut bytes = b"\x89PNG\r\n\x1a\n\0\0\0\rIHDR".to_vec(); + bytes.extend(width.to_be_bytes()); + bytes.extend(height.to_be_bytes()); + bytes.resize(40 + padding, 0); + Item::ToolResult { + id: None, + call_id: id.into(), + summary: "Attached image".into(), + content: None, + disposition: ToolResultDisposition::Success, + is_error: false, + attachments: vec![Attachment::Image(ImageAttachment::new("image/png", bytes))], + } +} +fn history() -> History { + History::from_items(vec![ + Item::user_message("inspect"), + Item::tool_call("small", "ViewImage", "{}"), + Item::tool_call("large", "ViewImage", "{}"), + image_result("small", 10, 10, 2000), + image_result("large", 1600, 33349, 0), + ]) +} +fn images(request: &Request) -> Vec { + request + .items + .iter() + .filter_map(|item| match item { + Item::ToolResult { + call_id, + attachments, + .. + } if !attachments.is_empty() => Some(call_id.clone()), + _ => None, + }) + .collect() +} +#[derive(Clone)] +struct Client { + requests: Arc>>, + failures: usize, + other_error: bool, +} +impl Client { + fn new(failures: usize) -> Self { + Self { + requests: Arc::default(), + failures, + other_error: false, + } + } +} +#[async_trait] +impl LlmClient for Client { + fn clone_boxed(&self) -> Box { + Box::new(self.clone()) + } + async fn stream(&self, request: Request) -> Result { + let mut requests = self.requests.lock().unwrap(); + requests.push(request); + if requests.len() <= self.failures { + return Err(rejection()); + } + if self.other_error { + return Err(ClientError::Api { + status: Some(400), + code: Some("invalid_value".into()), + message: "invalid tool arguments".into(), + retry_after: None, + }); + } + Ok(Box::pin(futures::stream::iter(vec![ + Ok(Event::text_block_start(0)), + Ok(Event::text_delta(0, "resize and retry")), + Ok(Event::text_block_stop(0, None)), + Ok(Event::Status(StatusEvent { + status: ResponseStatus::Completed, + })), + ]))) + } +} + +#[tokio::test] +async fn removes_largest_pixels_not_largest_file_and_keeps_pairing() { + let client = Client::new(1); + let observer = client.clone(); + let mut history = history(); + let mut engine = Engine::new(client); + let committed = Arc::new(Mutex::new(Vec::new())); + let commits = committed.clone(); + engine.set_image_rejection_handler(move |original, replacement| { + commits + .lock() + .unwrap() + .push((original.item.clone(), replacement.clone())); + Ok(()) + }); + let output = engine.run(&mut history, "continue").await; + assert!( + matches!(output.result, EngineRunExit::Finished), + "{:?}", + output.result + ); + let requests = observer.requests.lock().unwrap(); + assert_eq!(requests.len(), 2); + assert_eq!(images(&requests[0]), ["small", "large"]); + assert_eq!(images(&requests[1]), ["small"]); + assert!(requests[1].items.iter().any(|item| matches!(item, Item::ToolResult { call_id, is_error: true, disposition: ToolResultDisposition::Error, summary, .. } if call_id == "large" && summary.contains("retry ViewImage")))); + assert_eq!(committed.lock().unwrap().len(), 1); + assert_eq!( + history + .items() + .filter(|item| matches!(item, Item::ToolResult {call_id, ..} if call_id == "large")) + .count(), + 1 + ); +} + +#[tokio::test] +async fn multiple_rejections_remove_one_image_at_a_time_and_stop_at_exhaustion() { + let client = Client::new(99); + let observer = client.clone(); + let mut history = history(); + let out = Engine::new(client).run(&mut history, "continue").await; + assert!(matches!(out.result, EngineRunExit::Interrupted(_))); + let requests = observer.requests.lock().unwrap(); + assert_eq!(requests.len(), 3); + assert_eq!(images(&requests[1]), ["small"]); + assert!(images(&requests[2]).is_empty()); +} + +#[tokio::test] +async fn unrelated_error_stops_recovery_without_removing_another_image() { + let mut client = Client::new(1); + client.other_error = true; + let observer = client.clone(); + let mut history = history(); + let out = Engine::new(client).run(&mut history, "continue").await; + assert!(matches!(out.result, EngineRunExit::Interrupted(_))); + assert_eq!(observer.requests.lock().unwrap().len(), 2); + assert!(history.items().any(|item| matches!(item, Item::ToolResult {call_id, attachments, ..} if call_id == "small" && !attachments.is_empty()))); +} + +#[tokio::test] +async fn persistence_failure_does_not_change_history_or_retry() { + let client = Client::new(99); + let observer = client.clone(); + let mut history = history(); + let original = history.items_cloned(); + let mut engine = Engine::new(client); + engine.set_image_rejection_handler(|_, _| Err("disk full".into())); + let out = engine.run(&mut history, "continue").await; + assert!(matches!(out.result, EngineRunExit::Interrupted(_))); + assert_eq!( + &history.items_cloned()[..original.len()], + original.as_slice() + ); + assert_eq!(observer.requests.lock().unwrap().len(), 1); +} + +#[tokio::test] +async fn previously_accepted_images_are_not_recovery_candidates() { + let client = Client::new(99); + let observer = client.clone(); + let mut history = history(); + history.push(Item::assistant_message("I saw those images")); + history.push(Item::tool_call("new", "ViewImage", "{}")); + history.push(image_result("new", 1, 1, 0)); + let out = Engine::new(client).run(&mut history, "continue").await; + assert!(matches!(out.result, EngineRunExit::Interrupted(_))); + let requests = observer.requests.lock().unwrap(); + assert_eq!(requests.len(), 2); + // Already-consumed images are projected out, but their durable records + // must not be rewritten as failures. + assert!(images(&requests[1]).is_empty()); + assert_eq!(history.items().filter(|item| matches!(item, Item::ToolResult {call_id, is_error: false, attachments, ..} if (call_id == "small" || call_id == "large") && !attachments.is_empty())).count(), 2); +} + +#[tokio::test] +async fn legacy_empty_failed_request_boundaries_do_not_hide_candidates() { + let client = Client::new(1); + let observer = client.clone(); + let mut history = history(); + history.push(Item::assistant_response_boundary("failed-1")); + history.push(Item::assistant_response_boundary("failed-2")); + let out = Engine::new(client).run(&mut history, "continue").await; + assert!(matches!(out.result, EngineRunExit::Finished)); + assert_eq!(images(&observer.requests.lock().unwrap()[1]), ["small"]); +} + +#[tokio::test] +async fn early_results_remain_candidates_despite_later_text_in_same_response() { + let client = Client::new(1); + let observer = client.clone(); + let mut history = History::from_items(vec![ + Item::user_message("inspect"), + Item::assistant_response_boundary("response-1"), + Item::tool_call("image", "ViewImage", "{}"), + image_result("image", 99, 99, 0), + Item::assistant_message("waiting for image"), + ]); + let out = Engine::new(client).run(&mut history, "continue").await; + assert!( + matches!(out.result, EngineRunExit::Finished), + "{:?}; requests: {:?}", + out.result, + observer.requests.lock().unwrap() + ); + assert!(images(&observer.requests.lock().unwrap()[1]).is_empty()); +} + +#[derive(Clone)] +struct CountingTool { + image: bool, + count: Arc, +} +#[async_trait] +impl agen::tool::Tool for CountingTool { + async fn execute( + &self, + _: &str, + _: agen::tool::ToolExecutionContext, + ) -> Result { + self.count.fetch_add(1, std::sync::atomic::Ordering::SeqCst); + let attachments = if self.image { + let Item::ToolResult { attachments, .. } = image_result("image", 1600, 33349, 0) else { + unreachable!() + }; + attachments + } else { + Vec::new() + }; + Ok(agen::tool::ToolOutput { + summary: "executed once".into(), + content: None, + attachments, + }) + } +} +#[derive(Clone)] +struct ToolClient(Client); +#[async_trait] +impl LlmClient for ToolClient { + fn clone_boxed(&self) -> Box { + Box::new(self.clone()) + } + fn tool_call_completion_support(&self) -> agen::llm_client::ToolCallCompletionSupport { + agen::llm_client::ToolCallCompletionSupport::PerBlock + } + async fn stream(&self, request: Request) -> Result { + if self.0.requests.lock().unwrap().is_empty() { + self.0.requests.lock().unwrap().push(request); + return Ok(Box::pin(futures::stream::iter( + vec![ + Event::tool_use_start(0, "image", "ViewImage"), + Event::tool_input_delta(0, "{}"), + Event::tool_use_stop(0), + Event::tool_use_start(1, "side-effect", "Bash"), + Event::tool_input_delta(1, "{}"), + Event::tool_use_stop(1), + Event::Status(StatusEvent { + status: ResponseStatus::Completed, + }), + ] + .into_iter() + .map(Ok), + ))); + } + self.0.stream(request).await + } +} + +#[tokio::test] +async fn fresh_image_recovery_does_not_repeat_tools_in_either_dispatch_mode() { + for mode in [ + agen::ToolCallDispatchMode::AfterResponse, + agen::ToolCallDispatchMode::OnToolCallComplete, + ] { + let client = Client::new(2); + let mut engine = Engine::new(ToolClient(client.clone())); + engine.set_tool_call_dispatch_mode(mode); + let mut counts = Vec::new(); + for (name, image) in [("ViewImage", true), ("Bash", false)] { + let tool = CountingTool { + image, + count: Arc::default(), + }; + counts.push(tool.count.clone()); + engine.register_tool(Arc::new(move || { + ( + agen::tool::ToolMeta::new(name) + .description("test") + .input_schema(serde_json::json!({"type":"object"})), + Arc::new(tool.clone()), + ) + })); + } + let mut history = History::new(); + let out = engine.run(&mut history, "inspect").await; + assert!( + matches!(out.result, EngineRunExit::Finished), + "{:?}", + out.result + ); + assert!( + counts + .iter() + .all(|count| count.load(std::sync::atomic::Ordering::SeqCst) == 1) + ); + let requests = client.requests.lock().unwrap(); + assert_eq!(requests.len(), 3); + assert_eq!(images(&requests[1]), ["image"]); + assert!(images(&requests[2]).is_empty()); + } +} + +#[test] +fn image_error_classifier_is_not_a_generic_400_retry() { + assert!(rejection().is_image_size_rejection()); + for (status, message) in [ + (400, "invalid image format"), + (400, "image URL not found"), + (400, "context window exceeds token limit"), + (401, "image is too large"), + (500, "image is too large"), + ] { + assert!( + !ClientError::Api { + status: Some(status), + code: Some("invalid_value".into()), + message: message.into(), + retry_after: None + } + .is_image_size_rejection() + ); + } +} diff --git a/crates/session-store/src/public_index.rs b/crates/session-store/src/public_index.rs index d569c09b7..1e1406d16 100644 --- a/crates/session-store/src/public_index.rs +++ b/crates/session-store/src/public_index.rs @@ -1071,6 +1071,7 @@ fn persisted_entry_units(entries: &[LogEntry]) -> usize { | LogEntry::AnnotatedUserInput { history, .. } => history.len(), LogEntry::AnnotatedAssistantItem { .. } | LogEntry::AnnotatedToolResult { .. } + | LogEntry::ToolResultCorrected { .. } | LogEntry::AnnotatedSystemItem { .. } => 1, _ => 0, }; diff --git a/crates/session-store/src/public_snapshot.rs b/crates/session-store/src/public_snapshot.rs index bc7659c12..3f7f8e53c 100644 --- a/crates/session-store/src/public_snapshot.rs +++ b/crates/session-store/src/public_snapshot.rs @@ -7,7 +7,7 @@ use protocol::{ use serde::{Deserialize, Serialize}; use sha2::{Digest, Sha256}; -use std::collections::HashSet; +use std::collections::{HashMap, HashSet}; use std::path::Path; use thiserror::Error; @@ -139,6 +139,17 @@ struct RetainedHistoryCursor { before_offset: u64, record_end_offset: Option, seed_entry_index: Option, + #[serde(default, skip_serializing_if = "Vec::is_empty")] + corrections: Vec, +} + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +#[serde(deny_unknown_fields)] +struct ToolCorrectionPosition { + entry_id: String, + segment_id: SegmentId, + start: u64, + end: u64, } #[derive(Debug)] @@ -445,6 +456,7 @@ pub fn read_retained_session_history_page( limits, &mut scanned_bytes, &mut scanned_entries, + &mut Vec::new(), )?; Some(SessionHistoryLineageBoundary { lineage_id: lineage_identity(active.session_id, &lineage[1..]), @@ -456,6 +468,10 @@ pub fn read_retained_session_history_page( } else { None }; + let mut correction_positions = cursor + .as_ref() + .map(|c| c.corrections.clone()) + .unwrap_or_default(); let mut turns = read_history_turns_backward( &store, active.session_id, @@ -465,11 +481,23 @@ pub fn read_retained_session_history_page( limits, &mut scanned_bytes, &mut scanned_entries, + &mut correction_positions, )?; let has_more = turns.len() > page_limit; if has_more { turns.truncate(page_limit); } + // Carry only corrections whose original results are still on older pages. + let returned_ids: HashSet<_> = turns + .iter() + .flat_map(|turn| { + turn.turn + .entries + .iter() + .map(|entry| entry.entry_id.as_str()) + }) + .collect(); + correction_positions.retain(|position| !returned_ids.contains(position.entry_id.as_str())); let next_cursor = if has_more { let oldest = turns .last() @@ -485,6 +513,7 @@ pub fn read_retained_session_history_page( before_offset: oldest.position.before_offset, record_end_offset: oldest.position.record_end_offset, seed_entry_index: oldest.position.seed_entry_index, + corrections: correction_positions, })?) } else { None @@ -702,6 +731,7 @@ fn persisted_entry_units(entries: &[LogEntry]) -> usize { | LogEntry::AnnotatedUserInput { history, .. } => history.len(), LogEntry::AnnotatedAssistantItem { .. } | LogEntry::AnnotatedToolResult { .. } + | LogEntry::ToolResultCorrected { .. } | LogEntry::AnnotatedSystemItem { .. } => 1, _ => 0, }; @@ -814,10 +844,56 @@ fn read_history_turns_backward( limits: RetainedHistoryReadLimits, scanned_bytes: &mut u64, scanned_entries: &mut usize, + correction_positions: &mut Vec, ) -> Result, RetainedHistoryReadError> { let mut turns = Vec::new(); let mut pending_entries = Vec::new(); let mut seen_entries = HashSet::new(); + let mut corrections = HashMap::::new(); + + // Carry only bounded log positions, never correction text or image bodies, + // across pages. Re-read and validate the referenced records from adopted + // lineage rather than rescanning an arbitrarily large newer suffix. + for position in correction_positions.iter() { + let segment = lineage + .iter() + .find(|s| s.segment_id == position.segment_id) + .ok_or(RetainedHistoryReadError::InvalidCursor)?; + let adopted_end = locate_adopted_segment_end( + store, + session_id, + segment, + limits, + scanned_bytes, + scanned_entries, + )?; + if position.start >= position.end || position.end > adopted_end { + return Err(RetainedHistoryReadError::InvalidCursor); + } + let record = store + .read_log_record_range_read_only_bounded( + session_id, + position.segment_id, + position.start, + position.end, + limits.max_scan_bytes.saturating_sub(*scanned_bytes), + ) + .map_err(map_history_cursor_store_error)?; + *scanned_bytes = scanned_bytes + .checked_add(record.end_offset - record.start_offset) + .ok_or(RetainedHistoryReadError::ResourceLimit)?; + *scanned_entries = scanned_entries + .checked_add(persisted_entry_units(std::slice::from_ref(&record.entry))) + .ok_or(RetainedHistoryReadError::ResourceLimit)?; + if *scanned_entries > limits.max_entries { + return Err(RetainedHistoryReadError::ResourceLimit); + } + if !matches!(&record.entry, LogEntry::ToolResultCorrected { entry, .. } if entry.metadata.entry_id.0 == position.entry_id) + || !capture_tool_correction(&record.entry, &mut corrections) + { + return Err(RetainedHistoryReadError::InvalidCursor); + } + } let (start_segment_index, start_position, cursor_turn_id) = match start { Some((position, turn_id)) => { @@ -878,32 +954,6 @@ fn read_history_turns_backward( for (segment_index, segment) in lineage.iter().enumerate().skip(start_segment_index) { let starting_here = segment_index == start_segment_index; - if starting_here - && start_position - .as_ref() - .and_then(|position| position.seed_entry_index) - .is_some() - { - let seed_index = start_position - .as_ref() - .and_then(|position| position.seed_entry_index) - .expect("checked seed cursor"); - if collect_history_entries_backward( - &segment.seed_entries[..seed_index], - segment.segment_id, - 0, - None, - true, - turn_limit, - &mut seen_entries, - &mut pending_entries, - &mut turns, - ) { - return Ok(turns); - } - continue; - } - let adopted_end = locate_adopted_segment_end( store, session_id, @@ -923,6 +973,32 @@ fn read_history_turns_backward( if before_offset > adopted_end { return Err(RetainedHistoryReadError::InvalidCursor); } + if starting_here + && start_position + .as_ref() + .and_then(|position| position.seed_entry_index) + .is_some() + { + let seed_index = start_position + .as_ref() + .and_then(|position| position.seed_entry_index) + .expect("checked seed cursor"); + if collect_history_entries_backward( + &corrected_history_entries(&segment.seed_entries[..seed_index], &corrections), + segment.segment_id, + 0, + None, + true, + turn_limit, + &mut seen_entries, + &mut pending_entries, + &mut turns, + ) { + return Ok(turns); + } + continue; + } + let mut reader = store .open_retained_segment_reader(session_id, segment.segment_id, before_offset) .map_err(map_history_store_error)?; @@ -949,7 +1025,25 @@ fn read_history_turns_backward( } continue; } - let entries = project_history_record(session_id, segment.segment_id, &record.entry); + if capture_tool_correction(&record.entry, &mut corrections) { + let LogEntry::ToolResultCorrected { entry, .. } = &record.entry else { + unreachable!() + }; + let position = ToolCorrectionPosition { + entry_id: entry.metadata.entry_id.0.clone(), + segment_id: segment.segment_id, + start: record.start_offset, + end: record.end_offset, + }; + if !correction_positions.contains(&position) { + correction_positions.push(position); + } + continue; + } + let entries = corrected_history_entries( + &project_history_record(session_id, segment.segment_id, &record.entry), + &corrections, + ); if collect_history_entries_backward( &entries, segment.segment_id, @@ -967,7 +1061,7 @@ fn read_history_turns_backward( if segment_index + 1 == lineage.len() && collect_history_entries_backward( - &segment.seed_entries, + &corrected_history_entries(&segment.seed_entries, &corrections), segment.segment_id, 0, None, @@ -984,6 +1078,38 @@ fn read_history_turns_backward( Ok(turns) } +fn capture_tool_correction( + record: &LogEntry, + corrections: &mut HashMap, +) -> bool { + let LogEntry::ToolResultCorrected { entry, .. } = record else { + return false; + }; + if let Some(data) = project_item(&entry.metadata.entry_id.0, &entry.item) { + // We read backwards: the first correction seen is the latest one. + corrections + .entry(entry.metadata.entry_id.0.clone()) + .or_insert(data); + } + true +} + +fn corrected_history_entries( + entries: &[SessionSnapshotEntry], + corrections: &HashMap, +) -> Vec { + entries + .iter() + .cloned() + .map(|mut entry| { + if let Some(data) = corrections.get(&entry.entry_id) { + entry.data = data.clone(); + } + entry + }) + .collect() +} + #[allow(clippy::too_many_arguments)] fn collect_history_entries_backward( entries: &[SessionSnapshotEntry], @@ -1052,9 +1178,13 @@ fn is_user_entry(entry: &SessionSnapshotEntry) -> bool { fn encode_history_cursor( cursor: &RetainedHistoryCursor, ) -> Result { - serde_json::to_vec(cursor) + let encoded = serde_json::to_vec(cursor) .map(|bytes| URL_SAFE_NO_PAD.encode(bytes)) - .map_err(|_| RetainedHistoryReadError::CorruptLog) + .map_err(|_| RetainedHistoryReadError::CorruptLog)?; + if encoded.len() > 4096 { + return Err(RetainedHistoryReadError::ResourceLimit); + } + Ok(encoded) } fn decode_history_cursor(cursor: &str) -> Result { @@ -1252,6 +1382,19 @@ pub(crate) fn project_session_snapshot_for_segment( entries.push(history_entry(entry, *ts, data)); } } + LogEntry::ToolResultCorrected { ts, entry } => { + if let Some(data) = project_item(&entry.metadata.entry_id.0, &entry.item) { + if let Some(existing) = entries + .iter_mut() + .find(|existing| existing.entry_id == entry.metadata.entry_id.0) + { + // The original result keeps its place/time in the conversation. + existing.data = data; + } else { + entries.push(history_entry(entry, *ts, data)); + } + } + } LogEntry::AnnotatedSystemItem { ts, entry, .. } => entries.push(system_entry( &entry.item, entry.metadata.entry_id.0.clone(), @@ -1782,6 +1925,129 @@ mod tests { .collect() } + #[test] + fn image_correction_keeps_original_position_across_history_pages() { + let session_id = crate::new_session_id(); + let segment_id = crate::new_segment_id(); + let mut image = assistant_message(1, "image"); + image.item = agen::Item::tool_result_item_with_disposition_and_attachments( + "image-call", + "original image", + None, + agen::ToolResultDisposition::Success, + vec![agen::tool::Attachment::Image( + agen::tool::ImageAttachment::new("image/png", vec![1_u8, 2, 3]), + )], + ) + .into(); + let mut corrected = image.clone(); + corrected.item = agen::Item::tool_result_item_with_disposition_and_attachments( + "image-call", + "Resize or crop and retry ViewImage", + None, + agen::ToolResultDisposition::Error, + Vec::new(), + ) + .into(); + let mut log = vec![segment_start(session_id, Vec::new(), None, None)]; + append_turn(&mut log, 1); + log.push(LogEntry::AnnotatedToolResult { + ts: 15, + entry: image.clone(), + }); + append_turn(&mut log, 2); + append_turn(&mut log, 3); + log.push(LogEntry::ToolResultCorrected { + ts: 35, + entry: corrected, + }); + let (_root, aggregate_root) = persist_history_fixture( + "corrected-image", + session_id, + segment_id, + vec![(segment_id, log.clone())], + ); + let first = read_retained_session_history_page( + &aggregate_root, + "corrected-image", + None, + Some(1), + RetainedHistoryReadLimits::default(), + ) + .unwrap(); + assert_eq!(turn_ids(&first), ["user-3"]); + assert!( + !first.turns[0] + .entries + .iter() + .any(|entry| entry.entry_id == image.metadata.entry_id.0) + ); + let second = read_retained_session_history_page( + &aggregate_root, + "corrected-image", + first.next_cursor.as_deref(), + Some(1), + RetainedHistoryReadLimits::default(), + ) + .unwrap(); + assert_eq!(turn_ids(&second), ["user-2"]); + let third = read_retained_session_history_page( + &aggregate_root, + "corrected-image", + second.next_cursor.as_deref(), + Some(1), + RetainedHistoryReadLimits::default(), + ) + .unwrap(); + assert_eq!(turn_ids(&third), ["user-1"]); + let entry = third.turns[0] + .entries + .iter() + .find(|entry| entry.entry_id == image.metadata.entry_id.0) + .unwrap(); + assert_eq!(entry.timestamp, 15); + assert!( + matches!(&entry.data, SessionSnapshotEntryData::ToolResult { is_error: true, attachments, .. } if attachments.is_empty()) + ); + assert_eq!( + *entry, + *project_session_snapshot(session_id, &log) + .entries + .iter() + .find(|entry| entry.entry_id == image.metadata.entry_id.0) + .unwrap() + ); + let restored = crate::collect_state(&log); + assert_eq!( + restored + .annotated_history + .iter() + .filter(|entry| entry.metadata.entry_id == image.metadata.entry_id) + .count(), + 1 + ); + assert!( + matches!(&restored.annotated_history.iter().find(|entry| entry.metadata.entry_id == image.metadata.entry_id).unwrap().item, LoggedItem::ToolResult { is_error: true, attachments, .. } if attachments.is_empty()) + ); + + // Cursor pointers are only hints: the referenced record is read under + // current lineage authority and must actually be a correction record. + let mut cursor = decode_history_cursor(first.next_cursor.as_deref().unwrap()).unwrap(); + cursor.corrections[0].segment_id = crate::new_segment_id(); + let forged = encode_history_cursor(&cursor).unwrap(); + assert_eq!( + read_retained_session_history_page( + &aggregate_root, + "corrected-image", + Some(&forged), + Some(1), + RetainedHistoryReadLimits::default() + ) + .unwrap_err(), + RetainedHistoryReadError::InvalidCursor + ); + } + #[test] fn retained_history_pages_twelve_turns_as_five_five_two_across_compaction() { let worker_name = "worker-history-12"; diff --git a/crates/session-store/src/segment_log.rs b/crates/session-store/src/segment_log.rs index fb61009b0..0201a0b3a 100644 --- a/crates/session-store/src/segment_log.rs +++ b/crates/session-store/src/segment_log.rs @@ -107,6 +107,9 @@ pub enum LogEntry { /// Canonical tool output and metadata committed as one journal record. AnnotatedToolResult { ts: u64, entry: LoggedHistoryEntry }, + /// Append-only correction of a provider-rejected image result. Retains the + /// original stable identity; replay replaces rather than appends the item. + ToolResultCorrected { ts: u64, entry: LoggedHistoryEntry }, /// Canonical typed system event and model-visible metadata committed /// together. @@ -399,6 +402,16 @@ pub fn collect_state(entries: &[LogEntry]) -> RestoredState { state.annotated_history.push(entry.clone()); state.history.push(Item::from(entry.item.clone())); } + LogEntry::ToolResultCorrected { entry, .. } => { + if let Some(index) = state + .annotated_history + .iter() + .position(|existing| existing.metadata.entry_id == entry.metadata.entry_id) + { + state.annotated_history[index] = entry.clone(); + state.history[index] = Item::from(entry.item.clone()); + } + } LogEntry::AnnotatedSystemItem { entry, extensions, .. } => { diff --git a/crates/worker/src/image_recovery_tests.rs b/crates/worker/src/image_recovery_tests.rs new file mode 100644 index 000000000..2c75374b2 --- /dev/null +++ b/crates/worker/src/image_recovery_tests.rs @@ -0,0 +1,221 @@ +//! Recovery tests use the normal Worker execution and append-log commit gate. +use super::*; +use agen::llm_client::event::{Event as LlmEvent, ResponseStatus, StatusEvent}; +use agen::llm_client::{ClientError, Request, ResponseStream}; +use agen::tool::{Attachment, ImageAttachment}; +use async_trait::async_trait; + +#[derive(Clone, Default)] +struct RejectImages { + requests: Arc>>, + cancel_on_retry: Arc>>>, +} +#[async_trait] +impl LlmClient for RejectImages { + fn clone_boxed(&self) -> Box { + Box::new(self.clone()) + } + async fn stream(&self, request: Request) -> Result { + let has_image = request.items.iter().any( + |item| matches!(item, Item::ToolResult { attachments, .. } if !attachments.is_empty()), + ); + self.requests.lock().unwrap().push(request); + if has_image { + return Err(ClientError::Api { + status: Some(400), + code: Some("invalid_value".into()), + message: "The image requires 52150 patches, exceeding the limit of 30000.".into(), + retry_after: None, + }); + } + if let Some(cancel) = self.cancel_on_retry.lock().unwrap().take() { + cancel.try_send(()).unwrap(); + return Ok(Box::pin(futures::stream::pending())); + } + Ok(Box::pin(futures::stream::iter(vec![ + Ok(LlmEvent::text_block_start(0)), + Ok(LlmEvent::text_delta(0, "Resize the source image")), + Ok(LlmEvent::text_block_stop(0, None)), + Ok(LlmEvent::Status(StatusEvent { + status: ResponseStatus::Completed, + })), + ]))) + } +} + +async fn image_worker() -> ( + tempfile::TempDir, + Worker, + RejectImages, +) { + let dir = tempfile::tempdir().unwrap(); + let store = session_store::FsStore::new(dir.path().join("session")).unwrap(); + let manifest = WorkerManifest::from_toml( + r#" +[worker] +name = "image-recovery-test" +[model] +scheme = "anthropic" +model_id = "test" +[engine] +[scope] +allow = [] +"#, + ) + .unwrap(); + let client = RejectImages::default(); + let mut worker = Worker::new( + manifest, + Engine::<_, Mutable, SessionHistoryMetadata>::new_annotated(client.clone()), + store, + WorkerWorkspaceContext::local_filesystem(None), + WorkerFilesystemAuthority::None, + Scope::writable(dir.path()).unwrap(), + ) + .await + .unwrap(); + worker.disable_manifest_lifecycle_features(); + worker.set_history_for_test(vec![ + Item::user_message("inspect this image"), + Item::tool_call("image-call", "ViewImage", "{}"), + Item::tool_result_item_with_disposition_and_attachments( + "image-call", + "Attached image", + None, + agen::ToolResultDisposition::Success, + vec![Attachment::Image(ImageAttachment::new( + "image/png", + vec![1_u8, 2, 3], + ))], + ), + ]); + worker.wire_history_persistence(); + (dir, worker, client) +} +fn rejected(item: &Item) -> bool { + matches!(item, Item::ToolResult { call_id, is_error: true, attachments, .. } if call_id == "image-call" && attachments.is_empty()) +} + +#[tokio::test] +async fn image_correction_is_committed_before_retry_and_survives_restore_and_reseed() { + let (_dir, mut worker, client) = image_worker().await; + assert_eq!( + worker.run_text("continue").await.unwrap(), + WorkerRunResult::Finished + ); + assert_eq!(client.requests.lock().unwrap().len(), 2); + let raw = worker + .store + .read_all(worker.session_id(), worker.segment_id()) + .unwrap(); + let correction = raw + .iter() + .find(|entry| matches!(entry, LogEntry::ToolResultCorrected { .. })) + .unwrap() + .clone(); + let restored = restore_history_entries(worker.session_id(), worker.segment_id(), &raw).unwrap(); + assert_eq!( + restored + .iter() + .filter(|entry| rejected(&entry.item)) + .count(), + 1 + ); + assert_eq!( + session_store::collect_state(&raw) + .history + .iter() + .filter(|item| rejected(item)) + .count(), + 1 + ); + let LogEntry::ToolResultCorrected { + entry: corrected, .. + } = &correction + else { + unreachable!() + }; + let original = match &raw[0] { + LogEntry::AnnotatedSegmentStart { history, .. } => history + .iter() + .find(|entry| entry.metadata.entry_id == corrected.metadata.entry_id) + .unwrap(), + _ => panic!("missing seed"), + }; + assert!( + matches!(&original.item, session_store::LoggedItem::ToolResult { attachments, is_error: false, .. } if attachments.len() == 1) + ); + assert_eq!( + restored + .iter() + .find(|entry| rejected(&entry.item)) + .unwrap() + .annotation, + original.metadata + ); + let snapshot = session_store::public_snapshot::project_current_session_snapshot(&raw); + let projected = snapshot + .entries + .iter() + .find(|entry| entry.entry_id == corrected.metadata.entry_id.0) + .unwrap(); + assert!( + matches!(&projected.data, protocol::SessionSnapshotEntryData::ToolResult { is_error: true, attachments, .. } if attachments.is_empty()) + ); + assert!(matches!( + crate::ipc::protocol_session::live_log_entry_event(correction), + Some(Event::SessionEntryCommitted { .. }) + )); + + // Fork/compact seed writers serialize the effective annotated history. + let seed = LogEntry::AnnotatedSegmentStart { + ts: 1, + session_id: worker.session_id(), + system_prompt: None, + config: agen::llm_client::RequestConfig::default(), + history: restored.iter().map(to_logged_history_entry).collect(), + forked_from: None, + compacted_from: None, + }; + let restored_again = + restore_history_entries(worker.session_id(), worker.segment_id(), &[seed]).unwrap(); + assert_eq!( + restored_again + .iter() + .filter(|entry| rejected(&entry.item)) + .count(), + 1 + ); + let mut history = History::from_entries(restored_again); + let out = Engine::<_, Mutable, SessionHistoryMetadata>::new_annotated(client.clone()) + .run_with_annotation(&mut history, "continue after restart", &mut |_| { + Ok(SessionHistoryMetadata::legacy_unknown()) + }) + .await; + assert!(matches!(out.result, EngineRunExit::Finished)); + assert_eq!(client.requests.lock().unwrap().len(), 3); +} + +#[tokio::test] +async fn cancellation_after_image_correction_does_not_rollback_the_correction() { + let (_dir, mut worker, client) = image_worker().await; + *client.cancel_on_retry.lock().unwrap() = Some(worker.engine().cancel_sender()); + let result = worker.run_text("continue").await.unwrap(); + assert_eq!(result, WorkerRunResult::Cancelled); + let raw = worker + .store + .read_all(worker.session_id(), worker.segment_id()) + .unwrap(); + assert!( + raw.iter() + .any(|entry| matches!(entry, LogEntry::ToolResultCorrected { .. })) + ); + let restored = restore_history_entries(worker.session_id(), worker.segment_id(), &raw).unwrap(); + assert_eq!( + restored + .iter() + .filter(|entry| rejected(&entry.item)) + .count(), + 1 + ); +} diff --git a/crates/worker/src/ipc/protocol_session.rs b/crates/worker/src/ipc/protocol_session.rs index 2e84facd2..a9dd18dd7 100644 --- a/crates/worker/src/ipc/protocol_session.rs +++ b/crates/worker/src/ipc/protocol_session.rs @@ -76,6 +76,7 @@ pub fn live_log_entry_event(entry: LogEntry) -> Option { } entry @ (LogEntry::AnnotatedAssistantItem { .. } | LogEntry::AnnotatedToolResult { .. } + | LogEntry::ToolResultCorrected { .. } | LogEntry::RunYielded { .. } | LogEntry::RunResumed { .. } | LogEntry::RunCancelled { .. } diff --git a/crates/worker/src/segment_log_sink.rs b/crates/worker/src/segment_log_sink.rs index 55da19660..93d0b0c73 100644 --- a/crates/worker/src/segment_log_sink.rs +++ b/crates/worker/src/segment_log_sink.rs @@ -121,6 +121,7 @@ impl SegmentLogSink { | LogEntry::AnnotatedUserInput { .. } | LogEntry::AnnotatedAssistantItem { .. } | LogEntry::AnnotatedToolResult { .. } + | LogEntry::ToolResultCorrected { .. } | LogEntry::AnnotatedSystemItem { .. } | LogEntry::Invoke { .. } | LogEntry::RunYielded { .. } diff --git a/crates/worker/src/session_history.rs b/crates/worker/src/session_history.rs index 420429185..5a5f60d8d 100644 --- a/crates/worker/src/session_history.rs +++ b/crates/worker/src/session_history.rs @@ -116,6 +116,17 @@ pub(crate) fn restore_history_entries( } LogEntry::AnnotatedAssistantItem { entry, .. } | LogEntry::AnnotatedToolResult { entry, .. } => history.push(from_logged(entry)), + LogEntry::ToolResultCorrected { entry, .. } => { + let target = history + .iter_mut() + .find(|existing| existing.annotation.entry_id == entry.metadata.entry_id) + .ok_or_else(|| "tool result correction target is missing".to_string())?; + if !matches!((&target.item, &entry.item), (Item::ToolResult { call_id, .. }, session_store::LoggedItem::ToolResult { call_id: replacement, .. }) if call_id == replacement) + { + return Err("tool result correction changed call identity".to_string()); + } + *target = from_logged(entry); + } LogEntry::AnnotatedSystemItem { entry, .. } => history.push(HistoryEntry::new( entry.item.to_history_item(), entry.metadata.clone(), diff --git a/crates/worker/src/worker.rs b/crates/worker/src/worker.rs index 45046af93..348fd6705 100644 --- a/crates/worker/src/worker.rs +++ b/crates/worker/src/worker.rs @@ -1,3 +1,7 @@ +#[cfg(test)] +#[path = "image_recovery_tests.rs"] +mod image_recovery_tests; + use std::collections::{HashMap, VecDeque}; #[cfg(test)] use std::path::Path; @@ -4046,6 +4050,24 @@ impl Worker { St: Clone + 'static, { self.materialize_durable_session_head().await?; + let writer = self.log_writer_handle(); + let activity = self.ai_activity_counter.clone(); + self.engine_mut() + .set_image_rejection_handler(move |original, replacement| { + writer + .append_entry(LogEntry::ToolResultCorrected { + ts: segment_log::now_millis(), + entry: session_store::LoggedHistoryEntry { + item: replacement.clone().into(), + metadata: original.annotation.clone(), + }, + }) + .map_err(|error| error.to_string())?; + // A later Cancel must not roll back this committed correction while + // leaving the in-memory result changed (the original may predate Run). + activity.fetch_add(1, Ordering::SeqCst); + Ok(()) + }); self.ensure_interceptor_installed(); if self.should_pre_run_compact() { self.try_pre_run_compact().await?; diff --git a/web/workspace/src/lib/workspace/console/model.test.ts b/web/workspace/src/lib/workspace/console/model.test.ts index 0f6e54dc4..9ba89bec5 100644 --- a/web/workspace/src/lib/workspace/console/model.test.ts +++ b/web/workspace/src/lib/workspace/console/model.test.ts @@ -3769,6 +3769,29 @@ Deno.test("committed tool identities reconcile one live call block", () => { ); }); +Deno.test("image rejection correction updates the existing tool result without duplicate rows", () => { + const result: SessionSnapshotEntry = { + kind: "tool_result", entry_id: "result-image", timestamp: 2, + provenance: "tool_output", call_id: "image", summary: "Attached image", is_error: false, + attachments: [{ attachment_id: "image-body", media_type: "image/png", byte_len: 100 }], + }; + const lines = projectConsole([ + { eventId: "call", event: { event: "session_entry_committed", data: { entry: { + kind: "tool_call", entry_id: "call-image", timestamp: 1, provenance: "model_output", + call_id: "image", name: "ViewImage", arguments: "{}", + } } } }, + { eventId: "original", event: { event: "session_entry_committed", data: { entry: result } } }, + { eventId: "correction", event: { event: "session_entry_committed", data: { entry: { + ...result, timestamp: 3, summary: "Resize or crop and retry ViewImage", is_error: true, attachments: [], + } } } }, + ]).lines; + assertEquals(lines.length, 1); + assertEquals(lines[0].entryId, "result-image"); + assertEquals(lines[0].toolCall?.state, "error"); + assertEquals(lines[0].attachments ?? [], []); + assert(lines[0].body.includes("Resize or crop"), "show the model recovery instruction"); +}); + function readHistory(ids: string[], summary: string): ConsoleLine[] { const entries = ids.flatMap((id): SessionSnapshotEntry[] => [{ kind: "tool_call", entry_id: `call-entry-${id}`, timestamp: 1,