fix: recover from provider image size rejections
This commit is contained in:
Generated
+7
@@ -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"
|
||||
|
||||
@@ -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 }
|
||||
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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
|
||||
}
|
||||
@@ -3,6 +3,7 @@
|
||||
mod engine;
|
||||
mod handler;
|
||||
mod history;
|
||||
mod image_recovery;
|
||||
mod message;
|
||||
|
||||
pub(crate) mod callback;
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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);
|
||||
|
||||
@@ -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()
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -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,
|
||||
};
|
||||
|
||||
@@ -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";
|
||||
|
||||
@@ -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, ..
|
||||
} => {
|
||||
|
||||
@@ -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 { .. }
|
||||
|
||||
@@ -121,6 +121,7 @@ impl SegmentLogSink {
|
||||
| LogEntry::AnnotatedUserInput { .. }
|
||||
| LogEntry::AnnotatedAssistantItem { .. }
|
||||
| LogEntry::AnnotatedToolResult { .. }
|
||||
| LogEntry::ToolResultCorrected { .. }
|
||||
| LogEntry::AnnotatedSystemItem { .. }
|
||||
| LogEntry::Invoke { .. }
|
||||
| LogEntry::RunYielded { .. }
|
||||
|
||||
@@ -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(),
|
||||
|
||||
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user