fix: integrate guarded SSE image recovery
This commit is contained in:
+82
-28
@@ -3199,6 +3199,55 @@ impl<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// 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<A>,
|
||||||
|
request: &mut Request,
|
||||||
|
) -> Result<bool, EngineError> {
|
||||||
|
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<A>,
|
||||||
|
) -> Result<ResponseStream, EngineError> {
|
||||||
|
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.
|
/// Open a stream, dispatch all events to the timeline, handle cancellation.
|
||||||
async fn stream_response(
|
async fn stream_response(
|
||||||
&mut self,
|
&mut self,
|
||||||
@@ -3217,37 +3266,14 @@ impl<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
|
|||||||
);
|
);
|
||||||
|
|
||||||
let mut request = request;
|
let mut request = request;
|
||||||
let stream = loop {
|
let stream = self
|
||||||
match self
|
.open_stream_with_image_recovery(&mut request, turn, llm_call, history)
|
||||||
.open_stream_with_retry(request.clone(), turn, llm_call)
|
.await?;
|
||||||
.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();
|
self.tool_call_collector.begin_response();
|
||||||
let mut stream = ResponseStreamPump::start(stream);
|
let mut stream = ResponseStreamPump::start(stream);
|
||||||
let mut early_tools: Option<EarlyToolExecutionBatch> = None;
|
let mut early_tools: Option<EarlyToolExecutionBatch> = None;
|
||||||
let mut response_completed = false;
|
let mut response_completed = false;
|
||||||
|
let mut response_progress = false;
|
||||||
|
|
||||||
let mut event_count: usize = 0;
|
let mut event_count: usize = 0;
|
||||||
loop {
|
loop {
|
||||||
@@ -3285,6 +3311,35 @@ impl<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
|
|||||||
json!({}),
|
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!(
|
if matches!(
|
||||||
&event,
|
&event,
|
||||||
Event::Status(StatusEvent {
|
Event::Status(StatusEvent {
|
||||||
@@ -3293,7 +3348,6 @@ impl<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
|
|||||||
) {
|
) {
|
||||||
response_completed = true;
|
response_completed = true;
|
||||||
}
|
}
|
||||||
self.emit_stream_event(turn, llm_call, &event);
|
|
||||||
self.timeline.dispatch(&event);
|
self.timeline.dispatch(&event);
|
||||||
if early_dispatch {
|
if early_dispatch {
|
||||||
let collected = self.tool_call_collector.take_collected_with_index();
|
let collected = self.tool_call_collector.take_collected_with_index();
|
||||||
|
|||||||
@@ -103,20 +103,7 @@ impl ClientError {
|
|||||||
else {
|
else {
|
||||||
return false;
|
return false;
|
||||||
};
|
};
|
||||||
if matches!(
|
is_image_size_diagnostic(code.as_deref(), message)
|
||||||
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<Duration> {
|
pub fn retry_after(&self) -> Option<Duration> {
|
||||||
@@ -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 な失敗としてリトライ対象になるかを判定する。
|
/// transient な失敗としてリトライ対象になるかを判定する。
|
||||||
///
|
///
|
||||||
/// 対象:
|
/// 対象:
|
||||||
|
|||||||
@@ -109,6 +109,15 @@ pub struct ErrorEvent {
|
|||||||
pub message: String,
|
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 イベントの観測用メタイベント。
|
/// 未対応 SSE イベントの観測用メタイベント。
|
||||||
///
|
///
|
||||||
/// `data_preview` は provider から受け取った raw SSE data の bounded preview、
|
/// `data_preview` は provider から受け取った raw SSE data の bounded preview、
|
||||||
|
|||||||
@@ -753,17 +753,31 @@ pub(crate) fn parse_sse(
|
|||||||
}
|
}
|
||||||
|
|
||||||
"error" => {
|
"error" => {
|
||||||
let ev = from_json::<TopLevelErrorEnvelope>(data).unwrap_or_else(|_| {
|
// Responses documents flat error frames (code/message/param),
|
||||||
TopLevelErrorEnvelope {
|
// while providers also send nested error envelopes. Normalize both
|
||||||
|
// structurally so classification never scans arbitrary raw JSON.
|
||||||
|
let ev = from_json::<TopLevelErrorEnvelope>(data)
|
||||||
|
.or_else(|_| {
|
||||||
|
from_json::<TopLevelError>(data).map(|error| TopLevelErrorEnvelope {
|
||||||
|
error,
|
||||||
|
extra: BTreeMap::new(),
|
||||||
|
})
|
||||||
|
})
|
||||||
|
.unwrap_or_else(|_| TopLevelErrorEnvelope {
|
||||||
error: TopLevelError {
|
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,
|
error_type: None,
|
||||||
code: None,
|
code: None,
|
||||||
extra: BTreeMap::new(),
|
extra: BTreeMap::from([(
|
||||||
|
"raw_data".to_string(),
|
||||||
|
Value::String(data.to_string()),
|
||||||
|
)]),
|
||||||
},
|
},
|
||||||
extra: BTreeMap::new(),
|
extra: BTreeMap::new(),
|
||||||
}
|
});
|
||||||
});
|
|
||||||
let (code, message) = top_level_error_diagnostic(ev);
|
let (code, message) = top_level_error_diagnostic(ev);
|
||||||
Ok(vec![Event::Error(ErrorEvent { code, message })])
|
Ok(vec![Event::Error(ErrorEvent { code, message })])
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -1,8 +1,11 @@
|
|||||||
use agen::llm_client::event::{Event, ResponseStatus, StatusEvent};
|
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::llm_client::{ClientError, LlmClient, Request, ResponseStream};
|
||||||
use agen::tool::{Attachment, ImageAttachment};
|
use agen::tool::{Attachment, ImageAttachment};
|
||||||
use agen::{Engine, EngineRunExit, History, Item, ToolResultDisposition};
|
use agen::{Engine, EngineRunExit, History, Item, ToolResultDisposition};
|
||||||
use async_trait::async_trait;
|
use async_trait::async_trait;
|
||||||
|
use futures::StreamExt;
|
||||||
|
use serde_json::{Value, json};
|
||||||
use std::sync::{Arc, Mutex};
|
use std::sync::{Arc, Mutex};
|
||||||
|
|
||||||
fn rejection() -> ClientError {
|
fn rejection() -> ClientError {
|
||||||
@@ -48,11 +51,28 @@ fn images(request: &Request) -> Vec<String> {
|
|||||||
})
|
})
|
||||||
.collect()
|
.collect()
|
||||||
}
|
}
|
||||||
|
#[derive(Clone, Copy)]
|
||||||
|
enum RejectionPath {
|
||||||
|
Http,
|
||||||
|
SseError,
|
||||||
|
SseNestedError,
|
||||||
|
ResponseFailed,
|
||||||
|
}
|
||||||
|
const STREAM_PATHS: [RejectionPath; 3] = [
|
||||||
|
RejectionPath::SseError,
|
||||||
|
RejectionPath::SseNestedError,
|
||||||
|
RejectionPath::ResponseFailed,
|
||||||
|
];
|
||||||
|
|
||||||
#[derive(Clone)]
|
#[derive(Clone)]
|
||||||
struct Client {
|
struct Client {
|
||||||
requests: Arc<Mutex<Vec<Request>>>,
|
requests: Arc<Mutex<Vec<Request>>>,
|
||||||
failures: usize,
|
failures: usize,
|
||||||
other_error: bool,
|
other_error: bool,
|
||||||
|
path: RejectionPath,
|
||||||
|
prefix: Vec<(&'static str, Value)>,
|
||||||
|
wait_before_error: Option<Arc<tokio::sync::Notify>>,
|
||||||
|
rejection_data: Option<&'static str>,
|
||||||
}
|
}
|
||||||
impl Client {
|
impl Client {
|
||||||
fn new(failures: usize) -> Self {
|
fn new(failures: usize) -> Self {
|
||||||
@@ -60,11 +80,69 @@ impl Client {
|
|||||||
requests: Arc::default(),
|
requests: Arc::default(),
|
||||||
failures,
|
failures,
|
||||||
other_error: false,
|
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<ResponseStream, ClientError> {
|
||||||
|
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]
|
#[async_trait]
|
||||||
impl LlmClient for Client {
|
impl LlmClient for Client {
|
||||||
|
fn tool_call_completion_support(&self) -> agen::llm_client::ToolCallCompletionSupport {
|
||||||
|
agen::llm_client::ToolCallCompletionSupport::PerBlock
|
||||||
|
}
|
||||||
fn clone_boxed(&self) -> Box<dyn LlmClient> {
|
fn clone_boxed(&self) -> Box<dyn LlmClient> {
|
||||||
Box::new(self.clone())
|
Box::new(self.clone())
|
||||||
}
|
}
|
||||||
@@ -72,10 +150,10 @@ impl LlmClient for Client {
|
|||||||
let mut requests = self.requests.lock().unwrap();
|
let mut requests = self.requests.lock().unwrap();
|
||||||
requests.push(request);
|
requests.push(request);
|
||||||
if requests.len() <= self.failures {
|
if requests.len() <= self.failures {
|
||||||
return Err(rejection());
|
return self.reject(rejection());
|
||||||
}
|
}
|
||||||
if self.other_error {
|
if self.other_error {
|
||||||
return Err(ClientError::Api {
|
return self.reject(ClientError::Api {
|
||||||
status: Some(400),
|
status: Some(400),
|
||||||
code: Some("invalid_value".into()),
|
code: Some("invalid_value".into()),
|
||||||
message: "invalid tool arguments".into(),
|
message: "invalid tool arguments".into(),
|
||||||
@@ -226,6 +304,7 @@ async fn early_results_remain_candidates_despite_later_text_in_same_response() {
|
|||||||
struct CountingTool {
|
struct CountingTool {
|
||||||
image: bool,
|
image: bool,
|
||||||
count: Arc<std::sync::atomic::AtomicUsize>,
|
count: Arc<std::sync::atomic::AtomicUsize>,
|
||||||
|
executed: Option<Arc<tokio::sync::Notify>>,
|
||||||
}
|
}
|
||||||
#[async_trait]
|
#[async_trait]
|
||||||
impl agen::tool::Tool for CountingTool {
|
impl agen::tool::Tool for CountingTool {
|
||||||
@@ -235,6 +314,9 @@ impl agen::tool::Tool for CountingTool {
|
|||||||
_: agen::tool::ToolExecutionContext,
|
_: agen::tool::ToolExecutionContext,
|
||||||
) -> Result<agen::tool::ToolOutput, agen::tool::ToolError> {
|
) -> Result<agen::tool::ToolOutput, agen::tool::ToolError> {
|
||||||
self.count.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
|
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 attachments = if self.image {
|
||||||
let Item::ToolResult { attachments, .. } = image_result("image", 1600, 33349, 0) else {
|
let Item::ToolResult { attachments, .. } = image_result("image", 1600, 33349, 0) else {
|
||||||
unreachable!()
|
unreachable!()
|
||||||
@@ -297,6 +379,7 @@ async fn fresh_image_recovery_does_not_repeat_tools_in_either_dispatch_mode() {
|
|||||||
let tool = CountingTool {
|
let tool = CountingTool {
|
||||||
image,
|
image,
|
||||||
count: Arc::default(),
|
count: Arc::default(),
|
||||||
|
executed: None,
|
||||||
};
|
};
|
||||||
counts.push(tool.count.clone());
|
counts.push(tool.count.clone());
|
||||||
engine.register_tool(Arc::new(move || {
|
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()
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|||||||
@@ -1,13 +1,37 @@
|
|||||||
//! Recovery tests use the normal Worker execution and append-log commit gate.
|
//! Recovery tests use the normal Worker execution and append-log commit gate.
|
||||||
use super::*;
|
use super::*;
|
||||||
use agen::llm_client::event::{Event as LlmEvent, ResponseStatus, StatusEvent};
|
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::llm_client::{ClientError, Request, ResponseStream};
|
||||||
use agen::tool::{Attachment, ImageAttachment};
|
use agen::tool::{Attachment, ImageAttachment};
|
||||||
use async_trait::async_trait;
|
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<Vec<LogEntry>>,
|
||||||
|
}
|
||||||
|
|
||||||
#[derive(Clone, Default)]
|
#[derive(Clone, Default)]
|
||||||
struct RejectImages {
|
struct RejectImages {
|
||||||
|
route: RejectionRoute,
|
||||||
requests: Arc<Mutex<Vec<Request>>>,
|
requests: Arc<Mutex<Vec<Request>>>,
|
||||||
|
retry_observation: Arc<Mutex<Option<RetryObservation>>>,
|
||||||
cancel_on_retry: Arc<Mutex<Option<tokio::sync::mpsc::Sender<()>>>>,
|
cancel_on_retry: Arc<Mutex<Option<tokio::sync::mpsc::Sender<()>>>>,
|
||||||
}
|
}
|
||||||
#[async_trait]
|
#[async_trait]
|
||||||
@@ -21,12 +45,49 @@ impl LlmClient for RejectImages {
|
|||||||
);
|
);
|
||||||
self.requests.lock().unwrap().push(request);
|
self.requests.lock().unwrap().push(request);
|
||||||
if has_image {
|
if has_image {
|
||||||
return Err(ClientError::Api {
|
if matches!(self.route, RejectionRoute::Http) {
|
||||||
status: Some(400),
|
return Err(ClientError::Api {
|
||||||
code: Some("invalid_value".into()),
|
status: Some(400),
|
||||||
message: "The image requires 52150 patches, exceeding the limit of 30000.".into(),
|
code: Some("invalid_value".into()),
|
||||||
retry_after: None,
|
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() {
|
if let Some(cancel) = self.cancel_on_retry.lock().unwrap().take() {
|
||||||
cancel.try_send(()).unwrap();
|
cancel.try_send(()).unwrap();
|
||||||
@@ -43,7 +104,9 @@ impl LlmClient for RejectImages {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
async fn image_worker() -> (
|
async fn image_worker(
|
||||||
|
route: RejectionRoute,
|
||||||
|
) -> (
|
||||||
tempfile::TempDir,
|
tempfile::TempDir,
|
||||||
Worker<RejectImages, session_store::FsStore>,
|
Worker<RejectImages, session_store::FsStore>,
|
||||||
RejectImages,
|
RejectImages,
|
||||||
@@ -63,7 +126,10 @@ allow = []
|
|||||||
"#,
|
"#,
|
||||||
)
|
)
|
||||||
.unwrap();
|
.unwrap();
|
||||||
let client = RejectImages::default();
|
let client = RejectImages {
|
||||||
|
route,
|
||||||
|
..RejectImages::default()
|
||||||
|
};
|
||||||
let mut worker = Worker::new(
|
let mut worker = Worker::new(
|
||||||
manifest,
|
manifest,
|
||||||
Engine::<_, Mutable, SessionHistoryMetadata>::new_annotated(client.clone()),
|
Engine::<_, Mutable, SessionHistoryMetadata>::new_annotated(client.clone()),
|
||||||
@@ -90,24 +156,81 @@ allow = []
|
|||||||
),
|
),
|
||||||
]);
|
]);
|
||||||
worker.wire_history_persistence();
|
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)
|
(dir, worker, client)
|
||||||
}
|
}
|
||||||
fn rejected(item: &Item) -> bool {
|
fn rejected(item: &Item) -> bool {
|
||||||
matches!(item, Item::ToolResult { call_id, is_error: true, attachments, .. } if call_id == "image-call" && attachments.is_empty())
|
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]
|
#[tokio::test]
|
||||||
async fn image_correction_is_committed_before_retry_and_survives_restore_and_reseed() {
|
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!(
|
assert_eq!(
|
||||||
worker.run_text("continue").await.unwrap(),
|
worker.run_text("continue").await.unwrap(),
|
||||||
WorkerRunResult::Finished
|
WorkerRunResult::Finished
|
||||||
);
|
);
|
||||||
assert_eq!(client.requests.lock().unwrap().len(), 2);
|
assert_correction_committed_before_retry(&client);
|
||||||
let raw = worker
|
let raw = worker
|
||||||
.store
|
.store
|
||||||
.read_all(worker.session_id(), worker.segment_id())
|
.read_all(worker.session_id(), worker.segment_id())
|
||||||
.unwrap();
|
.unwrap();
|
||||||
|
assert_eq!(
|
||||||
|
raw.iter()
|
||||||
|
.filter(|entry| matches!(entry, LogEntry::ToolResultCorrected { .. }))
|
||||||
|
.count(),
|
||||||
|
1,
|
||||||
|
"{route:?}: recovery must append exactly one correction"
|
||||||
|
);
|
||||||
let correction = raw
|
let correction = raw
|
||||||
.iter()
|
.iter()
|
||||||
.find(|entry| matches!(entry, LogEntry::ToolResultCorrected { .. }))
|
.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,
|
forked_from: None,
|
||||||
compacted_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 =
|
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!(
|
assert_eq!(
|
||||||
restored_again
|
restored_again
|
||||||
.iter()
|
.iter()
|
||||||
@@ -186,6 +318,12 @@ async fn image_correction_is_committed_before_retry_and_survives_restore_and_res
|
|||||||
.count(),
|
.count(),
|
||||||
1
|
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 mut history = History::from_entries(restored_again);
|
||||||
let out = Engine::<_, Mutable, SessionHistoryMetadata>::new_annotated(client.clone())
|
let out = Engine::<_, Mutable, SessionHistoryMetadata>::new_annotated(client.clone())
|
||||||
.run_with_annotation(&mut history, "continue after restart", &mut |_| {
|
.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;
|
.await;
|
||||||
assert!(matches!(out.result, EngineRunExit::Finished));
|
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]
|
#[tokio::test]
|
||||||
async fn cancellation_after_image_correction_does_not_rollback_the_correction() {
|
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());
|
*client.cancel_on_retry.lock().unwrap() = Some(worker.engine().cancel_sender());
|
||||||
let result = worker.run_text("continue").await.unwrap();
|
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
|
let raw = worker
|
||||||
.store
|
.store
|
||||||
.read_all(worker.session_id(), worker.segment_id())
|
.read_all(worker.session_id(), worker.segment_id())
|
||||||
.unwrap();
|
.unwrap();
|
||||||
assert!(
|
assert_eq!(
|
||||||
raw.iter()
|
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();
|
let restored = restore_history_entries(worker.session_id(), worker.segment_id(), &raw).unwrap();
|
||||||
assert_eq!(
|
assert_eq!(
|
||||||
|
|||||||
Reference in New Issue
Block a user