diff --git a/Cargo.lock b/Cargo.lock index 0f17a976..e43df6f5 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3370,8 +3370,10 @@ dependencies = [ "async-trait", "aws-config", "aws-sdk-bedrockruntime", + "aws-smithy-runtime-api", "aws-smithy-types", "base64", + "futures-util", "infinity-provider-protocol", "serde_json", "tokio", diff --git a/crates/infinity-agent-core/Cargo.toml b/crates/infinity-agent-core/Cargo.toml index 212db464..eed60270 100644 --- a/crates/infinity-agent-core/Cargo.toml +++ b/crates/infinity-agent-core/Cargo.toml @@ -30,7 +30,7 @@ rhai = { workspace = true } insta = { version = "1", features = ["json", "redactions"] } libc = "0.2" infinity-provider-protocol = { path = "../infinity-provider-protocol", version = "^0.1.0", features = ["mock"] } -tokio = { workspace = true, features = ["macros", "rt-multi-thread", "time"] } +tokio = { workspace = true, features = ["macros", "rt-multi-thread", "time", "test-util"] } futures-util = { workspace = true, features = ["sink"] } [lints] diff --git a/crates/infinity-agent-core/src/event_processor.rs b/crates/infinity-agent-core/src/event_processor.rs index e16f5027..95b38cb8 100644 --- a/crates/infinity-agent-core/src/event_processor.rs +++ b/crates/infinity-agent-core/src/event_processor.rs @@ -19,7 +19,7 @@ use crate::message::{ use crate::system::AgentEvent; use crate::tools::{Tool, ToolContext}; use crate::traits::{ConversationStore, InputSender, StateStore}; -use infinity_provider_protocol::{FinalResponse, ModelProvider}; +use infinity_provider_protocol::{ErrorClass, FinalResponse, ModelProvider}; // ── Public types ── @@ -147,7 +147,27 @@ pub struct HistoryManager { pub history: RefCell>, processed_message_ids: RefCell>, metadata: RefCell>, + // Un-persisted content moves through three phases: + // + // 1. `unvalidated_items` — inputs (user text, tool results, injected + // synthetic results) that the model has not produced output for + // yet. They are part of the in-memory `history` (so completions + // include them) but are **not** persisted by [`Self::sync`]: one + // of them could be the oversized input that blows up the model's + // context window, and persisting it would permanently wedge the + // thread on a poison message. + // 2. `pending_items` — known-safe content awaiting persistence. + // Inputs are promoted here by + // [`Self::mark_inputs_model_validated`] as soon as the model + // streams any output for them (proof the context did not + // overflow); model-produced content is appended here directly. + // 3. Synced — persisted to the conversation store by [`Self::sync`]. + // + // Sequentiality invariant: known-safe content is never appended while + // unvalidated items exist (enforced by an assert), so `pending_items` + // always precedes `unvalidated_items` in history order. pending_items: RefCell>, + unvalidated_items: RefCell>, /// until a turn is complete the data lives here. If errors occur, /// it'll get discarded, if the turn completes, then it will get flushed to /// _both_ pending_items and history. @@ -210,6 +230,7 @@ impl HistoryManager { processed_message_ids: RefCell::new(processed_message_ids), metadata: RefCell::new(metadata), pending_items: RefCell::new(Vec::new()), + unvalidated_items: RefCell::new(Vec::new()), turn_buffer: RefCell::new(Vec::new()), interrupted_tool_calls: RefCell::new(Vec::new()), compacted_up_to: RefCell::new(compacted_up_to), @@ -264,7 +285,7 @@ impl HistoryManager { self.interrupt_pending_tool_call(); } - self.append_pending(message, message_id.clone()); + self.append_unvalidated(message, message_id.clone()); self.processed_message_ids.borrow_mut().insert(message_id); Ok(true) } @@ -319,13 +340,13 @@ impl HistoryManager { call_id: tool_call.call_id.clone(), content: vec![ToolResultContent::Text( infinity_provider_protocol::message::Text { - text: "Tool call interrupted by user".to_owned(), + text: TOOL_CALL_INTERRUPTED_TEXT.to_owned(), }, )], }, display_segments: None, }; - self.append_pending(synthetic_result, format!("{}-interrupted", tool_call.id)); + self.append_unvalidated(synthetic_result, format!("{}-interrupted", tool_call.id)); } } @@ -378,10 +399,15 @@ impl HistoryManager { /// `pending_items`. Called at a flush point: a completed turn (`Final`) or /// a turn-ending tool call. After this, the buffered messages are part of /// committed history and will be persisted by the next [`Self::sync`]. + /// + /// Callers must have promoted any unvalidated inputs first (see + /// [`Self::mark_inputs_model_validated`]): the buffered content is model + /// output, and model output for a request that included unvalidated + /// inputs is exactly the proof required to validate them. pub fn flush_turn(&self) { let drained = std::mem::take(&mut *self.turn_buffer.borrow_mut()); for item in drained { - self.append_pending(item.message, item.message_id); + self.append_known_safe(item.message, item.message_id); } } @@ -435,10 +461,38 @@ impl HistoryManager { .collect() } - fn append_pending(&self, message: InfinityMessage, message_id: String) { + /// Append an *input* (user text, tool result, injected synthetic + /// result) to the in-memory history. The item is not persistable yet: + /// it stays in `unvalidated_items` until the model produces output for + /// it (see [`Self::mark_inputs_model_validated`]). + fn append_unvalidated(&self, message: InfinityMessage, message_id: String) { assert!( self.turn_buffer.borrow().is_empty(), - "bug: append_pending called with un-flushed turn_buffer content" + "bug: append_unvalidated called with un-flushed turn_buffer content" + ); + + self.history.borrow_mut().push(message.clone()); + self.unvalidated_items.borrow_mut().push(PendingItem { + message, + message_id, + }); + } + + /// Append *model-produced* content (assistant text/reasoning, tool + /// calls) to the in-memory history and the known-safe persistence + /// queue. + fn append_known_safe(&self, message: InfinityMessage, message_id: String) { + assert!( + self.turn_buffer.borrow().is_empty(), + "bug: append_known_safe called with un-flushed turn_buffer content" + ); + // Sequentiality safeguard: known-safe content must never be + // committed while unvalidated inputs exist, otherwise sync() would + // persist the model's output without the inputs it answers. + assert!( + self.unvalidated_items.borrow().is_empty(), + "bug: committing model output while unvalidated inputs exist \ + (mark_inputs_model_validated must run first)" ); self.history.borrow_mut().push(message.clone()); @@ -448,6 +502,109 @@ impl HistoryManager { }); } + /// Promote all unvalidated inputs to the known-safe persistence queue. + /// + /// Called as soon as the model streams any output for a request that + /// included them: output means the request was accepted, i.e. the + /// context window did not overflow, so the inputs are safe to persist. + pub fn mark_inputs_model_validated(&self) { + let drained = std::mem::take(&mut *self.unvalidated_items.borrow_mut()); + self.pending_items.borrow_mut().extend(drained); + } + + /// Number of inputs still awaiting model validation. + pub fn unvalidated_len(&self) -> usize { + self.unvalidated_items.borrow().len() + } + + /// Drop all unvalidated *user* inputs: items that are neither tool + /// results nor subscription events (those are placeholder-replaceable, + /// see [`Self::replace_unvalidated_tool_results`], and answering a tool + /// call is always better than stranding it). Dropped items are removed + /// from the in-memory history and their dedup IDs forgotten so a + /// redelivery is not silently ignored. Returns how many items were + /// dropped. + /// + /// Used when the model reports a context overflow: an oversized user + /// input has no safe substitute, and it must not be persisted (or kept + /// in memory) or the thread would be permanently wedged on it. + pub fn drop_unvalidated_user_inputs(&self) -> usize { + let mut unvalidated = self.unvalidated_items.borrow_mut(); + if unvalidated.is_empty() { + return 0; + } + let mut history = self.history.borrow_mut(); + // By the sequentiality invariant the unvalidated items are exactly + // the in-memory history tail. + let tail_start = history.len() - unvalidated.len(); + let mut processed = self.processed_message_ids.borrow_mut(); + + let mut kept = Vec::new(); + let mut dropped = 0; + for item in unvalidated.drain(..) { + if matches!( + item.message, + InfinityMessage::ToolResult { .. } | InfinityMessage::SubscriptionEvent { .. } + ) { + kept.push(item); + } else { + processed.remove(&item.message_id); + dropped += 1; + } + } + history.truncate(tail_start); + history.extend(kept.iter().map(|item| item.message.clone())); + *unvalidated = kept; + dropped + } + + /// Replace the content of every unvalidated tool result — including the + /// bodies of subscription events, which carry a tool result — with + /// `placeholder` (in both the persistence queue and the in-memory + /// history). Returns `true` if at least one was replaced. + /// + /// Used on context overflow: unlike user text, a tool result cannot + /// simply be dropped without stranding its tool call (and a + /// subscription event should still record that an event arrived), but + /// both *can* be answered with a fixed placeholder. Bodies already + /// equal to the placeholder are not counted, so a second overflow + /// reports `false` and the caller falls back to dropping the inputs. + pub fn replace_unvalidated_tool_results(&self, placeholder: &str) -> bool { + let mut unvalidated = self.unvalidated_items.borrow_mut(); + let mut history = self.history.borrow_mut(); + // By the sequentiality invariant the unvalidated items are exactly + // the in-memory history tail. + let tail_start = history.len() - unvalidated.len(); + let mut replaced = false; + for (i, item) in unvalidated.iter_mut().enumerate() { + let result = match &mut item.message { + InfinityMessage::ToolResult { + result, + display_segments, + } => { + *display_segments = None; + result + } + InfinityMessage::SubscriptionEvent { result, .. } => result.as_mut(), + _ => continue, + }; + if matches!( + result.content.first(), + Some(ToolResultContent::Text(t)) if t.text == placeholder + ) { + continue; + } + result.content = vec![ToolResultContent::Text( + infinity_provider_protocol::message::Text { + text: placeholder.to_owned(), + }, + )]; + history[tail_start + i] = item.message.clone(); + replaced = true; + } + replaced + } + /// If the last buffered turn entry is an assistant text message, append /// `text` to it and return `true`. Otherwise return `false` so the caller /// pushes a new buffer entry. This coalesces consecutive text chunks within @@ -468,9 +625,10 @@ impl HistoryManager { } pub async fn sync(&self) -> Result<(), BoxError> { - // `sync` only persists committed (`pending_items`) content. Any in-flight - // turn must have been flushed or discarded before this point; otherwise a - // flush point was missed and buffered content would be silently lost. + // `sync` only persists known-safe (`pending_items`) content. + // Unvalidated inputs deliberately stay in memory (see the field + // docs); any in-flight turn must have been flushed or discarded + // before this point. assert!( self.turn_buffer.borrow().is_empty(), "bug: sync() called with un-flushed turn_buffer content" @@ -600,9 +758,10 @@ impl HistoryManager { std::mem::take(&mut *self.interrupted_tool_calls.borrow_mut()) } - /// Compute a safe spawn point that excludes trailing unanswered tool calls. - /// Returns an absolute store order (accounting for prior compaction offset - /// and ancestor prefix) suitable for use as `spawn_order_override`. + /// Compute a safe spawn point that excludes trailing unanswered tool calls + /// and any unvalidated (not yet persistable) inputs. Returns an absolute + /// store order (accounting for prior compaction offset and ancestor + /// prefix) suitable for use as `spawn_order_override`. pub fn safe_spawn_point(&self) -> usize { let history = self.history.borrow(); // Walk the trailing run of tool calls / tool results (future-proofing @@ -626,6 +785,11 @@ impl HistoryManager { _ => break, } } + // Child threads inherit history *from the store*, so the spawn + // point must also stay before any unvalidated inputs: they occupy + // the in-memory history tail but have not been persisted yet. + let validated_len = history.len() - self.unvalidated_items.borrow().len(); + safe = safe.min(validated_len); // Convert in-memory index to absolute store order by adding the offset // from any prior compaction. The -1 accounts for the compaction summary // message occupying slot 0 in the in-memory history. @@ -1148,6 +1312,76 @@ where pub const IMAGE_OMITTED_PLACEHOLDER: &str = "[image omitted: the current model does not support image inputs]"; +/// Fixed text substituted for a pending tool result when the model reports +/// a context overflow ([`ErrorClass::ContextOverflow`]): the oversized +/// result cannot be sent, but its tool call must still be answered, so the +/// call is resolved with this placeholder instead. +pub const TOOL_RESULT_TOO_LARGE_PLACEHOLDER: &str = + "[tool result omitted: it was too large for the model's remaining context window]"; + +/// Text used to settle a tool call as interrupted. Injected when a user +/// message arrives while the call is unanswered, and also substituted for a +/// pending tool result during overflow recovery when user input had to be +/// dropped: the next thing the model sees is a fresh user message, and +/// "interrupted by user" describes that situation accurately, whereas a +/// "too large" placeholder could mislead the model into re-running the +/// tool. +pub const TOOL_CALL_INTERRUPTED_TEXT: &str = "Tool call interrupted by user"; + +/// Maximum number of retries for one completion round (transient and +/// throttled provider errors). +const MAX_COMPLETION_RETRIES: u32 = 10; +/// Backoff before retrying a [`ErrorClass::Throttled`] provider error. +const THROTTLED_RETRY_DELAY: Duration = Duration::from_secs(30); +/// Backoff before retrying a [`ErrorClass::Transient`] request-initiation +/// error. +const TRANSIENT_RETRY_DELAY: Duration = Duration::from_secs(5); + +/// How [`run_completion`] recovered from a provider-declared context +/// overflow (see [`ErrorClass::ContextOverflow`]). +enum OverflowRecovery { + /// The not-known-safe queue was exactly one replaceable item (tool + /// result / subscription event) — the culprit is unambiguous, so it was + /// replaced with [`TOOL_RESULT_TOO_LARGE_PLACEHOLDER`] and the + /// completion should be retried. + RetryWithPlaceholder, + /// Anything else: replaceable items were settled with + /// [`TOOL_CALL_INTERRUPTED_TEXT`], user inputs were dropped (count + /// `usize`, they have no safe substitute), and the completion must + /// fail. Dropping is never followed by a retry — the user clearly + /// wanted to say something, and silently re-running the round without + /// their words would act on stale instructions. + DroppedInputs(usize), +} + +/// Recover from a context overflow. +/// +/// * If exactly one unvalidated item is pending and it is replaceable (a +/// tool result or subscription event), it must be what overflowed: +/// replace its body with the "too large" placeholder and retry. +/// * Otherwise the culprit is ambiguous (or is user text, which has no safe +/// substitute): settle replaceable items with the same "interrupted by +/// user" text a user interruption would inject — the next thing the model +/// sees is a fresh user message, and a "too large" note could mislead it +/// into re-running the tool — drop the user inputs, and stop. +/// +/// TODO(deferral): when the *committed* history is what overflows (shrinking +/// the fresh inputs does not help), the thread should block on a +/// lower-threshold compaction instead of erroring. That needs deferral +/// support for "wait for the in-flight compaction child", which does not +/// exist yet. +fn recover_from_overflow( + history: &HistoryManager, +) -> OverflowRecovery { + if history.unvalidated_len() == 1 + && history.replace_unvalidated_tool_results(TOOL_RESULT_TOO_LARGE_PLACEHOLDER) + { + return OverflowRecovery::RetryWithPlaceholder; + } + history.replace_unvalidated_tool_results(TOOL_CALL_INTERRUPTED_TEXT); + OverflowRecovery::DroppedInputs(history.drop_unvalidated_user_inputs()) +} + /// Replace image tool-result content with a text placeholder, in place. Used /// to sanitize the chat history before invoking a model that does not declare /// image input support (see `ModelEntry::supports_image_input`). @@ -1226,60 +1460,79 @@ where additional_params: None, }); + // No initiation timeout here: request timeouts are the + // provider's responsibility (e.g. infinity-provider-bedrock + // applies its own 60s initiation timeout and reports it as an + // `ErrorClass::Transient` error). Only cancellation ends the + // wait. let stream_result = tokio::select! { - r = stream_result => { - Ok(r) - } + r = stream_result => r, _ = &mut cancel_rx => { tracing::info!("Completion cancelled during request initiation"); + // The model never accepted this request, so this + // round's inputs remain unvalidated: they stay in + // memory for the next round but are not persisted. return; } - _ = tokio::time::sleep(Duration::from_secs(60)) => { - if retry_count < 10 { - yield CompletionEvent::Info("Stream error (timeout initiating request), retrying...".to_owned()); - retry_count += 1; - continue 'outer; - } else { - Err(Into::::into("Timed out initiating request")) - } - } - }?; + }; let mut llm_stream = match stream_result { Ok(s) => s, Err(e) => { - let err_str = format!("{}", e); - tracing::error!(error = %e, "Completion stream initiation failed"); - - if (err_str.contains("please wait before trying again") || err_str.contains("please try again")) && retry_count < 10 { - tracing::warn!("Stream error (rate limit), retrying..."); - - yield CompletionEvent::Info("Stream error (rate limit), retrying after 30 seconds...".to_owned()); - tokio::select! { - _ = tokio::time::sleep(Duration::from_secs(30)) => {} - _ = &mut cancel_rx => { - tracing::info!("Completion cancelled during retry wait"); - return; + tracing::error!(error = %e, class = ?e.class(), "Completion stream initiation failed"); + + match e.class() { + ErrorClass::Throttled if retry_count < MAX_COMPLETION_RETRIES => { + tracing::warn!("Stream error (rate limit), retrying..."); + yield CompletionEvent::Info(format!( + "Stream error (rate limit), retrying after {} seconds...", + THROTTLED_RETRY_DELAY.as_secs() + )); + tokio::select! { + _ = tokio::time::sleep(THROTTLED_RETRY_DELAY) => {} + _ = &mut cancel_rx => { + tracing::info!("Completion cancelled during retry wait"); + return; + } } + retry_count += 1; + continue 'outer; } - retry_count += 1; - continue 'outer; - } else if (err_str.contains("unexpected end of stream") || err_str.contains("unexpected error when processing the request") || err_str.contains("is unable to process your request")) && retry_count < 10 { - tracing::warn!("Stream error ({err_str}), retrying..."); - - yield CompletionEvent::Info(format!("Stream error ({err_str}), retrying...")); - tokio::select! { - _ = tokio::time::sleep(Duration::from_secs(5)) => {} - _ = &mut cancel_rx => { - tracing::info!("Completion cancelled during retry wait"); - return; + ErrorClass::Transient if retry_count < MAX_COMPLETION_RETRIES => { + tracing::warn!("Stream error ({e}), retrying..."); + yield CompletionEvent::Info(format!("Stream error ({e}), retrying...")); + tokio::select! { + _ = tokio::time::sleep(TRANSIENT_RETRY_DELAY) => {} + _ = &mut cancel_rx => { + tracing::info!("Completion cancelled during retry wait"); + return; + } } + retry_count += 1; + continue 'outer; + } + ErrorClass::ContextOverflow => { + match recover_from_overflow(history) { + OverflowRecovery::RetryWithPlaceholder => { + tracing::warn!("Context overflow: replaced oversized tool result with a placeholder, retrying..."); + yield CompletionEvent::Info("Tool result too large for the model's context window; replaced it with a placeholder and retrying...".to_owned()); + retry_count += 1; + continue 'outer; + } + OverflowRecovery::DroppedInputs(dropped) => { + if dropped > 0 { + tracing::warn!("Context overflow: dropped {dropped} oversized input message(s)"); + yield CompletionEvent::Info("The last input was too large for the model's context window and has been discarded.".to_owned()); + } + Err(Into::::into(e))?; + unreachable!() + } + } + } + _ => { + Err(Into::::into(e))?; + unreachable!() } - retry_count += 1; - continue 'outer; - } else { - Err(Into::::into(e))?; - unreachable!() } } }; @@ -1291,40 +1544,23 @@ where // Race between LLM output and cancellation signal. // We avoid `yield` inside `select!` (async_stream limitation) // by capturing the result into locals first. + // + // Deliberately no inactivity timer here: once a stream is + // live it is never artificially cut off by us. Stall + // handling, if any, belongs to the provider. let cancelled; let llm_next = tokio::select! { - res = llm_stream.next() => { cancelled = false; Ok(res) }, - _ = &mut cancel_rx => { cancelled = true; Ok(None) }, - _ = tokio::time::sleep(Duration::from_secs(120)) => { - cancelled = false; - if retry_count < 10 { - yield CompletionEvent::Info("Stream error (timeout), retrying...".to_owned()); - tracing::warn!("Stream stalled, discarding partial turn and retrying..."); - // Retry rebuilds the request from committed history, so - // drop the abandoned partial turn. - history.discard_turn(); - if is_thinking { - is_thinking = false; - yield CompletionEvent::ThinkingEnd; - } - retry_count += 1; - continue 'outer; - } else { - // Giving up: preserve whatever visible text streamed, - // but trim trailing reasoning — the next turn appends a - // user message and reasoning-then-user is rejected by - // some providers. - history.flush_turn_trimming_reasoning(); - Err(Into::::into("Stream timed out")) - } - }, - }?; + res = llm_stream.next() => { cancelled = false; res }, + _ = &mut cancel_rx => { cancelled = true; None }, + }; if cancelled { tracing::info!("Completion cancelled"); // Terminal: keep the visible partial text, but trim trailing // reasoning. This path fires on user interruption, so the next // message is a user turn and must not follow a reasoning block. + // If the model produced no output yet, this is a no-op and the + // round's inputs stay unvalidated (in memory, unpersisted). history.flush_turn_trimming_reasoning(); if is_thinking { yield CompletionEvent::ThinkingEnd; @@ -1337,7 +1573,7 @@ where is_thinking = false; yield CompletionEvent::ThinkingEnd; } - if retry_count < 10 { + if retry_count < MAX_COMPLETION_RETRIES { // Retry rebuilds the request from committed history, so // drop the abandoned partial turn. history.discard_turn(); @@ -1350,7 +1586,7 @@ where // Giving up: keep visible text, trim trailing reasoning // (the next appended message is a user turn). history.flush_turn_trimming_reasoning(); - Err(Into::::into("Stream timed out"))?; + Err(Into::::into("Stream ended unexpectedly"))?; unreachable!() } }; @@ -1358,6 +1594,10 @@ where let chunk = match res { Ok(c) => { retry_count = 0; + // The model produced output for this request, so its + // context accepted the inputs: they are now safe to + // persist. + history.mark_inputs_model_validated(); c }, Err(e) => { @@ -1365,21 +1605,61 @@ where is_thinking = false; yield CompletionEvent::ThinkingEnd; } - let err_str = format!("{}", e); - if (err_str.contains("unexpected end of stream") || err_str.contains("unexpected error when processing the request")) && retry_count < 10 { - // Retry rebuilds from committed history. - history.discard_turn(); - yield CompletionEvent::Info("Stream error (unexpected end), retrying...".to_owned()); - tracing::warn!("Stream error (unexpected end), discarding partial turn and retrying..."); - tokio::time::sleep(Duration::from_secs(1)).await; - retry_count += 1; - continue 'outer; - } else { - // Giving up: keep visible text, trim trailing reasoning - // (the next appended message is a user turn). - history.flush_turn_trimming_reasoning(); - Err(Into::::into(e))?; - unreachable!() + tracing::error!(error = %e, class = ?e.class(), "Completion stream error"); + match e.class() { + ErrorClass::Transient if retry_count < MAX_COMPLETION_RETRIES => { + // Retry rebuilds from committed history. + history.discard_turn(); + yield CompletionEvent::Info(format!("Stream error ({e}), retrying...")); + tracing::warn!("Stream error (transient), discarding partial turn and retrying..."); + tokio::time::sleep(Duration::from_secs(1)).await; + retry_count += 1; + continue 'outer; + } + ErrorClass::Throttled if retry_count < MAX_COMPLETION_RETRIES => { + history.discard_turn(); + yield CompletionEvent::Info(format!( + "Stream error (rate limit), retrying after {} seconds...", + THROTTLED_RETRY_DELAY.as_secs() + )); + tokio::select! { + _ = tokio::time::sleep(THROTTLED_RETRY_DELAY) => {} + _ = &mut cancel_rx => { + tracing::info!("Completion cancelled during retry wait"); + return; + } + } + retry_count += 1; + continue 'outer; + } + ErrorClass::ContextOverflow => { + // The request never fit the model's context; + // any partial turn is unusable. + history.discard_turn(); + match recover_from_overflow(history) { + OverflowRecovery::RetryWithPlaceholder => { + tracing::warn!("Context overflow: replaced oversized tool result with a placeholder, retrying..."); + yield CompletionEvent::Info("Tool result too large for the model's context window; replaced it with a placeholder and retrying...".to_owned()); + retry_count += 1; + continue 'outer; + } + OverflowRecovery::DroppedInputs(dropped) => { + if dropped > 0 { + tracing::warn!("Context overflow: dropped {dropped} oversized input message(s)"); + yield CompletionEvent::Info("The last input was too large for the model's context window and has been discarded.".to_owned()); + } + Err(Into::::into(e))?; + unreachable!() + } + } + } + _ => { + // Giving up: keep visible text, trim trailing reasoning + // (the next appended message is a user turn). + history.flush_turn_trimming_reasoning(); + Err(Into::::into(e))?; + unreachable!() + } } } }; @@ -4127,4 +4407,721 @@ mod tests { assert_eq!(response_url, "https://example.com/choice"); assert!(hm.history.into_inner().is_empty()); } + + // ═══════════════════════════════════════════════════════════════════ + // Context overflow & input-persistence safety + // + // Inputs must not be persisted to the conversation store until the + // model has produced output for them (which proves the context did not + // overflow). Provider errors carry an `ErrorClass`; core reacts to the + // classification instead of parsing message strings, and never applies + // its own timeouts to model requests. + // ═══════════════════════════════════════════════════════════════════ + + use infinity_provider_protocol::{CompletionError, ErrorClass, ModelProvider}; + + /// Feed a user text input through the same path a step uses + /// (`handle_content`), as opposed to `make_history` which fabricates + /// already-committed history. + fn add_user_input( + hm: &HistoryManager, + text: &str, + message_id: &str, + ) { + let accepted = hm + .handle_content( + InfinityMessage::User { + content: UserContent::text(text), + }, + message_id.to_owned(), + ) + .expect("handle user input"); + assert!(accepted, "input should be accepted"); + } + + /// Feed a tool-result input through the same path a step uses. + fn add_tool_result_input( + hm: &HistoryManager, + tool_call_id: &str, + text: &str, + message_id: &str, + ) { + let accepted = hm + .handle_content( + InfinityMessage::ToolResult { + result: ToolResult { + id: tool_call_id.to_owned(), + call_id: None, + content: vec![ToolResultContent::Text( + infinity_provider_protocol::message::Text { + text: text.to_owned(), + }, + )], + }, + display_segments: None, + }, + message_id.to_owned(), + ) + .expect("handle tool result input"); + assert!(accepted, "tool result should be accepted"); + } + + /// Run one completion to termination, summarizing the yielded events as + /// strings (`text:`, `info:`, `error:`, `done`). + async fn collect_completion_events( + provider: &P, + hm: &HistoryManager, + cancel_rx: tokio::sync::oneshot::Receiver<()>, + ) -> Vec { + let (tool_names, tool_defs, tool_registry) = no_tools(); + let ctx = tool_context(); + let thread_id = ThreadId::from("thread-1"); + let stream = run_completion( + provider, + "mock", + false, + hm, + &tool_names, + &tool_defs, + &tool_registry, + &ctx, + &thread_id, + "msg-1", + None, + cancel_rx, + ); + tokio::pin!(stream); + let mut events = Vec::new(); + while let Some(ev) = stream.next().await { + match ev { + Ok(CompletionEvent::TextChunk(t)) => events.push(format!("text:{t}")), + Ok(CompletionEvent::Info(t)) => events.push(format!("info:{t}")), + Ok(CompletionEvent::Action(CompletionAction::Done(_))) => { + events.push("done".to_owned()); + } + Ok(_) => {} + Err(e) => events.push(format!("error:{e}")), + } + } + events + } + + /// The messages persisted for `thread-1`, debug-formatted for asserts. + fn persisted(store: &InMemoryConversationStore) -> String { + format!( + "{:?}", + store + .thread_messages(&ThreadId::from("thread-1")) + .unwrap_or_default() + ) + } + + fn overflow_error() -> CompletionError { + CompletionError::provider( + ErrorClass::ContextOverflow, + "input is too long for the model", + ) + } + + /// An oversized *user input* cannot be shrunk: on a context-overflow + /// error it must be dropped from the in-memory history and never + /// persisted, so the thread does not permanently hang on a poison + /// message. + #[tokio::test(flavor = "current_thread", start_paused = true)] + async fn oversized_user_input_is_dropped_and_not_persisted() { + let local = tokio::task::LocalSet::new(); + local + .run_until(async { + let (provider, mut ctrl) = mock_provider(); + let convo_store = InMemoryConversationStore::new(); + let hm = make_history(&convo_store, vec![]).await; + add_user_input(&hm, "HUGE INPUT", "msg-huge"); + let (_cancel_tx, cancel_rx) = tokio::sync::oneshot::channel(); + + let handle = tokio::task::spawn_local(async move { + let events = collect_completion_events(&provider, &hm, cancel_rx).await; + (hm, events) + }); + + let _req = ctrl.next_request().await; + ctrl.send_error(overflow_error()); + ctrl.drop_stream(); + + let (hm, events) = handle.await.expect("join completion task"); + assert!( + events.iter().any(|e| e.starts_with("error:")), + "the overflow must surface as a terminal error, got {events:?}" + ); + // The poison input is gone from the in-memory history... + assert!( + !format!("{:?}", hm.history.borrow()).contains("HUGE INPUT"), + "oversized input must be dropped from in-memory history" + ); + // ...and the commit that ends the step persists nothing. + hm.sync().await.expect("sync"); + assert!( + !persisted(&convo_store).contains("HUGE INPUT"), + "oversized input must not be persisted" + ); + }) + .await; + } + + /// An oversized *tool result* can be shrunk: on a context-overflow + /// error it is replaced with a fixed placeholder (the tool call must + /// still be answered) and the completion is retried. + #[tokio::test(flavor = "current_thread", start_paused = true)] + async fn oversized_tool_result_is_replaced_with_placeholder_and_retried() { + let local = tokio::task::LocalSet::new(); + local + .run_until(async { + let (provider, mut ctrl) = mock_provider(); + let convo_store = InMemoryConversationStore::new(); + let hm = make_history( + &convo_store, + vec![ + Message::User { + content: vec![UserContent::text("run it")], + }, + tool_call_msg("tc-1", "some_tool", serde_json::json!({})), + ], + ) + .await; + add_tool_result_input(&hm, "tc-1", "HUGE RESULT", "msg-tr"); + let (_cancel_tx, cancel_rx) = tokio::sync::oneshot::channel(); + + let handle = tokio::task::spawn_local(async move { + let events = collect_completion_events(&provider, &hm, cancel_rx).await; + (hm, events) + }); + + let _req1 = ctrl.next_request().await; + ctrl.send_error(overflow_error()); + ctrl.drop_stream(); + + // The retry must present the placeholder instead of the + // oversized result. + let req2 = tokio::time::timeout(Duration::from_secs(300), ctrl.next_request()) + .await + .expect("core should retry after replacing the oversized tool result"); + let history_debug = format!("{:?}", req2.chat_history); + assert!( + !history_debug.contains("HUGE RESULT"), + "retry must not include the oversized tool result" + ); + assert!( + history_debug.contains(TOOL_RESULT_TOO_LARGE_PLACEHOLDER), + "retry must include the placeholder tool result" + ); + ctrl.send_text("recovered"); + ctrl.finish(); + + let (hm, events) = handle.await.expect("join completion task"); + assert!(events.contains(&"done".to_owned()), "events: {events:?}"); + hm.sync().await.expect("sync"); + let stored = persisted(&convo_store); + assert!( + stored.contains(TOOL_RESULT_TOO_LARGE_PLACEHOLDER), + "the placeholder result must be persisted" + ); + assert!( + !stored.contains("HUGE RESULT"), + "the oversized result must not be persisted" + ); + }) + .await; + } + + /// An oversized *subscription event* is also replaced with the + /// placeholder (its body is a tool result) rather than dropped: the + /// agent should still learn that an event arrived. + #[tokio::test(flavor = "current_thread", start_paused = true)] + async fn oversized_subscription_event_is_replaced_with_placeholder_and_retried() { + let local = tokio::task::LocalSet::new(); + local + .run_until(async { + let (provider, mut ctrl) = mock_provider(); + let convo_store = InMemoryConversationStore::new(); + let hm = make_history( + &convo_store, + vec![Message::User { + content: vec![UserContent::text("subscribe to things")], + }], + ) + .await; + let accepted = hm + .handle_content( + InfinityMessage::SubscriptionEvent { + result: Box::new(ToolResult { + id: "evt-1".to_owned(), + call_id: None, + content: vec![ToolResultContent::Text( + infinity_provider_protocol::message::Text { + text: "HUGE EVENT BODY".to_owned(), + }, + )], + }), + tool_call_id: "sub-1".to_owned(), + child_thread_id: None, + invocation: Some(Box::new(ToolCall::new( + "evt-1", + "receive_event__injected", + serde_json::json!({}), + ))), + }, + "msg-evt".to_owned(), + ) + .expect("handle subscription event"); + assert!(accepted); + let (_cancel_tx, cancel_rx) = tokio::sync::oneshot::channel(); + + let handle = tokio::task::spawn_local(async move { + let events = collect_completion_events(&provider, &hm, cancel_rx).await; + (hm, events) + }); + + let _req1 = ctrl.next_request().await; + ctrl.send_error(overflow_error()); + ctrl.drop_stream(); + + let req2 = tokio::time::timeout(Duration::from_secs(300), ctrl.next_request()) + .await + .expect("core should retry after replacing the oversized event body"); + let history_debug = format!("{:?}", req2.chat_history); + assert!( + !history_debug.contains("HUGE EVENT BODY"), + "retry must not include the oversized event body" + ); + assert!( + history_debug.contains(TOOL_RESULT_TOO_LARGE_PLACEHOLDER), + "retry must include the placeholder event body" + ); + ctrl.send_text("noted"); + ctrl.finish(); + + let (hm, events) = handle.await.expect("join completion task"); + assert!(events.contains(&"done".to_owned()), "events: {events:?}"); + hm.sync().await.expect("sync"); + let stored = persisted(&convo_store); + assert!( + stored.contains(TOOL_RESULT_TOO_LARGE_PLACEHOLDER) + && !stored.contains("HUGE EVENT BODY"), + "the placeholder event must be persisted, not the oversized body; store: {stored}" + ); + }) + .await; + } + + /// If the unvalidated inputs include *user text* alongside a tool + /// result, an overflow must not trigger a retry: the culprit is + /// ambiguous and the user's words have to be dropped, so the tool + /// result is settled with the same "interrupted by user" text a user + /// interruption would inject (a "too large" note could mislead the + /// model into re-running the tool) and the round stops so the user can + /// re-send. + #[tokio::test(flavor = "current_thread", start_paused = true)] + async fn overflow_with_pending_tool_result_and_user_input_does_not_retry() { + let local = tokio::task::LocalSet::new(); + local + .run_until(async { + let (provider, mut ctrl) = mock_provider(); + let convo_store = InMemoryConversationStore::new(); + let hm = make_history( + &convo_store, + vec![ + Message::User { + content: vec![UserContent::text("run it")], + }, + tool_call_msg("tc-1", "some_tool", serde_json::json!({})), + ], + ) + .await; + // The tool result arrives first, then a huge user input + // interrupts before the model produced any output. + add_tool_result_input(&hm, "tc-1", "normal tool result", "msg-tr"); + add_user_input(&hm, "HUGE USER INPUT", "msg-huge"); + let (_cancel_tx, cancel_rx) = tokio::sync::oneshot::channel(); + + let handle = tokio::task::spawn_local(async move { + let events = collect_completion_events(&provider, &hm, cancel_rx).await; + (hm, events) + }); + + let _req1 = ctrl.next_request().await; + ctrl.send_error(overflow_error()); + ctrl.drop_stream(); + + // Time-boxed so a wrong retry attempt fails fast instead of + // hanging the test (the consumer would wait forever on the + // retry round's mock stream). + let (hm, events) = tokio::time::timeout(Duration::from_secs(600), handle) + .await + .expect("the round must stop, not retry") + .expect("join completion task"); + assert!( + events.iter().any(|e| e.starts_with("error:")), + "the overflow must surface as a terminal error, got {events:?}" + ); + assert!( + ctrl.try_next_request().is_none(), + "dropping user input must stop the round, not retry it" + ); + + let history = format!("{:?}", hm.history.borrow()); + // The user text was dropped... + assert!( + !history.contains("HUGE USER INPUT"), + "oversized user input must be dropped; history: {history}" + ); + // ...and the tool call stays answered — settled as + // interrupted, not "too large" (the culprit is ambiguous). + assert!( + history.contains(TOOL_CALL_INTERRUPTED_TEXT), + "the tool result must be settled as interrupted; history: {history}" + ); + assert!( + !history.contains(TOOL_RESULT_TOO_LARGE_PLACEHOLDER), + "ambiguous overflows must not blame the tool result; history: {history}" + ); + assert!(!history.contains("normal tool result")); + + // Nothing from this round persists (the settled result is + // still unvalidated). + hm.sync().await.expect("sync"); + let stored = persisted(&convo_store); + assert!( + !stored.contains("HUGE USER INPUT") + && !stored.contains("normal tool result") + && !stored.contains(TOOL_CALL_INTERRUPTED_TEXT), + "nothing unvalidated may persist; store: {stored}" + ); + }) + .await; + } + + /// If the *placeholder* tool result still overflows, give up: surface a + /// terminal error and persist nothing. Only user input is ever dropped, + /// so the tool call stays answered in memory — re-settled as + /// "interrupted by user", since the round is over and the next thing + /// the model sees will be a fresh user message (a "too large" note + /// could mislead it into re-running the tool). After a process restart + /// the store ends on the unanswered call and the next input settles it + /// the same way (see + /// `rebooted_session_settles_persisted_unanswered_tool_call`). + #[tokio::test(flavor = "current_thread", start_paused = true)] + async fn second_overflow_after_placeholder_gives_up_without_persisting_result() { + let local = tokio::task::LocalSet::new(); + local + .run_until(async { + let (provider, mut ctrl) = mock_provider(); + let convo_store = InMemoryConversationStore::new(); + let hm = make_history( + &convo_store, + vec![ + Message::User { + content: vec![UserContent::text("run it")], + }, + tool_call_msg("tc-1", "some_tool", serde_json::json!({})), + ], + ) + .await; + add_tool_result_input(&hm, "tc-1", "HUGE RESULT", "msg-tr"); + let (_cancel_tx, cancel_rx) = tokio::sync::oneshot::channel(); + + let handle = tokio::task::spawn_local(async move { + let events = collect_completion_events(&provider, &hm, cancel_rx).await; + (hm, events) + }); + + let _req1 = ctrl.next_request().await; + ctrl.send_error(overflow_error()); + ctrl.drop_stream(); + let _req2 = tokio::time::timeout(Duration::from_secs(300), ctrl.next_request()) + .await + .expect("core should retry once with the placeholder result"); + ctrl.send_error(overflow_error()); + ctrl.drop_stream(); + + let (hm, events) = handle.await.expect("join completion task"); + assert!( + events.iter().any(|e| e.starts_with("error:")), + "the second overflow must surface as a terminal error, got {events:?}" + ); + assert!( + ctrl.try_next_request().is_none(), + "the placeholder round must not be retried again" + ); + // The tool call stays answered in memory — re-settled as + // "interrupted" now that the round is abandoned — but + // nothing about the result persists. + let last = format!("{:?}", hm.history.borrow().last()); + assert!( + last.contains("tc-1") && last.contains(TOOL_CALL_INTERRUPTED_TEXT), + "the call must stay answered as interrupted, got {last}" + ); + assert!( + !last.contains(TOOL_RESULT_TOO_LARGE_PLACEHOLDER), + "a 'too large' note would mislead the next round, got {last}" + ); + hm.sync().await.expect("sync"); + let stored = persisted(&convo_store); + assert!(!stored.contains("HUGE RESULT"), "stored: {stored}"); + assert!( + !stored.contains(TOOL_RESULT_TOO_LARGE_PLACEHOLDER) + && !stored.contains(TOOL_CALL_INTERRUPTED_TEXT), + "stored: {stored}" + ); + + // A later user input does not need to interrupt anything — + // the call is already answered in memory. + add_user_input(&hm, "hello again", "msg-next"); + let history = format!("{:?}", hm.history.borrow()); + assert!(history.contains(TOOL_CALL_INTERRUPTED_TEXT)); + assert!(history.contains("hello again")); + }) + .await; + } + + /// Interrupting a completion *before the model produced any output* + /// must not persist the batch's inputs — the model never validated them + /// (they could be the poison input) — but they stay in memory so the + /// next round still sends them. + #[tokio::test(flavor = "current_thread")] + async fn interrupt_before_model_output_keeps_input_in_memory_unpersisted() { + let local = tokio::task::LocalSet::new(); + local + .run_until(async { + let (provider, mut ctrl) = mock_provider(); + let convo_store = InMemoryConversationStore::new(); + let hm = make_history(&convo_store, vec![]).await; + add_user_input(&hm, "hello there", "msg-1"); + let (cancel_tx, cancel_rx) = tokio::sync::oneshot::channel(); + + let handle = tokio::task::spawn_local(async move { + let events = collect_completion_events(&provider, &hm, cancel_rx).await; + (hm, events) + }); + + // Request is in flight; no output yet. Interrupt. + let _req = ctrl.next_request().await; + cancel_tx.send(()).expect("send cancel"); + + let (hm, _events) = handle.await.expect("join completion task"); + // The input survives in memory (the next step resends it)... + assert!( + format!("{:?}", hm.history.borrow()).contains("hello there"), + "interrupted input must stay in the in-memory history" + ); + // ...but must not be persisted: the model never accepted it. + hm.sync().await.expect("sync"); + assert!( + !persisted(&convo_store).contains("hello there"), + "input must not be persisted before the model produced output" + ); + }) + .await; + } + + /// Interrupting *after* the model produced output persists both the + /// input and the partial output: streamed output proves the context did + /// not overflow. + #[tokio::test(flavor = "current_thread")] + async fn interrupt_after_model_output_persists_input_and_partial_text() { + let local = tokio::task::LocalSet::new(); + local + .run_until(async { + let (provider, mut ctrl) = mock_provider(); + let convo_store = InMemoryConversationStore::new(); + let hm = make_history(&convo_store, vec![]).await; + add_user_input(&hm, "hello there", "msg-1"); + let (cancel_tx, cancel_rx) = tokio::sync::oneshot::channel(); + + let handle = tokio::task::spawn_local(async move { + let events = collect_completion_events(&provider, &hm, cancel_rx).await; + (hm, events) + }); + + let _req = ctrl.next_request().await; + ctrl.send_text("partial answer"); + tokio::task::yield_now().await; + tokio::task::yield_now().await; + cancel_tx.send(()).expect("send cancel"); + + let (hm, events) = handle.await.expect("join completion task"); + assert!(events.contains(&"text:partial answer".to_owned())); + hm.sync().await.expect("sync"); + let stored = persisted(&convo_store); + assert!( + stored.contains("hello there"), + "model output validated the input: it must persist" + ); + assert!( + stored.contains("partial answer"), + "the partial output must persist" + ); + }) + .await; + } + + /// A throttled error retries after a longer backoff — driven by the + /// provider's classification, not by message-string matching. + #[tokio::test(flavor = "current_thread", start_paused = true)] + async fn throttled_error_waits_and_retries() { + let local = tokio::task::LocalSet::new(); + local + .run_until(async { + let (provider, mut ctrl) = mock_provider(); + let convo_store = InMemoryConversationStore::new(); + let hm = make_history(&convo_store, vec![]).await; + add_user_input(&hm, "hi", "msg-1"); + let (_cancel_tx, cancel_rx) = tokio::sync::oneshot::channel(); + + let handle = tokio::task::spawn_local(async move { + collect_completion_events(&provider, &hm, cancel_rx).await + }); + + let started = tokio::time::Instant::now(); + let _req1 = ctrl.next_request().await; + ctrl.send_error(CompletionError::provider( + ErrorClass::Throttled, + "simulated throttle", + )); + ctrl.drop_stream(); + + let _req2 = tokio::time::timeout(Duration::from_secs(600), ctrl.next_request()) + .await + .expect("core should retry after a throttled error"); + assert!( + started.elapsed() >= Duration::from_secs(10), + "throttled retries should back off" + ); + ctrl.send_text("ok"); + ctrl.finish(); + + let events = handle.await.expect("join completion task"); + assert!(events.contains(&"done".to_owned()), "events: {events:?}"); + }) + .await; + } + + /// A transient error retries quickly — again from the classification, + /// regardless of the message text. + #[tokio::test(flavor = "current_thread", start_paused = true)] + async fn transient_error_discards_turn_and_retries() { + let local = tokio::task::LocalSet::new(); + local + .run_until(async { + let (provider, mut ctrl) = mock_provider(); + let convo_store = InMemoryConversationStore::new(); + let hm = make_history(&convo_store, vec![]).await; + add_user_input(&hm, "hi", "msg-1"); + let (_cancel_tx, cancel_rx) = tokio::sync::oneshot::channel(); + + let handle = tokio::task::spawn_local(async move { + collect_completion_events(&provider, &hm, cancel_rx).await + }); + + let _req1 = ctrl.next_request().await; + ctrl.send_error(CompletionError::provider( + ErrorClass::Transient, + "connection reset by peer", + )); + ctrl.drop_stream(); + + let _req2 = tokio::time::timeout(Duration::from_secs(600), ctrl.next_request()) + .await + .expect("core should retry after a transient error"); + ctrl.send_text("ok"); + ctrl.finish(); + + let events = handle.await.expect("join completion task"); + assert!(events.contains(&"done".to_owned()), "events: {events:?}"); + }) + .await; + } + + /// A fatal error terminates immediately: no retry request is made. + #[tokio::test(flavor = "current_thread", start_paused = true)] + async fn fatal_error_gives_up_immediately() { + let local = tokio::task::LocalSet::new(); + local + .run_until(async { + let (provider, mut ctrl) = mock_provider(); + let convo_store = InMemoryConversationStore::new(); + let hm = make_history(&convo_store, vec![]).await; + add_user_input(&hm, "hi", "msg-1"); + let (_cancel_tx, cancel_rx) = tokio::sync::oneshot::channel(); + + let handle = tokio::task::spawn_local(async move { + collect_completion_events(&provider, &hm, cancel_rx).await + }); + + let _req1 = ctrl.next_request().await; + ctrl.send_error(CompletionError::provider( + ErrorClass::Fatal, + "the model does not exist", + )); + ctrl.drop_stream(); + + let events = handle.await.expect("join completion task"); + assert!( + events.iter().any(|e| e.starts_with("error:")), + "fatal errors must surface, got {events:?}" + ); + assert!( + ctrl.try_next_request().is_none(), + "fatal errors must not be retried" + ); + }) + .await; + } + + // ── HistoryManager three-phase pending safeguards ── + + /// Inputs folded into history are not persisted by `sync()` until the + /// model validates them. + #[tokio::test] + async fn sync_does_not_persist_inputs_before_model_output() { + let store = InMemoryConversationStore::new(); + let hm = make_history(&store, vec![]).await; + add_user_input(&hm, "not yet validated", "msg-1"); + hm.sync().await.expect("sync"); + assert!( + !persisted(&store).contains("not yet validated"), + "unvalidated inputs must not be persisted by sync()" + ); + } + + /// Sequentiality safeguard: committing model output while unvalidated + /// inputs exist would let `sync()` persist the output *without* the + /// input it answers. That is a bug and must panic. + #[tokio::test] + #[should_panic(expected = "unvalidated")] + async fn flushing_model_output_with_unvalidated_inputs_panics() { + let store = InMemoryConversationStore::new(); + let hm = make_history(&store, vec![]).await; + add_user_input(&hm, "pending input", "msg-1"); + hm.handle_completion( + &StreamChunk::Text("model output".to_owned()), + "completion-1".to_owned(), + None, + ); + // Flushing without validating the pending input first violates + // sequentiality. + hm.flush_turn(); + } + + /// Child threads inherit history *from the store*, so a spawn point + /// must not extend past inputs that have not been persisted yet. + #[tokio::test] + async fn safe_spawn_point_excludes_unvalidated_inputs() { + let store = InMemoryConversationStore::new(); + let hm = make_history(&store, vec![]).await; + add_user_input(&hm, "not persisted yet", "msg-1"); + assert_eq!( + hm.safe_spawn_point(), + 0, + "spawn point must not cover unvalidated (unpersisted) inputs" + ); + } } diff --git a/crates/infinity-agent-core/src/system/local/driver.rs b/crates/infinity-agent-core/src/system/local/driver.rs index 3cfc0773..201a2afd 100644 --- a/crates/infinity-agent-core/src/system/local/driver.rs +++ b/crates/infinity-agent-core/src/system/local/driver.rs @@ -44,7 +44,8 @@ struct InFlightStep<'a> { impl InFlightStep<'_> { /// Interrupt the step and wait for it to wind down. The cancellation /// path flushes whatever streamed so far to the store before returning, - /// so no partial turn is lost. + /// so no partial turn is lost. Inputs the model produced no output for + /// stay in memory (unpersisted) and are re-sent on the next round. async fn cancel(self) -> Result { let _ = self.cancel_tx.send(()); self.fut.await diff --git a/crates/infinity-agent-core/src/system/tests.rs b/crates/infinity-agent-core/src/system/tests.rs index 5b231775..25a7a325 100644 --- a/crates/infinity-agent-core/src/system/tests.rs +++ b/crates/infinity-agent-core/src/system/tests.rs @@ -445,8 +445,10 @@ async fn thread_report_deferred_during_async_tool_wait() { /// Shutting the system down while a completion is in flight interrupts the /// completion (stripping trailing reasoning) and waits for every thread /// driver to flush pending history items before the router task returns. +/// A tool result the model has already produced output for is validated and +/// must be persisted together with the partial output. #[tokio::test(flavor = "current_thread")] -async fn shutdown_persists_in_flight_tool_result() { +async fn shutdown_persists_model_validated_tool_result() { let local = tokio::task::LocalSet::new(); local .run_until(async { @@ -462,9 +464,8 @@ async fn shutdown_persists_in_flight_tool_result() { collect_until_finished(&mut rx).await; // 2. The tool result arrives → the driver starts a new - // completion. Waiting for the model request guarantees the - // completion is in flight and the tool result is sitting in - // the history manager's pending (unsynced) items. + // completion, and the model streams some output for it + // (validating the result for persistence). running .send( tool_result_input("t1", "tc-1", "tool execution result").0, @@ -472,6 +473,12 @@ async fn shutdown_persists_in_flight_tool_result() { ) .await; let _req2 = ctrl.next_request().await; + ctrl.send_text("partial reply"); + loop { + if let Evt::E(AgentEvent::TextChunk { .. }) = next_evt(&mut rx).await { + break; + } + } // 3. Shut down while the model is mid-response. let active_threads = running.active_threads(); @@ -484,25 +491,61 @@ async fn shutdown_persists_in_flight_tool_result() { "no thread drivers should remain after shutdown" ); - // 4. The tool result must have been synced to the store. - use crate::traits::ConversationStore; - let history = conv - .load_history_up_to(ThreadId::from_ref("t1"), None, None) - .await - .expect("load history"); - let has_tool_result = history.iter().any(|m| { - if let crate::message::InfinityMessage::ToolResult { result, .. } = m - && let Some(infinity_provider_protocol::message::ToolResultContent::Text(t)) = - result.content.first() - { - result.id == "tc-1" && t.text.contains("tool execution result") - } else { - false - } - }); + // 4. The tool result and the partial output must have been + // synced to the store. + let stored = persisted(&conv, "t1"); + assert!( + stored.contains("tool execution result"), + "validated in-flight tool result should be persisted on shutdown; history: {stored}" + ); + assert!( + stored.contains("partial reply"), + "partial model output should be persisted on shutdown; history: {stored}" + ); + }) + .await; +} + +/// Shutting down while the model has produced *no* output for a freshly +/// arrived tool result must not persist that result: the model never +/// validated it, and it could be the oversized input that permanently +/// wedges the thread. The durable state keeps the tool call unanswered +/// (a later user input settles it with a synthetic "interrupted" result). +#[tokio::test(flavor = "current_thread")] +async fn shutdown_does_not_persist_unvalidated_tool_result() { + let local = tokio::task::LocalSet::new(); + local + .run_until(async { + let (running, mut rx, mut ctrl, conv) = start_system(vec![Box::new(AsyncTool)], None); + + running + .send_user_text(ThreadId::from_ref("t1"), "do something") + .await; + let _req = ctrl.next_request().await; + ctrl.send_tool_call("tc-1", "async_tool", serde_json::json!({})); + ctrl.finish(); + collect_until_finished(&mut rx).await; + + // The tool result arrives; the completion is in flight but the + // model has not produced any output yet. + running + .send( + tool_result_input("t1", "tc-1", "tool execution result").0, + "res-1", + ) + .await; + let _req2 = ctrl.next_request().await; + + running.shutdown().await; + + let stored = persisted(&conv, "t1"); + assert!( + stored.contains("tc-1"), + "the tool call itself is model output and must be persisted; history: {stored}" + ); assert!( - has_tool_result, - "in-flight tool result should be persisted on shutdown; history: {history:#?}" + !stored.contains("tool execution result"), + "an unvalidated tool result must not be persisted on shutdown; history: {stored}" ); }) .await; @@ -1114,3 +1157,352 @@ async fn compaction_inside_child_thread_does_not_panic() { }) .await; } + +// ═══════════════════════════════════════════════════════════════════════ +// Oversized-input / persistence-safety driver tests +// ═══════════════════════════════════════════════════════════════════════ + +/// The messages persisted for a thread, debug-formatted for asserts. +fn persisted(conv: &crate::stores::InMemoryConversationStore, thread_id: &str) -> String { + format!( + "{:?}", + conv.thread_messages(ThreadId::from_ref(thread_id)) + .unwrap_or_default() + ) +} + +/// Interrupting a request that produced no model output yet must not +/// persist the interrupted input — but the input stays in memory and is +/// re-sent (together with the interrupting message) on the next round. +/// Everything persists once the model produces output for it. +#[tokio::test(flavor = "current_thread")] +async fn interrupt_before_output_persists_nothing_until_model_output() { + let local = tokio::task::LocalSet::new(); + local + .run_until(async { + let (mut running, mut rx, mut ctrl, conv) = start_system(vec![], None); + + // Message A: request goes out, model never answers. + running + .send_user_text(ThreadId::from_ref("t1"), "first message") + .await; + let _req1 = ctrl.next_request().await; + + // Message B interrupts before any output. + running + .send_user_text(ThreadId::from_ref("t1"), "second message") + .await; + let req2 = ctrl.next_request().await; + + // Nothing persisted yet: the model never validated A (it could + // be the poison input blowing up the context window). + let stored = persisted(&conv, "t1"); + assert!( + !stored.contains("first message"), + "interrupted input must not be persisted before model output; store: {stored}" + ); + + // But A is still in memory: the retry request carries A and B. + let history = format!("{:?}", req2.chat_history); + assert!( + history.contains("first message") && history.contains("second message"), + "the retry must include both inputs; history: {history}" + ); + + ctrl.send_text("model reply"); + ctrl.finish(); + collect_until_finished(&mut rx).await; // interrupted round 1 + collect_until_finished(&mut rx).await; // round 2 + wait_idle(&mut running).await; + + // Model output validated everything: now it all persists. + let stored = persisted(&conv, "t1"); + assert!( + stored.contains("first message") + && stored.contains("second message") + && stored.contains("model reply"), + "validated inputs and output must be persisted; store: {stored}" + ); + }) + .await; +} + +/// End-to-end oversized tool result recovery: the result is replaced with a +/// placeholder and retried; if the placeholder still overflows, the step +/// errors out *without* persisting the poison result, keeping the call +/// answered in memory as "interrupted by user" — so a later user input just +/// works and the thread never hangs. +#[tokio::test(flavor = "current_thread")] +async fn oversized_tool_result_never_wedges_the_thread() { + let local = tokio::task::LocalSet::new(); + local + .run_until(async { + use infinity_provider_protocol::{CompletionError, ErrorClass}; + + let (mut running, mut rx, mut ctrl, conv) = + start_system(vec![Box::new(AsyncTool)], None); + + // Round 1: the model calls an async tool. + running + .send_user_text(ThreadId::from_ref("t1"), "use the tool") + .await; + let _req1 = ctrl.next_request().await; + ctrl.send_tool_call("tc-1", "async_tool", serde_json::json!({})); + ctrl.finish(); + collect_until_finished(&mut rx).await; + + // The tool result is enormous: the model reports an overflow. + running + .send( + tool_result_input("t1", "tc-1", "HUGE TOOL RESULT").0, + "res-1", + ) + .await; + let _req2 = ctrl.next_request().await; + ctrl.send_error(CompletionError::provider( + ErrorClass::ContextOverflow, + "input is too long for the model", + )); + ctrl.drop_stream(); + + // Core retries with the placeholder result... + let req3 = tokio::time::timeout(std::time::Duration::from_secs(5), ctrl.next_request()) + .await + .expect("core should retry with a placeholder tool result"); + let history = format!("{:?}", req3.chat_history); + assert!( + !history.contains("HUGE TOOL RESULT"), + "retry must not resend the oversized result; history: {history}" + ); + assert!( + history.contains(crate::event_processor::TOOL_RESULT_TOO_LARGE_PLACEHOLDER), + "retry must answer the call with the placeholder; history: {history}" + ); + + // ...which still overflows: the step gives up. + ctrl.send_error(CompletionError::provider( + ErrorClass::ContextOverflow, + "input is too long for the model", + )); + ctrl.drop_stream(); + collect_until_finished(&mut rx).await; + + // The poison result must not have been persisted. + let stored = persisted(&conv, "t1"); + assert!( + !stored.contains("HUGE TOOL RESULT"), + "the oversized result must not be persisted; store: {stored}" + ); + + // A later user input settles the stranded call and the thread + // keeps working. + running + .send_user_text(ThreadId::from_ref("t1"), "are you alive?") + .await; + let req4 = tokio::time::timeout(std::time::Duration::from_secs(5), ctrl.next_request()) + .await + .expect("a later user input must reach the model (no permanent hang)"); + let history = format!("{:?}", req4.chat_history); + assert!( + history.contains(crate::event_processor::TOOL_CALL_INTERRUPTED_TEXT), + "the call must stay answered (settled as interrupted); history: {history}" + ); + assert!( + !history.contains(crate::event_processor::TOOL_RESULT_TOO_LARGE_PLACEHOLDER), + "a 'too large' note would mislead the model; history: {history}" + ); + assert!( + !history.contains("HUGE TOOL RESULT"), + "the poison result must be gone from the model-facing history" + ); + ctrl.send_text("alive and well"); + ctrl.finish(); + let texts = collect_until_finished(&mut rx).await; + assert_eq!(texts, vec!["alive and well"]); + wait_idle(&mut running).await; + }) + .await; +} + +/// Reboot recovery: a session dies while a tool result was still +/// unvalidated, so the durable store ends on the unanswered tool *call*. +/// When a fresh session boots from that store and user input arrives, the +/// stranded call must be settled with a synthetic "interrupted" result and +/// the thread must keep working. +#[tokio::test(flavor = "current_thread")] +async fn rebooted_session_settles_persisted_unanswered_tool_call() { + let local = tokio::task::LocalSet::new(); + local + .run_until(async { + use super::AgentSystemBuilder; + use crate::stores::{InMemoryConversationStore, InMemoryStateStore}; + + let conv = InMemoryConversationStore::new(); + let state = InMemoryStateStore::new(); + + // ── Session 1: the model calls a tool; its (oversized) result + // arrives but the session shuts down before the model produces + // any output for it. ── + let (model1, mut ctrl1) = model_source(None); + let (tx1, mut rx1) = mpsc::unbounded_channel(); + let running1 = AgentSystemBuilder::new_local(conv.clone(), state.clone(), model1) + .tools(vec![Box::new(AsyncTool) as Box>]) + .start_with_observer(move |_thread_id| TestObserver { tx: tx1.clone() }); + running1 + .send_user_text(ThreadId::from_ref("t1"), "do something") + .await; + let _req1 = ctrl1.next_request().await; + ctrl1.send_tool_call("tc-1", "async_tool", serde_json::json!({})); + ctrl1.finish(); + collect_until_finished(&mut rx1).await; + + running1 + .send( + tool_result_input("t1", "tc-1", "HUGE TOOL RESULT").0, + "res-1", + ) + .await; + let _req2 = ctrl1.next_request().await; + running1.shutdown().await; + + // Durable state ends on the unanswered tool call; the + // unvalidated result was not persisted. + let stored = persisted(&conv, "t1"); + assert!( + stored.contains("tc-1") && !stored.contains("HUGE TOOL RESULT"), + "store must end on the unanswered call; store: {stored}" + ); + + // ── Session 2: boots from the same stores. ── + let (model2, mut ctrl2) = model_source(None); + let (tx2, mut rx2) = mpsc::unbounded_channel(); + let mut running2 = AgentSystemBuilder::new_local(conv.clone(), state.clone(), model2) + .tools(vec![Box::new(AsyncTool) as Box>]) + .start_with_observer(move |_thread_id| TestObserver { tx: tx2.clone() }); + running2 + .send_user_text(ThreadId::from_ref("t1"), "hello again") + .await; + + // The stranded call is settled before the new input reaches the + // model. + let req = ctrl2.next_request().await; + let history = format!("{:?}", req.chat_history); + assert!( + history.contains("Tool call interrupted by user"), + "reboot + user input must settle the stranded call; history: {history}" + ); + assert!(history.contains("hello again"), "history: {history}"); + + ctrl2.send_text("back online"); + ctrl2.finish(); + let texts = collect_until_finished(&mut rx2).await; + assert_eq!(texts, vec!["back online"]); + wait_idle(&mut running2).await; + + // Everything from the recovered round is validated + persisted. + let stored = persisted(&conv, "t1"); + assert!( + stored.contains("Tool call interrupted by user") + && stored.contains("hello again") + && stored.contains("back online"), + "recovered round must persist; store: {stored}" + ); + }) + .await; +} + +/// Variant of the interrupt test where the retry round overflows the +/// context. User text has no safe placeholder substitute, so the only +/// possible reaction is to drop the pending inputs — and when inputs are +/// dropped the step must *stop* (no blind retry: the user clearly wants to +/// say something, and silently re-running without their words would be +/// worse). The thread stays usable for the next input. +#[tokio::test(flavor = "current_thread")] +async fn overflow_after_interrupt_drops_user_inputs_and_stops() { + let local = tokio::task::LocalSet::new(); + local + .run_until(async { + use infinity_provider_protocol::{CompletionError, ErrorClass}; + + let (mut running, mut rx, mut ctrl, conv) = start_system(vec![], None); + + // Message A: request goes out, model never answers. + running + .send_user_text(ThreadId::from_ref("t1"), "first message") + .await; + let _req1 = ctrl.next_request().await; + + // Message B interrupts; the retry carries A and B... + running + .send_user_text(ThreadId::from_ref("t1"), "second message") + .await; + let req2 = ctrl.next_request().await; + let history = format!("{:?}", req2.chat_history); + assert!(history.contains("first message") && history.contains("second message")); + + // ...and overflows the context. + ctrl.send_error(CompletionError::provider( + ErrorClass::ContextOverflow, + "input is too long for the model", + )); + ctrl.drop_stream(); + + // The step surfaces the drop and stops — dropping inputs must + // never be followed by a silent retry. (The first + // CompletionFinished belongs to the interrupted round; the + // overflow happens in the second.) + let mut saw_drop_info = false; + let mut finished_rounds = 0; + while finished_rounds < 2 { + match next_evt(&mut rx).await { + Evt::E(AgentEvent::Info { text }) if text.contains("too large") => { + saw_drop_info = true; + } + Evt::E(AgentEvent::CompletionFinished { .. }) => finished_rounds += 1, + _ => {} + } + } + assert!( + saw_drop_info, + "the dropped input must be surfaced to the user" + ); + wait_idle(&mut running).await; + assert!( + ctrl.try_next_request().is_none(), + "dropping user input must stop the round, not retry it" + ); + + // Nothing was persisted. + let stored = persisted(&conv, "t1"); + assert!( + !stored.contains("first message") && !stored.contains("second message"), + "dropped inputs must not be persisted; store: {stored}" + ); + + // The thread is not wedged: a fresh (smaller) message works, and + // the dropped inputs are gone from the model-facing history too. + running + .send_user_text(ThreadId::from_ref("t1"), "third message") + .await; + let req3 = ctrl.next_request().await; + let history = format!("{:?}", req3.chat_history); + assert!( + history.contains("third message") + && !history.contains("first message") + && !history.contains("second message"), + "dropped inputs must not be resent; history: {history}" + ); + ctrl.send_text("hello!"); + ctrl.finish(); + let texts = collect_until_finished(&mut rx).await; + assert_eq!(texts, vec!["hello!"]); + wait_idle(&mut running).await; + + let stored = persisted(&conv, "t1"); + assert!( + stored.contains("third message") && stored.contains("hello!"), + "the recovered round must persist; store: {stored}" + ); + }) + .await; +} diff --git a/crates/infinity-agent-core/src/system/thread.rs b/crates/infinity-agent-core/src/system/thread.rs index ac2a75be..7b19e4bf 100644 --- a/crates/infinity-agent-core/src/system/thread.rs +++ b/crates/infinity-agent-core/src/system/thread.rs @@ -470,8 +470,10 @@ where } if !any_ready { - // Commit anything prepare persisted (processed IDs, interruption - // results) even though no completion runs. + // Commit anything already known-safe (e.g. processed IDs from + // deduped inputs). Interruption results and other fresh inputs + // stay unvalidated — and unpersisted — until a completion + // produces model output for them. self.history.sync().await?; return Ok(StepOutcome::Skipped); } diff --git a/crates/infinity-provider-bedrock/Cargo.toml b/crates/infinity-provider-bedrock/Cargo.toml index 9aba584e..57079a00 100644 --- a/crates/infinity-provider-bedrock/Cargo.toml +++ b/crates/infinity-provider-bedrock/Cargo.toml @@ -23,3 +23,13 @@ tracing-subscriber = { workspace = true } [lints] workspace = true + +[features] +# Tests that talk to the real Bedrock service using your local AWS +# credentials: `cargo test -p infinity-provider-bedrock --features live-tests`. +live-tests = [] + +[dev-dependencies] +futures-util = { workspace = true } +aws-smithy-runtime-api = { workspace = true } +tokio = { workspace = true, features = ["test-util"] } diff --git a/crates/infinity-provider-bedrock/src/convert.rs b/crates/infinity-provider-bedrock/src/convert.rs index 47be898c..00a8cc9c 100644 --- a/crates/infinity-provider-bedrock/src/convert.rs +++ b/crates/infinity-provider-bedrock/src/convert.rs @@ -6,8 +6,8 @@ use aws_smithy_types::{Document, Number}; use base64::Engine; use base64::prelude::BASE64_STANDARD; use infinity_provider_protocol::{ - AssistantContent, CompletionError, Image, ImageMediaType, ImageSource, Message, Reasoning, - ReasoningContent, ToolDefinition, ToolResultContent, UserContent, + AssistantContent, CompletionError, ErrorClass, Image, ImageMediaType, ImageSource, Message, + Reasoning, ReasoningContent, ToolDefinition, ToolResultContent, UserContent, }; /// Convert a JSON value into a smithy [`Document`] (the representation @@ -165,23 +165,24 @@ fn reasoning_block( }) .count(); if signed_text_count > 1 { - return Err(CompletionError::ProviderError( - "AWS Bedrock does not support multiple signed reasoning text blocks".to_owned(), + return Err(CompletionError::provider( + ErrorClass::Fatal, + "AWS Bedrock does not support multiple signed reasoning text blocks", )); } if signed_text_count == 1 && reasoning.content.len() > 1 { - return Err(CompletionError::ProviderError( + return Err(CompletionError::provider( + ErrorClass::Fatal, "AWS Bedrock requires a single signed reasoning text block without additional \ - reasoning parts" - .to_owned(), + reasoning parts", )); } let text = reasoning.display_text(); if text.is_empty() { - return Err(CompletionError::ProviderError( - "AWS Bedrock reasoning conversion requires at least one text or summary block" - .to_owned(), + return Err(CompletionError::provider( + ErrorClass::Fatal, + "AWS Bedrock reasoning conversion requires at least one text or summary block", )); } @@ -200,14 +201,18 @@ fn image_block(image: Image) -> Result { Some(ImageMediaType::GIF) => bedrock::ImageFormat::Gif, Some(ImageMediaType::WEBP) => bedrock::ImageFormat::Webp, Some(other) => { - return Err(CompletionError::ProviderError(format!( - "AWS Bedrock does not support {} images", - other.to_mime_type() - ))); + return Err(CompletionError::provider( + ErrorClass::Fatal, + format!( + "AWS Bedrock does not support {} images", + other.to_mime_type() + ), + )); } None => { - return Err(CompletionError::ProviderError( - "image content requires a media type for AWS Bedrock".to_owned(), + return Err(CompletionError::provider( + ErrorClass::Fatal, + "image content requires a media type for AWS Bedrock", )); } }; @@ -217,9 +222,9 @@ fn image_block(image: Image) -> Result { "only base64-encoded image data is supported by AWS Bedrock".into(), )); }; - let bytes = BASE64_STANDARD - .decode(data) - .map_err(|e| CompletionError::ProviderError(format!("invalid base64 image data: {e}")))?; + let bytes = BASE64_STANDARD.decode(data).map_err(|e| { + CompletionError::provider(ErrorClass::Fatal, format!("invalid base64 image data: {e}")) + })?; bedrock::ImageBlock::builder() .format(format) diff --git a/crates/infinity-provider-bedrock/src/lib.rs b/crates/infinity-provider-bedrock/src/lib.rs index cbd8340b..ebcf5167 100644 --- a/crates/infinity-provider-bedrock/src/lib.rs +++ b/crates/infinity-provider-bedrock/src/lib.rs @@ -14,12 +14,20 @@ mod stream; use async_trait::async_trait; use aws_sdk_bedrockruntime::error::{DisplayErrorContext, ProvideErrorMetadata, SdkError}; use infinity_provider_protocol::{ - CompletionError, CompletionRequest, ModelEntry, ModelProvider, ModelStream, + CompletionError, CompletionRequest, ErrorClass, ModelEntry, ModelProvider, ModelStream, }; use tokio::sync::OnceCell; type BoxError = Box; +/// How long to wait for Bedrock to accept a `ConverseStream` request before +/// giving up. Bedrock occasionally black-holes a request (it neither +/// responds nor fails); classifying the timeout as [`ErrorClass::Transient`] +/// lets the caller retry. This deliberately only covers request +/// *initiation* — once a response stream is live it is never artificially +/// cut off by us. +const REQUEST_INITIATION_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(60); + /// Extract a useful message from an AWS SDK error: the service error message /// when present (the SDK's plain `Display` omits it), otherwise the full /// error chain. @@ -34,6 +42,79 @@ where } } +/// Classify a Bedrock error into the retry classification declared to +/// callers, from the service error code (when the failure reached the +/// service) and the error message. +/// +/// The message heuristics live *here*, in the provider — the agent runtime +/// only ever sees the resulting [`ErrorClass`]. +pub(crate) fn classify_bedrock_error(code: Option<&str>, message: &str) -> ErrorClass { + let msg = message.to_ascii_lowercase(); + if matches!( + code, + Some("ThrottlingException" | "ServiceQuotaExceededException") + ) || msg.contains("please wait before trying again") + || msg.contains("too many requests") + || msg.contains("please try again") + { + return ErrorClass::Throttled; + } + // Context overflow: Bedrock reports it as a ValidationException + // ("Input is too long for requested model."), anthropic models as + // "input length and `max_tokens` exceed context limit". + if msg.contains("too long") + || msg.contains("too large") + || msg.contains("input length") + || (msg.contains("exceed") && msg.contains("context")) + { + return ErrorClass::ContextOverflow; + } + if matches!( + code, + Some( + "InternalServerException" + | "ServiceUnavailableException" + | "ModelTimeoutException" + | "ModelNotReadyException" + | "ModelStreamErrorException" + | "ModelErrorException" + ) + ) || msg.contains("unexpected end of stream") + || msg.contains("unexpected error when processing the request") + || msg.contains("is unable to process your request") + { + return ErrorClass::Transient; + } + ErrorClass::Fatal +} + +/// Classify a full [`SdkError`]: service errors go through +/// [`classify_bedrock_error`]; transport-level failures (dispatch, timeout, +/// unparsable response) are transient; request construction failures are +/// ours and fatal. +pub(crate) fn classify_sdk_error(err: &SdkError) -> ErrorClass +where + E: ProvideErrorMetadata + std::error::Error + 'static, + R: std::fmt::Debug, +{ + match err { + SdkError::ConstructionFailure(_) => ErrorClass::Fatal, + SdkError::TimeoutError(_) | SdkError::DispatchFailure(_) | SdkError::ResponseError(_) => { + ErrorClass::Transient + } + _ => classify_bedrock_error(err.code(), &sdk_error_message(err)), + } +} + +/// Convert an [`SdkError`] into a classified [`CompletionError`]. +pub(crate) fn completion_error(err: &SdkError) -> CompletionError +where + E: ProvideErrorMetadata + std::error::Error + 'static, + R: std::fmt::Debug, +{ + CompletionError::provider(classify_sdk_error(err), sdk_error_message(err)) +} + /// A model offered by the Bedrock provider, along with the Bedrock-specific /// invocation configuration that stays internal to this crate. struct BedrockModel { @@ -244,7 +325,7 @@ impl ModelProvider for BedrockProvider { ) -> Result { let prepared = prepare_request(&self.models, model_id, request)?; - let response = self + let send = self .client() .await .converse_stream() @@ -254,11 +335,28 @@ impl ModelProvider for BedrockProvider { .set_tool_config(prepared.tool_config) .set_inference_config(Some(prepared.inference_config)) .set_additional_model_request_fields(prepared.additional_params) - .send() + .send(); + + // Guard request *initiation* only (see REQUEST_INITIATION_TIMEOUT); + // the returned stream itself is never timed out. + let response = tokio::time::timeout(REQUEST_INITIATION_TIMEOUT, send) .await + .map_err(|_| { + tracing::error!( + "Bedrock ConverseStream request initiation timed out after {:?}", + REQUEST_INITIATION_TIMEOUT + ); + CompletionError::provider( + ErrorClass::Transient, + format!( + "timed out waiting {}s for Bedrock to accept the request", + REQUEST_INITIATION_TIMEOUT.as_secs() + ), + ) + })? .map_err(|e| { tracing::error!(error = %DisplayErrorContext(&e), "Bedrock ConverseStream SDK error"); - CompletionError::ProviderError(sdk_error_message(&e)) + completion_error(&e) })?; Ok(stream::convert_stream(response)) @@ -382,4 +480,63 @@ mod tests { }; assert!(obj.contains_key("anthropic_beta")); } + + // ── Error classification ── + // + // Classification is mostly string/code matching against real Bedrock + // responses, so asserting the match table here would be tautological. + // The real assertions live in `tests/live.rs` (feature `live-tests`), + // which classifies actual Bedrock service errors using local AWS + // credentials. + + /// A black-holed `ConverseStream` request must fail with a transient + /// (retryable) error after the initiation timeout instead of hanging + /// forever. Uses a Bedrock client whose HTTP connector never responds. + #[tokio::test(start_paused = true)] + async fn request_initiation_times_out_with_transient_error() { + #[derive(Debug)] + struct NeverRespond; + impl aws_smithy_runtime_api::client::http::HttpConnector for NeverRespond { + fn call( + &self, + _request: aws_smithy_runtime_api::client::orchestrator::HttpRequest, + ) -> aws_smithy_runtime_api::client::http::HttpConnectorFuture { + aws_smithy_runtime_api::client::http::HttpConnectorFuture::new( + std::future::pending(), + ) + } + } + impl aws_smithy_runtime_api::client::http::HttpClient for NeverRespond { + fn http_connector( + &self, + _settings: &aws_smithy_runtime_api::client::http::HttpConnectorSettings, + _components: &aws_smithy_runtime_api::client::runtime_components::RuntimeComponents, + ) -> aws_smithy_runtime_api::client::http::SharedHttpConnector { + aws_smithy_runtime_api::client::http::SharedHttpConnector::new(NeverRespond) + } + } + + let config = aws_sdk_bedrockruntime::Config::builder() + .behavior_version(aws_config::BehaviorVersion::latest()) + .region(aws_sdk_bedrockruntime::config::Region::new("us-east-1")) + .credentials_provider(aws_sdk_bedrockruntime::config::Credentials::for_tests()) + .http_client(NeverRespond) + .build(); + let provider = BedrockProvider::new(aws_sdk_bedrockruntime::Client::from_conf(config)); + + let started = tokio::time::Instant::now(); + let Err(err) = provider + .invoke_model("global.anthropic.claude-sonnet-4-6", request("hi")) + .await + else { + panic!("black-holed request must time out") + }; + assert_eq!(err.class(), ErrorClass::Transient); + assert!( + err.to_string().contains("timed out"), + "unexpected message: {err}" + ); + // The timeout must be the initiation timeout, not some other layer. + assert!(started.elapsed() >= REQUEST_INITIATION_TIMEOUT); + } } diff --git a/crates/infinity-provider-bedrock/src/stream.rs b/crates/infinity-provider-bedrock/src/stream.rs index e0549fe9..be1aad32 100644 --- a/crates/infinity-provider-bedrock/src/stream.rs +++ b/crates/infinity-provider-bedrock/src/stream.rs @@ -4,8 +4,8 @@ use aws_sdk_bedrockruntime::operation::converse_stream::ConverseStreamOutput as ConverseStreamResponse; use aws_sdk_bedrockruntime::types as bedrock; use infinity_provider_protocol::{ - CompletionError, FinalResponse, ModelStream, Reasoning, ReasoningContent, StreamChunk, - ToolCall, ToolCallDeltaContent, Usage, + CompletionError, ErrorClass, FinalResponse, ModelStream, Reasoning, ReasoningContent, + StreamChunk, ToolCall, ToolCallDeltaContent, Usage, }; /// An in-progress tool-use content block. @@ -58,10 +58,10 @@ pub(crate) fn convert_stream(response: ConverseStreamResponse) -> ModelStream { Ok(None) => break, Err(e) => { // Forward mid-stream transport/service errors instead of - // silently ending the stream; agent-core retries on some - // of these messages. + // silently ending the stream; the retry classification + // tells agent-core how to react. tracing::error!(error = ?e, "Bedrock ConverseStream receive error"); - yield Err(CompletionError::ProviderError(crate::sdk_error_message(&e))); + yield Err(crate::completion_error(&e)); break; } }; @@ -81,16 +81,18 @@ pub(crate) fn convert_stream(response: ConverseStreamResponse) -> ModelStream { }); } _ => { - yield Err(CompletionError::ProviderError( - "AWS Bedrock sent an unsupported ContentBlockStart".to_owned(), + yield Err(CompletionError::provider( + ErrorClass::Fatal, + "AWS Bedrock sent an unsupported ContentBlockStart", )); } } } bedrock::ConverseStreamOutput::ContentBlockDelta(event) => { let Some(delta) = event.delta else { - yield Err(CompletionError::ProviderError( - "The delta for a content block is missing".to_owned(), + yield Err(CompletionError::provider( + ErrorClass::Fatal, + "The delta for a content block is missing", )); continue; }; @@ -161,8 +163,9 @@ pub(crate) fn convert_stream(response: ConverseStreamResponse) -> ModelStream { } } bedrock::StopReason::MaxTokens => { - yield Err(CompletionError::ProviderError( - "Exceeded max tokens".to_owned(), + yield Err(CompletionError::provider( + ErrorClass::Fatal, + "Exceeded max tokens", )); } _ => {} diff --git a/crates/infinity-provider-bedrock/tests/live.rs b/crates/infinity-provider-bedrock/tests/live.rs new file mode 100644 index 00000000..0f61cb7c --- /dev/null +++ b/crates/infinity-provider-bedrock/tests/live.rs @@ -0,0 +1,95 @@ +//! Tests against the real Bedrock service, gated behind the non-default +//! `live-tests` feature because they need working AWS credentials (and cost +//! a few tokens): +//! +//! ```sh +//! cargo test -p infinity-provider-bedrock --features live-tests +//! ``` +//! +//! These exist because the error *classification* is a table of string/code +//! heuristics matched against real Bedrock responses — unit-testing that +//! table against copies of the same strings would be tautological. Instead, +//! we provoke real service errors and assert the classification. +#![cfg(feature = "live-tests")] + +use futures_util::StreamExt; +use infinity_provider_bedrock::BedrockProvider; +use infinity_provider_protocol::{ + CompletionRequest, ErrorClass, Message, ModelProvider, StreamChunk, +}; + +const MODEL: &str = "global.anthropic.claude-sonnet-4-6"; + +fn request(prompt: impl Into) -> CompletionRequest { + CompletionRequest { + preamble: None, + chat_history: vec![Message::user(prompt.into())], + tools: vec![], + max_tokens: Some(64), + additional_params: Some(serde_json::json!({ "thinking": { "type": "disabled" } })), + } +} + +/// An input far beyond the model's 200k-token context window must be +/// classified as [`ErrorClass::ContextOverflow`] — this is what lets +/// agent-core recover from oversized inputs instead of hanging. +#[tokio::test] +async fn oversized_input_is_classified_as_context_overflow() { + let provider = BedrockProvider::from_env(); + // ~12M characters ≈ >2M tokens, comfortably past any context window + // Bedrock currently offers for this model. + let huge = "lorem ipsum dolor sit amet consectetur ".repeat(300_000); + let err = match provider.invoke_model(MODEL, request(huge)).await { + Err(e) => e, + // Some backends only report the overflow once the stream starts. + Ok(mut stream) => loop { + match stream.next().await { + Some(Err(e)) => break e, + Some(Ok(_)) => continue, + None => panic!("oversized request unexpectedly succeeded"), + } + }, + }; + assert_eq!( + err.class(), + ErrorClass::ContextOverflow, + "expected ContextOverflow, got {:?} for: {err}", + err.class() + ); +} + +/// A nonexistent model id is a permanent error: retrying can never help. +#[tokio::test] +async fn unknown_model_is_classified_as_fatal() { + let provider = BedrockProvider::from_env(); + let err = provider + .invoke_model("anthropic.does-not-exist-v0", request("hi")) + .await + .err() + .expect("invoking a nonexistent model must fail"); + assert_eq!( + err.class(), + ErrorClass::Fatal, + "expected Fatal, got {:?} for: {err}", + err.class() + ); +} + +/// Happy-path sanity check: a small request streams text and finishes, +/// proving the initiation timeout wrapper doesn't interfere with normal +/// operation. +#[tokio::test] +async fn small_request_streams_text() { + let provider = BedrockProvider::from_env(); + let mut stream = provider + .invoke_model(MODEL, request("Reply with the single word: ok")) + .await + .expect("invoke model"); + let mut text = String::new(); + while let Some(item) = stream.next().await { + if let StreamChunk::Text(t) = item.expect("stream item") { + text.push_str(&t); + } + } + assert!(!text.is_empty(), "expected some streamed text"); +} diff --git a/crates/infinity-provider-protocol/src/completion.rs b/crates/infinity-provider-protocol/src/completion.rs index 8bf97435..67473816 100644 --- a/crates/infinity-provider-protocol/src/completion.rs +++ b/crates/infinity-provider-protocol/src/completion.rs @@ -54,6 +54,27 @@ pub struct Usage { pub cached_input_tokens: u64, } +/// Provider-declared classification of a completion error, telling callers +/// how to react. Classification is the provider's job — it knows its own +/// failure modes — so the agent runtime never has to parse error message +/// strings. +#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub enum ErrorClass { + /// Transient failure (dropped stream, backend hiccup, request timeout); + /// retrying the same request may succeed. + Transient, + /// Rate limited / throttled; retrying the same request may succeed + /// after a longer backoff. + Throttled, + /// The request's input does not fit the model's context window; + /// retrying the same request can never succeed. Callers must shrink the + /// input (drop or truncate oversized messages, compact history). + ContextOverflow, + /// Permanent failure (bad request, unknown model, access denied); do + /// not retry. + Fatal, +} + /// Error invoking a completion model. #[derive(Debug, thiserror::Error)] pub enum CompletionError { @@ -65,15 +86,43 @@ pub enum CompletionError { #[error("ResponseError: {0}")] ResponseError(String), - /// Error returned by the completion model provider. - #[error("ProviderError: {0}")] - ProviderError(String), + /// Error returned by the completion model provider, with the provider's + /// retry classification. + #[error("ProviderError: {message}")] + ProviderError { message: String, class: ErrorClass }, /// JSON (de)serialization error. #[error("JsonError: {0}")] JsonError(#[from] serde_json::Error), } +impl CompletionError { + /// A provider error with an explicit retry classification. + pub fn provider(class: ErrorClass, message: impl Into) -> Self { + Self::ProviderError { + message: message.into(), + class, + } + } + + /// How callers should treat this error. + /// + /// * [`RequestError`](Self::RequestError) and + /// [`JsonError`](Self::JsonError) are our side failing to build or + /// parse — retrying the same request cannot help ([`ErrorClass::Fatal`]). + /// * [`ResponseError`](Self::ResponseError) is a transport/framing + /// failure while reading the response — typically transient. + /// * [`ProviderError`](Self::ProviderError) carries the provider's own + /// classification. + pub fn class(&self) -> ErrorClass { + match self { + Self::RequestError(_) | Self::JsonError(_) => ErrorClass::Fatal, + Self::ResponseError(_) => ErrorClass::Transient, + Self::ProviderError { class, .. } => *class, + } + } +} + /// One streamed item of a completion response. /// /// Deltas stream incremental UI-facing content; the non-delta variants diff --git a/crates/infinity-provider-protocol/src/lib.rs b/crates/infinity-provider-protocol/src/lib.rs index 5bc976c6..e5e32552 100644 --- a/crates/infinity-provider-protocol/src/lib.rs +++ b/crates/infinity-provider-protocol/src/lib.rs @@ -27,7 +27,7 @@ pub mod mock; pub mod remote; pub use completion::{ - CompletionError, CompletionRequest, FinalResponse, ModelStream, StreamChunk, + CompletionError, CompletionRequest, ErrorClass, FinalResponse, ModelStream, StreamChunk, ToolCallDeltaContent, ToolDefinition, Usage, }; pub use message::{ diff --git a/crates/infinity-provider-protocol/src/remote.rs b/crates/infinity-provider-protocol/src/remote.rs index 97eb03f0..70bca597 100644 --- a/crates/infinity-provider-protocol/src/remote.rs +++ b/crates/infinity-provider-protocol/src/remote.rs @@ -32,7 +32,8 @@ use tokio::net::{UnixListener, UnixStream}; use tokio_util::codec::{Framed, LinesCodec}; use crate::{ - CompletionError, CompletionRequest, ModelEntry, ModelProvider, ModelStream, StreamChunk, + CompletionError, CompletionRequest, ErrorClass, ModelEntry, ModelProvider, ModelStream, + StreamChunk, }; type BoxError = Box; @@ -42,6 +43,37 @@ type JsonLines = Framed; // ── Wire types ── +/// A completion error as it travels over the wire: the stringified message +/// plus the provider's [`ErrorClass`], so retry classification survives the +/// process boundary. +#[derive(Debug, Serialize, Deserialize)] +pub struct WireError { + pub message: String, + /// Missing in messages from older providers; default to the safe + /// (non-retrying) classification. + #[serde(default = "fatal_class")] + pub class: ErrorClass, +} + +fn fatal_class() -> ErrorClass { + ErrorClass::Fatal +} + +impl From<&CompletionError> for WireError { + fn from(e: &CompletionError) -> Self { + Self { + message: e.to_string(), + class: e.class(), + } + } +} + +impl From for CompletionError { + fn from(e: WireError) -> Self { + CompletionError::provider(e.class, e.message) + } +} + /// A request sent from the client to the provider server. One request is /// sent per connection. #[derive(Debug, Serialize, Deserialize)] @@ -64,12 +96,13 @@ pub enum ProviderResponse { /// The invocation succeeded; `Chunk`s follow. InvokeStarted, /// One streamed item of an invocation. `Err` carries a mid-stream - /// provider error (stringified) and does not end the stream. - Chunk(Result), + /// provider error (message plus retry classification) and does not end + /// the stream. + Chunk(Result), /// The invocation stream finished; the connection will close. StreamEnd, /// The request failed (sent in place of `Models` / `InvokeStarted`). - Error(String), + Error(WireError), } // ── Framing ── @@ -154,7 +187,10 @@ async fn handle_connection( Err(e) => { send_json( &mut framed, - &ProviderResponse::Error(format!("invalid request: {e}")), + &ProviderResponse::Error(WireError { + message: format!("invalid request: {e}"), + class: ErrorClass::Fatal, + }), ) .await?; return Ok(()); @@ -165,7 +201,10 @@ async fn handle_connection( ProviderRequest::ListModels => { let response = match provider.list_models().await { Ok(models) => ProviderResponse::Models(models), - Err(e) => ProviderResponse::Error(e.to_string()), + Err(e) => ProviderResponse::Error(WireError { + message: e.to_string(), + class: ErrorClass::Fatal, + }), }; send_json(&mut framed, &response).await?; } @@ -186,13 +225,13 @@ async fn handle_invoke( let mut stream = match provider.invoke_model(model_id, request).await { Ok(stream) => stream, Err(e) => { - send_json(framed, &ProviderResponse::Error(e.to_string())).await?; + send_json(framed, &ProviderResponse::Error(WireError::from(&e))).await?; return Ok(()); } }; send_json(framed, &ProviderResponse::InvokeStarted).await?; while let Some(item) = stream.next().await { - let wire = item.map_err(|e| e.to_string()); + let wire = item.map_err(|e| WireError::from(&e)); send_json(framed, &ProviderResponse::Chunk(wire)).await?; } send_json(framed, &ProviderResponse::StreamEnd).await @@ -232,7 +271,7 @@ impl ModelProvider for RemoteModelProvider { let mut framed = self.connect_and_send(&ProviderRequest::ListModels).await?; match recv_json::(&mut framed).await? { Some(ProviderResponse::Models(models)) => Ok(models), - Some(ProviderResponse::Error(e)) => Err(e.into()), + Some(ProviderResponse::Error(e)) => Err(e.message.into()), Some(_) => Err("unexpected response from model provider".into()), None => Err("model provider closed the connection without responding".into()), } @@ -250,10 +289,13 @@ impl ModelProvider for RemoteModelProvider { }) .await .map_err(|e| { - CompletionError::ProviderError(format!( - "failed to reach model provider at {}: {e}", - self.socket_path.display() - )) + CompletionError::provider( + ErrorClass::Transient, + format!( + "failed to reach model provider at {}: {e}", + self.socket_path.display() + ), + ) })?; match recv_json::(&mut framed) @@ -262,7 +304,7 @@ impl ModelProvider for RemoteModelProvider { CompletionError::ResponseError(format!("failed reading from model provider: {e}")) })? { Some(ProviderResponse::InvokeStarted) => {} - Some(ProviderResponse::Error(e)) => return Err(CompletionError::ProviderError(e)), + Some(ProviderResponse::Error(e)) => return Err(CompletionError::from(e)), Some(_) => { return Err(CompletionError::ResponseError( "unexpected response from model provider".to_owned(), @@ -279,7 +321,7 @@ impl ModelProvider for RemoteModelProvider { loop { match recv_json::(&mut framed).await { Ok(Some(ProviderResponse::Chunk(item))) => { - yield item.map_err(CompletionError::ProviderError); + yield item.map_err(CompletionError::from); } Ok(Some(ProviderResponse::StreamEnd)) => break, Ok(Some(_)) => { @@ -388,10 +430,44 @@ mod tests { async fn connecting_to_missing_socket_fails_cleanly() { let remote = RemoteModelProvider::new("/nonexistent/provider.sock"); match remote.invoke_model("mock", test_request("hi")).await { - Err(CompletionError::ProviderError(_)) => {} + Err(CompletionError::ProviderError { .. }) => {} Err(other) => panic!("unexpected error variant: {other}"), Ok(_) => panic!("connect should fail"), } assert!(remote.list_models().await.is_err()); } + + /// A provider's retry classification must survive the wire round-trip + /// so core can react to it (e.g. drop oversized inputs on + /// [`ErrorClass::ContextOverflow`]). + #[tokio::test] + async fn error_class_survives_the_wire() { + let (model, mut ctrl) = mock_model(); + let provider = Arc::new(SingleModelProvider::new(test_entry(), model)); + let (path, server) = serve_provider(provider).expect("bind provider socket"); + tokio::spawn(server); + + let remote = RemoteModelProvider::new(&path); + let mut response = remote + .invoke_model("mock", test_request("hello")) + .await + .expect("invoke model"); + + let _request = ctrl.next_request().await; + ctrl.send_error(CompletionError::provider( + ErrorClass::ContextOverflow, + "input is too long", + )); + ctrl.drop_stream(); + + let item = response + .next() + .await + .expect("one stream item") + .expect_err("stream item should be an error"); + assert_eq!(item.class(), ErrorClass::ContextOverflow); + assert!(item.to_string().contains("input is too long")); + + std::fs::remove_file(&path).ok(); + } } diff --git a/crates/infinity-provider-rig/src/convert.rs b/crates/infinity-provider-rig/src/convert.rs index 7fc7ad73..096c6225 100644 --- a/crates/infinity-provider-rig/src/convert.rs +++ b/crates/infinity-provider-rig/src/convert.rs @@ -313,19 +313,50 @@ pub fn usage_from_rig(usage: rig::completion::Usage) -> proto::Usage { /// Convert a rig completion error. Variants the protocol models map 1:1; /// transport-level errors (HTTP, URL parsing) are stringified into /// `ProviderError`. +/// +/// Rig does not expose structured retry information, so provider errors are +/// classified from their message: known rate-limit and overflow phrasings +/// used by the common OpenAI-compatible backends are recognized, everything +/// else defaults to [`proto::ErrorClass::Fatal`]. pub fn error_from_rig(error: rig::completion::CompletionError) -> proto::CompletionError { use rig::completion::CompletionError as RigError; match error { RigError::RequestError(e) => proto::CompletionError::RequestError(e), RigError::ResponseError(e) => proto::CompletionError::ResponseError(e), - RigError::ProviderError(e) => proto::CompletionError::ProviderError(e), + RigError::ProviderError(e) => { + let class = classify_provider_message(&e); + proto::CompletionError::provider(class, e) + } RigError::JsonError(e) => proto::CompletionError::JsonError(e), other @ (RigError::HttpError(_) | RigError::UrlError(_)) => { - proto::CompletionError::ProviderError(other.to_string()) + proto::CompletionError::provider(proto::ErrorClass::Transient, other.to_string()) } } } +fn classify_provider_message(message: &str) -> proto::ErrorClass { + let msg = message.to_ascii_lowercase(); + if msg.contains("rate limit") + || msg.contains("too many requests") + || msg.contains("please try again") + || msg.contains("overloaded") + { + return proto::ErrorClass::Throttled; + } + if msg.contains("context length") + || msg.contains("context window") + || msg.contains("too long") + || msg.contains("too large") + || msg.contains("maximum context") + { + return proto::ErrorClass::ContextOverflow; + } + if msg.contains("internal server error") || msg.contains("service unavailable") { + return proto::ErrorClass::Transient; + } + proto::ErrorClass::Fatal +} + #[cfg(test)] mod tests { use super::*;