diff --git a/crates/agen/src/engine.rs b/crates/agen/src/engine.rs index 8638332d4..16f76d790 100644 --- a/crates/agen/src/engine.rs +++ b/crates/agen/src/engine.rs @@ -3199,6 +3199,55 @@ impl Engine { } } + /// Correct one candidate actually sent in the rejected request. Commit the + /// durable correction before changing either in-memory projection or retrying. + fn recover_rejected_image( + &self, + history: &mut History, + request: &mut Request, + ) -> Result { + let Some((history_index, request_index)) = + crate::image_recovery::largest_candidate(history, request) + else { + return Ok(false); + }; + 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(); + // Retain the original request's prune/interceptor decisions. + request.items[request_index] = replacement; + self.emit_warning(crate::image_recovery::REJECTION_MESSAGE); + Ok(true) + } + + /// Opening a corrected request can itself yield an HTTP image rejection. + /// Each successful correction removes one candidate, bounding this loop. + async fn open_stream_with_image_recovery( + &mut self, + request: &mut Request, + turn: usize, + llm_call: usize, + history: &mut History, + ) -> Result { + loop { + match self + .open_stream_with_retry(request.clone(), turn, llm_call) + .await + { + Ok(stream) => return Ok(stream), + Err(EngineError::Client(error)) if error.is_image_size_rejection() => { + if !self.recover_rejected_image(history, request)? { + return Err(EngineError::Client(error)); + } + } + Err(error) => return Err(error), + } + } + } + /// Open a stream, dispatch all events to the timeline, handle cancellation. async fn stream_response( &mut self, @@ -3217,37 +3266,14 @@ impl Engine { ); 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), - } - }; + let stream = self + .open_stream_with_image_recovery(&mut request, turn, llm_call, history) + .await?; self.tool_call_collector.begin_response(); let mut stream = ResponseStreamPump::start(stream); let mut early_tools: Option = None; let mut response_completed = false; + let mut response_progress = false; let mut event_count: usize = 0; loop { @@ -3285,6 +3311,35 @@ impl Engine { json!({}), ); } + self.emit_stream_event(turn, llm_call, &event); + // Only pre-output provider rejection can safely replay the + // original request. Even an empty block start closes this + // window, before collectors or early tool admission run. + // Unknown events are conservatively not replayable either. + match &event { + Event::Ping(_) | Event::Usage(_) | Event::Error(_) => {} + Event::Status(StatusEvent { + status: crate::llm_client::event::ResponseStatus::Started, + }) => {} + _ => response_progress = true, + } + if let Event::Error(error) = &event { + if !response_progress && error.is_image_size_rejection() { + self.timeline.flush_usage(); + if self.recover_rejected_image(history, &mut request)? { + // Drop the old pump (including queued failed status) + // before opening the corrected request. No collector or + // tool state exists for this pre-output attempt. + drop(stream); + stream = ResponseStreamPump::start(self.open_stream_with_image_recovery( + &mut request, turn, llm_call, history, + ).await?); + self.tool_call_collector.begin_response(); + event_count = 0; + continue; + } + } + } if matches!( &event, Event::Status(StatusEvent { @@ -3293,7 +3348,6 @@ impl Engine { ) { response_completed = true; } - self.emit_stream_event(turn, llm_call, &event); self.timeline.dispatch(&event); if early_dispatch { let collected = self.tool_call_collector.take_collected_with_index(); diff --git a/crates/agen/src/llm_client/error.rs b/crates/agen/src/llm_client/error.rs index 69334b8a2..2566e6c26 100644 --- a/crates/agen/src/llm_client/error.rs +++ b/crates/agen/src/llm_client/error.rs @@ -103,20 +103,7 @@ impl ClientError { 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"))) + is_image_size_diagnostic(code.as_deref(), message) } pub fn retry_after(&self) -> Option { @@ -127,6 +114,29 @@ impl ClientError { } } +/// Shared diagnostic policy for HTTP rejections and provider Error events. +/// The caller must establish the rejection path and whether replay is safe; +/// this predicate alone never authorizes a retry. +pub(crate) fn is_image_size_diagnostic(code: Option<&str>, message: &str) -> bool { + if matches!(code, Some("image_too_large" | "image_size_exceeded")) { + return true; + } + // OpenAI Responses appends structured diagnostic details for observation. + // Extra fields (including echoed input) are not the error message itself. + let message = message + .split(" | diagnostic=") + .next() + .unwrap_or(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"))) +} + /// transient な失敗としてリトライ対象になるかを判定する。 /// /// 対象: diff --git a/crates/agen/src/llm_client/event.rs b/crates/agen/src/llm_client/event.rs index d31c035d1..9f1e65bfe 100644 --- a/crates/agen/src/llm_client/event.rs +++ b/crates/agen/src/llm_client/event.rs @@ -109,6 +109,15 @@ pub struct ErrorEvent { pub message: String, } +impl ErrorEvent { + /// A provider Error event has no HTTP status. Classify its diagnostic + /// explicitly, rather than making every status-less ClientError eligible. + /// The engine must separately establish that no response output was seen. + pub(crate) fn is_image_size_rejection(&self) -> bool { + super::error::is_image_size_diagnostic(self.code.as_deref(), &self.message) + } +} + /// 未対応 SSE イベントの観測用メタイベント。 /// /// `data_preview` は provider から受け取った raw SSE data の bounded preview、 diff --git a/crates/agen/src/llm_client/scheme/openai_responses/events.rs b/crates/agen/src/llm_client/scheme/openai_responses/events.rs index cb6ead57e..8b9fbb3ed 100644 --- a/crates/agen/src/llm_client/scheme/openai_responses/events.rs +++ b/crates/agen/src/llm_client/scheme/openai_responses/events.rs @@ -753,17 +753,31 @@ pub(crate) fn parse_sse( } "error" => { - let ev = from_json::(data).unwrap_or_else(|_| { - TopLevelErrorEnvelope { + // Responses documents flat error frames (code/message/param), + // while providers also send nested error envelopes. Normalize both + // structurally so classification never scans arbitrary raw JSON. + let ev = from_json::(data) + .or_else(|_| { + from_json::(data).map(|error| TopLevelErrorEnvelope { + error, + extra: BTreeMap::new(), + }) + }) + .unwrap_or_else(|_| TopLevelErrorEnvelope { error: TopLevelError { - message: Some(data.to_string()), + // Raw data is observation, not a provider diagnostic that + // can authorize request recovery. Keep it out of the base + // message even when it contains image-size wording. + message: Some("OpenAI Responses error with unparseable data".to_string()), error_type: None, code: None, - extra: BTreeMap::new(), + extra: BTreeMap::from([( + "raw_data".to_string(), + Value::String(data.to_string()), + )]), }, extra: BTreeMap::new(), - } - }); + }); let (code, message) = top_level_error_diagnostic(ev); Ok(vec![Event::Error(ErrorEvent { code, message })]) } diff --git a/crates/agen/tests/image_recovery_test.rs b/crates/agen/tests/image_recovery_test.rs index 2fa11ee06..276450d8e 100644 --- a/crates/agen/tests/image_recovery_test.rs +++ b/crates/agen/tests/image_recovery_test.rs @@ -1,8 +1,11 @@ use agen::llm_client::event::{Event, ResponseStatus, StatusEvent}; +use agen::llm_client::scheme::{Scheme, openai_responses::OpenAIResponsesScheme}; 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 futures::StreamExt; +use serde_json::{Value, json}; use std::sync::{Arc, Mutex}; fn rejection() -> ClientError { @@ -48,11 +51,28 @@ fn images(request: &Request) -> Vec { }) .collect() } +#[derive(Clone, Copy)] +enum RejectionPath { + Http, + SseError, + SseNestedError, + ResponseFailed, +} +const STREAM_PATHS: [RejectionPath; 3] = [ + RejectionPath::SseError, + RejectionPath::SseNestedError, + RejectionPath::ResponseFailed, +]; + #[derive(Clone)] struct Client { requests: Arc>>, failures: usize, other_error: bool, + path: RejectionPath, + prefix: Vec<(&'static str, Value)>, + wait_before_error: Option>, + rejection_data: Option<&'static str>, } impl Client { fn new(failures: usize) -> Self { @@ -60,11 +80,69 @@ impl Client { requests: Arc::default(), failures, other_error: false, + path: RejectionPath::Http, + prefix: Vec::new(), + wait_before_error: None, + rejection_data: None, } } + fn streamed(failures: usize, path: RejectionPath) -> Self { + Self { + path, + ..Self::new(failures) + } + } + fn reject(&self, error: ClientError) -> Result { + if matches!(self.path, RejectionPath::Http) { + return Err(error); + } + let ClientError::Api { code, message, .. } = error else { + unreachable!() + }; + let error = json!({"type":"invalid_request_error", "code":code, "message":message, "param":"input"}); + let (kind, data) = match self.path { + RejectionPath::SseError => ( + "error", + json!({"type":"error", "code":code, "message":message, "param":"input", "error_type":"invalid_request_error"}), + ), + RejectionPath::SseNestedError => ("error", json!({"type":"error", "error":error})), + RejectionPath::ResponseFailed => ( + "response.failed", + json!({"type":"response.failed", "response":{"id":"rejected", "error":error}}), + ), + RejectionPath::Http => unreachable!(), + }; + let scheme = OpenAIResponsesScheme::new(); + let mut state = Default::default(); + let mut events = scheme.parse_sse("response.created", r#"{"response":{}}"#, &mut state)?; + for (kind, data) in &self.prefix { + events.extend(scheme.parse_sse(kind, &data.to_string(), &mut state)?); + } + let encoded = data.to_string(); + events.extend(scheme.parse_sse( + kind, + self.rejection_data.unwrap_or(&encoded), + &mut state, + )?); + let wait = self.wait_before_error.clone(); + Ok(Box::pin(futures::stream::iter(events).then(move |event| { + let wait = wait.clone(); + async move { + if matches!(event, Event::Error(_)) + && let Some(wait) = wait + { + wait.notified().await; + } + Ok(event) + } + }))) + } } #[async_trait] impl LlmClient for Client { + fn tool_call_completion_support(&self) -> agen::llm_client::ToolCallCompletionSupport { + agen::llm_client::ToolCallCompletionSupport::PerBlock + } fn clone_boxed(&self) -> Box { Box::new(self.clone()) } @@ -72,10 +150,10 @@ impl LlmClient for Client { let mut requests = self.requests.lock().unwrap(); requests.push(request); if requests.len() <= self.failures { - return Err(rejection()); + return self.reject(rejection()); } if self.other_error { - return Err(ClientError::Api { + return self.reject(ClientError::Api { status: Some(400), code: Some("invalid_value".into()), message: "invalid tool arguments".into(), @@ -226,6 +304,7 @@ async fn early_results_remain_candidates_despite_later_text_in_same_response() { struct CountingTool { image: bool, count: Arc, + executed: Option>, } #[async_trait] impl agen::tool::Tool for CountingTool { @@ -235,6 +314,9 @@ impl agen::tool::Tool for CountingTool { _: agen::tool::ToolExecutionContext, ) -> Result { self.count.fetch_add(1, std::sync::atomic::Ordering::SeqCst); + if let Some(executed) = &self.executed { + executed.notify_one(); + } let attachments = if self.image { let Item::ToolResult { attachments, .. } = image_result("image", 1600, 33349, 0) else { unreachable!() @@ -297,6 +379,7 @@ async fn fresh_image_recovery_does_not_repeat_tools_in_either_dispatch_mode() { let tool = CountingTool { image, count: Arc::default(), + executed: None, }; counts.push(tool.count.clone()); engine.register_tool(Arc::new(move || { @@ -348,3 +431,308 @@ fn image_error_classifier_is_not_a_generic_400_retry() { ); } } + +#[tokio::test] +async fn streamed_image_rejections_correct_largest_candidate_before_next_request() { + for path in STREAM_PATHS { + let client = Client::streamed(1, path); + let mut history = history(); + let mut engine = Engine::new(client.clone()); + 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 out = engine.run(&mut history, "continue").await; + assert!( + matches!(out.result, EngineRunExit::Finished), + "{:?}", + out.result + ); + let requests = client.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("Resize or crop") && summary.contains("retry ViewImage")))); + assert!(requests[1].items.iter().any(|item| matches!(item, Item::ToolCall { call_id, name, .. } if call_id == "large" && name == "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 streamed_rejections_stop_when_images_are_absent_or_candidates_exhausted() { + for path in STREAM_PATHS { + for (mut history, expected_requests) in [(History::new(), 1), (history(), 3)] { + let client = Client::streamed(99, path); + let out = Engine::new(client.clone()) + .run(&mut history, "continue") + .await; + assert!( + matches!(out.result, EngineRunExit::Interrupted(_)), + "{:?}", + out.result + ); + let requests = client.requests.lock().unwrap(); + assert_eq!(requests.len(), expected_requests); + assert!(images(requests.last().unwrap()).is_empty()); + } + } +} + +#[tokio::test] +async fn unrelated_streamed_invalid_value_preserves_images_and_diagnostic() { + for path in STREAM_PATHS { + let mut client = Client::streamed(0, path); + client.other_error = true; + let mut history = history(); + let original = history.items_cloned(); + let out = Engine::new(client.clone()) + .run(&mut history, "continue") + .await; + let EngineRunExit::Interrupted(reason) = out.result else { + panic!("{:?}", out.result) + }; + let reason = format!("{reason:?}"); + assert!( + reason.contains("invalid_value") + && reason.contains("invalid tool arguments") + && reason.contains("diagnostic=") + && reason.contains("input"), + "{reason}" + ); + assert_eq!(client.requests.lock().unwrap().len(), 1); + assert_eq!( + &history.items_cloned()[..original.len()], + original.as_slice() + ); + } +} + +#[tokio::test] +async fn streamed_correction_commit_failure_does_not_mutate_or_resend() { + for path in STREAM_PATHS { + let client = Client::streamed(99, path); + let mut history = history(); + let original = history.items_cloned(); + let mut engine = Engine::new(client.clone()); + 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!(client.requests.lock().unwrap().len(), 1); + } +} + +#[tokio::test] +async fn streamed_rejection_after_output_never_replays_or_corrects_input() { + for path in STREAM_PATHS { + // Start alone closes the replay window, as do partial/completed text, + // reasoning and unrecognized events. None may change accepted input. + for prefix in [ + vec![( + "response.content_part.added", + json!({"output_index":0,"content_index":0,"part":{"type":"output_text"}}), + )], + vec![( + "response.output_text.delta", + json!({"output_index":0,"content_index":0,"delta":"partial"}), + )], + vec![ + ( + "response.output_text.delta", + json!({"output_index":0,"content_index":0,"delta":"completed"}), + ), + ( + "response.content_part.done", + json!({"output_index":0,"content_index":0,"part":{"type":"output_text"}}), + ), + ], + vec![( + "response.reasoning_text.delta", + json!({"output_index":0,"content_index":0,"delta":"thinking"}), + )], + vec![("future.unknown", json!({"value":"unknown semantics"}))], + ] { + let mut client = Client::streamed(1, path); + client.prefix = prefix; + let mut history = history(); + let original = history.items_cloned(); + let out = Engine::new(client.clone()) + .run(&mut history, "continue") + .await; + assert!( + matches!(out.result, EngineRunExit::Interrupted(_)), + "{:?}", + out.result + ); + assert_eq!(client.requests.lock().unwrap().len(), 1); + assert_eq!( + &history.items_cloned()[..original.len()], + original.as_slice() + ); + } + } +} + +#[test] +fn statusless_api_errors_are_not_image_retry_authority() { + let ClientError::Api { code, message, .. } = rejection() else { + unreachable!() + }; + assert!( + !ClientError::Api { + status: None, + code, + message, + retry_after: None + } + .is_image_size_rejection() + ); +} + +#[tokio::test] +async fn streamed_rejection_after_tool_completion_never_reexecutes_side_effects() { + for path in STREAM_PATHS { + for mode in [ + agen::ToolCallDispatchMode::AfterResponse, + agen::ToolCallDispatchMode::OnToolCallComplete, + ] { + let executed = Arc::new(tokio::sync::Notify::new()); + let mut client = Client::streamed(1, path); + client.prefix = vec![ + ( + "response.output_item.added", + json!({"output_index":0,"item":{"type":"function_call","call_id":"side-effect","name":"Bash","arguments":"{}"}}), + ), + ( + "response.output_item.done", + json!({"output_index":0,"item":{"type":"function_call","call_id":"side-effect","name":"Bash","arguments":"{}"}}), + ), + ]; + if mode == agen::ToolCallDispatchMode::OnToolCallComplete { + client.wait_before_error = Some(executed.clone()); + } + let tool = CountingTool { + image: false, + count: Arc::default(), + executed: Some(executed), + }; + let count = tool.count.clone(); + let mut engine = Engine::new(client.clone()); + engine.set_tool_call_dispatch_mode(mode); + engine.register_tool(Arc::new(move || { + ( + agen::tool::ToolMeta::new("Bash") + .description("side effect") + .input_schema(json!({"type":"object"})), + Arc::new(tool.clone()), + ) + })); + let mut history = history(); + let original = history.items_cloned(); + let out = tokio::time::timeout( + std::time::Duration::from_secs(5), + engine.run(&mut history, "continue"), + ) + .await + .expect("tool admission/error cleanup must terminate"); + assert!( + matches!(out.result, EngineRunExit::Interrupted(_)), + "{:?}", + out.result + ); + assert_eq!(client.requests.lock().unwrap().len(), 1); + assert_eq!( + &history.items_cloned()[..original.len()], + original.as_slice() + ); + if mode == agen::ToolCallDispatchMode::OnToolCallComplete { + assert_eq!(count.load(std::sync::atomic::Ordering::SeqCst), 1); + assert_eq!(history.items().filter(|item| matches!(item, Item::ToolCall { call_id, .. } if call_id == "side-effect")).count(), 1); + assert_eq!(history.items().filter(|item| matches!(item, Item::ToolResult { call_id, .. } if call_id == "side-effect")).count(), 1); + } else { + assert_eq!(count.load(std::sync::atomic::Ordering::SeqCst), 0); + assert!(!history.items().any(|item| matches!(item, Item::ToolCall { call_id, .. } | Item::ToolResult { call_id, .. } if call_id == "side-effect"))); + } + } + } +} + +#[test] +fn streamed_error_extra_fields_do_not_turn_unrelated_message_into_image_rejection() { + let scheme = OpenAIResponsesScheme::new(); + let mut state = Default::default(); + let events = scheme.parse_sse("error", r#"{"type":"error","code":"invalid_value","message":"invalid tool arguments","param":"input","input":"image too large"}"#, &mut state).unwrap(); + let Event::Error(error) = &events[0] else { + panic!("{events:?}") + }; + assert_eq!(error.code.as_deref(), Some("invalid_value")); + assert!( + error + .message + .starts_with("invalid tool arguments | diagnostic=") + ); + // HTTP uses the same diagnostic policy, but still requires status authority. + assert!( + !ClientError::Api { + status: Some(400), + code: error.code.clone(), + message: error.message.clone(), + retry_after: None + } + .is_image_size_rejection() + ); +} + +#[tokio::test] +async fn unparseable_stream_error_data_preserves_images_and_terminates_with_diagnostic() { + for data in [ + r#"{"type":"error","code":"invalid_value","message":123,"input":"image too large"}"#, + "unparseable provider error with echoed input: image too large", + ] { + let mut client = Client::streamed(1, RejectionPath::SseError); + client.rejection_data = Some(data); + let mut history = history(); + let original = history.items_cloned(); + let out = Engine::new(client.clone()) + .run(&mut history, "continue") + .await; + let EngineRunExit::Interrupted(agen::RunInterruptionReason::Unexpected( + agen::EngineError::Client(ClientError::Api { code, message, .. }), + )) = out.result + else { + panic!( + "malformed provider error must terminate, got {:?}", + out.result + ); + }; + assert_eq!(code, None); + let (base, diagnostic) = message + .split_once(" | diagnostic=") + .expect("raw data remains diagnostic-only"); + assert!(!base.contains("image")); + let diagnostic: Value = serde_json::from_str(diagnostic).unwrap(); + assert_eq!(diagnostic["error_extra"]["raw_data"], data); + assert_eq!(client.requests.lock().unwrap().len(), 1); + assert_eq!( + &history.items_cloned()[..original.len()], + original.as_slice() + ); + } +} diff --git a/crates/worker/src/image_recovery_tests.rs b/crates/worker/src/image_recovery_tests.rs index 2c75374b2..179631bce 100644 --- a/crates/worker/src/image_recovery_tests.rs +++ b/crates/worker/src/image_recovery_tests.rs @@ -1,13 +1,37 @@ //! 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::scheme::{Scheme, openai_responses::OpenAIResponsesScheme}; use agen::llm_client::{ClientError, Request, ResponseStream}; use agen::tool::{Attachment, ImageAttachment}; use async_trait::async_trait; +#[derive(Clone, Copy, Debug, Default)] +enum RejectionRoute { + #[default] + Http, + SseError, + ResponseFailed, +} + +const REJECTION_ROUTES: [RejectionRoute; 3] = [ + RejectionRoute::Http, + RejectionRoute::SseError, + RejectionRoute::ResponseFailed, +]; + +struct RetryObservation { + store: session_store::FsStore, + session_id: SessionId, + segment_id: SegmentId, + committed: Option>, +} + #[derive(Clone, Default)] struct RejectImages { + route: RejectionRoute, requests: Arc>>, + retry_observation: Arc>>, cancel_on_retry: Arc>>>, } #[async_trait] @@ -21,12 +45,49 @@ impl LlmClient for RejectImages { ); 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 matches!(self.route, RejectionRoute::Http) { + 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, + }); + } + // Exercise the public Responses parser, not a hand-built ErrorEvent. + // Started is metadata; no generated block precedes this rejection. + let scheme = OpenAIResponsesScheme::new(); + let mut state = Default::default(); + let mut events = scheme.parse_sse( + "response.created", + r#"{"response":{"id":"image-response"}}"#, + &mut state, + )?; + let (event_type, data) = match self.route { + RejectionRoute::Http => unreachable!(), + RejectionRoute::SseError => ( + "error", + r#"{"type":"error","code":"invalid_value","message":"The image requires 52150 patches, exceeding the limit of 30000."}"#, + ), + RejectionRoute::ResponseFailed => ( + "response.failed", + r#"{"type":"response.failed","response":{"id":"image-response","status":"failed","error":{"code":"invalid_value","message":"The image requires 52150 patches, exceeding the limit of 30000."}}}"#, + ), + }; + events.extend(scheme.parse_sse(event_type, data, &mut state)?); + return Ok(Box::pin(futures::stream::iter(events.into_iter().map(Ok)))); + } + // Capture committed filesystem state at the retry boundary, before any + // successful output or cancellation can cause a later append. + if let Some(observation) = self.retry_observation.lock().unwrap().as_mut() + && observation.committed.is_none() + { + observation.committed = Some( + observation + .store + .read_all(observation.session_id, observation.segment_id) + .unwrap(), + ); } if let Some(cancel) = self.cancel_on_retry.lock().unwrap().take() { cancel.try_send(()).unwrap(); @@ -43,7 +104,9 @@ impl LlmClient for RejectImages { } } -async fn image_worker() -> ( +async fn image_worker( + route: RejectionRoute, +) -> ( tempfile::TempDir, Worker, RejectImages, @@ -63,7 +126,10 @@ allow = [] "#, ) .unwrap(); - let client = RejectImages::default(); + let client = RejectImages { + route, + ..RejectImages::default() + }; let mut worker = Worker::new( manifest, Engine::<_, Mutable, SessionHistoryMetadata>::new_annotated(client.clone()), @@ -90,24 +156,81 @@ allow = [] ), ]); worker.wire_history_persistence(); + *client.retry_observation.lock().unwrap() = Some(RetryObservation { + // Reopen the store so this observes filesystem records, not engine state. + store: session_store::FsStore::new(worker.store.root_dir()).unwrap(), + session_id: worker.session_id(), + segment_id: worker.segment_id(), + committed: None, + }); (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()) } +fn assert_correction_committed_before_retry(client: &RejectImages) { + let observation = client.retry_observation.lock().unwrap(); + let observation = observation.as_ref().unwrap(); + let raw = observation + .committed + .as_ref() + .expect("the retry must observe the committed correction"); + assert_eq!( + raw.iter() + .filter(|entry| matches!(entry, LogEntry::ToolResultCorrected { .. })) + .count(), + 1, + "{:?}: correction must already be appended at retry", + client.route + ); + let restored = + restore_history_entries(observation.session_id, observation.segment_id, raw).unwrap(); + let corrected = restored + .iter() + .find(|entry| rejected(&entry.item)) + .expect("committed retry history must contain the attachment-free error"); + let requests = client.requests.lock().unwrap(); + assert_eq!(requests.len(), 2, "{:?}", client.route); + assert!(requests[0].items.iter().any( + |item| matches!(item, Item::ToolResult { attachments, .. } if !attachments.is_empty()) + )); + assert_eq!( + requests[1].items.iter().find(|item| rejected(item)), + Some(&corrected.item), + "{:?}: retry must use the durable corrected result", + client.route + ); + assert!(requests[1].items.iter().all( + |item| !matches!(item, Item::ToolResult { attachments, .. } if !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; + for route in REJECTION_ROUTES { + assert_image_correction_survives_restore_and_reseed(route).await; + } +} + +async fn assert_image_correction_survives_restore_and_reseed(route: RejectionRoute) { + let (_dir, mut worker, client) = image_worker(route).await; assert_eq!( worker.run_text("continue").await.unwrap(), WorkerRunResult::Finished ); - assert_eq!(client.requests.lock().unwrap().len(), 2); + assert_correction_committed_before_retry(&client); let raw = worker .store .read_all(worker.session_id(), worker.segment_id()) .unwrap(); + assert_eq!( + raw.iter() + .filter(|entry| matches!(entry, LogEntry::ToolResultCorrected { .. })) + .count(), + 1, + "{route:?}: recovery must append exactly one correction" + ); let correction = raw .iter() .find(|entry| matches!(entry, LogEntry::ToolResultCorrected { .. })) @@ -177,8 +300,17 @@ async fn image_correction_is_committed_before_retry_and_survives_restore_and_res forked_from: None, compacted_from: None, }; + let reseeded_segment_id = session_store::new_segment_id(); + worker + .store + .create_segment(worker.session_id(), reseeded_segment_id, &[seed]) + .unwrap(); + let reopened_store = session_store::FsStore::new(worker.store.root_dir()).unwrap(); + let reseeded_raw = reopened_store + .read_all(worker.session_id(), reseeded_segment_id) + .unwrap(); let restored_again = - restore_history_entries(worker.session_id(), worker.segment_id(), &[seed]).unwrap(); + restore_history_entries(worker.session_id(), reseeded_segment_id, &reseeded_raw).unwrap(); assert_eq!( restored_again .iter() @@ -186,6 +318,12 @@ async fn image_correction_is_committed_before_retry_and_survives_restore_and_res .count(), 1 ); + let effective_correction = restored_again + .iter() + .find(|entry| rejected(&entry.item)) + .unwrap(); + assert_eq!(effective_correction.annotation, original.metadata); + let expected_result = effective_correction.item.clone(); 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 |_| { @@ -193,22 +331,41 @@ async fn image_correction_is_committed_before_retry_and_survives_restore_and_res }) .await; assert!(matches!(out.result, EngineRunExit::Finished)); - assert_eq!(client.requests.lock().unwrap().len(), 3); + let requests = client.requests.lock().unwrap(); + assert_eq!(requests.len(), 3, "{route:?}"); + assert_eq!( + requests[2].items.iter().find(|item| rejected(item)), + Some(&expected_result), + "{route:?}: next request must retain the restored correction" + ); + assert!(requests[2].items.iter().all( + |item| !matches!(item, Item::ToolResult { attachments, .. } if !attachments.is_empty()) + )); } #[tokio::test] async fn cancellation_after_image_correction_does_not_rollback_the_correction() { - let (_dir, mut worker, client) = image_worker().await; + for route in REJECTION_ROUTES { + assert_cancellation_preserves_image_correction(route).await; + } +} + +async fn assert_cancellation_preserves_image_correction(route: RejectionRoute) { + let (_dir, mut worker, client) = image_worker(route).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); + assert_eq!(result, WorkerRunResult::Cancelled, "{route:?}"); + assert_correction_committed_before_retry(&client); let raw = worker .store .read_all(worker.session_id(), worker.segment_id()) .unwrap(); - assert!( + assert_eq!( raw.iter() - .any(|entry| matches!(entry, LogEntry::ToolResultCorrected { .. })) + .filter(|entry| matches!(entry, LogEntry::ToolResultCorrected { .. })) + .count(), + 1, + "{route:?}: cancellation must preserve exactly one correction" ); let restored = restore_history_entries(worker.session_id(), worker.segment_id(), &raw).unwrap(); assert_eq!(