From 0fcaaf99f2fccad0ec364c3937aca62262b494e5 Mon Sep 17 00:00:00 2001 From: daniel Date: Mon, 19 Oct 2026 14:31:24 +0100 Subject: [PATCH 1/3] feat: recover from typed context-length rejections --- Cargo.lock | 1 + crates/agentkit-loop/Cargo.toml | 1 + .../src/context_recovery_tests.rs | 521 ++++++++++++++++++ crates/agentkit-loop/src/lib.rs | 458 +++++++++++---- crates/agentkit-loop/src/retry.rs | 15 + crates/agentkit-provider-anthropic/src/lib.rs | 117 ++++ .../agentkit-provider-anthropic/src/stream.rs | 22 +- crates/agentkit-provider-openai/src/lib.rs | 64 +++ .../agentkit-provider-openai/src/responses.rs | 58 +- .../src/responses/retry.rs | 91 ++- .../responses/tests/retry_observability.rs | 94 ++++ .../src/responses/websocket/tests.rs | 48 ++ docs/architecture.md | 25 + 13 files changed, 1398 insertions(+), 117 deletions(-) create mode 100644 crates/agentkit-loop/src/context_recovery_tests.rs diff --git a/Cargo.lock b/Cargo.lock index 8b7893f..4d09e61 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -221,6 +221,7 @@ dependencies = [ "agentkit-task-manager", "agentkit-tools-core", "async-trait", + "futures-timer", "opentelemetry", "serde", "serde_json", diff --git a/crates/agentkit-loop/Cargo.toml b/crates/agentkit-loop/Cargo.toml index 8b93faf..d77cb4a 100644 --- a/crates/agentkit-loop/Cargo.toml +++ b/crates/agentkit-loop/Cargo.toml @@ -18,6 +18,7 @@ agentkit-core = { version = "0.10.5", path = "../agentkit-core" } agentkit-task-manager = { version = "0.10.7", path = "../agentkit-task-manager" } agentkit-tools-core = { version = "0.10.5", path = "../agentkit-tools-core" } async-trait.workspace = true +futures-timer.workspace = true serde.workspace = true serde_json.workspace = true thiserror.workspace = true diff --git a/crates/agentkit-loop/src/context_recovery_tests.rs b/crates/agentkit-loop/src/context_recovery_tests.rs new file mode 100644 index 0000000..bb01b41 --- /dev/null +++ b/crates/agentkit-loop/src/context_recovery_tests.rs @@ -0,0 +1,521 @@ +mod context_recovery { + use super::*; + + #[derive(Clone, Copy)] + enum FailureSite { + Begin, + EarlyStream, + AfterDelta, + AfterSuperseded, + AfterTool, + NonContext, + Always, + AfterTools, + SlowRetry, + SlowStream, + SlowInitial, + } + + struct RecoveryAdapter { + executions: StdArc, + site: FailureSite, + requests: StdArc>>, + } + struct RecoverySession { + executions: StdArc, + site: FailureSite, + requests: StdArc>>, + } + struct RecoveryTurn { + hold_when_empty: bool, + events: VecDeque>, + } + + fn rejection() -> LoopError { + LoopError::ProviderFailure(Box::new(ProviderFailure { + route: ProviderRoute::Unknown, + reason: ProviderFailureReason::InvalidRequest, + last_attempt_reason: None, + upstream: ProviderClassification { + code: UpstreamErrorKind::ContextLengthExceeded, + ..Default::default() + }, + accounting: RetryAccounting::default(), + })) + } + fn finished() -> ModelTurnEvent { + ModelTurnEvent::Finished(ModelTurnResult { + finish_reason: FinishReason::Completed, + output_items: vec![Item::text(ItemKind::Assistant, "ok")], + usage: None, + metadata: MetadataMap::new(), + model: None, + response_id: None, + }) + } + #[async_trait] + impl ModelAdapter for RecoveryAdapter { + type Session = RecoverySession; + async fn start_session(&self, _: SessionConfig) -> Result { + Ok(RecoverySession { + executions: self.executions.clone(), + site: self.site, + requests: self.requests.clone(), + }) + } + } + #[async_trait] + impl ModelSession for RecoverySession { + type Turn = RecoveryTurn; + async fn begin_turn( + &mut self, + request: TurnRequest, + cancellation: Option, + ) -> Result { + let attempt = { + let mut requests = self.requests.lock().unwrap(); + requests.push(request.clone()); + requests.len() + }; + if matches!(self.site, FailureSite::AfterTools) { + if attempt == 2 { + return Err(rejection()); + } + let mut turn = FakeSession.begin_turn(request, cancellation).await?; + let mut events = VecDeque::new(); + while let Some(event) = turn.next_event(None).await? { + events.push_back(Ok(event)); + } + return Ok(RecoveryTurn { + events, + hold_when_empty: false, + }); + } + if attempt > 1 && !matches!(self.site, FailureSite::Always) { + if matches!(self.site, FailureSite::SlowRetry) { + std::future::pending::<()>().await; + } + return Ok(RecoveryTurn { + hold_when_empty: matches!(self.site, FailureSite::SlowStream), + events: if matches!(self.site, FailureSite::SlowStream) { + VecDeque::from([Ok(ModelTurnEvent::Usage(Usage::default()))]) + } else { + VecDeque::from([Ok(finished())]) + }, + }); + } + if matches!(self.site, FailureSite::SlowInitial) { + tokio::time::sleep(Duration::from_millis(40)).await; + } + let events = match self.site { + FailureSite::Begin + | FailureSite::Always + | FailureSite::SlowRetry + | FailureSite::SlowStream + | FailureSite::SlowInitial => { + return Err(rejection()); + } + FailureSite::NonContext => { + return Err(LoopError::Provider("context_length_exceeded".into())); + } + FailureSite::EarlyStream => VecDeque::from([ + Ok(ModelTurnEvent::Usage(Usage::default())), + Err(rejection()), + ]), + FailureSite::AfterDelta => VecDeque::from([ + Ok(ModelTurnEvent::Delta(Delta::AppendText { + part_id: PartId::new("text"), + chunk: "visible".into(), + })), + Err(rejection()), + ]), + FailureSite::AfterSuperseded => VecDeque::from([ + Ok(ModelTurnEvent::Delta(Delta::AppendText { + part_id: PartId::new("text"), + chunk: "visible".into(), + })), + Ok(ModelTurnEvent::ResponseAttemptSuperseded), + Err(rejection()), + ]), + FailureSite::AfterTool => VecDeque::from([ + Ok(ModelTurnEvent::ToolCall(ToolCallPart { + id: ToolCallId::new("unsafe"), + name: "echo".into(), + input: json!({}), + metadata: MetadataMap::new(), + })), + Err(rejection()), + ]), + FailureSite::AfterTools => unreachable!(), + }; + Ok(RecoveryTurn { + events, + hold_when_empty: false, + }) + } + } + #[async_trait] + impl ModelTurn for RecoveryTurn { + async fn next_event( + &mut self, + _: Option, + ) -> Result, LoopError> { + if self.hold_when_empty && self.events.is_empty() { + std::future::pending::<()>().await; + } + self.events.pop_front().transpose() + } + } + + #[derive(Clone, Copy)] + enum Edit { + Change, + Noop, + Hang, + Invalid, + Fail, + } + struct RecoveryMutator { + deadline: StdMutex>, + edit: Edit, + calls: StdArc, + } + #[async_trait] + impl LoopMutator for RecoveryMutator { + async fn mutate( + &self, + cursor: &mut TranscriptCursor<'_>, + ctx: LoopCtx<'_>, + ) -> Result<(), LoopError> { + if ctx.point != MutationPoint::ContextLengthRejected { + assert!(ctx.recovery.is_none()); + return Ok(()); + } + let attempt = self.calls.fetch_add(1, Ordering::SeqCst) + 1; + let info = ctx.recovery.unwrap(); + assert_eq!(info.attempt as usize, attempt); + { + let mut deadline = self.deadline.lock().unwrap(); + match *deadline { + Some(previous) => assert_eq!( + info.deadline, previous, + "all recovery attempts share one absolute deadline" + ), + None => *deadline = Some(info.deadline), + } + } + if matches!(self.edit, Edit::Noop) { + let copy = cursor.to_vec(); + **cursor = copy; + return Ok(()); + } + cursor.insert( + 0, + Item::text(ItemKind::System, format!("recovered {attempt}")), + ); + match self.edit { + Edit::Hang => std::future::pending::<()>().await, + Edit::Fail => return Err(LoopError::Mutator("failed".into())), + Edit::Invalid => cursor.push(Item { + parts: vec![Part::ToolResult(ToolResultPart { + call_id: ToolCallId::new("orphan"), + output: ToolOutput::Text("bad".into()), + is_error: false, + metadata: MetadataMap::new(), + })], + ..Item::text(ItemKind::Tool, "") + }), + _ => {} + } + Ok(()) + } + } + struct CountingEcho { + inner: EchoTool, + count: StdArc, + } + #[async_trait] + impl Tool for CountingEcho { + fn spec(&self) -> &ToolSpec { + self.inner.spec() + } + async fn invoke( + &self, + request: ToolRequest, + ctx: &mut ToolContext<'_>, + ) -> Result { + self.count.fetch_add(1, Ordering::SeqCst); + self.inner.invoke(request, ctx).await + } + } + struct Replacements(StdArc>>>); + impl TranscriptObserver for Replacements { + fn on_transcript_event(&self, _: TranscriptEvent<'_>) {} + fn on_transcript_replaced(&self, _: &SessionId, items: &[Item]) { + self.0.lock().unwrap().push(items.to_vec()); + } + } + async fn driver( + site: FailureSite, + edit: Edit, + attempts: u32, + budget: Duration, + ) -> ( + LoopDriver, + StdArc>>, + StdArc, + StdArc>>>, + CancellationController, + ) { + let requests = StdArc::new(StdMutex::new(Vec::new())); + let calls = StdArc::new(AtomicUsize::new(0)); + let replacements = StdArc::new(StdMutex::new(Vec::new())); + let controller = CancellationController::default(); + let mut tools = ToolRegistry::default(); + let executions = StdArc::new(AtomicUsize::new(0)); + tools.register(CountingEcho { + inner: EchoTool::default(), + count: executions.clone(), + }); + let agent = Agent::builder() + .model(RecoveryAdapter { + executions, + site, + requests: requests.clone(), + }) + .add_tool_source(tools) + .input(vec![Item::text(ItemKind::User, "hello")]) + .context_recovery(ContextRecoveryPolicy { + max_attempts: attempts, + budget, + }) + .mutator(RecoveryMutator { + deadline: StdMutex::new(None), + edit, + calls: calls.clone(), + }) + .transcript_observer(Replacements(replacements.clone())) + .cancellation(controller.handle()) + .build() + .unwrap(); + ( + agent + .start(SessionConfig::new("recovery").with_cache(PromptCacheRequest::automatic())) + .await + .unwrap(), + requests, + calls, + replacements, + controller, + ) + } + + #[tokio::test] + async fn begin_and_early_stream_recover_canonical_request() { + for site in [FailureSite::Begin, FailureSite::EarlyStream] { + let (mut driver, requests, calls, replacements, _) = + driver(site, Edit::Change, 2, Duration::from_secs(1)).await; + assert!(matches!( + driver.next().await.unwrap(), + LoopStep::Finished(_) + )); + let requests = requests.lock().unwrap(); + assert_eq!(requests.len(), 2); + assert_eq!(calls.load(Ordering::SeqCst), 1); + assert_eq!(requests[0].turn_id, requests[1].turn_id); + assert_eq!(requests[0].cache, requests[1].cache); + assert_eq!(requests[0].available_tools, requests[1].available_tools); + assert_eq!(replacements.lock().unwrap()[0], requests[1].transcript); + assert!(driver.transcript.starts_with(&requests[1].transcript)); + } + } + #[tokio::test] + async fn refuses_noop_noncontext_and_visible_output() { + for (site, edit, attempts, expected_calls) in [ + (FailureSite::Begin, Edit::Noop, 2, 1), + (FailureSite::NonContext, Edit::Change, 2, 0), + (FailureSite::AfterDelta, Edit::Change, 2, 0), + (FailureSite::AfterSuperseded, Edit::Change, 2, 0), + (FailureSite::AfterTool, Edit::Change, 2, 0), + (FailureSite::Begin, Edit::Change, 0, 0), + ] { + let (mut driver, requests, calls, replacements, _) = + driver(site, edit, attempts, Duration::from_secs(1)).await; + assert!( + timeout(Duration::from_secs(2), driver.next()) + .await + .expect("recovery hung") + .is_err() + ); + assert_eq!(requests.lock().unwrap().len(), 1); + assert_eq!(calls.load(Ordering::SeqCst), expected_calls); + assert!(replacements.lock().unwrap().is_empty()); + } + } + #[tokio::test] + async fn bounded_exhaustion_and_transactional_failures() { + let (mut exhausted, requests, calls, _, _) = + driver(FailureSite::Always, Edit::Change, 2, Duration::from_secs(1)).await; + assert!( + exhausted + .next() + .await + .unwrap_err() + .is_context_length_exceeded() + ); + assert_eq!(requests.lock().unwrap().len(), 3); + assert_eq!(calls.load(Ordering::SeqCst), 2); + for edit in [Edit::Invalid, Edit::Fail, Edit::Hang] { + let (mut driver, requests, _, replacements, _) = + driver(FailureSite::Begin, edit, 2, Duration::from_millis(20)).await; + assert!( + timeout(Duration::from_secs(2), driver.next()) + .await + .expect("recovery hung") + .is_err() + ); + assert_eq!(driver.transcript, requests.lock().unwrap()[0].transcript); + assert!(replacements.lock().unwrap().is_empty()); + } + } + #[tokio::test] + async fn deadline_bounds_replacement_generation() { + for site in [FailureSite::SlowRetry, FailureSite::SlowStream] { + let (mut driver, requests, _, _, _) = + driver(site, Edit::Change, 2, Duration::from_millis(20)).await; + assert!(matches!( + timeout(Duration::from_secs(2), driver.next()) + .await + .expect("replacement generation hung"), + Err(LoopError::ContextRecoveryDeadlineExceeded) + )); + assert_eq!(requests.lock().unwrap().len(), 2); + } + } + #[tokio::test] + async fn cancellation_interrupts_mutation_without_commit() { + let (mut driver, requests, calls, replacements, controller) = + driver(FailureSite::Begin, Edit::Hang, 2, Duration::from_secs(5)).await; + let cancel = async { + while calls.load(Ordering::SeqCst) == 0 { + tokio::task::yield_now().await; + } + controller.interrupt(); + }; + let (result, _) = timeout(Duration::from_secs(2), async { + tokio::join!(driver.next(), cancel) + }) + .await + .expect("cancellation hung"); + assert!(matches!( + result, + Err(LoopError::Cancelled) | Ok(LoopStep::Finished(_)) + )); + assert!( + driver + .transcript + .starts_with(&requests.lock().unwrap()[0].transcript) + ); + assert!( + !driver + .transcript + .iter() + .any(|item| item.kind == ItemKind::System) + ); + assert!(replacements.lock().unwrap().is_empty()); + } + #[tokio::test] + async fn recovery_after_tools_does_not_replay_execution() { + let (mut driver, requests, calls, _, _) = driver( + FailureSite::AfterTools, + Edit::Change, + 2, + Duration::from_secs(1), + ) + .await; + let _ = run_until_finished(&mut driver).await; + assert_eq!(calls.load(Ordering::SeqCst), 1); + let requests = requests.lock().unwrap(); + assert_eq!(requests.len(), 3); + assert_eq!( + driver + .session + .as_ref() + .unwrap() + .executions + .load(Ordering::SeqCst), + 1 + ); + for request in &requests[1..] { + assert_eq!( + request + .transcript + .iter() + .filter(|item| item.kind == ItemKind::Tool) + .count(), + 1 + ); + } + assert_eq!( + driver + .transcript + .iter() + .filter(|item| item.kind == ItemKind::Tool) + .count(), + 1 + ); + } + #[tokio::test] + async fn budget_starts_at_rejection_and_default_remains_disabled() { + let (mut delayed, _, _, _, _) = driver( + FailureSite::SlowInitial, + Edit::Change, + 2, + Duration::from_millis(20), + ) + .await; + assert!(matches!( + delayed.next().await.unwrap(), + LoopStep::Finished(_) + )); + let (mut disabled, requests, calls, _, _) = + driver(FailureSite::Begin, Edit::Change, 2, Duration::from_secs(1)).await; + disabled.context_recovery = None; + assert!( + disabled + .next() + .await + .unwrap_err() + .is_context_length_exceeded() + ); + assert_eq!(requests.lock().unwrap().len(), 1); + assert_eq!(calls.load(Ordering::SeqCst), 0); + } + #[test] + fn classifier_accepts_only_typed_context_evidence() { + assert!(rejection().is_context_length_exceeded()); + let mut failure = *rejection().provider_failure().unwrap(); + failure.upstream.code = UpstreamErrorKind::Unknown; + failure.upstream.error_type = UpstreamErrorKind::ContextLengthExceeded; + assert!(!LoopError::ProviderFailure(Box::new(failure)).is_context_length_exceeded()); + failure.upstream.error_type = UpstreamErrorKind::InvalidRequestError; + failure.upstream.http_status = Some(400); + assert!(!LoopError::ProviderFailure(Box::new(failure)).is_context_length_exceeded()); + assert!( + !LoopError::Provider("maximum context length exceeded".into()) + .is_context_length_exceeded() + ); + failure.upstream.code = UpstreamErrorKind::ContextLengthExceeded; + for reason in [ + ProviderFailureReason::Cancelled, + ProviderFailureReason::ReplayUnsafe, + ] { + failure.reason = reason; + assert!(!LoopError::ProviderFailure(Box::new(failure)).is_context_length_exceeded()); + } + assert_eq!( + serde_json::to_string(&MutationPoint::ContextLengthRejected).unwrap(), + "\"ContextLengthRejected\"" + ); + } +} diff --git a/crates/agentkit-loop/src/lib.rs b/crates/agentkit-loop/src/lib.rs index 33f7c43..daeb3ed 100644 --- a/crates/agentkit-loop/src/lib.rs +++ b/crates/agentkit-loop/src/lib.rs @@ -820,10 +820,15 @@ pub trait TranscriptObserver: Send + Sync { /// state behind interior mutability so the driver can share an /// `Arc`. fn on_transcript_event(&self, event: TranscriptEvent<'_>); + + /// A validated recovery transaction replaced the canonical transcript. + /// Consumers retaining transcript state should replace it with `items`. + fn on_transcript_replaced(&self, _session_id: &SessionId, _items: &[Item]) {} } /// Session-addressed transcript append event delivered to -/// [`TranscriptObserver`]s. +/// [`TranscriptObserver`]s. Replacement transactions use +/// [`TranscriptObserver::on_transcript_replaced`] instead. #[derive(Clone, Debug)] pub struct TranscriptEvent<'a> { /// Session this transcript append belongs to. @@ -845,6 +850,28 @@ pub enum MutationPoint { /// A turn has fully ended (assistant final, interrupt, or cancellation) /// and any new user input has not yet been dispatched. AfterTurnEnded, + /// A confirmed context-length rejection before any assistant output. + ContextLengthRejected, +} + +/// Opt-in limits for transcript recovery after confirmed context rejection. +/// Zero attempts or a zero budget disables recovery. The budget begins at the +/// first rejection, not at the start of an ordinary model request. +#[derive(Clone, Copy, Debug)] +pub struct ContextRecoveryPolicy { + /// Maximum mutation/retry attempts per inference boundary. + pub max_attempts: u32, + /// Total time allowed for mutation and replacement model generation. + pub budget: std::time::Duration, +} + +/// Bounds for the current context-rejection mutation pass. +#[derive(Clone, Copy, Debug)] +pub struct ContextRecoveryInfo { + /// One-based recovery attempt. + pub attempt: u32, + /// Absolute deadline shared by all recovery attempts. + pub deadline: std::time::Instant, } /// Sink for emitting [`AgentEvent`]s from inside a [`LoopMutator`]. @@ -863,6 +890,8 @@ pub struct LoopCtx<'a> { pub turn_id: Option<&'a agentkit_core::TurnId>, /// Where in the loop the mutator is running. pub point: MutationPoint, + /// Present only during context-rejection recovery. + pub recovery: Option, /// Cancellation handle for the active turn, if any. pub cancellation: Option, /// Sink for emitting events from the mutator (telemetry, progress). @@ -905,6 +934,11 @@ impl<'a> std::ops::DerefMut for TranscriptCursor<'a> { #[async_trait] pub trait LoopMutator: Send + Sync { /// Run this mutator. Returning without writing to `cursor` is a no-op. + /// At [`MutationPoint::ContextLengthRejected`], edits are staged and only + /// committed after the entire pass succeeds and validates. The future may + /// be dropped on cancellation or deadline. Events emitted through `ctx` + /// are progress notifications, not evidence that edits were committed; + /// transcript observers receive replacement only after successful commit. /// Errors abort the loop; protocol-violating mutations (orphaned tool /// uses or results) are detected by validation and turned into /// [`LoopError::Mutator`]. @@ -1371,6 +1405,7 @@ where resources: Arc, cancellation: Option, mutators: Vec>, + context_recovery: Option, observers: Vec>, transcript_observers: Vec>, transcript: Vec, @@ -1436,6 +1471,7 @@ where resources: self.resources.clone(), cancellation: self.cancellation.clone(), mutators: self.mutators.clone(), + context_recovery: self.context_recovery, observers: self.observers.clone(), transcript_observers: self.transcript_observers.clone(), transcript: self.transcript.clone(), @@ -1474,6 +1510,7 @@ where resources: Arc, cancellation: Option, mutators: Vec>, + context_recovery: Option, observers: Vec>, transcript_observers: Vec>, transcript: Vec, @@ -1495,6 +1532,7 @@ where resources: Arc::new(()), cancellation: None, mutators: Vec::new(), + context_recovery: None, observers: Vec::new(), transcript_observers: Vec::new(), transcript: Vec::new(), @@ -1570,6 +1608,13 @@ where self } + /// Enable bounded recovery of confirmed context-length rejections. + /// Only model generation is retried; tool execution is never replayed. + pub fn context_recovery(mut self, policy: ContextRecoveryPolicy) -> Self { + self.context_recovery = Some(policy); + self + } + /// Register a [`LoopMutator`] that runs at every [`MutationPoint`]. /// /// Multiple mutators may be registered; they run in registration order @@ -1655,6 +1700,7 @@ where resources: self.resources, cancellation: self.cancellation, mutators: self.mutators, + context_recovery: self.context_recovery, observers: self.observers, transcript_observers: self.transcript_observers, transcript: self.transcript, @@ -1711,6 +1757,7 @@ where resources: Arc, cancellation: Option, mutators: Vec>, + context_recovery: Option, observers: Vec>, transcript_observers: Vec>, transcript: Vec, @@ -2059,6 +2106,7 @@ where session_id: &session_id, turn_id, point, + recovery: None, cancellation: cancellation.clone(), emitter: &emitter, }; @@ -2070,6 +2118,89 @@ where Ok(()) } + /// Stage recovery edits separately: a failed, cancelled, or timed-out + /// mutator must never leave an unvalidated canonical transcript behind. + async fn recover_context( + &mut self, + error: &LoopError, + turn_id: &agentkit_core::TurnId, + cancellation: Option, + recovery: &mut Option, + ) -> Result { + let Some(policy) = self.context_recovery else { + return Ok(false); + }; + if !error.is_context_length_exceeded() + || policy.max_attempts == 0 + || policy.budget.is_zero() + || self.mutators.is_empty() + { + return Ok(false); + } + let previous_attempt = recovery.map_or(0, |info| info.attempt); + if previous_attempt >= policy.max_attempts { + return Ok(false); + } + let deadline = match *recovery { + Some(info) => info.deadline, + None => std::time::Instant::now() + .checked_add(policy.budget) + .ok_or_else(|| { + LoopError::InvalidState("context recovery budget is too large".into()) + })?, + }; + let info = ContextRecoveryInfo { + attempt: previous_attempt + 1, + deadline, + }; + *recovery = Some(info); + let mut candidate = self.transcript.clone(); + let emitter = DriverEmitter { + session_id: &self.observed_session_id, + observers: &self.observers, + }; + let mutation = async { + let mut cursor = TranscriptCursor { + items: &mut candidate, + dirty: false, + }; + for mutator in &self.mutators { + mutator + .mutate( + &mut cursor, + LoopCtx { + session_id: &self.session_id, + turn_id: Some(turn_id), + point: MutationPoint::ContextLengthRejected, + recovery: Some(info), + cancellation: cancellation.clone(), + emitter: &emitter, + }, + ) + .await?; + } + validate_transcript_invariants(cursor.items) + }; + recovery_bound(mutation, Some(info), cancellation.clone()).await?; + if cancellation + .as_ref() + .is_some_and(TurnCancellation::is_cancelled) + { + return Err(LoopError::Cancelled); + } + if std::time::Instant::now() >= deadline { + return Err(LoopError::ContextRecoveryDeadlineExceeded); + } + if candidate == self.transcript { + return Ok(false); + } + self.transcript = candidate; + for observer in &self.transcript_observers { + observer.on_transcript_replaced(&self.session_id, &self.transcript); + } + Ok(true) + } + async fn continue_active_tool_round(&mut self) -> Result, LoopError> { let Some((presentation_turn_id, task_turn_id, cancellation)) = self.active_tool_round.as_ref().map(|active| { @@ -2363,122 +2494,187 @@ where } use tracing::Instrument; - let mut turn = match session - .begin_turn(request, cancellation.clone()) - .instrument(chat_span.clone()) + let mut recovery = None; + let (mut result, saw_tool_call, latest_usage) = 'generation: loop { + let mut request = request.clone(); + request.transcript = self.transcript.clone(); + let mut turn = match recovery_bound( + self.session + .as_mut() + .expect("session checked above") + .begin_turn(request, cancellation.clone()) + .instrument(chat_span.clone()), + recovery, + cancellation.clone(), + ) .await - { - Ok(turn) => turn, - Err(LoopError::Cancelled) => { - self.task_manager - .on_turn_interrupted(&turn_id) - .await - .map_err(|error| LoopError::Tool(ToolError::Internal(error.to_string())))?; - return self.finish_cancelled(turn_id, interrupted_assistant_items()); - } - Err(error) => return Err(error), - }; + { + Ok(turn) => turn, + Err(LoopError::Cancelled) => { + self.task_manager + .on_turn_interrupted(&turn_id) + .await + .map_err(|error| LoopError::Tool(ToolError::Internal(error.to_string())))?; + return self.finish_cancelled(turn_id, interrupted_assistant_items()); + } + Err(error) => { + match self + .recover_context(&error, &turn_id, cancellation.clone(), &mut recovery) + .await + { + Ok(true) => continue 'generation, + Ok(false) => return Err(error), + Err(LoopError::Cancelled) => { + self.task_manager + .on_turn_interrupted(&turn_id) + .await + .map_err(|error| { + LoopError::Tool(ToolError::Internal(error.to_string())) + })?; + return self.finish_cancelled(turn_id, interrupted_assistant_items()); + } + Err(error) => return Err(error), + } + } + }; - // begin_turn may apply per-turn routing. Sample the effective selection - // only after that work, while the chat span still wraps begin_turn. - let provider_name = - effective_provider_name(session.provider_name(), self.provider_name.as_deref()); - if let Some(provider) = &provider_name { - chat_span.record("gen_ai.provider.name", provider.as_str()); - tracing::Span::current().record("gen_ai.provider.name", provider.as_str()); - } - match session.model_name() { - Some(model) => { - chat_span.record("gen_ai.request.model", model); - chat_span.record("otel.name", format!("chat {model}").as_str()); + // begin_turn may apply per-turn routing. Sample the effective selection + // only after that work, while the chat span still wraps begin_turn. + let session = self.session.as_ref().expect("session checked above"); + let provider_name = + effective_provider_name(session.provider_name(), self.provider_name.as_deref()); + if let Some(provider) = &provider_name { + chat_span.record("gen_ai.provider.name", provider.as_str()); + tracing::Span::current().record("gen_ai.provider.name", provider.as_str()); } - None => { - chat_span.record("otel.name", "chat"); + match session.model_name() { + Some(model) => { + chat_span.record("gen_ai.request.model", model); + chat_span.record("otel.name", format!("chat {model}").as_str()); + } + None => { + chat_span.record("otel.name", "chat"); + } } - } - let mut saw_tool_call = false; - let mut finished_result = None; - let mut latest_usage = None; - let mut streamed_content = cancellation - .is_some() - .then(StreamedAssistantContent::default); - - while let Some(event) = match turn - .next_event(cancellation.clone()) - .instrument(chat_span.clone()) + let mut emitted_output = false; + let mut saw_tool_call = false; + let mut finished_result = None; + let mut latest_usage = None; + let mut streamed_content = cancellation + .is_some() + .then(StreamedAssistantContent::default); + + while let Some(event) = match recovery_bound( + turn.next_event(cancellation.clone()) + .instrument(chat_span.clone()), + recovery, + cancellation.clone(), + ) .await - { - Ok(event) => event, - Err(LoopError::Cancelled) => { - self.task_manager - .on_turn_interrupted(&turn_id) - .await - .map_err(|error| LoopError::Tool(ToolError::Internal(error.to_string())))?; - return self.finish_cancelled( - turn_id, - interrupted_stream_items(streamed_content.as_ref()), - ); - } - Err(error) => return Err(error), - } { - let attempt_superseded = matches!(event, ModelTurnEvent::ResponseAttemptSuperseded); - if attempt_superseded { - saw_tool_call = false; - latest_usage = None; - if let Some(content) = &mut streamed_content { - content.reset(); - } - self.emit(AgentEvent::ResponseAttemptSuperseded); - } - if cancellation - .as_ref() - .is_some_and(TurnCancellation::is_cancelled) { - turn.on_cancelled(); - self.task_manager - .on_turn_interrupted(&turn_id) - .await - .map_err(|error| LoopError::Tool(ToolError::Internal(error.to_string())))?; - return self.finish_cancelled( - turn_id, - interrupted_stream_items(streamed_content.as_ref()), - ); - } - if attempt_superseded { - continue; - } - match event { - ModelTurnEvent::Delta(delta) => { - if let Some(content) = &mut streamed_content { - content.apply_delta(&delta); - } - self.emit(AgentEvent::ContentDelta(delta)); + Ok(event) => event, + Err(LoopError::Cancelled) => { + self.task_manager + .on_turn_interrupted(&turn_id) + .await + .map_err(|error| LoopError::Tool(ToolError::Internal(error.to_string())))?; + return self.finish_cancelled( + turn_id, + interrupted_stream_items(streamed_content.as_ref()), + ); } - ModelTurnEvent::Usage(usage) => { - latest_usage = Some(usage.clone()); - self.emit(AgentEvent::UsageUpdated(usage)); + Err(error) => { + drop(turn); + if emitted_output { + return Err(error); + } + match self + .recover_context(&error, &turn_id, cancellation.clone(), &mut recovery) + .await + { + Ok(true) => continue 'generation, + Ok(false) => return Err(error), + Err(LoopError::Cancelled) => { + self.task_manager + .on_turn_interrupted(&turn_id) + .await + .map_err(|error| { + LoopError::Tool(ToolError::Internal(error.to_string())) + })?; + return self.finish_cancelled(turn_id, interrupted_assistant_items()); + } + Err(error) => return Err(error), + } } - ModelTurnEvent::ToolCall(call) => { - saw_tool_call = true; + } { + // Every delta and tool-call event is output, including empty deltas. + // Supersession never resets this guard. + emitted_output |= matches!( + event, + ModelTurnEvent::Delta(_) + | ModelTurnEvent::ToolCall(_) + | ModelTurnEvent::Finished(_) + ); + let attempt_superseded = matches!(event, ModelTurnEvent::ResponseAttemptSuperseded); + if attempt_superseded { + saw_tool_call = false; + latest_usage = None; if let Some(content) = &mut streamed_content { - content.commit_tool_call(&call); + content.reset(); } - self.emit(AgentEvent::ToolCallRequested(call)); + self.emit(AgentEvent::ResponseAttemptSuperseded); + } + if cancellation + .as_ref() + .is_some_and(TurnCancellation::is_cancelled) + { + turn.on_cancelled(); + self.task_manager + .on_turn_interrupted(&turn_id) + .await + .map_err(|error| LoopError::Tool(ToolError::Internal(error.to_string())))?; + return self.finish_cancelled( + turn_id, + interrupted_stream_items(streamed_content.as_ref()), + ); } - ModelTurnEvent::ResponseAttemptSuperseded => { - unreachable!("response-attempt supersession is handled before cancellation") + if attempt_superseded { + continue; } - ModelTurnEvent::Finished(result) => { - finished_result = Some(result); - break; + match event { + ModelTurnEvent::Delta(delta) => { + if let Some(content) = &mut streamed_content { + content.apply_delta(&delta); + } + self.emit(AgentEvent::ContentDelta(delta)); + } + ModelTurnEvent::Usage(usage) => { + latest_usage = Some(usage.clone()); + self.emit(AgentEvent::UsageUpdated(usage)); + } + ModelTurnEvent::ToolCall(call) => { + saw_tool_call = true; + if let Some(content) = &mut streamed_content { + content.commit_tool_call(&call); + } + self.emit(AgentEvent::ToolCallRequested(call)); + } + ModelTurnEvent::ResponseAttemptSuperseded => { + unreachable!("response-attempt supersession is handled before cancellation") + } + ModelTurnEvent::Finished(result) => { + finished_result = Some(result); + break; + } } } - } - let mut result = finished_result.ok_or_else(|| { - LoopError::Provider("model turn ended without a Finished event".into()) - })?; + let result = finished_result.ok_or_else(|| { + LoopError::Provider("model turn ended without a Finished event".into()) + })?; + break 'generation (result, saw_tool_call, latest_usage); + }; result.usage = merge_usage(result.usage, latest_usage); if let Some(model) = &result.model { chat_span.record("gen_ai.response.model", model.as_str()); @@ -4819,6 +5015,9 @@ pub enum LoopError { /// The current turn was cancelled via the [`CancellationHandle`]. #[error("turn cancelled")] Cancelled, + /// The absolute context-recovery budget expired. + #[error("context recovery deadline exceeded")] + ContextRecoveryDeadlineExceeded, /// An error originating from the model provider. #[error("provider error: {0}")] Provider(String), @@ -4833,7 +5032,63 @@ pub enum LoopError { Unsupported(String), } +/// Ordinary inference remains unbounded; only recovery has an absolute budget. +async fn recovery_bound( + future: impl std::future::Future>, + recovery: Option, + cancellation: Option, +) -> Result { + let Some(info) = recovery else { + return future.await; + }; + let cancelled = async { + match cancellation { + Some(cancellation) => cancellation.cancelled().await, + None => std::future::pending::<()>().await, + } + }; + use std::future::Future; + use std::task::Poll; + + let mut cancelled = std::pin::pin!(cancelled); + let mut deadline = std::pin::pin!(futures_timer::Delay::new( + info.deadline + .saturating_duration_since(std::time::Instant::now()), + )); + let mut future = std::pin::pin!(future); + std::future::poll_fn(|cx| { + // Cancellation wins over timeout, and both win over model/mutator + // completion. Check the clock as well as the timer so an expired + // deadline never depends on the timer thread having woken us yet. + if cancelled.as_mut().poll(cx).is_ready() { + return Poll::Ready(Err(LoopError::Cancelled)); + } + if std::time::Instant::now() >= info.deadline || deadline.as_mut().poll(cx).is_ready() { + return Poll::Ready(Err(LoopError::ContextRecoveryDeadlineExceeded)); + } + future.as_mut().poll(cx) + }) + .await +} + impl LoopError { + /// True only for structured, confirmed upstream context-length rejection. + /// Only the exact typed code is accepted; error types, messages, generic + /// client errors, cancellation, and replay-unsafe failures are ignored. + pub fn is_context_length_exceeded(&self) -> bool { + self.provider_failure().is_some_and(|failure| { + failure.upstream.code == UpstreamErrorKind::ContextLengthExceeded + && !matches!( + failure.reason, + ProviderFailureReason::Cancelled | ProviderFailureReason::ReplayUnsafe + ) + && !matches!( + failure.last_attempt_reason, + Some(ProviderFailureReason::Cancelled | ProviderFailureReason::ReplayUnsafe) + ) + }) + } + /// Returns structured provider metadata without parsing a rendered error. pub fn provider_failure(&self) -> Option<&ProviderFailure> { match self { @@ -4927,6 +5182,7 @@ fn validate_transcript_invariants(transcript: &[Item]) -> Result<(), LoopError> #[cfg(test)] mod tests { + include!("context_recovery_tests.rs"); use std::collections::VecDeque; use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; use std::sync::{Arc as StdArc, Mutex as StdMutex}; diff --git a/crates/agentkit-loop/src/retry.rs b/crates/agentkit-loop/src/retry.rs index d92a55a..a563cb9 100644 --- a/crates/agentkit-loop/src/retry.rs +++ b/crates/agentkit-loop/src/retry.rs @@ -34,6 +34,8 @@ pub enum UpstreamErrorKind { PermissionDenied, InsufficientQuota, ContentPolicyViolation, + /// Explicit provider rejection because the input exceeds the context window. + ContextLengthExceeded, #[default] #[serde(other)] Unknown, @@ -146,6 +148,19 @@ mod tests { use super::*; use crate::AgentEvent; + #[test] + fn context_length_kind_uses_stable_wire_name() { + let kind = UpstreamErrorKind::ContextLengthExceeded; + assert_eq!( + serde_json::to_string(&kind).unwrap(), + "\"context_length_exceeded\"" + ); + assert_eq!( + serde_json::from_str::("\"context_length_exceeded\"").unwrap(), + kind + ); + } + #[test] fn event_roundtrip_and_legacy_shape_remain_compatible() { let old = r#"{"RunStarted":{"session_id":"session"}}"#; diff --git a/crates/agentkit-provider-anthropic/src/lib.rs b/crates/agentkit-provider-anthropic/src/lib.rs index bd328cc..5fbedbc 100644 --- a/crates/agentkit-provider-anthropic/src/lib.rs +++ b/crates/agentkit-provider-anthropic/src/lib.rs @@ -492,6 +492,9 @@ impl ModelSession for AnthropicSession { }; if !status.is_success() { + if let Some(error) = context_length_error(&body_text, Some(status.as_u16())) { + return Err(error); + } return Err(LoopError::Provider(format!( "Anthropic request failed with status {status}: {body_text}" ))); @@ -703,6 +706,54 @@ mod tests { use super::*; + struct ContextRejection; + + #[async_trait] + impl HttpClient for ContextRejection { + async fn execute(&self, request: HttpRequest) -> Result { + Ok(HttpResponse::new( + StatusCode::BAD_REQUEST, + HeaderMap::new(), + request.url, + Box::pin(stream::once(async { + Ok(Bytes::from_static(br#"{"type":"error","error":{"type":"invalid_request_error","message":"prompt is too long: SECRET-PROMPT"}}"#)) + })), + )) + } + } + + #[tokio::test] + async fn http_context_rejection_is_typed_for_streaming_and_buffered_requests() { + for streaming in [false, true] { + let config = AnthropicConfig::new("key", "model", 1024) + .unwrap() + .with_streaming(streaming); + let adapter = + AnthropicAdapter::with_client(config, Http::from_arc(Arc::new(ContextRejection))); + let mut session = adapter + .start_session(SessionConfig::new("session")) + .await + .unwrap(); + let request = TurnRequest { + session_id: agentkit_core::SessionId::new("session"), + turn_id: agentkit_core::TurnId::new("turn"), + transcript: vec![], + available_tools: vec![], + cache: None, + metadata: agentkit_core::MetadataMap::new(), + }; + let error = match session.begin_turn(request, None).await { + Err(error) => error, + Ok(_) => panic!("expected rejection"), + }; + assert_eq!( + error.provider_failure().unwrap().upstream.code, + agentkit_loop::UpstreamErrorKind::ContextLengthExceeded + ); + assert!(!format!("{error:?} {error}").contains("SECRET-PROMPT")); + } + } + struct AlwaysUnauthorized { authorizations: std::sync::Mutex>, } @@ -1081,3 +1132,69 @@ mod tests { assert!(matches!(err, LoopError::Cancelled)); } } + +/// Recognize only Anthropic's structured invalid-request prompt-size rejection. +/// Message text is inspected here but never retained in the typed failure. +fn context_length_error(body: &str, status: Option) -> Option { + use agentkit_loop::{ + ProviderClassification, ProviderFailure, ProviderFailureReason, ProviderRoute, + RetryAccounting, UpstreamErrorKind, + }; + if status.is_some_and(|status| status != 400) { + return None; + } + let value: serde_json::Value = serde_json::from_str(body).ok()?; + if value.get("type")?.as_str()? != "error" + || value.pointer("/error/type")?.as_str()? != "invalid_request_error" + || !value + .pointer("/error/message")? + .as_str()? + .starts_with("prompt is too long: ") + { + return None; + } + Some(LoopError::ProviderFailure(Box::new(ProviderFailure { + route: ProviderRoute::Unknown, + reason: ProviderFailureReason::InvalidRequest, + last_attempt_reason: Some(ProviderFailureReason::InvalidRequest), + upstream: ProviderClassification { + error_type: UpstreamErrorKind::InvalidRequestError, + code: UpstreamErrorKind::ContextLengthExceeded, + http_status: status, + }, + accounting: RetryAccounting::default(), + }))) +} + +#[cfg(test)] +mod context_length_tests { + use super::*; + + #[test] + fn only_structured_prompt_too_long_is_classified_without_payload() { + let body = r#"{"type":"error","error":{"type":"invalid_request_error","message":"prompt is too long: SECRET"}}"#; + let error = context_length_error(body, Some(400)).unwrap(); + assert_eq!( + error.provider_failure().unwrap().upstream.code, + agentkit_loop::UpstreamErrorKind::ContextLengthExceeded + ); + assert_eq!( + error.provider_failure().unwrap().route, + agentkit_loop::ProviderRoute::Unknown + ); + assert!(!format!("{error:?} {error}").contains("SECRET")); + assert!(context_length_error(body, None).is_some()); + for status in [401, 403, 429, 500] { + assert!(context_length_error(body, Some(status)).is_none()); + } + for body in [ + "prompt is too long: SECRET", + "{}", + r#"{"type":"error","error":{"type":"invalid_request_error","message":"too many tokens"}}"#, + r#"{"type":"error","error":{"type":"authentication_error","message":"prompt is too long: SECRET"}}"#, + r#"{"type":"message","error":{"type":"invalid_request_error","message":"prompt is too long: SECRET"}}"#, + ] { + assert!(context_length_error(body, Some(400)).is_none()); + } + } +} diff --git a/crates/agentkit-provider-anthropic/src/stream.rs b/crates/agentkit-provider-anthropic/src/stream.rs index d8143d0..b5bb5c0 100644 --- a/crates/agentkit-provider-anthropic/src/stream.rs +++ b/crates/agentkit-provider-anthropic/src/stream.rs @@ -103,10 +103,11 @@ impl EventTranslator { "content_block_stop" => self.on_block_stop(&event.data), "message_delta" => self.on_message_delta(&event.data), "message_stop" => self.on_message_stop(), - "error" => Err(LoopError::Provider(format!( - "Anthropic stream error: {}", - event.data - ))), + "error" => Err( + crate::context_length_error(&event.data, None).unwrap_or_else(|| { + LoopError::Provider(format!("Anthropic stream error: {}", event.data)) + }), + ), "ping" | "" => Ok(Vec::new()), _ => Ok(Vec::new()), } @@ -752,6 +753,19 @@ mod tests { } } + #[test] + fn prompt_size_stream_error_is_typed_and_payload_free() { + let mut decoder = SseDecoder::new(); + let mut translator = EventTranslator::new(); + let events = decoder.feed("event: error\ndata: {\"type\":\"error\",\"error\":{\"type\":\"invalid_request_error\",\"message\":\"prompt is too long: SECRET-PROMPT\"}}\n\n"); + let error = translator.handle(&events[0]).unwrap_err(); + assert_eq!( + error.provider_failure().unwrap().upstream.code, + agentkit_loop::UpstreamErrorKind::ContextLengthExceeded + ); + assert!(!format!("{error:?} {error}").contains("SECRET-PROMPT")); + } + #[test] fn error_event_becomes_provider_error() { let mut decoder = SseDecoder::new(); diff --git a/crates/agentkit-provider-openai/src/lib.rs b/crates/agentkit-provider-openai/src/lib.rs index 0916972..702c9b0 100644 --- a/crates/agentkit-provider-openai/src/lib.rs +++ b/crates/agentkit-provider-openai/src/lib.rs @@ -356,6 +356,35 @@ impl CompletionsProvider for OpenAIProvider { Ok(()) } + fn preprocess_response( + &self, + status: agentkit_http::StatusCode, + body: &str, + ) -> Result<(), LoopError> { + use agentkit_loop::{ + ProviderClassification, ProviderFailure, ProviderFailureReason, ProviderRoute, + RetryAccounting, UpstreamErrorKind, + }; + if status == agentkit_http::StatusCode::BAD_REQUEST + && let Ok(value) = serde_json::from_str::(body) + && value.pointer("/error/code").and_then(Value::as_str) + == Some("context_length_exceeded") + { + return Err(LoopError::ProviderFailure(Box::new(ProviderFailure { + route: ProviderRoute::Unknown, + reason: ProviderFailureReason::InvalidRequest, + last_attempt_reason: Some(ProviderFailureReason::InvalidRequest), + upstream: ProviderClassification { + code: UpstreamErrorKind::ContextLengthExceeded, + http_status: Some(status.as_u16()), + ..ProviderClassification::default() + }, + accounting: RetryAccounting::default(), + }))); + } + Ok(()) + } + fn postprocess_response( &self, _usage: &mut Option, @@ -456,6 +485,41 @@ pub enum OpenAIError { #[cfg(test)] mod tests { + #[test] + fn completions_context_rejection_requires_exact_structured_code() { + use agentkit_http::StatusCode; + use agentkit_loop::UpstreamErrorKind; + let provider = OpenAIProvider::from(OpenAIConfig::new("key", "model")); + let body = r#"{"error":{"code":"context_length_exceeded","message":"SECRET-PROMPT"}}"#; + let error = provider + .preprocess_response(StatusCode::BAD_REQUEST, body) + .unwrap_err(); + assert_eq!( + error.provider_failure().unwrap().upstream.code, + UpstreamErrorKind::ContextLengthExceeded + ); + assert!(!format!("{error:?} {error}").contains("SECRET-PROMPT")); + for status in [ + StatusCode::UNAUTHORIZED, + StatusCode::FORBIDDEN, + StatusCode::TOO_MANY_REQUESTS, + ] { + assert!(provider.preprocess_response(status, body).is_ok()); + } + for body in [ + "context_length_exceeded", + r#"{"error":{"type":"invalid_request_error","message":"too many tokens"}}"#, + r#"{"error":{"code":"CONTEXT_LENGTH_EXCEEDED"}}"#, + r#"{"code":"context_length_exceeded"}"#, + ] { + assert!( + provider + .preprocess_response(StatusCode::BAD_REQUEST, body) + .is_ok() + ); + } + } + use agentkit_core::{MetadataMap, SessionId, TurnId}; use super::*; diff --git a/crates/agentkit-provider-openai/src/responses.rs b/crates/agentkit-provider-openai/src/responses.rs index 251f04c..3f0c298 100644 --- a/crates/agentkit-provider-openai/src/responses.rs +++ b/crates/agentkit-provider-openai/src/responses.rs @@ -2212,16 +2212,50 @@ async fn send_live_attempt( }); } if !status.is_success() { + let headers = retry_headers(response.headers()); + let mut upstream = ProviderClassification { + http_status: Some(status.as_u16()), + ..ProviderClassification::default() + }; + // Only inspect a bounded structured rejection body, never retain its text. + // Other status paths preserve their existing retry/authentication behavior. + if status == StatusCode::BAD_REQUEST { + let mut body = response.bytes_stream(); + let mut bytes = Zeroizing::new(Vec::new()); + let started = Instant::now(); + loop { + let Some(remaining) = Duration::from_secs(2).checked_sub(started.elapsed()) else { + break; + }; + let chunk = cancellable( + run_bounded_http( + async { futures_util::StreamExt::next(&mut body).await.transpose() }, + Some(remaining), + context.deadline.as_ref(), + "context rejection body", + ), + cancellation, + ) + .await + .map_err(nonretryable)?; + match chunk { + Ok(Some(chunk)) if bytes.len().saturating_add(chunk.len()) <= 64 * 1024 => { + bytes.extend_from_slice(&chunk); + } + Ok(None) => { + if let Ok(value) = serde_json::from_slice::(&bytes) { + upstream = retry::http_classification(&value, status.as_u16()); + } + break; + } + _ => break, + } + } + } return Err(AttemptFailure { - error: Box::new(provider_error( - ProviderFailureReason::HttpStatus, - ProviderClassification { - http_status: Some(status.as_u16()), - ..ProviderClassification::default() - }, - )), + error: Box::new(provider_error(ProviderFailureReason::HttpStatus, upstream)), retryable: retryable_response_status(status, context.config.profile), - headers: retry_headers(response.headers()), + headers, }); } if response @@ -2572,7 +2606,8 @@ impl ResponsesSseDecoder { classification.http_status = Some(status.as_u16()); *failure.error = provider_error(ProviderFailureReason::ResponseFailed, classification); - failure.retryable = retryable_response_status(status, self.state.profile); + failure.retryable = classification.code != UpstreamErrorKind::ContextLengthExceeded + && retryable_response_status(status, self.state.profile); } if let Some(headers) = value.get("headers").and_then(Value::as_object) { let mut parsed = HeaderMap::new(); @@ -3841,6 +3876,11 @@ fn protocol_failure(message: &str) -> AttemptFailure { } fn stream_failure_retryable(profile: OpenAIResponsesProfile, value: &Value, kind: &str) -> bool { + // A confirmed context rejection cannot succeed by resending the same input, + // including private code-only envelopes with no HTTP status or error type. + if stream_classification(value, kind).code == UpstreamErrorKind::ContextLengthExceeded { + return false; + } let error = if kind == "response.failed" { value.pointer("/response/error").unwrap_or(&Value::Null) } else { diff --git a/crates/agentkit-provider-openai/src/responses/retry.rs b/crates/agentkit-provider-openai/src/responses/retry.rs index 61f7da1..5e3a299 100644 --- a/crates/agentkit-provider-openai/src/responses/retry.rs +++ b/crates/agentkit-provider-openai/src/responses/retry.rs @@ -178,29 +178,114 @@ pub(super) fn upstream_kind(value: Option<&Value>) -> UpstreamErrorKind { } } +// Context overflow is an allowlisted code, not a free-form type or message. +fn upstream_code(value: Option<&Value>) -> UpstreamErrorKind { + if value.and_then(Value::as_str) == Some("context_length_exceeded") { + UpstreamErrorKind::ContextLengthExceeded + } else { + upstream_kind(value) + } +} + +// Conflicting authentication/throttling/status signals do not authorize recovery. +fn context_consistent(mut classification: ProviderClassification) -> ProviderClassification { + if classification.code == UpstreamErrorKind::ContextLengthExceeded + && (classification + .http_status + .is_some_and(|status| status != 400) + || matches!( + classification.error_type, + UpstreamErrorKind::AuthenticationError + | UpstreamErrorKind::InvalidApiKey + | UpstreamErrorKind::InvalidAuthentication + | UpstreamErrorKind::Unauthorized + | UpstreamErrorKind::PermissionDenied + | UpstreamErrorKind::RateLimitError + | UpstreamErrorKind::RateLimitExceeded + | UpstreamErrorKind::InsufficientQuota + )) + { + classification.code = UpstreamErrorKind::Unknown; + } + classification +} + +pub(super) fn http_classification(value: &Value, status: u16) -> ProviderClassification { + let error = value.get("error").unwrap_or(&Value::Null); + context_consistent(ProviderClassification { + error_type: upstream_kind(error.get("type")), + code: upstream_code(error.get("code")), + http_status: Some(status), + }) +} + pub(super) fn stream_classification(value: &Value, kind: &str) -> ProviderClassification { let error = if kind == "response.failed" { value.pointer("/response/error").unwrap_or(&Value::Null) } else { value.get("error").unwrap_or(value) }; - ProviderClassification { + context_consistent(ProviderClassification { error_type: upstream_kind(error.get("type")), - code: upstream_kind(error.get("code").or_else(|| value.get("code"))), + code: upstream_code(error.get("code").or_else(|| value.get("code"))), http_status: error .get("status") .or_else(|| value.get("status")) + .or_else(|| value.get("status_code")) .or_else(|| value.pointer("/response/status_code")) .and_then(Value::as_u64) .filter(|status| (100..=599).contains(status)) .map(|status| status as u16), - } + }) } #[cfg(test)] mod tests { use super::*; + #[test] + fn contradictory_context_signals_do_not_authorize_recovery() { + for value in [ + json!({"status":401,"error":{"code":"context_length_exceeded"}}), + json!({"status":429,"error":{"code":"context_length_exceeded"}}), + json!({"status_code":401,"error":{"code":"context_length_exceeded"}}), + json!({"error":{"type":"authentication_error","code":"context_length_exceeded"}}), + json!({"error":{"type":"rate_limit_error","code":"context_length_exceeded"}}), + json!({"error":{"type":"context_length_exceeded","message":"context_length_exceeded"}}), + ] { + assert_ne!( + stream_classification(&value, "error").code, + UpstreamErrorKind::ContextLengthExceeded + ); + } + } + + #[test] + fn structured_stream_errors_require_exact_context_code() { + for kind in ["response.failed", "error"] { + for (code, expected) in [ + ("context_length_exceeded", true), + ("invalid_request_error", false), + ("rate_limit_exceeded", false), + ("invalid_api_key", false), + ("CONTEXT_LENGTH_EXCEEDED", false), + ] { + let error = json!({"type":"invalid_request_error", "code":code, "message":"SECRET-PROMPT context_length_exceeded"}); + let value = if kind == "response.failed" { + json!({"response":{"error":error}}) + } else { + json!({"error":error}) + }; + let classification = stream_classification(&value, kind); + assert_eq!( + classification.code == UpstreamErrorKind::ContextLengthExceeded, + expected + ); + assert!(!format!("{classification:?}").contains("SECRET-PROMPT")); + } + } + } + #[test] fn throttling_keeps_exact_accounting_and_never_suppresses_terminal() { let events = Arc::new(Mutex::new(Vec::new())); diff --git a/crates/agentkit-provider-openai/src/responses/tests/retry_observability.rs b/crates/agentkit-provider-openai/src/responses/tests/retry_observability.rs index 9fb3756..7add406 100644 --- a/crates/agentkit-provider-openai/src/responses/tests/retry_observability.rs +++ b/crates/agentkit-provider-openai/src/responses/tests/retry_observability.rs @@ -874,3 +874,97 @@ async fn invalid_local_headers_preflight_have_zero_attempts_and_no_retry() { assert!(!format!("{error:?} {error}").contains("SECRET")); } } + +#[tokio::test] +async fn http_context_length_rejection_is_typed_and_payload_free() { + for (status, body, expected) in [ + ( + StatusCode::BAD_REQUEST, + r#"{"error":{"type":"invalid_request_error","code":"context_length_exceeded","message":"SECRET-PROMPT"}}"#, + true, + ), + ( + StatusCode::BAD_REQUEST, + r#"{"error":{"type":"invalid_request_error","message":"context_length_exceeded SECRET-PROMPT"}}"#, + false, + ), + ( + StatusCode::BAD_REQUEST, + r#"{"code":"context_length_exceeded"}"#, + false, + ), + ( + StatusCode::BAD_REQUEST, + "context_length_exceeded SECRET-PROMPT", + false, + ), + ( + StatusCode::UNAUTHORIZED, + r#"{"error":{"code":"invalid_api_key"}}"#, + false, + ), + ( + StatusCode::TOO_MANY_REQUESTS, + r#"{"error":{"code":"rate_limit_exceeded"}}"#, + false, + ), + ] { + let (mut session, events, _) = observed( + OpenAIResponsesConfig::new("key", "model").with_resilience(policy(0)), + // Authentication refresh may resend once even with max_retries = 0. + vec![wire(status, body), wire(status, body)], + ) + .await; + let error = match session.begin_turn(request(), None).await { + Err(error) => error, + Ok(_) => panic!("expected rejection"), + }; + assert_eq!( + error.provider_failure().unwrap().upstream.code + == UpstreamErrorKind::ContextLengthExceeded, + expected + ); + assert!( + !format!("{error:?} {error} {:?}", events.lock().unwrap()).contains("SECRET-PROMPT") + ); + } +} + +#[tokio::test] +async fn code_only_context_rejection_never_retries_private_or_public_sse() { + const BODY: &str = r#"event: response.created +data: {"type":"response.created","sequence_number":1,"response":{"id":"private-id","model":"private-model"}} + +event: response.failed +data: {"type":"response.failed","sequence_number":2,"response":{"error":{"code":"context_length_exceeded","message":"SECRET-PROMPT"}}} + +"#; + for private in [false, true] { + let config = if private { + OpenAIResponsesConfig::chatgpt_private("model", Authentication::bearer("key")) + } else { + OpenAIResponsesConfig::new("key", "model") + } + .with_resilience(policy(2)); + let (mut session, events, client) = + observed(config, vec![wire(StatusCode::OK, BODY)]).await; + let mut turn = session.begin_turn(request(), None).await.unwrap(); + let error = finish(&mut turn).await.unwrap_err(); + let failure = error.provider_failure().unwrap(); + assert_eq!( + failure.upstream.code, + UpstreamErrorKind::ContextLengthExceeded + ); + assert_eq!(failure.reason, ProviderFailureReason::ResponseFailed); + assert_eq!(failure.accounting.attempts, 1); + assert_eq!(client.requests.lock().unwrap().len(), 1); + assert!( + !events + .lock() + .unwrap() + .iter() + .any(|event| matches!(event, ProviderRetryEvent::Scheduled(_))) + ); + assert!(!format!("{error:?} {error}").contains("SECRET-PROMPT")); + } +} diff --git a/crates/agentkit-provider-openai/src/responses/websocket/tests.rs b/crates/agentkit-provider-openai/src/responses/websocket/tests.rs index 8ce5dec..088be96 100644 --- a/crates/agentkit-provider-openai/src/responses/websocket/tests.rs +++ b/crates/agentkit-provider-openai/src/responses/websocket/tests.rs @@ -2398,3 +2398,51 @@ async fn slow_progressing_stream_outlives_progress_gap() { assert!(start.elapsed() > STALL * 2); peer.join().unwrap(); } + +#[tokio::test] +async fn websocket_context_length_rejection_is_typed_and_payload_free() { + for private in [false, true] { + let (endpoint, peer) = server(|listener| { + let mut ws = socket(&listener); + receive(&mut ws); + ws.send(Message::Text(r#"{"type":"error","error":{"code":"context_length_exceeded","message":"SECRET-PROMPT"}}"#.into())).unwrap(); + assert!( + matches!(ws.read(), Err(_) | Ok(Message::Close(_))), + "context rejection must not resend on the socket" + ); + no_connection_within(&listener, Duration::from_millis(100)); + }); + let mut cfg = if private { + private_config(&endpoint) + } else { + config(&endpoint, OpenAIResponsesTransport::WebSocket) + }; + cfg.resilience.as_mut().unwrap().max_retries = 2; + let adapter = OpenAIResponsesAdapter::new(cfg).unwrap(); + let mut session = adapter + .start_session(SessionConfig::new("session")) + .await + .unwrap(); + let mut turn = session.begin_turn(request(), None).await.unwrap(); + let error = drain(&mut turn).await.unwrap_err(); + let failure = error.provider_failure().unwrap(); + assert_eq!(failure.accounting.attempts, 1); + assert_eq!(failure.reason, ProviderFailureReason::ResponseFailed); + assert_eq!( + failure.upstream.code, + UpstreamErrorKind::ContextLengthExceeded + ); + assert_eq!( + failure.route, + if private { + ProviderRoute::OpenAiChatGptResponses + } else { + ProviderRoute::OpenAiResponses + } + ); + assert!(!format!("{error:?} {error}").contains("SECRET-PROMPT")); + drop(turn); + drop(session); + peer.join().unwrap(); + } +} diff --git a/docs/architecture.md b/docs/architecture.md index f55d433..cf9173d 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -149,3 +149,28 @@ For each registered mutator the loop: If any mutator in the pass dirtied the cursor, the loop validates transcript invariants (tool_use ↔ tool_result pairing) and hard-fails with `LoopError::Mutator` if a mutator left the transcript protocol-invalid. Compaction is the canonical mutator: `agentkit-compaction` provides a `Compactor` trait, `StrategyCompactor`, and trigger helpers (`item_count_trigger`, `context_window_trigger`) that wire into this seam through the `AgentBuilderCompactorExt::compactor` extension. Semantic compaction is provided through an injected `CompactionBackend` (e.g. `AgentCompactor`, which runs a nested sub-agent) rather than a built-in model client. + +## Recovering from context-length rejection + +Enable context recovery when your host can compact or otherwise reduce conversation history after a provider rejects an oversized request: + +```rust,ignore +use agentkit_loop::ContextRecoveryPolicy; +use std::time::Duration; + +let builder = Agent::builder() + .context_recovery(ContextRecoveryPolicy { + max_attempts: 3, + budget: Duration::from_secs(60), + }); +``` + +Register a `LoopMutator` that handles `MutationPoint::ContextLengthRejected`. At that point, `LoopCtx::recovery` provides the one-based attempt and an absolute deadline. The mutator chooses what history to summarize or shorten; AgentKit does not estimate token counts or implement a degradation policy. Preserve applicable instructions, the current user request, and complete tool-call/result relationships. + +Recovery is disabled unless configured. It starts only on a typed context-length rejection, never a generic HTTP 400, authentication error, or rate limit. The provider adapter must preserve that classification; `LoopError::is_context_length_exceeded()` exposes it without parsing rendered error messages. Supported classification includes OpenAI Responses HTTP/SSE/WebSocket errors, OpenAI chat-completions HTTP errors, and Anthropic HTTP/SSE prompt-length errors. Other compatible-provider routes may still return an unclassified error. + +Recovery mutators edit a staged transcript. AgentKit validates and commits the candidate only after the mutation pass succeeds within its cancellation and deadline bounds. An unchanged candidate stops recovery. Failed, cancelled, invalid, or timed-out mutation passes do not replace canonical history. Mutators must not persist staged candidates themselves: implement `TranscriptObserver::on_transcript_replaced` to update a persisted or mirrored transcript after commitment. + +The driver retries only model generation, preserving the logical turn and avoiding repeated tool execution. Recovery is refused after model response content has been exposed. All recovery attempts and replacement generation share the configured budget, starting at the first rejection; zero attempts or a zero budget disables recovery. The deadline is cooperative and cannot preempt blocking synchronous code. + +If the retry limit is reached or protected context cannot be reduced, surface the failure rather than discarding instructions or repeatedly submitting the same request. A successful summary does not guarantee the resumed request fits: instructions, tools, attachments, or the current request can themselves exceed the provider's limit. From b22ef30592bce0eb9a5de3da48246bea05f62f97 Mon Sep 17 00:00:00 2001 From: daniel Date: Mon, 19 Oct 2026 15:03:26 +0100 Subject: [PATCH 2/3] fix: preserve recovery input and observer contracts --- .../src/context_recovery_tests.rs | 182 +++++++++++++++++- crates/agentkit-loop/src/lib.rs | 73 ++++--- crates/agentkit-provider-anthropic/src/lib.rs | 169 +++++++++++++++- docs/architecture.md | 2 +- 4 files changed, 391 insertions(+), 35 deletions(-) diff --git a/crates/agentkit-loop/src/context_recovery_tests.rs b/crates/agentkit-loop/src/context_recovery_tests.rs index bb01b41..127d057 100644 --- a/crates/agentkit-loop/src/context_recovery_tests.rs +++ b/crates/agentkit-loop/src/context_recovery_tests.rs @@ -22,11 +22,15 @@ mod context_recovery { requests: StdArc>>, } struct RecoverySession { + cancellations: StdArc, + polls: StdArc, executions: StdArc, site: FailureSite, requests: StdArc>>, } struct RecoveryTurn { + cancellations: StdArc, + polls: StdArc, hold_when_empty: bool, events: VecDeque>, } @@ -58,6 +62,8 @@ mod context_recovery { type Session = RecoverySession; async fn start_session(&self, _: SessionConfig) -> Result { Ok(RecoverySession { + cancellations: StdArc::new(AtomicUsize::new(0)), + polls: StdArc::new(AtomicUsize::new(0)), executions: self.executions.clone(), site: self.site, requests: self.requests.clone(), @@ -87,6 +93,8 @@ mod context_recovery { events.push_back(Ok(event)); } return Ok(RecoveryTurn { + cancellations: self.cancellations.clone(), + polls: self.polls.clone(), events, hold_when_empty: false, }); @@ -96,6 +104,8 @@ mod context_recovery { std::future::pending::<()>().await; } return Ok(RecoveryTurn { + cancellations: self.cancellations.clone(), + polls: self.polls.clone(), hold_when_empty: matches!(self.site, FailureSite::SlowStream), events: if matches!(self.site, FailureSite::SlowStream) { VecDeque::from([Ok(ModelTurnEvent::Usage(Usage::default()))]) @@ -149,6 +159,8 @@ mod context_recovery { FailureSite::AfterTools => unreachable!(), }; Ok(RecoveryTurn { + cancellations: self.cancellations.clone(), + polls: self.polls.clone(), events, hold_when_empty: false, }) @@ -156,10 +168,14 @@ mod context_recovery { } #[async_trait] impl ModelTurn for RecoveryTurn { + fn on_cancelled(&mut self) { + self.cancellations.fetch_add(1, Ordering::SeqCst); + } async fn next_event( &mut self, _: Option, ) -> Result, LoopError> { + self.polls.fetch_add(1, Ordering::SeqCst); if self.hold_when_empty && self.events.is_empty() { std::future::pending::<()>().await; } @@ -171,6 +187,8 @@ mod context_recovery { enum Edit { Change, Noop, + Empty, + AssistantTail, Hang, Invalid, Fail, @@ -214,6 +232,8 @@ mod context_recovery { Item::text(ItemKind::System, format!("recovered {attempt}")), ); match self.edit { + Edit::Empty => cursor.clear(), + Edit::AssistantTail => **cursor = vec![Item::text(ItemKind::Assistant, "summary")], Edit::Hang => std::future::pending::<()>().await, Edit::Fail => return Err(LoopError::Mutator("failed".into())), Edit::Invalid => cursor.push(Item { @@ -300,7 +320,11 @@ mod context_recovery { .unwrap(); ( agent - .start(SessionConfig::new("recovery").with_cache(PromptCacheRequest::automatic())) + .start( + SessionConfig::new("recovery") + .with_cache(PromptCacheRequest::automatic()) + .with_response_attempt_supersession(), + ) .await .unwrap(), requests, @@ -518,4 +542,160 @@ mod context_recovery { "\"ContextLengthRejected\"" ); } + #[tokio::test] + async fn recovery_without_pending_input_rolls_back() { + for edit in [Edit::Empty, Edit::AssistantTail] { + let (mut driver, requests, calls, replacements, _) = + driver(FailureSite::Begin, edit, 2, Duration::from_secs(1)).await; + assert!( + driver + .next() + .await + .unwrap_err() + .is_context_length_exceeded() + ); + let requests = requests.lock().unwrap(); + assert_eq!(requests.len(), 1); + assert_eq!(calls.load(Ordering::SeqCst), 1); + assert_eq!(driver.transcript, requests[0].transcript); + assert!(replacements.lock().unwrap().is_empty()); + } + } + + #[tokio::test] + async fn cancellation_finalizes_held_provider_turn() { + let (mut driver, _, _, _, controller) = driver( + FailureSite::SlowStream, + Edit::Change, + 2, + Duration::from_secs(5), + ) + .await; + let session = driver.session.as_ref().unwrap(); + let cancellations = session.cancellations.clone(); + let polls = session.polls.clone(); + let cancel = async { + while polls.load(Ordering::SeqCst) < 2 { + tokio::task::yield_now().await; + } + controller.interrupt(); + }; + let (result, _) = timeout(Duration::from_secs(2), async { + tokio::join!(driver.next(), cancel) + }) + .await + .unwrap(); + assert!(matches!(result, Ok(LoopStep::Finished(_)))); + assert_eq!(cancellations.load(Ordering::SeqCst), 1); + } + + #[tokio::test] + async fn usage_recovery_requires_supersession_and_invalidates_only_on_commit() { + for (allow_supersession, edit, succeeds) in [ + (false, Edit::Change, false), + (true, Edit::Noop, false), + (true, Edit::Fail, false), + (true, Edit::Change, true), + ] { + let (mut driver, requests, calls, _, _) = + driver(FailureSite::EarlyStream, edit, 2, Duration::from_secs(1)).await; + driver.allow_supersession = allow_supersession; + let events = StdArc::new(StdMutex::new(Vec::new())); + driver.observers.push(StdArc::new(RecordingObserver { + events: events.clone(), + })); + let result = driver.next().await; + if succeeds { + let LoopStep::Finished(result) = result.unwrap() else { + panic!("expected finished") + }; + assert!( + result.usage.is_none(), + "rejected attempt usage must not leak into result" + ); + } else { + assert!(result.is_err()); + } + assert_eq!(requests.lock().unwrap().len(), if succeeds { 2 } else { 1 }); + if !allow_supersession { + assert_eq!(calls.load(Ordering::SeqCst), 0); + } + let sequence: Vec<_> = events + .lock() + .unwrap() + .iter() + .filter_map(|event| match event { + AgentEvent::UsageUpdated(_) => Some("usage"), + AgentEvent::ResponseAttemptSuperseded => Some("superseded"), + _ => None, + }) + .collect(); + assert_eq!( + sequence, + if succeeds { + vec!["usage", "superseded"] + } else { + vec!["usage"] + } + ); + } + } + + #[cfg(not(feature = "otel"))] + #[tokio::test] + async fn input_capture_tracks_each_dispatched_transcript() { + use tracing::instrument::WithSubscriber; + use tracing_subscriber::prelude::*; + struct InputRecords(StdArc>>); + impl tracing_subscriber::Layer for InputRecords { + fn on_record( + &self, + _: &tracing::span::Id, + values: &tracing::span::Record<'_>, + _: tracing_subscriber::layer::Context<'_, S>, + ) { + struct Visitor<'a>(&'a StdMutex>); + impl tracing::field::Visit for Visitor<'_> { + fn record_debug( + &mut self, + field: &tracing::field::Field, + value: &dyn std::fmt::Debug, + ) { + if field.name() == "gen_ai.input.messages" { + self.0.lock().unwrap().push(format!("{value:?}")); + } + } + } + values.record(&mut Visitor(&self.0)); + } + } + let records = StdArc::new(StdMutex::new(Vec::new())); + let subscriber = tracing_subscriber::registry().with(InputRecords(records.clone())); + async { + let (mut driver, requests, _, _, _) = + driver(FailureSite::Begin, Edit::Change, 2, Duration::from_secs(1)).await; + driver.telemetry = TelemetryConfig::default() + .with_input_messages(MessageCapture::new(10, 10000).unwrap()); + driver.next().await.unwrap(); + let records = records.lock().unwrap(); + let requests = requests.lock().unwrap(); + assert_eq!(records.len(), requests.len()); + for (record, request) in records.iter().zip(requests.iter()) { + assert_eq!( + *record, + format!( + "{:?}", + capture_messages( + &request.transcript, + driver.telemetry.input_messages.unwrap(), + CaptureOrder::NewestTail + ) + ) + ); + } + assert_ne!(records[0], records[1]); + } + .with_subscriber(subscriber) + .await; + } } diff --git a/crates/agentkit-loop/src/lib.rs b/crates/agentkit-loop/src/lib.rs index daeb3ed..08f6aa0 100644 --- a/crates/agentkit-loop/src/lib.rs +++ b/crates/agentkit-loop/src/lib.rs @@ -1439,6 +1439,7 @@ where /// Returns [`LoopError`] if the model adapter fails to create a session. pub async fn start(&self, config: SessionConfig) -> Result, LoopError> { let session_id = config.session_id.clone(); + let allow_supersession = config.consumer_capabilities.response_attempt_supersession; let default_cache = config.cache.clone(); let mut session = self.model.start_session(config).await?; if !self.observers.is_empty() { @@ -1464,6 +1465,7 @@ where telemetry: self.telemetry, default_cache, next_turn_cache: None, + allow_supersession, session: Some(session), tool_executor, task_manager: self.task_manager.clone(), @@ -1744,6 +1746,7 @@ pub struct LoopDriver where S: ModelSession, { + allow_supersession: bool, session_id: SessionId, observed_session_id: Arc, provider_name: Option, @@ -2191,7 +2194,7 @@ where if std::time::Instant::now() >= deadline { return Err(LoopError::ContextRecoveryDeadlineExceeded); } - if candidate == self.transcript { + if candidate == self.transcript || !transcript_has_pending_input(&candidate) { return Ok(false); } self.transcript = candidate; @@ -2426,17 +2429,11 @@ where let catalog_events = self.tool_executor.drain_catalog_events(); self.emit_tool_catalog_events(catalog_events); - let request = TurnRequest { - session_id: self.session_id.clone(), - turn_id: turn_id.clone(), - transcript: self.transcript.clone(), - available_tools: self.tool_executor.specs(), - cache: self - .next_turn_cache - .take() - .or_else(|| self.default_cache.clone()), - metadata: MetadataMap::new(), - }; + let available_tools = self.tool_executor.specs(); + let cache = self + .next_turn_cache + .take() + .or_else(|| self.default_cache.clone()); let session = self .session @@ -2466,14 +2463,6 @@ where "gen_ai.usage.output_tokens" = tracing::field::Empty, "gen_ai.usage.cost" = tracing::field::Empty, ); - if let Some(capture) = self.telemetry.input_messages { - record_string_array_attribute( - &chat_span, - "gen_ai.input.messages", - capture_messages(&request.transcript, capture, CaptureOrder::NewestTail), - ); - } - // Seed known identity before begin_turn so request setup failures and // cancellation remain attributable. Successful per-turn routing below // overwrites these values with the effective selection. @@ -2496,14 +2485,34 @@ where use tracing::Instrument; let mut recovery = None; let (mut result, saw_tool_call, latest_usage) = 'generation: loop { - let mut request = request.clone(); - request.transcript = self.transcript.clone(); + let request = TurnRequest { + session_id: self.session_id.clone(), + turn_id: turn_id.clone(), + transcript: self.transcript.clone(), + available_tools: available_tools.clone(), + cache: cache.clone(), + metadata: MetadataMap::new(), + }; let mut turn = match recovery_bound( - self.session - .as_mut() - .expect("session checked above") - .begin_turn(request, cancellation.clone()) - .instrument(chat_span.clone()), + async { + if let Some(capture) = self.telemetry.input_messages { + record_string_array_attribute( + &chat_span, + "gen_ai.input.messages", + capture_messages( + &request.transcript, + capture, + CaptureOrder::NewestTail, + ), + ); + } + self.session + .as_mut() + .expect("session checked above") + .begin_turn(request, cancellation.clone()) + .await + } + .instrument(chat_span.clone()), recovery, cancellation.clone(), ) @@ -2575,6 +2584,7 @@ where { Ok(event) => event, Err(LoopError::Cancelled) => { + turn.on_cancelled(); self.task_manager .on_turn_interrupted(&turn_id) .await @@ -2586,14 +2596,19 @@ where } Err(error) => { drop(turn); - if emitted_output { + if emitted_output || (latest_usage.is_some() && !self.allow_supersession) { return Err(error); } match self .recover_context(&error, &turn_id, cancellation.clone(), &mut recovery) .await { - Ok(true) => continue 'generation, + Ok(true) => { + if latest_usage.is_some() { + self.emit(AgentEvent::ResponseAttemptSuperseded); + } + continue 'generation; + } Ok(false) => return Err(error), Err(LoopError::Cancelled) => { self.task_manager diff --git a/crates/agentkit-provider-anthropic/src/lib.rs b/crates/agentkit-provider-anthropic/src/lib.rs index 5fbedbc..a18a8eb 100644 --- a/crates/agentkit-provider-anthropic/src/lib.rs +++ b/crates/agentkit-provider-anthropic/src/lib.rs @@ -301,6 +301,44 @@ impl ReplayRequest { .await } + /// Inspect only a complete, bounded HTTP 400 envelope. Failure to inspect is + /// not evidence of context overflow; never expose the rejection body. + async fn rejection_error(&self, response: HttpResponse) -> LoopError { + let status = response.status(); + if status == StatusCode::BAD_REQUEST { + let timeout = self + .resilience + .as_ref() + .and_then(|config| config.attempt_timeout) + .unwrap_or(Duration::from_secs(2)) + .min(Duration::from_secs(2)); + let inspect = async { + let mut body = response.bytes_stream(); + let mut bytes = Vec::new(); + while let Some(chunk) = next_body_chunk_bounded(&mut body, None, None).await? { + if bytes.len().saturating_add(chunk.len()) > 64 * 1024 { + return Ok(None); + } + bytes.extend_from_slice(&chunk); + } + Ok(std::str::from_utf8(&bytes) + .ok() + .and_then(|body| context_length_error(body, Some(400)))) + }; + if let Ok(Some(error)) = run_bounded( + inspect, + Some(timeout), + self.deadline.as_ref(), + "HTTP rejection body", + ) + .await + { + return error; + } + } + LoopError::Provider(format!("Anthropic request failed with status {status}")) + } + async fn open_response(&self) -> Result { loop { let authentication = self.authentication_attempt().await?; @@ -375,10 +413,7 @@ impl ReplayRequest { self.wait_before_retry(retry_number, None).await?; let response = self.open_response().await?; if !response.status().is_success() { - return Err(LoopError::Provider(format!( - "Anthropic stream replay failed with status {}", - response.status() - ))); + return Err(self.rejection_error(response).await); } let integrity = TruncatedStreamDetector::from_headers(response.headers()); Ok((response.bytes_stream(), integrity)) @@ -440,6 +475,9 @@ impl ModelSession for AnthropicSession { loop { let response = replay_request.open_response().await?; let status = response.status(); + if status == StatusCode::BAD_REQUEST { + return Err(replay_request.rejection_error(response).await); + } if config.streaming && status.is_success() { let integrity = TruncatedStreamDetector::from_headers(response.headers()); let idle_timeout = self @@ -754,6 +792,129 @@ mod tests { } } + struct ReplayRejection { + calls: AtomicUsize, + body: Vec, + pending: bool, + } + + #[async_trait] + impl HttpClient for ReplayRejection { + async fn execute(&self, request: HttpRequest) -> Result { + let first = self.calls.fetch_add(1, Ordering::SeqCst) == 0; + let body: BodyStream = if first { + Box::pin(stream::once(async { + Err(HttpError::body(std::io::Error::new( + std::io::ErrorKind::ConnectionReset, + "stream reset", + ))) + })) + } else if self.pending { + Box::pin(stream::pending()) + } else { + let bytes = Bytes::copy_from_slice(&self.body); + Box::pin(stream::once(async move { Ok(bytes) })) + }; + Ok(HttpResponse::new( + if first { + StatusCode::OK + } else { + StatusCode::BAD_REQUEST + }, + HeaderMap::new(), + request.url, + body, + )) + } + } + + async fn replay_rejection_turn( + body: Vec, + pending: bool, + budget: Duration, + ) -> (AnthropicTurn, Arc) { + let client = Arc::new(ReplayRejection { + calls: AtomicUsize::new(0), + body, + pending, + }); + let adapter = AnthropicAdapter::with_client( + AnthropicConfig::new("key", "model", 1024).unwrap(), + Http::from_arc(client.clone()), + ) + .with_resilience(ResilienceConfig { + max_retries: 2, + retry_budget: budget, + attempt_timeout: None, + stream_idle_timeout: None, + initial_backoff: Duration::ZERO, + max_backoff: Duration::ZERO, + }); + let mut session = adapter + .start_session(SessionConfig::new("session")) + .await + .unwrap(); + let request = TurnRequest { + session_id: agentkit_core::SessionId::new("session"), + turn_id: agentkit_core::TurnId::new("turn"), + transcript: vec![], + available_tools: vec![], + cache: None, + metadata: agentkit_core::MetadataMap::new(), + }; + (session.begin_turn(request, None).await.unwrap(), client) + } + + #[tokio::test] + async fn replay_http_context_rejection_is_typed_bounded_and_payload_free() { + let valid = br#"{"type":"error","error":{"type":"invalid_request_error","message":"prompt is too long: SECRET-PROMPT"}}"#; + let oversized = format!( + r#"{{"type":"error","error":{{"type":"invalid_request_error","message":"prompt is too long: {}"}}}}"#, + "x".repeat(64 * 1024) + ); + for (body, expected) in [ + (valid.to_vec(), true), + (br#"{"type":"error","error":{"type":"invalid_request_error","message":"SECRET-PROMPT too many tokens"}}"#.to_vec(), false), + (b"prompt is too long: SECRET-PROMPT".to_vec(), false), + (oversized.into_bytes(), false), + ] { + let (mut turn, client) = replay_rejection_turn(body, false, Duration::from_secs(10)).await; + let error = turn.next_event(None).await.unwrap_err(); + assert_eq!(error.is_context_length_exceeded(), expected); + assert!(!format!("{error:?} {error}").contains("SECRET-PROMPT")); + assert_eq!(client.calls.load(Ordering::SeqCst), 2); + } + } + + #[tokio::test] + async fn replay_rejection_body_respects_logical_deadline() { + let (mut turn, client) = + replay_rejection_turn(vec![], true, Duration::from_millis(50)).await; + let start = std::time::Instant::now(); + let error = turn.next_event(None).await.unwrap_err(); + assert!(!error.is_context_length_exceeded()); + assert!(start.elapsed() < Duration::from_secs(1)); + assert_eq!(client.calls.load(Ordering::SeqCst), 2); + } + + #[tokio::test] + async fn replay_rejection_body_is_cancellable() { + let (mut turn, client) = replay_rejection_turn(vec![], true, Duration::from_secs(10)).await; + let controller = CancellationController::new(); + let future = turn.next_event(Some(TurnCancellation::new(controller.handle()))); + futures_util::pin_mut!(future); + for _ in 0..100 { + assert!(futures_util::poll!(future.as_mut()).is_pending()); + if client.calls.load(Ordering::SeqCst) == 2 { + break; + } + resilience_sleep(Duration::from_millis(1)).await; + } + assert_eq!(client.calls.load(Ordering::SeqCst), 2); + controller.interrupt(); + assert!(matches!(future.await, Err(LoopError::Cancelled))); + } + struct AlwaysUnauthorized { authorizations: std::sync::Mutex>, } diff --git a/docs/architecture.md b/docs/architecture.md index cf9173d..4d1dc3a 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -171,6 +171,6 @@ Recovery is disabled unless configured. It starts only on a typed context-length Recovery mutators edit a staged transcript. AgentKit validates and commits the candidate only after the mutation pass succeeds within its cancellation and deadline bounds. An unchanged candidate stops recovery. Failed, cancelled, invalid, or timed-out mutation passes do not replace canonical history. Mutators must not persist staged candidates themselves: implement `TranscriptObserver::on_transcript_replaced` to update a persisted or mirrored transcript after commitment. -The driver retries only model generation, preserving the logical turn and avoiding repeated tool execution. Recovery is refused after model response content has been exposed. All recovery attempts and replacement generation share the configured budget, starting at the first rejection; zero attempts or a zero budget disables recovery. The deadline is cooperative and cannot preempt blocking synchronous code. +The driver retries only model generation, preserving the logical turn and avoiding repeated tool execution. Recovery is refused after model response content has been exposed. If a rejected attempt emitted usage, recovery also requires `SessionConfig::with_response_attempt_supersession()`: the host must discard that attempt's usage when it receives `ResponseAttemptSuperseded`, before handling replacement output. Without that capability, the rejection terminates without a recovery mutation. All recovery attempts and replacement generation share the configured budget, starting at the first rejection; zero attempts or a zero budget disables recovery. The deadline is cooperative and cannot preempt blocking synchronous code. If the retry limit is reached or protected context cannot be reduced, surface the failure rather than discarding instructions or repeatedly submitting the same request. A successful summary does not guarantee the resumed request fits: instructions, tools, attachments, or the current request can themselves exceed the provider's limit. From 891f43f4347c520e468b5895df193826f6698ab5 Mon Sep 17 00:00:00 2001 From: daniel Date: Mon, 19 Oct 2026 15:34:52 +0100 Subject: [PATCH 3/3] fix: isolate telemetry spans for recovered requests --- .../src/context_recovery_otel_harness.rs | 223 ++++++++++++++++++ .../src/context_recovery_tests.rs | 32 ++- crates/agentkit-loop/src/lib.rs | 149 +++++++----- 3 files changed, 346 insertions(+), 58 deletions(-) create mode 100644 crates/agentkit-loop/src/context_recovery_otel_harness.rs diff --git a/crates/agentkit-loop/src/context_recovery_otel_harness.rs b/crates/agentkit-loop/src/context_recovery_otel_harness.rs new file mode 100644 index 0000000..763ac9b --- /dev/null +++ b/crates/agentkit-loop/src/context_recovery_otel_harness.rs @@ -0,0 +1,223 @@ +// Appended to TRUE_OTEL_HARNESS, using its real tracing-opentelemetry layer +// and capturing OpenTelemetry Tracer. No SDK dependency or production hook. +struct RecoveryOtelAdapter { + requests: Arc>>, + reject_first: bool, + early_stream: bool, +} +struct RecoveryOtelSession { + requests: Arc>>, + reject_first: bool, + early_stream: bool, +} +struct RecoveryOtelTurn { + reject: bool, +} +struct RecoveryOtelMutator; + +fn context_rejection() -> LoopError { + LoopError::ProviderFailure(Box::new(agentkit_loop::ProviderFailure { + route: agentkit_loop::ProviderRoute::Unknown, + reason: agentkit_loop::ProviderFailureReason::InvalidRequest, + last_attempt_reason: None, + upstream: agentkit_loop::ProviderClassification { + code: agentkit_loop::UpstreamErrorKind::ContextLengthExceeded, + ..Default::default() + }, + accounting: Default::default(), + })) +} +#[async_trait] +impl ModelAdapter for RecoveryOtelAdapter { + type Session = RecoveryOtelSession; + async fn start_session(&self, _: SessionConfig) -> Result { + Ok(RecoveryOtelSession { + requests: self.requests.clone(), + reject_first: self.reject_first, + early_stream: self.early_stream, + }) + } +} +#[async_trait] +impl ModelSession for RecoveryOtelSession { + type Turn = RecoveryOtelTurn; + fn model_name(&self) -> Option<&str> { + Some("recovery-model") + } + fn provider_name(&self) -> Option<&str> { + Some("recovery-provider") + } + async fn begin_turn( + &mut self, + request: TurnRequest, + _: Option, + ) -> Result { + let reject = { + let mut requests = self.requests.lock().unwrap(); + requests.push(request); + self.reject_first && requests.len() == 1 + }; + if reject && !self.early_stream { + return Err(context_rejection()); + } + Ok(RecoveryOtelTurn { reject }) + } +} +#[async_trait] +impl ModelTurn for RecoveryOtelTurn { + async fn next_event( + &mut self, + _: Option, + ) -> Result, LoopError> { + if self.reject { + return Err(context_rejection()); + } + Ok(Some(ModelTurnEvent::Finished(ModelTurnResult { + finish_reason: FinishReason::Completed, + output_items: vec![Item::text(ItemKind::Assistant, "replacement result")], + usage: Some(Usage { + tokens: Some(TokenUsage::new(7, 3)), + cost: None, + metadata: MetadataMap::new(), + }), + metadata: MetadataMap::new(), + model: Some("response-model".into()), + response_id: Some("response-1".into()), + }))) + } +} +#[async_trait] +impl agentkit_loop::LoopMutator for RecoveryOtelMutator { + async fn mutate( + &self, + cursor: &mut agentkit_loop::TranscriptCursor<'_>, + ctx: agentkit_loop::LoopCtx<'_>, + ) -> Result<(), LoopError> { + if ctx.point == agentkit_loop::MutationPoint::ContextLengthRejected { + cursor.insert(0, Item::text(ItemKind::System, "recovery summary")); + } + Ok(()) + } +} +fn recovery_attempts_have_distinct_otel_spans() { + for (reject_first, early_stream) in [(true, false), (true, true), (false, false)] { + let tracer = MemoryTracer::default(); + let subscriber = tracing_subscriber::registry() + .with(tracing_opentelemetry::layer().with_tracer(tracer.clone())); + let requests = Arc::new(Mutex::new(Vec::new())); + let runtime = tokio::runtime::Builder::new_current_thread() + .build() + .unwrap(); + let result = tracing::subscriber::with_default(subscriber, || { + runtime.block_on(async { + let agent = Agent::builder() + .model(RecoveryOtelAdapter { + requests: requests.clone(), + reject_first, + early_stream, + }) + .input(vec![Item::text(ItemKind::User, "original request")]) + .context_recovery(agentkit_loop::ContextRecoveryPolicy { + max_attempts: 1, + budget: std::time::Duration::from_secs(5), + }) + .mutator(RecoveryOtelMutator) + .telemetry( + TelemetryConfig::default() + .with_input_messages(MessageCapture::new(10, 10000).unwrap()), + ) + .build() + .unwrap(); + let mut driver = agent + .start( + SessionConfig::new("recovery-otel") + .with_cache(agentkit_loop::PromptCacheRequest::automatic()), + ) + .await + .unwrap(); + driver.next().await.unwrap() + }) + }); + let LoopStep::Finished(result) = result else { + panic!("expected finished turn") + }; + assert_eq!(result.usage.unwrap().tokens.unwrap().input_tokens, 7); + let exported = tracer.exported.lock().unwrap(); + let chats: Vec<_> = exported + .iter() + .filter(|s| operation(s) == Some("chat")) + .collect(); + let requests = requests.lock().unwrap(); + assert_eq!(chats.len(), if reject_first { 2 } else { 1 }); + assert_eq!(chats.len(), requests.len()); + for (index, chat) in chats.iter().enumerate() { + // Count every occurrence, not a last-value lookup that hides duplicates. + assert_eq!( + chat.attributes + .iter() + .filter(|a| a.key.as_str() == "gen_ai.input.messages") + .count(), + 1 + ); + let input = json_array(chat, "gen_ai.input.messages"); + assert_eq!(input.len(), requests[index].transcript.len()); + assert!( + input + .last() + .unwrap() + .to_string() + .contains("original request") + ); + assert_eq!( + input + .iter() + .any(|item| item.to_string().contains("recovery summary")), + reject_first && index == 1 + ); + assert_eq!( + attr(chat, "gen_ai.conversation.id"), + Some(&Value::String("recovery-otel".into())) + ); + assert_eq!( + attr(chat, "gen_ai.provider.name"), + Some(&Value::String("recovery-provider".into())) + ); + assert_eq!( + attr(chat, "gen_ai.request.model"), + Some(&Value::String("recovery-model".into())) + ); + if reject_first && index == 0 { + assert_eq!( + attr(chat, "error.type"), + Some(&Value::String("context_length_exceeded".into())) + ); + assert!(matches!(chat.status, Status::Error { .. })); + assert!(attr(chat, "gen_ai.usage.input_tokens").is_none()); + assert!(attr(chat, "gen_ai.response.id").is_none()); + } else { + assert!(attr(chat, "error.type").is_none()); + assert_eq!( + attr(chat, "gen_ai.usage.input_tokens"), + Some(&Value::I64(7)) + ); + assert_eq!( + attr(chat, "gen_ai.usage.output_tokens"), + Some(&Value::I64(3)) + ); + assert_eq!( + attr(chat, "gen_ai.response.id"), + Some(&Value::String("response-1".into())) + ); + } + } + if reject_first { + assert_ne!( + chats[0].capture_id, chats[1].capture_id, + "each attempt must build a distinct OTel span" + ); + assert_eq!(requests[0].turn_id, requests[1].turn_id); + assert_eq!(requests[0].cache, requests[1].cache); + assert_eq!(requests[0].available_tools, requests[1].available_tools); + } + } +} diff --git a/crates/agentkit-loop/src/context_recovery_tests.rs b/crates/agentkit-loop/src/context_recovery_tests.rs index 127d057..839b36d 100644 --- a/crates/agentkit-loop/src/context_recovery_tests.rs +++ b/crates/agentkit-loop/src/context_recovery_tests.rs @@ -644,7 +644,33 @@ mod context_recovery { #[cfg(not(feature = "otel"))] #[tokio::test] async fn input_capture_tracks_each_dispatched_transcript() { - use tracing::instrument::WithSubscriber; + // Tracing callsite interest and max-level caches are process-global. + // Keep this subscriber assertion isolated from concurrently running + // tests that register/drop dispatchers for the same inference callsite. + const CHILD: &str = "AGENTKIT_TEST_RECOVERY_INPUT_CAPTURE_CHILD"; + if std::env::var_os(CHILD).is_none() { + let output = std::process::Command::new(std::env::current_exe().unwrap()) + .args([ + "--exact", + "tests::context_recovery::input_capture_tracks_each_dispatched_transcript", + "--nocapture", + "--test-threads=1", + ]) + .env(CHILD, "1") + .output() + .expect("launch isolated input capture regression"); + let stdout = String::from_utf8_lossy(&output.stdout); + let stderr = String::from_utf8_lossy(&output.stderr); + assert!( + output.status.success(), + "isolated regression failed:\n{stdout}\n{stderr}" + ); + assert!( + stdout.contains("1 passed"), + "child did not execute the regression:\n{stdout}" + ); + return; + } use tracing_subscriber::prelude::*; struct InputRecords(StdArc>>); impl tracing_subscriber::Layer for InputRecords { @@ -671,6 +697,9 @@ mod context_recovery { } let records = StdArc::new(StdMutex::new(Vec::new())); let subscriber = tracing_subscriber::registry().with(InputRecords(records.clone())); + // This process runs only this exact test. Install the subscriber before + // any driver work and keep it registered for the child's whole lifetime. + tracing::subscriber::set_global_default(subscriber).unwrap(); async { let (mut driver, requests, _, _, _) = driver(FailureSite::Begin, Edit::Change, 2, Duration::from_secs(1)).await; @@ -695,7 +724,6 @@ mod context_recovery { } assert_ne!(records[0], records[1]); } - .with_subscriber(subscriber) .await; } } diff --git a/crates/agentkit-loop/src/lib.rs b/crates/agentkit-loop/src/lib.rs index 08f6aa0..ee97be4 100644 --- a/crates/agentkit-loop/src/lib.rs +++ b/crates/agentkit-loop/src/lib.rs @@ -2435,56 +2435,60 @@ where .take() .or_else(|| self.default_cache.clone()); - let session = self - .session - .as_mut() - .ok_or_else(|| LoopError::InvalidState("model session is not available".into()))?; - - // Inference span per the OTel GenAI semantic conventions. It wraps the - // model request and the full event drain rather than just `begin_turn`, - // so attributes that streaming adapters only learn mid-stream (usage, - // stop reason, response identity) still land before the span closes. - // `otel.name` carries the dynamic `chat {model}` span name for - // OpenTelemetry bridges since tracing span names are static. - let chat_span = tracing::info_span!( - "chat", - "otel.name" = tracing::field::Empty, - "otel.kind" = "client", - "gen_ai.operation.name" = "chat", - "gen_ai.provider.name" = tracing::field::Empty, - "gen_ai.conversation.id" = %self.session_id, - "gen_ai.request.model" = tracing::field::Empty, - "gen_ai.response.model" = tracing::field::Empty, - "gen_ai.response.id" = tracing::field::Empty, - "gen_ai.response.finish_reasons" = tracing::field::Empty, - "gen_ai.input.messages" = tracing::field::Empty, - "gen_ai.output.messages" = tracing::field::Empty, - "gen_ai.usage.input_tokens" = tracing::field::Empty, - "gen_ai.usage.output_tokens" = tracing::field::Empty, - "gen_ai.usage.cost" = tracing::field::Empty, - ); - // Seed known identity before begin_turn so request setup failures and - // cancellation remain attributable. Successful per-turn routing below - // overwrites these values with the effective selection. - let initial_provider_name = - effective_provider_name(session.provider_name(), self.provider_name.as_deref()); - if let Some(provider) = &initial_provider_name { - chat_span.record("gen_ai.provider.name", provider.as_str()); - tracing::Span::current().record("gen_ai.provider.name", provider.as_str()); - } - match session.model_name() { - Some(model) => { - chat_span.record("gen_ai.request.model", model); - chat_span.record("otel.name", format!("chat {model}").as_str()); + use tracing::Instrument; + let mut recovery = None; + let (mut result, saw_tool_call, latest_usage, chat_span) = 'generation: loop { + let session = self + .session + .as_mut() + .ok_or_else(|| LoopError::InvalidState("model session is not available".into()))?; + + // One inference span per dispatched generation attempt. In particular, + // OTel array attributes append rather than replace, so recovery must + // not reuse a rejected attempt's span. This span wraps the + // model request and the full event drain rather than just `begin_turn`, + // so attributes that streaming adapters only learn mid-stream (usage, + // stop reason, response identity) still land before the span closes. + // `otel.name` carries the dynamic `chat {model}` span name for + // OpenTelemetry bridges since tracing span names are static. + let chat_span = tracing::info_span!( + "chat", + "otel.name" = tracing::field::Empty, + "otel.kind" = "client", + "otel.status_code" = tracing::field::Empty, + "error.type" = tracing::field::Empty, + "gen_ai.operation.name" = "chat", + "gen_ai.provider.name" = tracing::field::Empty, + "gen_ai.conversation.id" = %self.session_id, + "gen_ai.request.model" = tracing::field::Empty, + "gen_ai.response.model" = tracing::field::Empty, + "gen_ai.response.id" = tracing::field::Empty, + "gen_ai.response.finish_reasons" = tracing::field::Empty, + "gen_ai.input.messages" = tracing::field::Empty, + "gen_ai.output.messages" = tracing::field::Empty, + "gen_ai.usage.input_tokens" = tracing::field::Empty, + "gen_ai.usage.output_tokens" = tracing::field::Empty, + "gen_ai.usage.cost" = tracing::field::Empty, + ); + // Seed known identity before begin_turn so request setup failures and + // cancellation remain attributable. Successful per-turn routing below + // overwrites these values with the effective selection. + let initial_provider_name = + effective_provider_name(session.provider_name(), self.provider_name.as_deref()); + if let Some(provider) = &initial_provider_name { + chat_span.record("gen_ai.provider.name", provider.as_str()); + tracing::Span::current().record("gen_ai.provider.name", provider.as_str()); } - None => { - chat_span.record("otel.name", "chat"); + match session.model_name() { + Some(model) => { + chat_span.record("gen_ai.request.model", model); + chat_span.record("otel.name", format!("chat {model}").as_str()); + } + None => { + chat_span.record("otel.name", "chat"); + } } - } - use tracing::Instrument; - let mut recovery = None; - let (mut result, saw_tool_call, latest_usage) = 'generation: loop { let request = TurnRequest { session_id: self.session_id.clone(), turn_id: turn_id.clone(), @@ -2520,6 +2524,7 @@ where { Ok(turn) => turn, Err(LoopError::Cancelled) => { + record_inference_error(&chat_span, &LoopError::Cancelled); self.task_manager .on_turn_interrupted(&turn_id) .await @@ -2527,6 +2532,8 @@ where return self.finish_cancelled(turn_id, interrupted_assistant_items()); } Err(error) => { + record_inference_error(&chat_span, &error); + drop(chat_span); match self .recover_context(&error, &turn_id, cancellation.clone(), &mut recovery) .await @@ -2584,6 +2591,7 @@ where { Ok(event) => event, Err(LoopError::Cancelled) => { + record_inference_error(&chat_span, &LoopError::Cancelled); turn.on_cancelled(); self.task_manager .on_turn_interrupted(&turn_id) @@ -2595,6 +2603,8 @@ where ); } Err(error) => { + record_inference_error(&chat_span, &error); + drop(chat_span); drop(turn); if emitted_output || (latest_usage.is_some() && !self.allow_supersession) { return Err(error); @@ -2644,6 +2654,7 @@ where .as_ref() .is_some_and(TurnCancellation::is_cancelled) { + record_inference_error(&chat_span, &LoopError::Cancelled); turn.on_cancelled(); self.task_manager .on_turn_interrupted(&turn_id) @@ -2686,9 +2697,11 @@ where } let result = finished_result.ok_or_else(|| { - LoopError::Provider("model turn ended without a Finished event".into()) + let error = LoopError::Provider("model turn ended without a Finished event".into()); + record_inference_error(&chat_span, &error); + error })?; - break 'generation (result, saw_tool_call, latest_usage); + break 'generation (result, saw_tool_call, latest_usage, chat_span); }; result.usage = merge_usage(result.usage, latest_usage); if let Some(model) = &result.model { @@ -4232,7 +4245,8 @@ tracing-subscriber = "={tracing_subscriber_version}" assert_eq!(encoded, r#""C:\\Users\\name\\quoted\\\"dir""#); } - const TRUE_OTEL_HARNESS: &str = r#" + const TRUE_OTEL_HARNESS: &str = concat!( + r#" use std::borrow::Cow; use std::collections::VecDeque; use std::sync::{Arc, Mutex}; @@ -4245,14 +4259,15 @@ use opentelemetry::{Array, Context, KeyValue, Value}; use tracing_subscriber::layer::SubscriberExt; #[derive(Clone, Debug)] -struct Exported { name: String, attributes: Vec } +struct Exported { name: String, attributes: Vec, capture_id: usize, status: Status } #[derive(Clone, Default)] struct MemoryTracer { exported: Arc>> } -struct MemorySpan { name: String, attributes: Vec, exported: Arc>>, ended: bool } +struct MemorySpan { name: String, attributes: Vec, exported: Arc>>, ended: bool, capture_id: usize, status: Status } impl Tracer for MemoryTracer { type Span = MemorySpan; fn build_with_context(&self, builder: SpanBuilder, _: &Context) -> MemorySpan { - MemorySpan { name: builder.name.into_owned(), attributes: builder.attributes.unwrap_or_default(), exported: self.exported.clone(), ended: false } + static NEXT_CAPTURE: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0); + MemorySpan { name: builder.name.into_owned(), attributes: builder.attributes.unwrap_or_default(), exported: self.exported.clone(), ended: false, capture_id: NEXT_CAPTURE.fetch_add(1, std::sync::atomic::Ordering::Relaxed), status: Status::Unset } } } impl Span for MemorySpan { @@ -4260,13 +4275,13 @@ impl Span for MemorySpan { fn span_context(&self) -> &SpanContext { &SpanContext::NONE } fn is_recording(&self) -> bool { !self.ended } fn set_attribute(&mut self, attribute: KeyValue) { self.attributes.push(attribute); } - fn set_status(&mut self, _: Status) {} + fn set_status(&mut self, status: Status) { self.status = status; } fn update_name(&mut self, name: T) where T: Into> { self.name = name.into().into_owned(); } fn add_link(&mut self, _: SpanContext, _: Vec) {} fn end_with_timestamp(&mut self, _: SystemTime) { if !self.ended { self.ended = true; - self.exported.lock().unwrap().push(Exported { name: self.name.clone(), attributes: self.attributes.clone() }); + self.exported.lock().unwrap().push(Exported { name: self.name.clone(), attributes: self.attributes.clone(), capture_id: self.capture_id, status: self.status.clone() }); } } } @@ -4370,6 +4385,7 @@ fn run(adapter: ScriptedAdapter) -> (Vec, agentkit_loop::TurnResult) { } fn main() { + recovery_attempts_have_distinct_otel_spans(); let (spans, result) = run(ScriptedAdapter { adapter_provider: "adapter", final_usage: true, overflow: true, begin_mode: BeginMode::Success }); let chat = spans.iter().find(|s| operation(s) == Some("chat")).unwrap(); let agent = spans.iter().find(|s| operation(s) == Some("invoke_agent")).unwrap(); @@ -4418,12 +4434,18 @@ fn main() { } let chat = spans.iter().find(|s| operation(s) == Some("chat")).unwrap(); let agent = spans.iter().find(|s| operation(s) == Some("invoke_agent")).unwrap(); + assert_eq!(spans.iter().filter(|s| operation(s) == Some("chat")).count(), 1); + assert!(matches!(chat.status, Status::Error { .. })); + let expected_error = match begin_mode { BeginMode::Cancelled => "cancelled", _ => "model_error" }; + assert_eq!(attr(chat, "error.type"), Some(&Value::String(expected_error.into()))); assert_eq!(attr(chat, "gen_ai.provider.name"), Some(&Value::String("before-begin".into()))); assert_eq!(attr(chat, "gen_ai.request.model"), Some(&Value::String("before-model".into()))); assert_eq!(attr(agent, "gen_ai.provider.name"), Some(&Value::String("before-begin".into()))); } } -"#; +"#, + include_str!("context_recovery_otel_harness.rs") + ); } #[derive(Default)] @@ -5047,6 +5069,21 @@ pub enum LoopError { Unsupported(String), } +/// Attribute only sanitized, bounded failure categories to an inference attempt. +fn record_inference_error(span: &tracing::Span, error: &LoopError) { + let kind = if error.is_context_length_exceeded() { + "context_length_exceeded" + } else { + match error { + LoopError::Cancelled => "cancelled", + LoopError::ContextRecoveryDeadlineExceeded => "context_recovery_deadline_exceeded", + _ => "model_error", + } + }; + span.record("error.type", kind); + span.record("otel.status_code", "ERROR"); +} + /// Ordinary inference remains unbounded; only recovery has an absolute budget. async fn recovery_bound( future: impl std::future::Future>,