From 01ba2c1aaa698d276e835fbdb2e3eebb68bcf6cb Mon Sep 17 00:00:00 2001 From: wsp Date: Tue, 15 Sep 2026 11:38:23 +0800 Subject: [PATCH] fix(agent): Unify model request retry policies Share a ten-attempt budget, error classification, and backoff between ordinary model requests and context compression. - Remove nested retries from compression requests. - Retry empty summaries and tool-call responses within the same budget, including malformed tool calls accompanied by text. - Replan immediately on context overflow, granting each new plan ten attempts and limiting compression to four plans. - Remove local summary reconstruction and its configuration. - Fail the turn when compression fails instead of continuing with uncompressed context or an incomplete local summary. - Preserve existing context and compression counts on failure. - Cover retry budgets, replanning, cancellation, and state retention. --- src/apps/desktop/src/lib.rs | 2 +- src/crates/adapters/ai-adapters/src/client.rs | 70 ++-- .../src/client/response_aggregator.rs | 20 +- src/crates/assembly/core/Cargo.toml | 2 +- .../src/agentic/coordination/coordinator.rs | 8 +- .../src/agentic/coordination/scheduler.rs | 4 +- .../core/src/agentic/execution/AGENTS.md | 9 + .../agentic/execution/compression_request.rs | 263 +++++++++++++++ .../agentic/execution/compression_tests.rs | 227 +++++++++++++ .../src/agentic/execution/execution_engine.rs | 261 ++++----------- .../core/src/agentic/execution/mod.rs | 1 + .../src/agentic/execution/round_executor.rs | 74 +---- .../agentic/session/compression/compressor.rs | 245 +++----------- .../session/compression/fallback/builder.rs | 183 ----------- .../session/compression/fallback/mod.rs | 41 --- .../session/compression/fallback/payload.rs | 192 ----------- .../session/compression/fallback/render.rs | 104 ------ .../session/compression/fallback/sanitize.rs | 305 ------------------ .../session/compression/fallback/tests.rs | 261 --------------- .../session/compression/fallback/types.rs | 39 --- .../src/agentic/session/compression/mod.rs | 2 - .../assembly/core/src/agentic/system.rs | 2 +- .../assembly/core/src/product_runtime.rs | 13 +- src/crates/execution/agent-stream/src/lib.rs | 1 + .../execution/agent-stream/src/retry.rs | 96 ++++++ 25 files changed, 779 insertions(+), 1646 deletions(-) create mode 100644 src/crates/assembly/core/src/agentic/execution/compression_request.rs create mode 100644 src/crates/assembly/core/src/agentic/execution/compression_tests.rs delete mode 100644 src/crates/assembly/core/src/agentic/session/compression/fallback/builder.rs delete mode 100644 src/crates/assembly/core/src/agentic/session/compression/fallback/mod.rs delete mode 100644 src/crates/assembly/core/src/agentic/session/compression/fallback/payload.rs delete mode 100644 src/crates/assembly/core/src/agentic/session/compression/fallback/render.rs delete mode 100644 src/crates/assembly/core/src/agentic/session/compression/fallback/sanitize.rs delete mode 100644 src/crates/assembly/core/src/agentic/session/compression/fallback/tests.rs delete mode 100644 src/crates/assembly/core/src/agentic/session/compression/fallback/types.rs create mode 100644 src/crates/execution/agent-stream/src/retry.rs diff --git a/src/apps/desktop/src/lib.rs b/src/apps/desktop/src/lib.rs index 0b40c85f0f..d59bcb3493 100644 --- a/src/apps/desktop/src/lib.rs +++ b/src/apps/desktop/src/lib.rs @@ -1976,7 +1976,7 @@ async fn init_agentic_system() -> anyhow::Result<( let persistence_manager = Arc::new(persistence::PersistenceManager::new(path_manager.clone())?); let context_store = Arc::new(session::SessionContextStore::new()); - let context_compressor = Arc::new(session::ContextCompressor::new(Default::default())); + let context_compressor = Arc::new(session::ContextCompressor::new()); let session_manager = Arc::new(session::SessionManager::new( context_store, diff --git a/src/crates/adapters/ai-adapters/src/client.rs b/src/crates/adapters/ai-adapters/src/client.rs index c4e601216f..fab5451090 100644 --- a/src/crates/adapters/ai-adapters/src/client.rs +++ b/src/crates/adapters/ai-adapters/src/client.rs @@ -22,18 +22,13 @@ use crate::types::*; use anyhow::Result; use format::ApiFormat; use log::warn; -use openbitfun_core_types::errors::{AiProviderError, ErrorCategory}; +use openbitfun_core_types::errors::AiProviderError; use reqwest::Client; use std::time::Duration; use tokio::sync::mpsc; -const SEND_MESSAGE_STREAM_ATTEMPTS: usize = 10; +const SEND_MESSAGE_STREAM_ATTEMPTS: usize = openbitfun_agent_stream::retry::MAX_MODEL_ATTEMPTS; const TEST_CONNECTION_STREAM_ATTEMPTS: usize = 5; -const SEND_MESSAGE_RETRY_BASE_DELAY_MS: u64 = 500; -const SEND_MESSAGE_RATE_LIMIT_RETRY_BASE_DELAY_MS: u64 = 2_000; -const SEND_MESSAGE_MAX_EXPONENTIAL_DELAY_MS: u64 = 30_000; -const SEND_MESSAGE_MAX_RATE_LIMIT_DELAY_MS: u64 = 60_000; -const SEND_MESSAGE_MAX_RETRY_EXPONENT_SHIFT: u32 = 6; /// Streamed response result with the parsed stream and optional raw SSE receiver. pub struct StreamResponse { @@ -360,6 +355,37 @@ impl AIClient { .await } + /// Aggregate exactly one provider attempt; the caller owns retries and validation. + pub async fn send_message_once_with_trace_and_request_context( + &self, + messages: Vec, + tools: Option>, + request_context: Option, + trace: Option, + ) -> Result { + let response = self + .send_message_stream_once_with_request_context( + messages, + tools, + request_context, + trace.clone(), + ) + .await?; + let handle = response.trace_handle.clone(); + match response_aggregator::aggregate_stream_response_preserving_tool_presence(response) + .await + { + Ok(response) => { + complete_aggregated_trace(trace.as_ref(), handle.as_ref(), &response).await; + Ok(response) + } + Err(error) => { + fail_aggregated_trace(trace.as_ref(), handle.as_ref(), &error.to_string()).await; + Err(error) + } + } + } + pub async fn send_message_with_extra_body_and_trace( &self, messages: Vec, @@ -578,35 +604,7 @@ fn send_message_retry_delay_ms_with_provider( error_message: &str, provider_error: Option<&AiProviderError>, ) -> u64 { - let shift = u32::try_from(attempt_index) - .unwrap_or(u32::MAX) - .min(SEND_MESSAGE_MAX_RETRY_EXPONENT_SHIFT); - let msg = error_message.to_lowercase(); - let is_rate_limit = provider_error - .is_some_and(|error| error.category == ErrorCategory::RateLimit) - || msg.contains("429") - || msg.contains("rate limit") - || msg.contains("too many requests"); - - let fallback = if is_rate_limit { - SEND_MESSAGE_RATE_LIMIT_RETRY_BASE_DELAY_MS - .saturating_mul(1u64 << shift) - .min(SEND_MESSAGE_MAX_RATE_LIMIT_DELAY_MS) - } else { - SEND_MESSAGE_RETRY_BASE_DELAY_MS - .saturating_mul(1u64 << shift) - .min(SEND_MESSAGE_MAX_EXPONENTIAL_DELAY_MS) - }; - - match provider_error.and_then(|error| error.retry_after_ms) { - Some(retry_after_ms) if is_rate_limit => retry_after_ms - .max(fallback) - .min(SEND_MESSAGE_MAX_RATE_LIMIT_DELAY_MS), - Some(retry_after_ms) if retry_after_ms > 0 => { - retry_after_ms.min(SEND_MESSAGE_MAX_RATE_LIMIT_DELAY_MS) - } - Some(_) | None => fallback, - } + openbitfun_agent_stream::retry::delay_ms(attempt_index, error_message, provider_error) } async fn complete_aggregated_trace( diff --git a/src/crates/adapters/ai-adapters/src/client/response_aggregator.rs b/src/crates/adapters/ai-adapters/src/client/response_aggregator.rs index aeb8868ed5..51bfc47026 100644 --- a/src/crates/adapters/ai-adapters/src/client/response_aggregator.rs +++ b/src/crates/adapters/ai-adapters/src/client/response_aggregator.rs @@ -10,6 +10,22 @@ use super::StreamResponse; pub(crate) async fn aggregate_stream_response( stream_response: StreamResponse, ) -> Result { + aggregate_stream_response_inner(stream_response, false).await +} + +/// Preserve even malformed tool-call presence so summary callers cannot accept +/// text accompanied by tool deltas that the aggregator could not reconstruct. +pub(crate) async fn aggregate_stream_response_preserving_tool_presence( + stream_response: StreamResponse, +) -> Result { + aggregate_stream_response_inner(stream_response, true).await +} + +async fn aggregate_stream_response_inner( + stream_response: StreamResponse, + preserve_tool_presence: bool, +) -> Result { + let mut saw_tool_call = false; let mut stream = stream_response.stream; let mut full_text = String::new(); @@ -46,6 +62,7 @@ pub(crate) async fn aggregate_stream_response( } if let Some(tool_call) = tool_call { + saw_tool_call = true; let crate::stream::UnifiedToolCall { tool_call_index, id, @@ -153,7 +170,8 @@ pub(crate) async fn aggregate_stream_response( Ok(GeminiResponse { text: full_text, reasoning_content: (!full_reasoning.is_empty()).then_some(full_reasoning), - tool_calls: (!tool_calls.is_empty()).then_some(tool_calls), + tool_calls: (!tool_calls.is_empty() || (preserve_tool_presence && saw_tool_call)) + .then_some(tool_calls), usage, finish_reason, provider_metadata, diff --git a/src/crates/assembly/core/Cargo.toml b/src/crates/assembly/core/Cargo.toml index d74adaf5fc..2d32c81754 100644 --- a/src/crates/assembly/core/Cargo.toml +++ b/src/crates/assembly/core/Cargo.toml @@ -468,7 +468,7 @@ ssh-remote = [ tokio-tungstenite = { workspace = true } openbitfun-runtime-services = { path = "../../execution/runtime-services", features = ["test-support"] } tempfile = { workspace = true } -tokio = { workspace = true, features = ["macros", "rt-multi-thread"] } +tokio = { workspace = true, features = ["macros", "rt-multi-thread", "test-util"] } [[test]] name = "git_contracts" diff --git a/src/crates/assembly/core/src/agentic/coordination/coordinator.rs b/src/crates/assembly/core/src/agentic/coordination/coordinator.rs index 1c7fbb3211..15c14b1c1f 100644 --- a/src/crates/assembly/core/src/agentic/coordination/coordinator.rs +++ b/src/crates/assembly/core/src/agentic/coordination/coordinator.rs @@ -15095,9 +15095,9 @@ mod tests { use crate::agentic::goal_mode::thread_goal_patch; use crate::agentic::persistence::PersistenceManager; use crate::agentic::session::{ - compression::{CompressionConfig, ContextCompressor}, - PromptCachePolicy, SessionContextStore, SessionManager, SessionManagerConfig, - SystemPromptCacheIdentity, UserContextCacheIdentity, TEST_MODEL_RESOLUTION_AI_CONFIG, + compression::ContextCompressor, PromptCachePolicy, SessionContextStore, SessionManager, + SessionManagerConfig, SystemPromptCacheIdentity, UserContextCacheIdentity, + TEST_MODEL_RESOLUTION_AI_CONFIG, }; use crate::agentic::skill_agent_snapshot::SkillSnapshotEntry; use crate::agentic::tools::framework::{ @@ -16722,7 +16722,7 @@ mod tests { )), event_queue.clone(), session_manager.clone(), - Arc::new(ContextCompressor::new(CompressionConfig::default())), + Arc::new(ContextCompressor::new()), ExecutionEngineConfig::default(), )); let coordinator = ConversationCoordinator::new_with_coordination_database_file( diff --git a/src/crates/assembly/core/src/agentic/coordination/scheduler.rs b/src/crates/assembly/core/src/agentic/coordination/scheduler.rs index f932f203a7..a71d498720 100644 --- a/src/crates/assembly/core/src/agentic/coordination/scheduler.rs +++ b/src/crates/assembly/core/src/agentic/coordination/scheduler.rs @@ -3148,7 +3148,7 @@ mod tests { }; use crate::agentic::persistence::PersistenceManager; use crate::agentic::session::{ - compression::{CompressionConfig, ContextCompressor}, + compression::ContextCompressor, revert::{SessionRevertPhase, SessionRevertState, SESSION_REVERT_SCHEMA_VERSION}, PromptCachePolicy, SessionContextStore, SessionManagerConfig, TEST_MODEL_RESOLUTION_AI_CONFIG, @@ -3259,7 +3259,7 @@ mod tests { )), event_queue.clone(), session_manager.clone(), - Arc::new(ContextCompressor::new(CompressionConfig::default())), + Arc::new(ContextCompressor::new()), ExecutionEngineConfig::default(), )); let coordinator = Arc::new(ConversationCoordinator::new( diff --git a/src/crates/assembly/core/src/agentic/execution/AGENTS.md b/src/crates/assembly/core/src/agentic/execution/AGENTS.md index f92420edfe..ffa3878545 100644 --- a/src/crates/assembly/core/src/agentic/execution/AGENTS.md +++ b/src/crates/assembly/core/src/agentic/execution/AGENTS.md @@ -24,6 +24,15 @@ to name a trusted Bash 4+ executable and uses only isolated synthetic commands. For automatic/manual context compaction cancellation, preparation, and commit races, use: +Each fixed summary plan has ten total model attempts, shared across transport, +stream aggregation, empty summaries, and tool-call responses. Context overflow +immediately starts a smaller plan with a fresh budget (at most four plans). +Both summary and ordinary model requests use the provider-neutral policy in +`openbitfun-agent-stream::retry` and single-attempt adapter entry points. +Failed compression terminates the turn without replacing context; there is no +locally reconstructed summary fallback. Historical compression payloads remain +readable. + ```bash cargo test --locked -p openbitfun-core --no-default-features --features agent-runtime,git --lib compression cargo test --locked -p openbitfun-core --no-default-features --features agent-runtime,git --lib compaction diff --git a/src/crates/assembly/core/src/agentic/execution/compression_request.rs b/src/crates/assembly/core/src/agentic/execution/compression_request.rs new file mode 100644 index 0000000000..28c6e2e495 --- /dev/null +++ b/src/crates/assembly/core/src/agentic/execution/compression_request.rs @@ -0,0 +1,263 @@ +//! Model summary attempts for one fixed compression plan. + +use crate::agentic::session::ContextCompressor; +use crate::infrastructure::ai::AIClient; +use crate::util::errors::{OpenBitFunError, OpenBitFunResult}; +use openbitfun_agent_stream::retry::{delay_ms, should_retry, MAX_MODEL_ATTEMPTS}; +use openbitfun_ai_adapters::{ + Message, ModelExchangeTraceConfig, ModelRequestContext, ToolDefinition, +}; +use openbitfun_core_types::errors::{AiProviderError, ErrorCategory}; +use std::sync::Arc; + +pub(super) async fn request_summary( + client: Arc, + messages: Vec, + tools: Option>, + context: &ModelRequestContext, + trace: Option, +) -> OpenBitFunResult { + request_summary_with(|| { + client.send_message_once_with_trace_and_request_context( + messages.clone(), + tools.clone(), + Some(context.clone()), + trace.clone(), + ) + }) + .await +} + +async fn request_summary_with(mut request: F) -> OpenBitFunResult +where + F: FnMut() -> Fut, + Fut: std::future::Future>, +{ + for attempt in 0..MAX_MODEL_ATTEMPTS { + let result = request().await; + let mut error = match result { + Ok(response) => { + if response.tool_calls.is_some() { + AiProviderError::classified( + "Compression request returned tool calls instead of a summary".to_string(), + ErrorCategory::ModelError, + ) + } else if let Some(summary) = + ContextCompressor::normalize_model_summary_output(&response.text) + { + return Ok(summary); + } else { + AiProviderError::classified( + "Compression request returned an empty summary".to_string(), + ErrorCategory::ModelError, + ) + } + } + Err(error) => error + .downcast_ref::() + .cloned() + .unwrap_or_else(|| { + AiProviderError::from_parts(format!("{error:#}"), None, None, None) + }), + }; + if !should_retry(&error.category) || attempt + 1 == MAX_MODEL_ATTEMPTS { + error.message = format!( + "Compression summary failed after {} attempts: {}", + attempt + 1, + error.message + ); + return Err(if error.category == ErrorCategory::ContextOverflow { + OpenBitFunError::RecoverableContextOverflow(error) + } else { + OpenBitFunError::AIProvider(error) + }); + } + let wait_ms = delay_ms(attempt, &error.message, Some(&error)); + log::warn!( + "Retrying compression summary: attempt={}/{}, delay_ms={}, category={:?}, error={}", + attempt + 1, + MAX_MODEL_ATTEMPTS, + wait_ms, + error.category, + error + ); + // The entire preparation future, including aggregation and waits, is + // cancelled by prepare_compression_cancellable before context commit. + tokio::time::sleep(std::time::Duration::from_millis(wait_ms)).await; + } + unreachable!("summary attempts return success or a terminal error") +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::agentic::execution::round_executor::tests::{retry_test_success, RetryTestServer}; + use std::collections::VecDeque; + + fn empty() -> anyhow::Result { + Ok(openbitfun_ai_adapters::GeminiResponse { + text: " \n\t".into(), + reasoning_content: None, + tool_calls: None, + usage: None, + finish_reason: None, + provider_metadata: None, + }) + } + + #[tokio::test(start_paused = true)] + async fn compression_invalid_results_and_transport_errors_share_ten_attempts() { + let mut responses = VecDeque::new(); + for _ in 0..3 { + responses.push_back(empty()); + } + for _ in 0..3 { + responses.push_back(Err(AiProviderError::classified( + "network failure".into(), + ErrorCategory::Network, + ) + .into())); + } + for _ in 0..4 { + let mut response = empty().unwrap(); + response.text = "Text accompanying a tool call must not become a summary".into(); + response.tool_calls = Some(vec![]); + responses.push_back(Ok(response)); + } + let mut requests = 0; + let error = request_summary_with(|| { + requests += 1; + std::future::ready(responses.pop_front().expect("must not exceed budget")) + }) + .await + .unwrap_err(); + assert_eq!(requests, 10); + assert!(error.to_string().contains("10 attempts")); + assert!(error.to_string().contains("tool calls")); + assert_eq!(error.error_category(), ErrorCategory::ModelError); + } + + #[tokio::test(start_paused = true)] + async fn compression_overflow_on_last_attempt_allows_fresh_request_budget() { + let mut requests = 0; + let error = request_summary_with(|| { + requests += 1; + std::future::ready(if requests == 10 { + Err( + AiProviderError::classified("too large".into(), ErrorCategory::ContextOverflow) + .into(), + ) + } else { + empty() + }) + }) + .await + .unwrap_err(); + assert!(error.is_recoverable_context_overflow()); + assert_eq!(requests, 10); + + let mut requests = 0; + let summary = request_summary_with(|| { + requests += 1; + let mut response = empty().unwrap(); + if requests == 10 { + response.text = " Recovered ".into(); + } + std::future::ready(Ok(response)) + }) + .await + .unwrap(); + assert_eq!(requests, 10); + assert_eq!(summary, "Recovered"); + } + + #[tokio::test(start_paused = true)] + async fn compression_cancellation_interrupts_backoff() { + let token = tokio_util::sync::CancellationToken::new(); + let mut requests = 0; + let preparation = request_summary_with(|| { + requests += 1; + std::future::ready(empty()) + }); + let work = crate::agentic::execution::execution_engine::prepare_compression_cancellable( + &token, + preparation, + ); + let cancel = async { + tokio::time::sleep(std::time::Duration::from_millis(100)).await; + token.cancel(); + }; + let (result, _) = tokio::join!(work, cancel); + assert!(matches!(result, Err(OpenBitFunError::Cancelled(_)))); + assert_eq!(requests, 1); + } + + #[tokio::test] + async fn compression_http_rejections_and_overflow_do_not_retry_in_adapter() { + for (status, code, category) in [ + (401, "invalid_api_key", ErrorCategory::Auth), + (403, "permission_error", ErrorCategory::Permission), + (413, "invalid_request_error", ErrorCategory::InvalidRequest), + (402, "insufficient_quota", ErrorCategory::ProviderQuota), + ( + 400, + "context_length_exceeded", + ErrorCategory::ContextOverflow, + ), + ] { + let server = RetryTestServer::new(vec![ + ( + status, + serde_json::json!({"error": {"code": code, "message": code}}).to_string(), + ), + retry_test_success(), + ]); + let error = request_summary( + server.client(), + vec![Message::user("Summarize".into())], + None, + &ModelRequestContext::default(), + None, + ) + .await + .unwrap_err(); + assert_eq!(error.error_category(), category); + assert_eq!(server.requests.lock().unwrap().len(), 1); + } + } + + #[tokio::test] + async fn compression_http_retries_empty_and_malformed_tool_output() { + let invalid_tool = ( + 200, + format!( + "data: {}\n\ndata: [DONE]\n\n", + serde_json::json!({ + "id":"invalid", "object":"chat.completion.chunk", "created":1, "model":"retry-test-model", + "choices":[{"index":0,"delta":{"content":"Ignore this text", "tool_calls":[ + {"index":0,"id":"call-1","type":"function","function":{"name":"Read","arguments":"{"}} + ]},"finish_reason":"tool_calls"}] + }) + ), + ); + let server = RetryTestServer::new(vec![ + (200, "data: [DONE]\n\n".into()), + invalid_tool, + retry_test_success(), + ]); + let summary = request_summary( + server.client(), + vec![Message::user("Summarize".into())], + None, + &ModelRequestContext::default(), + None, + ) + .await + .unwrap(); + assert_eq!(summary, "Recovered"); + let requests = server.requests.lock().unwrap(); + assert_eq!(requests.len(), 3); + assert_eq!(requests[0], requests[1]); + assert_eq!(requests[1], requests[2]); + } +} diff --git a/src/crates/assembly/core/src/agentic/execution/compression_tests.rs b/src/crates/assembly/core/src/agentic/execution/compression_tests.rs new file mode 100644 index 0000000000..b321fbc16f --- /dev/null +++ b/src/crates/assembly/core/src/agentic/execution/compression_tests.rs @@ -0,0 +1,227 @@ +use super::*; +use crate::agentic::events::EventQueueConfig; +use crate::agentic::execution::round_executor::tests::{ + retry_test_success, test_round_executor, RetryTestServer, +}; +use crate::agentic::persistence::PersistenceManager; +use crate::agentic::session::{SessionContextStore, SessionManagerConfig}; + +fn engine() -> (tempfile::TempDir, ExecutionEngine) { + let temp = tempfile::Builder::new() + .prefix("compression-") + .tempdir() + .unwrap(); + let manager = SessionManager::new( + Arc::new(SessionContextStore::new()), + Arc::new( + PersistenceManager::new(Arc::new( + crate::infrastructure::PathManager::with_user_root_for_tests(temp.path().into()), + )) + .unwrap(), + ), + SessionManagerConfig { + enable_persistence: false, + ..Default::default() + }, + ); + ( + temp, + ExecutionEngine::new( + Arc::new(test_round_executor()), + Arc::new(EventQueue::new(EventQueueConfig::default())), + Arc::new(manager), + Arc::new(ContextCompressor::new()), + ExecutionEngineConfig::default(), + ), + ) +} + +fn messages() -> Vec { + let mut messages = vec![Message::system("system".into())]; + for index in 0..8 { + messages.push(Message::user(format!( + "Request {index}: {}", + "old context ".repeat(4_000) + ))); + messages.push(Message::assistant(format!( + "Answer {index}: {}", + "earlier work ".repeat(4_000) + ))); + } + messages +} + +fn overflow() -> (u16, String) { + ( + 400, + serde_json::json!({"error": { + "code": "context_length_exceeded", "message": "context length exceeded" + }}) + .to_string(), + ) +} + +#[tokio::test] +async fn compression_overflow_replans_with_smaller_input() { + let (_temp, engine) = engine(); + let server = RetryTestServer::new(vec![overflow(), retry_test_success()]); + let result = engine + .build_planned_compression_result( + "session", + "turn", + &messages(), + 128_000, + None, + server.client(), + &ModelRequestContext::default(), + &None, + &PrependedPromptReminders::default(), + false, + None, + None, + None, + ) + .await + .unwrap() + .unwrap(); + assert!(result + .messages + .iter() + .any(|message| message.content.to_string().contains("Recovered"))); + let requests = server.requests.lock().unwrap(); + assert_eq!(requests.len(), 2); + assert!( + requests[1]["messages"].as_array().unwrap().len() + < requests[0]["messages"].as_array().unwrap().len() + ); +} + +#[tokio::test] +async fn compression_overflow_exhausts_four_plans_without_local_summary() { + let (_temp, engine) = engine(); + let server = RetryTestServer::new(vec![overflow()]); + let error = engine + .build_planned_compression_result( + "session", + "turn", + &messages(), + 128_000, + None, + server.client(), + &ModelRequestContext::default(), + &None, + &PrependedPromptReminders::default(), + false, + None, + None, + None, + ) + .await + .unwrap_err(); + assert!(error.to_string().contains("plan 4"), "{error}"); + assert_eq!(server.requests.lock().unwrap().len(), 4); +} + +#[tokio::test] +async fn compression_failure_preserves_context_and_success_count() { + let (_temp, engine) = engine(); + let session = engine + .session_manager + .create_session( + "compression test".into(), + "Standard".into(), + crate::agentic::core::SessionConfig { + workspace_path: Some(_temp.path().to_string_lossy().into_owned()), + ..Default::default() + }, + ) + .await + .unwrap(); + let messages = messages(); + engine + .session_manager + .replace_context_messages(&session.session_id, messages.clone()) + .await; + let before = serde_json::to_value( + engine + .session_manager + .get_context_messages(&session.session_id) + .await + .unwrap(), + ) + .unwrap(); + let before_state = serde_json::to_value( + engine + .session_manager + .get_session(&session.session_id) + .unwrap() + .compression_state, + ) + .unwrap(); + let server = RetryTestServer::new(vec![( + 401, + serde_json::json!({ + "error": {"code": "invalid_api_key", "message": "invalid api key"} + }) + .to_string(), + )]); + let pressure = ExecutionEngine::estimate_auto_compression_pressure( + &messages, + None, + 128_000, + ExecutionEngine::compression_trigger_budget(128_000, None), + 0, + ); + let error = engine + .compress_messages( + &session.session_id, + "turn", + "auto", + messages, + pressure, + 128_000, + server.client(), + &ModelRequestContext::default(), + &None, + Message::system("system".into()), + &PrependedPromptReminders::default(), + false, + 10_000, + None, + None, + ) + .await + .unwrap_err(); + assert!(error.to_string().contains("plan 1")); + assert_eq!(server.requests.lock().unwrap().len(), 1); + assert_eq!( + before, + serde_json::to_value( + engine + .session_manager + .get_context_messages(&session.session_id) + .await + .unwrap() + ) + .unwrap() + ); + assert_eq!( + before_state, + serde_json::to_value( + engine + .session_manager + .get_session(&session.session_id) + .unwrap() + .compression_state + ) + .unwrap() + ); + let events = engine.event_queue.dequeue_batch(20).await; + assert!(events + .iter() + .any(|event| matches!(event.event, AgenticEvent::ContextCompressionFailed { .. }))); + assert!(!events.iter().any(|event| matches!( + event.event, + AgenticEvent::ContextCompressionCompleted { .. } + ))); +} 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 55a7f0ddca..c04d2d0baa 100644 --- a/src/crates/assembly/core/src/agentic/execution/execution_engine.rs +++ b/src/crates/assembly/core/src/agentic/execution/execution_engine.rs @@ -269,6 +269,20 @@ pub(crate) async fn prepare_compression_cancellable( result } +fn compression_plan_error(error: OpenBitFunError, plan: usize) -> OpenBitFunError { + match error { + OpenBitFunError::AIProvider(mut error) + | OpenBitFunError::RecoverableContextOverflow(mut error) => { + error.message = format!( + "Context compression failed on plan {plan}: {}", + error.message + ); + OpenBitFunError::AIProvider(error) + } + error => error, + } +} + fn manual_compaction_terminal_error(error: OpenBitFunError) -> OpenBitFunError { match error { error @ OpenBitFunError::Cancelled(_) => error, @@ -2314,100 +2328,10 @@ impl ExecutionEngine { Ok(compression_messages) } - async fn request_compression_summary_with_retry( - &self, - ai_client: Arc, - request_messages: Vec, - tool_definitions: Option>, - model_request_context: &ModelRequestContext, - trace_config: Option, - max_tries: usize, - ) -> OpenBitFunResult { - let mut last_error = None; - let base_wait_time_ms = 500; - - for attempt in 0..max_tries { - let result = ai_client - .send_message_with_trace_and_request_context( - request_messages.clone(), - tool_definitions.clone(), - Some(model_request_context.clone()), - trace_config.clone(), - ) - .await; - - match result { - Ok(response) => { - if response.tool_calls.is_some() { - return Err(OpenBitFunError::AIClient( - "Compression request returned tool calls instead of a summary" - .to_string(), - )); - } - if attempt > 0 { - debug!( - "Compression summary generation succeeded (attempt {}/{})", - attempt + 1, - max_tries - ); - } - 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_msg - ); - let category = provider_error - .as_ref() - .map(|error| error.category.clone()) - .unwrap_or_else(|| { - openbitfun_core_types::errors::classify_ai_error_message(&err_msg) - }); - if category == openbitfun_core_types::errors::ErrorCategory::ContextOverflow { - return Err(OpenBitFunError::RecoverableContextOverflow( - provider_error.unwrap_or_else(|| { - openbitfun_core_types::errors::AiProviderError::classified( - err_msg, - openbitfun_core_types::errors::ErrorCategory::ContextOverflow, - ) - }), - )); - } - last_error = Some(err); - - if attempt < max_tries - 1 { - let delay_ms = base_wait_time_ms * (1 << attempt.min(3)); - debug!( - "Waiting {}ms before compression summary retry {}...", - delay_ms, - attempt + 2 - ); - tokio::time::sleep(std::time::Duration::from_millis(delay_ms)).await; - } - } - } - } - - Err(OpenBitFunError::AIClient(format!( - "Compression summary generation failed after {} attempts: {}", - max_tries, - last_error - .map(|err| err.to_string()) - .unwrap_or_else(|| "Unknown error".to_string()) - ))) - } - async fn generate_compression_model_summary( &self, input: CompressionModelSummaryInput<'_>, - ) -> OpenBitFunResult> { + ) -> OpenBitFunResult { let request_messages = self .build_compression_request_messages( input.runtime_messages, @@ -2420,23 +2344,14 @@ impl ExecutionEngine { ) .await?; - let raw_summary = self - .request_compression_summary_with_retry( - input.ai_client, - request_messages, - input.tool_definitions.clone(), - input.model_request_context, - input.trace_config, - 2, - ) - .await?; - let summary = - ContextCompressor::normalize_model_summary_output(&raw_summary).ok_or_else(|| { - OpenBitFunError::AIClient( - "Model-based compression returned an empty summary".to_string(), - ) - })?; - Ok(Some(summary)) + super::compression_request::request_summary( + input.ai_client, + request_messages, + input.tool_definitions.clone(), + input.model_request_context, + input.trace_config, + ) + .await } #[allow(clippy::too_many_arguments)] @@ -2459,8 +2374,7 @@ impl ExecutionEngine { 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 selected_plan = None; - let mut model_summary = None; + let mut previous_cutoff = None; for attempt in 0..Self::MAX_COMPRESSION_OVERFLOW_ATTEMPTS { let Some(plan) = self.context_compressor.plan_compression_for_turn( @@ -2471,8 +2385,16 @@ impl ExecutionEngine { recent_target, )? else { - break; + return Err(OpenBitFunError::AIClient( + "Context compression has no eligible plan".to_string(), + )); }; + if previous_cutoff.is_some_and(|cutoff| plan.cutoff_message_index >= cutoff) { + return Err(OpenBitFunError::AIClient( + "Context compression cannot reduce the summary input further".to_string(), + )); + } + previous_cutoff = Some(plan.cutoff_message_index); info!( "Compression context plan: session_id={}, turn_id={}, attempt={}/{}, retained_user_token_budget={}, retained_user_tokens={}, retained_user_messages={}, recent_target_tokens={}, recent_tail_tokens={}, cutoff_message_index={}, summary_messages={}, recent_tail_messages={}", session_id, @@ -2506,9 +2428,15 @@ impl ExecutionEngine { match summary_result { Ok(summary) => { - selected_plan = Some(plan); - model_summary = summary; - break; + return self + .context_compressor + .compress_plan_with_contract( + session_id, + plan, + compression_contract, + summary, + ) + .map(Some); } Err(err) if err.is_recoverable_context_overflow() => { warn!( @@ -2525,39 +2453,19 @@ impl ExecutionEngine { let can_retry = attempt + 1 < Self::MAX_COMPRESSION_OVERFLOW_ATTEMPTS && plan.next_recent_target_tokens.is_some(); let next_recent_target = plan.next_recent_target_tokens; - selected_plan = Some(plan); if can_retry { recent_target = recent_target .saturating_add(ContextCompressor::RECENT_CONTEXT_RETRY_STEP_TOKENS) .max(next_recent_target.expect("retry target checked above")); continue; } - break; - } - Err(err @ OpenBitFunError::Cancelled(_)) => return Err(err), - Err(err) => { - warn!( - "Model-based compression failed, falling back to structured local compression: {}", - err - ); - selected_plan = Some(plan); - break; + return Err(compression_plan_error(err, attempt + 1)); } + Err(err) => return Err(compression_plan_error(err, attempt + 1)), } } - let Some(selected_plan) = selected_plan else { - return Ok(None); - }; - self.context_compressor - .compress_plan_with_contract( - session_id, - context_window, - selected_plan, - compression_contract, - model_summary, - ) - .map(Some) + unreachable!("compression planning returns a result or terminal error") } async fn resolve_compression_runtime_scaffold( @@ -2989,11 +2897,7 @@ impl ExecutionEngine { prepended_reminder_tokens, ); let compressed_tokens = after_pressure.total_tokens; - let summary_source = if compression_result.has_model_summary { - "model" - } else { - "local_fallback" - }; + let summary_source = "model"; info!( "Compression completed: session_id={}, turn_id={}, messages {} -> {}, total_tokens {} -> {}, system_tokens {} -> {}, tool_tokens {} -> {}, prepended_reminder_tokens {} -> {}, conversation_tokens {} -> {}, context_window={}, input_limit={}, output_reserve={}, safety_reserve={}, usage {:.3} -> {:.3}, compression_count={}, duration_ms={}, summary_source={}", @@ -3037,7 +2941,7 @@ impl ExecutionEngine { (compressed_tokens as f64) / (before_pressure.total_tokens as f64) }, duration_ms, - has_summary: compression_result.has_model_summary, + has_summary: true, summary_source: summary_source.to_string(), applied: true, }, @@ -3315,11 +3219,7 @@ impl ExecutionEngine { after_pressure.usage_ratio, compression_count, duration_ms, - if compression_result.has_model_summary { - "model" - } else { - "local_fallback" - } + "model" ); self.emit_event( @@ -3332,12 +3232,8 @@ impl ExecutionEngine { tokens_after, compression_ratio, duration_ms, - has_summary: compression_result.has_model_summary, - summary_source: if compression_result.has_model_summary { - "model".to_string() - } else { - "local_fallback".to_string() - }, + has_summary: true, + summary_source: "model".to_string(), applied: true, }, EventPriority::Normal, @@ -3362,12 +3258,8 @@ impl ExecutionEngine { tokens_after, compression_ratio, duration_ms, - has_summary: compression_result.has_model_summary, - summary_source: if compression_result.has_model_summary { - "model".to_string() - } else { - "local_fallback".to_string() - }, + has_summary: true, + summary_source: "model".to_string(), applied: true, }) } @@ -3816,8 +3708,6 @@ impl ExecutionEngine { let mut total_tools = 0; let mut last_partial_recovery_reason: Option = None; 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; @@ -3829,7 +3719,7 @@ impl ExecutionEngine { const MAX_FAILED_TOOL_RECOVERY_ATTEMPTS: usize = 3; const MAX_PARTIAL_CONTINUATION_ATTEMPTS: usize = 3; let mut full_compression_count = 0usize; - let mut compression_failure_count = 0u32; + let compression_failure_count = 0u32; // Save the last token usage statistics let mut last_usage: Option = None; @@ -4027,11 +3917,6 @@ impl ExecutionEngine { && token_pressure.total_tokens >= token_pressure.input_limit; let mut send_pressure_reusable = true; - // Circuit breaker: skip full compression if it has failed too many - // consecutive times. Microcompact and emergency truncation still run. - let circuit_breaker_open = - consecutive_compression_failures >= MAX_CONSECUTIVE_COMPRESSION_FAILURES; - if !should_compress { debug!( "No compression needed: session={}, total_tokens={}, input_limit={}, context_window={}, output_reserve={}, safety_reserve={}, usage={:.1}%", @@ -4043,11 +3928,6 @@ impl ExecutionEngine { token_pressure.safety_reserve_tokens, token_pressure.usage_ratio * 100.0 ); - } else if circuit_breaker_open { - warn!( - "Compression circuit breaker open ({} consecutive failures), skipping full compression for round {}", - consecutive_compression_failures, round_index - ); } else { info!( "Triggering context compression: session={}, total_tokens={}, input_limit={}, context_window={}, output_reserve={}, safety_reserve={}, usage={:.1}%", @@ -4113,25 +3993,15 @@ impl ExecutionEngine { &turn_prompt_scaffold, ); full_compression_count += 1; - consecutive_compression_failures = 0; send_pressure_reusable = false; } Ok(None) => { - debug!("No eligible multi-turn context available for compression"); - consecutive_compression_failures = 0; + return Err(OpenBitFunError::AIClient( + "Context compression has no eligible plan".to_string(), + )); } Err(err @ OpenBitFunError::Cancelled(_)) => return Err(err), - Err(e) => { - consecutive_compression_failures += 1; - compression_failure_count += 1; - error!( - "Round {} compression failed ({}/{}): {}, continuing with uncompressed context", - round_index, - consecutive_compression_failures, - MAX_CONSECUTIVE_COMPRESSION_FAILURES, - e - ); - } + Err(e) => return Err(e), } } @@ -4380,7 +4250,6 @@ impl ExecutionEngine { ) .await; full_compression_count += 1; - consecutive_compression_failures = 0; continue; } Ok(None) => { @@ -4399,7 +4268,7 @@ impl ExecutionEngine { round_index, compression_error ); - return Err(err); + return Err(compression_error); } } } @@ -5301,6 +5170,10 @@ impl ExecutionEngine { } } +#[cfg(test)] +#[path = "compression_tests.rs"] +mod compression_tests; + #[cfg(test)] mod tests { use super::{ @@ -6056,19 +5929,13 @@ mod tests { .to_string() .contains("Old rule")); - let compressor = ContextCompressor::new(Default::default()); + let compressor = ContextCompressor::new(); let plan = compressor .plan_compression(&context.session_id, &persisted, 128_000, 100) .expect("compression plan") .expect("compressible context"); let compressed = compressor - .compress_plan_with_contract( - &context.session_id, - 128_000, - plan, - None, - Some("summary".to_string()), - ) + .compress_plan_with_contract(&context.session_id, plan, None, "summary".to_string()) .expect("compression result") .messages; session_manager diff --git a/src/crates/assembly/core/src/agentic/execution/mod.rs b/src/crates/assembly/core/src/agentic/execution/mod.rs index 917066f0dd..03a675ba9a 100644 --- a/src/crates/assembly/core/src/agentic/execution/mod.rs +++ b/src/crates/assembly/core/src/agentic/execution/mod.rs @@ -2,6 +2,7 @@ //! //! Responsible for AI interaction and model round control +pub(super) mod compression_request; #[cfg(feature = "agent-runtime")] pub(crate) mod conditional_instructions; pub mod edit_constraint_guard; 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 b5d2885823..4023a8928b 100644 --- a/src/crates/assembly/core/src/agentic/execution/round_executor.rs +++ b/src/crates/assembly/core/src/agentic/execution/round_executor.rs @@ -116,32 +116,10 @@ impl ModelRoundLifecycle { } impl RoundExecutor { - const MAX_STREAM_ATTEMPTS: usize = 10; - const RETRY_BASE_DELAY_MS: u64 = 500; - const RATE_LIMIT_RETRY_BASE_DELAY_MS: u64 = 2_000; - const MAX_EXPONENTIAL_DELAY_MS: u64 = 30_000; - const MAX_RATE_LIMIT_DELAY_MS: u64 = 60_000; - const MAX_RETRY_EXPONENT_SHIFT: u32 = 6; - - /// Unknown and malformed provider responses retain the bounded recovery - /// introduced with the unified attempt budget. Only classified rejections - /// leave this loop immediately; context overflow is recovered by the caller. + const MAX_STREAM_ATTEMPTS: usize = openbitfun_agent_stream::retry::MAX_MODEL_ATTEMPTS; + fn should_retry_provider_error(category: &ErrorCategory) -> bool { - match category { - ErrorCategory::Auth - | ErrorCategory::Permission - | ErrorCategory::ProviderQuota - | ErrorCategory::ProviderBilling - | ErrorCategory::InvalidRequest - | ErrorCategory::ContentPolicy - | ErrorCategory::ContextOverflow => false, - ErrorCategory::Network - | ErrorCategory::RateLimit - | ErrorCategory::Timeout - | ErrorCategory::ProviderUnavailable - | ErrorCategory::ModelError - | ErrorCategory::Unknown => true, - } + openbitfun_agent_stream::retry::should_retry(category) } fn terminal_request_error(error: &anyhow::Error, attempts: u32) -> OpenBitFunError { @@ -1578,35 +1556,7 @@ impl RoundExecutor { error_message: &str, provider_error: Option<&AiProviderError>, ) -> u64 { - let shift = u32::try_from(attempt_index) - .unwrap_or(u32::MAX) - .min(Self::MAX_RETRY_EXPONENT_SHIFT); - let msg = error_message.to_lowercase(); - let is_rate_limit = provider_error - .is_some_and(|error| error.category == ErrorCategory::RateLimit) - || msg.contains("429") - || msg.contains("rate limit") - || msg.contains("too many requests"); - - let fallback = if is_rate_limit { - Self::RATE_LIMIT_RETRY_BASE_DELAY_MS - .saturating_mul(1u64 << shift) - .min(Self::MAX_RATE_LIMIT_DELAY_MS) - } else { - Self::RETRY_BASE_DELAY_MS - .saturating_mul(1u64 << shift) - .min(Self::MAX_EXPONENTIAL_DELAY_MS) - }; - - match provider_error.and_then(|error| error.retry_after_ms) { - Some(retry_after_ms) if is_rate_limit => retry_after_ms - .max(fallback) - .min(Self::MAX_RATE_LIMIT_DELAY_MS), - Some(retry_after_ms) if retry_after_ms > 0 => { - retry_after_ms.min(Self::MAX_RATE_LIMIT_DELAY_MS) - } - Some(_) | None => fallback, - } + openbitfun_agent_stream::retry::delay_ms(attempt_index, error_message, provider_error) } } @@ -1638,7 +1588,7 @@ fn token_details_from_usage( } #[cfg(test)] -mod tests { +pub(super) mod tests { use super::{ normalize_deferred_tool_calls_for_replay, ModelRoundLifecycle, RoundExecutor, StreamProcessor, @@ -1666,7 +1616,7 @@ mod tests { use std::time::Duration; use tokio_util::sync::CancellationToken; - fn test_round_executor() -> RoundExecutor { + pub(in crate::agentic::execution) fn test_round_executor() -> RoundExecutor { let event_queue = Arc::new(EventQueue::new(EventQueueConfig::default())); RoundExecutor { stream_processor: Arc::new(StreamProcessor::new(event_queue.clone())), @@ -1801,15 +1751,15 @@ mod tests { } } - struct RetryTestServer { + pub(in crate::agentic::execution) struct RetryTestServer { url: String, - requests: Arc>>, + pub(in crate::agentic::execution) requests: Arc>>, stop: Arc, thread: Option>, } impl RetryTestServer { - fn new(replies: Vec<(u16, String)>) -> Self { + pub(in crate::agentic::execution) fn new(replies: Vec<(u16, String)>) -> Self { Self::with_open_stream(replies, false) } @@ -1886,7 +1836,9 @@ mod tests { } } - fn client(&self) -> Arc { + pub(in crate::agentic::execution) fn client( + &self, + ) -> Arc { Arc::new(crate::infrastructure::ai::AIClient::new( openbitfun_core_types::AIConfig { name: "retry-test".to_string(), @@ -1923,7 +1875,7 @@ mod tests { } } - fn retry_test_success() -> (u16, String) { + pub(in crate::agentic::execution) fn retry_test_success() -> (u16, String) { ( 200, format!( 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 cb9a37e6bd..e333689491 100644 --- a/src/crates/assembly/core/src/agentic/session/compression/compressor.rs +++ b/src/crates/assembly/core/src/agentic/session/compression/compressor.rs @@ -2,10 +2,6 @@ //! //! Responsible only for transforming a session context into a compressed one. -use super::fallback::{ - build_structured_compression_summary_with_contract, CompressionFallbackOptions, - CompressionSummaryArtifact, -}; use crate::agentic::core::{ render_system_reminder, CompressedTodoSnapshot, CompressionContract, CompressionEntry, CompressionPayload, InternalReminderKind, Message, MessageContent, MessageHelper, MessageRole, @@ -16,43 +12,14 @@ use crate::util::errors::OpenBitFunResult; use log::{debug, trace}; use std::borrow::Cow; -/// Context compressor configuration -#[derive(Debug, Clone)] -pub struct CompressionConfig { - pub fallback_max_tokens_ratio: f32, - pub fallback_user_chars: usize, - pub fallback_assistant_chars: usize, - pub fallback_tool_arg_chars: usize, - pub fallback_tool_command_chars: usize, -} - -impl Default for CompressionConfig { - fn default() -> Self { - Self { - fallback_max_tokens_ratio: 0.25, - fallback_user_chars: 1000, - fallback_assistant_chars: 1000, - fallback_tool_arg_chars: 100, - fallback_tool_command_chars: 100, - } - } -} - -#[derive(Debug, Clone)] -struct TurnWithTokens { - messages: Vec, -} - -impl TurnWithTokens { - fn new(messages: Vec) -> Self { - Self { messages } - } +struct CompressionSummaryArtifact { + summary_text: String, + payload: CompressionPayload, } #[derive(Debug, Clone)] pub struct CompressionResult { pub messages: Vec, - pub has_model_summary: bool, } #[derive(Debug, Clone)] @@ -83,9 +50,8 @@ struct AtomicMessageUnit { } /// Stateless context compression service. -pub struct ContextCompressor { - config: CompressionConfig, -} +#[derive(Default)] +pub struct ContextCompressor {} impl ContextCompressor { pub const DEFAULT_RECENT_CONTEXT_TOKENS: usize = 10_000; @@ -94,8 +60,8 @@ impl ContextCompressor { const COMPRESSION_CONTINUATION_REMINDER: &'static str = "This conversation was compacted. Re-establish the working state from the retained user messages, summary, and recent context. Continue any unfinished work; otherwise use this context for the user's next request. Do not ask the user to repeat information already retained here. If a required detail is missing and a pre-compaction transcript is available, inspect the relevant part before proceeding."; - pub fn new(config: CompressionConfig) -> Self { - Self { config } + pub fn new() -> Self { + Self::default() } pub fn plan_compression( @@ -375,21 +341,19 @@ impl ContextCompressor { pub fn compress_plan_with_contract( &self, session_id: &str, - context_window: usize, plan: CompressionPlan, contract: Option, - model_summary: Option, + model_summary: String, ) -> OpenBitFunResult { - let turns = MessageHelper::group_messages_by_turns(plan.summary_messages); - let turns = turns.into_iter().map(TurnWithTokens::new).collect(); - let mut summary_artifact = match model_summary { - Some(summary) => self.build_model_summary_artifact(summary, contract), - None => self.build_fallback_summary_artifact(turns, context_window, contract), - }; + let summary = Self::normalize_model_summary_output(&model_summary).ok_or_else(|| { + crate::OpenBitFunError::AIClient( + "Context compression requires a non-empty model summary".to_string(), + ) + })?; + let mut summary_artifact = self.build_model_summary_artifact(summary, contract); if let Some(checkpoint) = plan.current_turn_todo_checkpoint { Self::append_current_turn_todo_checkpoint(&mut summary_artifact, checkpoint); } - let has_model_summary = summary_artifact.used_model_summary; let summary_message = self.create_summary_message(summary_artifact); let mut messages = plan.retained_user_messages; messages.push(summary_message); @@ -404,10 +368,7 @@ impl ContextCompressor { session_id, messages.len() ); - Ok(CompressionResult { - messages, - has_model_summary, - }) + Ok(CompressionResult { messages }) } fn append_current_turn_todo_checkpoint( @@ -493,7 +454,7 @@ impl ContextCompressor { } fn create_summary_message(&self, summary_artifact: CompressionSummaryArtifact) -> Message { - let boundary_text = Self::render_boundary_marker_text(summary_artifact.used_model_summary); + let boundary_text = Self::render_boundary_marker_text(); let content = render_system_reminder(&format!( "{}\n\n{}", boundary_text, summary_artifact.summary_text @@ -503,13 +464,8 @@ impl ContextCompressor { .with_compression_payload(summary_artifact.payload) } - fn render_boundary_marker_text(used_model_summary: bool) -> String { - let mut msg = "Some earlier user messages were retained verbatim, and the remaining earlier conversation has been summarized below. Use both as prior context." - .to_string(); - if !used_model_summary { - msg.push_str(" This is a partial reconstructed record. Message text, tool arguments, task lists, and tool results may be truncated or omitted."); - } - msg + fn render_boundary_marker_text() -> &'static str { + "Some earlier user messages were retained verbatim, and the remaining earlier conversation has been summarized below. Use both as prior context." } fn render_transcript_reference_text( @@ -552,34 +508,6 @@ impl ContextCompressor { CompressionSummaryArtifact { summary_text, payload, - used_model_summary: true, - } - } - - fn build_fallback_summary_artifact( - &self, - turns_to_compress: Vec, - context_window: usize, - contract: Option, - ) -> CompressionSummaryArtifact { - build_structured_compression_summary_with_contract( - turns_to_compress - .into_iter() - .map(|turn| turn.messages) - .collect(), - &self.build_fallback_options(context_window), - contract, - ) - } - - fn build_fallback_options(&self, context_window: usize) -> CompressionFallbackOptions { - CompressionFallbackOptions { - max_tokens: ((context_window as f32 * self.config.fallback_max_tokens_ratio) as usize) - .max(256), - user_chars: self.config.fallback_user_chars, - assistant_chars: self.config.fallback_assistant_chars, - tool_arg_chars: self.config.fallback_tool_arg_chars, - tool_command_chars: self.config.fallback_tool_command_chars, } } @@ -683,7 +611,7 @@ mod tests { #[test] fn recent_context_keeps_the_exact_suffix_that_fits_the_budget() { - let compressor = ContextCompressor::new(Default::default()); + let compressor = ContextCompressor::new(); let current_user = Message::user("Current request".to_string()); let current_assistant = Message::assistant("Current answer".to_string()); let recent_target = current_user.estimate_tokens_with_reasoning(true) @@ -709,7 +637,7 @@ mod tests { #[test] fn cutoff_can_split_a_turn_without_adding_an_anchor_or_boundary_reminder() { - let compressor = ContextCompressor::new(Default::default()); + let compressor = ContextCompressor::new(); 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".to_string()) @@ -740,13 +668,7 @@ mod tests { .any(|message| message.id == current_user.id)); let result = compressor - .compress_plan_with_contract( - "session", - 128_000, - plan, - None, - Some("Earlier work summary".to_string()), - ) + .compress_plan_with_contract("session", plan, None, "Earlier work summary".to_string()) .expect("compression succeeds"); assert_eq!(result.messages.len(), 5); assert_eq!(result.messages[0].content.to_string(), "Older request"); @@ -769,7 +691,7 @@ mod tests { #[test] fn recent_context_never_splits_tool_results_from_their_assistant_call() { - let compressor = ContextCompressor::new(Default::default()); + let compressor = ContextCompressor::new(); let assistant = Message::assistant_with_tools("Planning".to_string(), vec![todo_call()]); let result = todo_result(); let messages = vec![ @@ -799,7 +721,7 @@ mod tests { #[test] fn current_turn_todo_outside_recent_tail_is_checkpointed() { - let compressor = ContextCompressor::new(Default::default()); + let compressor = ContextCompressor::new(); let todos = serde_json::json!([ {"id": "todo-1", "content": "Implement checkpoint", "status": "in_progress"}, {"id": "todo-2", "content": "Run focused tests", "status": "pending"} @@ -834,13 +756,7 @@ mod tests { assert_eq!(plan.recent_tail_messages[0].id, recent.id); let result = compressor - .compress_plan_with_contract( - "session", - 128_000, - plan, - None, - Some("Earlier work summary".to_string()), - ) + .compress_plan_with_contract("session", plan, None, "Earlier work summary".to_string()) .expect("compression succeeds"); let summary = result .messages @@ -874,7 +790,7 @@ mod tests { #[test] fn historical_turn_todo_is_not_checkpointed_for_current_turn() { - let compressor = ContextCompressor::new(Default::default()); + let compressor = ContextCompressor::new(); let todos = serde_json::json!([ {"id": "old-todo", "content": "Historical task", "status": "pending"} ]); @@ -907,7 +823,7 @@ mod tests { #[test] fn current_turn_todo_already_in_recent_tail_is_not_duplicated() { - let compressor = ContextCompressor::new(Default::default()); + let compressor = ContextCompressor::new(); let todos = serde_json::json!([ {"id": "todo-tail", "content": "Stay in tail", "status": "in_progress"} ]); @@ -940,7 +856,7 @@ mod tests { #[test] fn current_turn_empty_todo_checkpoint_prevents_historical_state_resurrection() { - let compressor = ContextCompressor::new(Default::default()); + let compressor = ContextCompressor::new(); let cleared_assistant = Message::assistant_with_tools( "Clearing completed work".to_string(), vec![todo_call_with("todo-clear", serde_json::json!([]))], @@ -966,10 +882,9 @@ mod tests { let result = compressor .compress_plan_with_contract( "session", - 128_000, plan, None, - Some("A historical task was pending.".to_string()), + "A historical task was pending.".to_string(), ) .expect("compression succeeds"); let summary = result @@ -1000,7 +915,7 @@ mod tests { #[test] fn failed_current_turn_todo_does_not_replace_the_last_successful_snapshot() { - let compressor = ContextCompressor::new(Default::default()); + let compressor = ContextCompressor::new(); let successful_todos = serde_json::json!([ {"id": "todo-good", "content": "Keep successful state", "status": "in_progress"} ]); @@ -1052,7 +967,7 @@ mod tests { #[test] fn current_turn_todo_checkpoint_survives_recompression_from_payload() { - let compressor = ContextCompressor::new(Default::default()); + let compressor = ContextCompressor::new(); let prior_summary = Message::user(render_system_reminder("Earlier compressed context")) .with_semantic_kind(MessageSemanticKind::CompressionSummary) .with_compression_payload(CompressionPayload { @@ -1099,7 +1014,7 @@ mod tests { #[test] fn increasing_the_budget_to_the_next_atomic_unit_moves_the_cutoff() { - let compressor = ContextCompressor::new(Default::default()); + let compressor = ContextCompressor::new(); let messages = vec![ Message::system("system".to_string()), Message::user("request".to_string()), @@ -1128,7 +1043,7 @@ mod tests { #[test] fn compression_orders_retained_users_before_summary_and_recent_context() { - let compressor = ContextCompressor::new(Default::default()); + let compressor = ContextCompressor::new(); let user1 = Message::user("user1".to_string()); let user2 = Message::user("user2".to_string()); let user3 = Message::user("user3".to_string()); @@ -1149,13 +1064,7 @@ mod tests { .expect("planning succeeds") .expect("plan exists"); let mut result = compressor - .compress_plan_with_contract( - "session", - 128_000, - plan, - None, - Some("Model summary".to_string()), - ) + .compress_plan_with_contract("session", plan, None, "Model summary".to_string()) .expect("compression succeeds"); assert_eq!(result.messages.len(), 6); @@ -1215,7 +1124,7 @@ mod tests { #[test] fn recompression_replaces_the_previous_continuation_reminder() { - let compressor = ContextCompressor::new(Default::default()); + let compressor = ContextCompressor::new(); let old_reminder = Message::internal_reminder( InternalReminderKind::CompressionContinuation, "old continuation", @@ -1244,13 +1153,7 @@ mod tests { .any(|message| message.id == old_reminder.id)); let result = compressor - .compress_plan_with_contract( - "session", - 128_000, - plan, - None, - Some("updated summary".to_string()), - ) + .compress_plan_with_contract("session", plan, None, "updated summary".to_string()) .expect("compression succeeds"); assert_eq!( @@ -1268,7 +1171,7 @@ mod tests { #[test] fn conditional_instruction_reminders_do_not_enter_any_compacted_context() { - let compressor = ContextCompressor::new(Default::default()); + let compressor = ContextCompressor::new(); let summarized_reminder = Message::internal_reminder( InternalReminderKind::ConditionalInstructions, "older path rule", @@ -1327,7 +1230,7 @@ mod tests { #[test] fn retained_user_token_budget_is_ten_percent_capped_at_twenty_thousand() { - let compressor = ContextCompressor::new(Default::default()); + let compressor = ContextCompressor::new(); let messages = vec![ Message::system("system".to_string()), Message::user("older request".to_string()), @@ -1378,80 +1281,6 @@ mod tests { assert_eq!(retained_tokens, 0); } - #[test] - fn synthetic_summary_turn_payload_remains_atomic_on_recompression() { - let marker = Message::user(render_system_reminder( - "Earlier conversation was compressed.", - )) - .with_semantic_kind(MessageSemanticKind::CompressionBoundaryMarker); - let summary = Message::assistant("Summary text".to_string()) - .with_semantic_kind(MessageSemanticKind::CompressionSummary) - .with_compression_payload(CompressionPayload::from_summary("Summary text".to_string())); - - let summary_artifact = - crate::agentic::session::compression::fallback::build_structured_compression_summary( - vec![vec![marker, summary]], - &crate::agentic::session::compression::fallback::CompressionFallbackOptions { - max_tokens: 10_000, - user_chars: 120, - assistant_chars: 120, - tool_arg_chars: 80, - tool_command_chars: 80, - }, - ); - - assert!(matches!( - &summary_artifact.payload.entries[0], - CompressionEntry::ModelSummary { text } if text == "Summary text" - )); - } - - #[test] - fn merged_compression_payload_remains_atomic_on_recompression() { - let compressor = ContextCompressor::new(Default::default()); - let messages = vec![ - Message::system("system".to_string()), - Message::user("Continue the refactor".to_string()), - Message::assistant("Work in progress".to_string()), - Message::user("Keep recent context".to_string()), - Message::assistant("Recent evidence".to_string()), - ]; - let plan = compressor - .plan_compression("session", &messages, 8_000, 1) - .expect("planning succeeds") - .expect("plan exists"); - let compressed = compressor - .compress_plan_with_contract( - "session", - 8_000, - plan, - None, - Some("Model summary".to_string()), - ) - .expect("compression succeeds"); - - let summary_artifact = - crate::agentic::session::compression::fallback::build_structured_compression_summary( - vec![compressed.messages], - &crate::agentic::session::compression::fallback::CompressionFallbackOptions { - max_tokens: 10_000, - user_chars: 120, - assistant_chars: 120, - tool_arg_chars: 80, - tool_command_chars: 80, - }, - ); - - assert!(summary_artifact - .payload - .entries - .iter() - .any(|entry| matches!( - entry, - CompressionEntry::ModelSummary { text } if text == "Model summary" - ))); - } - #[test] fn model_summary_output_trims_plain_text() { let normalized = ContextCompressor::normalize_model_summary_output(" Plain summary\n"); diff --git a/src/crates/assembly/core/src/agentic/session/compression/fallback/builder.rs b/src/crates/assembly/core/src/agentic/session/compression/fallback/builder.rs deleted file mode 100644 index b0e021f1a8..0000000000 --- a/src/crates/assembly/core/src/agentic/session/compression/fallback/builder.rs +++ /dev/null @@ -1,183 +0,0 @@ -use super::sanitize::{ - sanitize_assistant_text, sanitize_todo_snapshot, sanitize_tool_arguments, sanitize_user_text, -}; -use super::types::CompressionFallbackOptions; -use crate::agentic::core::{ - strip_prompt_markup, CompressedMessage, CompressedMessageRole, CompressedTodoSnapshot, - CompressedToolCall, CompressionEntry, Message, MessageContent, MessageRole, - MessageSemanticKind, -}; - -pub(super) fn build_entries_from_turns( - turns: Vec>, - options: &CompressionFallbackOptions, -) -> Vec { - let mut entries = Vec::new(); - - for turn in turns { - build_entries_from_messages(turn, options, &mut entries); - } - - entries -} - -fn build_entries_from_messages( - messages: Vec, - options: &CompressionFallbackOptions, - output: &mut Vec, -) { - let turn_id = messages - .first() - .and_then(|message| message.metadata.turn_id.clone()); - let mut turn_messages = Vec::new(); - let mut latest_todo = None; - - for message in messages { - if let Some(entries) = extract_nested_compression_entries(&message) { - flush_turn_entry( - output, - turn_id.clone(), - &mut turn_messages, - &mut latest_todo, - ); - output.extend(entries); - continue; - } - - match message.content { - MessageContent::Text(text) => match message.role { - MessageRole::User => { - if let Some(text) = sanitize_user_text(&text, options) { - turn_messages.push(CompressedMessage { - role: CompressedMessageRole::User, - text: Some(text), - tool_calls: Vec::new(), - }); - } - } - MessageRole::Assistant => { - if let Some(text) = sanitize_assistant_text(&text, options) { - turn_messages.push(CompressedMessage { - role: CompressedMessageRole::Assistant, - text: Some(text), - tool_calls: Vec::new(), - }); - } - } - MessageRole::System | MessageRole::Tool => {} - }, - MessageContent::Multimodal { text, images } => { - if message.role == MessageRole::User { - let mut rendered = sanitize_user_text(&text, options).unwrap_or_default(); - if !images.is_empty() { - if !rendered.is_empty() { - rendered.push('\n'); - } - rendered.push_str(&format!("[{} image(s) omitted]", images.len())); - } - if !rendered.trim().is_empty() { - turn_messages.push(CompressedMessage { - role: CompressedMessageRole::User, - text: Some(rendered), - tool_calls: Vec::new(), - }); - } - } - } - MessageContent::Mixed { - text, tool_calls, .. - } => { - if message.role != MessageRole::Assistant { - continue; - } - - let mut compressed_tool_calls = Vec::new(); - - for tool_call in tool_calls { - if tool_call.tool_name == "TodoWrite" { - latest_todo = sanitize_todo_snapshot(&tool_call.arguments); - continue; - } - - let compressed_tool_call = CompressedToolCall { - tool_name: tool_call.tool_name.clone(), - arguments: sanitize_tool_arguments( - &tool_call.tool_name, - &tool_call.arguments, - options, - ), - is_error: tool_call.is_error, - }; - compressed_tool_calls.push(compressed_tool_call); - } - - let sanitized_text = sanitize_assistant_text(&text, options); - if sanitized_text.is_some() || !compressed_tool_calls.is_empty() { - turn_messages.push(CompressedMessage { - role: CompressedMessageRole::Assistant, - text: sanitized_text, - tool_calls: compressed_tool_calls, - }); - } - } - MessageContent::ToolResult { .. } => {} - } - } - - flush_turn_entry(output, turn_id, &mut turn_messages, &mut latest_todo); -} - -fn flush_turn_entry( - output: &mut Vec, - turn_id: Option, - turn_messages: &mut Vec, - latest_todo: &mut Option, -) { - if turn_messages.is_empty() && latest_todo.is_none() { - return; - } - - output.push(CompressionEntry::Turn { - turn_id, - messages: std::mem::take(turn_messages), - todo: latest_todo.take(), - }); -} - -fn extract_nested_compression_entries(message: &Message) -> Option> { - match message.metadata.semantic_kind { - Some(MessageSemanticKind::CompressionBoundaryMarker) => return Some(Vec::new()), - Some(MessageSemanticKind::CompressionSummary) - | Some(MessageSemanticKind::InternalReminder) => {} - _ => return None, - } - - if let Some(payload) = message.metadata.compression_payload.clone() { - if !payload.is_empty() { - return Some(payload.entries); - } - } - - if message.metadata.semantic_kind == Some(MessageSemanticKind::CompressionSummary) { - return None; - } - - if message - .internal_reminder_kind() - .is_some_and(|kind| kind.should_drop_during_compaction()) - { - return Some(Vec::new()); - } - - let raw_text = match &message.content { - MessageContent::Text(text) => text.clone(), - MessageContent::Multimodal { text, .. } => text.clone(), - _ => String::new(), - }; - let stripped = strip_prompt_markup(&raw_text); - if stripped.is_empty() { - return None; - } - - Some(vec![CompressionEntry::ModelSummary { text: stripped }]) -} diff --git a/src/crates/assembly/core/src/agentic/session/compression/fallback/mod.rs b/src/crates/assembly/core/src/agentic/session/compression/fallback/mod.rs deleted file mode 100644 index dc0aa0b626..0000000000 --- a/src/crates/assembly/core/src/agentic/session/compression/fallback/mod.rs +++ /dev/null @@ -1,41 +0,0 @@ -mod builder; -mod payload; -mod render; -mod sanitize; -mod types; - -use crate::agentic::core::{CompressionContract, CompressionEntry}; -use builder::build_entries_from_turns; -use payload::trim_payload_to_budget; -use render::render_payload_for_model; - -pub use types::{CompressionFallbackOptions, CompressionSummaryArtifact}; - -pub fn build_structured_compression_summary( - turns: Vec>, - options: &CompressionFallbackOptions, -) -> CompressionSummaryArtifact { - build_structured_compression_summary_with_contract(turns, options, None) -} - -pub fn build_structured_compression_summary_with_contract( - turns: Vec>, - options: &CompressionFallbackOptions, - contract: Option, -) -> CompressionSummaryArtifact { - let mut entries = build_entries_from_turns(turns, options); - if let Some(contract) = contract.filter(|contract| !contract.is_empty()) { - entries.insert(0, CompressionEntry::Contract { contract }); - } - let trimmed_payload = trim_payload_to_budget(entries, options); - let summary_text = render_payload_for_model(&trimmed_payload); - - CompressionSummaryArtifact { - summary_text, - payload: trimmed_payload, - used_model_summary: false, - } -} - -#[cfg(test)] -mod tests; diff --git a/src/crates/assembly/core/src/agentic/session/compression/fallback/payload.rs b/src/crates/assembly/core/src/agentic/session/compression/fallback/payload.rs deleted file mode 100644 index 4721cae6ea..0000000000 --- a/src/crates/assembly/core/src/agentic/session/compression/fallback/payload.rs +++ /dev/null @@ -1,192 +0,0 @@ -use super::render::render_payload_for_model; -use super::types::{CompressionFallbackOptions, CompressionUnit}; -use crate::agentic::core::{ - render_system_reminder, CompressedMessage, CompressedTodoSnapshot, CompressionEntry, - CompressionPayload, Message, -}; - -pub(super) fn trim_payload_to_budget( - entries: Vec, - options: &CompressionFallbackOptions, -) -> CompressionPayload { - if entries.is_empty() { - return CompressionPayload::default(); - } - - let units = flatten_entries_to_units(entries); - let mut selected_units: Vec = units - .iter() - .filter_map(|unit| match unit { - CompressionUnit::Contract { .. } => Some(unit.clone()), - _ => None, - }) - .collect(); - let history_units: Vec = units - .into_iter() - .filter(|unit| !matches!(unit, CompressionUnit::Contract { .. })) - .collect(); - - for unit in history_units.into_iter().rev() { - let mut candidate_units = vec![unit.clone()]; - candidate_units.extend(selected_units.clone()); - - let candidate_payload = rebuild_payload_from_units(candidate_units); - if estimate_payload_tokens(&candidate_payload) <= options.max_tokens { - let history_insert_index = selected_units - .iter() - .take_while(|selected| matches!(selected, CompressionUnit::Contract { .. })) - .count(); - selected_units.insert(history_insert_index, unit); - } - } - - rebuild_payload_from_units(selected_units) -} - -fn flatten_entries_to_units(entries: Vec) -> Vec { - let mut units = Vec::new(); - - for (entry_id, entry) in entries.into_iter().enumerate() { - match entry { - CompressionEntry::Contract { contract } => { - units.push(CompressionUnit::Contract { contract }); - } - CompressionEntry::ModelSummary { text } => { - units.push(CompressionUnit::ModelSummary { text }); - } - CompressionEntry::Turn { - turn_id, - messages, - todo, - } => { - for message in messages { - units.push(CompressionUnit::TurnMessage { - entry_id, - turn_id: turn_id.clone(), - message, - }); - } - if let Some(todo) = todo { - units.push(CompressionUnit::TurnTodo { - entry_id, - turn_id, - todo, - }); - } - } - } - } - - units -} - -fn rebuild_payload_from_units(units: Vec) -> CompressionPayload { - let mut entries = Vec::new(); - let mut current_turn_entry_id: Option = None; - let mut current_turn_id: Option = None; - let mut current_messages = Vec::new(); - let mut current_todo = None; - - for unit in units { - match unit { - CompressionUnit::Contract { contract } => { - flush_rebuilt_turn( - &mut entries, - &mut current_turn_entry_id, - &mut current_turn_id, - &mut current_messages, - &mut current_todo, - ); - entries.push(CompressionEntry::Contract { contract }); - } - CompressionUnit::ModelSummary { text } => { - flush_rebuilt_turn( - &mut entries, - &mut current_turn_entry_id, - &mut current_turn_id, - &mut current_messages, - &mut current_todo, - ); - entries.push(CompressionEntry::ModelSummary { text }); - } - CompressionUnit::TurnMessage { - entry_id, - turn_id, - message, - } => { - if current_turn_entry_id != Some(entry_id) { - flush_rebuilt_turn( - &mut entries, - &mut current_turn_entry_id, - &mut current_turn_id, - &mut current_messages, - &mut current_todo, - ); - current_turn_entry_id = Some(entry_id); - current_turn_id = turn_id; - } - current_messages.push(message); - } - CompressionUnit::TurnTodo { - entry_id, - turn_id, - todo, - } => { - if current_turn_entry_id != Some(entry_id) { - flush_rebuilt_turn( - &mut entries, - &mut current_turn_entry_id, - &mut current_turn_id, - &mut current_messages, - &mut current_todo, - ); - current_turn_entry_id = Some(entry_id); - current_turn_id = turn_id; - } - current_todo = Some(todo); - } - } - } - - flush_rebuilt_turn( - &mut entries, - &mut current_turn_entry_id, - &mut current_turn_id, - &mut current_messages, - &mut current_todo, - ); - - CompressionPayload { entries } -} - -fn flush_rebuilt_turn( - entries: &mut Vec, - current_turn_entry_id: &mut Option, - current_turn_id: &mut Option, - current_messages: &mut Vec, - current_todo: &mut Option, -) { - if current_turn_entry_id.is_none() { - return; - } - - if current_messages.is_empty() && current_todo.is_none() { - *current_turn_entry_id = None; - *current_turn_id = None; - return; - } - - entries.push(CompressionEntry::Turn { - turn_id: current_turn_id.clone(), - messages: std::mem::take(current_messages), - todo: current_todo.take(), - }); - *current_turn_entry_id = None; - *current_turn_id = None; -} - -fn estimate_payload_tokens(payload: &CompressionPayload) -> usize { - let rendered = render_payload_for_model(payload); - let mut synthetic_message = Message::user(render_system_reminder(&rendered)); - synthetic_message.get_tokens() -} diff --git a/src/crates/assembly/core/src/agentic/session/compression/fallback/render.rs b/src/crates/assembly/core/src/agentic/session/compression/fallback/render.rs deleted file mode 100644 index fd0335226a..0000000000 --- a/src/crates/assembly/core/src/agentic/session/compression/fallback/render.rs +++ /dev/null @@ -1,104 +0,0 @@ -use crate::agentic::core::{ - CompressedMessage, CompressedMessageRole, CompressionContract, CompressionEntry, - CompressionPayload, -}; -use serde_json::{json, Value}; - -pub(super) fn render_payload_for_model(payload: &CompressionPayload) -> String { - if payload.entries.is_empty() { - return "No detailed historical entries fit within the remaining context budget." - .to_string(); - } - - let mut contract_sections = Vec::new(); - let mut history_sections = Vec::new(); - - for (index, entry) in payload.entries.iter().enumerate() { - match entry { - CompressionEntry::Contract { contract } => { - contract_sections.push(render_contract(contract)); - } - CompressionEntry::ModelSummary { text } => { - history_sections.push(format!( - "Earlier summarized history {}:\n{}", - index + 1, - text - )); - } - CompressionEntry::Turn { messages, todo, .. } => { - let mut lines = vec![format!("Historical turn {}:", index + 1)]; - let mut previous_role = None; - for message in messages { - render_compressed_message(&mut lines, message, &mut previous_role); - } - if let Some(todo) = todo { - lines.push("Latest task list for this turn:".to_string()); - if todo.todos.is_empty() { - if let Some(summary) = todo.summary.as_ref() { - lines.push(format!("- {}", summary)); - } - } else { - for todo_item in &todo.todos { - lines.push(format!("- [{}] {}", todo_item.status, todo_item.content)); - } - if let Some(summary) = todo.summary.as_ref() { - lines.push(format!("Task list note: {}", summary)); - } - } - } - history_sections.push(lines.join("\n")); - } - } - } - - let mut sections = contract_sections; - sections.extend(history_sections); - sections.join("\n\n") -} - -fn render_contract(contract: &CompressionContract) -> String { - contract.render_for_model() -} - -fn render_compressed_message( - lines: &mut Vec, - message: &CompressedMessage, - previous_role: &mut Option, -) { - let role_label = match message.role { - CompressedMessageRole::User => "User", - CompressedMessageRole::Assistant => "Assistant", - }; - let is_new_role_segment = *previous_role != Some(message.role); - - if let Some(text) = message.text.as_ref() { - if is_new_role_segment { - lines.push(format!("{role_label}: {text}")); - } else { - lines.push(text.clone()); - } - } else if is_new_role_segment { - lines.push(format!("{role_label}:")); - } - - for tool_call in &message.tool_calls { - let mut rendered = tool_call.tool_name.clone(); - if let Some(arguments) = tool_call.arguments.as_ref() { - rendered.push(' '); - rendered.push_str(&render_tool_arguments(arguments)); - } - if tool_call.is_error { - rendered.push_str(" [error]"); - } - lines.push(format!("Tool call: {}", rendered)); - } - - *previous_role = Some(message.role); -} - -fn render_tool_arguments(arguments: &Value) -> String { - if arguments.is_null() { - return "{}".to_string(); - } - serde_json::to_string(arguments).unwrap_or_else(|_| json!({}).to_string()) -} diff --git a/src/crates/assembly/core/src/agentic/session/compression/fallback/sanitize.rs b/src/crates/assembly/core/src/agentic/session/compression/fallback/sanitize.rs deleted file mode 100644 index b5ee3be266..0000000000 --- a/src/crates/assembly/core/src/agentic/session/compression/fallback/sanitize.rs +++ /dev/null @@ -1,305 +0,0 @@ -use super::types::CompressionFallbackOptions; -use crate::agentic::core::{strip_prompt_markup, CompressedTodoItem, CompressedTodoSnapshot}; -use serde_json::{Map, Value}; - -pub(super) fn sanitize_user_text( - text: &str, - options: &CompressionFallbackOptions, -) -> Option { - let normalized = strip_prompt_markup(text); - sanitize_text(&normalized, options.user_chars) -} - -pub(super) fn sanitize_assistant_text( - text: &str, - options: &CompressionFallbackOptions, -) -> Option { - sanitize_text(text, options.assistant_chars) -} - -pub(super) fn sanitize_todo_snapshot(arguments: &Value) -> Option { - let todos = arguments.get("todos")?.as_array()?; - let mut compressed_todos = Vec::new(); - - for todo in todos { - let Some(todo_object) = todo.as_object() else { - continue; - }; - let Some(content) = todo_object - .get("content") - .and_then(Value::as_str) - .map(str::trim) - .filter(|content| !content.is_empty()) - else { - continue; - }; - let status = todo_object - .get("status") - .and_then(Value::as_str) - .unwrap_or("pending"); - let id = todo_object - .get("id") - .and_then(Value::as_str) - .map(str::to_string); - - compressed_todos.push(CompressedTodoItem { - id, - content: content.to_string(), - status: status.to_string(), - }); - } - - if compressed_todos.is_empty() { - return None; - } - - Some(CompressedTodoSnapshot { - todos: compressed_todos, - summary: None, - }) -} - -pub(super) fn sanitize_tool_arguments( - tool_name: &str, - arguments: &Value, - options: &CompressionFallbackOptions, -) -> Option { - let Some(object) = arguments.as_object() else { - return sanitize_generic_value(arguments, options); - }; - - let sanitized = match tool_name { - "Read" => { - let mut result = Map::new(); - copy_field(object, &mut result, "file_path"); - copy_field(object, &mut result, "start_line"); - copy_field(object, &mut result, "limit"); - result - } - "Write" => { - let mut result = Map::new(); - copy_field(object, &mut result, "file_path"); - insert_cleared_field(object, &mut result, "content"); - result - } - "Edit" => { - let mut result = Map::new(); - copy_field(object, &mut result, "file_path"); - copy_field(object, &mut result, "replace_all"); - insert_cleared_field(object, &mut result, "old_string"); - insert_cleared_field(object, &mut result, "new_string"); - result - } - "Grep" => { - let mut result = Map::new(); - for key in [ - "pattern", - "path", - "glob", - "type", - "head_limit", - "multiline", - "-A", - "-B", - "-C", - "-i", - "-n", - "output_mode", - ] { - copy_field(object, &mut result, key); - } - result - } - "Glob" => { - let mut result = Map::new(); - copy_field(object, &mut result, "pattern"); - copy_field(object, &mut result, "path"); - copy_field(object, &mut result, "limit"); - result - } - "LS" => { - let mut result = Map::new(); - copy_field(object, &mut result, "path"); - copy_field(object, &mut result, "ignore"); - copy_field(object, &mut result, "limit"); - result - } - "GetFileDiff" => { - let mut result = Map::new(); - copy_field(object, &mut result, "file_path"); - result - } - "DeleteFile" => { - let mut result = Map::new(); - copy_field(object, &mut result, "path"); - copy_field(object, &mut result, "recursive"); - result - } - "Git" => { - let mut result = Map::new(); - copy_field(object, &mut result, "operation"); - copy_field(object, &mut result, "working_directory"); - if let Some(args) = object.get("args") { - if let Some(value) = sanitize_generic_value(args, options) { - result.insert("args".to_string(), value); - } - } - result - } - "Bash" => { - let mut result = Map::new(); - insert_sanitize_text(object, &mut result, "command", options.tool_command_chars); - result - } - "TerminalControl" => { - let mut result = Map::new(); - copy_field(object, &mut result, "action"); - copy_field(object, &mut result, "terminal_session_id"); - result - } - "Skill" => { - let mut result = Map::new(); - copy_field(object, &mut result, "command"); - result - } - "CreatePlan" => { - let mut result = Map::new(); - copy_field(object, &mut result, "name"); - copy_field(object, &mut result, "overview"); - insert_cleared_field(object, &mut result, "plan"); - insert_cleared_field(object, &mut result, "todos"); - result - } - "WebSearch" => { - let mut result = Map::new(); - copy_field(object, &mut result, "query"); - result - } - "WebFetch" => { - let mut result = Map::new(); - copy_field(object, &mut result, "url"); - result - } - _ => sanitize_generic_object(object, options), - }; - - if sanitized.is_empty() { - None - } else { - Some(Value::Object(sanitized)) - } -} - -pub(super) fn sanitize_generic_object( - object: &Map, - options: &CompressionFallbackOptions, -) -> Map { - let mut sanitized = Map::new(); - - for (key, value) in object { - let heavy_string = matches!( - key.as_str(), - "content" - | "contents" - | "old_string" - | "new_string" - | "text" - | "output" - | "stdout" - | "stderr" - | "diff" - | "file_diff" - | "original_content" - | "new_content" - | "data_url" - | "data_base64" - ); - if heavy_string { - if let Some(text) = value.as_str() { - sanitized.insert( - format!("{key}_chars"), - Value::Number(serde_json::Number::from(text.chars().count() as u64)), - ); - } - continue; - } - - if let Some(value) = sanitize_generic_value(value, options) { - sanitized.insert(key.clone(), value); - } - } - - sanitized -} - -pub(super) fn sanitize_generic_value( - value: &Value, - options: &CompressionFallbackOptions, -) -> Option { - match value { - Value::Null => None, - Value::Bool(_) | Value::Number(_) => Some(value.clone()), - Value::String(text) => sanitize_text(text, options.tool_arg_chars).map(Value::String), - Value::Array(values) => { - let sanitized_values: Vec = values - .iter() - .take(5) - .filter_map(|value| sanitize_generic_value(value, options)) - .collect(); - if sanitized_values.is_empty() { - None - } else { - Some(Value::Array(sanitized_values)) - } - } - Value::Object(object) => { - let sanitized_object = sanitize_generic_object(object, options); - if sanitized_object.is_empty() { - None - } else { - Some(Value::Object(sanitized_object)) - } - } - } -} - -fn sanitize_text(text: &str, limit: usize) -> Option { - let trimmed = text.trim(); - if trimmed.is_empty() { - return None; - } - - let text_len = trimmed.chars().count(); - if text_len <= limit { - return Some(trimmed.to_string()); - } - - let mut truncated: String = trimmed.chars().take(limit).collect(); - truncated.push_str(" ... [truncated]"); - Some(truncated) -} - -fn copy_field(source: &Map, target: &mut Map, key: &str) { - if let Some(value) = source.get(key) { - target.insert(key.to_string(), value.clone()); - } -} - -fn insert_sanitize_text( - source: &Map, - target: &mut Map, - key: &str, - limit: usize, -) { - if let Some(value) = source.get(key).and_then(Value::as_str) { - if let Some(text) = sanitize_text(value, limit) { - target.insert(key.to_string(), Value::String(text)); - } - } -} - -fn insert_cleared_field(source: &Map, target: &mut Map, key: &str) { - if source.get(key).is_some() { - target.insert(key.to_string(), Value::String("[cleared]".to_string())); - } -} diff --git a/src/crates/assembly/core/src/agentic/session/compression/fallback/tests.rs b/src/crates/assembly/core/src/agentic/session/compression/fallback/tests.rs deleted file mode 100644 index 1d60861c5f..0000000000 --- a/src/crates/assembly/core/src/agentic/session/compression/fallback/tests.rs +++ /dev/null @@ -1,261 +0,0 @@ -use super::{ - build_structured_compression_summary, build_structured_compression_summary_with_contract, - CompressionFallbackOptions, -}; -use crate::agentic::core::{ - render_system_reminder, render_user_query, CompressedMessageRole, CompressionContract, - CompressionContractItem, CompressionEntry, CompressionPayload, InternalReminderKind, Message, - MessageSemanticKind, ToolCall, ToolResult, -}; -use serde_json::json; - -fn default_options() -> CompressionFallbackOptions { - CompressionFallbackOptions { - max_tokens: 10_000, - user_chars: 120, - assistant_chars: 120, - tool_arg_chars: 80, - tool_command_chars: 80, - } -} - -#[test] -fn clears_tool_results_from_compressed_history() { - let assistant = Message::assistant_with_tools( - "Checking file".to_string(), - vec![ToolCall { - tool_id: "tool_1".to_string(), - tool_name: "Read".to_string(), - arguments: json!({ - "file_path": "/tmp/demo.rs", - "start_line": 1, - "limit": 20 - }), - raw_arguments: None, - is_error: false, - parse_error: None, - recovered_from_truncation: false, - repair_kind: Default::default(), - }], - ); - let tool_result = Message::tool_result(ToolResult { - tool_id: "tool_1".to_string(), - tool_name: "Read".to_string(), - effective_tool_name: None, - result: json!({"content": "ignored"}), - result_for_assistant: Some("Read succeeded with file preview".to_string()), - is_error: false, - duration_ms: None, - image_attachments: None, - }); - - let summary_artifact = build_structured_compression_summary( - vec![vec![ - Message::user("inspect".to_string()), - assistant, - tool_result, - ]], - &default_options(), - ); - - let turn = match &summary_artifact.payload.entries[0] { - CompressionEntry::Turn { messages, .. } => messages, - _ => panic!("expected turn entry"), - }; - let assistant_message = turn - .iter() - .find(|message| message.role == CompressedMessageRole::Assistant) - .expect("assistant message"); - - assert_eq!(assistant_message.tool_calls.len(), 1); - assert!(!summary_artifact.summary_text.contains("Tool result:")); - assert!(!summary_artifact - .summary_text - .contains("All tool results have been cleared")); - assert!(summary_artifact.summary_text.contains("Historical turn 1:")); -} - -#[test] -fn reuses_existing_compression_payload_atomically() { - let prior_summary = "Previous conversation summary".to_string(); - let reminder_message = Message::user(render_system_reminder(&prior_summary)) - .with_semantic_kind(MessageSemanticKind::InternalReminder) - .with_compression_payload(CompressionPayload::from_summary(prior_summary.clone())); - - let summary_artifact = - build_structured_compression_summary(vec![vec![reminder_message]], &default_options()); - - assert!(matches!( - &summary_artifact.payload.entries[0], - CompressionEntry::ModelSummary { text } if text == &prior_summary - )); -} - -#[test] -fn strips_user_query_markup_from_fallback_user_messages() { - let raw = format!( - "{}\n{}", - render_user_query("Implement manual /compact"), - render_system_reminder("Keep responses concise") - ); - - let summary_artifact = - build_structured_compression_summary(vec![vec![Message::user(raw)]], &default_options()); - - let turn = match &summary_artifact.payload.entries[0] { - CompressionEntry::Turn { messages, .. } => messages, - _ => panic!("expected turn entry"), - }; - let user_message = turn - .iter() - .find(|message| message.role == CompressedMessageRole::User) - .expect("user message"); - - assert_eq!( - user_message.text.as_deref(), - Some("Implement manual /compact") - ); - assert!(!summary_artifact.summary_text.contains("")); - assert!(!summary_artifact.summary_text.contains("")); -} - -#[test] -fn drops_system_reminder_only_user_messages_from_fallback_summary() { - let summary_artifact = build_structured_compression_summary( - vec![vec![Message::user(render_system_reminder( - "Summarized context boundary marker", - ))]], - &default_options(), - ); - - assert!(summary_artifact.payload.entries.is_empty()); - assert_eq!( - summary_artifact.summary_text, - "No detailed historical entries fit within the remaining context budget." - ); -} - -#[test] -fn drops_listing_diff_internal_reminders_from_fallback_summary() { - let summary_artifact = build_structured_compression_summary( - vec![vec![Message::internal_reminder( - InternalReminderKind::SkillListingDiff, - "# Skill Listing Update\n\nChanged", - )]], - &default_options(), - ); - - assert!(summary_artifact.payload.entries.is_empty()); - assert_eq!( - summary_artifact.summary_text, - "No detailed historical entries fit within the remaining context budget." - ); -} - -#[test] -fn groups_consecutive_assistant_messages_under_single_role_header() { - let summary_artifact = build_structured_compression_summary( - vec![vec![ - Message::user("Update the component styling.".to_string()), - Message::assistant_with_tools( - "".to_string(), - vec![ToolCall { - tool_id: "tool_1".to_string(), - tool_name: "Read".to_string(), - arguments: json!({ - "file_path": "/workspace/example.txt" - }), - raw_arguments: None, - is_error: false, - parse_error: None, - recovered_from_truncation: false, - repair_kind: Default::default(), - }], - ), - Message::assistant_with_tools( - "".to_string(), - vec![ToolCall { - tool_id: "tool_2".to_string(), - tool_name: "Edit".to_string(), - arguments: json!({ - "file_path": "/workspace/example.txt", - "old_string": "before", - "new_string": "after" - }), - raw_arguments: None, - is_error: false, - parse_error: None, - recovered_from_truncation: false, - repair_kind: Default::default(), - }], - ), - Message::assistant("Updated the styling changes.".to_string()), - ]], - &default_options(), - ); - - let assistant_headers = summary_artifact.summary_text.matches("Assistant:").count(); - assert_eq!(assistant_headers, 1); - assert!(summary_artifact - .summary_text - .contains("Assistant:\nTool call: Read {\"file_path\":\"/workspace/example.txt\"}")); - assert!(summary_artifact - .summary_text - .contains("Updated the styling changes.")); -} - -#[test] -fn renders_contract_facts_even_when_tool_results_are_cleared() { - let contract = CompressionContract { - touched_files: vec!["src/main.rs".to_string()], - verification_commands: vec![CompressionContractItem { - target: "cargo test".to_string(), - status: "succeeded".to_string(), - summary: "Verification command completed.".to_string(), - error_kind: None, - }], - blocking_failures: vec![CompressionContractItem { - target: "pnpm run type-check:web".to_string(), - status: "failed".to_string(), - summary: "Type check failed before compression.".to_string(), - error_kind: Some("exit_code:2".to_string()), - }], - subagent_statuses: vec![CompressionContractItem { - target: "ReviewSecurity".to_string(), - status: "partial_timeout".to_string(), - summary: "Security reviewer timed out after partial output.".to_string(), - error_kind: Some("timeout".to_string()), - }], - }; - - let summary_artifact = build_structured_compression_summary_with_contract( - vec![vec![Message::tool_result(ToolResult { - tool_id: "tool_1".to_string(), - tool_name: "Read".to_string(), - effective_tool_name: None, - result: json!({"content": "large output omitted"}), - result_for_assistant: Some("large output omitted".to_string()), - is_error: false, - duration_ms: None, - image_attachments: None, - })]], - &default_options(), - Some(contract), - ); - - assert!(summary_artifact - .summary_text - .contains("The following facts were retained during compression.")); - assert!(summary_artifact.summary_text.contains("src/main.rs")); - assert!(summary_artifact.summary_text.contains("cargo test")); - assert!(summary_artifact - .summary_text - .contains("pnpm run type-check:web")); - assert!(summary_artifact.summary_text.contains("exit_code:2")); - assert!(summary_artifact.summary_text.contains("ReviewSecurity")); - assert!(summary_artifact.summary_text.contains("partial_timeout")); - assert!(matches!( - &summary_artifact.payload.entries[0], - CompressionEntry::Contract { .. } - )); -} diff --git a/src/crates/assembly/core/src/agentic/session/compression/fallback/types.rs b/src/crates/assembly/core/src/agentic/session/compression/fallback/types.rs deleted file mode 100644 index e90edaf42c..0000000000 --- a/src/crates/assembly/core/src/agentic/session/compression/fallback/types.rs +++ /dev/null @@ -1,39 +0,0 @@ -use crate::agentic::core::{ - CompressedMessage, CompressedTodoSnapshot, CompressionContract, CompressionPayload, -}; - -#[derive(Debug, Clone)] -pub struct CompressionFallbackOptions { - pub max_tokens: usize, - pub user_chars: usize, - pub assistant_chars: usize, - pub tool_arg_chars: usize, - pub tool_command_chars: usize, -} - -#[derive(Debug, Clone)] -pub struct CompressionSummaryArtifact { - pub summary_text: String, - pub payload: CompressionPayload, - pub used_model_summary: bool, -} - -#[derive(Debug, Clone)] -pub(super) enum CompressionUnit { - Contract { - contract: CompressionContract, - }, - ModelSummary { - text: String, - }, - TurnMessage { - entry_id: usize, - turn_id: Option, - message: CompressedMessage, - }, - TurnTodo { - entry_id: usize, - turn_id: Option, - todo: CompressedTodoSnapshot, - }, -} diff --git a/src/crates/assembly/core/src/agentic/session/compression/mod.rs b/src/crates/assembly/core/src/agentic/session/compression/mod.rs index 0da008d6e0..4f70429f6b 100644 --- a/src/crates/assembly/core/src/agentic/session/compression/mod.rs +++ b/src/crates/assembly/core/src/agentic/session/compression/mod.rs @@ -11,7 +11,5 @@ //! `execution_engine.rs`. pub mod compressor; -pub mod fallback; pub use compressor::*; -pub use fallback::*; diff --git a/src/crates/assembly/core/src/agentic/system.rs b/src/crates/assembly/core/src/agentic/system.rs index d9c6dd6772..8f874e9d58 100644 --- a/src/crates/assembly/core/src/agentic/system.rs +++ b/src/crates/assembly/core/src/agentic/system.rs @@ -97,7 +97,7 @@ pub async fn init_agentic_system_for_profile_with_runtime_ownership( set_global_token_usage_service(token_usage_service.clone()); let context_store = Arc::new(session::SessionContextStore::new()); - let context_compressor = Arc::new(session::ContextCompressor::new(Default::default())); + let context_compressor = Arc::new(session::ContextCompressor::new()); let session_manager = Arc::new(session::SessionManager::new( context_store, diff --git a/src/crates/assembly/core/src/product_runtime.rs b/src/crates/assembly/core/src/product_runtime.rs index 3492133959..e0352bcccb 100644 --- a/src/crates/assembly/core/src/product_runtime.rs +++ b/src/crates/assembly/core/src/product_runtime.rs @@ -2793,9 +2793,8 @@ mod tests { }; use crate::agentic::persistence::PersistenceManager; use crate::agentic::session::{ - compression::{CompressionConfig, ContextCompressor}, - PromptCachePolicy, SessionContextStore, SessionManager, SessionManagerConfig, - UserContextCacheIdentity, + compression::ContextCompressor, PromptCachePolicy, SessionContextStore, SessionManager, + SessionManagerConfig, UserContextCacheIdentity, }; use crate::agentic::tools::registry::ToolRegistry; use crate::agentic::tools::{ToolPipeline, ToolStateManager}; @@ -3460,7 +3459,7 @@ mod tests { )), event_queue.clone(), session_manager.clone(), - Arc::new(ContextCompressor::new(CompressionConfig::default())), + Arc::new(ContextCompressor::new()), ExecutionEngineConfig::default(), )); let coordinator = Arc::new(ConversationCoordinator::new( @@ -3720,7 +3719,7 @@ mod tests { )), event_queue.clone(), session_manager.clone(), - Arc::new(ContextCompressor::new(CompressionConfig::default())), + Arc::new(ContextCompressor::new()), ExecutionEngineConfig::default(), )); let coordinator = Arc::new(ConversationCoordinator::new( @@ -4135,7 +4134,7 @@ mod tests { )), event_queue.clone(), session_manager.clone(), - Arc::new(ContextCompressor::new(CompressionConfig::default())), + Arc::new(ContextCompressor::new()), ExecutionEngineConfig::default(), )); let coordinator = Arc::new(ConversationCoordinator::new( @@ -4624,7 +4623,7 @@ mod tests { )), event_queue.clone(), session_manager.clone(), - Arc::new(ContextCompressor::new(CompressionConfig::default())), + Arc::new(ContextCompressor::new()), ExecutionEngineConfig::default(), )); let coordinator = Arc::new(ConversationCoordinator::new( diff --git a/src/crates/execution/agent-stream/src/lib.rs b/src/crates/execution/agent-stream/src/lib.rs index 3c4bfd1846..25243930d8 100644 --- a/src/crates/execution/agent-stream/src/lib.rs +++ b/src/crates/execution/agent-stream/src/lib.rs @@ -3,6 +3,7 @@ //! Processes AI streaming responses, supports tool pre-detection and parameter streaming mod hidden_text; +pub mod retry; pub mod tool_call_accumulator; mod unified; diff --git a/src/crates/execution/agent-stream/src/retry.rs b/src/crates/execution/agent-stream/src/retry.rs new file mode 100644 index 0000000000..118263e009 --- /dev/null +++ b/src/crates/execution/agent-stream/src/retry.rs @@ -0,0 +1,96 @@ +//! Shared retry policy for one immutable model request. Input recovery starts a new budget. + +use openbitfun_core_types::errors::{AiProviderError, ErrorCategory}; + +pub const MAX_MODEL_ATTEMPTS: usize = 10; + +pub fn should_retry(category: &ErrorCategory) -> bool { + match category { + ErrorCategory::Auth + | ErrorCategory::Permission + | ErrorCategory::ProviderQuota + | ErrorCategory::ProviderBilling + | ErrorCategory::InvalidRequest + | ErrorCategory::ContentPolicy + | ErrorCategory::ContextOverflow => false, + ErrorCategory::Network + | ErrorCategory::RateLimit + | ErrorCategory::Timeout + | ErrorCategory::ProviderUnavailable + | ErrorCategory::ModelError + | ErrorCategory::Unknown => true, + } +} + +pub fn delay_ms(attempt_index: usize, message: &str, error: Option<&AiProviderError>) -> u64 { + let shift = attempt_index.min(6) as u32; + let message = message.to_lowercase(); + let rate_limit = error.is_some_and(|error| error.category == ErrorCategory::RateLimit) + || message.contains("429") + || message.contains("rate limit") + || message.contains("too many requests"); + let fallback = if rate_limit { + (2_000u64 << shift).min(60_000) + } else { + (500u64 << shift).min(30_000) + }; + match error.and_then(|error| error.retry_after_ms) { + Some(hint) if rate_limit => hint.max(fallback).min(60_000), + Some(hint) if hint > 0 => hint.min(60_000), + _ => fallback, + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn deterministic_rejections_and_overflow_never_enter_retry() { + for category in [ + ErrorCategory::Auth, + ErrorCategory::Permission, + ErrorCategory::ProviderQuota, + ErrorCategory::ProviderBilling, + ErrorCategory::InvalidRequest, + ErrorCategory::ContentPolicy, + ErrorCategory::ContextOverflow, + ] { + assert!(!should_retry(&category), "{category:?}"); + } + for category in [ + ErrorCategory::Network, + ErrorCategory::RateLimit, + ErrorCategory::Timeout, + ErrorCategory::ProviderUnavailable, + ErrorCategory::ModelError, + ErrorCategory::Unknown, + ] { + assert!(should_retry(&category), "{category:?}"); + } + } + + #[test] + fn backoff_and_provider_hints_have_shared_caps() { + let ordinary: Vec<_> = (0..9).map(|i| delay_ms(i, "", None)).collect(); + assert_eq!( + ordinary, + vec![500, 1000, 2000, 4000, 8000, 16000, 30000, 30000, 30000] + ); + let rate: Vec<_> = (0..9).map(|i| delay_ms(i, "429", None)).collect(); + assert_eq!( + rate, + vec![2000, 4000, 8000, 16000, 32000, 60000, 60000, 60000, 60000] + ); + let rate = AiProviderError::classified("slow down".into(), ErrorCategory::RateLimit) + .with_retry_after_ms(Some(1000)); + assert_eq!(delay_ms(3, "", Some(&rate)), 16000); + let unavailable = + AiProviderError::classified("offline".into(), ErrorCategory::ProviderUnavailable) + .with_retry_after_ms(Some(90000)); + assert_eq!(delay_ms(0, "", Some(&unavailable)), 60000); + let denied = AiProviderError::classified("denied".into(), ErrorCategory::Permission) + .with_retry_after_ms(Some(1000)); + assert!(!should_retry(&denied.category)); + } +}