diff --git a/src/crates/adapters/ai-adapters/src/client/sse.rs b/src/crates/adapters/ai-adapters/src/client/sse.rs index 2b4080d7cc..07eff41b74 100644 --- a/src/crates/adapters/ai-adapters/src/client/sse.rs +++ b/src/crates/adapters/ai-adapters/src/client/sse.rs @@ -3,6 +3,7 @@ use crate::client::StreamResponse; use crate::stream::UnifiedResponse; use crate::trace::{ModelExchangeRequestAttempt, ModelExchangeTraceConfig}; use anyhow::{anyhow, Result}; +use bitfun_core_types::errors::{AiProviderError, ErrorCategory}; use chrono::{DateTime, Utc}; use futures::Stream; use log::{debug, error, warn}; @@ -104,6 +105,32 @@ fn is_retryable_http_status(status: StatusCode) -> bool { status.is_server_error() || matches!(status.as_u16(), 408 | 409 | 425 | 429) } +fn provider_error_code(body: &str) -> Option { + let value: serde_json::Value = serde_json::from_str(body).ok()?; + let error = value.get("error").unwrap_or(&value); + ["code", "type", "status"].iter().find_map(|field| { + error.get(field).and_then(|value| match value { + serde_json::Value::String(value) => Some(value.clone()), + serde_json::Value::Number(value) => Some(value.to_string()), + _ => None, + }) + }) +} + +fn http_provider_error( + label: &str, + status: StatusCode, + error_text: &str, + error_kind: &str, +) -> AiProviderError { + AiProviderError::from_parts( + format!("{} {} {}: {}", label, error_kind, status, error_text), + Some(label.to_string()), + provider_error_code(error_text), + Some(status.as_u16()), + ) +} + fn exponential_retry_delay_ms(attempt: usize) -> u64 { let shift = u32::try_from(attempt) .unwrap_or(u32::MAX) @@ -228,6 +255,7 @@ where request_url: url.to_string(), request_body: trace.capture_request_body.then(|| request_body.clone()), attempt_number: attempt + 1, + round_attempt: trace.round_attempt().cloned(), }) .await } else { @@ -248,22 +276,24 @@ where .text() .await .unwrap_or_else(|e| format!("Failed to read error response: {}", e)); + let provider_error = + http_provider_error(label, status, &error_text, "client error"); if let Some(trace) = trace.as_ref() { trace .sink .request_attempt_failed( trace_handle.as_ref(), - &format!("{} client error {}: {}", label, status, error_text), + &provider_error.to_string(), ) .await; } - error!("{} client error {}: {}", label, status, error_text); - return Err(anyhow!("{} client error {}: {}", label, status, error_text)); + error!("{}", provider_error); + return Err(anyhow!(provider_error)); } if status.is_success() { debug!( - "{} request connected: {}ms, status: {}, protocol: {:?}, attempt: {}/{}", + "{} request connected: {}ms, status: {}, protocol: {:?}, transport_attempt: {}/{}", label, connect_time, status, @@ -277,9 +307,23 @@ where .text() .await .unwrap_or_else(|e| format!("Failed to read error response: {}", e)); - let error = anyhow!("{} error {}: {}", label, status, error_text); + let provider_error = http_provider_error(label, status, &error_text, "error"); + if provider_error.category == ErrorCategory::ContextOverflow { + if let Some(trace) = trace.as_ref() { + trace + .sink + .request_attempt_failed( + trace_handle.as_ref(), + &provider_error.to_string(), + ) + .await; + } + error!("{}", provider_error); + return Err(anyhow!(provider_error)); + } + let error = anyhow!(provider_error); warn!( - "{} request failed: {}ms, attempt {}/{}, error: {}", + "{} request failed: {}ms, transport_attempt {}/{}, error: {}", label, connect_time, attempt + 1, @@ -303,7 +347,7 @@ where if attempt < max_tries - 1 { let delay_ms = retry_delay_ms(attempt, &headers, status); debug!( - "Retrying {} after {}ms (attempt {}, status {})", + "Retrying {} after {}ms (transport_attempt {}, status {})", label, delay_ms, attempt + 2, @@ -319,7 +363,7 @@ where let error_msg = format_transport_error(label, &e); let error = anyhow!("{}", error_msg); warn!( - "{} request failed: {}ms, attempt {}/{}, error: {}", + "{} request failed: {}ms, transport_attempt {}/{}, error: {}", label, connect_time, attempt + 1, @@ -337,7 +381,7 @@ where if attempt < max_tries - 1 { let delay_ms = exponential_retry_delay_ms(attempt); debug!( - "Retrying {} after {}ms (attempt {})", + "Retrying {} after {}ms (transport_attempt {})", label, delay_ms, attempt + 2 @@ -351,7 +395,7 @@ where let error_msg = format_ttft_timeout_error(label, ttft_timeout); let error = anyhow!("{}", error_msg); warn!( - "{} request failed: {}ms, attempt {}/{}, error: {}", + "{} request failed: {}ms, transport_attempt {}/{}, error: {}", label, connect_time, attempt + 1, @@ -369,7 +413,7 @@ where if attempt < max_tries - 1 { let delay_ms = exponential_retry_delay_ms(attempt); debug!( - "Retrying {} after {}ms (attempt {})", + "Retrying {} after {}ms (transport_attempt {})", label, delay_ms, attempt + 2 @@ -419,6 +463,35 @@ mod tests { Arc, }; + #[test] + fn http_error_uses_structured_code_before_generic_message() { + let error = http_provider_error( + "OpenAI Responses API", + StatusCode::BAD_REQUEST, + r#"{"error":{"code":"context_length_exceeded","message":"Request failed"}}"#, + "client error", + ); + + assert_eq!(error.category, ErrorCategory::ContextOverflow); + assert_eq!( + error.provider_code.as_deref(), + Some("context_length_exceeded") + ); + assert_eq!(error.http_status, Some(400)); + } + + #[test] + fn no_body_bad_request_is_not_assumed_to_be_context_overflow() { + let error = http_provider_error( + "OpenAI Responses API", + StatusCode::BAD_REQUEST, + "400 status code (no body)", + "client error", + ); + + assert_eq!(error.category, ErrorCategory::InvalidRequest); + } + #[test] fn format_ttft_timeout_error_includes_timeout_seconds() { let message = format_ttft_timeout_error( diff --git a/src/crates/adapters/ai-adapters/src/lib.rs b/src/crates/adapters/ai-adapters/src/lib.rs index 6733f20e5a..80fbaefe84 100644 --- a/src/crates/adapters/ai-adapters/src/lib.rs +++ b/src/crates/adapters/ai-adapters/src/lib.rs @@ -19,7 +19,7 @@ pub use model_selector::{ pub use stream::{UnifiedResponse, UnifiedTokenUsage, UnifiedToolCall}; pub use trace::{ ModelExchangeRequestAttempt, ModelExchangeRequestTraceHandle, ModelExchangeResponseTrace, - ModelExchangeTraceConfig, ModelExchangeTraceSink, + ModelExchangeRoundAttempt, ModelExchangeTraceConfig, ModelExchangeTraceSink, }; pub use types::{ resolve_request_url, AIConfig, ConnectionTestMessageCode, ConnectionTestResult, GeminiResponse, diff --git a/src/crates/adapters/ai-adapters/src/stream/stream_handler/anthropic.rs b/src/crates/adapters/ai-adapters/src/stream/stream_handler/anthropic.rs index 49e2595edd..d1638d2103 100644 --- a/src/crates/adapters/ai-adapters/src/stream/stream_handler/anthropic.rs +++ b/src/crates/adapters/ai-adapters/src/stream/stream_handler/anthropic.rs @@ -7,6 +7,7 @@ use crate::stream::types::anthropic::{ }; use crate::stream::types::unified::UnifiedResponse; use anyhow::{anyhow, Result}; +use bitfun_core_types::errors::AiProviderError; use eventsource_stream::Eventsource; use log::{error, trace}; use reqwest::Response; @@ -100,11 +101,11 @@ pub async fn handle_anthropic_stream( let _ = tx.send(format!("[{}] {}", event_type, data)); } - if let Some(error_msg) = format_provider_error_from_sse_message(&event_type, &data) { + if let Some(provider_error) = provider_error_from_sse_message(&event_type, &data) { stats.increment("error:provider_message"); stats.log_summary("provider_error_message_received"); - error!("{}", error_msg); - let _ = tx_event.send(Err(anyhow!(error_msg))); + error!("{}", provider_error); + let _ = tx_event.send(Err(anyhow!(provider_error))); return; } @@ -223,7 +224,14 @@ pub async fn handle_anthropic_stream( }; stats.increment("error:api"); stats.log_summary("error_event_received"); - let _ = tx_event.send(Err(anyhow!(String::from(sse_error.error)))); + let code = sse_error.error.error_type.clone(); + let provider_error = AiProviderError::from_parts( + String::from(sse_error.error), + Some("anthropic".to_string()), + Some(code), + None, + ); + let _ = tx_event.send(Err(anyhow!(provider_error))); return; } "message_stop" => { @@ -240,7 +248,7 @@ pub async fn handle_anthropic_stream( } } -fn format_provider_error_from_sse_message(event_type: &str, data: &str) -> Option { +fn provider_error_from_sse_message(event_type: &str, data: &str) -> Option { if event_type != "message" { return None; } @@ -269,7 +277,12 @@ fn format_provider_error_from_sse_message(event_type: &str, data: &str) -> Optio formatted.push_str(&format!(", request_id={}", request_id)); } - Some(formatted) + Some(AiProviderError::from_parts( + formatted, + Some("anthropic_compatible".to_string()), + Some(code), + None, + )) } fn should_trace_anthropic_sse_event(event_type: &str, _data: &str) -> bool { @@ -376,28 +389,31 @@ fn emit_normalized_response( #[cfg(test)] mod tests { use super::{ - format_provider_error_from_sse_message, should_log_full_stream_events, + provider_error_from_sse_message, should_log_full_stream_events, should_trace_anthropic_sse_event, should_trace_unified_response, }; use crate::stream::types::unified::{UnifiedResponse, UnifiedToolCall}; + use bitfun_core_types::errors::ErrorCategory; #[test] fn extracts_glm_business_error_from_message_event() { let raw = r#"{"error":{"code":"1113","message":"余额不足或无可用资源包,请充值。"},"request_id":"20260425142416"}"#; - let formatted = format_provider_error_from_sse_message("message", raw).unwrap(); + let error = provider_error_from_sse_message("message", raw).unwrap(); + let formatted = error.message; assert!(formatted.contains("Provider error")); assert!(formatted.contains("code=1113")); assert!(formatted.contains("余额不足或无可用资源包")); assert!(formatted.contains("request_id=20260425142416")); + assert_eq!(error.category, ErrorCategory::ProviderQuota); } #[test] fn ignores_regular_anthropic_delta_events() { let raw = r#"{"type":"message_delta","delta":{"stop_reason":null}}"#; - assert!(format_provider_error_from_sse_message("message_delta", raw).is_none()); + assert!(provider_error_from_sse_message("message_delta", raw).is_none()); } #[test] diff --git a/src/crates/adapters/ai-adapters/src/stream/stream_handler/gemini.rs b/src/crates/adapters/ai-adapters/src/stream/stream_handler/gemini.rs index a274db0505..d374f725e7 100644 --- a/src/crates/adapters/ai-adapters/src/stream/stream_handler/gemini.rs +++ b/src/crates/adapters/ai-adapters/src/stream/stream_handler/gemini.rs @@ -3,6 +3,7 @@ use super::{next_stream_item, StreamTimeoutController, StreamTimeoutStage, Timed use crate::stream::types::gemini::GeminiSSEData; use crate::stream::types::unified::UnifiedResponse; use anyhow::{anyhow, Result}; +use bitfun_core_types::errors::AiProviderError; use eventsource_stream::Eventsource; use log::{error, trace}; use reqwest::Response; @@ -76,6 +77,24 @@ fn extract_api_error_message(event_json: &Value) -> Option { Some("Gemini streaming request failed".to_string()) } +fn extract_api_error(event_json: &Value) -> Option { + let error = event_json.get("error")?; + let code = error + .get("status") + .or_else(|| error.get("code")) + .and_then(|value| match value { + Value::String(value) => Some(value.clone()), + Value::Number(value) => Some(value.to_string()), + _ => None, + }); + Some(AiProviderError::from_parts( + extract_api_error_message(event_json)?, + Some("gemini".to_string()), + code, + None, + )) +} + pub async fn handle_gemini_stream( response: Response, tx_event: mpsc::UnboundedSender>, @@ -157,12 +176,15 @@ pub async fn handle_gemini_stream( } }; - if let Some(message) = extract_api_error_message(&event_json) { - let error_msg = format!("Gemini SSE API error: {}, data: {}", message, raw); + if let Some(mut provider_error) = extract_api_error(&event_json) { + provider_error.message = format!( + "Gemini SSE API error: {}, data: {}", + provider_error.message, raw + ); stats.increment("error:api"); stats.log_summary("sse_api_error"); - error!("{}", error_msg); - let _ = tx_event.send(Err(anyhow!(error_msg))); + error!("{}", provider_error); + let _ = tx_event.send(Err(anyhow!(provider_error))); return; } @@ -211,8 +233,9 @@ pub async fn handle_gemini_stream( #[cfg(test)] mod tests { - use super::GeminiToolCallState; + use super::{extract_api_error, GeminiToolCallState}; use crate::stream::types::unified::UnifiedToolCall; + use bitfun_core_types::errors::ErrorCategory; #[test] fn reuses_active_tool_id_by_omitting_follow_up_ids() { @@ -325,4 +348,19 @@ mod tests { assert_ne!(first.id, second.id); } + + #[test] + fn classifies_context_overflow_from_gemini_error_message() { + let event = serde_json::json!({ + "error": { + "code": 400, + "status": "INVALID_ARGUMENT", + "message": "The input token count exceeds the maximum number of tokens allowed" + } + }); + + let error = extract_api_error(&event).expect("provider error"); + assert_eq!(error.category, ErrorCategory::ContextOverflow); + assert_eq!(error.provider_code.as_deref(), Some("INVALID_ARGUMENT")); + } } diff --git a/src/crates/adapters/ai-adapters/src/stream/stream_handler/openai.rs b/src/crates/adapters/ai-adapters/src/stream/stream_handler/openai.rs index cc4ab202bd..5ee23471a5 100644 --- a/src/crates/adapters/ai-adapters/src/stream/stream_handler/openai.rs +++ b/src/crates/adapters/ai-adapters/src/stream/stream_handler/openai.rs @@ -4,6 +4,7 @@ use super::{next_stream_item, StreamTimeoutController, StreamTimeoutStage, Timed use crate::stream::types::openai::OpenAISSEData; use crate::stream::types::unified::UnifiedResponse; use anyhow::{anyhow, Result}; +use bitfun_core_types::errors::AiProviderError; use eventsource_stream::Eventsource; use log::{error, trace, warn}; use reqwest::Response; @@ -62,6 +63,22 @@ fn extract_sse_api_error_message(event_json: &Value) -> Option { Some("An error occurred during streaming".to_string()) } +fn extract_sse_api_error(event_json: &Value) -> Option { + let error = event_json.get("error")?; + let code = error + .get("code") + .or_else(|| error.get("type")) + .and_then(Value::as_str) + .map(str::to_string); + let message = extract_sse_api_error_message(event_json)?; + Some(AiProviderError::from_parts( + message, + Some("openai".to_string()), + code, + None, + )) +} + /// Convert a byte stream into a structured response stream /// /// # Arguments @@ -160,12 +177,13 @@ pub async fn handle_openai_stream( } }; - if let Some(api_error_message) = extract_sse_api_error_message(&event_json) { - let error_msg = format!("SSE API error: {}, data: {}", api_error_message, raw); + if let Some(mut provider_error) = extract_sse_api_error(&event_json) { + provider_error.message = + format!("SSE API error: {}, data: {}", provider_error.message, raw); stats.increment("error:api"); stats.log_summary("sse_api_error"); - error!("{}", error_msg); - let _ = tx_event.send(Err(anyhow!(error_msg))); + error!("{}", provider_error); + let _ = tx_event.send(Err(anyhow!(provider_error))); return; } @@ -250,7 +268,10 @@ pub async fn handle_openai_stream( #[cfg(test)] mod tests { - use super::{extract_sse_api_error_message, is_valid_chat_completion_chunk_weak}; + use super::{ + extract_sse_api_error, extract_sse_api_error_message, is_valid_chat_completion_chunk_weak, + }; + use bitfun_core_types::errors::ErrorCategory; #[test] fn weak_filter_accepts_chat_completion_chunk() { @@ -302,6 +323,23 @@ mod tests { ); } + #[test] + fn preserves_context_overflow_code_from_stream_error() { + let event = serde_json::json!({ + "error": { + "code": "context_length_exceeded", + "message": "Request failed" + } + }); + + let error = extract_sse_api_error(&event).expect("provider error"); + assert_eq!(error.category, ErrorCategory::ContextOverflow); + assert_eq!( + error.provider_code.as_deref(), + Some("context_length_exceeded") + ); + } + #[test] fn extracts_api_error_message_from_string_shape() { let event = serde_json::json!({ diff --git a/src/crates/adapters/ai-adapters/src/stream/stream_handler/responses.rs b/src/crates/adapters/ai-adapters/src/stream/stream_handler/responses.rs index 7dce63cbbd..ea6e4382ef 100644 --- a/src/crates/adapters/ai-adapters/src/stream/stream_handler/responses.rs +++ b/src/crates/adapters/ai-adapters/src/stream/stream_handler/responses.rs @@ -6,6 +6,7 @@ use crate::stream::types::responses::{ use crate::stream::types::unified::UnifiedResponse; use anyhow::{anyhow, Result}; use bitfun_agent_stream::ToolCallCompletion; +use bitfun_core_types::errors::AiProviderError; use eventsource_stream::Eventsource; use log::{error, trace}; use reqwest::Response; @@ -232,8 +233,10 @@ fn handle_function_call_output_item_done( } fn extract_api_error_message(event_json: &Value) -> Option { - let response = event_json.get("response")?; - let error = response.get("error")?; + let error = event_json + .get("response") + .and_then(|response| response.get("error")) + .or_else(|| event_json.get("error"))?; if error.is_null() { return None; @@ -249,6 +252,39 @@ fn extract_api_error_message(event_json: &Value) -> Option { Some("An error occurred during responses streaming".to_string()) } +fn extract_api_error(event_json: &Value) -> Option { + let error = event_json + .get("response") + .and_then(|response| response.get("error")) + .or_else(|| event_json.get("error")); + let code = event_json + .get("code") + .and_then(Value::as_str) + .or_else(|| { + error.and_then(|error| { + error + .get("code") + .or_else(|| error.get("type")) + .and_then(Value::as_str) + }) + }) + .map(str::to_string); + if error.is_none() && code.is_none() { + return None; + } + let message = event_json + .get("message") + .and_then(Value::as_str) + .map(str::to_string) + .or_else(|| extract_api_error_message(event_json))?; + Some(AiProviderError::from_parts( + message, + Some("openai_responses".to_string()), + code, + None, + )) +} + pub async fn handle_responses_stream( response: Response, tx_event: mpsc::UnboundedSender>, @@ -332,15 +368,15 @@ pub async fn handle_responses_stream( } }; - if let Some(api_error_message) = extract_api_error_message(&event_json) { - let error_msg = format!( + if let Some(mut provider_error) = extract_api_error(&event_json) { + provider_error.message = format!( "Responses SSE API error: {}, data: {}", - api_error_message, raw + provider_error.message, raw ); stats.increment("error:api"); stats.log_summary("sse_api_error"); - error!("{}", error_msg); - let _ = tx_event.send(Err(anyhow!(error_msg))); + error!("{}", provider_error); + let _ = tx_event.send(Err(anyhow!(provider_error))); return; } @@ -686,11 +722,12 @@ pub async fn handle_responses_stream( #[cfg(test)] mod tests { use super::{ - super::stream_stats::StreamStats, extract_api_error_message, + super::stream_stats::StreamStats, extract_api_error, extract_api_error_message, handle_function_call_arguments_delta, handle_function_call_output_item_done, responses_completed_tool_call_completion, InProgressToolCall, StreamTimeoutController, }; use bitfun_agent_stream::ToolCallCompletion; + use bitfun_core_types::errors::ErrorCategory; use serde_json::json; use std::collections::HashMap; use tokio::sync::mpsc; @@ -712,6 +749,26 @@ mod tests { ); } + #[test] + fn preserves_context_overflow_code_from_failed_response() { + let event = json!({ + "type": "response.failed", + "response": { + "error": { + "code": "context_length_exceeded", + "message": "Request failed" + } + } + }); + + let error = extract_api_error(&event).expect("provider error"); + assert_eq!(error.category, ErrorCategory::ContextOverflow); + assert_eq!( + error.provider_code.as_deref(), + Some("context_length_exceeded") + ); + } + #[test] fn returns_none_when_no_response_error_exists() { let event = json!({ diff --git a/src/crates/adapters/ai-adapters/src/trace.rs b/src/crates/adapters/ai-adapters/src/trace.rs index 0f858de725..495a408f2f 100644 --- a/src/crates/adapters/ai-adapters/src/trace.rs +++ b/src/crates/adapters/ai-adapters/src/trace.rs @@ -3,12 +3,23 @@ use serde::{Deserialize, Serialize}; use serde_json::Value; use std::sync::Arc; +/// Identity of the logical model-round attempt that owns one or more +/// adapter-level transport attempts. +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +pub struct ModelExchangeRoundAttempt { + pub attempt_id: String, + pub attempt_index: u32, +} + #[derive(Debug, Clone, Serialize, Deserialize)] pub struct ModelExchangeRequestAttempt { pub request_url: String, #[serde(default, skip_serializing_if = "Option::is_none")] pub request_body: Option, + /// One-based retry number inside a single adapter request invocation. pub attempt_number: usize, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub round_attempt: Option, } #[derive(Debug, Clone, Serialize, Deserialize)] @@ -59,4 +70,76 @@ pub trait ModelExchangeTraceSink: Send + Sync { pub struct ModelExchangeTraceConfig { pub sink: Arc, pub capture_request_body: bool, + round_attempt: Option, +} + +impl ModelExchangeTraceConfig { + pub fn new(sink: Arc, capture_request_body: bool) -> Self { + Self { + sink, + capture_request_body, + round_attempt: None, + } + } + + pub fn with_round_attempt(mut self, attempt_id: String, attempt_index: u32) -> Self { + self.round_attempt = Some(ModelExchangeRoundAttempt { + attempt_id, + attempt_index, + }); + self + } + + pub fn round_attempt(&self) -> Option<&ModelExchangeRoundAttempt> { + self.round_attempt.as_ref() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + struct NoopTraceSink; + + #[async_trait] + impl ModelExchangeTraceSink for NoopTraceSink { + async fn request_attempt_started( + &self, + _attempt: &ModelExchangeRequestAttempt, + ) -> Option { + None + } + + async fn request_attempt_failed( + &self, + _handle: Option<&ModelExchangeRequestTraceHandle>, + _error: &str, + ) { + } + + async fn request_attempt_completed( + &self, + _handle: &ModelExchangeRequestTraceHandle, + _response: &ModelExchangeResponseTrace, + ) { + } + } + + #[test] + fn round_attempt_is_scoped_without_mutating_the_base_trace_config() { + let base = ModelExchangeTraceConfig::new(Arc::new(NoopTraceSink), true); + let scoped = base + .clone() + .with_round_attempt("round-1:attempt:2".to_string(), 2); + + assert!(base.round_attempt().is_none()); + assert_eq!( + scoped.round_attempt(), + Some(&ModelExchangeRoundAttempt { + attempt_id: "round-1:attempt:2".to_string(), + attempt_index: 2, + }) + ); + assert!(Arc::ptr_eq(&base.sink, &scoped.sink)); + } } diff --git a/src/crates/assembly/core/src/agentic/core/message.rs b/src/crates/assembly/core/src/agentic/core/message.rs index da8de36512..bbdb27c814 100644 --- a/src/crates/assembly/core/src/agentic/core/message.rs +++ b/src/crates/assembly/core/src/agentic/core/message.rs @@ -110,6 +110,8 @@ pub enum InternalReminderKind { InterruptedContinue, ThinkingOnlyRescue, FinalizeCacheAnchor, + /// Marks a retained suffix whose earlier exchanges in the same turn were compacted. + RecentContextBoundary, /// A Stop hook blocked the end of a turn and asked the agent to continue. StopHookBlock, /// Model-visible context contributed by a SessionStart or diff --git a/src/crates/assembly/core/src/agentic/execution/execution_engine.rs b/src/crates/assembly/core/src/agentic/execution/execution_engine.rs index ab54e9291f..4a260d5794 100644 --- a/src/crates/assembly/core/src/agentic/execution/execution_engine.rs +++ b/src/crates/assembly/core/src/agentic/execution/execution_engine.rs @@ -5,7 +5,7 @@ use super::model_exchange_trace::{ prepare_model_exchange_trace_for_workspace, ModelExchangeTraceOperation, }; -use super::round_executor::RoundExecutor; +use super::round_executor::{ModelRoundLifecycle, RoundExecutor}; use super::types::{ExecutionContext, ExecutionResult, RoundContext, RoundResult}; use crate::agentic::agents::{ build_prompt_context_for_workspace, get_agent_registry, PrependedPromptReminders, @@ -348,6 +348,8 @@ pub struct ExecutionEngine { impl ExecutionEngine { const AUTO_COMPRESSION_SAFETY_RESERVE_TOKENS: usize = 10_000; + const MAX_COMPRESSION_OVERFLOW_ATTEMPTS: usize = 4; + const MAX_MAIN_CONTEXT_OVERFLOW_RECOVERIES: usize = 2; const FINALIZE_AFTER_REPEATED_TOOL_FAILURES_REMINDER: &'static str = "This turn must end now because repeated tool failures have prevented further progress. Ignore any unfinished work. Your task now is to give the user a final answer. Do not call any more tools; any tool call will fail. Respond in plain text only. Summarize what was completed, what failed, the evidence available from the tool results, and the single best next step for the user."; const FINALIZE_AFTER_MAX_ROUNDS_REMINDER: &'static str = "This turn must end now because it has reached the round limit. Ignore any unfinished work. Your task now is to give the user a final answer. Do not call any more tools; any tool call will fail. Respond in plain text only. Summarize the most useful completed work and evidence collected so far, and clearly distinguish resolved items from anything still unresolved."; const FINALIZE_TOOL_DENIED_MESSAGE: &'static str = @@ -1785,12 +1787,32 @@ impl ExecutionEngine { return Ok(response.text); } Err(err) => { + let provider_error = err + .downcast_ref::() + .cloned(); + let err_msg = err.to_string(); warn!( "Compression summary generation failed (attempt {}/{}): {}", attempt + 1, max_tries, - err + err_msg ); + let category = provider_error + .as_ref() + .map(|error| error.category.clone()) + .unwrap_or_else(|| { + bitfun_core_types::errors::classify_ai_error_message(&err_msg) + }); + if category == bitfun_core_types::errors::ErrorCategory::ContextOverflow { + return Err(BitFunError::RecoverableContextOverflow( + provider_error.unwrap_or_else(|| { + bitfun_core_types::errors::AiProviderError::classified( + err_msg, + bitfun_core_types::errors::ErrorCategory::ContextOverflow, + ) + }), + )); + } last_error = Some(err); if attempt < max_tries - 1 { @@ -2053,6 +2075,7 @@ impl ExecutionEngine { &self, session_id: &str, dialog_turn_id: &str, + trigger: &str, runtime_messages: Vec, before_pressure: TokenPressureSnapshot, context_window: usize, @@ -2073,13 +2096,12 @@ impl ExecutionEngine { let start_time = std::time::Instant::now(); let old_messages_len = runtime_messages.len(); - let turns = self - .context_compressor - .collect_turns_for_auto_compression(session_id, runtime_messages.clone())?; - if turns.is_empty() { + if !runtime_messages + .iter() + .any(|message| message.role != MessageRole::System) + { return Ok(None); } - // Generate compression ID let compression_id = format!("compression_{}", uuid::Uuid::new_v4()); // Captured before `ai_client` is consumed by summary generation. @@ -2087,7 +2109,7 @@ impl ExecutionEngine { native_hooks::dispatch_pre_compact( Self::native_hook_facts(session_id, dialog_turn_id, workspace, &ai_client_model), - "auto", + trigger, ) .await; @@ -2097,7 +2119,7 @@ impl ExecutionEngine { session_id: session_id.to_string(), turn_id: dialog_turn_id.to_string(), compression_id: compression_id.clone(), - trigger: "auto".to_string(), + trigger: trigger.to_string(), tokens_before: before_pressure.total_tokens, context_window, }, @@ -2121,38 +2143,105 @@ impl ExecutionEngine { ModelExchangeTraceOperation { kind: "context_compression", id: &compression_id, - trigger: Some("auto"), + trigger: Some(trigger), }, ai_client.as_ref(), ) .await; - let model_summary = match self - .generate_compression_model_summary(CompressionModelSummaryInput { - ai_client, - runtime_messages: &runtime_messages, + let max_initial_recent = context_window.saturating_div(2).max(1); + let mut recent_target = + ContextCompressor::DEFAULT_RECENT_CONTEXT_TOKENS.min(max_initial_recent); + let mut previous_cutoff = None; + let mut selected_plan = None; + let mut model_summary = None; + + for attempt in 0..Self::MAX_COMPRESSION_OVERFLOW_ATTEMPTS { + let Some(plan) = self.context_compressor.plan_auto_compression( + session_id, + &runtime_messages, + recent_target, + previous_cutoff, + )? + else { + break; + }; + info!( + "Compression context plan: session_id={}, turn_id={}, trigger={}, attempt={}/{}, recent_target_tokens={}, recent_tail_tokens={}, recent_anchor_tokens={}, cutoff_message_index={}, summary_messages={}, recent_tail_messages={}, last_turn_complete={}", + session_id, dialog_turn_id, - workspace, - tool_definitions, - prepended_prompt_reminders, - primary_supports_image_understanding, - trace_config, - }) - .await - { - Ok(summary) => summary, - Err(err) => { - warn!( - "Model-based compression failed, falling back to structured local compression: {}", - err - ); - None + trigger, + attempt + 1, + Self::MAX_COMPRESSION_OVERFLOW_ATTEMPTS, + plan.recent_target_tokens, + plan.recent_tail_tokens, + plan.recent_anchor_tokens, + plan.cutoff_message_index, + plan.summary_messages.len(), + plan.recent_tail_messages.len(), + plan.last_turn_complete + ); + + let summary_result = self + .generate_compression_model_summary(CompressionModelSummaryInput { + ai_client: ai_client.clone(), + runtime_messages: &plan.summary_request_messages, + dialog_turn_id, + workspace, + tool_definitions, + prepended_prompt_reminders, + primary_supports_image_understanding, + trace_config: trace_config.clone(), + }) + .await; + + match summary_result { + Ok(summary) => { + selected_plan = Some(plan); + model_summary = summary; + break; + } + Err(err) if err.is_recoverable_context_overflow() => { + warn!( + "Compression request exceeded provider context: session_id={}, turn_id={}, trigger={}, attempt={}/{}, recent_target_tokens={}, cutoff_message_index={}, can_shorten_summary_prefix={}, error={}", + session_id, + dialog_turn_id, + trigger, + attempt + 1, + Self::MAX_COMPRESSION_OVERFLOW_ATTEMPTS, + plan.recent_target_tokens, + plan.cutoff_message_index, + plan.can_shorten_summary_prefix, + err + ); + let can_retry = attempt + 1 < Self::MAX_COMPRESSION_OVERFLOW_ATTEMPTS + && plan.can_shorten_summary_prefix; + previous_cutoff = Some(plan.cutoff_message_index); + selected_plan = Some(plan); + if can_retry { + recent_target = recent_target + .saturating_add(ContextCompressor::RECENT_CONTEXT_RETRY_STEP_TOKENS); + continue; + } + break; + } + Err(err) => { + warn!( + "Model-based compression failed, falling back to structured local compression: {}", + err + ); + selected_plan = Some(plan); + break; + } } + } + + let Some(selected_plan) = selected_plan else { + return Ok(None); }; - match self.context_compressor.compress_turns_with_contract( + match self.context_compressor.compress_auto_plan_with_contract( session_id, context_window, - turns, - CompressionMode::Auto, + selected_plan, compression_contract, model_summary, ) { @@ -2167,7 +2256,7 @@ impl ExecutionEngine { session_id, boundary_turn_index, &compression_id, - "auto", + trigger, ) .await { @@ -2296,7 +2385,7 @@ impl ExecutionEngine { workspace, &ai_client_model, ), - "auto", + trigger, ) .await; @@ -3008,6 +3097,8 @@ impl ExecutionEngine { let mut finalization_reason: Option<&'static str> = None; let mut consecutive_compression_failures: u32 = 0; const MAX_CONSECUTIVE_COMPRESSION_FAILURES: u32 = 3; + let mut main_context_overflow_recoveries = 0usize; + let mut active_round_lifecycle: Option = None; // Track tool-call patterns for context health, but only use rounds with // actual failed tool results for no-progress recovery decisions. @@ -3229,6 +3320,7 @@ impl ExecutionEngine { .compress_messages( &context.session_id, &context.dialog_turn_id, + "auto", messages.clone(), token_pressure, context_window, @@ -3416,16 +3508,118 @@ impl ExecutionEngine { ) .await?; - let round_result = self + let round_lifecycle = + active_round_lifecycle.get_or_insert_with(ModelRoundLifecycle::new); + let round_result = match self .round_executor - .execute_round( + .execute_round_with_lifecycle( ai_client.clone(), round_context, ai_messages, tool_definitions.clone(), Some(context_window), + round_lifecycle, ) - .await?; + .await + { + Ok(result) => result, + Err(err) + if enable_context_compression + && err.is_recoverable_context_overflow() + && main_context_overflow_recoveries + < Self::MAX_MAIN_CONTEXT_OVERFLOW_RECOVERIES => + { + main_context_overflow_recoveries += 1; + warn!( + "Main model request exceeded provider context; starting recovery compression: session_id={}, turn_id={}, round_index={}, recovery={}/{}, error={}", + context.session_id, + context.dialog_turn_id, + round_index, + main_context_overflow_recoveries, + Self::MAX_MAIN_CONTEXT_OVERFLOW_RECOVERIES, + err + ); + match self + .compress_messages( + &context.session_id, + &context.dialog_turn_id, + "context_overflow_recovery", + messages.clone(), + send_pressure, + context_window, + ai_client.clone(), + &tool_definitions, + turn_prompt_scaffold.system_prompt_message.clone(), + &turn_prompt_scaffold.prepended_prompt_reminders, + primary_supports_image_understanding, + context_profile_policy.compression_contract_limit, + context.workspace.as_ref(), + ) + .await + { + Ok(Some((compressed_tokens, compressed_messages))) => { + info!( + "Context-overflow recovery compression completed: session_id={}, turn_id={}, round_index={}, recovery={}, messages {} -> {}, tokens {} -> {}", + context.session_id, + context.dialog_turn_id, + round_index, + main_context_overflow_recoveries, + messages.len(), + compressed_messages.len(), + send_pressure.total_tokens, + compressed_tokens + ); + messages = compressed_messages; + turn_prompt_scaffold = self + .resolve_turn_prompt_scaffold(TurnPromptScaffoldInput { + context: &context, + current_agent: current_agent.as_ref(), + model_name: &ai_client.config.model, + supports_image_understanding: + primary_supports_image_understanding, + tool_listing_sections: tool_listing_sections.clone(), + runtime_context_needs, + stage: "after_context_overflow_recovery", + }) + .await?; + Self::apply_turn_prompt_scaffold_to_messages( + &mut messages, + &turn_prompt_scaffold, + ); + self.round_executor + .record_context_overflow_recovery( + &context.session_id, + &context.dialog_turn_id, + round_lifecycle, + err.to_string(), + ) + .await; + full_compression_count += 1; + consecutive_compression_failures = 0; + continue; + } + Ok(None) => { + warn!( + "Context-overflow recovery found no compressible context: session_id={}, turn_id={}, round_index={}", + context.session_id, context.dialog_turn_id, round_index + ); + return Err(err); + } + Err(compression_error) => { + error!( + "Context-overflow recovery compression failed: session_id={}, turn_id={}, round_index={}, error={}", + context.session_id, + context.dialog_turn_id, + round_index, + compression_error + ); + return Err(err); + } + } + } + Err(err) => return Err(err), + }; + active_round_lifecycle = None; debug!( "Model round completed: round_index={}, has_more_rounds={}, tool_calls={}", diff --git a/src/crates/assembly/core/src/agentic/execution/model_exchange_trace.rs b/src/crates/assembly/core/src/agentic/execution/model_exchange_trace.rs index 2b20872c7d..a9be0df1fd 100644 --- a/src/crates/assembly/core/src/agentic/execution/model_exchange_trace.rs +++ b/src/crates/assembly/core/src/agentic/execution/model_exchange_trace.rs @@ -8,7 +8,7 @@ use crate::service::workspace_runtime::get_workspace_runtime_service_arc; use async_trait::async_trait; use bitfun_ai_adapters::{ ModelExchangeRequestAttempt, ModelExchangeRequestTraceHandle, ModelExchangeResponseTrace, - ModelExchangeTraceConfig, ModelExchangeTraceSink, + ModelExchangeRoundAttempt, ModelExchangeTraceConfig, ModelExchangeTraceSink, }; use chrono::{DateTime, Utc}; use dashmap::DashMap; @@ -55,7 +55,11 @@ struct ModelExchangeTraceRequestRecord { request_url: String, #[serde(default, skip_serializing_if = "Option::is_none")] body: Option, + /// Adapter-local HTTP/SSE retry number. This restarts for each logical + /// model-round attempt. attempt_number: usize, + #[serde(default, skip_serializing_if = "Option::is_none")] + round_attempt: Option, } #[derive(Debug, Clone, Copy)] @@ -288,6 +292,7 @@ impl ModelExchangeTraceSink for WorkspaceModelExchangeTraceSink { request_url: attempt.request_url.clone(), body: attempt.request_body.clone(), attempt_number: attempt.attempt_number, + round_attempt: attempt.round_attempt.clone(), }, }; @@ -404,8 +409,8 @@ pub(super) async fn prepare_model_exchange_trace_for_workspace( }, }; - Some(ModelExchangeTraceConfig { - sink: Arc::new(WorkspaceModelExchangeTraceSink::new( + Some(ModelExchangeTraceConfig::new( + Arc::new(WorkspaceModelExchangeTraceSink::new( WorkspaceModelExchangeTraceInput { trace_session_dir, policy, @@ -419,8 +424,8 @@ pub(super) async fn prepare_model_exchange_trace_for_workspace( model_id: ai_client.config.model.clone(), }, )), - capture_request_body: policy.capture_request_body, - }) + policy.capture_request_body, + )) } async fn current_model_exchange_trace_policy() -> Option { @@ -527,6 +532,10 @@ mod tests { request_url: "https://example.invalid/model".to_string(), request_body: Some(serde_json::json!({"request": "body"})), attempt_number: 2, + round_attempt: Some(ModelExchangeRoundAttempt { + attempt_id: "round-1:attempt:3".to_string(), + attempt_index: 3, + }), }) .await .expect("trace request should be recorded"); @@ -550,6 +559,13 @@ mod tests { assert_eq!(record.request.api_format, "api-format"); assert_eq!(record.request.model_id, "model-identity"); assert_eq!(record.request.attempt_number, 2); + assert_eq!( + record.request.round_attempt, + Some(ModelExchangeRoundAttempt { + attempt_id: "round-1:attempt:3".to_string(), + attempt_index: 3, + }) + ); assert_eq!( record.request.body, Some(serde_json::json!({"request": "body"})) diff --git a/src/crates/assembly/core/src/agentic/execution/round_executor.rs b/src/crates/assembly/core/src/agentic/execution/round_executor.rs index 2958d5235a..c8dd94c8b3 100644 --- a/src/crates/assembly/core/src/agentic/execution/round_executor.rs +++ b/src/crates/assembly/core/src/agentic/execution/round_executor.rs @@ -39,6 +39,7 @@ use bitfun_agent_runtime::turn_cancellation::DialogTurnCancellationTokenStore; use bitfun_ai_adapters::{ ModelExchangeRequestTraceHandle, ModelExchangeResponseTrace, ModelExchangeTraceConfig, }; +use bitfun_core_types::errors::{AiProviderError, ErrorCategory}; use bitfun_runtime_ports::PermissionRule; use log::{debug, error, warn}; use std::sync::Arc; @@ -53,6 +54,45 @@ pub struct RoundExecutor { cancellation_tokens: DialogTurnCancellationTokenStore, } +/// Mutable lifecycle shared by all provider attempts that belong to one +/// logical model round, including attempts made after overflow recovery. +#[derive(Debug)] +pub(super) struct ModelRoundLifecycle { + round_id: String, + started_at: Instant, + started_event_emitted: bool, + attempts_started: u32, +} + +impl ModelRoundLifecycle { + pub(super) fn new() -> Self { + Self { + round_id: uuid::Uuid::new_v4().to_string(), + started_at: Instant::now(), + started_event_emitted: false, + attempts_started: 0, + } + } + + fn take_started_event(&mut self) -> bool { + if self.started_event_emitted { + false + } else { + self.started_event_emitted = true; + true + } + } + + fn begin_attempt(&mut self) -> u32 { + self.attempts_started = self.attempts_started.saturating_add(1); + self.attempts_started + } + + fn attempts_started(&self) -> u32 { + self.attempts_started + } +} + impl RoundExecutor { const MAX_STREAM_ATTEMPTS: usize = 10; const RETRY_BASE_DELAY_MS: u64 = 500; @@ -115,6 +155,35 @@ impl RoundExecutor { .await; } + pub(super) async fn record_context_overflow_recovery( + &self, + session_id: &str, + turn_id: &str, + lifecycle: &ModelRoundLifecycle, + raw_error: String, + ) { + let attempt_number = lifecycle.attempts_started(); + if attempt_number == 0 { + return; + } + self.emit_event( + AgenticEvent::ModelRoundAttemptSuperseded { + session_id: session_id.to_string(), + turn_id: turn_id.to_string(), + round_id: lifecycle.round_id.clone(), + diagnostic: Self::retry_diagnostic( + format!("{}:attempt:{attempt_number}", lifecycle.round_id), + attempt_number, + "context_overflow", + Some(raw_error), + &[], + ), + }, + EventPriority::Normal, + ) + .await; + } + fn parsed_memory_citation_from_stream_result( stream_result: &StreamResult, ) -> Option { @@ -209,31 +278,55 @@ impl RoundExecutor { tool_definitions: Option>, context_window: Option, ) -> BitFunResult { - let round_started_at = Instant::now(); + let mut lifecycle = ModelRoundLifecycle::new(); + self.execute_round_with_lifecycle( + ai_client, + context, + ai_messages, + tool_definitions, + context_window, + &mut lifecycle, + ) + .await + } + + pub(super) async fn execute_round_with_lifecycle( + &self, + ai_client: Arc, + context: RoundContext, + ai_messages: Vec, + tool_definitions: Option>, + context_window: Option, + lifecycle: &mut ModelRoundLifecycle, + ) -> BitFunResult { + let round_started_at = lifecycle.started_at; let subagent_parent_info = context.subagent_parent_info.clone(); let is_subagent = subagent_parent_info.is_some(); - let round_id = uuid::Uuid::new_v4().to_string(); + let round_id = lifecycle.round_id.clone(); // Create or reuse cancellation token let cancel_token = self .cancellation_tokens .get_or_insert_new(&context.dialog_turn_id); - // Emit model round started event - self.emit_event( - AgenticEvent::ModelRoundStarted { - session_id: context.session_id.clone(), - turn_id: context.dialog_turn_id.clone(), - round_id: round_id.clone(), - round_group_id: context.round_group_id.clone(), - round_index: context.round_number, - model_config_id: context.model_config_id.clone(), - effective_model_name: context.effective_model_name.clone(), - }, - EventPriority::High, - ) - .await; + // Overflow recovery re-enters this executor with the same lifecycle. + // The logical round starts once even though it may contain many attempts. + if lifecycle.take_started_event() { + self.emit_event( + AgenticEvent::ModelRoundStarted { + session_id: context.session_id.clone(), + turn_id: context.dialog_turn_id.clone(), + round_id: round_id.clone(), + round_group_id: context.round_group_id.clone(), + round_index: context.round_number, + model_config_id: context.model_config_id.clone(), + effective_model_name: context.effective_model_name.clone(), + }, + EventPriority::High, + ) + .await; + } let trace_config = prepare_model_exchange_trace(&context, &round_id, ai_client.as_ref()).await; @@ -247,9 +340,9 @@ impl RoundExecutor { }; let allow_normal_tool_json_repair = global_config.ai.allow_tool_json_repair; let max_attempts = Self::MAX_STREAM_ATTEMPTS; - let mut attempt_index = 0usize; + let mut local_attempt_index = 0usize; let (stream_result, send_to_stream_ms, stream_processing_ms, final_trace_handle) = loop { - let attempt_number = (attempt_index + 1) as u32; + let attempt_number = lifecycle.begin_attempt(); let attempt_id = format!("{round_id}:attempt:{attempt_number}"); // Check cancellation before opening a model stream. This catches // early cancellation registered before the first round starts. @@ -263,18 +356,22 @@ impl RoundExecutor { let request_started_at = Instant::now(); debug!( - "Sending request: model={}, messages={}, tools={}, attempt={}/{}", + "Sending request: model={}, messages={}, tools={}, round_attempt={}, local_retry={}/{}", context.effective_model_name, ai_messages.len(), tool_definitions.as_ref().map(|t| t.len()).unwrap_or(0), - attempt_index + 1, + attempt_number, + local_attempt_index + 1, max_attempts ); // Use dynamically obtained client for call + let request_trace_config = trace_config + .clone() + .map(|config| config.with_round_attempt(attempt_id.clone(), attempt_number)); let send_future = ai_client.send_message_stream( ai_messages.clone(), tool_definitions.clone(), - trace_config.clone(), + request_trace_config, ); let send_result = tokio::select! { _ = cancel_token.cancelled() => { @@ -286,10 +383,11 @@ impl RoundExecutor { Ok(response) => { let send_to_stream_ms = elapsed_ms_u64(request_started_at); debug!( - "AI stream opened: session_id={}, round_id={}, attempt={}/{}, send_to_stream_ms={}", + "AI stream opened: session_id={}, round_id={}, round_attempt={}, local_retry={}/{}, send_to_stream_ms={}", context.session_id, round_id, - attempt_index + 1, + attempt_number, + local_attempt_index + 1, max_attempts, send_to_stream_ms ); @@ -297,9 +395,14 @@ impl RoundExecutor { } Err(e) => { error!("AI request failed: {}", e); + let provider_error = e.downcast_ref::().cloned(); let err_msg = e.to_string(); - if Self::is_transient_network_error(&err_msg) - && attempt_index < max_attempts - 1 + let is_structured_context_overflow = provider_error + .as_ref() + .is_some_and(|error| error.category == ErrorCategory::ContextOverflow); + if !is_structured_context_overflow + && Self::is_transient_network_error(&err_msg) + && local_attempt_index < max_attempts - 1 { self.record_retry_diagnostic( &context, @@ -311,21 +414,24 @@ impl RoundExecutor { &[], ) .await; - let delay_ms = Self::retry_delay_ms_for_error(attempt_index, &err_msg); + let delay_ms = + Self::retry_delay_ms_for_error(local_attempt_index, &err_msg); warn!( - "Retrying AI request after connection failure: session_id={}, round_id={}, attempt={}/{}, delay_ms={}, error={}", + "Retrying AI request after connection failure: session_id={}, round_id={}, round_attempt={}, local_retry={}/{}, delay_ms={}, error={}", context.session_id, round_id, - attempt_index + 1, + attempt_number, + local_attempt_index + 1, max_attempts, delay_ms, err_msg ); Self::sleep_with_cancellation(delay_ms, &cancel_token).await?; - attempt_index += 1; + local_attempt_index += 1; continue; } - if Self::is_transient_network_error(&err_msg) { + if !is_structured_context_overflow && Self::is_transient_network_error(&err_msg) + { return Err(BitFunError::AIClient(format!( "Stream retry budget exhausted after {} attempts: {}", max_attempts, err_msg @@ -337,7 +443,21 @@ impl RoundExecutor { // `BitFunError::error_category()` into `ErrorCategory` for // frontend recovery actions (wait_and_retry, switch_model, // etc.). - let error = BitFunError::AIClient(err_msg); + let category = provider_error + .as_ref() + .map(|error| error.category.clone()) + .unwrap_or_else(|| { + bitfun_core_types::errors::classify_ai_error_message(&err_msg) + }); + let error = if category == ErrorCategory::ContextOverflow { + BitFunError::RecoverableContextOverflow(provider_error.unwrap_or_else( + || AiProviderError::classified(err_msg, ErrorCategory::ContextOverflow), + )) + } else if let Some(error) = provider_error { + BitFunError::AIProvider(error) + } else { + BitFunError::AIClient(err_msg) + }; warn!( "AI request terminal failure: session_id={}, round_id={}, category={:?}, error={}", context.session_id, @@ -370,11 +490,12 @@ impl RoundExecutor { } debug!( - "Starting AI stream processing: session={}, round={}, thread={:?}, attempt={}/{}", + "Starting AI stream processing: session={}, round={}, thread={:?}, round_attempt={}, local_retry={}/{}", context.session_id, round_id, std::thread::current().id(), - attempt_index + 1, + attempt_number, + local_attempt_index + 1, max_attempts ); @@ -407,7 +528,7 @@ impl RoundExecutor { }); if !Self::has_user_visible_assistant_text(&result.full_text) - && attempt_index < max_attempts - 1 + && local_attempt_index < max_attempts - 1 && Self::is_transient_network_error(&err_msg) { self.record_retry_diagnostic( @@ -426,12 +547,14 @@ impl RoundExecutor { Self::trace_response_from_stream_result("partial", &result), ) .await; - let delay_ms = Self::retry_delay_ms_for_error(attempt_index, &err_msg); + let delay_ms = + Self::retry_delay_ms_for_error(local_attempt_index, &err_msg); warn!( - "Retrying stream because tool arguments were interrupted before valid JSON completed: session_id={}, round_id={}, attempt={}/{}, delay_ms={}, invalid_tool_calls={}, error={}", + "Retrying stream because tool arguments were interrupted before valid JSON completed: session_id={}, round_id={}, round_attempt={}, local_retry={}/{}, delay_ms={}, invalid_tool_calls={}, error={}", context.session_id, round_id, - attempt_index + 1, + attempt_number, + local_attempt_index + 1, max_attempts, delay_ms, result @@ -442,7 +565,7 @@ impl RoundExecutor { err_msg ); Self::sleep_with_cancellation(delay_ms, &cancel_token).await?; - attempt_index += 1; + local_attempt_index += 1; continue; } @@ -509,7 +632,7 @@ impl RoundExecutor { && !Self::has_user_visible_assistant_text(&result.full_text) && !result.tool_calls.is_empty() && Self::is_transient_network_error(partial_recovery_reason) - && attempt_index < max_attempts - 1 + && local_attempt_index < max_attempts - 1 { self.record_retry_diagnostic( &context, @@ -527,26 +650,29 @@ impl RoundExecutor { Self::trace_response_from_stream_result("partial", &result), ) .await; - let delay_ms = - Self::retry_delay_ms_for_error(attempt_index, partial_recovery_reason); + let delay_ms = Self::retry_delay_ms_for_error( + local_attempt_index, + partial_recovery_reason, + ); warn!( - "Retrying stream because tool calls arrived on an interrupted network stream without assistant text: session_id={}, round_id={}, attempt={}/{}, delay_ms={}, tool_calls={}, reason={}", + "Retrying stream because tool calls arrived on an interrupted network stream without assistant text: session_id={}, round_id={}, round_attempt={}, local_retry={}/{}, delay_ms={}, tool_calls={}, reason={}", context.session_id, round_id, - attempt_index + 1, + attempt_number, + local_attempt_index + 1, max_attempts, delay_ms, result.tool_calls.len(), partial_recovery_reason ); Self::sleep_with_cancellation(delay_ms, &cancel_token).await?; - attempt_index += 1; + local_attempt_index += 1; continue; } if Self::is_invalid_tool_only_without_text(&result) { let err_msg = "Provider returned only invalid tool arguments".to_string(); - if attempt_index < max_attempts - 1 { + if local_attempt_index < max_attempts - 1 { self.record_retry_diagnostic( &context, &round_id, @@ -567,18 +693,19 @@ impl RoundExecutor { ), ) .await; - let delay_ms = Self::retry_delay_ms(attempt_index); + let delay_ms = Self::retry_delay_ms(local_attempt_index); warn!( - "Retrying stream because provider returned only invalid tool arguments: session_id={}, round_id={}, attempt={}/{}, delay_ms={}, tool_calls={}", + "Retrying stream because provider returned only invalid tool arguments: session_id={}, round_id={}, round_attempt={}, local_retry={}/{}, delay_ms={}, tool_calls={}", context.session_id, round_id, - attempt_index + 1, + attempt_number, + local_attempt_index + 1, max_attempts, delay_ms, result.tool_calls.len() ); Self::sleep_with_cancellation(delay_ms, &cancel_token).await?; - attempt_index += 1; + local_attempt_index += 1; continue; } @@ -605,7 +732,7 @@ impl RoundExecutor { ))); } - if no_effective_output && attempt_index < max_attempts - 1 { + if no_effective_output && local_attempt_index < max_attempts - 1 { self.record_retry_diagnostic( &context, &round_id, @@ -625,26 +752,28 @@ impl RoundExecutor { ), ) .await; - let delay_ms = Self::retry_delay_ms(attempt_index); + let delay_ms = Self::retry_delay_ms(local_attempt_index); warn!( - "Retrying stream because no effective output was received: session_id={}, round_id={}, attempt={}/{}, delay_ms={}", + "Retrying stream because no effective output was received: session_id={}, round_id={}, round_attempt={}, local_retry={}/{}, delay_ms={}", context.session_id, round_id, - attempt_index + 1, + attempt_number, + local_attempt_index + 1, max_attempts, delay_ms ); Self::sleep_with_cancellation(delay_ms, &cancel_token).await?; - attempt_index += 1; + local_attempt_index += 1; continue; } if is_partial_recovery { warn!( - "Accepting stream partial recovery without retry: session_id={}, round_id={}, attempt={}/{}, reason={}", + "Accepting stream partial recovery without retry: session_id={}, round_id={}, round_attempt={}, local_retry={}/{}, reason={}", context.session_id, round_id, - attempt_index + 1, + attempt_number, + local_attempt_index + 1, max_attempts, result .partial_recovery_reason @@ -662,8 +791,10 @@ impl RoundExecutor { } Err(stream_err) => { let err_msg = stream_err.error.to_string(); + let stream_error_category = stream_err.error.error_category(); let can_retry = !stream_err.has_effective_output - && attempt_index < max_attempts - 1 + && stream_error_category != ErrorCategory::ContextOverflow + && local_attempt_index < max_attempts - 1 && Self::is_transient_network_error(&err_msg); Self::complete_model_exchange_trace( trace_config.as_ref(), @@ -682,26 +813,42 @@ impl RoundExecutor { &[], ) .await; - let delay_ms = Self::retry_delay_ms_for_error(attempt_index, &err_msg); + let delay_ms = + Self::retry_delay_ms_for_error(local_attempt_index, &err_msg); warn!( - "Retrying stream after transient error with no effective output: session_id={}, round_id={}, attempt={}/{}, delay_ms={}, error={}", + "Retrying stream after transient error with no effective output: session_id={}, round_id={}, round_attempt={}, local_retry={}/{}, delay_ms={}, error={}", context.session_id, round_id, - attempt_index + 1, + attempt_number, + local_attempt_index + 1, max_attempts, delay_ms, err_msg ); Self::sleep_with_cancellation(delay_ms, &cancel_token).await?; - attempt_index += 1; + local_attempt_index += 1; continue; } - if Self::is_transient_network_error(&err_msg) { + if stream_error_category != ErrorCategory::ContextOverflow + && Self::is_transient_network_error(&err_msg) + { return Err(BitFunError::AIClient(format!( "Stream retry budget exhausted after {} attempts: {}", max_attempts, err_msg ))); } + if !stream_err.has_effective_output + && stream_error_category == ErrorCategory::ContextOverflow + { + let provider_error = match stream_err.error { + BitFunError::AIProvider(error) + | BitFunError::RecoverableContextOverflow(error) => error, + _ => { + AiProviderError::classified(err_msg, ErrorCategory::ContextOverflow) + } + }; + return Err(BitFunError::RecoverableContextOverflow(provider_error)); + } return Err(stream_err.error); } } @@ -782,7 +929,7 @@ impl RoundExecutor { first_chunk_ms: stream_result.first_chunk_ms, first_visible_output_ms: stream_result.first_visible_output_ms, stream_duration_ms: Some(stream_processing_ms), - attempt_count: Some((attempt_index + 1) as u32), + attempt_count: Some(lifecycle.attempts_started()), failure_category: None, token_details: stream_result .usage @@ -879,8 +1026,11 @@ impl RoundExecutor { session_id: context.session_id.clone(), dialog_turn_id: context.dialog_turn_id.clone(), round_id: round_id.clone(), - attempt_id: Some(format!("{round_id}:attempt:{}", attempt_index + 1)), - attempt_index: Some((attempt_index + 1) as u32), + attempt_id: Some(format!( + "{round_id}:attempt:{}", + lifecycle.attempts_started() + )), + attempt_index: Some(lifecycle.attempts_started()), agent_type: context.agent_type.clone(), workspace: context.workspace.clone(), primary_model_facts: context.primary_model_facts.clone(), @@ -1483,9 +1633,9 @@ fn token_details_from_usage( #[cfg(test)] mod tests { - use super::{RoundExecutor, StreamProcessor}; + use super::{ModelRoundLifecycle, RoundExecutor, StreamProcessor}; use crate::agentic::core::ToolCall; - use crate::agentic::events::{EventQueue, EventQueueConfig}; + use crate::agentic::events::{AgenticEvent, EventQueue, EventQueueConfig}; use crate::agentic::execution::stream_processor::StreamResult; use crate::agentic::execution::types::RoundContext; use crate::agentic::tools::ToolRuntimeRestrictions; @@ -1514,6 +1664,55 @@ mod tests { } } + #[test] + fn model_round_lifecycle_reuses_identity_and_counts_recovery_attempts() { + let mut lifecycle = ModelRoundLifecycle::new(); + let round_id = lifecycle.round_id.clone(); + + assert!(lifecycle.take_started_event()); + assert!(!lifecycle.take_started_event()); + assert_eq!(lifecycle.begin_attempt(), 1); + assert_eq!(lifecycle.begin_attempt(), 2); + assert_eq!(lifecycle.attempts_started(), 2); + assert_eq!(lifecycle.round_id, round_id); + } + + #[tokio::test] + async fn context_overflow_recovery_supersedes_the_current_attempt() { + let executor = test_round_executor(); + let mut lifecycle = ModelRoundLifecycle::new(); + let round_id = lifecycle.round_id.clone(); + assert_eq!(lifecycle.begin_attempt(), 1); + + executor + .record_context_overflow_recovery( + "session-1", + "turn-1", + &lifecycle, + "request exceeds context window".to_string(), + ) + .await; + + let events = executor.event_queue.dequeue_batch(10).await; + assert_eq!(events.len(), 1); + match &events[0].event { + AgenticEvent::ModelRoundAttemptSuperseded { + session_id, + turn_id, + round_id: event_round_id, + diagnostic, + } => { + assert_eq!(session_id, "session-1"); + assert_eq!(turn_id, "turn-1"); + assert_eq!(event_round_id, &round_id); + assert_eq!(diagnostic.attempt_id, format!("{round_id}:attempt:1")); + assert_eq!(diagnostic.attempt_index, 1); + assert_eq!(diagnostic.category, "context_overflow"); + } + event => panic!("unexpected event: {event:?}"), + } + } + fn test_round_context() -> RoundContext { RoundContext { session_id: "session-1".to_string(), diff --git a/src/crates/assembly/core/src/agentic/session/compression/compressor.rs b/src/crates/assembly/core/src/agentic/session/compression/compressor.rs index 420e26e5b4..823a844a9e 100644 --- a/src/crates/assembly/core/src/agentic/session/compression/compressor.rs +++ b/src/crates/assembly/core/src/agentic/session/compression/compressor.rs @@ -54,6 +54,26 @@ pub struct CompressionResult { pub has_model_summary: bool, } +#[derive(Debug, Clone)] +pub struct AutoCompressionPlan { + pub summary_request_messages: Vec, + pub summary_messages: Vec, + pub recent_anchor_messages: Vec, + pub recent_tail_messages: Vec, + pub recent_target_tokens: usize, + pub recent_tail_tokens: usize, + pub recent_anchor_tokens: usize, + pub cutoff_message_index: usize, + pub can_shorten_summary_prefix: bool, + pub last_turn_complete: bool, +} + +#[derive(Debug, Clone, Copy)] +struct AtomicMessageUnit { + start: usize, + tokens: usize, +} + #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum CompressionMode { Auto, @@ -66,8 +86,12 @@ pub struct ContextCompressor { } impl ContextCompressor { + pub const DEFAULT_RECENT_CONTEXT_TOKENS: usize = 10_000; + pub const RECENT_CONTEXT_RETRY_STEP_TOKENS: usize = 10_000; const AUTO_COMPRESSION_CONTINUATION_REMINDER: &'static str = "The conversation context above was automatically compacted while work was still in progress. Re-establish the current working state from this summary and the latest available evidence, then continue working on the user's most recent request. Do not stop merely because compaction occurred, and do not ask the user to repeat information already captured here. If a required detail is missing and a pre-compaction transcript is available above, inspect the relevant part of that transcript before proceeding."; + const RECENT_CONTEXT_BOUNDARY_REMINDER: &'static str = + "The user message above started the current turn. Some intermediate assistant and tool exchanges from this turn were compacted into the preceding summary and are not repeated below. Continue from the retained recent context without repeating completed work."; pub fn new(config: CompressionConfig) -> Self { Self { config } @@ -149,6 +173,191 @@ impl ContextCompressor { self.collect_conversation_turns(session_id, messages) } + pub fn plan_auto_compression( + &self, + session_id: &str, + runtime_messages: &[Message], + recent_target_tokens: usize, + previous_cutoff: Option, + ) -> BitFunResult> { + let system_message_count = runtime_messages + .iter() + .take_while(|message| message.role == MessageRole::System) + .count(); + let conversation = &runtime_messages[system_message_count..]; + if conversation.is_empty() { + debug!( + "No conversation messages available for automatic compression planning: session_id={}", + session_id + ); + return Ok(None); + } + + let units = Self::atomic_message_units(conversation); + if units.is_empty() { + return Ok(None); + } + + let minimum_cutoff = if units.len() > 1 { + units[1].start + } else { + conversation.len() + }; + let mut cutoff = conversation.len(); + let mut accumulated_tokens = 0usize; + + for unit in units.iter().rev() { + if accumulated_tokens >= recent_target_tokens { + break; + } + if unit.start < minimum_cutoff { + break; + } + cutoff = unit.start; + accumulated_tokens = accumulated_tokens.saturating_add(unit.tokens); + } + + if let Some(previous_cutoff) = previous_cutoff { + if cutoff >= previous_cutoff { + if let Some(earlier_unit) = units + .iter() + .rev() + .find(|unit| unit.start < previous_cutoff && unit.start >= minimum_cutoff) + { + cutoff = earlier_unit.start; + } + } + } + + let summary_messages = conversation[..cutoff].to_vec(); + if summary_messages.is_empty() { + debug!( + "Automatic compression plan has no summary prefix: session_id={}, recent_target_tokens={}", + session_id, recent_target_tokens + ); + return Ok(None); + } + + let recent_tail_messages = conversation[cutoff..].to_vec(); + let recent_tail_tokens = recent_tail_messages + .iter() + .map(|message| message.estimate_tokens_with_reasoning(true)) + .sum(); + let mut summary_request_messages = runtime_messages[..system_message_count].to_vec(); + summary_request_messages.extend(summary_messages.clone()); + + let last_user_index = conversation + .iter() + .rposition(Message::is_actual_user_message); + let last_turn_complete = last_user_index.is_some_and(|index| cutoff <= index); + let mut recent_anchor_messages = Vec::new(); + + if let Some(last_user_index) = last_user_index.filter(|_| !last_turn_complete) { + let user_message = conversation[last_user_index].clone(); + let latest_todo = Self::latest_todo_snapshot_with_source( + &conversation[last_user_index..], + last_user_index, + ); + let missing_todo = latest_todo + .filter(|(source_index, _)| *source_index < cutoff) + .map(|(_, snapshot)| snapshot); + let mut reminder_text = Self::RECENT_CONTEXT_BOUNDARY_REMINDER.to_string(); + if let Some(todo) = missing_todo.as_ref() { + reminder_text + .push_str("\n\nLatest task list before the retained recent context:\n"); + reminder_text.push_str(&Self::render_todo_snapshot(todo)); + } + + let turn_id = user_message.metadata.turn_id.clone(); + recent_anchor_messages.push(user_message); + let mut reminder = Message::internal_reminder( + crate::agentic::core::InternalReminderKind::RecentContextBoundary, + reminder_text, + ); + reminder.metadata.turn_id = turn_id.clone(); + if let Some(todo) = missing_todo { + reminder = reminder.with_compression_payload(CompressionPayload { + entries: vec![CompressionEntry::Turn { + turn_id, + messages: Vec::new(), + todo: Some(todo), + }], + }); + } + recent_anchor_messages.push(reminder); + } + + let recent_anchor_tokens = recent_anchor_messages + .iter() + .map(|message| message.estimate_tokens_with_reasoning(true)) + .sum(); + let can_shorten_summary_prefix = units + .iter() + .any(|unit| unit.start < cutoff && unit.start >= minimum_cutoff); + + debug!( + "Automatic compression plan: session_id={}, recent_target_tokens={}, recent_tail_tokens={}, recent_anchor_tokens={}, cutoff_message_index={}, summary_messages={}, recent_tail_messages={}, last_turn_complete={}, can_shorten_summary_prefix={}", + session_id, + recent_target_tokens, + recent_tail_tokens, + recent_anchor_tokens, + cutoff, + summary_messages.len(), + recent_tail_messages.len(), + last_turn_complete, + can_shorten_summary_prefix + ); + + Ok(Some(AutoCompressionPlan { + summary_request_messages, + summary_messages, + recent_anchor_messages, + recent_tail_messages, + recent_target_tokens, + recent_tail_tokens, + recent_anchor_tokens, + cutoff_message_index: cutoff, + can_shorten_summary_prefix, + last_turn_complete, + })) + } + + fn atomic_message_units(messages: &[Message]) -> Vec { + let mut units = Vec::new(); + let mut index = 0usize; + + while index < messages.len() { + let start = index; + index += 1; + if messages[start].role == MessageRole::Assistant { + while index < messages.len() && messages[index].role == MessageRole::Tool { + index += 1; + } + } + let tokens = messages[start..index] + .iter() + .map(|message| message.estimate_tokens_with_reasoning(true)) + .sum(); + units.push(AtomicMessageUnit { start, tokens }); + } + + units + } + + fn latest_todo_snapshot_with_source( + messages: &[Message], + base_index: usize, + ) -> Option<(usize, CompressedTodoSnapshot)> { + messages + .iter() + .enumerate() + .rev() + .find_map(|(index, message)| { + MessageHelper::get_last_todo_snapshot(std::slice::from_ref(message)) + .map(|snapshot| (base_index + index, snapshot)) + }) + } + pub fn compress_turns( &self, session_id: &str, @@ -175,6 +384,28 @@ impl ContextCompressor { mode: CompressionMode, contract: Option, model_summary: Option, + ) -> BitFunResult { + self.compress_turns_internal( + session_id, + context_window, + turns, + mode, + contract, + model_summary, + matches!(mode, CompressionMode::Auto), + ) + } + + #[allow(clippy::too_many_arguments)] + fn compress_turns_internal( + &self, + session_id: &str, + context_window: usize, + turns: Vec, + mode: CompressionMode, + contract: Option, + model_summary: Option, + append_live_boundary_context: bool, ) -> BitFunResult { if turns.is_empty() { debug!("No turns need compression: session_id={}", session_id); @@ -205,7 +436,7 @@ impl ContextCompressor { Some(summary) => self.build_model_summary_artifact(summary, contract), None => self.build_fallback_summary_artifact(turns, context_window, contract), }; - if matches!(mode, CompressionMode::Auto) { + if append_live_boundary_context { self.append_live_boundary_context( &mut summary_artifact, last_user_message.as_ref(), @@ -228,6 +459,30 @@ impl ContextCompressor { }) } + pub fn compress_auto_plan_with_contract( + &self, + session_id: &str, + context_window: usize, + plan: AutoCompressionPlan, + contract: Option, + model_summary: Option, + ) -> BitFunResult { + let turns = MessageHelper::group_messages_by_turns(plan.summary_messages); + let turns = turns.into_iter().map(TurnWithTokens::new).collect(); + let mut result = self.compress_turns_internal( + session_id, + context_window, + turns, + CompressionMode::Auto, + contract, + model_summary, + false, + )?; + result.messages.extend(plan.recent_anchor_messages); + result.messages.extend(plan.recent_tail_messages); + Ok(result) + } + pub fn append_transcript_reference( &self, result: &mut CompressionResult, @@ -235,10 +490,10 @@ impl ContextCompressor { transcript_index_range: &TranscriptLineRange, ) -> bool { let has_model_summary = result.has_model_summary; - let summary_is_inline = result.messages.len() == 1 - && result.messages[0].role == MessageRole::User - && result.messages[0].metadata.semantic_kind - == Some(MessageSemanticKind::CompressionSummary); + let summary_is_inline = result.messages.first().is_some_and(|message| { + message.role == MessageRole::User + && message.metadata.semantic_kind == Some(MessageSemanticKind::CompressionSummary) + }); let old_boundary = render_system_reminder(&Self::render_boundary_marker_text( has_model_summary, summary_is_inline, @@ -587,7 +842,8 @@ fn extract_tag_content<'a>(text: &'a str, tag: &str) -> Option<&'a str> { mod tests { use super::{CompressionMode, ContextCompressor, TurnWithTokens}; use crate::agentic::core::{ - render_system_reminder, CompressionEntry, CompressionPayload, Message, MessageSemanticKind, + render_system_reminder, CompressionEntry, CompressionPayload, InternalReminderKind, + Message, MessageContent, MessageSemanticKind, ToolCall, ToolResult, }; use crate::service::session::TranscriptLineRange; @@ -619,6 +875,180 @@ mod tests { ]) } + fn todo_call() -> ToolCall { + ToolCall { + tool_id: "todo_recent".to_string(), + tool_name: "TodoWrite".to_string(), + arguments: serde_json::json!({ + "todos": [ + {"content": "Keep recent context", "status": "in_progress"}, + {"content": "Retry compression", "status": "pending"} + ] + }), + raw_arguments: None, + is_error: false, + parse_error: None, + recovered_from_truncation: false, + repair_kind: Default::default(), + } + } + + fn todo_result() -> Message { + Message::tool_result(ToolResult { + tool_id: "todo_recent".to_string(), + tool_name: "TodoWrite".to_string(), + effective_tool_name: None, + result: serde_json::json!({"success": true}), + result_for_assistant: None, + is_error: false, + duration_ms: None, + image_attachments: None, + }) + } + + #[test] + fn recent_context_keeps_complete_last_turn_without_duplicate_anchors() { + let compressor = ContextCompressor::new(Default::default()); + let messages = vec![ + Message::system("system".to_string()), + Message::user("Older request".repeat(200)), + Message::assistant("Older answer".repeat(200)), + Message::user("Current request".to_string()), + Message::assistant("Current answer".to_string()), + ]; + + let plan = compressor + .plan_auto_compression("session", &messages, 1_000, None) + .expect("planning succeeds") + .expect("plan exists"); + + assert!(plan.last_turn_complete); + assert!(plan.recent_anchor_messages.is_empty()); + assert!(plan + .recent_tail_messages + .iter() + .any(|message| message.is_actual_user_message())); + } + + #[test] + fn partial_last_turn_restores_user_and_missing_todo() { + let compressor = ContextCompressor::new(Default::default()); + let current_user = Message::user("Continue the current task".to_string()) + .with_turn_id("turn-current".to_string()); + let retained_assistant = Message::assistant("Latest evidence".repeat(100)) + .with_turn_id("turn-current".to_string()); + let messages = vec![ + Message::system("system".to_string()), + Message::user("Older request".to_string()), + Message::assistant("Older answer".to_string()), + current_user.clone(), + Message::assistant_with_tools("Planning".to_string(), vec![todo_call()]) + .with_turn_id("turn-current".to_string()), + todo_result().with_turn_id("turn-current".to_string()), + retained_assistant.clone(), + ]; + + let plan = compressor + .plan_auto_compression("session", &messages, 1, None) + .expect("planning succeeds") + .expect("plan exists"); + + assert!(!plan.last_turn_complete); + assert_eq!(plan.recent_anchor_messages.len(), 2); + assert_eq!(plan.recent_anchor_messages[0].id, current_user.id); + assert_eq!( + plan.recent_anchor_messages[1].internal_reminder_kind(), + Some(InternalReminderKind::RecentContextBoundary) + ); + assert!(matches!( + plan.recent_anchor_messages[1] + .metadata + .compression_payload + .as_ref() + .and_then(|payload| payload.entries.first()), + Some(CompressionEntry::Turn { todo: Some(todo), .. }) + if todo.todos.len() == 2 + )); + assert_eq!(plan.recent_tail_messages[0].id, retained_assistant.id); + + let mut result = compressor + .compress_auto_plan_with_contract( + "session", + 128_000, + plan, + None, + Some("Earlier work summary".to_string()), + ) + .expect("compression succeeds"); + assert_eq!(result.messages.len(), 4); + let MessageContent::Text(summary) = &result.messages[0].content else { + panic!("expected summary text"); + }; + assert!(!summary.contains("Most recent user message before this summary")); + assert_eq!(result.messages[1].id, current_user.id); + assert_eq!(result.messages[3].id, retained_assistant.id); + assert!(compressor.append_transcript_reference( + &mut result, + "bitfun://current-session/artifacts/compression-transcripts/3-recent.txt", + &TranscriptLineRange { + start_line: 1, + end_line: 20, + }, + )); + let MessageContent::Text(summary) = &result.messages[0].content else { + panic!("expected summary text"); + }; + assert!(summary.contains("3-recent.txt")); + } + + #[test] + fn recent_context_never_splits_tool_results_from_their_assistant_call() { + let compressor = ContextCompressor::new(Default::default()); + let assistant = Message::assistant_with_tools("Planning".to_string(), vec![todo_call()]); + let result = todo_result(); + let messages = vec![ + Message::system("system".to_string()), + Message::user("Older request".to_string()), + Message::assistant("Older answer".repeat(500)), + Message::user("Current request".to_string()), + assistant.clone(), + result.clone(), + ]; + + let plan = compressor + .plan_auto_compression("session", &messages, 1, None) + .expect("planning succeeds") + .expect("plan exists"); + + assert_eq!(plan.recent_tail_messages[0].id, assistant.id); + assert_eq!(plan.recent_tail_messages[1].id, result.id); + } + + #[test] + fn overflow_retry_forces_cutoff_to_move_across_large_atomic_units() { + let compressor = ContextCompressor::new(Default::default()); + let messages = vec![ + Message::system("system".to_string()), + Message::user("request".to_string()), + Message::assistant("first".repeat(2_000)), + Message::assistant("second".repeat(2_000)), + Message::assistant("third".repeat(2_000)), + ]; + + let first = compressor + .plan_auto_compression("session", &messages, 1, None) + .expect("planning succeeds") + .expect("first plan exists"); + let second = compressor + .plan_auto_compression("session", &messages, 2, Some(first.cutoff_message_index)) + .expect("planning succeeds") + .expect("second plan exists"); + + assert!(second.cutoff_message_index < first.cutoff_message_index); + assert!(second.summary_messages.len() < first.summary_messages.len()); + assert!(second.recent_tail_messages.len() > first.recent_tail_messages.len()); + } + #[test] fn manual_compression_creates_closed_compression_turn() { let compressor = ContextCompressor::new(Default::default()); diff --git a/src/crates/assembly/core/src/util/errors.rs b/src/crates/assembly/core/src/util/errors.rs index 3b91aff8fe..14fa351bbb 100644 --- a/src/crates/assembly/core/src/util/errors.rs +++ b/src/crates/assembly/core/src/util/errors.rs @@ -3,7 +3,8 @@ //! Provide unified error types and handling for the whole application use bitfun_core_types::errors::{ - ai_error_detail_from_message, classify_ai_error_message, AiErrorDetail, ErrorCategory, + ai_error_detail_from_message, classify_ai_error_message, AiErrorDetail, AiProviderError, + ErrorCategory, }; use serde::Serialize; use thiserror::Error; @@ -23,6 +24,15 @@ pub enum BitFunError { #[error("AI client error: {0}")] AIClient(String), + #[error("AI provider error: {0}")] + AIProvider(AiProviderError), + + /// A provider rejected the request before any effective model output because + /// the input exceeded its context window. The execution engine may compact + /// the session and retry the same logical model round. + #[error("AI client error: {0}")] + RecoverableContextOverflow(AiProviderError), + #[error("Session error: {0}")] Session(String), @@ -165,6 +175,8 @@ impl BitFunError { pub fn error_category(&self) -> ErrorCategory { match self { BitFunError::AIClient(msg) => classify_ai_error_message(msg), + BitFunError::AIProvider(error) => error.category.clone(), + BitFunError::RecoverableContextOverflow(_) => ErrorCategory::ContextOverflow, BitFunError::Timeout(_) => ErrorCategory::Timeout, BitFunError::Cancelled(_) => ErrorCategory::Unknown, _ => ErrorCategory::Unknown, @@ -173,16 +185,26 @@ impl BitFunError { /// Build a structured, provider-agnostic AI error detail for UI recovery. pub fn error_detail(&self) -> AiErrorDetail { + if let BitFunError::AIProvider(error) | BitFunError::RecoverableContextOverflow(error) = + self + { + return error.detail(); + } let category = self.error_category(); let message = self.to_string(); ai_error_detail_from_message(&message, category) } + + pub fn is_recoverable_context_overflow(&self) -> bool { + matches!(self, Self::RecoverableContextOverflow(_)) + } } impl From for BitFunError { fn from(error: bitfun_agent_stream::StreamProcessorError) -> Self { match error { bitfun_agent_stream::StreamProcessorError::AiClient(msg) => Self::AIClient(msg), + bitfun_agent_stream::StreamProcessorError::AiProvider(error) => Self::AIProvider(error), bitfun_agent_stream::StreamProcessorError::Cancelled(msg) => Self::Cancelled(msg), } } diff --git a/src/crates/contracts/core-types/src/errors.rs b/src/crates/contracts/core-types/src/errors.rs index ff9a17b92b..36d1d8dc6c 100644 --- a/src/crates/contracts/core-types/src/errors.rs +++ b/src/crates/contracts/core-types/src/errors.rs @@ -57,6 +57,123 @@ pub struct AiErrorDetail { pub action_hints: Vec, } +/// Provider failure normalized before it crosses adapter/runtime boundaries. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct AiProviderError { + pub message: String, + pub category: ErrorCategory, + #[serde(skip_serializing_if = "Option::is_none")] + pub provider: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub provider_code: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub http_status: Option, +} + +impl AiProviderError { + pub fn from_parts( + message: String, + provider: Option, + provider_code: Option, + http_status: Option, + ) -> Self { + let category = classify_ai_error_parts(&message, provider_code.as_deref(), http_status); + Self { + message, + category, + provider, + provider_code, + http_status, + } + } + + pub fn classified(message: String, category: ErrorCategory) -> Self { + Self { + message, + category, + provider: None, + provider_code: None, + http_status: None, + } + } + + pub fn detail(&self) -> AiErrorDetail { + AiErrorDetail { + category: self.category.clone(), + provider: self.provider.clone(), + provider_code: self.provider_code.clone(), + provider_message: Some(self.message.clone()), + request_id: extract_error_field(&self.message, "request_id"), + http_status: self.http_status, + retryable: Some(is_retryable_category(&self.category)), + action_hints: action_hints_for_category(&self.category), + } + } +} + +impl std::fmt::Display for AiProviderError { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.write_str(&self.message) + } +} + +impl std::error::Error for AiProviderError {} + +/// Prefer structured provider facts over text. Text classification remains a +/// fallback for providers that expose only opaque error strings. +pub fn classify_ai_error_parts( + message: &str, + provider_code: Option<&str>, + http_status: Option, +) -> ErrorCategory { + match http_status { + Some(401) => return ErrorCategory::Auth, + Some(402) => return ErrorCategory::ProviderQuota, + Some(403) => return ErrorCategory::Permission, + Some(429) => return ErrorCategory::RateLimit, + Some(500..=599) => return ErrorCategory::ProviderUnavailable, + _ => {} + } + + let code = provider_code + .unwrap_or_default() + .trim() + .to_ascii_lowercase(); + if matches!( + code.as_str(), + "context_length_exceeded" | "model_context_window_exceeded" | "context_window_exceeded" + ) { + return ErrorCategory::ContextOverflow; + } + match code.as_str() { + "insufficient_quota" | "quota_exceeded" => return ErrorCategory::ProviderQuota, + "rate_limit_exceeded" | "throttling_error" => return ErrorCategory::RateLimit, + "overloaded_error" | "server_is_overloaded" | "server_error" => { + return ErrorCategory::ProviderUnavailable; + } + "authentication_error" | "invalid_api_key" => return ErrorCategory::Auth, + "permission_error" => return ErrorCategory::Permission, + "content_filter" | "safety" => return ErrorCategory::ContentPolicy, + _ => {} + } + + let category = classify_ai_error_message(message); + if category != ErrorCategory::ModelError { + return category; + } + + if matches!(code.as_str(), "invalid_request_error" | "invalid_prompt") { + return ErrorCategory::InvalidRequest; + } + + if matches!(http_status, Some(400 | 404 | 409 | 413 | 422)) { + ErrorCategory::InvalidRequest + } else { + category + } +} + /// Classify an AI client error message into a structured category. pub fn classify_ai_error_message(msg: &str) -> ErrorCategory { let m = msg.to_lowercase(); @@ -92,9 +209,9 @@ pub fn classify_ai_error_message(msg: &str) -> ErrorCategory { "subscription expired", "plan expired", "套餐已到期", - "1309", ], - ) { + ) || contains_provider_code(&m, "1309") + { ErrorCategory::ProviderBilling } else if contains_any( &m, @@ -108,9 +225,9 @@ pub fn classify_ai_error_message(msg: &str) -> ErrorCategory { "error 503", "http 529", "error 529", - "1305", ], - ) { + ) || contains_provider_code(&m, "1305") + { ErrorCategory::ProviderUnavailable } else if contains_any( &m, @@ -120,26 +237,26 @@ pub fn classify_ai_error_message(msg: &str) -> ErrorCategory { "safety", "sensitive", "content_filter", - "1301", "api 调用被策略阻止", ], - ) { + ) || contains_provider_code(&m, "1301") + { ErrorCategory::ContentPolicy } else if m.contains("rate limit") - || m.contains("429") + || contains_http_status(&m, 429) || m.contains("too many requests") - || m.contains("1302") + || contains_provider_code(&m, "1302") || m.contains("concurrency") || m.contains("请求并发超额") { ErrorCategory::RateLimit } else if m.contains("authentication") - || m.contains("401") + || contains_http_status(&m, 401) || m.contains("invalid api key") || m.contains("incorrect api key") || m.contains("unauthorized") - || m.contains("1000") - || m.contains("1002") + || contains_provider_code(&m, "1000") + || contains_provider_code(&m, "1002") { ErrorCategory::Auth } else if contains_any( @@ -151,15 +268,11 @@ pub fn classify_ai_error_message(msg: &str) -> ErrorCategory { "not authorized", "no permission", "无权访问", - "1220", ], - ) { - ErrorCategory::Permission - } else if m.contains("context window") - || m.contains("token limit") - || m.contains("max_tokens") - || m.contains("context length") + ) || contains_provider_code(&m, "1220") { + ErrorCategory::Permission + } else if is_context_overflow_message(&m) { ErrorCategory::ContextOverflow } else if contains_any( &m, @@ -178,11 +291,11 @@ pub fn classify_ai_error_message(msg: &str) -> ErrorCategory { "error 413", "http 422", "error 422", - "1210", - "1211", - "435", ], - ) { + ) || contains_provider_code(&m, "1210") + || contains_provider_code(&m, "1211") + || contains_provider_code(&m, "435") + { ErrorCategory::InvalidRequest } else if m.contains("timeout") || m.contains("timed out") { ErrorCategory::Timeout @@ -215,6 +328,87 @@ fn contains_any(value: &str, needles: &[&str]) -> bool { needles.iter().any(|needle| value.contains(needle)) } +fn contains_provider_code(message: &str, code: &str) -> bool { + if message.trim() == code { + return true; + } + [ + format!("code={code}"), + format!("code: {code}"), + format!("\"code\":\"{code}"), + format!("\"code\":{code}"), + ] + .iter() + .any(|marker| contains_numeric_marker(message, marker)) +} + +fn contains_http_status(message: &str, status: u16) -> bool { + let status = status.to_string(); + message.trim_start().starts_with(&format!("{status} ")) + || [ + format!("http {status}"), + format!("error {status}"), + format!("status {status}"), + format!("{status} status code"), + ] + .iter() + .any(|marker| contains_numeric_marker(message, marker)) +} + +fn contains_numeric_marker(message: &str, marker: &str) -> bool { + message.match_indices(marker).any(|(start, _)| { + message[start + marker.len()..] + .chars() + .next() + .is_none_or(|character| !character.is_ascii_digit()) + }) +} + +fn is_context_overflow_message(message: &str) -> bool { + if contains_any( + message, + &[ + "finish_reason=max_tokens", + "max_output_tokens", + "maximum output token", + "output token limit", + "completion token limit", + "response truncated by model", + ], + ) { + return false; + } + + contains_any( + message, + &[ + "context_length_exceeded", + "model_context_window_exceeded", + "context window exceeded", + "context window exceeds limit", + "context length exceeded", + "context length is only", + "maximum context length", + "maximum prompt length is", + "maximum allowed input length", + "exceeds the context window", + "exceeds the available context size", + "greater than the context length", + "prompt is too long", + "prompt too long; exceeded", + "request_too_large", + "input is too long for requested model", + "tokens in request more than max tokens allowed", + "reduce the length of the messages", + "request entity too large", + ], + ) || ((message.contains("input") || message.contains("prompt")) + && message.contains("token") + && contains_any(message, &["exceed", "too long", "maximum", "limit"])) + || (message.contains("input length") && message.contains("context length")) + || (message.contains("prompt has") && message.contains("configured context size")) +} + fn is_retryable_category(category: &ErrorCategory) -> bool { matches!( category, @@ -303,7 +497,10 @@ fn extract_http_status(message: &str) -> Option { #[cfg(test)] mod tests { - use super::{ai_error_detail_from_message, classify_ai_error_message, ErrorCategory}; + use super::{ + ai_error_detail_from_message, classify_ai_error_message, classify_ai_error_parts, + AiProviderError, ErrorCategory, + }; #[test] fn classifies_quota_and_provider_unavailable_errors() { @@ -347,4 +544,107 @@ mod tests { vec!["wait_and_retry", "switch_model", "copy_diagnostics"] ); } + + #[test] + fn distinguishes_input_context_overflow_from_output_token_limits() { + assert_eq!( + classify_ai_error_message( + "context_length_exceeded: maximum context length is 200000 tokens" + ), + ErrorCategory::ContextOverflow + ); + assert_eq!( + classify_ai_error_message( + "The input token count exceeds the maximum number of tokens allowed" + ), + ErrorCategory::ContextOverflow + ); + assert_eq!( + classify_ai_error_message( + "response truncated by model output token limit (finish_reason=max_tokens)" + ), + ErrorCategory::ModelError + ); + } + + #[test] + fn classifies_common_provider_context_overflow_messages() { + for message in [ + "request_too_large", + "Input is too long for requested model", + "tokens in request more than max tokens allowed", + "Please reduce the length of the messages or completion", + "request entity too large", + "model_context_window_exceeded", + "context window exceeds limit", + "maximum prompt length is 200000", + "input length 210000 exceeds context length 200000", + "prompt has 210,000 tokens, but the configured context size is 200,000 tokens", + ] { + assert_eq!( + classify_ai_error_message(message), + ErrorCategory::ContextOverflow, + "message: {message}" + ); + } + } + + #[test] + fn structured_status_and_code_take_precedence_over_ambiguous_text() { + assert_eq!( + classify_ai_error_parts("Request failed", Some("context_length_exceeded"), Some(400)), + ErrorCategory::ContextOverflow + ); + assert_eq!( + classify_ai_error_parts( + "The prompt is too long for this model", + Some("invalid_request_error"), + Some(400) + ), + ErrorCategory::ContextOverflow + ); + assert_eq!( + classify_ai_error_parts("429 status code (no body)", None, Some(429)), + ErrorCategory::RateLimit + ); + assert_eq!( + classify_ai_error_parts("400 status code (no body)", None, Some(400)), + ErrorCategory::InvalidRequest + ); + assert_eq!( + classify_ai_error_parts("Service unavailable: token limit exceeded", None, Some(503)), + ErrorCategory::ProviderUnavailable + ); + assert_eq!( + classify_ai_error_message("Processed 429000 input tokens successfully"), + ErrorCategory::ModelError + ); + assert_eq!( + classify_ai_error_message("Observed status 401000 in token accounting"), + ErrorCategory::ModelError + ); + assert_eq!( + classify_ai_error_message("Provider error: code=1302, message=concurrency exceeded"), + ErrorCategory::RateLimit + ); + } + + #[test] + fn provider_error_preserves_structured_diagnostics() { + let error = AiProviderError::from_parts( + "Request failed".to_string(), + Some("openai".to_string()), + Some("context_length_exceeded".to_string()), + Some(400), + ); + + assert_eq!(error.category, ErrorCategory::ContextOverflow); + let detail = error.detail(); + assert_eq!(detail.provider.as_deref(), Some("openai")); + assert_eq!( + detail.provider_code.as_deref(), + Some("context_length_exceeded") + ); + assert_eq!(detail.http_status, Some(400)); + } } diff --git a/src/crates/execution/agent-stream/Cargo.toml b/src/crates/execution/agent-stream/Cargo.toml index ac23009ee1..f8c4771232 100644 --- a/src/crates/execution/agent-stream/Cargo.toml +++ b/src/crates/execution/agent-stream/Cargo.toml @@ -13,6 +13,7 @@ crate-type = ["rlib"] anyhow = { workspace = true } async-trait = { workspace = true } bitfun-events = { path = "../../contracts/events" } +bitfun-core-types = { path = "../../contracts/core-types" } futures = { workspace = true } bitfun-tool-call-jsonrepair = { path = "../tool-call-jsonrepair" } log = { workspace = true } diff --git a/src/crates/execution/agent-stream/src/lib.rs b/src/crates/execution/agent-stream/src/lib.rs index 7e8e08ea02..f10b2cad9c 100644 --- a/src/crates/execution/agent-stream/src/lib.rs +++ b/src/crates/execution/agent-stream/src/lib.rs @@ -12,6 +12,7 @@ use crate::tool_call_accumulator::{ FinalizedToolCall, PendingToolCalls, ToolCallBoundary, ToolCallFinalizeOptions, ToolCallStreamKey, }; +use bitfun_core_types::errors::AiProviderError; use bitfun_events::{AgenticEvent, AgenticEventPriority as EventPriority, ToolEventData}; use futures::{Stream, StreamExt}; pub use hidden_text::{HiddenTextBlock, HiddenTextStreamParser, HiddenTextTag}; @@ -57,6 +58,7 @@ impl ToolCall { #[derive(Debug, Clone, PartialEq, Eq)] pub enum StreamProcessorError { AiClient(String), + AiProvider(AiProviderError), Cancelled(String), } @@ -64,6 +66,7 @@ impl fmt::Display for StreamProcessorError { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { match self { Self::AiClient(msg) => write!(f, "AI client error: {}", msg), + Self::AiProvider(error) => write!(f, "AI provider error: {}", error), Self::Cancelled(msg) => write!(f, "Operation cancelled: {}", msg), } } @@ -1050,6 +1053,7 @@ impl StreamProcessor { break; } TimedStreamItem::Item(Err(e)) => { + let provider_error = e.downcast_ref::().cloned(); let error_msg = format!("Stream processing error: {}", e); error!("{}", error_msg); let non_recoverable_stream_error = @@ -1067,7 +1071,9 @@ impl StreamProcessor { flush_sse_on_error(&sse_collector, &error_msg).await; self.graceful_shutdown_from_ctx(&mut ctx, error_msg.clone()).await; return Err(StreamProcessError::new( - StreamProcessorError::AiClient(error_msg), + provider_error + .map(StreamProcessorError::AiProvider) + .unwrap_or(StreamProcessorError::AiClient(error_msg)), ctx.has_effective_output, )); } @@ -1218,10 +1224,11 @@ impl StreamProcessor { mod tests { use super::{ is_token_limit_finish_reason, GracefulShutdownInput, HiddenTextTag, SseLogCollector, - SseLogConfig, StreamEventSink, StreamProcessOptions, StreamProcessor, + SseLogConfig, StreamEventSink, StreamProcessOptions, StreamProcessor, StreamProcessorError, ToolArgumentRepairKind, ToolCall, ToolCallCompletion, }; use super::{UnifiedResponse, UnifiedTokenUsage, UnifiedToolCall}; + use bitfun_core_types::errors::{AiProviderError, ErrorCategory}; use bitfun_events::{AgenticEvent, AgenticEventPriority as EventPriority, ToolEventData}; use futures::StreamExt; use serde_json::json; @@ -1253,6 +1260,47 @@ mod tests { StreamProcessor::new(Arc::new(NoopEventSink)) } + #[tokio::test] + async fn preserves_structured_provider_error_from_adapter_stream() { + let processor = build_processor(); + let provider_error = AiProviderError::from_parts( + "Request failed".to_string(), + Some("openai".to_string()), + Some("context_length_exceeded".to_string()), + None, + ); + let stream = iter(vec![Err::(anyhow::Error::new( + provider_error, + ))]) + .boxed(); + + let result = processor + .process_stream( + stream, + None, + None, + "session_1".to_string(), + "turn_1".to_string(), + "round_1".to_string(), + "round_1:attempt:1".to_string(), + 1, + &CancellationToken::new(), + ) + .await + .expect_err("provider error should fail the stream"); + + match result.error { + StreamProcessorError::AiProvider(error) => { + assert_eq!(error.category, ErrorCategory::ContextOverflow); + assert_eq!( + error.provider_code.as_deref(), + Some("context_length_exceeded") + ); + } + error => panic!("unexpected stream error: {error:?}"), + } + } + #[tokio::test] async fn graceful_shutdown_emits_tool_cleanup_before_turn_cancellation() { let sink = Arc::new(RecordingEventSink::default());