fix: recover from provider image size rejections

This commit is contained in:
2026-10-04 23:17:43 +09:00
parent 08881fa27a
commit 74a00f9edc
17 changed files with 1135 additions and 39 deletions
Generated
+7
View File
@@ -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"
+1
View File
@@ -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 }
+45 -1
View File
@@ -626,6 +626,10 @@ pub struct Engine<C: LlmClient, S: EngineState = Mutable, A: Send + Sync = ()> {
/// 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<Box<dyn Fn(&Item) -> 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<Box<dyn Fn(&HistoryEntry<A>, &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<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
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<A>, &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<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
"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<EarlyToolExecutionBatch> = None;
@@ -3567,6 +3608,7 @@ impl<C: LlmClient, A: Send + Sync> Engine<C, Mutable, A> {
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<C: LlmClient, A: Send + Sync> Engine<C, Mutable, A> {
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<C: LlmClient, A: Send + Sync> Engine<C, Locked, A> {
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,
+106
View File
@@ -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<Item = &'a Item>) -> Option<usize> {
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<A>(
history: &History<A>,
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
}
+1
View File
@@ -3,6 +3,7 @@
mod engine;
mod handler;
mod history;
mod image_recovery;
mod message;
pub(crate) mod callback;
+29
View File
@@ -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<Duration> {
match self {
ClientError::Api { retry_after, .. } => *retry_after,
+6 -7
View File
@@ -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);
+350
View File
@@ -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<String> {
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<Mutex<Vec<Request>>>,
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<dyn LlmClient> {
Box::new(self.clone())
}
async fn stream(&self, request: Request) -> Result<ResponseStream, ClientError> {
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<std::sync::atomic::AtomicUsize>,
}
#[async_trait]
impl agen::tool::Tool for CountingTool {
async fn execute(
&self,
_: &str,
_: agen::tool::ToolExecutionContext,
) -> Result<agen::tool::ToolOutput, agen::tool::ToolError> {
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<dyn LlmClient> {
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<ResponseStream, ClientError> {
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()
);
}
}
+1
View File
@@ -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,
};
+297 -31
View File
@@ -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<u64>,
seed_entry_index: Option<usize>,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
corrections: Vec<ToolCorrectionPosition>,
}
#[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<ToolCorrectionPosition>,
) -> Result<Vec<PositionedHistoryTurn>, RetainedHistoryReadError> {
let mut turns = Vec::new();
let mut pending_entries = Vec::new();
let mut seen_entries = HashSet::new();
let mut corrections = HashMap::<String, SessionSnapshotEntryData>::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<String, SessionSnapshotEntryData>,
) -> 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<String, SessionSnapshotEntryData>,
) -> Vec<SessionSnapshotEntry> {
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<String, RetainedHistoryReadError> {
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<RetainedHistoryCursor, RetainedHistoryReadError> {
@@ -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";
+13
View File
@@ -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, ..
} => {
+221
View File
@@ -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<Mutex<Vec<Request>>>,
cancel_on_retry: Arc<Mutex<Option<tokio::sync::mpsc::Sender<()>>>>,
}
#[async_trait]
impl LlmClient for RejectImages {
fn clone_boxed(&self) -> Box<dyn LlmClient> {
Box::new(self.clone())
}
async fn stream(&self, request: Request) -> Result<ResponseStream, ClientError> {
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, session_store::FsStore>,
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
);
}
@@ -76,6 +76,7 @@ pub fn live_log_entry_event(entry: LogEntry) -> Option<Event> {
}
entry @ (LogEntry::AnnotatedAssistantItem { .. }
| LogEntry::AnnotatedToolResult { .. }
| LogEntry::ToolResultCorrected { .. }
| LogEntry::RunYielded { .. }
| LogEntry::RunResumed { .. }
| LogEntry::RunCancelled { .. }
+1
View File
@@ -121,6 +121,7 @@ impl SegmentLogSink {
| LogEntry::AnnotatedUserInput { .. }
| LogEntry::AnnotatedAssistantItem { .. }
| LogEntry::AnnotatedToolResult { .. }
| LogEntry::ToolResultCorrected { .. }
| LogEntry::AnnotatedSystemItem { .. }
| LogEntry::Invoke { .. }
| LogEntry::RunYielded { .. }
+11
View File
@@ -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(),
+22
View File
@@ -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<C: LlmClient + 'static, St: Store> Worker<C, St> {
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?;
@@ -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,