diff --git a/Cargo.lock b/Cargo.lock index ab81366b0c..d87fb5b96f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1106,6 +1106,7 @@ dependencies = [ "futures-util", "globset", "htmd", + "http", "ignore", "image", "jsonschema", diff --git a/crates/tui/Cargo.toml b/crates/tui/Cargo.toml index 357c376198..5f0e2590a5 100644 --- a/crates/tui/Cargo.toml +++ b/crates/tui/Cargo.toml @@ -28,6 +28,7 @@ name = "codewhale-tui" path = "src/main.rs" [dependencies] +http = "1.5" ahash = "0.8" anyhow.workspace = true codewhale-config = { path = "../config", version = "0.9.12" } @@ -118,6 +119,7 @@ mimalloc.workspace = true codewhale-build-support = { path = "../build-support", version = "0.9.12" } [dev-dependencies] +tokio = { workspace = true, features = ["test-util"] } cucumber = "0.23.0" jsonschema.workspace = true wiremock = "0.6" diff --git a/crates/tui/src/client.rs b/crates/tui/src/client.rs index 4bdb1e64e6..ac4d61d8a3 100644 --- a/crates/tui/src/client.rs +++ b/crates/tui/src/client.rs @@ -14,10 +14,7 @@ use futures_util::StreamExt; use reqwest::header::{AUTHORIZATION, CONTENT_TYPE, HeaderMap, HeaderName, HeaderValue}; use serde::{Deserialize, Serialize}; use serde_json::{Value, json}; -use tokio::sync::{ - Mutex as AsyncMutex, OwnedRwLockReadGuard, OwnedRwLockWriteGuard, OwnedSemaphorePermit, RwLock, - Semaphore, -}; +use tokio::sync::{Mutex as AsyncMutex, OwnedRwLockReadGuard, OwnedRwLockWriteGuard, RwLock}; use codewhale_config::catalog::{ CatalogOffering, CatalogRefreshError, CatalogSnapshot, CatalogSource, CatalogStatus, @@ -255,6 +252,14 @@ pub struct DeepSeekClient { pub(super) http1_client: reqwest::Client, request_idempotency_header: Option, operation_id: Option, + operation_recovery: Arc< + std::sync::Mutex< + HashMap< + uuid::Uuid, + Arc>, + >, + >, + >, // In-memory only: includes frozen auth/context headers without exposing // their values in receipts, logs, prompts, or persisted session state. operation_route_identity: String, @@ -296,6 +301,8 @@ pub struct DeepSeekClient { connection_health: Arc>, rate_limiter: Arc>, request_concurrency: Option, + shared_quota: Arc, + operation_permit: Arc>>, path_suffix: Option, /// Unit tests keep the semantic route exact while sending the actual /// production request through a local capture server. This field is @@ -363,9 +370,8 @@ const RATE_LIMIT_PAUSE_RECHECK_INTERVAL: Duration = Duration::from_millis(250); /// Retry-After included). Streaming paths never carry it: their opens go /// through `send_stream_open_with_retry`, which sets no per-request total /// (a total would ride on the returned body), so a stream stays bounded by -/// its open cap and per-chunk idle checks only. Also applied to the -/// Anthropic dialect's non-streaming Messages request, which has no retry -/// loop of its own. +/// its open cap and per-chunk idle checks only. This also covers +/// non-streaming Anthropic Messages and isolated requests. pub(crate) const NON_STREAMING_REQUEST_ENVELOPE: Duration = Duration::from_secs(1800); #[cfg(test)] @@ -420,50 +426,312 @@ struct TokenBucket { capacity: f64, tokens: f64, refill_per_sec: f64, + normal_refill_per_sec: f64, last_refill: Instant, } +struct SharedQuota { + key: String, + last_used: std::sync::Mutex, + storage: Option, + pause: std::sync::Mutex>>, + loaded: tokio::sync::OnceCell<()>, + persistence: AsyncMutex<()>, + limiter: Arc>, + concurrency: Option, + episode: std::sync::Mutex<(Option, u32, u32)>, + adaptive: AsyncMutex<()>, +} + +const QUOTA_SCOPE_RETENTION: Duration = Duration::from_secs(24 * 60 * 60); +const MAX_QUOTA_SCOPES: usize = 65_536; + +fn quota_registry() -> &'static std::sync::Mutex>> { + static QUOTAS: std::sync::OnceLock>>> = + std::sync::OnceLock::new(); + QUOTAS.get_or_init(Default::default) +} + +fn shared_quota( + key: String, + limit: Option, + storage: Option, +) -> Result> { + let mut quotas = quota_registry() + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + // Retain recent episodes across short-lived auxiliary clients. Never + // retire an active holder or a still-live server minimum. The idle window + // covers the maximum configurable recovery budget; the map is bounded. + quotas.retain(|_, quota| { + Arc::strong_count(quota) > 1 + || quota + .last_used + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .elapsed() + < QUOTA_SCOPE_RETENTION + || quota.remaining().is_some() + }); + if let Some(quota) = quotas.get(&key) { + quota.touch(); + if let (Some(limiter), Some(limit)) = (quota.concurrency.as_ref(), limit) { + limiter.restrict(limit); + } + return Ok(quota.clone()); + } + anyhow::ensure!( + quotas.len() < MAX_QUOTA_SCOPES, + "Provider quota scope capacity exhausted; inactive scopes must retire before adding another route" + ); + let quota = Arc::new(SharedQuota { + key: key.clone(), + last_used: std::sync::Mutex::new(Instant::now()), + storage, + pause: std::sync::Mutex::new(None), + loaded: tokio::sync::OnceCell::new(), + persistence: AsyncMutex::new(()), + limiter: Arc::new(AsyncMutex::new(TokenBucket::from_env())), + concurrency: Some(ProviderConcurrencyLimiter::new(limit.unwrap_or(16))), + episode: std::sync::Mutex::new((None, 0, 0)), + adaptive: AsyncMutex::new(()), + }); + quotas.insert(key, quota.clone()); + Ok(quota) +} + +impl SharedQuota { + fn touch(&self) { + *self + .last_used + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) = Instant::now(); + } + fn pause(&self, delay: Duration) { + self.touch(); + let deadline = chrono::Duration::from_std(delay) + .ok() + .and_then(|d| chrono::Utc::now().checked_add_signed(d)) + .unwrap_or(chrono::DateTime::::MAX_UTC); + let mut pause = self + .pause + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + *pause = Some(pause.map_or(deadline, |old| old.max(deadline))); + } + fn remaining(&self) -> Option { + self.pause + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .and_then(|deadline| (deadline - chrono::Utc::now()).to_std().ok()) + .filter(|delay| !delay.is_zero()) + } + async fn initialize(&self) -> Result<()> { + self.touch(); + self.loaded.get_or_try_init(|| async { + if cfg!(test) && self.storage.is_none() { return Ok::<(), String>(()); } + let storage = self.storage.clone().or_else(|| codewhale_paths::codewhale_home().ok().flatten().map(|p| p.join("rate-limit-recovery"))); + let Some(storage) = storage else { return Ok(()); }; + let path = storage.join(format!("{}.json", self.key)); + let file = match tokio::fs::File::open(&path).await { + Ok(file) => file, + Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(()), + Err(_) => return Err("Cannot read saved provider recovery deadline; repair local recovery storage before continuing".into()), + }; + let mut limited = tokio::io::AsyncReadExt::take(file, 4097); + let mut bytes = Vec::new(); + tokio::io::AsyncReadExt::read_to_end(&mut limited, &mut bytes).await + .map_err(|_| "Cannot read saved provider recovery deadline; repair local recovery storage before continuing".to_owned())?; + if bytes.len() > 4096 { return Err("Saved provider recovery deadline is invalid; repair local recovery storage before continuing".into()); } + let reset: chrono::DateTime = serde_json::from_slice(&bytes) + .map_err(|_| "Saved provider recovery deadline is invalid; repair local recovery storage before continuing".to_owned())?; + if let Ok(delay) = (reset - chrono::Utc::now()).to_std() { self.pause(delay); } + Ok(()) + }).await.map(|_| ()).map_err(anyhow::Error::msg) + } + async fn persist(&self) -> std::result::Result<(), String> { + if cfg!(test) && self.storage.is_none() { + return Ok(()); + } + let _io = self.persistence.lock().await; + let reset = *self + .pause + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + let Some(reset) = reset else { + return Ok(()); + }; + let Some(dir) = self.storage.clone().or_else(|| { + codewhale_paths::codewhale_home() + .ok() + .flatten() + .map(|p| p.join("rate-limit-recovery")) + }) else { + return Ok(()); + }; + let temporary = dir.join(format!("{}.{}.tmp", self.key, uuid::Uuid::new_v4())); + let result = async { + tokio::fs::create_dir_all(&dir).await?; + let path = dir.join(format!("{}.json", self.key)); + tokio::fs::write(&temporary, serde_json::to_vec(&reset)?).await?; + tokio::fs::rename(&temporary, &path).await?; + Ok::<_, anyhow::Error>(()) + } + .await; + if result.is_err() { + let _ = tokio::fs::remove_file(&temporary).await; + logging::warn( + "Could not persist provider recovery deadline; in-memory pause remains active", + ); + return Err("Provider recovery deadline could not be saved; repair local recovery storage before continuing".into()); + } + Ok(()) + } + fn admit_recovery(&self, policy: &crate::config::RateLimitRetryPolicy) -> Option { + self.touch(); + let mut episode = self + .episode + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + let started = episode.0?; + let remaining = + Duration::from_secs(policy.recovery_budget_secs).saturating_sub(started.elapsed()); + if remaining.is_zero() { + *episode = (None, 0, 0); + return None; + } + if episode.1 >= policy.max_attempts { + return Some(remaining); + } + episode.1 += 1; + None + } + async fn refused(&self) { + self.touch(); + let _adaptive = self.adaptive.lock().await; + { + let mut episode = self + .episode + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if episode.0.is_none() { + episode.0 = Some(Instant::now()); + episode.1 = 1; + } + episode.2 = 0; + } + if let Some(limiter) = &self.concurrency { + limiter + .gate + .set_capacity((limiter.gate.capacity() / 2).max(1)); + } + let mut bucket = self.limiter.lock().await; + if bucket.enabled { + bucket.refill_per_sec = (bucket.refill_per_sec / 2.0).max(0.1); + bucket.tokens = bucket.tokens.min(1.0); + } + } + async fn recovered(&self) { + self.touch(); + let _adaptive = self.adaptive.lock().await; + { + let mut episode = self + .episode + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + // A real HTTP success ends the shared refusal episode. Logical + // operation budgets remain owned by their operation, while AIMD + // restoration still requires three successful admissions per step. + episode.0 = None; + episode.1 = 0; + episode.2 += 1; + if episode.2 < 3 { + return; + } + episode.2 = 0; + } + if let Some(limiter) = &self.concurrency { + limiter + .gate + .set_capacity((limiter.gate.capacity() + 1).min(limiter.limit())); + } + let mut bucket = self.limiter.lock().await; + bucket.refill_per_sec = (bucket.refill_per_sec + 1.0).min(bucket.normal_refill_per_sec); + } +} + #[derive(Debug, Clone)] struct ProviderConcurrencyLimiter { - semaphore: Arc, + gate: Arc, active: Arc, - limit: usize, + limit: Arc, } struct ProviderRequestPermit { - _permit: OwnedSemaphorePermit, + lease: std::sync::Mutex>, + limiter: ProviderConcurrencyLimiter, +} + +struct ProviderRequestLease { + _permit: crate::core::capacity::DynamicGatePermit, active: Arc, } +impl ProviderRequestPermit { + fn suspend(&self) { + self.lease + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .take(); + } + async fn resume(&self) { + if self + .lease + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .is_some() + { + return; + } + { + let permit = self.limiter.gate.acquire().await; + self.limiter.active.fetch_add(1, Ordering::AcqRel); + *self + .lease + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) = Some(ProviderRequestLease { + _permit: permit, + active: self.limiter.active.clone(), + }); + } + } +} + impl ProviderConcurrencyLimiter { fn new(limit: usize) -> Self { let limit = limit.max(1); Self { - semaphore: Arc::new(Semaphore::new(limit)), + gate: Arc::new(crate::core::capacity::DynamicGate::new(limit)), active: Arc::new(AtomicUsize::new(0)), - limit, + limit: Arc::new(AtomicUsize::new(limit)), } } - async fn acquire(&self) -> Option { - let permit = Arc::clone(&self.semaphore).acquire_owned().await.ok()?; - self.active.fetch_add(1, Ordering::AcqRel); - Some(ProviderRequestPermit { - _permit: permit, - active: Arc::clone(&self.active), - }) - } - fn active(&self) -> usize { self.active.load(Ordering::Acquire) } fn limit(&self) -> usize { - self.limit + self.limit.load(Ordering::Acquire) + } + fn restrict(&self, limit: usize) { + self.limit.fetch_min(limit.max(1), Ordering::AcqRel); + self.gate + .set_capacity(self.gate.capacity().min(self.limit())); } } -impl Drop for ProviderRequestPermit { +impl Drop for ProviderRequestLease { fn drop(&mut self) { self.active.fetch_sub(1, Ordering::AcqRel); } @@ -489,6 +757,7 @@ impl TokenBucket { capacity: burst, tokens: burst, refill_per_sec: rps, + normal_refill_per_sec: rps, last_refill: Instant::now(), } } @@ -591,6 +860,7 @@ impl Clone for DeepSeekClient { http1_client: self.http1_client.clone(), request_idempotency_header: self.request_idempotency_header.clone(), operation_id: self.operation_id, + operation_recovery: self.operation_recovery.clone(), operation_route_identity: self.operation_route_identity.clone(), api_key: self.api_key.clone(), model_bound_secret_values: Arc::clone(&self.model_bound_secret_values), @@ -610,6 +880,8 @@ impl Clone for DeepSeekClient { connection_health: self.connection_health.clone(), rate_limiter: self.rate_limiter.clone(), request_concurrency: self.request_concurrency.clone(), + shared_quota: self.shared_quota.clone(), + operation_permit: self.operation_permit.clone(), path_suffix: self.path_suffix.clone(), #[cfg(test)] test_chat_transport_base_url: self.test_chat_transport_base_url.clone(), @@ -792,15 +1064,29 @@ pub(super) const NON_STREAMING_HTTP_TIMEOUT: Duration = Duration::from_secs(30); /// Read an error response body with a size limit to prevent unbounded allocation. pub(super) async fn bounded_error_text(response: reqwest::Response, max_bytes: usize) -> String { + bounded_error_text_until(response, max_bytes, None).await +} + +async fn bounded_error_text_until( + response: reqwest::Response, + max_bytes: usize, + deadline: Option, +) -> String { use futures_util::StreamExt; let mut stream = response.bytes_stream(); let mut buf = Vec::with_capacity(max_bytes.min(8192)); - while let Some(chunk) = stream.next().await { - let Ok(chunk) = chunk else { break }; - let remaining = max_bytes.saturating_sub(buf.len()); - if remaining == 0 { + while buf.len() < max_bytes { + let next = match deadline { + Some(deadline) => tokio::time::timeout_at(deadline, stream.next()) + .await + .ok() + .flatten(), + None => stream.next().await, + }; + let Some(Ok(chunk)) = next else { break; - } + }; + let remaining = max_bytes - buf.len(); buf.extend_from_slice(&chunk[..chunk.len().min(remaining)]); } String::from_utf8_lossy(&buf).into_owned() @@ -1236,6 +1522,24 @@ impl DeepSeekClient { wire_format: WireFormat, route_limits: Option, config: &Config, + ) -> Result { + Self::from_parts_with_exact_key( + base_url, + default_model, + wire_format, + route_limits, + config, + None, + ) + } + + fn from_parts_with_exact_key( + base_url: String, + default_model: String, + wire_format: WireFormat, + route_limits: Option, + config: &Config, + exact_api_key: Option<&str>, ) -> Result { let api_provider = config.api_provider(); let provider_identity = config.provider_identity_for(api_provider); @@ -1249,7 +1553,11 @@ impl DeepSeekClient { if api_provider == ApiProvider::OpencodeGo { validate_route(api_provider, &default_model).map_err(anyhow::Error::msg)?; } - let (api_key, codex_account_id) = if api_provider == ApiProvider::OpenaiCodex { + let (api_key, codex_account_id) = if let Some(api_key) = exact_api_key { + // A dedicated auxiliary endpoint owns this already-resolved key. + // Re-reading main CLI overrides here could cross a credential boundary. + (api_key.to_owned(), None) + } else if api_provider == ApiProvider::OpenaiCodex { // The official endpoint requires Codex OAuth credentials. A custom // endpoint prefers its own configured key, but an explicit // `OPENAI_CODEX_ACCESS_TOKEN` still wins (`codex_credentials` @@ -1281,7 +1589,7 @@ impl DeepSeekClient { &provider_identity, &base_url, &api_key, - sorted_headers, + &sorted_headers, ))?); let request_idempotency_header = config .request_idempotency_header @@ -1317,6 +1625,50 @@ impl DeepSeekClient { .provider_config_for(api_provider) .and_then(|p| p.reasoning_stream_style.clone()); let request_concurrency_limit = config.provider_max_concurrency(api_provider); + // All models on a credential share a conservative quota group until a provider declares finer groups. + let effective_headers = build_default_headers( + &api_key, + &http_headers, + api_provider, + &base_url, + wire_format, + auth_disabled, + )?; + let mut quota_headers = std::collections::BTreeMap::>::new(); + for (name, value) in &effective_headers { + let name = name.as_str(); + if [ + "key", + "token", + "auth", + "account", + "project", + "organization", + "workspace", + "secret", + "credential", + "cookie", + ] + .iter() + .any(|part| name.contains(part)) + { + quota_headers.insert(name.to_owned(), value.as_bytes().to_vec()); + } + } + if let Some(account_id) = &codex_account_id { + quota_headers.insert("chatgpt-account-id".into(), account_id.as_bytes().to_vec()); + } + let quota_endpoint = reqwest::Url::parse(&base_url)? + .to_string() + .trim_end_matches('/') + .to_owned(); + let quota_key = + crate::hashing::sha256_hex(&serde_json::to_vec(&("a_endpoint, "a_headers))?); + let shared_quota = shared_quota( + quota_key, + request_concurrency_limit, + config.rate_limit_recovery_dir.clone(), + )?; logging::info(format!("API provider: {}", api_provider.as_str())); logging::info(format!( @@ -1380,6 +1732,7 @@ impl DeepSeekClient { http1_client, request_idempotency_header, operation_id: None, + operation_recovery: Arc::new(std::sync::Mutex::new(HashMap::new())), operation_route_identity, api_key, model_bound_secret_values, @@ -1398,8 +1751,10 @@ impl DeepSeekClient { && !config.runtime_thread_inference_unrelated, default_model, connection_health: Arc::new(AsyncMutex::new(ConnectionHealth::default())), - rate_limiter: Arc::new(AsyncMutex::new(TokenBucket::from_env())), - request_concurrency: request_concurrency_limit.map(ProviderConcurrencyLimiter::new), + rate_limiter: shared_quota.limiter.clone(), + request_concurrency: shared_quota.concurrency.clone(), + shared_quota, + operation_permit: Arc::default(), path_suffix, #[cfg(test)] test_chat_transport_base_url: None, @@ -2411,11 +2766,17 @@ impl DeepSeekClient { .map_or(0, ProviderConcurrencyLimiter::active) } - async fn acquire_provider_request_permit(&self) -> Option { - match self.request_concurrency.as_ref() { - Some(limiter) => limiter.acquire().await, - None => None, - } + async fn acquire_provider_request_permit(&self) -> Option> { + let limiter = self.request_concurrency.as_ref()?; + let permit = Arc::new(ProviderRequestPermit { + lease: std::sync::Mutex::new(None), + limiter: limiter.clone(), + }); + *self + .operation_permit + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) = Arc::downgrade(&permit); + Some(permit) } pub(crate) async fn acquire_remote_control_inference_permit( @@ -2429,7 +2790,7 @@ impl DeepSeekClient { fn hold_provider_request_permit_for_stream( stream: crate::llm_client::StreamEventBox, - permit: Option, + permit: Option>, ) -> crate::llm_client::StreamEventBox { Box::pin(async_stream::stream! { let _permit = permit; @@ -2485,10 +2846,10 @@ impl DeepSeekClient { target_language: &str, ) -> Result { let _inference = self.acquire_remote_control_inference_permit().await; - let _permit = self.acquire_provider_request_permit().await; let model = wire_model_for_provider_route(self.api_provider, &self.base_url, model); let max_tokens = self.effective_max_output_tokens(&model); let call = self.for_operation(uuid::Uuid::new_v4()); + let _permit = call.acquire_provider_request_permit().await; if self.wire_format != WireFormat::ChatCompletions { // Non-Chat dialects reuse the prepared-request seam so translation // cannot drift from production shaping. Translation is still an @@ -2844,7 +3205,6 @@ impl DeepSeekClient { request: SpeechSynthesisRequest, ) -> Result { let _inference = self.acquire_remote_control_inference_permit().await; - let _permit = self.acquire_provider_request_permit().await; if self.api_provider != crate::config::ApiProvider::XiaomiMimo { anyhow::bail!( "speech synthesis requires provider 'xiaomi-mimo' (current: {})", @@ -2898,6 +3258,7 @@ impl DeepSeekClient { let url = api_url(&self.base_url, "chat/completions"); let call = self.for_operation(uuid::Uuid::new_v4()); + let _permit = call.acquire_provider_request_permit().await; let response = call.send_json_with_retry(&url, &body).await?; let status = response.status(); if !status.is_success() { @@ -2937,6 +3298,11 @@ impl DeepSeekClient { } } + pub(crate) async fn provider_quota_pause_remaining(&self) -> Result> { + self.shared_quota.initialize().await?; + Ok(self.shared_quota.remaining()) + } + async fn mark_request_success(&self) { let mut health = self.connection_health.lock().await; if apply_request_success(&mut health, Instant::now()) { @@ -2995,10 +3361,7 @@ impl DeepSeekClient { where F: FnMut() -> reqwest::RequestBuilder, { - if self.isolated_request_state { - return self.send_with_isolated_retry(build).await; - } - self.send_retry_loop(build, Some(non_streaming_request_envelope())) + self.send_with_retry_total(non_streaming_request_envelope(), build) .await } @@ -3014,10 +3377,8 @@ impl DeepSeekClient { where F: FnMut() -> reqwest::RequestBuilder, { - if self.isolated_request_state { - return self.send_with_isolated_retry(build).await; - } - self.send_retry_loop(build, Some(total)).await + self.send_with_retry_policy(build, Some(total), self.retry.clone().into()) + .await } /// The streaming-open twin of [`Self::send_with_retry`]: the same retry @@ -3028,28 +3389,157 @@ impl DeepSeekClient { /// `stream_open_timeout` around the open and per-chunk idle checks on /// the returned body instead. pub(super) async fn send_stream_open_with_retry(&self, build: F) -> Result + where + F: FnMut() -> reqwest::RequestBuilder, + { + self.send_with_retry_policy(build, None, self.retry.clone().into()) + .await + } + + async fn send_with_retry_policy( + &self, + build: F, + attempt_total: Option, + retry_cfg: LlmRetryConfig, + ) -> Result where F: FnMut() -> reqwest::RequestBuilder, { if self.isolated_request_state { - return self.send_with_isolated_retry(build).await; + return self + .send_with_isolated_retry(build, attempt_total, &retry_cfg) + .await; + } + self.send_retry_loop(build, attempt_total, retry_cfg).await + } + + async fn remember_retry_minimum(&self, error: &LlmError) -> std::result::Result<(), LlmError> { + if matches!( + error, + LlmError::RateLimited { .. } | LlmError::QuotaWindow { .. } + ) && let Some(delay) = error.suggested_retry_delay() + { + self.shared_quota.pause(delay); + self.shared_quota + .persist() + .await + .map_err(|message| LlmError::QuotaWindow { + message, + retry_after: Some(delay), + })?; + } + Ok(()) + } + + fn disclose_http_error(&self, status: u16, raw: &str) -> String { + fn redact_strings(value: &mut serde_json::Value, secrets: &[String]) { + match value { + serde_json::Value::String(text) => *text = redact_model_bound_text(text, secrets), + serde_json::Value::Array(values) => { + for value in values { + redact_strings(value, secrets); + } + } + serde_json::Value::Object(fields) => { + for value in fields.values_mut() { + redact_strings(value, secrets); + } + } + _ => {} + } + } + let raw = match serde_json::from_str::(raw) { + Ok(mut value) => { + redact_strings(&mut value, &self.model_bound_secret_values); + value.to_string() + } + Err(_) => raw.to_owned(), + }; + let raw = redact_model_bound_text(&raw, &self.model_bound_secret_values); + let safe = sanitize_http_error_body(Some(self.api_provider.display_name()), status, &raw); + let safe = if self.wire_format == WireFormat::AnthropicMessages && status >= 400 { + let (error_type, _) = anthropic::parse_anthropic_error_envelope(&raw); + let error_type = if error_type.len() <= 64 + && error_type + .bytes() + .all(|byte| byte.is_ascii_alphanumeric() || byte == b'_') + { + error_type.as_str() + } else { + "unknown" + }; + format!("Anthropic API error (HTTP {status} {error_type}): {safe}") + } else { + safe + }; + redact_model_bound_text(&safe, &self.model_bound_secret_values) + } + + fn classify_http_error( + &self, + status: u16, + raw: &str, + safe: &str, + hint: Option, + ) -> LlmError { + let hint = hint + .max(crate::llm_client::rate_limit::body_delay(raw)) + .max(crate::llm_client::retry_delay_from_error_body(raw)); + let route = LlmError::from_http_response_with_retry_after(status, safe, hint); + match crate::llm_client::rate_limit::provider_error( + self.retry_provider_surface(), + status, + raw, + safe, + hint, + ) { + Some(LlmError::QuotaExhausted(_)) if matches!(route, LlmError::QuotaExhausted(_)) => { + route + } + Some(native) => native, + None => route, } - self.send_retry_loop(build, None).await + } + + fn retry_provider_surface(&self) -> &str { + if self.api_provider == ApiProvider::Custom + && let Ok(url) = reqwest::Url::parse(&self.base_url) + && url.host_str().is_some_and(|host| { + host == "dashscope.aliyuncs.com" + || host.ends_with(".dashscope.aliyuncs.com") + || host == "dashscope-intl.aliyuncs.com" + }) + { + return "dashscope"; + } + self.api_provider.as_str() } async fn send_retry_loop( &self, mut build: F, attempt_total: Option, + retry_cfg: LlmRetryConfig, ) -> Result where F: FnMut() -> reqwest::RequestBuilder, { - let retry_cfg: LlmRetryConfig = self.retry.clone().into(); + self.shared_quota.initialize().await?; + let operation_recovery = match self.operation_id { + Some(id) => self.recovery_for_operation(id)?, + None => crate::llm_client::rate_limit::operation(), + }; + operation_recovery + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .wire_attempts + .get_or_insert(0); + let callback_quota = self.shared_quota.clone(); + let callback_permit = self.operation_permit.clone(); // Two bounded layers around the non-streaming completion // (`attempt_total` = `Some`): the per-attempt request total (connect // through body end) and an envelope around the whole retry loop (all - // attempts + backoff + honored Retry-After). The shared client + // active attempts, excluding backoff and honored Retry-After). The shared client // intentionally has no client-level total timeout — streaming is // protected per-chunk instead — so without these nothing bounds a // non-streaming completion: a provider that accepts the connection @@ -3059,61 +3549,293 @@ impl DeepSeekClient { // any total here would be inherited by the returned body and // truncate the stream, so those calls keep only the caller's own // open budget. - let retry_future = with_retry( - &retry_cfg, - || { - // Per-attempt total: unlike the loop envelope below, - // reqwest's per-request timeout also covers the response - // body, so a slow-drip body cannot outlive the budget. - let request = match attempt_total { - Some(total) => self.with_operation_header(build()).timeout(total), - None => self.with_operation_header(build()), - }; - async move { - // Sleep in bounded slices rather than the full remaining - // window: the pause is process-global, so a concurrent - // `clear_rate_limit()` (or a shortened deadline) must - // release requests that are already waiting instead of - // stranding them for the whole original window. - while let Some(delay) = crate::retry_status::rate_limit_remaining() { - tokio::time::sleep(delay.min(RATE_LIMIT_PAUSE_RECHECK_INTERVAL)).await; + // Auxiliary callers also need one recovery/receipt owner throughout admission and retry. + let retry_future = crate::llm_client::rate_limit::with_operation( + operation_recovery.clone(), + with_retry( + &retry_cfg, + || { + // Per-attempt total: unlike the loop envelope below, + // reqwest's per-request timeout also covers the response + // body, so a slow-drip body cannot outlive the budget. + let request = match attempt_total { + Some(total) => self.with_operation_header(build()).timeout(total), + None => self.with_operation_header(build()), + }; + async move { + let recovery = crate::llm_client::rate_limit::operation(); + let started = recovery + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .started; + let recovery_budget = + Duration::from_secs(self.retry.rate_limit.recovery_budget_secs); + let admission_deadline = + started.unwrap_or_else(tokio::time::Instant::now) + recovery_budget; + let permit = self + .operation_permit + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .upgrade(); + let mut observed_deadline = None; + loop { + while let Some(delay) = self.shared_quota.remaining() { + let deadline = *self + .shared_quota + .pause + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if let Some(permit) = &permit { + permit.suspend(); + } + if delay + >= admission_deadline + .saturating_duration_since(tokio::time::Instant::now()) + { + return Err(LlmError::QuotaWindow { message: "Shared quota recovery paused; continue after the provider reset".into(), retry_after: Some(delay) }); + } + if deadline != observed_deadline { + observed_deadline = deadline; + let attempts = recovery + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .attempts; + crate::llm_client::rate_limit::observe( + attempts, + self.retry.rate_limit.max_attempts, + Some(delay), + "waiting", + ) + .await; + } + crate::llm_client::rate_limit::wait( + delay.min(RATE_LIMIT_PAUSE_RECHECK_INTERVAL), + ) + .await; + } + let admission = async { + self.wait_for_rate_limit().await; + if let Some(permit) = &permit { + crate::llm_client::rate_limit::without_active_time( + permit.resume(), + ) + .await; + } + }; + if started.is_some() || observed_deadline.is_some() { + let admitted = + tokio::time::timeout_at(admission_deadline, admission).await; + if admitted.is_err() + || tokio::time::Instant::now() >= admission_deadline + { + return Err(LlmError::QuotaWindow { + message: "Model recovery budget exhausted while waiting for admission".into(), + retry_after: self.shared_quota.remaining(), + }); + } + } else { + admission.await; + } + // A peer may have refused while this request queued for tokens or a slot. + if self.shared_quota.remaining().is_some() { + if let Some(permit) = &permit { + permit.suspend(); + } + continue; + } + if let Some(remaining) = + self.shared_quota.admit_recovery(&self.retry.rate_limit) + { + return Err(LlmError::QuotaWindow { + message: "Shared quota recovery attempt budget exhausted" + .into(), + retry_after: Some(remaining), + }); + } + break; + } + { + let recovery = crate::llm_client::rate_limit::operation(); + let mut recovery = recovery + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + let attempts = recovery.wire_attempts.get_or_insert(0); + *attempts = attempts.saturating_add(1); + } + crate::llm_client::rate_limit::set_response_headers_received(false); + let response = request + .send() + .await + .map_err(|err| LlmError::from_reqwest(&err))?; + crate::llm_client::rate_limit::set_response_headers_received(true); + let open_timeout = stream_entry::stream_open_timeout(); + let inspection = + crate::llm_client::rate_limit::inspect_json_envelope_with_disclosure( + self.api_provider.as_str(), + response, + Some((open_timeout, self.stream_idle_timeout)), + |status, raw| self.disclose_http_error(status, raw), + ); + let inspected = if attempt_total.is_none() { + let envelope_timeout = + open_timeout.saturating_add(self.stream_idle_timeout); + crate::llm_client::rate_limit::without_active_time(async { + tokio::time::timeout(envelope_timeout, inspection) + .await + .unwrap_or(Err(LlmError::Timeout(envelope_timeout))) + }) + .await + .map_err(|error| match error { + // The provider already accepted this POST. An + // incomplete JSON envelope is a response error, + // never a reason to open another stream. + LlmError::Timeout(_) + | LlmError::NetworkError(_) + | LlmError::Other(_) => LlmError::ParseError(format!( + "Could not read the provider JSON envelope: {error}" + )), + error => error, + }) + } else { + inspection.await + }; + let start_recovery = || { + let mut recovery = recovery + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + let started = *recovery + .started + .get_or_insert_with(tokio::time::Instant::now); + started + recovery_budget + }; + let response = match inspected { + Ok(response) => response, + Err(error) => { + crate::llm_client::rate_limit::record_stream_response_error(&error); + if matches!( + error, + LlmError::RateLimited { .. } | LlmError::QuotaWindow { .. } + ) { + let deadline = start_recovery(); + let retry_after = error.suggested_retry_delay(); + if let Some(delay) = retry_after { + self.shared_quota.pause(delay); + } + let record_refusal = async { + self.shared_quota.refused().await; + self.remember_retry_minimum(&error).await + }; + match tokio::time::timeout_at(deadline, record_refusal).await { + Ok(result) => result?, + Err(_) => return Err(LlmError::QuotaWindow { + message: "Model recovery deadline reached while recording the provider quota response".into(), + retry_after, + }), + } + } + return Err(error); + } + }; + let status = response.status(); + if status.is_success() { + self.shared_quota.recovered().await; + return Ok(response); + } + let retry_after = extract_retry_after(response.headers()); + let quota_response = status == reqwest::StatusCode::TOO_MANY_REQUESTS; + let recovery_deadline = if quota_response { + crate::llm_client::rate_limit::record_stream_open_quota(retry_after); + // Honor the explicit header for this quota before + // diagnostic reads or another operation's admission. + if let Some(delay) = retry_after { + self.shared_quota.pause(delay); + } + Some(start_recovery()) + } else { + None + }; + let classify = async { + if quota_response { + if let Some(delay) = retry_after { + self.shared_quota.persist().await.map_err(|message| { + LlmError::QuotaWindow { + message, + retry_after: Some(delay), + } + })?; + } + self.shared_quota.refused().await; + } + record_provider_response(self.api_provider, status.as_u16()); + let deadline = if attempt_total.is_none() { + let diagnostic_deadline = tokio::time::Instant::now() + + self + .stream_idle_timeout + .min(stream_entry::stream_open_timeout()); + Some(recovery_deadline.map_or(diagnostic_deadline, |deadline| { + deadline.min(diagnostic_deadline) + })) + } else { + recovery_deadline + }; + let body_read = + bounded_error_text_until(response, ERROR_BODY_MAX_BYTES, deadline); + let raw_body = if attempt_total.is_none() { + crate::llm_client::rate_limit::without_active_time(body_read).await + } else { + body_read.await + }; + let body = self.disclose_http_error(status.as_u16(), &raw_body); + let error = self.classify_http_error( + status.as_u16(), + &raw_body, + &body, + retry_after, + ); + crate::llm_client::rate_limit::record_stream_response_error(&error); + self.remember_retry_minimum(&error).await?; + Err(error) + }; + if let Some(deadline) = recovery_deadline { + match tokio::time::timeout_at(deadline, classify).await { + Ok(result) => result, + Err(_) => { + let error = LlmError::QuotaWindow { + message: "Model recovery deadline reached while processing the provider response".into(), + retry_after, + }; + crate::llm_client::rate_limit::record_stream_response_error( + &error, + ); + Err(error) + } + } + } else { + classify.await + } } - self.wait_for_rate_limit().await; - let response = request - .send() - .await - .map_err(|err| LlmError::from_reqwest(&err))?; - let status = response.status(); - if status.is_success() { - return Ok(response); + }, + Some(Box::new(move |err: &LlmError, attempt, delay| { + if let Some(permit) = callback_permit + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .upgrade() + { + permit.suspend(); } - let retry_after = extract_retry_after(response.headers()); - let body = bounded_error_text(response, ERROR_BODY_MAX_BYTES).await; - let body = sanitize_http_error_body( - Some(self.api_provider.display_name()), - status.as_u16(), - &body, - ); - Err(LlmError::from_http_response_with_retry_after( - status.as_u16(), - &body, - retry_after, - )) - } - }, - Some(Box::new(|err: &LlmError, attempt, delay| { - let (reason_label, human_reason) = retry_reason_label_and_human(err); - logging::warn(format!( - "HTTP retry reason={} attempt={} delay={:.2}s", - reason_label, - attempt + 1, - delay.as_secs_f64(), - )); - if matches!(err, LlmError::RateLimited { .. }) { - crate::retry_status::note_rate_limit(delay); - } - crate::retry_status::start(attempt + 1, delay, human_reason); - })), + let (reason_label, human_reason) = retry_reason_label_and_human(err); + logging::warn(format!( + "HTTP retry reason={} attempt={} delay={:.2}s", + reason_label, + attempt + 1, + delay.as_secs_f64(), + )); + if matches!(err, LlmError::RateLimited { .. }) { + callback_quota.pause(delay); + } + crate::retry_status::start(attempt + 1, delay, human_reason); + })), + ), ); let request_result = if let Some(total) = attempt_total { // The loop envelope must dominate the per-attempt total it @@ -3121,7 +3843,9 @@ impl DeepSeekClient { // the shared envelope, and the envelope must never strangle // its own attempts. let loop_envelope = total.max(non_streaming_request_envelope()); - match tokio::time::timeout(loop_envelope, retry_future).await { + match crate::llm_client::rate_limit::active_timeout_total(loop_envelope, retry_future) + .await + { Ok(result) => result, Err(_elapsed) => { let last = LlmError::Timeout(loop_envelope); @@ -3147,11 +3871,27 @@ impl DeepSeekClient { Ok(response) } Err(err) => { - if let LlmError::RateLimited { retry_after, .. } = &err.last_error { - crate::retry_status::note_rate_limit( + if let LlmError::RateLimited { retry_after, .. } + | LlmError::QuotaWindow { retry_after, .. } = &err.last_error + { + self.shared_quota.pause( retry_after .unwrap_or_else(|| retry_cfg.delay_for_attempt(retry_cfg.max_retries)), ); + let deadline = operation_recovery + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .started + .map(|started| { + started + + Duration::from_secs(self.retry.rate_limit.recovery_budget_secs) + }); + if let Some(deadline) = deadline { + let _ = + tokio::time::timeout_at(deadline, self.shared_quota.persist()).await; + } else { + let _ = self.shared_quota.persist().await; + } } let last = err.last_error.to_string(); if err.attempts > 1 { @@ -3159,8 +3899,38 @@ impl DeepSeekClient { } else { crate::retry_status::clear(); } - self.mark_request_failure(&last).await; - self.maybe_probe_recovery().await; + logging::warn(format!( + "model_recovery_stopped provider={} attempts={} elapsed_ms={} resumable={}", + self.api_provider.as_str(), + err.attempts, + err.total_time.as_millis(), + matches!( + err.last_error, + LlmError::RateLimited { .. } | LlmError::QuotaWindow { .. } + ) + )); + self.mark_request_failure(retry_reason_label_and_human(&err.last_error).0) + .await; + // Definitive HTTP stream errors must reach the caller without + // a /models probe that could turn them into a header timeout. + let transport_failure = matches!( + err.last_error, + LlmError::NetworkError(_) | LlmError::Timeout(_) + ); + let ordinary_wire = self.wire_format != WireFormat::AnthropicMessages; + let probe = if attempt_total.is_some() { + ordinary_wire || transport_failure + } else { + ordinary_wire && transport_failure + }; + if probe + && !matches!( + err.last_error, + LlmError::RateLimited { .. } | LlmError::QuotaWindow { .. } + ) + { + self.maybe_probe_recovery().await; + } // Keep the structured `LlmError` downcastable so failure // surfaces can classify auth/rate-limit/invalid-request // instead of reporting an opaque string (#3884). @@ -3172,37 +3942,50 @@ impl DeepSeekClient { /// The same bounded transport retry policy without process-global retry /// banners, provider-wide pause cells, or shared connection-health writes. /// Used only by the Auto classifier during read-only request inspection. - async fn send_with_isolated_retry(&self, mut build: F) -> Result + async fn send_with_isolated_retry( + &self, + mut build: F, + attempt_total: Option, + retry_cfg: &LlmRetryConfig, + ) -> Result where F: FnMut() -> reqwest::RequestBuilder, { - let retry_cfg: LlmRetryConfig = self.retry.clone().into(); - let request_result = with_retry( - &retry_cfg, + let retry_future = with_retry( + retry_cfg, || { let request = self.with_operation_header(build()); + let request = match attempt_total { + Some(total) => request.timeout(total), + None => request, + }; async move { self.wait_for_rate_limit().await; + crate::llm_client::rate_limit::set_response_headers_received(false); let response = request .send() .await .map_err(|err| LlmError::from_reqwest(&err))?; + crate::llm_client::rate_limit::set_response_headers_received(true); + let response = + crate::llm_client::rate_limit::inspect_json_envelope_with_disclosure( + self.api_provider.as_str(), + response, + Some(( + stream_entry::stream_open_timeout(), + self.stream_idle_timeout, + )), + |status, raw| self.disclose_http_error(status, raw), + ) + .await?; let status = response.status(); if status.is_success() { return Ok(response); } let retry_after = extract_retry_after(response.headers()); - let body = bounded_error_text(response, ERROR_BODY_MAX_BYTES).await; - let body = sanitize_http_error_body( - Some(self.api_provider.display_name()), - status.as_u16(), - &body, - ); - Err(LlmError::from_http_response_with_retry_after( - status.as_u16(), - &body, - retry_after, - )) + let raw_body = bounded_error_text(response, ERROR_BODY_MAX_BYTES).await; + let body = self.disclose_http_error(status.as_u16(), &raw_body); + Err(self.classify_http_error(status.as_u16(), &raw_body, &body, retry_after)) } }, Some(Box::new(|err, attempt, delay| { @@ -3214,12 +3997,72 @@ impl DeepSeekClient { delay.as_secs_f64(), )); })), - ) - .await; - + ); + let request_result = if let Some(total) = attempt_total { + let envelope = total.max(non_streaming_request_envelope()); + crate::llm_client::rate_limit::active_timeout_total(envelope, retry_future) + .await + .map_err(|_| LlmError::Timeout(envelope))? + } else { + retry_future.await + }; request_result.map_err(|err| anyhow::Error::new(err.last_error)) } + /// Reuse the frozen route only when endpoint and credential match. A + /// dedicated vision endpoint must not inherit the main route's headers. + pub(crate) fn auxiliary_json_route( + parent: Option<&Self>, + base_url: &str, + api_key: &str, + model: &str, + ) -> Result { + if let Some(parent) = parent.filter(|p| p.matches_auxiliary_chat_route(base_url, api_key)) { + return Ok(parent.clone()); + } + let config = Config { + provider: Some("custom".into()), + api_key: Some(api_key.into()), + base_url: Some(base_url.into()), + default_text_model: Some(model.into()), + rate_limit_recovery_dir: parent.and_then(|p| p.shared_quota.storage.clone()), + ..Default::default() + }; + let mut client = Self::from_parts_with_exact_key( + base_url.to_owned(), + model.to_owned(), + WireFormat::ChatCompletions, + None, + &config, + Some(api_key), + )?; + if let Some(parent) = parent { + client.retry = parent.retry.clone(); + } + Ok(client) + } + + pub(crate) async fn request_auxiliary_json( + &self, + url: &str, + body: &Value, + operation_id: uuid::Uuid, + ) -> Result { + crate::llm_client::rate_limit::with_operation( + self.recovery_for_operation(operation_id)?, + async { + let call = self.for_operation(operation_id); + let _permit = call.acquire_provider_request_permit().await; + let response = call.send_json_with_retry(url, body).await?; + response + .json() + .await + .context("Failed to parse auxiliary model response") + }, + ) + .await + } + pub(super) async fn send_json_with_retry( &self, url: &str, @@ -3241,17 +4084,29 @@ impl DeepSeekClient { /// [`Self::send_stream_open_with_retry`]). pub(super) async fn open_stream_json_with_retry( &self, + client: &reqwest::Client, url: &str, body: &serde_json::Value, + policy: stream_entry::StreamHttpPolicy, ) -> Result { let request_body = serde_json::to_vec(body).context("Failed to serialize JSON request body")?; - self.send_stream_open_with_retry(|| { - self.http_client - .post(url) - .header(CONTENT_TYPE, "application/json") - .body(request_body.clone()) - }) + let mut retry: LlmRetryConfig = self.retry.clone().into(); + if policy == stream_entry::StreamHttpPolicy::Http1Only { + // Retain the single H1 send for ordinary failures. Provider quota + // recovery uses its own policy and still applies to this client. + retry.max_retries = 0; + } + self.send_with_retry_policy( + || { + client + .post(url) + .header(CONTENT_TYPE, "application/json") + .body(request_body.clone()) + }, + None, + retry, + ) .await } } @@ -3291,7 +4146,7 @@ fn retry_reason_label_and_human(err: &LlmError) -> (&'static str, String) { .bump_error(codewhale_telemetry::ErrorCounter::NetworkError); } match err { - LlmError::RateLimited { retry_after, .. } => { + LlmError::RateLimited { retry_after, .. } | LlmError::QuotaWindow { retry_after, .. } => { let human = if let Some(after) = retry_after { format!("rate limited (Retry-After {}s)", after.as_secs()) } else { @@ -3325,6 +4180,9 @@ impl DeepSeekClient { // auxiliary classifier call, however: it must neither consume nor // inherit that mutable foreground state. isolated.rate_limiter = Arc::new(AsyncMutex::new(TokenBucket::from_env())); + isolated.request_concurrency = None; + isolated.operation_recovery = Arc::new(std::sync::Mutex::new(HashMap::new())); + isolated.operation_permit = Arc::default(); let _inference = isolated.acquire_remote_control_inference_permit().await; let _permit = isolated.acquire_provider_request_permit().await; let prepared = isolated.prepare_outbound_request(request, false)?; @@ -3339,9 +4197,44 @@ impl DeepSeekClient { impl DeepSeekClient { // A request-local clone owns exactly one outbound logical call. Do not // reuse it for an auxiliary request with a different body. + fn recovery_for_operation( + &self, + id: uuid::Uuid, + ) -> Result>> { + let mut operations = self + .operation_recovery + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + operations.retain(|_, recovery| { + Arc::strong_count(recovery) > 1 + || recovery + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .created + .elapsed() + < Duration::from_secs(86400) + }); + anyhow::ensure!( + operations.len() < 65536 || operations.contains_key(&id), + "Too many retained model recovery operations" + ); + Ok(operations + .entry(id) + .or_insert_with(|| { + Arc::new(std::sync::Mutex::new( + crate::llm_client::rate_limit::OperationRecovery { + operation_id: Some(id.to_string()), + ..Default::default() + }, + )) + }) + .clone()) + } + fn for_operation(&self, operation_id: uuid::Uuid) -> Self { let mut client = self.clone(); client.operation_id = Some(operation_id); + client.operation_permit = Arc::default(); client } @@ -3375,9 +4268,8 @@ impl DeepSeekClient { None => builder, } } - pub(crate) fn idempotency_matches_route(&self, base_url: &str, api_key: &str) -> bool { - self.request_idempotency_header.is_some() - && self.wire_format == WireFormat::ChatCompletions + pub(crate) fn matches_auxiliary_chat_route(&self, base_url: &str, api_key: &str) -> bool { + self.wire_format == WireFormat::ChatCompletions && !self.auth_disabled && !(self.api_provider == ApiProvider::XiaomiMimo && (xiaomi_mimo_base_url_uses_token_plan(&self.base_url) @@ -3422,7 +4314,7 @@ impl DeepSeekClient { &self, prepared: PreparedOutboundRequest, inference: Option, - permit: Option, + permit: Option>, ) -> Result { let projection_warning = (!prepared.omitted_tool_names.is_empty()).then(|| { let omitted_tool_count = prepared.omitted_tool_names.len(); @@ -3535,18 +4427,23 @@ impl LlmClient for DeepSeekClient { request: MessageRequest, operation_id: uuid::Uuid, ) -> Result { - self.for_operation(operation_id) - .perform_message(request) - .await + crate::llm_client::rate_limit::with_operation( + self.recovery_for_operation(operation_id)?, + self.for_operation(operation_id).perform_message(request), + ) + .await } async fn create_message_stream_for_operation( &self, request: MessageRequest, operation_id: uuid::Uuid, ) -> Result { - self.for_operation(operation_id) - .perform_message_stream(request) - .await + crate::llm_client::rate_limit::with_operation( + self.recovery_for_operation(operation_id)?, + self.for_operation(operation_id) + .perform_message_stream(request), + ) + .await } fn stream_operation_identity(&self, request: &MessageRequest) -> Result { if self.request_idempotency_header.is_none() { @@ -3615,9 +4512,11 @@ impl LlmClient for DeepSeekClient { .prepare_outbound_request(call.request, true) .map_err(crate::llm_client::DeferredStreamPreparationError)?, }; - client - .finish_prepared_stream(prepared, inference, permit) - .await + crate::llm_client::rate_limit::with_operation( + client.recovery_for_operation(operation_id)?, + client.finish_prepared_stream(prepared, inference, permit), + ) + .await } } @@ -4382,7 +5281,6 @@ impl DeepSeekClient { max_tokens: u32, ) -> anyhow::Result { let _inference = self.acquire_remote_control_inference_permit().await; - let _permit = self.acquire_provider_request_permit().await; if self.api_provider == ApiProvider::OpencodeZen || self.wire_format != WireFormat::ChatCompletions { @@ -4402,6 +5300,7 @@ impl DeepSeekClient { "max_tokens": max_tokens, }); let call = self.for_operation(uuid::Uuid::new_v4()); + let _permit = call.acquire_provider_request_permit().await; let response = call.send_json_with_retry(&url, &body).await?; let status = response.status(); if !status.is_success() { @@ -8022,6 +8921,8 @@ mod tests { .acquire_provider_request_permit() .await .expect("zai default should install provider request limiter"); + assert_eq!(client.active_provider_requests(), 0); + permit.resume().await; assert_eq!(client.active_provider_requests(), 1); assert_eq!(clone.active_provider_requests(), 1); @@ -8038,6 +8939,8 @@ mod tests { .acquire_provider_request_permit() .await .expect("zai default should install provider request limiter"); + assert_eq!(client.active_provider_requests(), 0); + permit.resume().await; let stream: crate::llm_client::StreamEventBox = Box::pin(futures_util::stream::iter( vec![Ok(StreamEvent::MessageStop)], )); @@ -11917,6 +12820,7 @@ mod tests { capacity: 1.0, tokens: 1.0, refill_per_sec: 2.0, + normal_refill_per_sec: 2.0, last_refill: now, }; @@ -11944,6 +12848,7 @@ mod tests { capacity: 1.0, tokens: 1.0, refill_per_sec: 1.0, + normal_refill_per_sec: 1.0, last_refill: now, }; @@ -12946,3 +13851,7 @@ mod tests { assert_eq!(body["model"], "custom-claude", "model id verbatim: {body}"); } } + +#[cfg(test)] +#[path = "client/rate_limit_tests.rs"] +mod rate_limit_tests; diff --git a/crates/tui/src/client/anthropic.rs b/crates/tui/src/client/anthropic.rs index 68e8b9460f..812f70148a 100644 --- a/crates/tui/src/client/anthropic.rs +++ b/crates/tui/src/client/anthropic.rs @@ -36,7 +36,7 @@ use crate::tools::schema_sanitize; use super::prepared::WireDialect; use super::role_placement::{RolePlacement, role_placement}; use super::wire::{extract_sse_data_value, next_sse_line}; -use super::{DeepSeekClient, ERROR_BODY_MAX_BYTES, bounded_error_text}; +use super::{DeepSeekClient, LlmRetryConfig, non_streaming_request_envelope}; /// Maximum `cache_control` breakpoints Anthropic accepts per request. const MAX_CACHE_BREAKPOINTS: usize = 4; @@ -273,70 +273,23 @@ impl DeepSeekClient { async fn send_anthropic_request(&self, url: &str, body: &Value) -> Result { let url = self.messages_transport_url(url); - self.wait_for_rate_limit().await; - // This request is consumed non-streaming (the body is decoded as one - // JSON document), so it gets the same total budget as every other - // non-streaming completion; without it a provider that accepts and - // stalls - or trickles the body - wedged the caller indefinitely. - let response = self - .with_operation_header(self.http_client.post(&url)) - .header("Accept", "text/event-stream") - .timeout(crate::client::non_streaming_request_envelope()) - .json(body) - .send() - .await; - let response = match response { - Ok(response) => response, - Err(error) => { - let rendered = if error.is_timeout() { - "anthropic request exceeded the envelope".to_string() - } else { - format!("anthropic transport failure: {error}") - }; - // A transport-level failure never reaches the status checks - // below, so it must mark health and probe on its own — - // otherwise a provider that stalls past the envelope leaves - // connection health stale until the next request retries. - // The isolated classifier request is the exception: a - // read-only inspection must not write the shared connection - // health (or fire a /models probe), matching the isolated - // dispatch contract of the generic send paths. - if !self.isolated_request_state { - self.mark_request_failure(&rendered).await; - self.maybe_probe_recovery().await; - } - return Err( - anyhow::Error::new(error).context("Anthropic Messages API request failed") - ); - } - }; - self.check_anthropic_response(response).await + self.send_with_retry_policy( + || { + self.http_client + .post(&url) + .header("Accept", "text/event-stream") + .json(body) + }, + Some(non_streaming_request_envelope()), + self.anthropic_retry_policy(), + ) + .await } - /// Shared status/error-envelope handling for streaming and - /// non-streaming Messages responses. - async fn check_anthropic_response( - &self, - response: reqwest::Response, - ) -> Result { - let status = response.status(); - if !status.is_success() { - let raw = bounded_error_text(response, ERROR_BODY_MAX_BYTES).await; - let (error_type, message) = parse_anthropic_error_envelope(&raw); - // Same isolation contract as the transport-failure arm above: a - // status-level failure on the isolated classifier request must - // not write the shared connection health (the generic dialect's - // isolated dispatch path writes nothing either). - if !self.isolated_request_state { - self.mark_request_failure(&format!("anthropic status={status}")) - .await; - } - anyhow::bail!("Anthropic API error (HTTP {status} {error_type}): {message}"); - } - if !self.isolated_request_state { - self.mark_request_success().await; - } - Ok(response) + fn anthropic_retry_policy(&self) -> LlmRetryConfig { + // Keep Messages' existing transport attempt policy. Rate-limit + // recovery has its own budget and does not consume these retries. + LlmRetryConfig::from(self.retry.clone()).with_max_retries(0) } /// Open the streaming Messages request through the shared stream-entry @@ -357,30 +310,35 @@ impl DeepSeekClient { let opened = super::stream_entry::open_sse_response(&open_req, |policy| { let url = url.clone(); async move { - self.wait_for_rate_limit().await; let client = super::stream_entry::client_for_policy( &self.http_client, self.http1_fallback_client(), policy, ); - self.with_operation_header(client.post(&url)) - .header("Accept", "text/event-stream") - .json(body) - .send() - .await - .context("Anthropic Messages API request failed") + self.send_with_retry_policy( + || { + client + .post(&url) + .header("Accept", "text/event-stream") + .json(body) + }, + None, + self.anthropic_retry_policy(), + ) + .await } }) .await; - let response = match opened { - Ok(response) => response, - Err(err) => { - self.mark_request_failure(&format!("anthropic stream open: {err}")) - .await; - return Err(err); - } - }; - self.check_anthropic_response(response).await + // The shared deadline can cancel a send before it returns an error + // to the sender's health accounting. Other errors were handled there. + if let Err(error) = &opened + && !self.isolated_request_state + && error.is::() + { + self.mark_request_failure(&format!("anthropic stream open: {error}")) + .await; + } + opened } /// Handle a streaming Messages API request. @@ -1092,7 +1050,7 @@ fn parse_anthropic_usage(usage: &Value) -> Usage { /// Extract `error.type` / `error.message` from an Anthropic error envelope /// (`{"type":"error","error":{"type":...,"message":...}}`), falling back to /// the raw body so nothing is swallowed. -fn parse_anthropic_error_envelope(raw: &str) -> (String, String) { +pub(super) fn parse_anthropic_error_envelope(raw: &str) -> (String, String) { let Ok(value) = serde_json::from_str::(raw) else { return ("unknown".to_string(), raw.to_string()); }; @@ -2782,33 +2740,38 @@ mod tests { let server = MockServer::start().await; // A definitive provider error before any stream body must fail fast: - // exactly one request, no H1 fallback, envelope preserved. + // one request per call, no H1 fallback or probe, envelope preserved. Mock::given(method("POST")) .and(path("/v1/messages")) .respond_with(ResponseTemplate::new(401).set_body_string( "{\"error\":{\"type\":\"authentication_error\",\"message\":\"bad key\"}}", )) - .expect(1) + .expect(2) .mount(&server) .await; - let client = deepseek_test_client(&server.uri()); - let err = match client - .handle_anthropic_stream( - &client - .prepare_outbound_request(request_with("deepseek-v4", None, None, None), true) - .expect("anthropic request prepares"), - ) - .await - { - Ok(_) => panic!("auth errors must fail fast"), - Err(err) => err, - }; - let text = err.to_string(); - assert!( - text.contains("HTTP 401") && text.contains("authentication_error"), - "error envelope should be preserved: {text}" - ); + Mock::given(method("GET")) + .respond_with(ResponseTemplate::new(500)) + .expect(0) + .mount(&server) + .await; + + let client = anthropic_test_client(Some(&server.uri())); + for _ in 0..2 { + let prepared = client + .prepare_outbound_request(request_with("claude-opus-4-6", None, None, None), true) + .expect("anthropic request prepares"); + let err = match client.handle_anthropic_stream(&prepared).await { + Ok(_) => panic!("auth errors must fail fast"), + Err(err) => err, + }; + let text = err.to_string(); + assert!( + text.contains("HTTP 401") && text.contains("authentication_error"), + "error envelope should be preserved: {text}" + ); + } + server.verify().await; } /// A `system`-role history message — what a compaction summary, a branch diff --git a/crates/tui/src/client/chat.rs b/crates/tui/src/client/chat.rs index 83d3cbee1c..42fe9e3500 100644 --- a/crates/tui/src/client/chat.rs +++ b/crates/tui/src/client/chat.rs @@ -1348,30 +1348,13 @@ impl DeepSeekClient { ); let idle_timeout = open_req.idle_timeout; let response = super::stream_entry::open_sse_response(&open_req, |policy| async move { - match policy { - // The prebuilt HTTP/1.1 twin carries the same default - // headers/auth; send once, without the JSON retry loop - // (matching the pre-seam H1-pin behavior). - super::stream_entry::StreamHttpPolicy::Http1Only => { - let client = super::stream_entry::client_for_policy( - &self.http_client, - self.http1_fallback_client(), - policy, - ); - Ok(self - .with_operation_header(client.post(url)) - .header(reqwest::header::CONTENT_TYPE, "application/json") - .json(body) - .send() - .await?) - } - super::stream_entry::StreamHttpPolicy::DualWithH1Fallback => { - // Stream open: no per-request total — the response body - // outlives the open, and a total would ride on it and - // truncate the stream. - self.open_stream_json_with_retry(url, body).await - } - } + let client = super::stream_entry::client_for_policy( + &self.http_client, + self.http1_fallback_client(), + policy, + ); + self.open_stream_json_with_retry(client, url, body, policy) + .await }) .await?; Ok((response, idle_timeout)) diff --git a/crates/tui/src/client/operation_tests.rs b/crates/tui/src/client/operation_tests.rs index 30ff6e708e..96e89e2892 100644 --- a/crates/tui/src/client/operation_tests.rs +++ b/crates/tui/src/client/operation_tests.rs @@ -691,3 +691,47 @@ async fn forkguard_model_operation_http1_stream_keeps_id() { == id.to_string() })); } + +#[tokio::test] +async fn forkguard_rate_limit_transport_preserves_logical_request() { + let server = MockServer::start().await; + let count = Arc::new(std::sync::atomic::AtomicUsize::new(0)); + let seen = count.clone(); + Mock::given(method("POST")).respond_with(move |_: &wiremock::Request| { + if seen.fetch_add(1, Ordering::SeqCst) == 0 { + ResponseTemplate::new(429).insert_header("retry-after", "0") + .set_body_json(json!({"error":{"message":"fixture rate limit"}})) + } else { + ResponseTemplate::new(200).insert_header("content-type", "text/event-stream") + .set_body_string("data: {\"choices\":[{\"index\":0,\"delta\":{\"content\":\"done\"},\"finish_reason\":null}]}\n\ndata: {\"choices\":[{\"index\":0,\"delta\":{},\"finish_reason\":\"stop\"}]}\n\ndata: [DONE]\n\n") + } + }).mount(&server).await; + let mut cfg = config(&server.uri(), Some("idempotency-key")); + let mut policy = crate::config::RateLimitRetryPolicy::default(); + policy.max_attempts = 2; + policy.initial_delay_secs = 1; + policy.max_delay_secs = 1; + cfg.retry = Some(crate::config::RetryConfig { + rate_limit: policy, + enabled: None, + max_retries: Some(0), + initial_delay: None, + max_delay: None, + exponential_base: None, + }); + let client = DeepSeekClient::new(&cfg).unwrap(); + let id = uuid::Uuid::new_v4(); + let mut stream = client + .create_message_stream_for_operation(request(), id) + .await + .unwrap(); + while let Some(event) = stream.next().await { + event.unwrap(); + } + let requests = server.received_requests().await.unwrap(); + assert_eq!(requests.len(), 2); + assert_eq!(requests[0].body, requests[1].body); + for request in requests { + assert_eq!(request.headers["idempotency-key"], id.to_string()); + } +} diff --git a/crates/tui/src/client/rate_limit_tests.rs b/crates/tui/src/client/rate_limit_tests.rs new file mode 100644 index 0000000000..8da984ac83 --- /dev/null +++ b/crates/tui/src/client/rate_limit_tests.rs @@ -0,0 +1,1072 @@ +use super::*; + +#[tokio::test] +async fn forkguard_rate_limit_native_and_raw_errors_reach_both_senders() { + use wiremock::matchers::method; + use wiremock::{Mock, MockServer, ResponseTemplate}; + for isolated in [false, true] { + for escaped in [false, true] { + for (status, raw, expected) in [ + ( + 200, + r#"{"base_resp":{"status_code":2045,"status_msg":"weekly plan quota exhausted fixture-config-secret"},"retry_after":7200}"#, + "window", + ), + ( + 429, + r#"{"base_resp":{"status_code":1002,"status_msg":"retry in 7200s fixture-config-secret"}}"#, + "window", + ), + ( + 429, + r#"{"error":{"message":"opaque fixture-config-secret"},"detail":"Billing quota has been exhausted"}"#, + "billing", + ), + ] { + let server = MockServer::start().await; + let secret = format!("fixture-secret-{}", uuid::Uuid::new_v4()); + let echoed = if escaped { + secret + .chars() + .map(|ch| format!("\\u{:04x}", u32::from(ch))) + .collect::() + } else { + secret.clone() + }; + let raw = raw.replace("fixture-config-secret", &echoed); + Mock::given(method("POST")) + .respond_with( + ResponseTemplate::new(status).set_body_raw(raw, "application/json"), + ) + .expect(1) + .mount(&server) + .await; + let mut client = DeepSeekClient::new(&Config { + provider: Some("minimax".into()), + providers: Some(crate::config::ProvidersConfig { + minimax: crate::config::ProviderConfig { + api_key: Some(secret.clone()), + base_url: Some(server.uri()), + ..Default::default() + }, + ..Default::default() + }), + ..Default::default() + }) + .unwrap(); + client.isolated_request_state = isolated; + let error = client + .for_operation(uuid::Uuid::new_v4()) + .send_json_with_retry(&server.uri(), &serde_json::json!({"fixture":true})) + .await + .unwrap_err(); + let typed = error.downcast_ref::().unwrap(); + assert!(!error.to_string().contains(&secret)); + assert!(!error.to_string().contains(&echoed)); + let diagnostic = if expected == "billing" { + "opaque" + } else if status == 200 { + "weekly plan quota" + } else { + "retry in 7200s" + }; + assert!(error.to_string().contains(diagnostic)); + if expected == "window" { + assert!( + matches!(typed, LlmError::QuotaWindow { retry_after: Some(delay), .. } if *delay == Duration::from_secs(7200)) + ); + if !isolated { + assert!( + client.shared_quota.remaining().unwrap() > Duration::from_secs(7190) + ); + } + } else { + assert!(matches!(typed, LlmError::QuotaExhausted(_))); + } + assert_eq!(server.received_requests().await.unwrap().len(), 1); + } + } + } +} + +#[tokio::test(start_paused = true)] +async fn forkguard_rate_limit_local_tokens_consume_active_budget() { + let client = DeepSeekClient::new(&Config { + provider: Some("custom".into()), + api_key: Some(uuid::Uuid::new_v4().to_string()), + base_url: Some("https://fixture.example.test/v1".into()), + ..Default::default() + }) + .unwrap(); + *client.rate_limiter.lock().await = TokenBucket { + enabled: true, + capacity: 1.0, + tokens: 0.0, + refill_per_sec: 0.0001, + normal_refill_per_sec: 0.0001, + last_refill: Instant::now(), + }; + assert!( + crate::llm_client::rate_limit::active_timeout_total( + Duration::from_secs(2), + client.wait_for_rate_limit(), + ) + .await + .is_err() + ); +} + +#[tokio::test] +async fn forkguard_rate_limit_unknown_quota_sender_and_billing_boundary() { + use wiremock::matchers::method; + use wiremock::{Mock, MockServer, Request, ResponseTemplate}; + for (provider, message, expected_wires) in [ + ("custom", "TPM exceeded", 2), + ("custom", "Credit balance exhausted", 1), + ("openai", "TPM exceeded", 1), + ] { + let server = MockServer::start().await; + let count = Arc::new(AtomicUsize::new(0)); + let seen = count.clone(); + let first_wire = Arc::new(std::sync::Mutex::new(None)); + let first = first_wire.clone(); + Mock::given(method("POST")) + .respond_with(move |_: &Request| { + if seen.fetch_add(1, Ordering::SeqCst) == 0 { + *first.lock().unwrap() = Some(Instant::now()); + ResponseTemplate::new(429).insert_header("retry-after", "1") + .set_body_json(serde_json::json!({"error":{"type":"insufficient_quota","message":message}})) + } else { + assert!(first.lock().unwrap().unwrap().elapsed() >= Duration::from_secs(1)); + ResponseTemplate::new(200).set_body_string("ok") + } + }) + .mount(&server).await; + let mut client = DeepSeekClient::new(&Config { + provider: Some(provider.into()), + api_key: Some(uuid::Uuid::new_v4().to_string()), + base_url: Some(server.uri()), + providers: Some(crate::config::ProvidersConfig { + openai: crate::config::ProviderConfig { + api_key: Some(uuid::Uuid::new_v4().to_string()), + base_url: Some(server.uri()), + ..Default::default() + }, + ..Default::default() + }), + ..Default::default() + }) + .unwrap(); + client.retry.rate_limit.initial_delay_secs = 1; + client.retry.rate_limit.max_delay_secs = 1; + let result = client + .for_operation(uuid::Uuid::new_v4()) + .send_json_with_retry(&server.uri(), &serde_json::json!({"fixture":true})) + .await; + assert_eq!(count.load(Ordering::SeqCst), expected_wires); + if expected_wires == 2 { + assert_eq!(result.unwrap().text().await.unwrap(), "ok"); + } else { + assert!(matches!( + result.unwrap_err().downcast_ref::(), + Some(LlmError::QuotaExhausted(_)) + )); + } + } +} + +#[test] +fn forkguard_rate_limit_identity_uses_effective_accounts_and_normalized_headers() { + let _env = crate::test_support::lock_test_env(); + let home = tempfile::tempdir().unwrap(); + let _home = crate::test_support::EnvVarGuard::set("CODEWHALE_HOME", home.path()); + let token = crate::test_support::future_test_jwt(&uuid::Uuid::new_v4().to_string()); + let _token = crate::test_support::EnvVarGuard::set("OPENAI_CODEX_ACCESS_TOKEN", token); + let _account = crate::test_support::EnvVarGuard::set("OPENAI_CODEX_ACCOUNT_ID", "account-a"); + let config = Config { + provider: Some("openai-codex".into()), + base_url: Some("https://relay.example.test/v1".into()), + ..Default::default() + }; + let first = DeepSeekClient::new(&config).unwrap(); + let _second_account = + crate::test_support::EnvVarGuard::set("OPENAI_CODEX_ACCOUNT_ID", "account-b"); + let second = DeepSeekClient::new(&config).unwrap(); + assert_ne!(first.shared_quota.key, second.shared_quota.key); + first.shared_quota.pause(Duration::from_secs(3600)); + assert!(second.shared_quota.remaining().is_none()); + + let mut upper = config.clone(); + upper.http_headers = Some(HashMap::from([ + (" OpenAI-Project ".into(), " project-fixture ".into()), + ("Authorization".into(), "ignored-auth-a".into()), + ("X-Request-Id".into(), "trace-a".into()), + ])); + let mut lower = config; + lower.http_headers = Some(HashMap::from([ + ("openai-project".into(), "project-fixture".into()), + ("Authorization".into(), "ignored-auth-b".into()), + ("X-Request-Id".into(), "trace-b".into()), + ])); + let upper = DeepSeekClient::new(&upper).unwrap(); + let lower = DeepSeekClient::new(&lower).unwrap(); + assert!(Arc::ptr_eq(&upper.shared_quota, &lower.shared_quota)); + upper.shared_quota.pause(Duration::from_secs(10)); + assert!(lower.shared_quota.remaining().is_some()); +} + +#[tokio::test] +async fn forkguard_rate_limit_restores_default_capacity_and_durable_deadlines() { + let directory = tempfile::tempdir().unwrap(); + let key = uuid::Uuid::new_v4().to_string(); + let quota = shared_quota(key.clone(), Some(16), Some(directory.path().to_owned())).unwrap(); + quota.initialize().await.unwrap(); + quota.refused().await; + assert_eq!(quota.concurrency.as_ref().unwrap().gate.capacity(), 8); + for _ in 0..24 { + quota.recovered().await; + } + assert_eq!(quota.concurrency.as_ref().unwrap().gate.capacity(), 16); + quota.limiter.lock().await.enabled = false; + quota.refused().await; + for _ in 0..24 { + quota.recovered().await; + } + assert_eq!(quota.concurrency.as_ref().unwrap().gate.capacity(), 16); + + quota.pause(Duration::from_secs(1000)); + quota.persist().await.unwrap(); + quota.pause(Duration::from_secs(3600)); + quota.persist().await.unwrap(); + drop(quota); + forget_quota_process_state_for_restart(&key); + let restored = shared_quota(key.clone(), Some(16), Some(directory.path().to_owned())).unwrap(); + restored.initialize().await.unwrap(); + assert!(restored.remaining().unwrap() > Duration::from_secs(3595)); + drop(restored); + forget_quota_process_state_for_restart(&key); + tokio::fs::write( + directory.path().join(format!("{key}.json")), + b"corrupt fixture", + ) + .await + .unwrap(); + let corrupt = shared_quota(key.clone(), Some(16), Some(directory.path().to_owned())).unwrap(); + assert!(corrupt.initialize().await.is_err()); + // Repairing the receipt permits the next manual call; a failed load is not cached forever. + tokio::fs::write( + directory.path().join(format!("{key}.json")), + serde_json::to_vec(&(chrono::Utc::now() - chrono::Duration::seconds(1))).unwrap(), + ) + .await + .unwrap(); + corrupt.initialize().await.unwrap(); + assert!(corrupt.remaining().is_none()); + let blocked_directory = tempfile::NamedTempFile::new().unwrap(); + let blocked = shared_quota( + uuid::Uuid::new_v4().to_string(), + Some(1), + Some(blocked_directory.path().to_owned()), + ) + .unwrap(); + blocked.pause(Duration::from_secs(3600)); + assert!(blocked.persist().await.is_err()); + assert!(blocked.remaining().unwrap() > Duration::from_secs(3595)); +} + +async fn hold_http_response( + reply: &'static [u8], +) -> (String, Arc, tokio::task::JoinHandle<()>) { + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let url = format!("http://{}", listener.local_addr().unwrap()); + let attempts = Arc::new(AtomicUsize::new(0)); + let observed = attempts.clone(); + let fixture = tokio::spawn(async move { + let mut connections = tokio::task::JoinSet::new(); + while let Ok((mut socket, _)) = listener.accept().await { + let observed = observed.clone(); + connections.spawn(async move { + let mut request = [0; 4096]; + if socket.read(&mut request).await.unwrap_or(0) == 0 { + return; + } + observed.fetch_add(1, Ordering::Relaxed); + let _ = socket.write_all(reply).await; + std::future::pending::<()>().await; + drop(socket); + }); + } + }); + (url, attempts, fixture) +} + +#[tokio::test] +async fn forkguard_rate_limit_stalled_error_body_preserves_the_retry_after() { + use super::stream_entry::{ + StreamHttpPolicy, StreamOpenRequest, client_for_policy, open_sse_response, + }; + + let (url, attempts, fixture) = hold_http_response( + b"HTTP/1.1 429 Too Many Requests\r\nRetry-After: 7200\r\nContent-Length: 128\r\n\r\n", + ) + .await; + let state = tempfile::tempdir().unwrap(); + let mut client = DeepSeekClient::new(&Config { + api_key: Some(format!("owned-fixture-{}", uuid::Uuid::new_v4())), + base_url: Some(url.clone()), + ..Default::default() + }) + .unwrap(); + client.shared_quota = shared_quota( + uuid::Uuid::new_v4().to_string(), + None, + Some(state.path().to_owned()), + ) + .unwrap(); + client.retry.rate_limit.recovery_budget_secs = 1; + client.stream_idle_timeout = Duration::from_secs(3); + let open = StreamOpenRequest { + policy: StreamHttpPolicy::DualWithH1Fallback, + open_timeout: Duration::from_secs(1), + idle_timeout: client.stream_idle_timeout, + }; + let body = serde_json::json!({"fixture": true}); + let open_stream = |policy| { + let transport = + client_for_policy(&client.http_client, client.http1_fallback_client(), policy); + client.open_stream_json_with_retry(transport, &url, &body, policy) + }; + let persisted = state + .path() + .join(format!("{}.json", client.shared_quota.key)); + let peer = async { + tokio::time::timeout(Duration::from_millis(700), async { + while !persisted.is_file() { + tokio::time::sleep(Duration::from_millis(10)).await; + } + open_stream(StreamHttpPolicy::DualWithH1Fallback).await + }) + .await + }; + let (result, peer_result) = tokio::join!( + tokio::time::timeout( + Duration::from_secs(2), + open_sse_response(&open, open_stream) + ), + peer, + ); + fixture.abort(); + + let error = result.expect("known status must finish").unwrap_err(); + assert!( + matches!(error.downcast_ref::(), Some(LlmError::QuotaWindow { retry_after: Some(delay), .. }) if *delay == Duration::from_secs(7200)) + ); + let peer_error = peer_result + .expect("persist and share the header minimum before reading the stalled body") + .unwrap_err(); + assert!( + matches!(peer_error.downcast_ref::(), Some(LlmError::QuotaWindow { retry_after: Some(delay), .. }) if *delay > Duration::from_secs(7195)) + ); + assert_eq!(attempts.load(Ordering::Relaxed), 1); +} + +#[tokio::test] +async fn forkguard_rate_limit_incomplete_native_success_body_never_reopens() { + use super::stream_entry::{ + StreamHttpPolicy, StreamOpenRequest, client_for_policy, open_sse_response, + }; + let (url, attempts, fixture) = hold_http_response( + b"HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: 128\r\n\r\n{", + ) + .await; + let mut client = DeepSeekClient::new(&Config { + provider: Some("minimax".into()), + providers: Some(crate::config::ProvidersConfig { + minimax: crate::config::ProviderConfig { + api_key: Some(format!("owned-fixture-{}", uuid::Uuid::new_v4())), + base_url: Some(url.clone()), + ..Default::default() + }, + ..Default::default() + }), + ..Default::default() + }) + .unwrap(); + client.stream_idle_timeout = Duration::from_millis(600); + let open = StreamOpenRequest { + policy: StreamHttpPolicy::DualWithH1Fallback, + open_timeout: Duration::from_millis(350), + idle_timeout: client.stream_idle_timeout, + }; + let body = serde_json::json!({"fixture": true}); + let result = tokio::time::timeout( + Duration::from_secs(3), + open_sse_response(&open, |policy| { + let transport = + client_for_policy(&client.http_client, client.http1_fallback_client(), policy); + client.open_stream_json_with_retry(transport, &url, &body, policy) + }), + ) + .await; + fixture.abort(); + + let error = result + .expect("native response inspection is bounded") + .unwrap_err(); + assert!( + matches!(error.downcast_ref::(), Some(LlmError::ParseError(message)) if message.contains("provider JSON envelope")) + ); + assert_eq!(attempts.load(Ordering::Relaxed), 1); +} + +#[tokio::test] +async fn forkguard_rate_limit_blocked_receipt_after_headers_never_reopens() { + use super::stream_entry::{ + StreamHttpPolicy, StreamOpenRequest, client_for_policy, open_sse_response, + }; + use wiremock::matchers::method; + use wiremock::{Mock, MockServer, ResponseTemplate}; + + let server = MockServer::start().await; + Mock::given(method("POST")) + .respond_with(ResponseTemplate::new(200).set_body_string("accepted")) + .expect(1) + .mount(&server) + .await; + let url = server.uri(); + let client = DeepSeekClient::new(&Config { + api_key: Some(format!("owned-fixture-{}", uuid::Uuid::new_v4())), + base_url: Some(url.clone()), + ..Default::default() + }) + .unwrap(); + let saw_recovered = Arc::new(std::sync::atomic::AtomicBool::new(false)); + let observed = saw_recovered.clone(); + let observer: crate::llm_client::rate_limit::RecoveryObserver = Arc::new(move |receipt| { + let observed = observed.clone(); + Box::pin(async move { + if receipt.state == "recovered" { + observed.store(true, Ordering::Relaxed); + std::future::pending::<()>().await; + } + }) + }); + let open = StreamOpenRequest { + policy: StreamHttpPolicy::DualWithH1Fallback, + open_timeout: Duration::from_millis(500), + idle_timeout: client.stream_idle_timeout, + }; + let body = serde_json::json!({"fixture": true}); + client.shared_quota.pause(Duration::from_millis(250)); + let result = crate::llm_client::rate_limit::with_observer( + observer, + open_sse_response(&open, |policy| { + let transport = + client_for_policy(&client.http_client, client.http1_fallback_client(), policy); + client.open_stream_json_with_retry(transport, &url, &body, policy) + }), + ) + .await; + + assert!(saw_recovered.load(Ordering::Relaxed)); + assert!( + matches!(result.unwrap_err().downcast_ref::(), Some(LlmError::ParseError(message)) if message.contains("processing the response headers")) + ); + server.verify().await; +} + +#[tokio::test] +async fn forkguard_rate_limit_blocked_terminal_receipt_preserves_quota_result() { + use super::stream_entry::{ + StreamHttpPolicy, StreamOpenRequest, client_for_policy, open_sse_response, + }; + use wiremock::matchers::method; + use wiremock::{Mock, MockServer, ResponseTemplate}; + + for wire_request in [true, false] { + let server = MockServer::start().await; + Mock::given(method("POST")) + .respond_with(ResponseTemplate::new(429).insert_header("Retry-After", "7200")) + .expect(u64::from(wire_request)) + .mount(&server) + .await; + let url = server.uri(); + let client = DeepSeekClient::new(&Config { + api_key: Some(format!("owned-fixture-{}", uuid::Uuid::new_v4())), + base_url: Some(url.clone()), + ..Default::default() + }) + .unwrap(); + if !wire_request { + client.shared_quota.pause(Duration::from_secs(7200)); + } + let saw_terminal = Arc::new(std::sync::atomic::AtomicBool::new(false)); + let observed = saw_terminal.clone(); + let observer: crate::llm_client::rate_limit::RecoveryObserver = Arc::new(move |receipt| { + let observed = observed.clone(); + Box::pin(async move { + if receipt.state == "interrupted" { + observed.store(true, Ordering::Relaxed); + std::future::pending::<()>().await; + } + }) + }); + let open = StreamOpenRequest { + policy: StreamHttpPolicy::DualWithH1Fallback, + open_timeout: Duration::from_millis(500), + idle_timeout: client.stream_idle_timeout, + }; + let body = serde_json::json!({"fixture": true}); + let result = crate::llm_client::rate_limit::with_observer( + observer, + open_sse_response(&open, |policy| { + let transport = + client_for_policy(&client.http_client, client.http1_fallback_client(), policy); + client.open_stream_json_with_retry(transport, &url, &body, policy) + }), + ) + .await; + + assert!(saw_terminal.load(Ordering::Relaxed)); + assert!( + matches!(result.unwrap_err().downcast_ref::(), Some(LlmError::QuotaWindow { retry_after: Some(delay), .. }) if *delay > Duration::from_secs(7195)) + ); + server.verify().await; + } +} + +#[tokio::test] +async fn forkguard_rate_limit_blocked_persistence_respects_recovery_deadline() { + use super::stream_entry::{ + StreamHttpPolicy, StreamOpenRequest, client_for_policy, open_sse_response, + }; + use wiremock::matchers::method; + use wiremock::{Mock, MockServer, ResponseTemplate}; + + let native = ResponseTemplate::new(200).set_body_json( + serde_json::json!({"base_resp":{"status_code":1002,"status_msg":"rate limit"}}), + ); + for (response, budget, open_timeout) in [ + (ResponseTemplate::new(429), 1, Duration::from_secs(4)), + (ResponseTemplate::new(429), 600, Duration::from_millis(500)), + (native, 1, Duration::from_secs(4)), + ] { + let server = MockServer::start().await; + Mock::given(method("POST")) + .respond_with(response.insert_header("Retry-After", "7200")) + .expect(1) + .mount(&server) + .await; + let url = server.uri(); + let state = tempfile::tempdir().unwrap(); + let mut client = DeepSeekClient::new(&Config { + provider: Some("minimax".into()), + providers: Some(crate::config::ProvidersConfig { + minimax: crate::config::ProviderConfig { + api_key: Some(format!("owned-fixture-{}", uuid::Uuid::new_v4())), + base_url: Some(url.clone()), + ..Default::default() + }, + ..Default::default() + }), + rate_limit_recovery_dir: Some(state.path().to_owned()), + ..Default::default() + }) + .unwrap(); + client.retry.rate_limit.recovery_budget_secs = budget; + let _blocked = client.shared_quota.persistence.lock().await; + let open = StreamOpenRequest { + policy: StreamHttpPolicy::DualWithH1Fallback, + open_timeout, + idle_timeout: client.stream_idle_timeout, + }; + let body = serde_json::json!({"fixture": true}); + let result = tokio::time::timeout( + Duration::from_secs(2), + open_sse_response(&open, |policy| { + let transport = + client_for_policy(&client.http_client, client.http1_fallback_client(), policy); + client.open_stream_json_with_retry(transport, &url, &body, policy) + }), + ) + .await + .expect("either deadline must bound a blocked quota store"); + let error = result.unwrap_err(); + assert!( + matches!(error.downcast_ref::(), Some(LlmError::QuotaWindow { retry_after: Some(delay), .. }) if *delay == Duration::from_secs(7200)), + "{error:#}" + ); + assert!(client.shared_quota.remaining().unwrap() > Duration::from_secs(7195)); + server.verify().await; + } +} + +#[tokio::test] +async fn forkguard_rate_limit_definitive_stream_status_does_not_probe_or_fallback() { + use super::stream_entry::{ + StreamHttpPolicy, StreamOpenRequest, client_for_policy, open_sse_response, + }; + use wiremock::matchers::method; + use wiremock::{Mock, MockServer, ResponseTemplate}; + + let server = MockServer::start().await; + Mock::given(method("POST")) + .respond_with(ResponseTemplate::new(401)) + .expect(2) + .mount(&server) + .await; + Mock::given(method("GET")) + .respond_with(ResponseTemplate::new(500).set_delay(Duration::from_secs(2))) + .expect(0) + .mount(&server) + .await; + let url = server.uri(); + let client = DeepSeekClient::new(&Config { + api_key: Some(format!("owned-fixture-{}", uuid::Uuid::new_v4())), + base_url: Some(url.clone()), + ..Default::default() + }) + .unwrap(); + let open = StreamOpenRequest { + policy: StreamHttpPolicy::DualWithH1Fallback, + open_timeout: Duration::from_secs(1), + idle_timeout: client.stream_idle_timeout, + }; + let body = serde_json::json!({"fixture": true}); + for _ in 0..2 { + let error = open_sse_response(&open, |policy| { + let transport = + client_for_policy(&client.http_client, client.http1_fallback_client(), policy); + client.open_stream_json_with_retry(transport, &url, &body, policy) + }) + .await + .unwrap_err(); + assert!(matches!( + error.downcast_ref::(), + Some(LlmError::AuthenticationError(_)) + )); + } + server.verify().await; +} + +#[tokio::test] +async fn forkguard_rate_limit_http1_does_not_retry_ordinary_server_errors() { + use super::stream_entry::StreamHttpPolicy; + use wiremock::matchers::method; + use wiremock::{Mock, MockServer, ResponseTemplate}; + + let server = MockServer::start().await; + Mock::given(method("POST")) + .respond_with(ResponseTemplate::new(503)) + .expect(1) + .mount(&server) + .await; + let client = DeepSeekClient::new(&Config { + api_key: Some(format!("owned-fixture-{}", uuid::Uuid::new_v4())), + base_url: Some(server.uri()), + ..Default::default() + }) + .unwrap(); + let result = client + .open_stream_json_with_retry( + client.http1_fallback_client(), + &server.uri(), + &serde_json::json!({"fixture": true}), + StreamHttpPolicy::Http1Only, + ) + .await; + assert!(result.is_err()); + assert_eq!(server.received_requests().await.unwrap().len(), 1); +} + +#[tokio::test(start_paused = true)] +async fn forkguard_rate_limit_recovery_expires_while_queued_for_a_provider_slot() { + let state = tempfile::tempdir().unwrap(); + let mut client = DeepSeekClient::new(&Config { + api_key: Some(format!("owned-fixture-{}", uuid::Uuid::new_v4())), + base_url: Some("http://127.0.0.1:1".into()), + ..Default::default() + }) + .unwrap(); + let quota = shared_quota( + uuid::Uuid::new_v4().to_string(), + Some(1), + Some(state.path().to_owned()), + ) + .unwrap(); + client.shared_quota = quota.clone(); + client.request_concurrency = quota.concurrency.clone(); + client.rate_limiter = quota.limiter.clone(); + client.retry.rate_limit.recovery_budget_secs = 2; + quota.initialize().await.unwrap(); + let blocker = ProviderRequestPermit { + lease: std::sync::Mutex::new(None), + limiter: quota.concurrency.as_ref().unwrap().clone(), + }; + blocker.resume().await; + + let operation_id = uuid::Uuid::new_v4(); + let call = client.for_operation(operation_id); + let _lease = call.acquire_provider_request_permit().await; + let recovery = call.recovery_for_operation(operation_id).unwrap(); + recovery.lock().unwrap().started = Some(tokio::time::Instant::now()); + let result = tokio::time::timeout( + Duration::from_secs(3), + call.send_stream_open_with_retry(|| call.http_client.get("http://127.0.0.1:1")), + ) + .await + .expect("recovery deadline must stop a queued stream open"); + assert!(matches!( + result.unwrap_err().downcast_ref::(), + Some(LlmError::QuotaWindow { + retry_after: None, + .. + }) + )); + assert_eq!(recovery.lock().unwrap().wire_attempts, Some(0)); + assert_eq!(quota.concurrency.as_ref().unwrap().active(), 1); +} + +#[tokio::test] +async fn forkguard_rate_limit_rechecks_cooldown_after_admission_queues() { + use wiremock::matchers::method; + use wiremock::{Mock, MockServer, Request, ResponseTemplate}; + for queued_on_tokens in [false, true] { + let server = MockServer::start().await; + let wire_time = Arc::new(std::sync::Mutex::new(None)); + let receipt = wire_time.clone(); + Mock::given(method("POST")) + .respond_with(move |_: &Request| { + *receipt.lock().unwrap() = Some(Instant::now()); + ResponseTemplate::new(200).set_body_string("ok") + }) + .mount(&server) + .await; + let mut client = DeepSeekClient::new(&Config { + api_key: Some("owned-fixture".into()), + base_url: Some(server.uri()), + ..Default::default() + }) + .unwrap(); + let quota = shared_quota(uuid::Uuid::new_v4().to_string(), Some(1), None).unwrap(); + client.shared_quota = quota.clone(); + client.request_concurrency = quota.concurrency.clone(); + client.rate_limiter = quota.limiter.clone(); + { + let mut bucket = quota.limiter.lock().await; + bucket.enabled = queued_on_tokens; + bucket.tokens = 0.0; + bucket.refill_per_sec = 10.0; + bucket.last_refill = Instant::now(); + } + let blocker = ProviderRequestPermit { + lease: std::sync::Mutex::new(None), + limiter: quota.concurrency.as_ref().unwrap().clone(), + }; + if !queued_on_tokens { + blocker.resume().await; + } + let url = server.uri(); + let receipts = Arc::new(std::sync::Mutex::new(Vec::new())); + let seen_receipts = receipts.clone(); + let observer: crate::llm_client::rate_limit::RecoveryObserver = Arc::new(move |receipt| { + seen_receipts.lock().unwrap().push(receipt); + Box::pin(async {}) + }); + let call = client.for_operation(uuid::Uuid::new_v4()); + let _lease = call.acquire_provider_request_permit().await; + let body = serde_json::json!({"fixture": true}); + let sender = crate::llm_client::rate_limit::with_operation( + Default::default(), + crate::llm_client::rate_limit::with_observer( + observer, + call.send_json_with_retry(&url, &body), + ), + ); + tokio::pin!(sender); + // Poll until the actual send path has reached the held slot or token + // sleep before publishing the peer's new cooldown. No scheduler race. + assert!(futures_util::poll!(sender.as_mut()).is_pending()); + assert!(wire_time.lock().unwrap().is_none()); + if queued_on_tokens { + assert!(quota.limiter.lock().await.tokens < 0.0); + } else { + assert_eq!(quota.concurrency.as_ref().unwrap().active(), 1); + } + let not_before = Instant::now() + Duration::from_millis(250); + quota.pause(Duration::from_millis(250)); + blocker.suspend(); + tokio::time::timeout(Duration::from_secs(10), sender) + .await + .unwrap() + .unwrap(); + assert!( + wire_time.lock().unwrap().unwrap() >= not_before, + "a queued request must obey a peer's new cooldown" + ); + let receipts = receipts.lock().unwrap(); + assert!( + receipts + .iter() + .any(|r| r.state == "waiting" && r.attempt == 0) + ); + assert!( + receipts + .iter() + .any(|r| r.state == "recovered" && r.attempt == 1) + ); + } +} + +#[tokio::test] +async fn forkguard_rate_limit_scope_admission_and_leases() { + let key = uuid::Uuid::new_v4().to_string(); + let first = shared_quota(key.clone(), Some(4), None).unwrap(); + let second = shared_quota(key, Some(2), None).unwrap(); + let other = shared_quota(uuid::Uuid::new_v4().to_string(), Some(4), None).unwrap(); + assert!(Arc::ptr_eq(&first, &second)); + assert!(!Arc::ptr_eq(&first, &other)); + first.pause(Duration::from_secs(3600)); + assert!(second.remaining().unwrap() > Duration::from_secs(3598)); + assert!(other.remaining().is_none()); + let limiter = first.concurrency.as_ref().unwrap(); + assert_eq!(limiter.limit(), 2); + let lease = ProviderRequestPermit { + lease: std::sync::Mutex::new(None), + limiter: limiter.clone(), + }; + lease.resume().await; + assert_eq!(limiter.active(), 1); + lease.suspend(); + assert_eq!(limiter.active(), 0); + first.refused().await; + assert_eq!(limiter.gate.capacity(), 1); + let policy = crate::config::RateLimitRetryPolicy { + max_attempts: 3, + ..Default::default() + }; + assert!(first.admit_recovery(&policy).is_none()); + assert!(second.admit_recovery(&policy).is_none()); + assert!(first.admit_recovery(&policy).is_some()); + assert!(other.admit_recovery(&policy).is_none()); + for _ in 0..3 { + first.recovered().await; + } + assert_eq!(limiter.gate.capacity(), 2); + lease.resume().await; + assert_eq!(limiter.active(), 1); + drop(lease); + assert_eq!(limiter.active(), 0); +} + +#[tokio::test] +async fn forkguard_rate_limit_isolated_classifier_normalizes_native_json_refusal() { + use wiremock::matchers::method; + use wiremock::{Mock, MockServer, ResponseTemplate}; + let server = MockServer::start().await; + Mock::given(method("POST")) + .respond_with(ResponseTemplate::new(200).set_body_json( + serde_json::json!({"base_resp":{"status_code":1002,"status_msg":"owned-fixture"}}), + )) + .expect(1) + .mount(&server) + .await; + let mut client = DeepSeekClient::new(&Config { + provider: Some("minimax".into()), + providers: Some(crate::config::ProvidersConfig { + minimax: crate::config::ProviderConfig { + api_key: Some("owned-fixture".into()), + base_url: Some(server.uri()), + ..Default::default() + }, + ..Default::default() + }), + ..Default::default() + }) + .unwrap(); + client.retry.enabled = false; + client.isolated_request_state = true; + let error = client + .send_with_retry(|| client.http_client.post(server.uri())) + .await + .unwrap_err(); + assert!(matches!( + error.downcast_ref::(), + Some(LlmError::RateLimited { .. }) + )); + assert!(!error.to_string().contains("owned-fixture")); +} + +// The environment lock must cover construction and wire dispatch: this is an +// exact-credential boundary under source-marked process-wide CLI overrides. +#[allow(clippy::await_holding_lock)] +#[tokio::test(flavor = "current_thread")] +async fn forkguard_rate_limit_auxiliary_route_keeps_its_explicit_credential() { + use wiremock::matchers::{header, method}; + use wiremock::{Mock, MockServer, ResponseTemplate}; + let _env = crate::test_support::lock_test_env(); + let _key = + crate::test_support::EnvVarGuard::set("CODEWHALE_CLI_API_KEY", "main-cli-fixture-key"); + let _source = crate::test_support::EnvVarGuard::set("CODEWHALE_CLI_API_KEY_SOURCE", "cli"); + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(header("authorization", "Bearer independent-vision-fixture")) + .respond_with( + ResponseTemplate::new(200).set_body_json(serde_json::json!({"fixture": true})), + ) + .expect(1) + .mount(&server) + .await; + let route = DeepSeekClient::auxiliary_json_route( + None, + &server.uri(), + "independent-vision-fixture", + "vision-model", + ) + .unwrap(); + assert_eq!(route.api_key, "independent-vision-fixture"); + route + .request_auxiliary_json(&server.uri(), &serde_json::json!({}), uuid::Uuid::new_v4()) + .await + .unwrap(); +} + +#[tokio::test] +async fn forkguard_rate_limit_auxiliary_apis_publish_distinct_runtime_identities() { + use wiremock::matchers::method; + use wiremock::{Mock, MockServer, Request, ResponseTemplate}; + let server = MockServer::start().await; + let calls = Arc::new(AtomicUsize::new(0)); + let seen_calls = calls.clone(); + Mock::given(method("POST")).respond_with(move |_: &Request| { + if seen_calls.fetch_add(1, Ordering::SeqCst) % 2 == 0 { + ResponseTemplate::new(429).insert_header("retry-after", "0") + } else { + ResponseTemplate::new(200).set_body_json(serde_json::json!({"choices":[{"message":{"content":"translated"},"text":"completed"}]})) + } + }).mount(&server).await; + let mut policy = crate::config::RateLimitRetryPolicy::default(); + policy.initial_delay_secs = 1; + policy.max_delay_secs = 1; + let client = DeepSeekClient::new(&Config { + api_key: Some("owned-auxiliary-key".into()), + base_url: Some(server.uri()), + retry: Some(crate::config::RetryConfig { + rate_limit: policy, + enabled: None, + max_retries: None, + initial_delay: None, + max_delay: None, + exponential_base: None, + }), + ..Default::default() + }) + .unwrap(); + let receipts = Arc::new(std::sync::Mutex::new(Vec::new())); + let seen = receipts.clone(); + let observer: crate::llm_client::rate_limit::RecoveryObserver = Arc::new(move |receipt| { + let seen = seen.clone(); + Box::pin(async move { + seen.lock().unwrap().push(receipt); + }) + }); + crate::llm_client::rate_limit::with_observer(observer, async { + assert_eq!( + client + .translate("fixture", "deepseek-chat", "English") + .await + .unwrap(), + "translated" + ); + assert_eq!( + client + .fim_completion("deepseek-chat", "fixture", "suffix", 16) + .await + .unwrap(), + "completed" + ); + }) + .await; + let receipts = receipts.lock().unwrap(); + let ids: std::collections::BTreeSet<_> = receipts + .iter() + .map(|r| r.operation_id.clone().unwrap()) + .collect(); + assert_eq!(ids.len(), 2); + for id in ids { + let states: Vec<_> = receipts + .iter() + .filter(|r| r.operation_id.as_deref() == Some(id.as_str())) + .map(|r| (r.state.as_str(), r.attempt)) + .collect(); + assert_eq!(states, [("waiting", 1), ("recovered", 2)]); + } + assert_eq!(calls.load(Ordering::SeqCst), 4); +} + +// Simulate a process boundary without bypassing the production file reader. +fn forget_quota_process_state_for_restart(key: &str) { + let mut registry = quota_registry().lock().unwrap(); + assert_eq!( + Arc::strong_count(registry.get(key).unwrap()), + 1, + "no live client may be forgotten" + ); + registry.remove(key); +} + +#[tokio::test] +async fn forkguard_rate_limit_episode_survives_last_auxiliary_client_drop() { + let key = uuid::Uuid::new_v4().to_string(); + let policy = crate::config::RateLimitRetryPolicy { + max_attempts: 2, + ..Default::default() + }; + let first = shared_quota(key.clone(), Some(16), None).unwrap(); + first.refused().await; + assert!(first.admit_recovery(&policy).is_none()); + assert!(first.admit_recovery(&policy).is_some()); + drop(first); + let next = shared_quota(key, Some(16), None).unwrap(); + assert!( + next.admit_recovery(&policy).is_some(), + "client turnover cannot reset the process recovery episode" + ); + assert_eq!(next.concurrency.as_ref().unwrap().gate.capacity(), 8); + next.recovered().await; + assert!( + next.admit_recovery(&policy).is_none(), + "a real success ends the shared refusal episode" + ); + let isolated = shared_quota(uuid::Uuid::new_v4().to_string(), Some(16), None).unwrap(); + assert_eq!(isolated.concurrency.as_ref().unwrap().gate.capacity(), 16); +} + +#[tokio::test] +async fn forkguard_rate_limit_isolated_unknown_quota_uses_route_classification() { + use wiremock::matchers::method; + use wiremock::{Mock, MockServer, ResponseTemplate}; + let server = MockServer::start().await; + Mock::given(method("POST")).respond_with(ResponseTemplate::new(429).set_body_json( + serde_json::json!({"error":{"code":null,"type":"insufficient_quota","message":"TPM exceeded"}}) + )).expect(1).mount(&server).await; + let mut client = DeepSeekClient::new(&Config { + provider: Some("custom".into()), + api_key: Some(uuid::Uuid::new_v4().to_string()), + base_url: Some(server.uri()), + ..Default::default() + }) + .unwrap(); + client.retry.enabled = false; + client.isolated_request_state = true; + let error = client + .send_with_retry(|| client.http_client.post(server.uri())) + .await + .unwrap_err(); + assert!(matches!( + error.downcast_ref::(), + Some(LlmError::RateLimited { .. }) + )); +} diff --git a/crates/tui/src/client/responses/tests.rs b/crates/tui/src/client/responses/tests.rs index f05d90d8aa..9104ff627a 100644 --- a/crates/tui/src/client/responses/tests.rs +++ b/crates/tui/src/client/responses/tests.rs @@ -108,6 +108,7 @@ fn test_codex_config(server: &MockServer) -> Config { Config { provider: Some("openai-codex".to_string()), retry: Some(RetryConfig { + rate_limit: crate::config::RateLimitRetryPolicy::default(), enabled: Some(true), max_retries: Some(1), initial_delay: Some(0.0), diff --git a/crates/tui/src/client/stream_entry.rs b/crates/tui/src/client/stream_entry.rs index 125c101540..851020599a 100644 --- a/crates/tui/src/client/stream_entry.rs +++ b/crates/tui/src/client/stream_entry.rs @@ -9,12 +9,16 @@ //! `docs/notes/post-0.9.1-thin-tui-and-stream.md`. use std::future::Future; +use std::sync::{Arc, Mutex}; use std::time::Duration; use anyhow::Result; use reqwest::Client; -use crate::llm_client::LlmError; +use crate::llm_client::{ + LlmError, + rate_limit::{STREAM_OPEN_OUTCOME, StreamOpenOutcome}, +}; /// Default bounded wait for SSE response headers. Intentionally shorter than /// the per-chunk idle timeout: it covers connection setup and upstream header @@ -148,6 +152,30 @@ fn h1_fallback_error(err: anyhow::Error) -> anyhow::Error { ))) } +#[derive(Debug)] +pub(super) struct HeaderDeadline; + +impl std::fmt::Display for HeaderDeadline { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str("stream response header deadline elapsed") + } +} + +impl std::error::Error for HeaderDeadline {} + +fn settled_open_deadline(response: StreamOpenOutcome) -> anyhow::Error { + anyhow::Error::new(match response { + StreamOpenOutcome::RecoverableQuota(retry_after) => LlmError::QuotaWindow { + message: "Stream open could not deliver the model quota result before its deadline" + .into(), + retry_after, + }, + _ => LlmError::ParseError( + "Stream open did not finish processing the response headers".into(), + ), + }) +} + /// Open an SSE response through the shared transport policy. /// /// `attempt` builds and sends one wire-specific request on the client @@ -168,22 +196,36 @@ where F: Fn(StreamHttpPolicy) -> Fut, Fut: Future>, { - let fallback_reason = match tokio::time::timeout( - open_req.open_timeout, - attempt(open_req.policy), - ) - .await - { + let response = Arc::new(Mutex::new(StreamOpenOutcome::Pending)); + let wire_response = || { + *response + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + }; + let outcome_known = || !matches!(wire_response(), StreamOpenOutcome::Pending); + let run = |policy| { + *response + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) = StreamOpenOutcome::Pending; + crate::llm_client::rate_limit::active_timeout( + open_req.open_timeout, + STREAM_OPEN_OUTCOME.scope(response.clone(), attempt(policy)), + ) + }; + let fallback_reason = match run(open_req.policy).await { Ok(Ok(response)) => return Ok(response), Ok(Err(err)) => { - if !should_retry_error_with_h1(open_req.policy, &err) { + if outcome_known() || !should_retry_error_with_h1(open_req.policy, &err) { return Err(err); } "transport error before response headers" } Err(_elapsed) => { + if outcome_known() { + return Err(settled_open_deadline(wire_response())); + } if open_req.policy == StreamHttpPolicy::Http1Only { - return Err(anyhow::Error::new(LlmError::NetworkError(format!( + return Err(anyhow::Error::new(HeaderDeadline).context(LlmError::NetworkError(format!( "SSE stream request did not receive response headers after {}s. \ `codewhale doctor` can still pass when non-streaming requests work; \ on Windows or proxy networks, try `CODEWHALE_FORCE_HTTP1=1` and rerun `codewhale`.", @@ -200,20 +242,24 @@ where crate::logging::warn(format!( "SSE stream {fallback_reason}; retrying once with HTTP/1.1" )); - match tokio::time::timeout(h1_req.open_timeout, attempt(h1_req.policy)).await { + match run(h1_req.policy).await { Ok(Ok(response)) => Ok(response), + Ok(Err(err)) if outcome_known() => Err(err), Ok(Err(err)) => Err(h1_fallback_error(err)), + Err(_elapsed) if outcome_known() => Err(settled_open_deadline(wire_response())), // Typed, not a bare string: a header stall is a transport // failure, and `LlmError::NetworkError` is what the shared // retry layer recognizes as retryable. As an untyped anyhow // error this killed the whole turn outright. - Err(_elapsed) => Err(anyhow::Error::new(LlmError::NetworkError(format!( - "SSE stream request did not receive response headers after {}s \ + Err(_elapsed) => Err( + anyhow::Error::new(HeaderDeadline).context(LlmError::NetworkError(format!( + "SSE stream request did not receive response headers after {}s \ (HTTP/2 and HTTP/1.1). `codewhale doctor` can still pass when \ non-streaming requests work; try `CODEWHALE_FORCE_HTTP1=1` and \ rerun `codewhale`.", - open_req.open_timeout.as_secs() - )))), + open_req.open_timeout.as_secs() + ))), + ), } } @@ -379,6 +425,7 @@ mod tests { assert_eq!(attempts.load(Ordering::SeqCst), 1, "no retry when pinned"); let text = err.to_string(); assert!(text.contains("did not receive response headers"), "{text}"); + assert!(err.is::()); // The whole point of the typed error: a header stall must reach the // shared retry layer as retryable. As an untyped anyhow error it // killed the turn outright, so a long root run lost all its work while diff --git a/crates/tui/src/compaction.rs b/crates/tui/src/compaction.rs index 72044cdf41..e9c89e613b 100644 --- a/crates/tui/src/compaction.rs +++ b/crates/tui/src/compaction.rs @@ -1079,7 +1079,7 @@ pub fn report_compaction_failure( "provider plan quota exhausted — switch provider/model or renew the provider plan" .to_string() } - Some(crate::llm_client::LlmError::RateLimited { .. }) => { + Some(crate::llm_client::LlmError::RateLimited { .. } | crate::llm_client::LlmError::QuotaWindow { .. }) => { "provider rate limit blocked compaction — retry after the limit resets or switch provider/model" .to_string() } diff --git a/crates/tui/src/config.rs b/crates/tui/src/config.rs index d78aea3fc6..2b38bc5954 100644 --- a/crates/tui/src/config.rs +++ b/crates/tui/src/config.rs @@ -1725,9 +1725,73 @@ pub fn model_completion_names_for_provider(provider: ApiProvider) -> Vec<&'stati // === Types === +/// Independent recovery limits for HTTP 429; other transport errors keep their policy. +#[derive(Debug, Clone, serde::Serialize, Deserialize)] +#[serde(try_from = "RawRateLimitRetryPolicy")] +pub struct RateLimitRetryPolicy { + pub max_attempts: u32, + pub recovery_budget_secs: u64, + pub initial_delay_secs: u64, + pub max_delay_secs: u64, +} + +impl Default for RateLimitRetryPolicy { + fn default() -> Self { + Self { + max_attempts: 20, + recovery_budget_secs: 600, + initial_delay_secs: 5, + max_delay_secs: 60, + } + } +} + +#[derive(Deserialize)] +#[serde(default)] +struct RawRateLimitRetryPolicy { + max_attempts: u32, + recovery_budget_secs: u64, + initial_delay_secs: u64, + max_delay_secs: u64, +} + +impl Default for RawRateLimitRetryPolicy { + fn default() -> Self { + let p = RateLimitRetryPolicy::default(); + Self { + max_attempts: p.max_attempts, + recovery_budget_secs: p.recovery_budget_secs, + initial_delay_secs: p.initial_delay_secs, + max_delay_secs: p.max_delay_secs, + } + } +} + +impl TryFrom for RateLimitRetryPolicy { + type Error = String; + fn try_from(p: RawRateLimitRetryPolicy) -> Result { + if !(1..=1000).contains(&p.max_attempts) + || !(1..=86400).contains(&p.recovery_budget_secs) + || p.initial_delay_secs == 0 + || p.initial_delay_secs > p.max_delay_secs + || p.max_delay_secs > 86400 + { + return Err("invalid retry.rate_limit policy: attempts 1..1000, budget 1..86400s, 1 <= initial <= max <= 86400s".into()); + } + Ok(Self { + max_attempts: p.max_attempts, + recovery_budget_secs: p.recovery_budget_secs, + initial_delay_secs: p.initial_delay_secs, + max_delay_secs: p.max_delay_secs, + }) + } +} + /// Raw retry configuration loaded from config files. #[derive(Debug, Clone, Deserialize)] pub struct RetryConfig { + #[serde(default)] + pub rate_limit: RateLimitRetryPolicy, pub enabled: Option, pub max_retries: Option, pub initial_delay: Option, @@ -2618,6 +2682,7 @@ impl HeaderItem { /// Resolved retry policy with defaults applied. #[derive(Debug, Clone)] pub struct RetryPolicy { + pub rate_limit: RateLimitRetryPolicy, pub enabled: bool, pub max_retries: u32, pub initial_delay: f64, @@ -2964,6 +3029,9 @@ impl TranscriptConfig { /// Resolved CLI configuration, including defaults and environment overrides. #[derive(Debug, Clone, Default, Deserialize)] pub struct Config { + /// Host-owned location for content-free provider reset receipts. + #[serde(default)] + pub rate_limit_recovery_dir: Option, /// Single-token inputs that cancel the active turn before dispatch. #[serde(default)] pub stop_words: Option>, @@ -7874,6 +7942,7 @@ impl Config { #[must_use] pub fn retry_policy(&self) -> RetryPolicy { let defaults = RetryPolicy { + rate_limit: RateLimitRetryPolicy::default(), enabled: true, max_retries: 3, initial_delay: 1.0, @@ -7886,6 +7955,7 @@ impl Config { }; RetryPolicy { + rate_limit: cfg.rate_limit.clone(), enabled: cfg.enabled.unwrap_or(defaults.enabled), max_retries: cfg.max_retries.unwrap_or(defaults.max_retries), initial_delay: cfg.initial_delay.unwrap_or(defaults.initial_delay), @@ -10628,6 +10698,9 @@ fn merge_config(base: Config, override_cfg: Config) -> Config { .or(base.managed_config_path), requirements_path: override_cfg.requirements_path.or(base.requirements_path), max_subagents: override_cfg.max_subagents.or(base.max_subagents), + rate_limit_recovery_dir: override_cfg + .rate_limit_recovery_dir + .or(base.rate_limit_recovery_dir), retry: override_cfg.retry.or(base.retry), auto_review: override_cfg.auto_review.or(base.auto_review), tui: override_cfg.tui.or(base.tui), diff --git a/crates/tui/src/core/capacity.rs b/crates/tui/src/core/capacity.rs new file mode 100644 index 0000000000..69dd24075b --- /dev/null +++ b/crates/tui/src/core/capacity.rs @@ -0,0 +1,164 @@ +//! Cancellation-safe dynamic semaphore shared by model admission and child scheduling. + +use std::sync::Mutex; + +// === DynamicGate === + +#[derive(Debug)] +struct GateState { + capacity: usize, + outstanding: usize, +} + +/// A launch gate with runtime-adjustable capacity (see module docs). +/// +/// A thin accounting layer over [`tokio::sync::Semaphore`]: the semaphore +/// owns the FIFO wait queue and cancellation safety, this struct owns the +/// `capacity`/`outstanding` ledger that lets capacity drop below the +/// number of outstanding holders. The ledger — not the semaphore token +/// count — is the admission authority: a granted token is registered only +/// while `outstanding < capacity`; one granted while the ledger is full is +/// surplus and is destroyed (`forget`) instead of admitted. Two paths +/// create such tokens while no slot is free, both invisible to +/// `forget_permits` (which sees free tokens only): a grant that races a +/// capacity change (handed out before the shrink can forget it), and a +/// grant handed to a queued waiter that is cancelled before polling again +/// (tokio re-injects it into the pool). Surplus tokens are harmless — they +/// cannot admit — and the next admission attempt destroys them; every slot +/// freed below capacity mints a fresh token on release, so no slot is ever +/// lost. Once no surplus token remains and every in-flight admission has +/// registered, `sem.available_permits() == +/// capacity.saturating_sub(outstanding)` holds exactly. +/// +/// `acquire` returns a [`DynamicGatePermit`] whose `Drop` returns the slot. +/// Reducing capacity below `outstanding` is allowed: the surplus holders +/// finish naturally and no new permit is granted until the outstanding +/// count drops under the new capacity. +#[derive(Debug)] +pub(crate) struct DynamicGate { + sem: tokio::sync::Semaphore, + inner: Mutex, +} + +impl DynamicGate { + pub(crate) fn new(capacity: usize) -> Self { + let capacity = capacity.max(1); + Self { + sem: tokio::sync::Semaphore::new(capacity), + inner: Mutex::new(GateState { + capacity, + outstanding: 0, + }), + } + } + + pub(crate) fn capacity(&self) -> usize { + self.inner.lock().expect("launch gate poisoned").capacity + } + + /// Free semaphore permits right now — `capacity - outstanding` at + /// rest, plus any surplus tokens not yet destroyed by an admission + /// attempt (see the struct docs). Diagnostics and tests only; racy by + /// design. + #[cfg_attr(not(test), allow(dead_code))] + pub(crate) fn available_permits(&self) -> usize { + self.sem.available_permits() + } + + /// Adjust the gate capacity. Raising it adds the new free headroom to + /// the semaphore immediately (queued acquirers wake in FIFO order); + /// lowering it forgets free slots — holders above the new capacity keep + /// running, and their releases are absorbed instead of re-admitted. + /// Grants already handed to queued waiters are invisible to the shrink; + /// if their acquirers are cancelled, admission destroys the returned + /// surplus token (see the struct docs). + pub(crate) fn set_capacity(&self, capacity: usize) { + let mut inner = self.inner.lock().expect("launch gate poisoned"); + let old = inner.capacity; + inner.capacity = capacity; + if capacity < old { + // `forget_permits` caps the reduction at the available count, so + // an over-subscribed shrink (capacity < outstanding) simply + // drains free slots to zero. + self.sem.forget_permits(old - capacity); + } else { + let free = |cap: usize| cap.saturating_sub(inner.outstanding); + self.sem.add_permits(free(capacity) - free(old)); + } + } + + /// Try to acquire a permit without waiting. + pub(crate) fn try_acquire(self: &std::sync::Arc) -> Option { + let sem_permit = self.sem.try_acquire().ok()?; + self.admit_grant(sem_permit) + } + + /// Register a granted semaphore token as a held slot. The ledger, not + /// the token count, is the admission authority: a token granted while + /// the ledger is full is surplus drift — a cancelled waiter's + /// re-injected grant, or a grant that raced a shrink — and is destroyed + /// with `forget` (never dropped: a drop would re-inject it). Since + /// `DynamicGatePermit::drop` mints one fresh token per slot freed below + /// capacity, destroying a token never loses a slot. + fn admit_grant( + self: &std::sync::Arc, + sem_permit: tokio::sync::SemaphorePermit<'_>, + ) -> Option { + let mut inner = self.inner.lock().expect("launch gate poisoned"); + if inner.outstanding < inner.capacity { + inner.outstanding += 1; + // Slot ownership moves from the semaphore permit into + // `DynamicGatePermit::drop`. + sem_permit.forget(); + return Some(DynamicGatePermit { + gate: std::sync::Arc::clone(self), + }); + } + sem_permit.forget(); + None + } + + /// Acquire a permit, waiting until capacity is available. The semaphore + /// queue is FIFO and cancel safe: a cancelled future dequeues itself and + /// never swallows a slot or loses a wakeup. A grant that turns out to be + /// surplus (see [`Self::admit_grant`]) is destroyed and the wait + /// continues for the release that mints a fresh one; each surplus round + /// consumes exactly one pool token, so the loop cannot spin. + pub(crate) async fn acquire(self: &std::sync::Arc) -> DynamicGatePermit { + loop { + // The semaphore is never closed, so `acquire` cannot fail. + let sem_permit = self + .sem + .acquire() + .await + .expect("launch gate semaphore closed"); + if let Some(permit) = self.admit_grant(sem_permit) { + return permit; + } + } + } +} + +/// One held launch slot. Released on drop. +pub(crate) struct DynamicGatePermit { + gate: std::sync::Arc, +} + +impl std::fmt::Debug for DynamicGatePermit { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("DynamicGatePermit").finish() + } +} + +impl Drop for DynamicGatePermit { + fn drop(&mut self) { + let mut inner = self.gate.inner.lock().expect("launch gate poisoned"); + inner.outstanding = inner.outstanding.saturating_sub(1); + if inner.outstanding < inner.capacity { + // Return the slot; a queued acquirer wakes in FIFO order. + self.gate.sem.add_permits(1); + } + // Otherwise the capacity was shrunk below the outstanding count: + // the release is absorbed so the surplus drain does not over-admit. + } +} diff --git a/crates/tui/src/core/engine.rs b/crates/tui/src/core/engine.rs index ac4d014cf6..4ba53e1f77 100644 --- a/crates/tui/src/core/engine.rs +++ b/crates/tui/src/core/engine.rs @@ -1309,6 +1309,7 @@ pub struct Engine { /// delivery. delivered_subagent_completion_ids: HashSet, cancel_token: CancellationToken, + model_recovery_interrupted: bool, shared_cancel_token: SharedCancelToken, /// Latched reason for the current cancellation, mirrored to /// `EngineHandle::cancel_reason`. Read by `approval.rs` when @@ -2537,6 +2538,7 @@ impl Engine { rx_subagent_completion, delivered_subagent_completion_ids: HashSet::new(), cancel_token: cancel_token.clone(), + model_recovery_interrupted: false, shared_cancel_token: shared_cancel_token.clone(), cancel_reason: cancel_reason.clone(), tool_exec_lock, @@ -6466,13 +6468,13 @@ impl Engine { // event can no longer arrive after `TurnComplete` and be mistaken for // the following turn (or lost by a runtime monitor that already // settled the record). - if let Some(barrier) = mailbox_for_runtime.take() { - if status == TurnOutcomeStatus::Completed { - barrier.park_and_flush().await; - } else { - barrier.cancel_and_flush().await; - } - } + settle_turn_mailbox( + mailbox_for_runtime.take(), + status, + std::mem::take(&mut self.model_recovery_interrupted), + self.cancel_token.is_cancelled(), + ) + .await; if let Some(owner) = interactive_runtime_cost_owner.as_deref() { crate::cost_status::finish_runtime_usage_owner(owner); } @@ -9467,6 +9469,25 @@ fn terminal_turn_status_at_settlement( } } +async fn settle_turn_mailbox( + barrier: Option, + status: TurnOutcomeStatus, + model_recovery_interrupted: bool, + cancelled: bool, +) { + if let Some(barrier) = barrier { + if status == TurnOutcomeStatus::Completed + || (status == TurnOutcomeStatus::Interrupted + && model_recovery_interrupted + && !cancelled) + { + barrier.park_and_flush().await; + } else { + barrier.cancel_and_flush().await; + } + } +} + impl TurnMailboxBarrier { /// Settle the foreground subtree before closing the turn's mailbox. The /// ordering is intentional: a terminal turn event must never be emitted diff --git a/crates/tui/src/core/engine/operation_tests.rs b/crates/tui/src/core/engine/operation_tests.rs index 1dbc507516..5bd040425e 100644 --- a/crates/tui/src/core/engine/operation_tests.rs +++ b/crates/tui/src/core/engine/operation_tests.rs @@ -431,3 +431,462 @@ async fn forkguard_model_operation_preserved_partial_changes_call() { }) })); } + +struct RateRefusalClient { + calls: AtomicUsize, +} +impl LlmClient for RateRefusalClient { + fn provider_name(&self) -> &'static str { + "deepseek" + } + fn model(&self) -> &str { + "deepseek-chat" + } + async fn create_message( + &self, + _: MessageRequest, + ) -> anyhow::Result { + anyhow::bail!("stream fixture only") + } + async fn create_message_stream( + &self, + request: MessageRequest, + ) -> anyhow::Result { + if self.calls.fetch_add(1, Ordering::SeqCst) == 0 { + MockLlmClient::new(vec![canned::tool_call_turn( + "call-read", + "read_file", + r#"{"path":"fixture.txt"}"#, + )]) + .create_message_stream(request) + .await + } else { + Err(crate::llm_client::LlmError::RateLimited { + message: "recovery exhausted".into(), + retry_after: Some(Duration::from_secs(3600)), + } + .into()) + } + } +} + +#[test] +fn forkguard_rate_limit_interrupt_preserves_executed_tool_history() { + crate::test_support::run_on_engine_thread(|| async { + let workspace = tempdir().unwrap(); + fs::write(workspace.path().join("fixture.txt"), "fixture").unwrap(); + let client = Arc::new(RateRefusalClient { + calls: AtomicUsize::new(0), + }); + let (mut engine, _handle) = Engine::new_with_model_client( + deterministic_engine_config(workspace.path()), + &Config::default(), + client.clone(), + ); + engine.session.messages.push(Message { + role: Role::User, + content: vec![ContentBlock::Text { + text: "same user prompt".into(), + cache_control: None, + }], + }); + let mut registry = + crate::tools::ToolRegistry::new(crate::tools::ToolContext::new(workspace.path())); + registry.register(Arc::new(crate::tools::file::ReadFileTool)); + let tools = Some(registry.to_api_tools_with_cache(true)); + let policy = test_tool_surface(&engine, registry, tools, AppMode::Agent); + let (status, error) = engine + .run_turn( + &mut crate::core::turn::TurnContext::new(6), + policy, + None, + None, + ) + .await; + assert_eq!(status, TurnOutcomeStatus::Interrupted, "{error:?}"); + assert!(error.is_some()); + assert_eq!(client.calls.load(Ordering::SeqCst), 2); + let results = engine.session.messages.iter().flat_map(|message| &message.content) + .filter(|block| matches!(block, ContentBlock::ToolResult { tool_use_id, .. } if tool_use_id == "call-read")).count(); + assert_eq!( + results, 1, + "an exhausted model recovery must not replay executed tools" + ); + }); +} + +#[test] +fn forkguard_rate_limit_terminal_barrier_parks_children_and_honors_cancel() { + crate::test_support::run_on_engine_thread(|| async { + for cancelled in [false, true] { + let workspace = tempdir().unwrap(); + let client = Arc::new(RateRefusalClient { + calls: AtomicUsize::new(1), + }); + let (mut engine, _handle) = Engine::new_with_model_client( + deterministic_engine_config(workspace.path()), + &Config::default(), + client, + ); + engine.session.messages.push(Message { + role: Role::User, + content: vec![ContentBlock::Text { + text: "terminal fixture".into(), + cache_control: None, + }], + }); + let registry = + crate::tools::ToolRegistry::new(crate::tools::ToolContext::new(workspace.path())); + let tools = Some(registry.to_api_tools_with_cache(true)); + let policy = test_tool_surface(&engine, registry, tools, AppMode::Agent); + let (status, _) = engine + .run_turn( + &mut crate::core::turn::TurnContext::new(2), + policy, + None, + None, + ) + .await; + assert_eq!(status, TurnOutcomeStatus::Interrupted); + let recovery_pause = std::mem::take(&mut engine.model_recovery_interrupted); + assert!(recovery_pause); + let token = CancellationToken::new(); + let (mailbox, _receiver) = Mailbox::new(token.clone()); + let children = Arc::new(ForegroundChildRegistry::new()); + let child_token = token.child_token(); + let registration = children + .register("recovery-owned-child", child_token.clone()) + .unwrap(); + let parked = registration.parking_signal(); + let child = tokio::spawn(async move { + child_token.cancelled().await; + drop(registration); + }); + let (flush_tx, flush_rx) = tokio::sync::oneshot::channel(); + let drain_handle = tokio::spawn(async move { + let _ = flush_rx.await; + }); + let barrier = TurnMailboxBarrier { + mailbox, + cancel_token: token, + foreground_children: children, + flush_tx, + drain_handle, + }; + if cancelled { + engine.cancel_token.cancel(); + } + settle_turn_mailbox( + Some(barrier), + status, + recovery_pause, + engine.cancel_token.is_cancelled(), + ) + .await; + child.await.unwrap(); + assert_eq!( + parked.load(Ordering::Acquire), + !cancelled, + "provider recovery preserves continuation, explicit cancellation does not" + ); + } + }); +} + +struct TransportDropClient { + inner: crate::client::DeepSeekClient, + calls: AtomicUsize, + drops: usize, + partial: bool, + ids: Mutex>, +} +impl LlmClient for TransportDropClient { + fn provider_name(&self) -> &'static str { + "deepseek" + } + fn model(&self) -> &str { + LlmClient::model(&self.inner) + } + async fn create_message( + &self, + request: MessageRequest, + ) -> anyhow::Result { + self.inner.create_message(request).await + } + async fn create_message_stream( + &self, + request: MessageRequest, + ) -> anyhow::Result { + self.create_message_stream_for_operation(request, uuid::Uuid::new_v4()) + .await + } + async fn create_message_stream_for_operation( + &self, + request: MessageRequest, + id: uuid::Uuid, + ) -> anyhow::Result { + self.ids.lock().unwrap().push(id); + let stream = self + .inner + .create_message_stream_for_operation(request, id) + .await?; + if self.calls.fetch_add(1, Ordering::SeqCst) < self.drops { + drop(stream); + let mut events = Vec::new(); + if self.partial { + events.extend([ + Ok(canned::message_start("partial-refusal-fixture")), + Ok(canned::text_block_start(0)), + Ok(canned::text_delta(0, "already delivered partial answer")), + ]); + } + events.push(Err(anyhow::anyhow!( + "Stream read error: connection reset before content" + ))); + Ok(Box::pin(futures_util::stream::iter(events))) + } else { + Ok(stream) + } + } +} + +#[test] +fn forkguard_rate_limit_stream_reopen_keeps_budget_identity_and_receipts() { + crate::test_support::run_on_engine_thread(|| async { + use wiremock::matchers::method; + use wiremock::{Mock, MockServer, Request, ResponseTemplate}; + for (max_attempts, drops, refusal_index, partial, expected_wires) in [ + (2, 1, 1, false, 2), + (3, 1, 1, false, 3), + (4, 3, 3, false, 4), + (2, 1, 0, true, 3), + (2, 1, 0, false, 2), + ] { + let server = MockServer::start().await; + let wires = Arc::new(AtomicUsize::new(0)); + let seen = wires.clone(); + Mock::given(method("POST")).respond_with(move |_: &Request| { + if seen.fetch_add(1, Ordering::SeqCst) == refusal_index { + ResponseTemplate::new(429).insert_header("retry-after", "0") + .set_body_json(serde_json::json!({"error":{"code":"insufficient_quota","message":"TPM exceeded"}})) + } else { + ResponseTemplate::new(200).insert_header("content-type", "text/event-stream") + .set_body_string(concat!( + "data: {\"choices\":[{\"index\":0,\"delta\":{\"content\":\"done\"},\"finish_reason\":null}]}\n\n", + "data: {\"choices\":[{\"index\":0,\"delta\":{},\"finish_reason\":\"stop\"}]}\n\n", + "data: [DONE]\n\n")) + } + }).mount(&server).await; + let mut policy = crate::config::RateLimitRetryPolicy::default(); + policy.max_attempts = max_attempts; + policy.initial_delay_secs = 1; + policy.max_delay_secs = 1; + let config = Config { + provider: Some("custom".into()), + // Wiremock may reuse a server address between independent cases. + // Each case owns a separate credential and recovery episode. + api_key: Some(format!( + "owned-{max_attempts}-{drops}-{refusal_index}-{partial}" + )), + base_url: Some(server.uri()), + retry: Some(crate::config::RetryConfig { + rate_limit: policy, + enabled: None, + max_retries: Some(0), + initial_delay: None, + max_delay: None, + exponential_base: None, + }), + ..Default::default() + }; + let client = Arc::new(TransportDropClient { + inner: crate::client::DeepSeekClient::new(&config).unwrap(), + calls: AtomicUsize::new(0), + drops, + partial, + ids: Mutex::new(Vec::new()), + }); + let workspace = tempdir().unwrap(); + let (mut engine, handle) = Engine::new_with_model_client( + deterministic_engine_config(workspace.path()), + &config, + client.clone(), + ); + engine.session.messages.push(Message { + role: Role::User, + content: vec![ContentBlock::Text { + text: "stream recovery fixture".into(), + cache_control: None, + }], + }); + let registry = + crate::tools::ToolRegistry::new(crate::tools::ToolContext::new(workspace.path())); + let tools = Some(registry.to_api_tools_with_cache(true)); + let surface = test_tool_surface(&engine, registry, tools, AppMode::Agent); + let (status, error) = engine + .run_turn( + &mut crate::core::turn::TurnContext::new(2), + surface, + None, + None, + ) + .await; + assert_eq!( + status, + if max_attempts == 3 || partial { + TurnOutcomeStatus::Completed + } else { + TurnOutcomeStatus::Interrupted + }, + "{error:?}" + ); + assert_eq!(wires.load(Ordering::SeqCst), expected_wires); + let ids = client.ids.lock().unwrap().clone(); + assert_eq!(ids.len(), drops + 1); + if partial { + assert_ne!( + ids[0], ids[1], + "persisted partial content changes the logical request" + ); + } else { + assert!( + ids.iter().all(|id| *id == ids[0]), + "ordinary routes retain a logical operation through inner and outer recovery" + ); + } + let expected = if max_attempts == 3 || partial { + "waiting" + } else { + "interrupted" + }; + let mut events = handle.rx_event.write().await; + let mut receipt_seen = false; + while let Ok(event) = events.try_recv() { + if let Event::ModelRecovery { receipt } = event { + receipt_seen |= receipt.state == expected + && receipt.operation_id.as_deref() == Some(ids[0].to_string().as_str()); + } + } + assert!( + receipt_seen, + "transparent reopen must publish its recovery state" + ); + let requests = server.received_requests().await.unwrap(); + if partial { + assert_eq!(requests[0].body, requests[1].body); + assert_ne!(requests[1].body, requests[2].body); + assert!( + String::from_utf8_lossy(&requests[2].body) + .contains("already delivered partial answer") + ); + } else { + assert!( + requests + .iter() + .all(|request| request.body == requests[0].body) + ); + } + assert!( + requests + .iter() + .all(|request| !request.headers.contains_key("idempotency-key")), + "runtime identity does not opt a provider into a wire header" + ); + } + }); +} + +#[test] +fn forkguard_rate_limit_vision_tool_dispatch_publishes_recovery() { + crate::test_support::run_on_engine_thread(|| async { + use base64::Engine as _; + use wiremock::matchers::method; + use wiremock::{Mock, MockServer, Request, ResponseTemplate}; + let server = MockServer::start().await; + let calls = Arc::new(AtomicUsize::new(0)); + let seen = calls.clone(); + Mock::given(method("POST")) + .respond_with(move |_: &Request| { + if seen.fetch_add(1, Ordering::SeqCst) == 0 { + ResponseTemplate::new(429).insert_header("retry-after", "0") + } else { + ResponseTemplate::new(200).set_body_json( + serde_json::json!({"choices":[{"message":{"content":"vision done"}}]}), + ) + } + }) + .mount(&server) + .await; + let workspace = tempdir().unwrap(); + let png = base64::engine::general_purpose::STANDARD.decode("iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mP8/x8AAwMCAO+jRZkAAAAASUVORK5CYII=").unwrap(); + fs::write(workspace.path().join("tiny.png"), png).unwrap(); + let base = server.uri(); + let mut policy = crate::config::RateLimitRetryPolicy::default(); + policy.initial_delay_secs = 1; + policy.max_delay_secs = 1; + let route = crate::client::DeepSeekClient::new(&Config { + api_key: Some("owned-vision-key".into()), + base_url: Some(base.clone()), + retry: Some(crate::config::RetryConfig { + rate_limit: policy, + enabled: None, + max_retries: None, + initial_delay: None, + max_delay: None, + exponential_base: None, + }), + ..Default::default() + }) + .unwrap(); + let context = crate::tools::ToolContext::new(workspace.path()) + .with_origin_turn_id("vision-tool-turn"); + let mut registry = crate::tools::ToolRegistry::new(context); + registry.register(Arc::new( + crate::vision::tools::ImageAnalyzeTool::new_with_route_client( + crate::config::VisionModelConfig { + model: "vision-fixture".into(), + api_key: Some("owned-vision-key".into()), + base_url: Some(base), + }, + Some(route), + ), + )); + let (tx, mut rx) = tokio::sync::mpsc::channel(64); + Engine::execute_tool_with_lock( + Arc::new(tokio::sync::RwLock::new(())), + true, + false, + tx, + None, + "image_analyze".into(), + serde_json::json!({"image_path":"tiny.png"}), + workspace.path().to_owned(), + Some(®istry), + None, + None, + None, + ) + .await + .unwrap(); + let mut receipts = Vec::new(); + while let Ok(event) = rx.try_recv() { + if let Event::ModelRecovery { receipt } = event { + receipts.push(receipt); + } + } + assert!(receipts.iter().any(|r| r.state == "waiting")); + assert!(receipts.iter().any(|r| r.state == "recovered")); + assert!( + receipts + .iter() + .all(|r| r.turn_id.as_deref() == Some("vision-tool-turn") + && r.operation_id.is_some()) + ); + assert!( + receipts + .iter() + .all(|r| r.operation_id == receipts[0].operation_id) + ); + assert_eq!(calls.load(Ordering::SeqCst), 2); + }); +} diff --git a/crates/tui/src/core/engine/tool_execution.rs b/crates/tui/src/core/engine/tool_execution.rs index 6517b3f2cf..06e167a693 100644 --- a/crates/tui/src/core/engine/tool_execution.rs +++ b/crates/tui/src/core/engine/tool_execution.rs @@ -624,9 +624,30 @@ impl Engine { .await .map(RichToolResult::plain) } else if let Some(registry) = registry { - registry - .execute_rich_full_with_context(&tool_name, tool_input, context_override.as_ref()) - .await + let context = context_override + .as_ref() + .unwrap_or_else(|| registry.context()); + let agent_id = context.owner_agent_id.clone(); + let turn_id = context.origin_turn_id.clone(); + let event_tx = tx_event.clone(); + let observer: crate::llm_client::rate_limit::RecoveryObserver = + Arc::new(move |mut receipt| { + receipt.agent_id = agent_id.clone(); + receipt.turn_id = turn_id.clone(); + let event_tx = event_tx.clone(); + Box::pin(async move { + let _ = event_tx.send(Event::ModelRecovery { receipt }).await; + }) + }); + crate::llm_client::rate_limit::with_observer( + observer, + registry.execute_rich_full_with_context( + &tool_name, + tool_input, + context_override.as_ref(), + ), + ) + .await } else { Err(ToolError::not_available(format!( "tool '{tool_name}' is not registered" diff --git a/crates/tui/src/core/engine/turn_loop.rs b/crates/tui/src/core/engine/turn_loop.rs index 4db39596f3..4e65907521 100644 --- a/crates/tui/src/core/engine/turn_loop.rs +++ b/crates/tui/src/core/engine/turn_loop.rs @@ -826,6 +826,7 @@ struct StreamOutcome { first_token_at: Option, request_dispatched_at: Instant, stream_error: Option, + rate_limit_interrupted: bool, } pub(super) fn initial_stream_error_user_message( @@ -1564,6 +1565,7 @@ impl Engine { // registry-derived field as unknown rather than guessing. inspection_surface: Option, ) -> (TurnOutcomeStatus, Option) { + self.model_recovery_interrupted = false; // R1: restart the cumulative per-turn wall-clock budget. This is the // only place it is started, so exactly one turn owns it at a time. self.turn_wall_clock = @@ -1861,12 +1863,12 @@ impl Engine { && snapshot.tokens_used >= u64::from(budget) { let _ = self - .tx_event - .send(Event::status(format!( - "Goal over token budget ({} / {budget} tokens) — continuing (unbounded); verify or /goal clear when done.", - snapshot.tokens_used - ))) - .await; + .tx_event + .send(Event::status(format!( + "Goal over token budget ({} / {budget} tokens) — continuing (unbounded); verify or /goal clear when done.", + snapshot.tokens_used + ))) + .await; } let auto_compaction_config = self.config.compaction.clone(); @@ -1960,11 +1962,11 @@ impl Engine { let Some(compaction_result) = compaction_result else { self.finish_compaction(&compaction_id); let message = if turn_was_canceled { - "Auto-compaction canceled with the active turn; conversation context was not changed" - } else { - "Auto-compaction canceled; conversation context was not changed" - } - .to_string(); + "Auto-compaction canceled with the active turn; conversation context was not changed" + } else { + "Auto-compaction canceled; conversation context was not changed" + } + .to_string(); self.emit_compaction_cancelled(compaction_id, true, message) .await; if turn_was_canceled { @@ -1996,11 +1998,11 @@ impl Engine { if turn_was_canceled || compaction_cancel.is_cancelled() { self.finish_compaction(&compaction_id); let message = if turn_was_canceled { - "Auto-compaction canceled with the active turn; conversation context was not changed" - } else { - "Auto-compaction canceled; conversation context was not changed" - } - .to_string(); + "Auto-compaction canceled with the active turn; conversation context was not changed" + } else { + "Auto-compaction canceled; conversation context was not changed" + } + .to_string(); self.emit_compaction_cancelled(compaction_id, true, message) .await; if turn_was_canceled { @@ -2122,7 +2124,7 @@ impl Engine { if context_recovery_attempts >= MAX_CONTEXT_RECOVERY_ATTEMPTS { let message = format!( "Context remains above model limit after {MAX_CONTEXT_RECOVERY_ATTEMPTS} recovery attempts \ - (~{estimated_input} token estimate, ~{input_budget} budget). Please run /compact or /clear." + (~{estimated_input} token estimate, ~{input_budget} budget). Please run /compact or /clear." ); turn_error = Some(TurnFailure::runtime(message.clone())); let _ = self @@ -2373,14 +2375,20 @@ impl Engine { let _ = self.tx_event.send(Event::status("Request cancelled")).await; return (TurnOutcomeStatus::Interrupted, None); } - result = async { client.prepare_stream_call(request) } => result, + result = async { + let call = client.prepare_stream_call(request)?; + let identity = match &call.operation_identity { + Some(identity) => identity.clone(), + None => crate::llm_client::caller_stream_operation_identity( + client.provider_name(), client.billing_base_url(), &call.request, + )?, + }; + Ok::<_, anyhow::Error>((call, identity)) + } => result, }; - let prepared_stream = preparation.as_ref().ok().cloned(); - let operation_id = match prepared_stream - .as_ref() - .and_then(|call| call.operation_identity.as_ref()) - { - Some(identity) => { + let prepared_stream = preparation.as_ref().ok().map(|(call, _)| call.clone()); + let operation_id = match preparation.as_ref().ok() { + Some((_, identity)) => { let id = match &model_operation { Some((previous, id)) if previous == identity => *id, _ => uuid::Uuid::new_v4(), @@ -2388,10 +2396,7 @@ impl Engine { model_operation = Some((identity.clone(), id)); id } - None => { - model_operation = None; - uuid::Uuid::new_v4() - } + None => uuid::Uuid::new_v4(), }; if prepared_stream.is_some() { let _ = self @@ -2417,6 +2422,17 @@ impl Engine { // instant (connection setup included), and time-to-first-token is // the gap to the first content-bearing stream event. let request_dispatched_at = Instant::now(); + let tx = self.tx_event.clone(); + let recovery_turn_id = turn.id.clone(); + let observer: crate::llm_client::rate_limit::RecoveryObserver = + std::sync::Arc::new(move |mut receipt| { + receipt.operation_id = Some(operation_id.to_string()); + receipt.turn_id = Some(recovery_turn_id.clone()); + let tx = tx.clone(); + Box::pin(async move { + let _ = tx.send(Event::ModelRecovery { receipt }).await; + }) + }); let request_error_source = if preparation.is_ok() { crate::core::events::TurnErrorSource::Model } else { @@ -2430,8 +2446,9 @@ impl Engine { return (TurnOutcomeStatus::Interrupted, None); } result = async { - let prepared = preparation?; - client.create_prepared_message_stream_for_operation(prepared, operation_id).await + let (prepared, _) = preparation?; + crate::llm_client::rate_limit::with_observer(observer.clone(), + client.create_prepared_message_stream_for_operation(prepared, operation_id)).await } => result, }; let stream = match stream_result { @@ -2476,6 +2493,14 @@ impl Engine { let _ = self.tx_event.send(Event::status(status)).await; continue; } + let rate_limited = matches!( + e.downcast_ref::(), + Some( + crate::llm_client::LlmError::RateLimited { .. } + | crate::llm_client::LlmError::QuotaWindow { .. } + ) + ); + self.model_recovery_interrupted |= rate_limited; turn_error = Some(TurnFailure { message: message.clone(), source: request_error_source, @@ -2486,6 +2511,15 @@ impl Engine { e, message, ))) .await; + if rate_limited { + // The recovery episode owns the resume, so the turn + // parks as interrupted instead of consuming the typed + // failure as a terminal model error. + return ( + TurnOutcomeStatus::Interrupted, + turn_error.map(|failure| failure.message), + ); + } return turn_error.expect("failure was assigned").finish(turn); } }; @@ -2509,18 +2543,23 @@ impl Engine { first_token_at, request_dispatched_at, stream_error, - } = self - .process_stream( - client.as_ref(), - stream, - prepared_stream - .as_ref() - .expect("successful stream has a prepared call"), - operation_id, - request_dispatched_at, - stream_retry_budget.spent(), - ) - .await; + rate_limit_interrupted, + } = crate::llm_client::rate_limit::with_observer( + observer, + Box::pin( + self.process_stream( + client.as_ref(), + stream, + prepared_stream + .as_ref() + .expect("successful stream has a prepared call"), + operation_id, + request_dispatched_at, + stream_retry_budget.spent(), + ), + ), + ) + .await; turn_error = turn_error.or(stream_error); // These belong to post-stream response assembly, not stream // consumption: blocks are built from the completed stream state, @@ -2564,6 +2603,14 @@ impl Engine { return (TurnOutcomeStatus::Interrupted, None); } + if rate_limit_interrupted { + self.model_recovery_interrupted = true; + return ( + TurnOutcomeStatus::Interrupted, + turn_error.map(|failure| failure.message), + ); + } + if is_incomplete_stop_reason(stop_reason.as_deref()) { let reason = stop_reason_detail(stop_reason.as_deref()); if is_output_limit_stop_reason(stop_reason.as_deref()) && stream_errors == 0 { @@ -2651,11 +2698,11 @@ impl Engine { "Resuming after system sleep (attempt {attempt}/{MAX_STREAM_RETRIES}); discarding partial output and retrying request" )); let _ = self - .tx_event - .send(Event::status(format!( - "System sleep detected; connection lost — retrying request ({attempt}/{MAX_STREAM_RETRIES})" - ))) - .await; + .tx_event + .send(Event::status(format!( + "System sleep detected; connection lost — retrying request ({attempt}/{MAX_STREAM_RETRIES})" + ))) + .await; // Finalize any partially-rendered assistant cell so // the retried stream renders fresh instead of // appending to the pre-sleep fragment. @@ -2736,21 +2783,21 @@ impl Engine { "Resuming interactive turn after mid-stream network drop (attempt {attempt}/{MAX_STREAM_RETRIES}); only hidden reasoning streamed — no partial reply to preserve, retrying request" )); let _ = self - .tx_event - .send(Event::status(format!( - "Connection interrupted; retrying ({attempt}/{MAX_STREAM_RETRIES})" - ))) - .await; + .tx_event + .send(Event::status(format!( + "Connection interrupted; retrying ({attempt}/{MAX_STREAM_RETRIES})" + ))) + .await; } else { crate::logging::warn(format!( "Resuming interactive turn after mid-stream network drop (attempt {attempt}/{MAX_STREAM_RETRIES}); preserving partial reply and retrying request" )); let _ = self - .tx_event - .send(Event::status(format!( - "Connection interrupted; preserving partial reply and retrying ({attempt}/{MAX_STREAM_RETRIES})" - ))) - .await; + .tx_event + .send(Event::status(format!( + "Connection interrupted; preserving partial reply and retrying ({attempt}/{MAX_STREAM_RETRIES})" + ))) + .await; // Finalize the partial text cell so the UI stops // streaming and the retried content lands in a // fresh cell instead of appending to an @@ -2906,11 +2953,11 @@ impl Engine { return TurnFailure::runtime(reason).finish(turn); } self.add_session_message(self.runtime_text_message_with_turn_metadata( - "All tools are disabled. The previous tool markers were ignored. Respond now with only the requested final answer in the exact required format." - .to_string(), - UserInputProvenance::Runtime, - )) - .await; + "All tools are disabled. The previous tool markers were ignored. Respond now with only the requested final answer in the exact required format." + .to_string(), + UserInputProvenance::Runtime, + )) + .await; turn.next_step(); continue; } @@ -3003,14 +3050,14 @@ impl Engine { .take() .expect("output_limit_truncated checked above"); self.add_session_message( - self.runtime_text_message_with_turn_metadata( - format!( - "[runtime] The provider stopped generation at its output limit (`{reason}`) before completing. Your last response was cut off. Continue from where you left off; do not repeat content already delivered." - ), - UserInputProvenance::Runtime, + self.runtime_text_message_with_turn_metadata( + format!( + "[runtime] The provider stopped generation at its output limit (`{reason}`) before completing. Your last response was cut off. Continue from where you left off; do not repeat content already delivered." ), - ) - .await; + UserInputProvenance::Runtime, + ), + ) + .await; let _ = self .tx_event .send(Event::status( @@ -3353,11 +3400,11 @@ impl Engine { // No FINAL — let the model iterate with the feedback. let _ = self - .tx_event - .send(Event::status(format!( - "Continuing — REPL round feedback (consecutive_empty={consecutive_empty_repl_rounds})" - ))) - .await; + .tx_event + .send(Event::status(format!( + "Continuing — REPL round feedback (consecutive_empty={consecutive_empty_repl_rounds})" + ))) + .await; turn.next_step(); continue; } @@ -3522,7 +3569,7 @@ impl Engine { { let message = if has_provider_reasoning { "Model returned reasoning but no answer or tool call; the provider response was incomplete." - .to_string() + .to_string() } else if let Some(reason) = stop_reason.as_deref() { format!( "Model returned terminal stop reason `{reason}` with no answer or tool call." @@ -3686,11 +3733,11 @@ impl Engine { self.session.pending_prefix_change_reason = Some("benchmark_final_only".to_string()); self.add_session_message(self.runtime_text_message_with_turn_metadata( - "Tool budget is exhausted. All tools are now disabled for this turn. Do not call any tool again. Use only the evidence already present in the conversation and provide the requested final answer now, in the exact output format requested by the user." - .to_string(), - UserInputProvenance::Runtime, - )) - .await; + "Tool budget is exhausted. All tools are now disabled for this turn. Do not call any tool again. Use only the evidence already present in the conversation and provide the requested final answer now, in the exact output format requested by the user." + .to_string(), + UserInputProvenance::Runtime, + )) + .await; let _ = self .tx_event .send(Event::status( @@ -3718,14 +3765,14 @@ impl Engine { // assistant `tool_use` directly) and the model can act on it. if let Some(reason) = output_limit_truncated.take() { self.add_session_message( - self.runtime_text_message_with_turn_metadata( - format!( - "[runtime] The provider stopped generation at its output limit (`{reason}`) before completing. Your last response was cut off. Continue from where you left off; do not repeat content already delivered." - ), - UserInputProvenance::Runtime, + self.runtime_text_message_with_turn_metadata( + format!( + "[runtime] The provider stopped generation at its output limit (`{reason}`) before completing. Your last response was cut off. Continue from where you left off; do not repeat content already delivered." ), - ) - .await; + UserInputProvenance::Runtime, + ), + ) + .await; } // A successful tool step is productive progress, not a runaway @@ -3748,11 +3795,11 @@ impl Engine { .map_or(0, |registry| registry.active_count()); if running > 0 { let _ = self - .tx_event - .send(Event::status(format!( - "Turn failed with {running} turn-owned sub-agent(s) still running; cancelling them." - ))) - .await; + .tx_event + .send(Event::status(format!( + "Turn failed with {running} turn-owned sub-agent(s) still running; cancelling them." + ))) + .await; } return err.finish(turn); } @@ -3761,11 +3808,11 @@ impl Engine { .map_or(0, |registry| registry.active_count()); if running > 0 { let _ = self - .tx_event - .send(Event::status(format!( - "Turn ending with {running} turn-owned sub-agent(s) still running; parking them as resumable work." - ))) - .await; + .tx_event + .send(Event::status(format!( + "Turn ending with {running} turn-owned sub-agent(s) still running; parking them as resumable work." + ))) + .await; self.add_session_message(self.runtime_text_message_with_turn_metadata( turn_owned_child_parking_runtime_text(running), UserInputProvenance::Runtime, @@ -3778,11 +3825,11 @@ impl Engine { }; if detached_running > 0 { let _ = self - .tx_event - .send(Event::status(format!( - "Turn ending with {detached_running} detached sub-agent(s) still running in the background; they'll report when done." - ))) - .await; + .tx_event + .send(Event::status(format!( + "Turn ending with {detached_running} detached sub-agent(s) still running in the background; they'll report when done." + ))) + .await; self.add_session_message(waiting_for_subagents_runtime_message(detached_running)) .await; } @@ -5797,6 +5844,7 @@ impl Engine { // breaking the existing pin invariants. let mut stream = stream; let mut stream_error: Option = None; + let mut rate_limit_interrupted = false; let mut current_text_raw = String::new(); let mut current_text_visible = String::new(); @@ -6018,6 +6066,13 @@ impl Engine { retry_err, crate::core::events::TurnErrorSource::Model, ); + rate_limit_interrupted = matches!( + retry_err.downcast_ref::(), + Some( + crate::llm_client::LlmError::RateLimited { .. } + | crate::llm_client::LlmError::QuotaWindow { .. } + ) + ); let retry_msg = self.decorate_auth_error_message(format!( "Stream retry failed: {retry_err}" )); @@ -6576,6 +6631,7 @@ impl Engine { first_token_at, request_dispatched_at, stream_error, + rate_limit_interrupted, } } diff --git a/crates/tui/src/core/events.rs b/crates/tui/src/core/events.rs index 4115b64b1a..d4841d1d14 100644 --- a/crates/tui/src/core/events.rs +++ b/crates/tui/src/core/events.rs @@ -3,6 +3,8 @@ //! These events flow from the engine to the TUI via a channel, //! enabling non-blocking, real-time updates. +pub use crate::llm_client::rate_limit::RecoveryReceipt; + use std::{path::PathBuf, sync::Arc}; use chrono::{DateTime, Utc}; @@ -171,6 +173,9 @@ impl AgentProgressEventMeta { /// Events emitted by the engine to update the UI. #[derive(Debug, Clone)] pub enum Event { + ModelRecovery { + receipt: RecoveryReceipt, + }, /// A route compatibility check omitted tools from the provider request. /// Names are wire-safe and bounded by the request catalog; schemas and /// validator details never cross this user-visible boundary. @@ -181,22 +186,36 @@ pub enum Event { }, // === Streaming Events === /// A new message block has started - MessageStarted { index: usize }, + MessageStarted { + index: usize, + }, /// Incremental text content delta - MessageDelta { index: usize, content: String }, + MessageDelta { + index: usize, + content: String, + }, /// Message block completed - MessageComplete { index: usize }, + MessageComplete { + index: usize, + }, /// Thinking block started - ThinkingStarted { index: usize }, + ThinkingStarted { + index: usize, + }, /// Incremental thinking content delta - ThinkingDelta { index: usize, content: String }, + ThinkingDelta { + index: usize, + content: String, + }, /// Thinking block completed - ThinkingComplete { index: usize }, + ThinkingComplete { + index: usize, + }, // === Tool Events === /// Tool call initiated @@ -249,7 +268,10 @@ pub enum Event { /// Immutable billing route captured immediately before the first provider /// request, after snapshots and other potentially slow pre-dispatch work. - RouteDispatched { turn_id: String, route: TurnRoute }, + RouteDispatched { + turn_id: String, + route: TurnRoute, + }, /// The turn is complete (no more tool calls) TurnComplete { @@ -297,17 +319,23 @@ pub enum Event { /// Runtime goal state changed inside the engine, usually from model-visible /// `create_goal` or `update_goal` tool calls. - GoalUpdated { snapshot: GoalSnapshot }, + GoalUpdated { + snapshot: GoalSnapshot, + }, /// The interactive engine is in the configured quiet period before one /// already-authorized goal continuation. This is lifecycle state, not a /// status string: Esc/Ctrl+C can cancel it without pretending a provider /// turn is still in flight. - GoalContinuationWaiting { delay_seconds: u64 }, + GoalContinuationWaiting { + delay_seconds: u64, + }, /// The between-turn quiet period ended. `interrupted` distinguishes a /// user/external cancel from normal expiry or a goal status control. - GoalContinuationWaitEnded { interrupted: bool }, + GoalContinuationWaitEnded { + interrupted: bool, + }, /// Context compaction started. CompactionStarted { @@ -368,7 +396,9 @@ pub enum Event { }, /// Context purge failed. - PurgeFailed { message: String }, + PurgeFailed { + message: String, + }, /// Context compaction failed. CompactionFailed { @@ -413,10 +443,14 @@ pub enum Event { }, /// A queued steer was committed into the session transcript. - SteerCommitted { steer_id: String }, + SteerCommitted { + steer_id: String, + }, /// A queued steer was retired without reaching the transcript. - SteerDropped { steer_id: String }, + SteerDropped { + steer_id: String, + }, /// Receipt for an operator follow-up sent to a child (`Op::FollowUpSubAgent`). /// `Ok` carries the delivery outcome (the target id may differ from the @@ -478,7 +512,9 @@ pub enum Event { }, /// Status message for UI display - Status { message: String }, + Status { + message: String, + }, /// Session-owned MCP + plugin boot progress. /// @@ -504,7 +540,9 @@ pub enum Event { /// prompt; it never carries runtime/system additions. There is no error /// variant: a manifest that cannot describe something says so in a typed /// unavailable section instead. - RequestManifestReady { rendered: String }, + RequestManifestReady { + rendered: String, + }, /// Pause terminal input events (for interactive subprocesses). PauseEvents { diff --git a/crates/tui/src/core/mod.rs b/crates/tui/src/core/mod.rs index 3ebb0dc00a..feed1e9788 100644 --- a/crates/tui/src/core/mod.rs +++ b/crates/tui/src/core/mod.rs @@ -15,6 +15,7 @@ #![deny(clippy::print_stderr)] pub mod authority; +pub(crate) mod capacity; // `crates/core` now owns the engine (issue #5261). The TUI's `engine` // consumes the core boundary directly so headless and TUI share one // `Op`-in / `EventMsg`-out API, one `ThreadId`/`SessionId` type, and one diff --git a/crates/tui/src/core/protocol_parity.rs b/crates/tui/src/core/protocol_parity.rs index 742b8ebb6a..ab365bd52a 100644 --- a/crates/tui/src/core/protocol_parity.rs +++ b/crates/tui/src/core/protocol_parity.rs @@ -396,6 +396,14 @@ pub fn event_to_protocol(event: &Event, ids: &ProtocolIds) -> wire::EventMsg { let thread_id = ids.thread_id.clone(); let session_id = ids.session_id.clone(); match event { + Event::ModelRecovery { receipt } => wire::EventMsg::Status { + thread_id, + session_id, + message: format!( + "Model recovery: {} ({}/{}) next_retry_at_ms={:?}", + receipt.state, receipt.attempt, receipt.max_attempts, receipt.next_retry_at_ms + ), + }, Event::ToolProjectionWarning { provider, omitted_tool_names, diff --git a/crates/tui/src/error_taxonomy.rs b/crates/tui/src/error_taxonomy.rs index ee655c41a8..1588018331 100644 --- a/crates/tui/src/error_taxonomy.rs +++ b/crates/tui/src/error_taxonomy.rs @@ -236,13 +236,15 @@ pub fn envelope_for_llm_error(error: anyhow::Error, display_message: String) -> impl From for ErrorEnvelope { fn from(value: LlmError) -> Self { match value { - LlmError::RateLimited { message, .. } => Self::new( - ErrorCategory::RateLimit, - ErrorSeverity::Warning, - true, - "llm_rate_limited", - message, - ), + LlmError::RateLimited { message, .. } | LlmError::QuotaWindow { message, .. } => { + Self::new( + ErrorCategory::RateLimit, + ErrorSeverity::Warning, + true, + "llm_rate_limited", + message, + ) + } // Keep the broad wire-compatible category while making the typed // code and recovery contract distinct from an ordinary 429. LlmError::QuotaExhausted(error) => Self::new( diff --git a/crates/tui/src/lib.rs b/crates/tui/src/lib.rs index 7741a034ea..bc583820da 100644 --- a/crates/tui/src/lib.rs +++ b/crates/tui/src/lib.rs @@ -1,6 +1,8 @@ //! Codewhale TUI library — single-binary entry point. #![allow(clippy::uninlined_format_args)] +// Scoped request/stream recovery futures exceed the default compiler query depth. +#![recursion_limit = "256"] use std::collections::{BTreeSet, HashMap}; use std::io::{self, IsTerminal, Read, Write}; diff --git a/crates/tui/src/llm_client/mod.rs b/crates/tui/src/llm_client/mod.rs index 45a6dd1697..39b6281d98 100644 --- a/crates/tui/src/llm_client/mod.rs +++ b/crates/tui/src/llm_client/mod.rs @@ -22,7 +22,9 @@ //! }, None).await; //! ``` -use crate::config::RetryPolicy; +pub mod rate_limit; + +use crate::config::{RateLimitRetryPolicy, RetryPolicy}; use crate::models::{MessageRequest, MessageResponse, StreamEvent}; use anyhow::Result; use serde_json::Value; @@ -84,7 +86,8 @@ impl PreparedStreamCall { } /// Compatibility identity for callers explicitly requesting a conservative -/// caller-request hash. The Engine's ordinary dispatch path does not invoke it. +/// caller-request hash. The Engine also uses it as a private recovery identity +/// when no prepared transport identity is available; it never enables a wire header. pub fn caller_stream_operation_identity( provider: &str, base_url: Option<&str>, @@ -290,6 +293,11 @@ pub enum LlmError { message: String, retry_after: Option, }, + /// A long-lived quota window. Preserve the checkpoint; resume requires task policy or user action. + QuotaWindow { + message: String, + retry_after: Option, + }, /// The provider explicitly reported that the account's plan quota is exhausted. /// @@ -335,6 +343,13 @@ pub enum LlmError { impl std::fmt::Display for LlmError { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { match self { + LlmError::QuotaWindow { + message, + retry_after, + } => write!( + f, + "Rate limit exceeded: {message}; quota window reset in {retry_after:?}" + ), LlmError::RateLimited { message, .. } => write!(f, "Rate limit exceeded: {message}"), LlmError::QuotaExhausted(error) => { write!(f, "Provider plan quota exhausted: {}", error.message) @@ -363,6 +378,19 @@ impl std::fmt::Display for LlmError { impl std::error::Error for LlmError {} impl LlmError { + fn for_recovery_budget(self, remaining: Duration) -> Self { + match self { + Self::RateLimited { + message, + retry_after: Some(minimum), + } if minimum >= remaining => Self::QuotaWindow { + message, + retry_after: Some(minimum), + }, + error => error, + } + } + /// Determines if this error is potentially transient and worth retrying. /// /// Retryable errors: @@ -393,7 +421,8 @@ impl LlmError { /// provides a Retry-After header. pub fn suggested_retry_delay(&self) -> Option { match self { - LlmError::RateLimited { retry_after, .. } => *retry_after, + LlmError::RateLimited { retry_after, .. } + | LlmError::QuotaWindow { retry_after, .. } => *retry_after, _ => None, } } @@ -840,6 +869,7 @@ impl From for LlmError { /// - `jitter_factor`: 0.1 (10% variation) #[derive(Debug, Clone)] pub struct RetryConfig { + pub rate_limit: RateLimitRetryPolicy, /// Whether retry logic is enabled pub enabled: bool, @@ -875,6 +905,7 @@ pub struct RetryConfig { impl Default for RetryConfig { fn default() -> Self { Self { + rate_limit: RateLimitRetryPolicy::default(), enabled: true, max_retries: 3, initial_delay: 1.0, @@ -967,6 +998,7 @@ impl RetryConfig { impl From for RetryConfig { fn from(policy: RetryPolicy) -> Self { Self { + rate_limit: policy.rate_limit, enabled: policy.enabled, max_retries: policy.max_retries, initial_delay: policy.initial_delay, @@ -1063,126 +1095,271 @@ where { // If retries are disabled, just run once if !config.enabled { - return operation().await.map_err(|e| RetryError { - last_error: e, - attempts: 1, - total_time: Duration::ZERO, + let budget = rate_limit::operation(); + let _observation = rate_limit::CancellationObservation::new(budget.clone()); + let attempts = { + let mut budget = budget + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + budget.attempts = budget.attempts.saturating_add(1); + budget.attempts + }; + let started = Instant::now(); + let result = operation().await; + let remaining = budget + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .started + .map_or( + Duration::from_secs(config.rate_limit.recovery_budget_secs), + |started| { + Duration::from_secs(config.rate_limit.recovery_budget_secs) + .saturating_sub(started.elapsed()) + }, + ); + let result = result.map_err(|error| error.for_recovery_budget(remaining)); + match &result { + Ok(_) + if budget + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .observed_wait => + { + rate_limit::observe(attempts, config.rate_limit.max_attempts, None, "recovered") + .await; + } + Err(error) + if matches!( + error, + LlmError::RateLimited { .. } | LlmError::QuotaWindow { .. } + ) => + { + rate_limit::record_terminal_quota(error); + rate_limit::observe( + attempts, + config.rate_limit.max_attempts, + error.suggested_retry_delay(), + "interrupted", + ) + .await; + } + Err(_) + if budget + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .observed_wait => + { + rate_limit::observe(attempts, config.rate_limit.max_attempts, None, "stopped") + .await; + } + _ => {} + } + return result.map_err(|last_error| RetryError { + last_error, + attempts, + total_time: started.elapsed(), }); } let start_time = Instant::now(); - let total_timeout = if config.total_timeout > 0.0 { - Some(Duration::from_secs_f64(config.total_timeout)) - } else { - None - }; - - let mut last_error: Option = None; - - // Attempt 0 is the first try, then up to max_retries additional attempts - for attempt in 0..=config.max_retries { - // Check total timeout - if let Some(timeout) = total_timeout - && start_time.elapsed() >= timeout - { + let operation_budget = rate_limit::operation(); + let _observation = rate_limit::CancellationObservation::new(operation_budget.clone()); + let mut transport_retries = 0u32; + let mut recovery_started = operation_budget + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .started; + let mut last_error = None; + loop { + let (attempts, exhausted) = { + let mut budget = operation_budget + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if budget.started.is_some_and(|started| { + budget.attempts >= config.rate_limit.max_attempts + || started.elapsed() + >= Duration::from_secs(config.rate_limit.recovery_budget_secs) + }) { + (budget.attempts, true) + } else { + budget.attempts = budget.attempts.saturating_add(1); + (budget.attempts, false) + } + }; + if exhausted { + let error = last_error.take().unwrap_or_else(|| LlmError::QuotaWindow { + message: "Logical operation recovery budget exhausted".into(), + retry_after: None, + }); + rate_limit::record_terminal_quota(&error); + rate_limit::observe( + attempts, + config.rate_limit.max_attempts, + None, + "interrupted", + ) + .await; return Err(RetryError { - last_error: last_error.unwrap_or(LlmError::Timeout(timeout)), - attempts: attempt, + last_error: error, + attempts, total_time: start_time.elapsed(), }); } - + rate_limit::new_attempt(); match operation().await { - Ok(result) => return Ok(result), + Ok(result) => { + if attempts > 1 + || recovery_started.is_some() + || operation_budget + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .observed_wait + { + rate_limit::observe( + attempts, + config.rate_limit.max_attempts, + None, + "recovered", + ) + .await; + } + return Ok(result); + } Err(err) => { - // Non-retryable errors fail immediately - if !err.is_retryable() { + let remaining = operation_budget + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .started + .map_or( + Duration::from_secs(config.rate_limit.recovery_budget_secs), + |started| { + Duration::from_secs(config.rate_limit.recovery_budget_secs) + .saturating_sub(started.elapsed()) + }, + ); + let err = err.for_recovery_budget(remaining); + let rate_limited = matches!( + err, + LlmError::RateLimited { .. } | LlmError::QuotaWindow { .. } + ); + let mut stop = !err.is_retryable(); + let delay = if rate_limited { + let (started, retry) = { + let mut budget = operation_budget + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + let started = *budget.started.get_or_insert_with(tokio::time::Instant::now); + let retry = budget.rate_retries; + budget.rate_retries = budget.rate_retries.saturating_add(1); + (started, retry) + }; + recovery_started = Some(started); + let local = rate_limit::retry_delay(&config.rate_limit, retry); + let minimum = err.suggested_retry_delay().unwrap_or_default(); + let delay = if minimum.is_zero() { + local + } else { + local.max(minimum).saturating_add(local.div_f64(2.0)) + }; + let remaining = Duration::from_secs(config.rate_limit.recovery_budget_secs) + .saturating_sub(started.elapsed()); + stop |= attempts >= config.rate_limit.max_attempts || delay >= remaining; + delay + } else { + stop |= transport_retries >= config.max_retries; + let delay = config.delay_for_attempt(transport_retries); + transport_retries = transport_retries.saturating_add(1); + delay.max(if config.respect_retry_after { + err.suggested_retry_delay().unwrap_or_default() + } else { + Duration::ZERO + }) + }; + stop |= recovery_started.is_some_and(|started| { + delay + >= Duration::from_secs(config.rate_limit.recovery_budget_secs) + .saturating_sub(started.elapsed()) + }); + stop |= recovery_started.is_some() && attempts >= config.rate_limit.max_attempts; + stop |= config.total_timeout > 0.0 + && start_time.elapsed().saturating_add(delay).as_secs_f64() + >= config.total_timeout; + if stop { + if rate_limited { + rate_limit::record_terminal_quota(&err); + rate_limit::observe( + attempts, + config.rate_limit.max_attempts, + err.suggested_retry_delay(), + "interrupted", + ) + .await; + } else if operation_budget + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .observed_wait + { + rate_limit::observe( + attempts, + config.rate_limit.max_attempts, + None, + "stopped", + ) + .await; + } return Err(RetryError { last_error: err, - attempts: attempt + 1, + attempts, total_time: start_time.elapsed(), }); } - - // Last attempt - no more retries - if attempt >= config.max_retries { - return Err(RetryError { - last_error: err, - attempts: attempt + 1, - total_time: start_time.elapsed(), - }); + if rate_limited { + rate_limit::observe( + attempts, + config.rate_limit.max_attempts, + Some(delay), + "waiting", + ) + .await; } - - // Calculate delay - // Use server's Retry-After if available and configured - let base_delay = config.delay_for_attempt(attempt); - let delay = if config.respect_retry_after { - err.suggested_retry_delay().unwrap_or(base_delay) - } else { - base_delay - }; - - // Notify callback if provided if let Some(ref cb) = callback { - cb(&err, attempt, delay); + cb(&err, attempts - 1, delay); } last_error = Some(err); - - // Wait before retrying - tokio::time::sleep(delay).await; + if rate_limited { + rate_limit::wait(delay).await; + } else { + tokio::time::sleep(delay).await; + } } } } - - // Should not reach here, but handle gracefully - Err(RetryError { - last_error: last_error.unwrap_or(LlmError::Other("Unknown retry error".to_string())), - attempts: config.max_retries + 1, - total_time: start_time.elapsed(), - }) } // === Utility Functions === -/// The longest a `Retry-After` value is ever believed. A server (or a proxy -/// in front of it) can send an arbitrarily large delay; without a ceiling a -/// single `Retry-After: 86400` would wedge the turn for a day. One hour is -/// well past any legitimate rate-limit window. -const RETRY_AFTER_MAX: Duration = Duration::from_secs(3600); +pub fn retry_delay_from_error_body(body: &str) -> Option { + static HINTS: std::sync::OnceLock<[regex::Regex; 2]> = std::sync::OnceLock::new(); + let [structured, prose] = HINTS.get_or_init(|| { + [ + regex::Regex::new(r#"retryDelay\\?"\s*:\s*\\?"\s*(\d+(?:\.\d+)?)\s*s"#).unwrap(), + regex::Regex::new(r"(?i)retry in\s+(\d+(?:\.\d+)?)\s*s").unwrap(), + ] + }); + [structured, prose] + .into_iter() + .flat_map(|pattern| pattern.captures_iter(body)) + .filter_map(|captures| captures.get(1)) + .filter_map(|value| value.as_str().parse::().ok()) + .filter_map(|seconds| Duration::try_from_secs_f64(seconds).ok()) + .max() +} -/// Parses the Retry-After header value into a Duration. -/// -/// Supports both: -/// - Seconds as integer: "120" -> 120 seconds -/// - HTTP-date format: "Wed, 21 Oct 2015 07:28:00 GMT" (not implemented, returns None) -/// -/// The value is server-controlled, so this never panics and never returns an -/// unbounded delay: negative / NaN / infinite / absurd floats are rejected -/// (`Duration::from_secs_f64` panics on a negative — a remote-triggerable -/// crash before this guard), and any result is clamped to [`RETRY_AFTER_MAX`]. +/// Parse numeric seconds or an HTTP-date. Long server minima are preserved; +/// the independent recovery budget parks the operation instead of retrying early. pub fn parse_retry_after(value: &str) -> Option { - // Try parsing as seconds - if let Ok(seconds) = value.parse::() { - return Some(Duration::from_secs(seconds).min(RETRY_AFTER_MAX)); - } - - // Try parsing as float seconds. Only a finite, non-negative value is a - // meaningful delay; everything else (`-5`, `nan`, `inf`) is "no usable - // hint". Clamp to the ceiling BEFORE `from_secs_f64` so an out-of-range - // float can never reach its overflow-panic path, while keeping the - // sub-second precision a legitimate `1.5` carries. - if let Ok(seconds) = value.parse::() - && seconds.is_finite() - && seconds >= 0.0 - { - let clamped = seconds.min(RETRY_AFTER_MAX.as_secs_f64()); - return Some(Duration::from_secs_f64(clamped)); - } - - // HTTP-date format not supported yet - // Could use chrono or httpdate crate if needed - None + rate_limit::server_delay(value) } /// Extracts Retry-After duration from response headers @@ -1430,12 +1607,12 @@ mod tests { assert_eq!(parse_retry_after("-inf"), None); // Absurdly large values clamp to the ceiling rather than overflowing // or wedging the turn for a day. - assert_eq!(parse_retry_after("1e300"), Some(RETRY_AFTER_MAX)); - assert_eq!(parse_retry_after("86400"), Some(RETRY_AFTER_MAX)); + assert_eq!(parse_retry_after("1e300"), None); + assert_eq!(parse_retry_after("86400"), Some(Duration::from_secs(86400))); assert_eq!( parse_retry_after("999999999999"), - Some(RETRY_AFTER_MAX), - "integer path is clamped too" + Some(Duration::from_secs(999999999999)), + "valid long minima are preserved" ); // A normal value still passes through untouched. assert_eq!(parse_retry_after("30"), Some(Duration::from_secs(30))); @@ -1444,6 +1621,7 @@ mod tests { #[test] fn test_retry_policy_conversion() { let policy = RetryPolicy { + rate_limit: RateLimitRetryPolicy::default(), enabled: true, max_retries: 5, initial_delay: 2.0, diff --git a/crates/tui/src/llm_client/rate_limit.rs b/crates/tui/src/llm_client/rate_limit.rs new file mode 100644 index 0000000000..72c146d852 --- /dev/null +++ b/crates/tui/src/llm_client/rate_limit.rs @@ -0,0 +1,722 @@ +//! Recovery policy for a refused model operation. Generation time is not a retry budget. + +use std::sync::{Arc, Mutex}; +use std::time::Duration; + +use super::{LlmError, RateLimitRetryPolicy}; + +#[derive(Clone, Copy)] +pub(crate) enum StreamOpenOutcome { + Pending, + Headers, + RecoverableQuota(Option), +} + +tokio::task_local! { + pub(crate) static STREAM_OPEN_OUTCOME: Arc>; +} + +fn set_stream_open_outcome(outcome: StreamOpenOutcome) { + let _ = STREAM_OPEN_OUTCOME.try_with(|state| { + *state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) = outcome; + }); +} + +/// Record the header boundary of the current wire attempt, if opening a stream. +pub(crate) fn set_response_headers_received(received: bool) { + set_stream_open_outcome(if received { + StreamOpenOutcome::Headers + } else { + StreamOpenOutcome::Pending + }); +} + +pub(crate) fn record_stream_open_quota(retry_after: Option) { + set_stream_open_outcome(StreamOpenOutcome::RecoverableQuota(retry_after)); +} + +pub(crate) fn record_stream_response_error(error: &LlmError) { + match error { + LlmError::RateLimited { retry_after, .. } | LlmError::QuotaWindow { retry_after, .. } => { + record_stream_open_quota(*retry_after); + } + _ => set_stream_open_outcome(StreamOpenOutcome::Headers), + } +} + +/// Preserve the classified result before an observer can block its delivery. +pub(super) fn record_terminal_quota(error: &LlmError) { + if let LlmError::RateLimited { retry_after, .. } | LlmError::QuotaWindow { retry_after, .. } = + error + { + record_stream_open_quota(*retry_after); + } +} + +#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] +pub struct RecoveryReceipt { + #[serde(default, skip_serializing_if = "Option::is_none")] + pub agent_id: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub operation_id: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub turn_id: Option, + pub attempt: u32, + pub max_attempts: u32, + pub next_retry_at_ms: Option, + pub state: String, +} + +pub struct OperationRecovery { + pub operation_id: Option, + pub created: tokio::time::Instant, + pub observed_wait: bool, + pub pending_observation: Option, + pub attempts: u32, + pub wire_attempts: Option, + pub rate_retries: u32, + pub started: Option, +} + +impl Default for OperationRecovery { + fn default() -> Self { + Self { + operation_id: None, + created: tokio::time::Instant::now(), + observed_wait: false, + pending_observation: None, + attempts: 0, + wire_attempts: None, + rate_retries: 0, + started: None, + } + } +} + +tokio::task_local! { + static OPERATION_RECOVERY: std::sync::Arc>; +} + +pub fn with_operation( + budget: std::sync::Arc>, + future: F, +) -> std::pin::Pin>> { + Box::pin(OPERATION_RECOVERY.scope(budget, future)) +} + +pub fn operation() -> std::sync::Arc> { + OPERATION_RECOVERY + .try_with(Clone::clone) + .unwrap_or_default() +} + +tokio::task_local! { + static OBSERVER: RecoveryObserver; +} + +/// Live event observers must attempt their FIFO event send before any other +/// suspension. CancellationObservation relies on that ordering for same-call +/// reopen/timeout retries; the full-channel regression protects this contract. +pub type RecoveryObserver = std::sync::Arc< + dyn Fn(RecoveryReceipt) -> std::pin::Pin + Send>> + + Send + + Sync, +>; + +pub fn with_observer( + observer: RecoveryObserver, + future: F, +) -> std::pin::Pin>> { + Box::pin(OBSERVER.scope(observer, future)) +} + +/// A dropped HTTP/backoff future only emits a terminal observation. It never +/// resumes the model request or owns a task checkpoint. +pub(super) struct CancellationObservation { + budget: std::sync::Arc>, + observer: Option, +} +impl CancellationObservation { + pub(super) fn new(budget: std::sync::Arc>) -> Self { + Self { + budget, + observer: OBSERVER.try_with(Clone::clone).ok(), + } + } +} +impl Drop for CancellationObservation { + fn drop(&mut self) { + let pending = { + let mut budget = self + .budget + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + let mut pending = budget.pending_observation.take(); + if let Some(receipt) = &mut pending { + receipt.attempt = budget + .wire_attempts + .unwrap_or(budget.attempts.max(receipt.attempt)); + } + pending + }; + if let (Some(mut receipt), Some(observer), Ok(runtime)) = ( + pending, + self.observer.as_ref(), + tokio::runtime::Handle::try_current(), + ) { + receipt.state = "cancelled".into(); + receipt.next_retry_at_ms = None; + // Every live event observer starts with the same FIFO channel + // send. Poll before returning to a same-operation retry: it either + // publishes now or reserves its queue position before a newer + // waiting receipt can overtake this terminal observation. + let mut future = observer(receipt); + let waker = futures_util::task::noop_waker_ref(); + if std::future::Future::poll( + future.as_mut(), + &mut std::task::Context::from_waker(waker), + ) + .is_pending() + { + runtime.spawn(future); + } + } + } +} + +pub async fn observe(attempt: u32, max_attempts: u32, delay: Option, state: &str) { + let (attempt, operation_id) = { + let budget = operation(); + let mut budget = budget + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if state == "waiting" { + budget.observed_wait = true; + } + ( + budget.wire_attempts.unwrap_or(attempt), + budget.operation_id.clone(), + ) + }; + let next_retry_at_ms = delay + .and_then(|d| chrono::Duration::from_std(d).ok()) + .and_then(|d| chrono::Utc::now().checked_add_signed(d)) + .map(|d| d.timestamp_millis()); + let receipt = RecoveryReceipt { + agent_id: None, + operation_id, + turn_id: None, + attempt, + max_attempts, + next_retry_at_ms, + state: state.into(), + }; + // A terminal observation remains pending until its async event send has + // completed. Dropping a blocked terminal publication must still settle + // the earlier waiting observation independently of the parent task. + let budget = operation(); + budget + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .pending_observation = Some(receipt.clone()); + if let Ok(emit) = OBSERVER.try_with(Clone::clone) { + emit(receipt).await; + } + if state != "waiting" { + budget + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .pending_observation = None; + } +} + +#[derive(Default)] +struct Clock { + phase: u64, + waited: Duration, + waiting: Option, +} + +tokio::task_local! { + static RECOVERY_CLOCK: std::sync::Arc>; +} + +struct PauseGuard(std::sync::Arc>); +impl Drop for PauseGuard { + fn drop(&mut self) { + let mut clock = self + .0 + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if let Some(started) = clock.waiting.take() { + clock.waited += started.elapsed(); + } + } +} + +pub fn new_attempt() { + let _ = RECOVERY_CLOCK.try_with(|clock| { + let mut clock = clock + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + clock.phase = clock.phase.wrapping_add(1); + }); +} + +pub async fn without_active_time(future: F) -> F::Output { + let _pause = RECOVERY_CLOCK + .try_with(|clock| { + clock + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .waiting = Some(tokio::time::Instant::now()); + PauseGuard(clock.clone()) + }) + .ok(); + future.await +} + +pub async fn wait(delay: Duration) { + without_active_time(tokio::time::sleep(delay)).await; +} + +/// Bound active work while a provider-directed backoff belongs to its own recovery budget. +pub fn active_timeout( + duration: Duration, + future: F, +) -> std::pin::Pin>>> { + active_timeout_policy(duration, future, true) +} + +pub fn active_timeout_total( + duration: Duration, + future: F, +) -> std::pin::Pin>>> { + active_timeout_policy(duration, future, false) +} + +fn active_timeout_policy( + duration: Duration, + future: F, + reset_on_attempt: bool, +) -> std::pin::Pin>>> { + let clock = RECOVERY_CLOCK.try_with(Clone::clone).unwrap_or_default(); + Box::pin(RECOVERY_CLOCK.scope(clock.clone(), async move { + let mut start = tokio::time::Instant::now(); + let mut baseline = clock + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .waited; + let mut phase = clock + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .phase; + tokio::pin!(future); + loop { + let waiting = { + let clock = clock + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if reset_on_attempt && clock.phase != phase { + phase = clock.phase; + start = tokio::time::Instant::now(); + baseline = clock.waited; + } + clock.waited.saturating_sub(baseline) + + clock.waiting.map_or(Duration::ZERO, |t| t.elapsed()) + }; + let remaining = duration.saturating_sub(start.elapsed().saturating_sub(waiting)); + if remaining.is_zero() { + return Err(()); + } + tokio::select! { + biased; + result = &mut future => return Ok(result), + () = tokio::time::sleep(remaining.min(Duration::from_millis(100))) => {} + } + } + })) +} + +pub fn server_delay(value: &str) -> Option { + let value = value.trim(); + if let Ok(seconds) = value.parse::() { + return Some(Duration::from_secs(seconds)); + } + if let Ok(seconds) = value.parse::() { + return Duration::try_from_secs_f64(seconds).ok(); + } + chrono::DateTime::parse_from_rfc2822(value) + .ok() + .map(|date| { + (date.with_timezone(&chrono::Utc) - chrono::Utc::now()) + .to_std() + .unwrap_or_default() + }) +} + +/// Some native providers report logical failures in a successful JSON envelope. +/// SSE bodies remain streaming and are owned by their existing recovery path. +#[cfg(test)] +async fn inspect_json_envelope( + provider: &str, + response: reqwest::Response, +) -> Result { + inspect_json_envelope_with_disclosure(provider, response, None, |status, raw| { + super::sanitize_http_error_body(Some(provider), status, raw) + }) + .await +} + +pub async fn inspect_json_envelope_with_disclosure( + provider: &str, + mut response: reqwest::Response, + body_timeouts: Option<(Duration, Duration)>, + disclose: F, +) -> Result +where + F: Fn(u16, &str) -> String + Send, +{ + if !response.status().is_success() + || !matches!(provider, "minimax" | "minimax-anthropic" | "zai") + || !response + .headers() + .get(reqwest::header::CONTENT_TYPE) + .and_then(|v| v.to_str().ok()) + .is_some_and(|value| { + value + .split(';') + .next() + .is_some_and(|kind| kind.trim().eq_ignore_ascii_case("application/json")) + }) + { + return Ok(response); + } + let status = response.status(); + let version = response.version(); + let headers = response.headers().clone(); + let retry_after = super::extract_retry_after(&headers); + const MAX_ENVELOPE_BYTES: usize = 32 * 1024 * 1024; + if response + .content_length() + .is_some_and(|length| length > MAX_ENVELOPE_BYTES as u64) + { + return Err(LlmError::ParseError( + "Provider JSON envelope exceeds 32 MiB".into(), + )); + } + let mut bytes = Vec::new(); + loop { + let chunk = if let Some((first_byte, idle)) = body_timeouts { + let timeout = if bytes.is_empty() { first_byte } else { idle }; + tokio::time::timeout(timeout, response.chunk()) + .await + .map_err(|_| LlmError::Timeout(timeout))? + } else { + response.chunk().await + } + .map_err(|error| LlmError::from_reqwest(&error))?; + let Some(chunk) = chunk else { break }; + if chunk.len() > MAX_ENVELOPE_BYTES.saturating_sub(bytes.len()) { + return Err(LlmError::ParseError( + "Provider JSON envelope exceeds 32 MiB".into(), + )); + } + bytes.extend_from_slice(&chunk); + } + let raw = String::from_utf8_lossy(&bytes); + if let Ok(body) = serde_json::from_str::(&raw) { + let failed = body + .pointer("/base_resp/status_code") + .and_then(serde_json::Value::as_i64) + .is_some_and(|code| code != 0) + || body.get("error").is_some_and(|error| !error.is_null()); + if failed { + let safe = disclose(status.as_u16(), &raw); + return Err( + provider_error(provider, status.as_u16(), &raw, &safe, retry_after).unwrap_or( + LlmError::InvalidRequest { + status: status.as_u16(), + message: safe, + }, + ), + ); + } + } + let mut rebuilt = http::Response::new(reqwest::Body::from(bytes)); + *rebuilt.status_mut() = status; + *rebuilt.version_mut() = version; + *rebuilt.headers_mut() = headers; + Ok(reqwest::Response::from(rebuilt)) +} + +pub fn body_delay(body: &str) -> Option { + let body: serde_json::Value = serde_json::from_str(body).ok()?; + let objects = [ + &body, + body.get("error").unwrap_or(&body), + body.get("base_resp").unwrap_or(&body), + ]; + let mut delay = None; + for object in objects { + for key in ["retry_after", "retry_after_seconds", "retryDelay"] { + if let Some(value) = object.get(key) { + let text = value + .as_str() + .map(str::to_owned) + .unwrap_or_else(|| value.to_string()); + delay = delay.max(server_delay(text.trim_end_matches('s'))); + } + } + for key in ["next_flush_time", "reset_time", "reset_at"] { + if let Some(value) = object.get(key) { + let text = value + .as_str() + .map(str::to_owned) + .unwrap_or_else(|| value.to_string()); + let reset = text + .parse::() + .ok() + .and_then(|n| { + if n > 10_000_000_000 { + chrono::DateTime::from_timestamp_millis(n) + } else { + chrono::DateTime::from_timestamp(n, 0) + } + }) + .or_else(|| { + chrono::DateTime::parse_from_rfc3339(&text) + .ok() + .map(|d| d.with_timezone(&chrono::Utc)) + }); + delay = + delay.max(reset.map(|d| (d - chrono::Utc::now()).to_std().unwrap_or_default())); + } + } + } + delay +} + +fn normalized_error_code(code: &str) -> String { + code.chars() + .filter(char::is_ascii_alphanumeric) + .map(|ch| ch.to_ascii_lowercase()) + .collect() +} + +fn independent_billing_evidence(body: &serde_json::Value, raw: &str) -> bool { + let coded = [ + "/error/code", + "/error/type", + "/error/error_code", + "/code", + "/type", + "/error_code", + ] + .into_iter() + .filter_map(|p| body.pointer(p).and_then(serde_json::Value::as_str)) + .any(|code| { + normalized_error_code(code) != "insufficientquota" && super::is_explicit_quota_code(code) + }); + // Reuse the legacy durable-message contract, excluding only the ambiguous + // code. Check every supported message field so a gateway's generic class + // cannot hide an explicit billing detail. Decode JSON Unicode escapes too. + coded + || std::iter::once(raw) + .chain( + [ + "/error/message", + "/error", + "/message", + "/detail", + "/error_description", + ] + .into_iter() + .filter_map(|p| body.pointer(p).and_then(serde_json::Value::as_str)), + ) + .any(|message| { + let message = message.to_ascii_lowercase(); + super::has_explicit_quota_phrase(&message) + || [ + "billing hard limit reached", + "billing limit reached", + "billing limit has been reached", + "credit balance too low", + "insufficient credits", + "out of credits", + "insufficient balance", + "insufficient funds", + "payment required", + "subscription expired", + "subscription has expired", + "\u{4f59}\u{989d}\u{4e0d}\u{8db3}", + "\u{4f59}\u{989d}\u{5df2}\u{8017}\u{5c3d}", + "\u{6b20}\u{8d39}", + "\u{8ba2}\u{9605}\u{5df2}\u{8fc7}\u{671f}", + ] + .iter() + .any(|phrase| message.contains(phrase)) + }) +} + +/// Interpret native codes before sanitization; unknown 429s remain bounded. +pub fn provider_error( + provider: &str, + status: u16, + raw: &str, + safe: &str, + hint: Option, +) -> Option { + let hint = hint + .max(body_delay(raw)) + .max(super::retry_delay_from_error_body(raw)); + let body: serde_json::Value = serde_json::from_str(raw).unwrap_or(serde_json::Value::Null); + let code = [ + "/error/code", + "/error/type", + "/error/error_code", + "/base_resp/status_code", + "/code", + "/type", + "/error_code", + ] + .into_iter() + .find_map(|p| body.pointer(p)); + let code = code + .map(|code| { + code.as_str() + .map(str::to_owned) + .unwrap_or_else(|| code.to_string()) + }) + .unwrap_or_default(); + let provider = provider.to_ascii_lowercase(); + let code_lower = code.to_ascii_lowercase(); + let window = match provider.as_str() { + "minimax" | "minimax-anthropic" | "minimax_anthropic" => { + matches!(code.as_str(), "2045" | "2046" | "2047") + } + "glm" | "zai" | "zhipu" => matches!( + code.as_str(), + "1308" | "1310" | "1316" | "1317" | "1318" | "1319" | "1320" | "1321" + ), + _ => false, + }; + if window { + return Some(LlmError::QuotaWindow { + message: safe.to_owned(), + retry_after: hint, + }); + } + if (matches!(status, 400 | 402 | 403 | 429) || (200..300).contains(&status)) + && independent_billing_evidence(&body, raw) + { + return Some(LlmError::QuotaExhausted( + super::QuotaExhaustionError::from_http_message(safe.to_owned()), + )); + } + // This code means billing on OpenAI, but RPM/TPM on other surfaces. + // An unknown relay must not turn that ambiguous code alone into a + // permanent account failure. Independent billing evidence still wins. + let raw_lower = raw.to_ascii_lowercase(); + let marked = raw_lower + .split_once("provider error code:") + .map(|(_, tail)| { + tail.trim_start() + .split(|ch: char| !(ch.is_ascii_alphanumeric() || matches!(ch, '_' | '-'))) + .next() + .unwrap_or_default() + }) + .unwrap_or_default(); + let ambiguous_quota = status == 429 + && !matches!( + provider.as_str(), + "openai" | "openai-codex" | "openai_codex" + ) + && ([ + "/error/code", + "/error/type", + "/error/error_code", + "/code", + "/type", + "/error_code", + ] + .into_iter() + .filter_map(|p| body.pointer(p).and_then(serde_json::Value::as_str)) + .any(|code| normalized_error_code(code) == "insufficientquota") + || normalized_error_code(marked) == "insufficientquota"); + if ambiguous_quota && independent_billing_evidence(&body, raw) { + return Some(LlmError::QuotaExhausted( + super::QuotaExhaustionError::from_http_message(safe.to_owned()), + )); + } + let transient = match provider.as_str() { + "qwen" + | "dashscope" + | "alibaba" + | "modelstudio" + | "modelstudio-token-plan" + | "modelstudio-token-plan-anthropic" + | "modelstudio_token_plan" + | "modelstudio_token_plan_anthropic" => { + status == 429 + && matches!( + code_lower.as_str(), + "insufficient_quota" | "throttling" | "throttling.ratequota" + ) + } + "minimax" | "minimax-anthropic" | "minimax_anthropic" => { + matches!(code.as_str(), "1002" | "1041") + } + "glm" | "zai" | "zhipu" => matches!( + code.as_str(), + "1302" | "1305" | "1308" | "1310" | "1316" | "1317" | "1318" | "1319" | "1320" | "1321" + ), + _ => false, + }; + if transient { + return Some(LlmError::RateLimited { + message: safe.to_owned(), + retry_after: hint, + }); + } + let user_action = match provider.as_str() { + "minimax" | "minimax-anthropic" | "minimax_anthropic" => { + matches!(code.as_str(), "1008" | "1039") + } + "glm" | "zai" | "zhipu" => matches!(code.as_str(), "1113" | "1309"), + _ => false, + }; + if user_action && code != "1039" { + return Some(LlmError::QuotaExhausted( + super::QuotaExhaustionError::from_http_message(safe.to_owned()), + )); + } + if user_action { + return Some(LlmError::InvalidRequest { + status, + message: safe.to_owned(), + }); + } + if ambiguous_quota { + return Some(LlmError::RateLimited { + message: safe.to_owned(), + retry_after: hint, + }); + } + None +} + +pub fn retry_delay(policy: &RateLimitRetryPolicy, retry: u32) -> Duration { + let multiplier = 2u32.saturating_pow(retry.min(31)); + let base = Duration::from_secs(policy.initial_delay_secs) + .saturating_mul(multiplier) + .min(Duration::from_secs(policy.max_delay_secs)); + let bytes = *uuid::Uuid::new_v4().as_bytes(); + let unit = f64::from(u16::from_le_bytes([bytes[0], bytes[1]])) / f64::from(u16::MAX); + base.mul_f64(0.5 + unit * 0.5) +} + +#[cfg(test)] +mod tests; diff --git a/crates/tui/src/llm_client/rate_limit/tests.rs b/crates/tui/src/llm_client/rate_limit/tests.rs new file mode 100644 index 0000000000..da546bc6f4 --- /dev/null +++ b/crates/tui/src/llm_client/rate_limit/tests.rs @@ -0,0 +1,789 @@ +use super::*; +use crate::llm_client::{RetryConfig, RetryResult, with_retry}; +use std::sync::{ + Arc, + atomic::{AtomicU32, Ordering}, +}; + +fn refused(hint: Option) -> LlmError { + LlmError::RateLimited { + message: "fixture refusal".into(), + retry_after: hint, + } +} + +#[test] +fn server_minima_and_provider_codes() { + assert_eq!(server_delay("86400"), Some(Duration::from_secs(86400))); + for invalid in ["-1", "NaN", "inf", "1e300", "bad"] { + assert_eq!(server_delay(invalid), None); + } + let date = chrono::Utc::now() + chrono::Duration::hours(1); + let parsed = server_delay(&date.to_rfc2822()).unwrap(); + assert!(parsed >= Duration::from_secs(3598) && parsed <= Duration::from_secs(3600)); + let raw = r#"{"error":{"code":"insufficient_quota","message":"TPM exceeded"}}"#; + assert!(matches!( + provider_error("dashscope", 429, raw, "safe", None), + Some(LlmError::RateLimited { .. }) + )); + assert!(provider_error("openai", 429, raw, "safe", None).is_none()); + assert!(matches!( + LlmError::from_http_response(429, raw), + LlmError::QuotaExhausted(_) + )); + for (provider, code, expected) in [ + ("minimax", 1002, "retry"), + ("minimax", 1008, "billing"), + ("minimax", 1039, "input"), + ("minimax", 2045, "window"), + ("zai", 1308, "window"), + ("zai", 1302, "retry"), + ] { + let raw = serde_json::json!({"error":{"code":code,"message":"fixture"}}).to_string(); + let error = provider_error(provider, 429, &raw, "safe", None).unwrap(); + assert!( + match expected { + "retry" => matches!(error, LlmError::RateLimited { .. }), + "billing" => matches!(error, LlmError::QuotaExhausted(_)), + "window" => matches!(error, LlmError::QuotaWindow { .. }), + _ => matches!(error, LlmError::InvalidRequest { .. }), + }, + "{provider}/{code}: {error:?}" + ); + } + assert_eq!( + body_delay(r#"{"error":{"retry_after":300}}"#), + Some(Duration::from_secs(300)) + ); + assert_eq!( + body_delay(r#"{"error":{"retryDelay":"120s"}}"#), + Some(Duration::from_secs(120)) + ); + let reset = (chrono::Utc::now() + chrono::Duration::hours(5)).timestamp(); + let hint = + body_delay(&serde_json::json!({"error":{"next_flush_time":reset}}).to_string()).unwrap(); + assert!(hint > Duration::from_secs(17998)); +} + +#[test] +fn policy_validation_and_equal_jitter() { + assert!(serde_json::from_str::(r#"{"max_attempts":0}"#).is_err()); + assert!( + serde_json::from_str::( + r#"{"initial_delay_secs":70,"max_delay_secs":60}"# + ) + .is_err() + ); + let p = RateLimitRetryPolicy::default(); + assert_eq!(p.max_attempts, 20); + assert_eq!(p.recovery_budget_secs, 600); + for n in 0..100 { + let delay = retry_delay(&p, n); + let base = (5u64.saturating_mul(2u64.saturating_pow(n.min(63)))).min(60); + assert!( + delay >= Duration::from_secs(base).div_f64(2.0) && delay <= Duration::from_secs(base) + ); + } +} + +#[tokio::test(start_paused = true)] +async fn rate_budget_is_independent_and_operation_scoped() { + let mut config = RetryConfig::default(); + config.max_retries = 0; + config.rate_limit.max_attempts = 3; + let calls = AtomicU32::new(0); + let budget = Arc::new(std::sync::Mutex::new(OperationRecovery::default())); + let first: RetryResult<()> = with_operation( + budget.clone(), + with_retry( + &config, + || async { + if calls.fetch_add(1, Ordering::SeqCst) == 0 { + Err(refused(None)) + } else { + Ok(()) + } + }, + None, + ), + ) + .await; + assert!(first.is_ok()); + let second: RetryResult<()> = with_operation( + budget, + with_retry( + &config, + || async { + calls.fetch_add(1, Ordering::SeqCst); + Err(refused(None)) + }, + None, + ), + ) + .await; + assert_eq!(second.unwrap_err().attempts, 3); + assert_eq!(calls.load(Ordering::SeqCst), 3); +} + +#[tokio::test(start_paused = true)] +async fn long_hints_and_quota_windows_park_without_early_retry() { + let config = RetryConfig::default(); + for window in [false, true] { + let calls = AtomicU32::new(0); + let start = tokio::time::Instant::now(); + let result: RetryResult<()> = with_retry( + &config, + || async { + calls.fetch_add(1, Ordering::SeqCst); + Err(if window { + LlmError::QuotaWindow { + message: "window".into(), + retry_after: Some(Duration::from_secs(2)), + } + } else { + refused(Some(Duration::from_secs(86400))) + }) + }, + None, + ) + .await; + assert!(result.is_err()); + assert_eq!(calls.load(Ordering::SeqCst), 1); + assert_eq!(start.elapsed(), Duration::ZERO); + } +} + +#[tokio::test(start_paused = true)] +async fn active_work_deadline_excludes_backoff_and_cancellation_drops_retries() { + assert!( + active_timeout(Duration::from_secs(3), async { + wait(Duration::from_secs(30)).await; + tokio::time::sleep(Duration::from_secs(2)).await; + }) + .await + .is_ok() + ); + assert!( + active_timeout( + Duration::from_secs(3), + tokio::time::sleep(Duration::from_secs(5)) + ) + .await + .is_err() + ); + let calls = Arc::new(AtomicU32::new(0)); + let count = calls.clone(); + let task = tokio::spawn(async move { + let _: RetryResult<()> = with_retry( + &RetryConfig::default(), + || async { + count.fetch_add(1, Ordering::SeqCst); + Err(refused(None)) + }, + None, + ) + .await; + }); + tokio::task::yield_now().await; + task.abort(); + let _ = task.await; + tokio::time::advance(Duration::from_secs(1000)).await; + assert_eq!(calls.load(Ordering::SeqCst), 1); +} + +#[tokio::test] +async fn native_json_envelopes_are_classified_before_model_parsing() { + let make = |body: &str| { + reqwest::Response::from( + http::Response::builder() + .header("content-type", "application/json") + .body(reqwest::Body::from(body.to_owned())) + .unwrap(), + ) + }; + let error = inspect_json_envelope( + "minimax", + make(r#"{"base_resp":{"status_code":1002,"status_msg":"throttled"}}"#), + ) + .await + .unwrap_err(); + assert!(matches!(error, LlmError::RateLimited { .. })); + let body = r#"{"base_resp":{"status_code":0},"choices":[]}"#; + let response = inspect_json_envelope("minimax", make(body)).await.unwrap(); + assert_eq!(response.text().await.unwrap(), body); +} + +#[tokio::test(start_paused = true)] +async fn successful_json_refusal_preserves_header_minimum_and_disclosure() { + let response = reqwest::Response::from(http::Response::builder() + .header("content-type", "Application/JSON; charset=utf-8") + .header("retry-after", "3600") + .body(reqwest::Body::from(r#"{"base_resp":{"status_code":1002,"status_msg":"fixture-secret"},"retry_after":7200}"#)).unwrap()); + let error = inspect_json_envelope_with_disclosure("minimax", response, None, |_, body| { + body.replace("fixture-secret", "[redacted]") + }) + .await + .unwrap_err(); + assert_eq!( + error.suggested_retry_delay(), + Some(Duration::from_secs(7200)) + ); + assert!(!error.to_string().contains("fixture-secret")); + // A larger/header-only minimum must also survive; the body-max case + // above alone would not detect a dropped Retry-After header. + let response = reqwest::Response::from( + http::Response::builder() + .header("content-type", "application/json") + .header("retry-after", "3600") + .body(reqwest::Body::from( + r#"{"base_resp":{"status_code":1002,"status_msg":"throttled"}}"#, + )) + .unwrap(), + ); + let header_only = inspect_json_envelope("minimax", response) + .await + .unwrap_err(); + assert_eq!( + header_only.suggested_retry_delay(), + Some(Duration::from_secs(3600)) + ); + let config = RetryConfig::default(); + let mut refusal = Some(error); + let calls = AtomicU32::new(0); + let result: RetryResult<()> = with_retry( + &config, + || { + calls.fetch_add(1, Ordering::SeqCst); + let error = refusal + .take() + .expect("long hint must not dispatch a second attempt"); + async move { Err(error) } + }, + None, + ) + .await; + assert_eq!(result.unwrap_err().attempts, 1); + assert_eq!(calls.load(Ordering::SeqCst), 1); +} + +#[tokio::test(start_paused = true)] +async fn long_backoff_does_not_consume_the_nonstream_active_envelope() { + let mut config = RetryConfig::default(); + config.rate_limit.recovery_budget_secs = 3600; + let calls = AtomicU32::new(0); + let result = active_timeout_total( + Duration::from_secs(1800), + with_retry( + &config, + || async { + if calls.fetch_add(1, Ordering::SeqCst) < 2 { + Err(refused(Some(Duration::from_secs(1000)))) + } else { + Ok(()) + } + }, + None, + ), + ) + .await; + assert!(result.unwrap().is_ok()); + assert_eq!(calls.load(Ordering::SeqCst), 3); + // Excluding waits must not disable a real active-work timeout. + assert!( + active_timeout_total( + Duration::from_secs(3), + tokio::time::sleep(Duration::from_secs(4)) + ) + .await + .is_err() + ); +} + +#[tokio::test(start_paused = true)] +async fn mixed_transport_failure_cannot_restart_expired_rate_recovery() { + let mut config = RetryConfig::default(); + config.rate_limit.initial_delay_secs = 1; + config.rate_limit.max_delay_secs = 1; + let calls = AtomicU32::new(0); + let result: RetryResult<()> = with_retry( + &config, + || async { + if calls.fetch_add(1, Ordering::SeqCst) == 0 { + Err(refused(None)) + } else { + tokio::time::sleep(Duration::from_secs(601)).await; + Err(LlmError::ServerError { + status: 503, + message: "fixture".into(), + }) + } + }, + None, + ) + .await; + assert!(matches!( + result.unwrap_err().last_error, + LlmError::ServerError { status: 503, ref message } if message == "fixture" + )); + assert_eq!(calls.load(Ordering::SeqCst), 2); +} + +#[tokio::test(start_paused = true)] +async fn cancelled_auxiliary_wait_emits_its_own_terminal_receipt() { + let receipts = std::sync::Arc::new(std::sync::Mutex::new(Vec::new())); + let seen = receipts.clone(); + let terminal = std::sync::Arc::new(tokio::sync::Notify::new()); + let notified = terminal.clone(); + let observer: RecoveryObserver = std::sync::Arc::new(move |receipt| { + let seen = seen.clone(); + let notified = notified.clone(); + Box::pin(async move { + let cancelled = receipt.state == "cancelled"; + seen.lock().unwrap().push(receipt); + if cancelled { + notified.notify_one(); + } + }) + }); + let budget = std::sync::Arc::new(std::sync::Mutex::new(OperationRecovery { + operation_id: Some("cancelled-auxiliary".into()), + ..Default::default() + })); + with_operation( + budget, + with_observer(observer, async { + let config = RetryConfig::default(); + let mut pending = Box::pin(with_retry( + &config, + || async { Err::<(), _>(refused(None)) }, + None, + )); + assert!(futures_util::poll!(pending.as_mut()).is_pending()); + drop(pending); + tokio::time::timeout(Duration::from_secs(1), terminal.notified()) + .await + .unwrap(); + }), + ) + .await; + let receipts = receipts.lock().unwrap(); + assert_eq!(receipts.len(), 2); + assert_eq!(receipts[0].state, "waiting"); + assert_eq!(receipts[1].state, "cancelled"); + assert_eq!( + receipts[1].operation_id.as_deref(), + Some("cancelled-auxiliary") + ); + assert_eq!(receipts[1].attempt, 1); + assert!(receipts[1].next_retry_at_ms.is_none()); +} + +#[tokio::test(flavor = "current_thread", start_paused = true)] +async fn cancelled_observation_cannot_overtake_same_operation_reopen() { + let (tx, mut rx) = tokio::sync::mpsc::channel(1); + let observer: RecoveryObserver = std::sync::Arc::new(move |receipt| { + let tx = tx.clone(); + Box::pin(async move { + tx.send(receipt).await.unwrap(); + }) + }); + let budget = std::sync::Arc::new(std::sync::Mutex::new(OperationRecovery { + operation_id: Some("same-operation".into()), + ..Default::default() + })); + with_operation( + budget.clone(), + with_observer(observer, async { + let guard = CancellationObservation::new(budget.clone()); + observe(1, 20, Some(Duration::from_secs(5)), "waiting").await; + budget.lock().unwrap().wire_attempts = Some(2); + drop(guard); + budget.lock().unwrap().wire_attempts = Some(3); + let mut newer = Box::pin(observe(3, 20, Some(Duration::from_secs(5)), "waiting")); + assert!(futures_util::poll!(newer.as_mut()).is_pending()); + let receive = async { + let mut events = Vec::new(); + for _ in 0..3 { + events.push( + tokio::time::timeout(Duration::from_secs(1), rx.recv()) + .await + .unwrap() + .unwrap(), + ); + } + events + }; + let (events, ()) = tokio::join!(receive, newer); + assert_eq!( + events + .iter() + .map(|r| (r.state.as_str(), r.attempt)) + .collect::>(), + [("waiting", 1), ("cancelled", 2), ("waiting", 3)] + ); + assert!( + events + .iter() + .all(|r| r.operation_id.as_deref() == Some("same-operation")) + ); + }), + ) + .await; +} + +#[tokio::test(flavor = "current_thread", start_paused = true)] +async fn cancelled_terminal_publication_settles_prior_waiting() { + let (tx, mut rx) = tokio::sync::mpsc::channel(1); + let observer: RecoveryObserver = std::sync::Arc::new(move |receipt| { + let tx = tx.clone(); + Box::pin(async move { + tx.send(receipt).await.unwrap(); + }) + }); + let budget = std::sync::Arc::new(std::sync::Mutex::new(OperationRecovery { + operation_id: Some("blocked-terminal".into()), + ..Default::default() + })); + with_operation( + budget.clone(), + with_observer(observer, async { + let guard = CancellationObservation::new(budget); + observe(1, 20, Some(Duration::from_secs(5)), "waiting").await; + let mut terminal = Box::pin(observe(2, 20, None, "recovered")); + assert!(futures_util::poll!(terminal.as_mut()).is_pending()); + drop(terminal); + drop(guard); + assert_eq!(rx.recv().await.unwrap().state, "waiting"); + let stopped = tokio::time::timeout(Duration::from_secs(1), rx.recv()) + .await + .unwrap() + .unwrap(); + assert_eq!(stopped.state, "cancelled"); + assert_eq!(stopped.attempt, 2); + assert_eq!(stopped.operation_id.as_deref(), Some("blocked-terminal")); + }), + ) + .await; +} + +#[test] +fn unknown_rate_limit_quota_code_requires_independent_billing_evidence() { + let hint = Some(Duration::from_secs(30)); + for raw in [ + r#"{"error":{"code":null,"type":"insufficient_quota","message":"TPM exceeded"},"retry_after":40}"#, + r#"{"error":{"code":"gateway_limit","error_code":"insufficient_quota","message":"TPM exceeded"},"retry_after":40}"#, + r#"{"error":{"code":"insufficient_quota","message":"TPM exceeded"},"retry_after":40}"#, + r#"{"error":{"type":"insufficient_quota","message":"quota exceeded"},"retry_after":40}"#, + r#"{"error_code":"insufficient_quota","message":"quota exceeded","retry_after":40}"#, + ] { + let error = provider_error("custom", 429, raw, "safe", hint).unwrap(); + assert!(matches!(error, LlmError::RateLimited { .. })); + assert_eq!(error.suggested_retry_delay(), Some(Duration::from_secs(40))); + assert!(!error.to_string().contains("TPM")); + } + let raw = "TPM exceeded (provider error code: insufficient_quota)"; + assert!(matches!( + provider_error("custom", 429, raw, "safe", hint), + Some(LlmError::RateLimited { .. }) + )); + for raw in [ + r#"{"error":{"code":"insufficient_quota","message":"Credit balance exhausted"}}"#, + r#"{"error":{"code":"insufficient_quota","message":"Billing quota has been exhausted"}}"#, + r#"{"error":{"code":"insufficient_quota","message":"Account quota exhausted"}}"#, + r#"{"error":{"code":"insufficient_quota","message":"Plan quota is exhausted"}}"#, + r#"{"error":{"code":"insufficient_quota","message":"TPM exceeded"},"detail":"Billing quota has been exhausted"}"#, + r#"{"error":{"code":"insufficient_quota","message":"You exceeded your current quota. Check plan and billing"}}"#, + r#"{"error":{"code":"insufficient_quota","type":"billing_hard_limit_reached","message":"fixture"}}"#, + r#"{"error":{"code":"insufficient_quota","message":"\u4f59\u989d\u4e0d\u8db3"}}"#, + ] { + for provider in ["custom", "dashscope", "qwen"] { + assert!(matches!( + provider_error(provider, 429, raw, "safe", hint), + Some(LlmError::QuotaExhausted(_)) + )); + } + } + let raw = r#"{"error":{"code":"insufficient_quota","message":"TPM exceeded"}}"#; + assert!(provider_error("openai", 429, raw, "safe", hint).is_none()); + assert!(matches!( + LlmError::from_http_response(429, raw), + LlmError::QuotaExhausted(_) + )); + assert!(provider_error("custom", 401, raw, "safe", hint).is_none()); +} + +#[test] +fn rate_limit_native_minima_and_billing_keep_safe_classification() { + for raw in [ + r#"{"base_resp":{"status_code":1002,"status_msg":"retry in 7200s"}}"#, + r#"{"base_resp":{"status_code":1002},"error":{"details":[{"retryDelay":"7200s"}]}}"#, + ] { + let error = + provider_error("minimax", 429, raw, "safe", Some(Duration::from_secs(30))).unwrap(); + assert_eq!( + error.suggested_retry_delay(), + Some(Duration::from_secs(7200)) + ); + } + let raw = r#"{"base_resp":{"status_code":2045,"status_msg":"weekly plan quota exhausted"},"retry_after":7200}"#; + assert!(matches!(provider_error("minimax", 200, raw, "safe", None), + Some(LlmError::QuotaWindow { retry_after: Some(delay), .. }) if delay == Duration::from_secs(7200))); + for status in [200, 429] { + let raw = r#"{"error":{"message":"opaque"},"detail":"Billing quota has been exhausted; fixture-secret"}"#; + let error = provider_error("custom", status, raw, "safe", None).unwrap(); + assert!(matches!(error, LlmError::QuotaExhausted(_))); + assert!(!error.to_string().contains("fixture-secret")); + } + assert_eq!( + super::super::retry_delay_from_error_body("retry in 30s; retry in 7200s"), + Some(Duration::from_secs(7200)) + ); +} + +#[tokio::test(start_paused = true)] +async fn rate_limit_long_minimum_is_terminal_with_retries_enabled_or_disabled() { + for enabled in [false, true] { + let config = RetryConfig { + enabled, + ..Default::default() + }; + let calls = AtomicU32::new(0); + let result: RetryResult<()> = with_retry( + &config, + || async { + calls.fetch_add(1, Ordering::SeqCst); + Err(refused(Some(Duration::from_secs(7200)))) + }, + None, + ) + .await; + assert!( + matches!(result.unwrap_err().last_error, LlmError::QuotaWindow { retry_after: Some(delay), .. } if delay == Duration::from_secs(7200)) + ); + assert_eq!(calls.load(Ordering::SeqCst), 1); + } +} + +#[tokio::test(start_paused = true)] +async fn rate_limit_mixed_failure_preserves_last_transport_error() { + let mut config = RetryConfig::default(); + config.rate_limit.initial_delay_secs = 1; + config.rate_limit.max_delay_secs = 1; + config.rate_limit.recovery_budget_secs = 3; + config.initial_delay = 5.0; + config.max_delay = 5.0; + config.jitter = false; + let calls = AtomicU32::new(0); + let result: RetryResult<()> = with_retry( + &config, + || async { + if calls.fetch_add(1, Ordering::SeqCst) == 0 { + Err(refused(Some(Duration::from_secs(2)))) + } else { + Err(LlmError::ServerError { + status: 503, + message: "actual transport failure".into(), + }) + } + }, + None, + ) + .await; + assert!( + matches!(result.unwrap_err().last_error, LlmError::ServerError { status: 503, ref message } if message == "actual transport failure") + ); + assert_eq!(calls.load(Ordering::SeqCst), 2); +} + +#[tokio::test(start_paused = true)] +async fn rate_limit_receipts_cover_ordinary_recovery_and_jitter_above_minimum() { + let receipts = Arc::new(std::sync::Mutex::new(Vec::new())); + let events = receipts.clone(); + let observer: RecoveryObserver = Arc::new(move |receipt| { + let events = events.clone(); + Box::pin(async move { + events.lock().unwrap().push(receipt); + }) + }); + let budget = Arc::new(std::sync::Mutex::new(OperationRecovery::default())); + let config = RetryConfig::default(); + let calls = AtomicU32::new(0); + let result: RetryResult<()> = with_operation( + budget, + with_observer( + observer, + with_retry( + &config, + || async { + if calls.fetch_add(1, Ordering::SeqCst) == 0 { + Err(LlmError::ServerError { + status: 503, + message: "fixture".into(), + }) + } else { + Ok(()) + } + }, + None, + ), + ), + ) + .await; + result.unwrap(); + { + let events = receipts.lock().unwrap(); + assert_eq!(events.len(), 1); + assert_eq!(events[0].state, "recovered"); + assert_eq!(events[0].attempt, 2); + } + let calls = AtomicU32::new(0); + let start = tokio::time::Instant::now(); + let result: RetryResult<()> = with_retry( + &config, + || async { + if calls.fetch_add(1, Ordering::SeqCst) == 0 { + Err(refused(Some(Duration::from_secs(30)))) + } else { + Ok(()) + } + }, + None, + ) + .await; + result.unwrap(); + assert!(start.elapsed() > Duration::from_secs(30)); + assert!(start.elapsed() <= Duration::from_millis(32500)); +} + +#[tokio::test(start_paused = true)] +async fn native_json_envelope_bounds_first_byte_idle_and_size() { + use futures_util::StreamExt; + for sent in [false, true] { + let prefix = futures_util::stream::iter(if sent { + vec![Ok::<_, std::io::Error>("{")] + } else { + vec![] + }); + let body = reqwest::Body::wrap_stream(prefix.chain(futures_util::stream::pending())); + let response = reqwest::Response::from( + http::Response::builder() + .header("content-type", "application/json") + .body(body) + .unwrap(), + ); + let error = inspect_json_envelope_with_disclosure( + "minimax", + response, + Some((Duration::from_secs(2), Duration::from_secs(3))), + |_, _| "safe".into(), + ) + .await + .unwrap_err(); + assert!( + matches!(error, LlmError::Timeout(delay) if delay == Duration::from_secs(if sent {3} else {2})) + ); + } + for declared in [false, true] { + let body = reqwest::Body::wrap_stream(futures_util::stream::iter( + (0..33).map(|_| Ok::<_, std::io::Error>(vec![b' '; 1024 * 1024])), + )); + let mut response = http::Response::builder().header("content-type", "application/json"); + if declared { + response = response.header("content-length", 33 * 1024 * 1024); + } + let response = reqwest::Response::from(response.body(body).unwrap()); + let error = inspect_json_envelope("minimax", response) + .await + .unwrap_err(); + assert!(matches!(error, LlmError::ParseError(ref message) if message.contains("32 MiB"))); + } +} + +#[tokio::test(start_paused = true)] +async fn rate_limit_small_minimum_preserves_local_backoff_floor() { + let mut config = RetryConfig::default(); + config.rate_limit.initial_delay_secs = 60; + config.rate_limit.max_delay_secs = 60; + let calls = AtomicU32::new(0); + let start = tokio::time::Instant::now(); + let result: RetryResult<()> = with_retry( + &config, + || async { + if calls.fetch_add(1, Ordering::SeqCst) == 0 { + Err(refused(Some(Duration::from_secs(1)))) + } else { + Ok(()) + } + }, + None, + ) + .await; + result.unwrap(); + assert_eq!(calls.load(Ordering::SeqCst), 2); + assert!(start.elapsed() >= Duration::from_secs(45)); + assert!(start.elapsed() <= Duration::from_secs(90)); +} + +#[tokio::test(start_paused = true)] +async fn rate_limit_ordinary_backoff_counts_toward_active_deadline() { + let config = RetryConfig::default() + .with_max_retries(1) + .with_initial_delay(5.0) + .with_jitter(false); + let calls = AtomicU32::new(0); + let result = active_timeout_total( + Duration::from_secs(2), + with_retry( + &config, + || async { + calls.fetch_add(1, Ordering::SeqCst); + Err::<(), _>(LlmError::ServerError { + status: 503, + message: "fixture".into(), + }) + }, + None, + ), + ) + .await; + assert!(result.is_err()); + assert_eq!(calls.load(Ordering::SeqCst), 1); +} + +#[tokio::test(start_paused = true)] +async fn rate_limit_budget_expiring_during_sleep_preserves_actual_error() { + let mut config = RetryConfig::default(); + config.rate_limit.initial_delay_secs = 1; + config.rate_limit.max_delay_secs = 1; + config.initial_delay = 1.0; + config.jitter = false; + let budget = Arc::new(std::sync::Mutex::new(OperationRecovery::default())); + let calls = AtomicU32::new(0); + let mut pending = Box::pin(with_operation( + budget, + with_retry( + &config, + || async { + if calls.fetch_add(1, Ordering::SeqCst) == 0 { + Err(refused(None)) + } else { + Err::<(), _>(LlmError::ServerError { + status: 503, + message: "actual delayed failure".into(), + }) + } + }, + None, + ), + )); + assert!(futures_util::poll!(pending.as_mut()).is_pending()); + tokio::time::advance(Duration::from_secs(2)).await; + assert!(futures_util::poll!(pending.as_mut()).is_pending()); + assert_eq!(calls.load(Ordering::SeqCst), 2); + tokio::time::advance(Duration::from_secs(601)).await; + let result = pending.await; + assert!( + matches!(result.unwrap_err().last_error, LlmError::ServerError { status: 503, ref message } if message == "actual delayed failure") + ); + assert_eq!(calls.load(Ordering::SeqCst), 2); +} diff --git a/crates/tui/src/retry_status.rs b/crates/tui/src/retry_status.rs index 622160f09f..aad350c0ab 100644 --- a/crates/tui/src/retry_status.rs +++ b/crates/tui/src/retry_status.rs @@ -99,16 +99,6 @@ fn with_state(f: impl FnOnce(&mut RetryState) -> R) -> R { f(&mut state) } -#[cfg(not(test))] -fn with_rate_limit(f: impl FnOnce(&mut Option) -> R) -> R { - static STATE: OnceLock>> = OnceLock::new(); - let mut state = STATE - .get_or_init(|| Mutex::new(None)) - .lock() - .unwrap_or_else(|error| error.into_inner()); - f(&mut state) -} - /// Under test, this state is per-thread. /// /// Production has exactly one foreground engine per process, so a global is the @@ -155,6 +145,7 @@ pub fn snapshot() -> RetryState { /// Extend the provider-wide rate-limit pause window. This is separate from /// the footer banner so one successful concurrent request cannot clear another /// request's active `Retry-After` window. +#[cfg(test)] pub fn note_rate_limit(delay: Duration) { let deadline = Instant::now() + delay; with_rate_limit(|current| { @@ -164,7 +155,8 @@ pub fn note_rate_limit(delay: Duration) { }); } -/// Remaining provider-wide rate-limit pause, if any. +/// Remaining synthetic pause used by legacy retry-banner fixtures. +#[cfg(test)] #[must_use] pub fn rate_limit_remaining() -> Option { let now = Instant::now(); diff --git a/crates/tui/src/test_support.rs b/crates/tui/src/test_support.rs index 4620e34291..8e0324a978 100644 --- a/crates/tui/src/test_support.rs +++ b/crates/tui/src/test_support.rs @@ -638,3 +638,23 @@ mod tests { writer.join().expect("writer thread"); } } + +/// Construct and drive Engine-sized fixtures with the production worker stack budget. +pub(crate) fn run_on_engine_thread(build: F) +where + F: FnOnce() -> Fut + Send + 'static, + Fut: std::future::Future, +{ + std::thread::Builder::new() + .stack_size(crate::CODEWHALE_MAIN_STACK_BYTES) + .spawn(move || { + tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap() + .block_on(build()); + }) + .unwrap() + .join() + .unwrap(); +} diff --git a/crates/tui/src/tools/subagent/governor.rs b/crates/tui/src/tools/subagent/governor.rs index 6a1f149e24..fcc910b060 100644 --- a/crates/tui/src/tools/subagent/governor.rs +++ b/crates/tui/src/tools/subagent/governor.rs @@ -85,166 +85,7 @@ pub(crate) fn rate_limit_retry_delay(retry_number: u32) -> Duration { Duration::from_secs_f64(base * (1.0 - RATE_LIMIT_BACKOFF_JITTER_FACTOR * random_unit_factor())) } -// === DynamicGate === - -#[derive(Debug)] -struct GateState { - capacity: usize, - outstanding: usize, -} - -/// A launch gate with runtime-adjustable capacity (see module docs). -/// -/// A thin accounting layer over [`tokio::sync::Semaphore`]: the semaphore -/// owns the FIFO wait queue and cancellation safety, this struct owns the -/// `capacity`/`outstanding` ledger that lets capacity drop below the -/// number of outstanding holders. The ledger — not the semaphore token -/// count — is the admission authority: a granted token is registered only -/// while `outstanding < capacity`; one granted while the ledger is full is -/// surplus and is destroyed (`forget`) instead of admitted. Two paths -/// create such tokens while no slot is free, both invisible to -/// `forget_permits` (which sees free tokens only): a grant that races a -/// capacity change (handed out before the shrink can forget it), and a -/// grant handed to a queued waiter that is cancelled before polling again -/// (tokio re-injects it into the pool). Surplus tokens are harmless — they -/// cannot admit — and the next admission attempt destroys them; every slot -/// freed below capacity mints a fresh token on release, so no slot is ever -/// lost. Once no surplus token remains and every in-flight admission has -/// registered, `sem.available_permits() == -/// capacity.saturating_sub(outstanding)` holds exactly. -/// -/// `acquire` returns a [`DynamicGatePermit`] whose `Drop` returns the slot. -/// Reducing capacity below `outstanding` is allowed: the surplus holders -/// finish naturally and no new permit is granted until the outstanding -/// count drops under the new capacity. -#[derive(Debug)] -pub(crate) struct DynamicGate { - sem: tokio::sync::Semaphore, - inner: Mutex, -} - -impl DynamicGate { - pub(crate) fn new(capacity: usize) -> Self { - let capacity = capacity.max(1); - Self { - sem: tokio::sync::Semaphore::new(capacity), - inner: Mutex::new(GateState { - capacity, - outstanding: 0, - }), - } - } - - pub(crate) fn capacity(&self) -> usize { - self.inner.lock().expect("launch gate poisoned").capacity - } - - /// Free semaphore permits right now — `capacity - outstanding` at - /// rest, plus any surplus tokens not yet destroyed by an admission - /// attempt (see the struct docs). Diagnostics and tests only; racy by - /// design. - #[cfg_attr(not(test), allow(dead_code))] - pub(crate) fn available_permits(&self) -> usize { - self.sem.available_permits() - } - - /// Adjust the gate capacity. Raising it adds the new free headroom to - /// the semaphore immediately (queued acquirers wake in FIFO order); - /// lowering it forgets free slots — holders above the new capacity keep - /// running, and their releases are absorbed instead of re-admitted. - /// Grants already handed to queued waiters are invisible to the shrink; - /// if their acquirers are cancelled, admission destroys the returned - /// surplus token (see the struct docs). - pub(crate) fn set_capacity(&self, capacity: usize) { - let mut inner = self.inner.lock().expect("launch gate poisoned"); - let old = inner.capacity; - inner.capacity = capacity; - if capacity < old { - // `forget_permits` caps the reduction at the available count, so - // an over-subscribed shrink (capacity < outstanding) simply - // drains free slots to zero. - self.sem.forget_permits(old - capacity); - } else { - let free = |cap: usize| cap.saturating_sub(inner.outstanding); - self.sem.add_permits(free(capacity) - free(old)); - } - } - - /// Try to acquire a permit without waiting. - pub(crate) fn try_acquire(self: &std::sync::Arc) -> Option { - let sem_permit = self.sem.try_acquire().ok()?; - self.admit_grant(sem_permit) - } - - /// Register a granted semaphore token as a held slot. The ledger, not - /// the token count, is the admission authority: a token granted while - /// the ledger is full is surplus drift — a cancelled waiter's - /// re-injected grant, or a grant that raced a shrink — and is destroyed - /// with `forget` (never dropped: a drop would re-inject it). Since - /// `DynamicGatePermit::drop` mints one fresh token per slot freed below - /// capacity, destroying a token never loses a slot. - fn admit_grant( - self: &std::sync::Arc, - sem_permit: tokio::sync::SemaphorePermit<'_>, - ) -> Option { - let mut inner = self.inner.lock().expect("launch gate poisoned"); - if inner.outstanding < inner.capacity { - inner.outstanding += 1; - // Slot ownership moves from the semaphore permit into - // `DynamicGatePermit::drop`. - sem_permit.forget(); - return Some(DynamicGatePermit { - gate: std::sync::Arc::clone(self), - }); - } - sem_permit.forget(); - None - } - - /// Acquire a permit, waiting until capacity is available. The semaphore - /// queue is FIFO and cancel safe: a cancelled future dequeues itself and - /// never swallows a slot or loses a wakeup. A grant that turns out to be - /// surplus (see [`Self::admit_grant`]) is destroyed and the wait - /// continues for the release that mints a fresh one; each surplus round - /// consumes exactly one pool token, so the loop cannot spin. - pub(crate) async fn acquire(self: &std::sync::Arc) -> DynamicGatePermit { - loop { - // The semaphore is never closed, so `acquire` cannot fail. - let sem_permit = self - .sem - .acquire() - .await - .expect("launch gate semaphore closed"); - if let Some(permit) = self.admit_grant(sem_permit) { - return permit; - } - } - } -} - -/// One held launch slot. Released on drop. -pub(crate) struct DynamicGatePermit { - gate: std::sync::Arc, -} - -impl std::fmt::Debug for DynamicGatePermit { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("DynamicGatePermit").finish() - } -} - -impl Drop for DynamicGatePermit { - fn drop(&mut self) { - let mut inner = self.gate.inner.lock().expect("launch gate poisoned"); - inner.outstanding = inner.outstanding.saturating_sub(1); - if inner.outstanding < inner.capacity { - // Return the slot; a queued acquirer wakes in FIFO order. - self.gate.sem.add_permits(1); - } - // Otherwise the capacity was shrunk below the outstanding count: - // the release is absorbed so the surplus drain does not over-admit. - } -} +pub(crate) use crate::core::capacity::{DynamicGate, DynamicGatePermit}; // === RateLimitGovernor === diff --git a/crates/tui/src/tools/subagent/mod.rs b/crates/tui/src/tools/subagent/mod.rs index 45e96577cf..8cc4e96f01 100644 --- a/crates/tui/src/tools/subagent/mod.rs +++ b/crates/tui/src/tools/subagent/mod.rs @@ -2046,7 +2046,7 @@ impl SubAgentTerminalDeliveryContext { /// Publish to every live sink without blocking or awaiting while the /// manager owns the terminal claim. The public agent/worker states remain - /// Running until all three sends have been attempted. + /// Running until every terminal sink has been attempted. fn deliver(&self, result: &SubAgentResult) { let report_ref = spill_subagent_final_report(&self.session_id, result); let completion = subagent_completion_from_result_with_ref_for_session( @@ -2071,6 +2071,26 @@ impl SubAgentTerminalDeliveryContext { id: result.agent_id.clone(), result: completion.payload, }); + // The parent mailbox may already be sealed for detached/off-turn + // work. Publish a typed terminal status on the independent event + // channel even when cancellation aborted the worker future. + let status = worker_status_from_subagent_status(&result.status); + let event = Event::AgentProgress { + owner_session_id: self.session_id.clone(), + id: result.agent_id.clone(), + status: agent_worker_status_name(status).into(), + activity: AgentProgressEventMeta::new(status).with_step(result.steps_taken), + parent_run_id: result.parent_run_id.clone(), + spawn_depth: result.spawn_depth, + }; + if let Err(mpsc::error::TrySendError::Full(event)) = event_tx.try_send(event) + && let Ok(runtime) = tokio::runtime::Handle::try_current() + { + let tx = event_tx.clone(); + runtime.spawn(async move { + let _ = tx.send(event).await; + }); + } } } } @@ -10494,14 +10514,6 @@ async fn spawn_subagent_from_input( ))); } - if let Some(remaining) = crate::retry_status::rate_limit_remaining() { - let seconds = remaining.as_secs() + u64::from(remaining.subsec_nanos() > 0); - return Err(ToolError::execution_failed(format!( - "Provider is rate-limiting; sub-agent spawning is paused for {seconds}s. \ - Wait for the current backoff window before starting new agent work." - ))); - } - let mut child_runtime = if spawn_request.detached { runtime.background_runtime() } else { @@ -10541,6 +10553,20 @@ async fn spawn_subagent_from_input( { child_runtime.client = rebound; } + if let Some(remaining) = child_runtime + .client + .provider_quota_pause_remaining() + .await + .map_err(|err| ToolError::execution_failed(err.to_string()))? + { + let seconds = remaining + .as_secs() + .saturating_add(u64::from(remaining.subsec_nanos() > 0)); + return Err(ToolError::execution_failed(format!( + "Provider is rate-limiting; sub-agent spawning is paused for {seconds}s. \ + Wait for the current backoff window before starting new agent work." + ))); + } child_runtime.reasoning_effort = route.reasoning_effort.clone(); child_runtime.reasoning_effort_auto = false; let model_route = route.model_route; @@ -12377,11 +12403,27 @@ async fn request_subagent_model_response_with_retries( if let Some(governor) = runtime.governor.as_ref() { governor.record_attempt(Instant::now()); } - match tokio::time::timeout( - runtime.step_api_timeout, - runtime - .client - .create_message_for_operation(request.clone(), operation_id), + let event_tx = runtime.event_tx.clone(); + let recovery_agent_id = agent_id.to_owned(); + let observer: crate::llm_client::rate_limit::RecoveryObserver = + std::sync::Arc::new(move |mut receipt| { + receipt.agent_id = Some(recovery_agent_id.clone()); + receipt.operation_id = Some(operation_id.to_string()); + let event_tx = event_tx.clone(); + Box::pin(async move { + if let Some(tx) = event_tx { + let _ = tx.send(Event::ModelRecovery { receipt }).await; + } + }) + }); + match crate::llm_client::rate_limit::with_observer( + observer, + crate::llm_client::rate_limit::active_timeout( + runtime.step_api_timeout, + runtime + .client + .create_message_for_operation(request.clone(), operation_id), + ), ) .await { @@ -12402,6 +12444,18 @@ async fn request_subagent_model_response_with_retries( { governor.record_rate_limited(Instant::now()); } + // The transport owns the complete 429 budget. An outer replay would reset it. + if matches!( + err.downcast_ref::(), + Some(LlmError::RateLimited { .. } | LlmError::QuotaWindow { .. }) + ) { + return Err(SubAgentApiRequestFailure::Interrupted { + reason: format!( + "Rate-limit recovery stopped: {err}; checkpoint preserved for continuation" + ), + checkpoint_reason: "api_rate_limited", + }); + } let retry_number = transient_failures.saturating_add(1); let Some(retryable) = retryable_subagent_provider_failure(&err, retry_number) else { @@ -17091,10 +17145,24 @@ impl SubAgentToolRegistry { .context() .clone() .with_owner_agent(self.owner_agent_id.clone(), self.owner_agent_name.clone()); - let result = self - .registry - .execute_rich_full_with_context(name, input, Some(&context)) - .await; + let event_tx = self.gate_runtime.event_tx.clone(); + let agent_id = self.owner_agent_id.clone(); + let observer: crate::llm_client::rate_limit::RecoveryObserver = + Arc::new(move |mut receipt| { + receipt.agent_id = Some(agent_id.clone()); + let event_tx = event_tx.clone(); + Box::pin(async move { + if let Some(tx) = event_tx { + let _ = tx.send(Event::ModelRecovery { receipt }).await; + } + }) + }); + let result = crate::llm_client::rate_limit::with_observer( + observer, + self.registry + .execute_rich_full_with_context(name, input, Some(&context)), + ) + .await; if let Some(guard) = reservation.take() { // Stamp the reservation with the owner's still-running shell // jobs: a background job keeps it alive through the shell @@ -17550,7 +17618,7 @@ fn build_allowed_tools( /// failures at a glance. fn subagent_failure_message(err: &anyhow::Error) -> String { let class = match err.downcast_ref::() { - Some(LlmError::RateLimited { .. }) => Some("rate_limited"), + Some(LlmError::RateLimited { .. } | LlmError::QuotaWindow { .. }) => Some("rate_limited"), Some(LlmError::QuotaExhausted(_)) => Some("quota_exhausted"), Some(LlmError::ServerError { .. }) => Some("server"), Some(LlmError::NetworkError(_)) | Some(LlmError::Timeout(_)) => Some("network"), diff --git a/crates/tui/src/tools/subagent/tests.rs b/crates/tui/src/tools/subagent/tests.rs index 6c89e4085c..d7a9d9c356 100644 --- a/crates/tui/src/tools/subagent/tests.rs +++ b/crates/tui/src/tools/subagent/tests.rs @@ -2526,6 +2526,7 @@ async fn always_rate_limited_chat_client() -> (DeepSeekClient, Arc) api_key: Some("test-key".to_string()), base_url: Some(format!("http://{addr}/v1")), retry: Some(crate::config::RetryConfig { + rate_limit: crate::config::RateLimitRetryPolicy::default(), enabled: Some(false), max_retries: Some(0), initial_delay: Some(0.0), @@ -2576,6 +2577,7 @@ async fn always_invalid_request_chat_client() -> (DeepSeekClient, Arc(); let (mailbox, mut mailbox_rx) = Mailbox::new(CancellationToken::new()); let (event_tx, mut event_rx) = mpsc::channel(8); + if event_queue_full { + for _ in 0..8 { + event_tx + .try_send(Event::status("fixture queue occupancy")) + .unwrap(); + } + } let mut runtime = runtime_with_depth(1, Some(completion_tx)); + if sealed_parent { + mailbox.seal(); + } runtime.mailbox = Some(mailbox); runtime.event_tx = Some(event_tx); agent.terminal_delivery = Some(SubAgentTerminalDeliveryContext::from_runtime(&runtime)); @@ -6598,16 +6620,33 @@ async fn model_wait_cancel_fans_in_once_and_preserves_checkpoint() { ) }) .collect::>(); - assert_eq!(terminal_mail.len(), 1); - assert!(matches!( - terminal_mail[0].message, - MailboxMessage::Cancelled { ref agent_id } if agent_id == "agent_model_wait_cancel" - )); + assert_eq!(terminal_mail.len(), usize::from(!sealed_parent)); + if !sealed_parent { + assert!(matches!( + terminal_mail[0].message, + MailboxMessage::Cancelled { ref agent_id } if agent_id == "agent_model_wait_cancel" + )); + } - let complete_events = std::iter::from_fn(|| event_rx.try_recv().ok()) - .filter(|event| matches!(event, Event::AgentComplete { .. })) - .count(); - assert_eq!(complete_events, 1); + let mut events: Vec<_> = std::iter::from_fn(|| event_rx.try_recv().ok()).collect(); + if event_queue_full { + events.push( + tokio::time::timeout(Duration::from_secs(1), event_rx.recv()) + .await + .expect("queued metadata terminal must be delivered") + .unwrap(), + ); + } + assert_eq!( + events + .iter() + .filter(|event| matches!(event, Event::AgentComplete { .. })) + .count(), + usize::from(!event_queue_full) + ); + assert_eq!(events.iter().filter(|event| matches!(event, + Event::AgentProgress { id, activity, .. } if id == &agent_id && activity.worker_status == AgentWorkerStatus::Cancelled)).count(), 1, + "a typed terminal event survives parent mailbox sealing and worker abortion"); let worker = manager.get_worker_record(&agent_id).expect("worker record"); assert_eq!(worker.status, AgentWorkerStatus::Cancelled); assert_eq!( @@ -9755,90 +9794,105 @@ async fn subagent_retries_transient_provider_header_timeout_before_succeeding() assert_eq!(snapshot.result.as_deref(), Some("recovered answer")); } -#[tokio::test] -async fn subagent_rate_limit_exhaustion_interrupts_with_checkpoint() { - let tmp = tempdir().expect("tempdir"); - let manager = Arc::new(RwLock::new(SubAgentManager::new( - tmp.path().to_path_buf(), - 2, - ))); - let agent_id = "agent_rate_limited_checkpoint".to_string(); - let (task_input_tx, task_input_rx) = mpsc::unbounded_channel(); - let agent = SubAgent::new( - agent_id.clone(), - FleetRole::Worker, - "Inspect rate-limit recovery".to_string(), - make_assignment(), - "deepseek-v4-flash".to_string(), - Some("Blue".to_string()), - Some(vec![]), - task_input_tx, - tmp.path().to_path_buf(), - "boot_test".to_string(), - ); - { - let mut manager = manager.write().await; - manager.agents.insert(agent_id.clone(), agent); - manager.register_worker(make_worker_spec(&agent_id, tmp.path().to_path_buf())); - } - - let (client, calls) = always_rate_limited_chat_client().await; - let mut runtime = stub_runtime().with_step_api_timeout(Duration::from_secs(5)); - runtime.client = client; - runtime.manager = Arc::clone(&manager); - runtime.context = ToolContext::new(tmp.path()); +#[test] +fn subagent_rate_limit_exhaustion_interrupts_with_checkpoint() { + // Exercise the full worker fixture with the product Engine thread's 16 MiB stack. + std::thread::Builder::new() + .stack_size(crate::CODEWHALE_MAIN_STACK_BYTES) + .spawn(|| { + tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap() + .block_on(async { + let tmp = tempdir().expect("tempdir"); + let manager = Arc::new(RwLock::new(SubAgentManager::new( + tmp.path().to_path_buf(), + 2, + ))); + let agent_id = "agent_rate_limited_checkpoint".to_string(); + let (task_input_tx, task_input_rx) = mpsc::unbounded_channel(); + let agent = SubAgent::new( + agent_id.clone(), + FleetRole::Worker, + "Inspect rate-limit recovery".to_string(), + make_assignment(), + "deepseek-v4-flash".to_string(), + Some("Blue".to_string()), + Some(vec![]), + task_input_tx, + tmp.path().to_path_buf(), + "boot_test".to_string(), + ); + { + let mut manager = manager.write().await; + manager.agents.insert(agent_id.clone(), agent); + manager + .register_worker(make_worker_spec(&agent_id, tmp.path().to_path_buf())); + } - let task = SubAgentTask { - manager_handle: Arc::clone(&manager), - runtime, - agent_id: agent_id.clone(), - agent_type: FleetRole::Worker, - prompt: "Inspect rate-limit recovery".to_string(), - assignment: make_assignment(), - allowed_tools: Some(vec![]), - fork_context: false, - started_at: Instant::now(), - max_steps: 3, - token_budget: None, - wall_time: DEFAULT_CHILD_WALL_TIME, - input_rx: task_input_rx, - launch_gate: None, - _foreground_child_registration: None, - }; + let (client, calls) = always_rate_limited_chat_client().await; + let mut runtime = stub_runtime().with_step_api_timeout(Duration::from_secs(5)); + runtime.client = client; + runtime.manager = Arc::clone(&manager); + runtime.context = ToolContext::new(tmp.path()); + + let task = SubAgentTask { + manager_handle: Arc::clone(&manager), + runtime, + agent_id: agent_id.clone(), + agent_type: FleetRole::Worker, + prompt: "Inspect rate-limit recovery".to_string(), + assignment: make_assignment(), + allowed_tools: Some(vec![]), + fork_context: false, + started_at: Instant::now(), + max_steps: 3, + token_budget: None, + wall_time: DEFAULT_CHILD_WALL_TIME, + input_rx: task_input_rx, + launch_gate: None, + _foreground_child_registration: None, + }; - tokio::time::timeout( - Duration::from_secs(10), - tokio::spawn(run_subagent_task(task)), - ) - .await - .expect("sub-agent task should finish") - .expect("sub-agent join should succeed"); + tokio::time::timeout( + Duration::from_secs(10), + tokio::spawn(run_subagent_task(task)), + ) + .await + .expect("sub-agent task should finish") + .expect("sub-agent join should succeed"); - assert_eq!( - calls.load(Ordering::SeqCst), - SUBAGENT_TRANSIENT_PROVIDER_MAX_RETRIES.saturating_add(1) as usize, - "rate-limit retries should be owned by the sub-agent retry loop" - ); - let snapshot = { - let manager = manager.read().await; - manager - .get_result(&agent_id) - .expect("agent should stay registered") - }; - let SubAgentStatus::Interrupted(reason) = &snapshot.status else { - panic!("expected interrupted sub-agent, got {:?}", snapshot.status); - }; - assert!( - reason.contains("rate-limited provider response"), - "reason should name the provider rate limit: {reason}" - ); - let checkpoint = snapshot - .checkpoint - .as_ref() - .expect("rate-limit interruption should preserve checkpoint"); - assert_eq!(checkpoint.reason, "api_rate_limited"); - assert!(checkpoint.continuable); - assert!(snapshot.needs_input.is_some()); + assert_eq!( + calls.load(Ordering::SeqCst), + 1, + "disabled transport retries must not be restarted by the child wrapper" + ); + let snapshot = { + let manager = manager.read().await; + manager + .get_result(&agent_id) + .expect("agent should stay registered") + }; + let SubAgentStatus::Interrupted(reason) = &snapshot.status else { + panic!("expected interrupted sub-agent, got {:?}", snapshot.status); + }; + assert!( + reason.contains("Rate-limit recovery stopped"), + "reason should name the provider rate limit: {reason}" + ); + let checkpoint = snapshot + .checkpoint + .as_ref() + .expect("rate-limit interruption should preserve checkpoint"); + assert_eq!(checkpoint.reason, "api_rate_limited"); + assert!(checkpoint.continuable); + assert!(snapshot.needs_input.is_some()); + }); + }) + .unwrap() + .join() + .unwrap(); } #[tokio::test] @@ -12497,19 +12551,32 @@ fn clamp_child_max_spawn_depth_enforces_absolute_ceiling() { } #[tokio::test] -#[allow(clippy::await_holding_lock)] async fn rate_limit_pause_blocks_subagent_spawn() { - let _guard = crate::retry_status::test_guard(); - // Drop-clear the window even if an assertion below panics: this state is - // process-global, and a leaked 30s pause strands every concurrently - // running test whose worker issues a model request. - let _clear = ClearRateLimitOnDrop; - crate::retry_status::clear(); - crate::retry_status::clear_rate_limit(); - crate::retry_status::note_rate_limit(Duration::from_secs(30)); - let tmp = tempdir().expect("tempdir"); + let server = wiremock::MockServer::start().await; + wiremock::Mock::given(wiremock::matchers::method("POST")) + .respond_with(wiremock::ResponseTemplate::new(429).insert_header("Retry-After", "7200")) + .expect(1) + .mount(&server) + .await; + let config = crate::config::Config { + api_key: Some("test-key".to_string()), + base_url: Some(server.uri()), + rate_limit_recovery_dir: Some(tmp.path().join("quota")), + ..Default::default() + }; let mut runtime = stub_runtime(); + runtime.client = DeepSeekClient::new(&config).expect("fixture client"); + runtime.api_config = Some(Arc::new(config)); + let error = runtime + .client + .send_json_with_retry(&server.uri(), &json!({})) + .await + .unwrap_err(); + assert!(matches!( + error.downcast_ref::(), + Some(crate::llm_client::LlmError::QuotaWindow { .. }) + )); runtime.context = ToolContext::new(tmp.path().to_path_buf()); let manager = new_shared_subagent_manager(tmp.path().to_path_buf(), 2); @@ -17244,6 +17311,7 @@ async fn tool_call_then_invalid_request_chat_client() -> (DeepSeekClient, Arc {} EngineEvent::MessageStarted { .. } => { // Assistant text starting after parallel tool work // means the tool group is done. Flush the active diff --git a/crates/tui/src/vision/tools.rs b/crates/tui/src/vision/tools.rs index 1653e437c6..f58554d3be 100644 --- a/crates/tui/src/vision/tools.rs +++ b/crates/tui/src/vision/tools.rs @@ -10,30 +10,19 @@ use serde_json::{Value, json}; use crate::client::DeepSeekClient; use crate::config::ApiProvider; use crate::config::VisionModelConfig; -use crate::llm_client::{LlmError, RetryConfig, sanitize_http_error_body, with_retry}; +use crate::llm_client::rate_limit::active_timeout_total; use crate::tools::spec::{ ToolCapability, ToolContext, ToolError, ToolResult, ToolSpec, required_str, }; pub struct ImageAnalyzeTool { config: VisionModelConfig, - client: reqwest::Client, route_client: Option, } -/// Total envelope for one image_analyze call, retry attempts and response -/// body consumption included. reqwest's `read_timeout` is *not* a per-read -/// idle bound for the request phase: its timer starts at `send()` and is -/// never reset until the response headers arrive, so it silently acts as a -/// total deadline on the multi-MB upload plus the full non-streaming vision -/// generation — precisely the healthy work a 120s cap used to kill. The -/// client therefore bounds only the connect handshake, and this envelope -/// (the same 30-minute budget as non-streaming model requests) is the sole -/// total bound; a stalled connection errors out through it instead of -/// hanging. The envelope is shared across retry attempts, not per attempt: -/// an attempt that fails slowly consumes its share of the one budget, so -/// only retries of fast failures (connect errors, quick 5xx) fit before -/// the shared envelope expires. +/// Active-work envelope for image upload, generation and response consumption. +/// Native HTTP transport owns per-attempt bounds and quota admission. Provider +/// backoff and shared admission waits use their separate recovery clock. const VISION_REQUEST_ENVELOPE: Duration = Duration::from_secs(1800); fn vision_request_envelope() -> Duration { @@ -56,18 +45,8 @@ impl ImageAnalyzeTool { config: VisionModelConfig, route_client: Option, ) -> Self { - let client = crate::tls::reqwest_client_builder() - // Bound only the connect handshake. A client- or request-level - // `read_timeout` would start counting at `send()` and never - // reset before the response headers, quietly re-introducing a - // total deadline on the upload + long non-streaming generation; - // the total bound lives in VISION_REQUEST_ENVELOPE instead. - .connect_timeout(Duration::from_secs(10)) - .build() - .expect("Failed to build HTTP client"); Self { config, - client, route_client, } } @@ -308,82 +287,26 @@ impl ToolSpec for ImageAnalyzeTool { let url = format!("{}/chat/completions", self.base_url()); let api_key = self.api_key(); - let operation_id = uuid::Uuid::new_v4(); - // The route client can be the parent's main model while vision has a - // separate provider. Only the exact matching route may contribute - // its frozen context headers or logical-call opt-in. - let idempotency_route = self - .route_client - .as_ref() - .filter(|route| route.idempotency_matches_route(&self.base_url(), &api_key)); - - let retry_config = RetryConfig { - max_retries: 3, - initial_delay: 1.0, - max_delay: 30.0, - enabled: true, - ..Default::default() - }; - let response_json = tokio::time::timeout(vision_request_envelope(), async { - // Acquire inside the envelope: the ownership window is part of - // the bounded call, so a contested permit cannot park this tool - // (and its share of Runtime Chat availability) outside every - // total bound. The guard is still held for the whole request. + let route = DeepSeekClient::auxiliary_json_route( + self.route_client.as_ref(), + &self.base_url(), + &api_key, + &self.config.model, + ) + .map_err(|e| { + ToolError::execution_failed(format!("Failed to configure vision route: {e}")) + })?; + let response_json = active_timeout_total(vision_request_envelope(), async { + // Ownership remains inside the tool's active envelope. The native + // HTTP adapter owns quota admission and releases slots in backoff. let _inference = match self.route_client.as_ref() { Some(client) => client.acquire_remote_control_inference_permit().await, None => Some(crate::client::acquire_remote_control_inference_participant().await), }; - let response = with_retry( - &retry_config, - || { - let client = idempotency_route - .map(|route| route.http_client.clone()) - .unwrap_or_else(|| self.client.clone()); - let url = url.clone(); - let api_key = api_key.clone(); - let payload = payload.clone(); - async move { - let builder = client - .post(&url) - .header("Content-Type", "application/json") - .header("Authorization", format!("Bearer {api_key}")) - .json(&payload); - let builder = match idempotency_route { - Some(route) => { - route.with_operation_header_for_id(builder, operation_id) - } - None => builder, - }; - let response = builder - .send() - .await - .map_err(|e| LlmError::from_reqwest(&e))?; - - let status = response.status(); - if !status.is_success() { - let error_text = response - .text() - .await - .unwrap_or_else(|_| "Unknown error".to_string()); - let error_text = sanitize_http_error_body( - Some("Vision provider"), - status.as_u16(), - &error_text, - ); - return Err(LlmError::from_http_response(status.as_u16(), &error_text)); - } - Ok(response) - } - }, - None, - ) - .await - .map_err(|e| ToolError::execution_failed(format!("Vision API request failed: {e}")))?; - - let json: Value = response.json().await.map_err(|e| { - ToolError::execution_failed(format!("Failed to parse response: {e}")) - })?; - Ok(json) + route + .request_auxiliary_json(&url, &payload, uuid::Uuid::new_v4()) + .await + .map_err(|e| ToolError::execution_failed(format!("Vision API request failed: {e}"))) }) .await .map_err(|_| ToolError::Timeout { @@ -864,7 +787,7 @@ mod tests { } #[test] - fn image_dimensions_omits_bmp_because_decoder_feature_is_off() { + fn image_dimensions_follows_available_bmp_decoder() { // Minimal well-formed 1x1 24-bit BMP: even valid BMP input must // degrade to `None` (no panic, no error) while the `bmp` feature of // the `image` dependency is not enabled. @@ -890,7 +813,11 @@ mod tests { let dir = tempdir().expect("tempdir"); let path = dir.path().join("tiny.bmp"); std::fs::write(&path, MINIMAL_BMP).expect("write fixture"); - assert_eq!(ImageAnalyzeTool::image_dimensions(&path), None); + let expected = + image::ImageFormat::Bmp + .reading_enabled() + .then_some((1, 1, String::from("bmp"))); + assert_eq!(ImageAnalyzeTool::image_dimensions(&path), expected); } #[tokio::test] @@ -1015,7 +942,7 @@ mod tests { "owned-context".into(), )])); let route = DeepSeekClient::new(&config).unwrap(); - assert!(!route.idempotency_matches_route(&server.uri(), key)); + assert!(!route.matches_auxiliary_chat_route(&server.uri(), key)); let tool = ImageAnalyzeTool::new_with_route_client( VisionModelConfig { model: "test-vision-model".into(), @@ -1033,7 +960,7 @@ mod tests { } let receipts = server.received_requests().await.unwrap(); assert_eq!(receipts.len(), 3); - for receipt in &receipts { + for (index, receipt) in receipts.iter().enumerate() { assert!(receipt.url.path().ends_with("/chat/completions")); for header in [ "idempotency-key", @@ -1047,15 +974,22 @@ mod tests { "incompatible parent header {header}" ); } - assert!( - receipt - .headers - .get("authorization") - .unwrap() - .to_str() - .unwrap() - .starts_with("Bearer") - ); + if index == 1 { + assert!( + !receipt.headers.contains_key("authorization"), + "an explicit empty key stays unauthenticated" + ); + } else { + assert!( + receipt + .headers + .get("authorization") + .unwrap() + .to_str() + .unwrap() + .starts_with("Bearer ") + ); + } } } @@ -1100,7 +1034,7 @@ mod tests { config.request_idempotency_header = Some("idempotency-key".into()); let route = DeepSeekClient::new(&config).unwrap(); assert_eq!( - route.idempotency_matches_route(&base, "owned-test-key"), + route.matches_auxiliary_chat_route(&base, "owned-test-key"), matching ); let tool = ImageAnalyzeTool::new_with_route_client( @@ -1130,6 +1064,68 @@ mod tests { } } + #[tokio::test] + async fn forkguard_rate_limit_vision_shared_policy_and_long_minimum() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .respond_with(ResponseTemplate::new(429).insert_header("retry-after", "3600")) + .expect(1) + .mount(&server) + .await; + let workspace = tempdir().unwrap(); + std::fs::write( + workspace.path().join("tiny.png"), + create_test_png(2, 2, [1, 2, 3, 255]), + ) + .unwrap(); + let context = ToolContext::new(workspace.path()); + let base = format!("{}/v1", server.uri()); + let route = DeepSeekClient::new(&crate::config::Config { + provider: Some("deepseek".into()), + api_key: Some("owned-test-key".into()), + base_url: Some(base.clone()), + http_headers: Some(std::collections::HashMap::from([( + "OpenAI-Project".into(), + "vision-project-fixture".into(), + )])), + ..Default::default() + }) + .unwrap(); + let tool = ImageAnalyzeTool::new_with_route_client( + VisionModelConfig { + model: "test-vision-model".into(), + api_key: Some("owned-test-key".into()), + base_url: Some(base.clone()), + }, + Some(route.clone()), + ); + let error = tool + .execute(json!({"image_path":"tiny.png"}), &context) + .await + .unwrap_err(); + assert!(error.to_string().contains("Rate limit")); + // The matching foreground client sees the same minimum before wire dispatch. + let error = route + .request_auxiliary_json( + &format!("{base}/chat/completions"), + &json!({}), + uuid::Uuid::new_v4(), + ) + .await + .unwrap_err(); + assert!(matches!( + error.downcast_ref::(), + Some(crate::llm_client::LlmError::QuotaWindow { .. }) + )); + let requests = server.received_requests().await.unwrap(); + assert_eq!(requests.len(), 1); + assert_eq!( + requests[0].headers.get("openai-project").unwrap(), + "vision-project-fixture" + ); + assert!(!requests[0].headers.contains_key("idempotency-key")); + } + #[tokio::test] async fn forkguard_model_operation_vision_http_retry_keeps_id() { use std::sync::Arc; diff --git a/docs/rate-limit-recovery.md b/docs/rate-limit-recovery.md new file mode 100644 index 0000000000..67151d7939 --- /dev/null +++ b/docs/rate-limit-recovery.md @@ -0,0 +1,13 @@ +# Model rate-limit recovery + +The transport policy uses 20 total attempts including the first and a 600-second recovery budget. Local exponential equal jitter starts from a 5-second base and stops growing at 60 seconds. With a nonzero server minimum, the larger of the local delay and the minimum receives an additional half-local jitter; requests never retry before that minimum. These defaults are engineering limits, not provider guarantees. + +Known quota windows and server minima that cannot fit the remaining operation budget return `QuotaWindow`, including with transport retries disabled. They do not enter short automatic retries. Actual transport failures remain typed when a preceding rate-limit episode exhausts the remaining budget. Ordinary transport recovery also emits the existing recovered receipt. + +MiniMax/Zai successful JSON envelopes are inspected before model parsing. Reads are limited to 32 MiB and bounded by the caller's stream-open first-byte deadline and configured stream idle deadline. Existing per-request totals still apply. SSE bodies remain with their existing stream adapters. Native classification combines header, structured JSON, prose, and nested retry hints, and valid JSON string values in the display copy are decoded and redacted before sanitization, with final exact-value redaction afterward. Known native quota windows take precedence over generic billing prose. + +The fork additionally owns shared effective-credential/account admission, stable logical-operation budgets, suspendable leases, persisted server reset times, and root/owned-child interruption. A long server reset remains persisted advice for later continuation; it is not shortened to the upstream transport's temporary pause cap. Local token-bucket waits and ordinary network backoff count toward active-work deadlines; provider-directed backoff and shared quota waits use the existing recovery clock. + +A transport error does not prove a checkpoint was saved. Checkpoint preservation and continued execution belong to the existing Engine and child lifecycle. No new scheduler, tool replay, credential resolver, or task UI protocol is introduced by these fixes. The older fork does not have the newer upstream HTTP-attempt transmission ledger or quota/account-guidance types, so those APIs are not imported independently. + +The branch still requires integration against its current target and applicable CI. Local loopback tests do not establish live-provider behavior, full host restart behavior, or non-Windows execution.