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_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 new file mode 100644 index 0000000..839b36d --- /dev/null +++ b/crates/agentkit-loop/src/context_recovery_tests.rs @@ -0,0 +1,729 @@ +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 { + cancellations: StdArc, + polls: StdArc, + executions: StdArc, + site: FailureSite, + requests: StdArc>>, + } + struct RecoveryTurn { + cancellations: StdArc, + polls: StdArc, + 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 { + cancellations: StdArc::new(AtomicUsize::new(0)), + polls: StdArc::new(AtomicUsize::new(0)), + 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 { + cancellations: self.cancellations.clone(), + polls: self.polls.clone(), + 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 { + 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()))]) + } 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 { + cancellations: self.cancellations.clone(), + polls: self.polls.clone(), + events, + hold_when_empty: false, + }) + } + } + #[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; + } + self.events.pop_front().transpose() + } + } + + #[derive(Clone, Copy)] + enum Edit { + Change, + Noop, + Empty, + AssistantTail, + 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::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 { + 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()) + .with_response_attempt_supersession(), + ) + .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\"" + ); + } + #[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() { + // 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 { + 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())); + // 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; + 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]); + } + .await; + } +} diff --git a/crates/agentkit-loop/src/lib.rs b/crates/agentkit-loop/src/lib.rs index 33f7c43..ee97be4 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, @@ -1404,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() { @@ -1429,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(), @@ -1436,6 +1473,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 +1512,7 @@ where resources: Arc, cancellation: Option, mutators: Vec>, + context_recovery: Option, observers: Vec>, transcript_observers: Vec>, transcript: Vec, @@ -1495,6 +1534,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 +1610,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 +1702,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, @@ -1698,6 +1746,7 @@ pub struct LoopDriver where S: ModelSession, { + allow_supersession: bool, session_id: SessionId, observed_session_id: Arc, provider_name: Option, @@ -1711,6 +1760,7 @@ where resources: Arc, cancellation: Option, mutators: Vec>, + context_recovery: Option, observers: Vec>, transcript_observers: Vec>, transcript: Vec, @@ -2059,6 +2109,7 @@ where session_id: &session_id, turn_id, point, + recovery: None, cancellation: cancellation.clone(), emitter: &emitter, }; @@ -2070,6 +2121,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 || !transcript_has_pending_input(&candidate) { + 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| { @@ -2295,190 +2429,280 @@ 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 - .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, - ); - 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), + 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()); - } - match session.model_name() { - Some(model) => { - chat_span.record("gen_ai.request.model", model); - chat_span.record("otel.name", format!("chat {model}").as_str()); + // 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 turn = match session - .begin_turn(request, cancellation.clone()) - .instrument(chat_span.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( + 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(), + ) .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) => { + record_inference_error(&chat_span, &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) => { + record_inference_error(&chat_span, &error); + drop(chat_span); + 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) => { + record_inference_error(&chat_span, &LoopError::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::Usage(usage) => { - latest_usage = Some(usage.clone()); - self.emit(AgentEvent::UsageUpdated(usage)); + 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); + } + match self + .recover_context(&error, &turn_id, cancellation.clone(), &mut recovery) + .await + { + Ok(true) => { + if latest_usage.is_some() { + self.emit(AgentEvent::ResponseAttemptSuperseded); + } + 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); } - ModelTurnEvent::ResponseAttemptSuperseded => { - unreachable!("response-attempt supersession is handled before cancellation") + if cancellation + .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) + .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::Finished(result) => { - finished_result = Some(result); - break; + 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)); + } + 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(|| { + 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, chat_span); + }; result.usage = merge_usage(result.usage, latest_usage); if let Some(model) = &result.model { chat_span.record("gen_ai.response.model", model.as_str()); @@ -4021,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}; @@ -4034,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 { @@ -4049,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() }); } } } @@ -4159,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(); @@ -4207,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)] @@ -4819,6 +5052,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 +5069,78 @@ 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>, + 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 +5234,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..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 @@ -492,6 +530,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 +744,177 @@ 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 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>, } @@ -1081,3 +1293,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..4d1dc3a 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. 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.