diff --git a/crates/tui/src/lib.rs b/crates/tui/src/lib.rs index 73e83b19c9..406c3adf9a 100644 --- a/crates/tui/src/lib.rs +++ b/crates/tui/src/lib.rs @@ -10481,6 +10481,7 @@ async fn build_direct_workflow_tool( } else { None }; + let fleet_governor = manager.read().await.rate_limit_governor(); let runtime = SubAgentRuntime::new( client, route.model.clone(), @@ -10489,6 +10490,7 @@ async fn build_direct_workflow_tool( Some(event_tx), manager.clone(), ) + .with_fleet_governor(fleet_governor) .with_locale_tag( crate::localization::resolve_locale( &crate::settings::Settings::load_persisted() diff --git a/crates/tui/src/tools/subagent/governor.rs b/crates/tui/src/tools/subagent/governor.rs new file mode 100644 index 0000000000..e9bb5320d5 --- /dev/null +++ b/crates/tui/src/tools/subagent/governor.rs @@ -0,0 +1,782 @@ +//! Rate-limit aware adaptive scheduling for sub-agent fan-out ("swarm mode"). +//! +//! A swarm can launch an unbounded number of sub-agents against one shared +//! LLM provider, so parallel 429s are the steady state rather than an edge +//! case. This module gives the sub-agent module two cooperating pieces: +//! +//! 1. [`DynamicGate`] — a launch gate with a *dynamically adjustable +//! capacity*. The previous gate was a `tokio::sync::Semaphore`, whose +//! capacity is fixed at construction; the only way to "shrink" it was to +//! replace the `Arc`, which silently fails while any child still holds a +//! permit (that is exactly why `update_runtime_limits` only applied +//! launch-concurrency changes when no sub-agent was running). A +//! custom gate can drop its capacity below the number of active holders: +//! existing children keep running to completion, while new admissions +//! block until `active < capacity`. +//! +//! 2. [`RateLimitGovernor`] — a sliding-window observer fed by the sub-agent +//! LLM call path. Every rate-limited attempt and every successful attempt +//! is reported; when the recent failure rate crosses a threshold the +//! governor shrinks the gate (multiplicative decrease), and under a +//! sustained burst it pauses new admissions entirely. Sustained success +//! recovers capacity additively (AIMD), which converges without the +//! oscillation a symmetric controller would show. +//! +//! Retries themselves stay in the LLM call path (see +//! `request_subagent_model_response_with_retries`): the governor never +//! delays an in-flight call, it only decides whether *new* launches may be +//! admitted. `QuotaExhausted` is deliberately not reported — quota is a +//! billing condition, not a transient throttle, and must keep following the +//! existing fatal/checkpoint path. + +use std::collections::VecDeque; +use std::sync::Mutex; +use std::time::{Duration, Instant}; + +use tokio::sync::oneshot; + +/// Observation window for rate-limit events. Events older than this are +/// pruned on every governor interaction. +const RATE_LIMIT_WINDOW: Duration = Duration::from_secs(60); + +/// Rate-limit events inside [`RATE_LIMIT_WINDOW`] at which the governor +/// starts shrinking launch concurrency (AIMD multiplicative decrease). +const THROTTLE_EVENT_THRESHOLD: usize = 2; + +/// Recent rate-limit *ratio* (limited attempts / attempts) at which the +/// governor also shrinks launch concurrency, even below the absolute count +/// threshold. With very few in-flight calls, two 429s may be 100% of traffic. +const THROTTLE_RATIO_THRESHOLD: f64 = 0.3; + +/// Rate-limit events inside the window at which the governor pauses new +/// admissions entirely (gate capacity 0). Held permits are unaffected. +const PAUSE_EVENT_THRESHOLD: usize = 4; + +/// Successful attempts required to add one unit of launch capacity back +/// (AIMD additive increase). Successes are counted per gate-holder, so a +/// shrunken fleet still recovers at a controlled pace. +const SUCCESS_PER_INCREASE_STEP: u32 = 3; + +/// Full-jitter exponential backoff for a rate-limited sub-agent API attempt +/// (`retry_number` is 1-based): the raw backoff is +/// `initial * 2^(n-1)` capped at [`RATE_LIMIT_MAX_BACKOFF`], and the actual +/// delay is drawn uniformly from `[0, backoff)` (AWS "full jitter"). Full +/// jitter de-synchronizes a fan-out of children that were all 429'd by the +/// same provider response; the cap keeps a retrying child inside its +/// wall-time budget instead of giving up. +const RATE_LIMIT_MAX_BACKOFF: Duration = Duration::from_secs(120); +const RATE_LIMIT_BACKOFF_JITTER_FACTOR: f64 = 1.0; // full jitter + +/// Uniformly random factor in `[0, 1)` derived from UUID v4 entropy, the +/// same idiom as `llm_client::RetryConfig::delay_for_attempt`. +fn random_unit_factor() -> f64 { + let bytes = *uuid::Uuid::new_v4().as_bytes(); + let sample = u16::from_le_bytes([bytes[0], bytes[1]]); + f64::from(sample) / f64::from(u16::MAX) +} + +/// Raw (pre-jitter) exponential backoff for a rate-limited attempt. +fn rate_limit_backoff_base(retry_number: u32) -> Duration { + let multiplier = 1u32 + .checked_shl(retry_number.saturating_sub(1)) + .unwrap_or(u32::MAX); + Duration::from_millis(250) + .saturating_mul(multiplier) + .min(RATE_LIMIT_MAX_BACKOFF) +} + +/// Full-jitter retry delay for a rate-limited attempt. +pub(crate) fn rate_limit_retry_delay(retry_number: u32) -> Duration { + let base = rate_limit_backoff_base(retry_number).as_secs_f64(); + // Full jitter: uniform in [0, base). Reaching exactly `base` is fine and + // only sharpens de-synchronization; the draw can never exceed it. + Duration::from_secs_f64(base * (1.0 - RATE_LIMIT_BACKOFF_JITTER_FACTOR * random_unit_factor())) +} + +// === DynamicGate === + +#[derive(Debug)] +struct GateWaiter { + sender: oneshot::Sender, +} + +#[derive(Debug)] +struct GateInner { + capacity: usize, + active: usize, + waiters: VecDeque, +} + +/// A launch gate with runtime-adjustable capacity (see module docs). +/// +/// `acquire` returns a [`DynamicGatePermit`] whose `Drop` releases the slot +/// and wakes one waiter. Reducing capacity below `active` is allowed: the +/// surplus holders finish naturally and no new permit is granted until the +/// active count drops under the new capacity. +/// +/// Waiters receive an *already granted* permit through a oneshot channel, so +/// a waiter future that is cancelled after the grant is dispatched simply +/// drops the permit, whose `Drop` hands the slot to the next waiter. (A +/// wake-and-recheck design would lose that wakeup — the cancelled waiter +/// never re-checks, and with no remaining holders there is no later release +/// to re-dispatch it.) +#[derive(Debug)] +pub(crate) struct DynamicGate { + inner: Mutex, +} + +impl DynamicGate { + pub(crate) fn new(capacity: usize) -> Self { + Self { + inner: Mutex::new(GateInner { + capacity: capacity.max(1), + active: 0, + waiters: VecDeque::new(), + }), + } + } + + pub(crate) fn capacity(&self) -> usize { + self.inner.lock().expect("launch gate poisoned").capacity + } + + /// Free admission slots right now (`capacity - active`). Diagnostics and + /// tests only; racy by design. + #[cfg_attr(not(test), allow(dead_code))] + pub(crate) fn available_permits(&self) -> usize { + let inner = self.inner.lock().expect("launch gate poisoned"); + inner.capacity.saturating_sub(inner.active) + } + + /// Adjust the gate capacity. Raising it grants queued waiters the new + /// headroom immediately; lowering it simply stops new admissions until + /// the active count drains below the new capacity. + pub(crate) fn set_capacity(self: &std::sync::Arc, capacity: usize) { + let mut inner = self.inner.lock().expect("launch gate poisoned"); + inner.capacity = capacity; + Self::wake_locked(self, &mut inner); + } + + /// Grant queued waiters while there is headroom. Called with the lock + /// held; each waiter receives an already-counted permit, so a cancelled + /// receiver's permit is disarmed (never `Drop`ped under the lock) and the + /// slot flows to the next waiter. + fn wake_locked(gate: &std::sync::Arc, inner: &mut GateInner) { + while inner.active < inner.capacity { + let Some(waiter) = inner.waiters.pop_front() else { + break; + }; + let permit = DynamicGatePermit { + gate: Some(std::sync::Arc::clone(gate)), + }; + match waiter.sender.send(permit) { + Ok(()) => inner.active += 1, + Err(mut returned) => { + // The waiter future was cancelled before receiving the + // grant. Disarm instead of dropping: `Drop` would call + // `release()` and re-enter the lock we are holding. + let _ = returned.disarm(); + } + } + } + } + + fn release(self: &std::sync::Arc) { + let mut inner = self.inner.lock().expect("launch gate poisoned"); + inner.active = inner.active.saturating_sub(1); + Self::wake_locked(self, &mut inner); + } + + /// Try to acquire a permit without waiting. + pub(crate) fn try_acquire(self: &std::sync::Arc) -> Option { + let mut inner = self.inner.lock().expect("launch gate poisoned"); + (inner.active < inner.capacity).then(|| { + inner.active += 1; + DynamicGatePermit { + gate: Some(std::sync::Arc::clone(self)), + } + }) + } + + /// Acquire a permit, waiting until capacity is available. Cancellation + /// safe: a dropped future either leaves a stale queue entry (skipped and + /// disarmed by the granter) or drops an already-dispatched permit (whose + /// `Drop` re-releases the slot). + pub(crate) async fn acquire(self: &std::sync::Arc) -> DynamicGatePermit { + loop { + let rx = { + let mut inner = self.inner.lock().expect("launch gate poisoned"); + if inner.active < inner.capacity { + inner.active += 1; + return DynamicGatePermit { + gate: Some(std::sync::Arc::clone(self)), + }; + } + let (tx, rx) = oneshot::channel(); + inner.waiters.push_back(GateWaiter { sender: tx }); + rx + }; + // A failed receive means the gate itself was dropped while we + // were queued; the loop re-queues under the lock. + if let Ok(permit) = rx.await { + return permit; + } + } + } +} + +/// One held launch slot. Released on drop. +/// +/// The gate is an `Option` so the wake path can disarm a permit whose +/// receiver vanished without running `Drop` (which would re-enter the locked +/// `release()`). +pub(crate) struct DynamicGatePermit { + gate: Option>, +} + +impl DynamicGatePermit { + fn disarm(&mut self) -> Option> { + self.gate.take() + } +} + +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) { + if let Some(gate) = self.gate.take() { + gate.release(); + } + } +} + +// === RateLimitGovernor === + +#[derive(Debug)] +struct GovernorState { + /// Ceiling additive increase may climb to (configured launch + /// concurrency). + max_capacity: usize, + /// Timestamps of rate-limited attempts inside the window. + limited: VecDeque, + /// Timestamps of all reported attempts inside the window (successes and + /// rate limits) — the denominator of the recent rate-limit ratio. + attempts: VecDeque, + consecutive_successes: u32, + paused: bool, +} + +/// Rate-limit aware scheduler over a [`DynamicGate`] (see module docs). +#[derive(Debug)] +pub(crate) struct RateLimitGovernor { + gate: std::sync::Arc, + state: Mutex, +} + +impl RateLimitGovernor { + pub(crate) fn new(max_capacity: usize) -> (std::sync::Arc, std::sync::Arc) { + let gate = std::sync::Arc::new(DynamicGate::new(max_capacity.max(1))); + let governor = std::sync::Arc::new(Self { + gate: std::sync::Arc::clone(&gate), + state: Mutex::new(GovernorState { + max_capacity: max_capacity.max(1), + limited: VecDeque::new(), + attempts: VecDeque::new(), + consecutive_successes: 0, + paused: false, + }), + }); + (governor, gate) + } + + /// The governor's launch gate. `SubAgentManager` hands this to spawned + /// tasks in place of the old fixed `Semaphore`. (Directly exercised by + /// governor unit tests.) + #[cfg_attr(not(test), allow(dead_code))] + pub(crate) fn gate(&self) -> std::sync::Arc { + std::sync::Arc::clone(&self.gate) + } + + /// Apply a new configured launch capacity: the AIMD ceiling and the gate + /// capacity while not throttled. Applies to the live gate immediately + /// (raising and lowering alike) unless the governor is paused — a pause + /// keeps capacity 0 until recovery, so an external limit change cannot + /// silently lift a rate-limit pause. + pub(crate) fn set_max_capacity(&self, max_capacity: usize) { + let mut state = self.state.lock().expect("rate limit governor poisoned"); + state.max_capacity = max_capacity.max(1); + if !state.paused { + self.gate.set_capacity(state.max_capacity); + } + } + + fn prune(state: &mut GovernorState, now: Instant) { + while state + .limited + .front() + .is_some_and(|at| now.duration_since(*at) > RATE_LIMIT_WINDOW) + { + state.limited.pop_front(); + } + while state + .attempts + .front() + .is_some_and(|at| now.duration_since(*at) > RATE_LIMIT_WINDOW) + { + state.attempts.pop_front(); + } + } + + /// Report that a sub-agent LLM attempt is starting. Contributes to the + /// recent-attempt denominator for the ratio heuristic. + pub(crate) fn record_attempt(&self, now: Instant) { + let mut state = self.state.lock().expect("rate limit governor poisoned"); + Self::prune(&mut state, now); + state.attempts.push_back(now); + } + + /// Lift a pause whose rate-limit events have all aged out of the window, + /// resuming at a conservative quarter of the configured capacity so + /// additive increase climbs the rest of the way. Callers must hold the + /// state lock; `prune` first. + fn unpause_if_window_drained(&self, state: &mut GovernorState) { + if !state.paused || !state.limited.is_empty() { + return; + } + state.paused = false; + let capacity = (state.max_capacity / 4).max(1); + self.gate.set_capacity(capacity); + tracing::info!( + target: "subagent", + launch_capacity = capacity, + max_capacity = state.max_capacity, + "rate-limit governor resumed launches after window drained" + ); + } + + /// Time-driven recovery probe for queued launches. A pause is normally + /// lifted by a successful LLM attempt from an in-flight child, but if the + /// entire in-flight fleet finishes while 429 events are still inside the + /// window, no success ever arrives — without this probe the queue would + /// freeze until each queued child hits its wall-time deadline. Once every + /// limit event has aged out, the next probe resumes launches. + pub(crate) fn recover_if_window_drained(&self, now: Instant) { + let mut state = self.state.lock().expect("rate limit governor poisoned"); + Self::prune(&mut state, now); + self.unpause_if_window_drained(&mut state); + } + + /// Report a successful sub-agent LLM attempt. Drives AIMD additive + /// increase and clears the pause once the window has drained. + pub(crate) fn record_success(&self, now: Instant) { + let mut state = self.state.lock().expect("rate limit governor poisoned"); + Self::prune(&mut state, now); + state.consecutive_successes = state.consecutive_successes.saturating_add(1); + + self.unpause_if_window_drained(&mut state); + + if !state.paused + && state.consecutive_successes >= SUCCESS_PER_INCREASE_STEP + && self.gate.capacity() < state.max_capacity + { + state.consecutive_successes = 0; + let capacity = (self.gate.capacity() + 1).min(state.max_capacity); + self.gate.set_capacity(capacity); + tracing::debug!( + target: "subagent", + launch_capacity = capacity, + "rate-limit governor additively increased launch capacity" + ); + } + } + + /// Report a rate-limited (429) sub-agent LLM attempt. May shrink or pause + /// the launch gate; never touches in-flight calls or retries. + pub(crate) fn record_rate_limited(&self, now: Instant) { + let mut state = self.state.lock().expect("rate limit governor poisoned"); + Self::prune(&mut state, now); + state.limited.push_back(now); + // The denominator (`attempts`) already contains this attempt — the + // call path reports `record_attempt` before every LLM call, retries + // included. Pushing again would double-count failures and skew the + // ratio. + state.consecutive_successes = 0; + + if state.paused { + return; + } + + let events = state.limited.len(); + let attempts = state.attempts.len().max(1); + let ratio = f64::from(events as u32) / f64::from(attempts as u32); + + if events >= PAUSE_EVENT_THRESHOLD { + state.paused = true; + // Capacity 0 blocks all *new* admissions; children already holding + // permits keep running to completion. + self.gate.set_capacity(0); + tracing::warn!( + target: "subagent", + window_events = events, + window_attempts = attempts, + "rate-limit governor paused new sub-agent launches (sustained provider 429s); \ + queued children wait for the window to drain" + ); + return; + } + + // The ratio heuristic only fires once the window has real volume + // (>= 2 observed attempts): with a single attempt every 429 is 100% + // and would shrink the gate on the first blip, fighting the absolute + // count threshold that is meant to own small-fleet behavior. + if events >= THROTTLE_EVENT_THRESHOLD + || (state.attempts.len() >= 2 && ratio > THROTTLE_RATIO_THRESHOLD) + { + let current = self.gate.capacity(); + if current > 1 { + let capacity = (current / 2).max(1); + self.gate.set_capacity(capacity); + tracing::warn!( + target: "subagent", + window_events = events, + window_ratio = format!("{ratio:.2}"), + previous_capacity = current, + launch_capacity = capacity, + "rate-limit governor multiplicatively decreased launch capacity" + ); + } + } + } + + /// Whether new launches are currently paused because of sustained 429s. + pub(crate) fn is_paused(&self, now: Instant) -> bool { + let mut state = self.state.lock().expect("rate limit governor poisoned"); + Self::prune(&mut state, now); + state.paused + } + + /// Observability snapshot: `(gate capacity, window limit events, paused)`. + /// (Unit-test/diagnostics surface; wired into status events by the parent + /// repo follow-up.) + #[cfg_attr(not(test), allow(dead_code))] + pub(crate) fn snapshot(&self, now: Instant) -> GovernorSnapshot { + let mut state = self.state.lock().expect("rate limit governor poisoned"); + Self::prune(&mut state, now); + GovernorSnapshot { + launch_capacity: self.gate.capacity(), + max_capacity: state.max_capacity, + window_limited: state.limited.len(), + window_attempts: state.attempts.len(), + paused: state.paused, + } + } +} + +/// Point-in-time view of the governor for tests and diagnostics. +#[cfg_attr(not(test), allow(dead_code))] +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) struct GovernorSnapshot { + pub(crate) launch_capacity: usize, + pub(crate) max_capacity: usize, + pub(crate) window_limited: usize, + pub(crate) window_attempts: usize, + pub(crate) paused: bool, +} + +#[cfg(test)] +mod tests { + use super::*; + + fn ms(n: u64) -> Duration { + Duration::from_millis(n) + } + + #[test] + fn window_counts_and_prunes_events() { + let (governor, _gate) = RateLimitGovernor::new(4); + let t0 = Instant::now(); + for i in 0..5 { + governor.record_attempt(t0 + ms(i * 10)); + governor.record_rate_limited(t0 + ms(i * 10)); + } + let snap = governor.snapshot(t0 + ms(60)); + assert_eq!(snap.window_limited, 5); + assert_eq!(snap.window_attempts, 5); + + // Events older than the 60s window drop out (strictly past the + // window edge: the newest event is at t0+40ms). + let snap = governor.snapshot(t0 + RATE_LIMIT_WINDOW + ms(50)); + assert_eq!(snap.window_limited, 0); + assert_eq!(snap.window_attempts, 0); + } + + #[test] + fn multiplicative_decrease_halves_capacity_on_threshold() { + let (governor, _gate) = RateLimitGovernor::new(8); + let t0 = Instant::now(); + // First event: below both thresholds, no change. + governor.record_attempt(t0); + governor.record_rate_limited(t0); + assert_eq!(governor.snapshot(t0).launch_capacity, 8); + // Second event: hits the count threshold, halve. + governor.record_attempt(t0 + ms(1)); + governor.record_rate_limited(t0 + ms(1)); + assert_eq!(governor.snapshot(t0).launch_capacity, 4); + // Third: halve again. + governor.record_attempt(t0 + ms(2)); + governor.record_rate_limited(t0 + ms(2)); + assert_eq!(governor.snapshot(t0).launch_capacity, 2); + // Fourth: hits the pause threshold. + governor.record_attempt(t0 + ms(3)); + governor.record_rate_limited(t0 + ms(3)); + let snap = governor.snapshot(t0); + assert!(snap.paused); + } + + #[test] + fn ratio_threshold_triggers_decrease_even_with_few_events() { + let (governor, _gate) = RateLimitGovernor::new(8); + let t0 = Instant::now(); + // One success then one 429: the absolute event count is below the + // threshold, but the 50% limit ratio must still shrink the gate. + governor.record_attempt(t0); + governor.record_success(t0); + governor.record_attempt(t0 + ms(1)); + governor.record_rate_limited(t0 + ms(1)); + assert!( + governor.snapshot(t0 + ms(2)).launch_capacity < 8, + "50% limit ratio should trigger a decrease" + ); + } + + #[test] + fn additive_increase_recovers_capacity_gradually() { + let (governor, _gate) = RateLimitGovernor::new(8); + let t0 = Instant::now(); + // Drive capacity down to 4 via two events. + governor.record_attempt(t0); + governor.record_rate_limited(t0); + governor.record_attempt(t0 + ms(1)); + governor.record_rate_limited(t0 + ms(1)); + assert_eq!(governor.snapshot(t0).launch_capacity, 4); + + // Three consecutive successes add exactly one unit of capacity. + for i in 0..3u32 { + governor.record_attempt(t0 + ms(10 + u64::from(i))); + governor.record_success(t0 + ms(10 + u64::from(i))); + } + assert_eq!(governor.snapshot(t0 + ms(20)).launch_capacity, 5); + for i in 0..3u32 { + governor.record_attempt(t0 + ms(30 + u64::from(i))); + governor.record_success(t0 + ms(30 + u64::from(i))); + } + assert_eq!(governor.snapshot(t0 + ms(40)).launch_capacity, 6); + + // A rate limit resets the success streak. + governor.record_attempt(t0 + ms(50)); + governor.record_rate_limited(t0 + ms(50)); + for i in 0..2u32 { + governor.record_attempt(t0 + ms(60 + u64::from(i))); + governor.record_success(t0 + ms(60 + u64::from(i))); + } + governor.record_attempt(t0 + ms(80)); + governor.record_success(t0 + ms(80)); + // 2 successes before the limit + 1 after = 3 successes, but the limit + // reset the streak, and the third event in the window halved again + // (6 -> 3) before successes could climb. + assert!(governor.snapshot(t0 + ms(90)).launch_capacity <= 6); + } + + #[test] + fn pause_releases_only_after_window_drains() { + let (governor, gate) = RateLimitGovernor::new(8); + let t0 = Instant::now(); + for i in 0..4 { + governor.record_attempt(t0 + ms(i)); + governor.record_rate_limited(t0 + ms(i)); + } + assert!(governor.is_paused(t0 + ms(10))); + assert_eq!(governor.snapshot(t0 + ms(10)).launch_capacity, 0); + + // Successes before the window drains do NOT unpause. + governor.record_success(t0 + ms(20)); + assert!(governor.is_paused(t0 + ms(30))); + + // Once every limit event ages out, the next success resumes at a + // quarter of capacity. + let late = t0 + RATE_LIMIT_WINDOW + ms(10); + governor.record_success(late); + assert!(!governor.is_paused(late)); + assert_eq!(governor.snapshot(late).launch_capacity, 2); + assert_eq!(gate.capacity(), 2); + } + + #[test] + fn capacity_increase_is_capped_at_max() { + let (governor, _gate) = RateLimitGovernor::new(2); + let t0 = Instant::now(); + for i in 0..12u32 { + governor.record_attempt(t0 + ms(u64::from(i))); + governor.record_success(t0 + ms(u64::from(i))); + } + assert_eq!(governor.snapshot(t0).launch_capacity, 2); + } + + #[test] + fn gate_blocks_when_full_and_releases_on_drop() { + let (governor, gate) = RateLimitGovernor::new(1); + let rt = tokio::runtime::Builder::new_current_thread() + .enable_time() + .build() + .expect("test runtime"); + rt.block_on(async move { + let first = governor.gate().try_acquire().expect("first permit"); + assert!(gate.try_acquire().is_none(), "capacity 1 must be full"); + + let g2 = std::sync::Arc::clone(&gate); + let waiter = tokio::spawn(async move { g2.acquire().await }); + + // Waiter stays blocked while the first permit is held. + tokio::time::sleep(ms(20)).await; + assert!(!waiter.is_finished()); + + drop(first); + let _second = waiter.await.expect("waiter task"); + }); + } + + #[test] + fn gate_set_capacity_shrinks_below_active_and_re_admits_later() { + let (governor, gate) = RateLimitGovernor::new(4); + let rt = tokio::runtime::Builder::new_current_thread() + .enable_time() + .build() + .expect("test runtime"); + rt.block_on(async move { + let mut held: Vec<_> = (0..4) + .map(|_| gate.try_acquire().expect("permit within capacity")) + .collect(); + assert_eq!(gate.capacity(), 4); + + // Shrink below the active count: no new permit is granted. + governor.gate().set_capacity(1); + assert_eq!(gate.capacity(), 1); + assert!(gate.try_acquire().is_none()); + + let g2 = std::sync::Arc::clone(&gate); + let waiter = tokio::spawn(async move { g2.acquire().await }); + tokio::time::sleep(ms(20)).await; + assert!(!waiter.is_finished(), "must wait while active >= capacity"); + + // Releasing holders drains `active` toward the new capacity; the + // waiter is admitted only once every held permit is released + // (active 4 -> 0 < capacity 1). + drop(held.swap_remove(0)); + drop(held.swap_remove(0)); + drop(held.swap_remove(0)); + drop(held); + let _permit = waiter.await.expect("waiter admitted after drain"); + assert!(gate.try_acquire().is_none(), "capacity 1 is now full"); + drop(_permit); + }); + } + + #[test] + fn rate_limit_retry_delay_is_full_jitter_within_base() { + for retry in 1..=12u32 { + let base = rate_limit_backoff_base(retry); + for _ in 0..64 { + let delay = rate_limit_retry_delay(retry); + assert!(delay <= base, "full jitter must not exceed the base"); + } + } + // The cap holds for absurd retry numbers. + assert_eq!(rate_limit_backoff_base(40), RATE_LIMIT_MAX_BACKOFF); + } + + /// A pause must lift via the time-driven probe even when no in-flight + /// child ever reports another success (the in-flight fleet drained before + /// the window did): otherwise queued children freeze until their + /// wall-time deadline. + #[test] + fn forkguard_rate_limit_governor_pauses_and_time_recovers_after_window_drains() { + let (governor, _gate) = RateLimitGovernor::new(8); + let t0 = Instant::now(); + for i in 0..4 { + governor.record_attempt(t0 + ms(i)); + governor.record_rate_limited(t0 + ms(i)); + } + assert!(governor.is_paused(t0 + ms(10))); + + // Probe while 429 events are still inside the window: stays paused. + governor.recover_if_window_drained(t0 + ms(20)); + assert!(governor.is_paused(t0 + ms(30))); + + // Once every limit event has aged out, the probe resumes launches at + // a quarter of the configured capacity — no success event required. + let late = t0 + RATE_LIMIT_WINDOW + ms(10); + governor.recover_if_window_drained(late); + assert!(!governor.is_paused(late)); + assert_eq!(governor.snapshot(late).launch_capacity, 2); + } + + /// A runtime launch-concurrency change must not silently lift a pause: + /// the gate stays at capacity 0 until the window drains, then resumes at + /// a quarter of the *new* configured capacity. + #[test] + fn forkguard_rate_limit_governor_limit_change_keeps_pause_capacity_zero() { + let (governor, gate) = RateLimitGovernor::new(8); + let t0 = Instant::now(); + for i in 0..4 { + governor.record_attempt(t0 + ms(i)); + governor.record_rate_limited(t0 + ms(i)); + } + assert!(governor.is_paused(t0 + ms(1))); + + governor.set_max_capacity(4); + assert_eq!(gate.capacity(), 0, "pause must keep capacity 0"); + + let late = t0 + RATE_LIMIT_WINDOW + ms(10); + governor.recover_if_window_drained(late); + assert_eq!( + gate.capacity(), + 1, + "resume at a quarter of the new capacity" + ); + } + + /// A waiter cancelled *after* its grant was dispatched must not swallow + /// the slot: the permit is dropped with the cancelled future and its + /// `Drop` re-releases it for the next waiter. + #[test] + fn forkguard_dynamic_gate_redispatches_grant_of_cancelled_waiter() { + let (_governor, gate) = RateLimitGovernor::new(1); + let rt = tokio::runtime::Builder::new_current_thread() + .enable_time() + .build() + .expect("test runtime"); + rt.block_on(async move { + let holder = gate.try_acquire().expect("holder"); + let g2 = std::sync::Arc::clone(&gate); + let waiter = tokio::spawn(async move { g2.acquire().await }); + tokio::time::sleep(ms(20)).await; + assert!(!waiter.is_finished(), "waiter must be queued"); + + // Releasing the holder dispatches the grant into the waiter's + // channel; on a current-thread runtime the waiter has not polled + // yet when we abort it, so the permit is dropped mid-flight. + drop(holder); + waiter.abort(); + tokio::time::sleep(ms(20)).await; + + assert!( + gate.try_acquire().is_some(), + "grant of cancelled waiter must be re-released, not leaked" + ); + }); + } +} diff --git a/crates/tui/src/tools/subagent/mod.rs b/crates/tui/src/tools/subagent/mod.rs index ec16708819..7c829271c6 100644 --- a/crates/tui/src/tools/subagent/mod.rs +++ b/crates/tui/src/tools/subagent/mod.rs @@ -17,7 +17,7 @@ use std::io::{Read, Write}; use std::path::{Component, Path, PathBuf}; use std::sync::Arc; use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH}; -use tokio::sync::{Mutex, RwLock, Semaphore}; +use tokio::sync::{Mutex, RwLock}; use anyhow::{Result, anyhow}; use async_trait::async_trait; @@ -67,6 +67,7 @@ use coord::{ pub mod advisor; pub mod coord; +mod governor; pub mod mailbox; mod naming; mod worktree; @@ -334,6 +335,9 @@ const SUBAGENT_RESTART_REASON: &str = "Interrupted by process restart"; #[cfg(test)] const SUBAGENT_MODEL_WAIT_REASON: &str = "waiting for model response"; const SUBAGENT_QUEUED_LAUNCH_REASON: &str = "queued: waiting for a sub-agent launch slot"; +/// Queued-reason variant used while the rate-limit governor has paused new +/// sub-agent launches after sustained provider 429s. +const SUBAGENT_QUEUED_RATE_LIMIT_REASON: &str = "queued: waiting for provider rate-limit recovery"; /// #freeze: minimum spacing between hot-path (per-step checkpoint) state /// persists. `update_checkpoint` fires on every step of every agent; at high /// fanout an unconditional full-fleet rewrite under the manager write lock @@ -2217,6 +2221,14 @@ pub struct SubAgentRuntime { pub todos: SharedTodoList, /// Session mode of the orchestrating parent at spawn time (Wave 7 M4/M5). pub parent_mode: AppMode, + /// Shared rate-limit governor for the fleet that spawned this runtime. + /// Stamped by the spawning manager in + /// `spawn_background_with_assignment_options` (the single chokepoint all + /// spawn variants funnel through), so every descendant LLM attempt + /// reports 429s/successes to the fleet's adaptive scheduler; cloned into + /// child runtimes. `None` for runtimes built outside a manager (tests, + /// tool-only runtimes). + pub(crate) governor: Option>, } impl SubAgentRuntime { @@ -2268,6 +2280,12 @@ impl SubAgentRuntime { speech_output_dir: None, todos: crate::tools::todo::new_shared_todo_list(), parent_mode: AppMode::Agent, + // Stamped by the spawning manager in + // `spawn_background_with_assignment_options`, so every descendant + // LLM attempt reports 429s/successes to the fleet's rate-limit + // governor. `None` for runtimes built outside a manager (tests, + // tool-only runtimes). + governor: None, } } @@ -2295,6 +2313,19 @@ impl SubAgentRuntime { self } + /// Stamp the fleet's rate-limit governor onto a root runtime. The manager + /// owns the governor (and its launch gate), but only runtimes carrying it + /// report 429s/successes; without this the whole descendant tree inherits + /// `None` and the AIMD scheduler never observes provider throttling. + #[must_use] + pub(crate) fn with_fleet_governor( + mut self, + governor: Arc, + ) -> Self { + self.governor = Some(governor); + self + } + /// Preserve the parent Agent-mode native tool surface for child registries. #[must_use] pub fn with_agent_tool_surface_options(mut self, options: AgentToolSurfaceOptions) -> Self { @@ -2554,6 +2585,9 @@ impl SubAgentRuntime { // siblings' progress. Parent todo state is still visible to an // opt-in forked child as immutable `fork_context` text. todos: crate::tools::todo::new_shared_todo_list(), + // Inherit the fleet's rate-limit governor so every descendant + // LLM attempt reports 429s/successes to the adaptive scheduler. + governor: self.governor.clone(), parent_mode: self.parent_mode, } } @@ -2950,7 +2984,20 @@ pub struct SubAgentManager { /// publishing a visible "queued" reason instead of bursting. Deeper /// descendants bypass the gate so a permit-holding parent waiting on /// its own children cannot deadlock the tree. - launch_gate: Arc, + /// + /// The gate is a [`governor::DynamicGate`] rather than a + /// `tokio::sync::Semaphore` so the rate-limit governor can shrink its + /// capacity at runtime (even below the number of active children) + /// without replacing the `Arc` — a semaphore swap silently fails while + /// any child still holds a permit, which is why + /// `update_runtime_limits` previously only applied launch-concurrency + /// changes to an idle fleet. + launch_gate: Arc, + /// Rate-limit aware scheduler feeding `launch_gate` (swarm-mode + /// adaptive throttling). Sub-agent LLM attempts report 429s and + /// successes through [`SubAgentRuntime::governor`]; the governor + /// shrinks/pauses admissions on sustained 429s and recovers via AIMD. + governor: Arc, /// #freeze: hot-path persist debounce bookkeeping (see /// `SUBAGENT_PERSIST_DEBOUNCE`). `last_persist_at` is the last time any /// state persist ran; `persist_pending` records that a hot-path write was @@ -2991,6 +3038,9 @@ impl SubAgentManager { /// separately from its execution workspace. #[must_use] pub fn new_with_state_root(workspace: PathBuf, state_root: PathBuf, max_agents: usize) -> Self { + // The governor owns the launch gate it schedules, so manager builders + // and the runtime limiter adjust capacity through the pair. + let (governor, launch_gate) = governor::RateLimitGovernor::new(max_agents.max(1)); Self { agents: HashMap::new(), worker_records: HashMap::new(), @@ -3014,7 +3064,8 @@ impl SubAgentManager { current_session_boot_id: format!("boot_{}", &Uuid::new_v4().to_string()[..12]), // Default launch concurrency = the full agent cap; the gate only // throttles when a lower `launch_concurrency` is configured. - launch_gate: Arc::new(Semaphore::new(max_agents.max(1))), + launch_gate, + governor, last_persist_at: None, persist_pending: false, last_cleanup_at: None, @@ -3027,12 +3078,25 @@ impl SubAgentManager { /// Set the number of direct children that may execute concurrently /// before further launches queue (#3095). Clamped to `1..=max_agents`. + /// Applied to the live gate capacity, so this also takes effect when + /// called after children have started. Routed through the governor so a + /// rate-limit pause is not silently lifted by a limit change. #[must_use] - pub fn with_launch_concurrency(mut self, limit: usize) -> Self { - self.launch_gate = Arc::new(Semaphore::new(limit.clamp(1, self.max_agents))); + pub fn with_launch_concurrency(self, limit: usize) -> Self { + let limit = limit.clamp(1, self.max_agents); + self.governor.set_max_capacity(limit); self } + /// The rate-limit governor backing [`Self::launch_gate`]; exposed so the + /// engine can stamp it onto root runtimes and tests can drive the + /// adaptive scheduler. (Surfacing governor state in status events is a + /// parent-repo follow-up.) + #[must_use] + pub(crate) fn rate_limit_governor(&self) -> Arc { + Arc::clone(&self.governor) + } + /// Set the total queued + running admission ceiling for this manager. /// The value is always at least the instantaneous concurrency cap. #[must_use] @@ -3580,9 +3644,11 @@ impl SubAgentManager { self } - /// Apply live runtime limits. The launch semaphore is replaced only when - /// no sub-agent is currently running, because active tasks may still hold - /// permits from the previous semaphore. + /// Apply live runtime limits. The launch gate is a + /// [`governor::DynamicGate`], so the new launch concurrency applies to + /// the live capacity immediately — children already holding permits keep + /// running, and no admission above the new capacity is granted until the + /// active count drains. Always returns `true`. pub fn update_runtime_limits( &mut self, max_agents: usize, @@ -3600,13 +3666,11 @@ impl SubAgentManager { } else { running_heartbeat_timeout }; - if self.running_count() == 0 { - self.launch_gate = - Arc::new(Semaphore::new(launch_concurrency.clamp(1, self.max_agents))); - true - } else { - false - } + let launch_concurrency = launch_concurrency.clamp(1, self.max_agents); + // Routed through the governor so a rate-limit pause (gate capacity 0) + // is not silently lifted by a runtime limit change. + self.governor.set_max_capacity(launch_concurrency); + true } /// Build the [`PersistedSubAgentState`] snapshot from the current fleet. @@ -5382,6 +5446,12 @@ impl SubAgentManager { allowed_tools: Option>, options: SubAgentSpawnOptions, ) -> Result { + // Every manager-spawned runtime carries the fleet governor, so the + // spawned agent and its whole descendant tree report 429s/successes + // to the adaptive scheduler. Runtimes built outside a manager (tests, + // tool-only runtimes) keep `governor: None`. + runtime.governor = Some(Arc::clone(&self.governor)); + self.cleanup(COMPLETED_AGENT_RETENTION); self.check_admission_capacity()?; @@ -8824,7 +8894,7 @@ struct SubAgentTask { /// children: the task acquires a permit before its first model step and /// holds it until completion, so a fanout burst beyond the limit queues /// with a visible reason instead of executing all at once. - launch_gate: Option>, + launch_gate: Option>, } #[allow(clippy::too_many_lines)] @@ -8856,9 +8926,9 @@ async fn run_subagent_task(task: SubAgentTask) { let mut _launch_permit = None; let mut launch_wait_timed_out = false; if let Some(gate) = task.launch_gate.as_ref() { - match Arc::clone(gate).try_acquire_owned() { - Ok(permit) => _launch_permit = Some(permit), - Err(tokio::sync::TryAcquireError::NoPermits) => { + match Arc::clone(gate).try_acquire() { + Some(permit) => _launch_permit = Some(permit), + None => { match tokio::time::timeout_at( deadline.into(), acquire_queued_launch_permit(&task, Arc::clone(gate)), @@ -8869,12 +8939,6 @@ async fn run_subagent_task(task: SubAgentTask) { Err(_) => launch_wait_timed_out = true, } } - Err(tokio::sync::TryAcquireError::Closed) => { - crate::logging::warn(format!( - "sub-agent launch gate closed for {}; proceeding without backpressure", - task.agent_id - )); - } } } @@ -8959,34 +9023,68 @@ async fn run_subagent_task(task: SubAgentTask) { async fn acquire_queued_launch_permit( task: &SubAgentTask, - gate: Arc, -) -> Option { - record_queued_launch_progress(task).await; - tokio::select! { - biased; - () = task.runtime.cancel_token.cancelled() => { - record_agent_progress( - &task.runtime, - &task.agent_id, - AgentProgressEventMeta::new(AgentWorkerStatus::Cancelled), - "cancelled while queued for a sub-agent launch slot".to_string(), - ); - None - } - permit = Arc::clone(&gate).acquire_owned() => { - permit.ok() + gate: Arc, +) -> Option { + // When the governor has paused launches over sustained provider 429s, + // surface the reason in the queued status instead of the generic + // "waiting for a launch slot" message. + let paused_for_rate_limit = task + .runtime + .governor + .as_ref() + .is_some_and(|governor| governor.is_paused(Instant::now())); + let queued_reason = if paused_for_rate_limit { + SUBAGENT_QUEUED_RATE_LIMIT_REASON + } else { + SUBAGENT_QUEUED_LAUNCH_REASON + }; + record_queued_launch_progress(task, queued_reason).await; + // While queued, periodically probe the governor: if a rate-limit pause + // outlives its window (the in-flight fleet finished before any success + // could lift the pause), the probe resumes launches instead of leaving + // the queue frozen until each child's wall-time deadline. + let mut recovery_probe = tokio::time::interval( + task.runtime + .governor + .as_ref() + .map(|_| std::time::Duration::from_secs(5)) + .unwrap_or(std::time::Duration::from_secs(3600)), + ); + loop { + tokio::select! { + biased; + () = task.runtime.cancel_token.cancelled() => { + record_agent_progress( + &task.runtime, + &task.agent_id, + AgentProgressEventMeta::new(AgentWorkerStatus::Cancelled), + "cancelled while queued for a sub-agent launch slot".to_string(), + ); + return None; + } + _ = recovery_probe.tick() => { + if let Some(governor) = task.runtime.governor.as_ref() { + governor.recover_if_window_drained(Instant::now()); + } + // If the probe lifted a pause it raised the gate capacity, + // which grants queued waiters; `gate.acquire` below is + // re-polled either way on the next loop iteration. + } + permit = gate.acquire() => { + return Some(permit); + } } } } -async fn record_queued_launch_progress(task: &SubAgentTask) { +async fn record_queued_launch_progress(task: &SubAgentTask, queued_reason: &'static str) { { let mut manager = task.runtime.manager.write().await; manager.touch(&task.agent_id); manager.record_worker_event( &task.agent_id, AgentWorkerStatus::Queued, - Some(SUBAGENT_QUEUED_LAUNCH_REASON.to_string()), + Some(queued_reason.to_string()), None, None, ); @@ -8994,16 +9092,13 @@ async fn record_queued_launch_progress(task: &SubAgentTask) { emit_agent_progress( task.runtime.event_tx.as_ref(), &task.agent_id, - SUBAGENT_QUEUED_LAUNCH_REASON.to_string(), + queued_reason.to_string(), AgentProgressEventMeta::new(AgentWorkerStatus::Queued), task.runtime.parent_agent_id.clone(), task.runtime.spawn_depth, ); if let Some(mailbox) = task.runtime.mailbox.as_ref() { - let _ = mailbox.send(MailboxMessage::progress( - &task.agent_id, - SUBAGENT_QUEUED_LAUNCH_REASON, - )); + let _ = mailbox.send(MailboxMessage::progress(&task.agent_id, queued_reason)); } } @@ -9611,8 +9706,11 @@ fn retryable_subagent_provider_failure( return Some(RetryableSubAgentProviderFailure { label: "rate-limited provider response", checkpoint_reason: "api_rate_limited", - delay: retry_after - .unwrap_or_else(|| subagent_transient_provider_retry_delay(retry_number)), + // Honor the provider's `Retry-After` when present. Without it, + // back off exponentially with full jitter (capped at 120s) so a + // fan-out of children 429'd by the same provider response does + // not retry in lockstep (thundering herd). + delay: retry_after.unwrap_or_else(|| governor::rate_limit_retry_delay(retry_number)), }); } @@ -9688,14 +9786,36 @@ async fn request_subagent_model_response_with_retries( let usage_route = runtime .client .effective_route_envelope(&runtime.model, chrono::Utc::now()); + // Report the attempt to the fleet's rate-limit governor; the ratio + // denominator for the AIMD heuristic counts retried attempts too. + 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(request.clone()), ) .await { - Ok(Ok(response)) => return Ok((response, usage_route)), + Ok(Ok(response)) => { + // A successful call signals recovery; drives AIMD additive + // increase and (once limits age out of the window) unpauses. + if let Some(governor) = runtime.governor.as_ref() { + governor.record_success(Instant::now()); + } + return Ok((response, usage_route)); + } Ok(Err(err)) => { + // A provider 429 feeds the governor's sliding window (AIMD + // multiplicative decrease / pause). `QuotaExhausted` and all + // other errors keep their existing paths untouched. + if matches!( + err.downcast_ref::(), + Some(LlmError::RateLimited { .. }) + ) && let Some(governor) = runtime.governor.as_ref() + { + governor.record_rate_limited(Instant::now()); + } let retry_number = transient_failures.saturating_add(1); let Some(retryable) = retryable_subagent_provider_failure(&err, retry_number) else { diff --git a/crates/tui/src/tools/subagent/tests.rs b/crates/tui/src/tools/subagent/tests.rs index f2c45cfb36..bc7c0a6f66 100644 --- a/crates/tui/src/tools/subagent/tests.rs +++ b/crates/tui/src/tools/subagent/tests.rs @@ -10761,6 +10761,9 @@ pub(crate) fn stub_runtime() -> SubAgentRuntime { tool_timeout: DEFAULT_TOOL_TIMEOUT, speech_output_dir: None, todos: crate::tools::todo::new_shared_todo_list(), + // Test stubs run without a manager-stamped governor; the LLM call + // path treats `None` as "report nothing". + governor: None, } } @@ -13007,7 +13010,6 @@ fn launch_gate_defaults_to_launch_concurrency_capped_by_max_agents() { #[tokio::test] async fn launch_gate_queues_extra_direct_children() { - use tokio::sync::Semaphore; use tokio_util::sync::CancellationToken; let tmp = tempdir().expect("tempdir"); @@ -13024,12 +13026,11 @@ async fn launch_gate_queues_extra_direct_children() { runtime.context = ToolContext::new(tmp.path()); runtime.mailbox = Some(mailbox); - let gate = Arc::new(Semaphore::new(1)); + let gate = Arc::new(governor::DynamicGate::new(1)); let held_launch_permit = Arc::clone(&gate) - .acquire_owned() - .await + .try_acquire() .expect("test holds the single launch permit"); - let spawn = |agent_id: &str, gate: Option>| { + let spawn = |agent_id: &str, gate: Option>| { let (input_tx, input_rx) = mpsc::unbounded_channel(); let agent = SubAgent::new( agent_id.to_string(), @@ -13156,7 +13157,6 @@ async fn launch_gate_queues_extra_direct_children() { #[tokio::test] async fn launch_gate_wait_counts_against_child_wall_timeout() { - use tokio::sync::Semaphore; use tokio_util::sync::CancellationToken; const WALL_TIME: Duration = Duration::from_millis(150); @@ -13188,10 +13188,9 @@ async fn launch_gate_wait_counts_against_child_wall_timeout() { runtime.context = ToolContext::new(tmp.path()); runtime.mailbox = Some(mailbox); - let gate = Arc::new(Semaphore::new(1)); + let gate = Arc::new(governor::DynamicGate::new(1)); let held_launch_permit = Arc::clone(&gate) - .acquire_owned() - .await + .try_acquire() .expect("test holds the single launch permit past the wall timeout"); let task = SubAgentTask { manager_handle: Arc::clone(&manager),