diff --git a/crates/execpolicy/src/lib.rs b/crates/execpolicy/src/lib.rs index 2edb6371e8..b589d9846f 100644 --- a/crates/execpolicy/src/lib.rs +++ b/crates/execpolicy/src/lib.rs @@ -6,6 +6,7 @@ pub mod shell_expand; pub use approval_mode::ApprovalMode; use std::collections::HashSet; +use std::sync::{Arc, RwLock}; use anyhow::Result; use bash_arity::BashArityDict; @@ -118,8 +119,11 @@ pub struct ToolAskRule { /// cannot silently authorize a later invocation with extra arguments. #[serde(default, skip_serializing_if = "std::ops::Not::not")] pub command_exact: bool, - /// Optional workspace-relative file path matched exactly after - /// normalization. + /// Optional file path matched exactly. A workspace-relative rule + /// normalizes against the call's workspace; a ROOTED rule (leading `/`, + /// `~/`, or a Windows drive) matches the call path exactly after + /// separator and case folding, so it can pin locations outside the + /// workspace. Traversal segments never match on either channel. #[serde(default, skip_serializing_if = "Option::is_none")] pub path: Option, /// Optional absolute workspace root that limits this rule to one repo. @@ -316,10 +320,20 @@ pub struct ExecPolicyContext<'a> { pub struct ExecPolicyEngine { /// Layered rulesets (builtin → agent → user). When non-empty, takes precedence /// over the legacy flat lists below. - rulesets: Vec, + /// + /// Shared behind an `Arc>` so that [`Self::set_ruleset`] applied + /// through one clone is observed by every clone. Hosts clone the engine + /// into long-lived side executors (nested sub-agent tool registries); a + /// plain `Vec` would leave those executors on a stale ruleset after a live + /// permission update, reopening an enforcement gap the parent no longer + /// has. + rulesets: Arc>>, /// Legacy flat lists kept for backward compatibility with `new()`. trusted_prefixes: Vec, denied_prefixes: Vec, + /// Deliberately clone-private, unlike `rulesets`: a remembered grant is a + /// decision the parent session made for its own calls, so it must not + /// silently authorize a delegated call in a cloned executor. approved_for_session: HashSet, /// Arity dictionary for command-prefix allow-rule matching. arity_dict: BashArityDict, @@ -329,7 +343,7 @@ impl ExecPolicyEngine { /// Legacy constructor: wraps the two vecs into a User-layer ruleset. pub fn new(trusted_prefixes: Vec, denied_prefixes: Vec) -> Self { Self { - rulesets: vec![], + rulesets: Arc::new(RwLock::new(vec![])), trusted_prefixes, denied_prefixes, approved_for_session: HashSet::new(), @@ -342,7 +356,7 @@ impl ExecPolicyEngine { pub fn with_rulesets(mut rulesets: Vec) -> Self { rulesets.sort_by_key(|r| r.layer); Self { - rulesets, + rulesets: Arc::new(RwLock::new(rulesets)), trusted_prefixes: vec![], denied_prefixes: vec![], approved_for_session: HashSet::new(), @@ -352,17 +366,39 @@ impl ExecPolicyEngine { /// Add a ruleset layer (re-sorts internally). pub fn add_ruleset(&mut self, ruleset: Ruleset) { - self.rulesets.push(ruleset); - self.rulesets.sort_by_key(|r| r.layer); + let mut guard = Self::lock_rulesets(&self.rulesets); + guard.push(ruleset); + guard.sort_by_key(|r| r.layer); } /// Replace the ruleset at one priority layer without clearing approvals /// remembered for the current session. pub fn set_ruleset(&mut self, ruleset: Ruleset) { + let mut guard = Self::lock_rulesets(&self.rulesets); + guard.retain(|existing| existing.layer != ruleset.layer); + guard.push(ruleset); + guard.sort_by_key(|existing| existing.layer); + } + + /// Lock the shared ruleset list for reading or writing. + /// + /// A poisoned lock (a panic held the guard mid-mutation) is recovered from: + /// the ruleset vec is plain data and a torn update sorts itself out on the + /// next `set_ruleset`, while refusing to answer policy checks would fail + /// closed for every command in the process. + fn lock_rulesets( + rulesets: &Arc>>, + ) -> std::sync::RwLockWriteGuard<'_, Vec> { + rulesets + .write() + .unwrap_or_else(std::sync::PoisonError::into_inner) + } + + /// Read-only snapshot of the shared ruleset list. + fn read_rulesets(&self) -> std::sync::RwLockReadGuard<'_, Vec> { self.rulesets - .retain(|existing| existing.layer != ruleset.layer); - self.rulesets.push(ruleset); - self.rulesets.sort_by_key(|existing| existing.layer); + .read() + .unwrap_or_else(std::sync::PoisonError::into_inner) } /// Resolve the effective trusted/denied prefix sets by merging all rulesets. @@ -372,17 +408,19 @@ impl ExecPolicyEngine { /// semantics: any matching deny prefix blocks the command regardless of layer. /// Trusted rules are only consulted after deny checks pass. fn resolve_prefixes(&self) -> (Vec, Vec) { - if self.rulesets.is_empty() { + let rulesets = self.read_rulesets(); + if rulesets.is_empty() { return (self.trusted_prefixes.clone(), self.denied_prefixes.clone()); } // Collect all trusted/denied across all layers, highest-priority last so they // shadow lower-priority entries with the same prefix. let mut trusted: Vec = vec![]; let mut denied: Vec = vec![]; - for rs in &self.rulesets { + for rs in rulesets.iter() { trusted.extend(rs.trusted_prefixes.iter().cloned()); denied.extend(rs.denied_prefixes.iter().cloned()); } + drop(rulesets); // Also merge legacy flat lists as user-layer. trusted.extend(self.trusted_prefixes.iter().cloned()); denied.extend(self.denied_prefixes.iter().cloned()); @@ -395,7 +433,8 @@ impl ExecPolicyEngine { .path .and_then(|path| normalize_workspace_relative_path(path, ctx.cwd)); - self.rulesets + let rulesets = self.read_rulesets(); + let matched = rulesets .iter() .flat_map(|ruleset| { ruleset @@ -415,18 +454,31 @@ impl ExecPolicyEngine { None => true, }) .filter(|(_, rule)| match (rule.path.as_deref(), ctx.path) { - (Some(pattern), Some(_)) => match ( - normalize_workspace_relative_path(pattern, ctx.cwd), - normalized_path.as_deref(), - ) { - (Some(pattern), Some(path)) => pattern == path, - _ => false, - }, + (Some(pattern), Some(call_path)) => { + let ws_rule = normalize_workspace_relative_path(pattern, ctx.cwd); + match (ws_rule, normalized_path.as_deref()) { + // Workspace-relative normalization fails for a call + // outside the workspace or a rule that names one, and + // on a POSIX host a Windows-spelled rule/call pair + // parses as unrelated relative forms. A rule spelling + // an ABSOLUTE path must still be able to match such a + // call exactly, or pinned locations (a real home, + // `/root`, a Windows profile) are unmatchable. The + // helper only fires for rooted rules, so relative + // semantics are unchanged. + (Some(ws_rule), Some(ws_call)) => { + ws_rule == ws_call || absolute_path_rule_matches(pattern, call_path) + } + _ => absolute_path_rule_matches(pattern, call_path), + } + } (Some(_), None) => false, (None, _) => true, }) .max_by_key(|(layer, rule)| (*layer, rule.action, ask_rule_specificity(rule))) - .map(|(_, rule)| rule.clone()) + .map(|(_, rule)| (*rule).clone()); + drop(rulesets); + matched } /// Records an approval key for the current session so subsequent checks skip approval. @@ -730,6 +782,16 @@ fn command_is_chained(command: &str) -> bool { /// direction. Matching stays anchored at the first positional token, so a /// non-flag token that isn't in the rule ends it — `git push` does not block /// `git checkout push`, and `rm` does not block `rmdir`. +/// +/// Two rule-side spellings widen what a rule can name. cmd.exe-style +/// single-letter `/` flags (`del /f /s /q`) in the *command* are skippable like +/// `-` flags, in any position. And a rule token of exactly `*` is a middle +/// wildcard matching zero or more consecutive command tokens regardless of +/// shape, so a rule can anchor on a tail (`grep * ~/.ssh/id_rsa`, +/// `dd * of=/dev/sda`) without enumerating every flag spelling. A wildcard +/// widens the deny face of a rule — each one must be justified by the rule +/// author. This engine is deliberately permissive; the rulesets that feed it +/// own the false-positive discipline of keeping wildcards narrow. fn denied_prefix_matches(rule: &str, command: &str) -> bool { let rule_tokens: Vec = normalize_command(rule) .split_whitespace() @@ -762,6 +824,24 @@ fn denied_prefix_matches(rule: &str, command: &str) -> bool { if j == rule_tokens.len() { return true; } + // A rule token of exactly `*` is a middle wildcard: it matches zero or + // more consecutive command tokens regardless of shape — that is its + // point, since `grep -i PATTERN ~/.ssh/id_rsa` interleaves flags and + // positionals no flag rule could enumerate. `(i, j+1)` lets it match + // nothing; `(i+1, j)` skips one more command token. `seen` keeps the + // run of states finite. This branch runs BEFORE the end-of-command + // bail below so a trailing `*` can still match zero tokens once the + // command is exhausted, degrading to plain prefix semantics, and a + // wildcard is never itself treated as a command word. + if rule_tokens[j] == "*" { + if seen.insert((i, j)) { + stack.push((i, j + 1)); + if i < command_tokens.len() { + stack.push((i + 1, j)); + } + } + continue; + } if i >= command_tokens.len() || !seen.insert((i, j)) { continue; } @@ -780,11 +860,16 @@ fn denied_prefix_matches(rule: &str, command: &str) -> bool { if matches_rule_token { stack.push((i + 1, j + 1)); } - if token.starts_with('-') { + if token.starts_with('-') || is_single_letter_slash_flag(token) { // An unrelated flag is skippable — alone, and (when it could take // a separate value) together with the token after it. Consuming it - // as a rule token above takes priority, so a rule that names a flag - // (`cargo test --danger`) still matches it. + // as a rule token above takes priority, so a rule that names a + // flag (`cargo test --danger`) still matches it. cmd.exe spells + // its flags the same way shells spell paths, so only the + // single-letter shape (`/f`, `/s`, `/q`, `/y`) may skip; anything + // longer is a POSIX path (`/tmp`, `/etc`, `/usr`, `/dev`) and must + // stay positional, or `cp /tmp/new_key ~/.ssh/authorized_keys` + // would slip past a rule guarding `~/.ssh/authorized_keys`. stack.push((i + 1, j)); if !token.contains('=') { stack.push((i + 2, j)); @@ -796,6 +881,18 @@ fn denied_prefix_matches(rule: &str, command: &str) -> bool { false } +/// True for a cmd.exe-style single-letter flag such as `/f`, `/s`, `/q`, `/y`. +/// +/// cmd.exe flags are a slash plus exactly one letter (`del /f /s /q`, `xcopy +/// /e /y`), so only that shape may skip like a `-` flag. The narrowness is +/// load-bearing: multi-character `/`-tokens are real POSIX paths (`/tmp`, +/// `/etc`, `/usr`, `/dev`) and must keep matching positionally. Case needs no +/// handling here — `normalize_command` has already lowercased the token. +fn is_single_letter_slash_flag(token: &str) -> bool { + let bytes = token.as_bytes(); + bytes.len() == 2 && bytes[0] == b'/' && bytes[1].is_ascii_alphabetic() +} + /// Whether a command word matches a deny rule's command word. /// /// Exact first, then the command's basename — `/bin/rm`, `./rm`, and @@ -803,6 +900,13 @@ fn denied_prefix_matches(rule: &str, command: &str) -> bool { /// direction only: a rule that spells a path (`/usr/bin/rm`) still requires /// that path, because the rule author asked for it specifically. Both /// separators are honored so a Windows spelling cannot slip past. +/// +/// A trailing `.exe` on the command's basename also folds: Windows spells the +/// same binary `cat.exe` or `C:\Windows\System32\cat.exe`, and a `cat +/// ~/.ssh/id_rsa` rule must hold against that spelling too. The fold is one +/// direction only — when the RULE itself ends in `.exe` (`control.exe`) it +/// keeps requiring that spelling, and `catalog` never matches `cat` because +/// only a whole `.exe` suffix strips, never a prefix. fn command_word_matches(rule_token: &str, command_token: &str) -> bool { if command_token == rule_token { return true; @@ -811,10 +915,15 @@ fn command_word_matches(rule_token: &str, command_token: &str) -> bool { if rule_token.contains('/') || rule_token.contains('\\') { return false; } - let basename = command_token + let mut basename = command_token .rsplit(['/', '\\']) .next() .unwrap_or(command_token); + if !rule_token.ends_with(".exe") + && let Some(stem) = basename.strip_suffix(".exe") + { + basename = stem; + } !basename.is_empty() && basename == rule_token } @@ -1029,6 +1138,33 @@ fn is_windows_absolute_path(value: &str) -> bool { bytes.len() >= 3 && bytes[0].is_ascii_alphabetic() && bytes[1] == b':' && bytes[2] == b'/' } +/// Exact-match fallback for a typed path rule that names an ABSOLUTE path. +/// +/// The primary match normalizes both sides to workspace-relative form, which +/// only succeeds when the call lives inside the workspace — so a rule pinning +/// a location outside it (a real home, `/root`, another user's home, or a +/// literal `~` spelling the tool passed through unexpanded) could never match. +/// This fallback fires only when workspace normalization failed on either +/// side, and only for a ROOTED rule (leading `/`, `~`, or a Windows drive): +/// separators fold to `/`, case folds on case-insensitive platforms, and the +/// comparison is plain equality. A relative rule never reaches it, so +/// workspace-relative semantics are unchanged, and because there are no +/// wildcards the deny direction keeps its precision while the allow direction +/// can only ever match the exact path the rule spells. +fn absolute_path_rule_matches(rule_path: &str, call_path: &str) -> bool { + let fold = |value: &str| { + let value = value.trim().replace('\\', "/"); + if platform_paths_are_case_insensitive() { + value.to_ascii_lowercase() + } else { + value + } + }; + let rule = fold(rule_path); + let rooted = rule.starts_with('/') || rule.starts_with("~/") || is_windows_absolute_path(&rule); + rooted && rule == fold(call_path) +} + fn has_windows_drive_prefix(value: &str) -> bool { let bytes = value.as_bytes(); bytes.len() >= 2 && bytes[0].is_ascii_alphabetic() && bytes[1] == b':' @@ -1398,6 +1534,232 @@ mod tests { assert!(allowed.allow, "rmdir must not be denied: {allowed:?}"); } + #[test] + fn denied_prefix_skips_cmd_exe_single_letter_slash_flags() { + // cmd.exe spells its flags `/f`, `/s`, `/q` — a slash plus exactly one + // letter, in any order and position. A deny rule must hold against + // every interleaving (`del /f /s /q`, `del /q /s /f`, ...); the app + // would otherwise have to enumerate canonical flag sequences, so the + // engine skips the shape itself, like `-` flags. + let engine = ExecPolicyEngine::new( + vec![], + vec![ + r"del c:\users\x\file".to_string(), + r"xcopy c:\src d:\dst".to_string(), + ], + ); + for command in [ + r"del c:\users\x\file", + r"del /f c:\users\x\file", + r"del /f /s /q c:\users\x\file", + r"del /q /s /f c:\users\x\file", + r"del /f c:\users\x\file /s /q", + r"xcopy /e /y c:\src d:\dst", + ] { + let decision = engine.check(ctx(command, AskForApproval::Never)).unwrap(); + assert!( + !decision.allow, + "cmd.exe flag spelling evaded deny: {command:?} -> {decision:?}" + ); + } + // A rule that NAMES a `/x` flag still consumes it as a rule token — + // the rule-token branch is tried before the skip branches. + let named = ExecPolicyEngine::new(vec![], vec![r"del /q c:\x".to_string()]); + let decision = named + .check(ctx(r"del /q c:\x", AskForApproval::Never)) + .unwrap(); + assert!( + !decision.allow, + "rule naming a slash flag missed: {decision:?}" + ); + } + + #[test] + fn denied_prefix_slash_skipping_keeps_multi_char_slash_tokens_positional() { + // The single-letter constraint is load-bearing: `/tmp` is a POSIX + // directory, not a flag. If multi-character `/`-tokens skipped, an + // exfil command could hide its real operand behind a skipped path and + // slip past a rule guarding the sensitive target. + let engine = ExecPolicyEngine::new(vec![], vec!["cp ~/.ssh/authorized_keys".to_string()]); + for command in [ + "cp /tmp/new_key ~/.ssh/authorized_keys", + "cp /etc/passwd ~/.ssh/authorized_keys", + ] { + let decision = engine + .check(ctx(command, AskForApproval::UnlessTrusted)) + .unwrap(); + assert!( + decision.allow, + "POSIX path argument wrongly treated as a flag: {command:?} -> {decision:?}" + ); + } + // The guarded target itself still denies, skip branches or not. + let denied = engine + .check(ctx( + "cp ~/.ssh/authorized_keys ~/.ssh/authorized_keys.bak", + AskForApproval::Never, + )) + .unwrap(); + assert!(!denied.allow, "guarded target must stay denied: {denied:?}"); + } + + #[test] + fn denied_prefix_middle_wildcard_matches_zero_or_more_tokens() { + // A rule token of exactly `*` matches zero or more consecutive command + // tokens REGARDLESS of shape — flags, flag values, extra positionals — + // so a rule can anchor on its sensitive tail without the app + // enumerating every flag spelling. + let engine = ExecPolicyEngine::new( + vec![], + vec![ + "grep * ~/.ssh/id_rsa".to_string(), + "dd * of=/dev/sda".to_string(), + ], + ); + for command in [ + "grep root ~/.ssh/id_rsa", + "grep -i root ~/.ssh/id_rsa", + "grep -r root ~/.ssh/id_rsa", + // The wildcard matches nothing at all. + "grep ~/.ssh/id_rsa", + // `dd` has no dash flags at all: its operands are `key=value`. + "dd if=/dev/zero of=/dev/sda", + "dd if=boot.img bs=1M of=/dev/sda", + // Deny rules are prefix matches: the anchored tail still denies + // when the command continues past it. + "grep -i root ~/.ssh/id_rsa > /tmp/out", + ] { + let decision = engine.check(ctx(command, AskForApproval::Never)).unwrap(); + assert!( + !decision.allow, + "wildcard rule missed {command:?}: {decision:?}" + ); + } + + // A rule whose LAST token is `*` still matches a shorter command — + // prefix semantics, not suffix equality. + let trailing = ExecPolicyEngine::new(vec![], vec!["grep * ~/.ssh/id_rsa *".to_string()]); + for command in [ + "grep root ~/.ssh/id_rsa", + "grep -i root ~/.ssh/id_rsa backup", + ] { + let decision = trailing.check(ctx(command, AskForApproval::Never)).unwrap(); + assert!( + !decision.allow, + "trailing-wildcard rule missed {command:?}: {decision:?}" + ); + } + } + + #[test] + fn denied_prefix_wildcard_stays_anchored_on_the_tail_token() { + // The wildcard bridges the MIDDLE of a rule; it does not relax the + // tail. A rule is still a prefix match: when the tail token never + // appears in the segment, there is no deny — here or inside a chain. + let engine = ExecPolicyEngine::new(vec![], vec!["grep * /home/z".to_string()]); + for command in ["grep x /etc/y", "ls && grep x /etc/y"] { + let decision = engine + .check(ctx(command, AskForApproval::UnlessTrusted)) + .unwrap(); + assert!( + decision.allow, + "wildcard rule over-matched {command:?}: {decision:?}" + ); + } + // Chained segments are still scanned individually: a wildcard rule + // denies when its anchor appears in ANY segment, and does not leak + // across the chain boundary in either direction. + let chain = ExecPolicyEngine::new(vec![], vec!["grep * ~/.ssh/id_rsa".to_string()]); + let denied = chain + .check(ctx( + "echo hi && grep root ~/.ssh/id_rsa", + AskForApproval::Never, + )) + .unwrap(); + assert!(!denied.allow, "chained segment must still deny: {denied:?}"); + let shielded = chain + .check(ctx("grep x /etc/y && echo done", AskForApproval::Never)) + .unwrap(); + assert!( + shielded.allow, + "wildcard must not reach into unrelated segments: {shielded:?}" + ); + } + + #[test] + fn denied_prefix_leading_wildcard_follows_generic_wildcard_semantics() { + // Rules in practice anchor their first token, but a leading `*` is not + // an error: the generic DFS gives it the same two branches and it is + // never treated as a command word. Documented consequence of keeping + // the anchor at the rule's literal first token: the command word after + // a leading wildcard is matched exactly, so `/bin/rm` is NOT folded to + // `rm` for it. Rule authors should not start rules with `*`; this test + // only pins the behavior the generic DFS produces. + let engine = ExecPolicyEngine::new(vec![], vec!["* rm -rf /".to_string()]); + let bare = engine + .check(ctx("rm -rf /", AskForApproval::Never)) + .unwrap(); + assert!( + !bare.allow, + "leading-wildcard rule must match its bare spelling: {bare:?}" + ); + let path = engine + .check(ctx("/bin/rm -rf /", AskForApproval::Never)) + .unwrap(); + assert!( + path.allow, + "leading wildcard must not gain command-word folding: {path:?}" + ); + } + + #[test] + fn denied_prefix_folds_windows_exe_suffix_on_the_command_word() { + // Windows spells the same binary `cat.exe` or + // `C:\Windows\System32\cat.exe`; a `cat ~/.ssh/id_rsa` rule must hold + // against those spellings. The fold is one-directional: a rule that + // names `.exe` itself keeps requiring it, and only a WHOLE `.exe` + // suffix strips — `catalog` never becomes `cat`. + let engine = ExecPolicyEngine::new(vec![], vec!["cat ~/.ssh/id_rsa".to_string()]); + for command in [ + "cat ~/.ssh/id_rsa", + "cat.exe ~/.ssh/id_rsa", + "cat.EXE ~/.ssh/id_rsa", + r"C:\Windows\System32\cat.exe ~/.ssh/id_rsa", + ] { + let decision = engine.check(ctx(command, AskForApproval::Never)).unwrap(); + assert!( + !decision.allow, + "`.exe` spelling evaded deny: {command:?} -> {decision:?}" + ); + } + + // A rule ending in `.exe` must still require that spelling: the bare + // `control` is a different binary and must not match `control.exe`. + let control = ExecPolicyEngine::new(vec![], vec!["control.exe".to_string()]); + let spelled = control + .check(ctx("control.exe", AskForApproval::Never)) + .unwrap(); + assert!(!spelled.allow, "control.exe must be denied: {spelled:?}"); + let bare = control + .check(ctx("control", AskForApproval::UnlessTrusted)) + .unwrap(); + assert!( + bare.allow, + "bare `control` must not match rule `control.exe`: {bare:?}" + ); + + // Only a whole `.exe` suffix folds, never a word prefix. + for command in ["catalog ~/.ssh/id_rsa", "catalog.exe ~/.ssh/id_rsa"] { + let decision = engine + .check(ctx(command, AskForApproval::UnlessTrusted)) + .unwrap(); + assert!( + decision.allow, + "prefix word must not fold into the rule word: {command:?} -> {decision:?}" + ); + } + } + #[test] fn path_rules_respect_filesystem_case_sensitivity() { // #4725: on a case-sensitive filesystem `config/allowed.toml` and @@ -1893,6 +2255,162 @@ mod tests { assert!(decision.requires_approval); } + #[test] + fn typed_ask_absolute_path_rule_matches_absolute_call_outside_workspace() { + let engine = + ExecPolicyEngine::with_rulesets(vec![Ruleset::user(vec![], vec![]).with_ask_rules( + vec![ToolAskRule { + tool: "read_file".into(), + command: None, + command_exact: false, + path: Some("/root/.ssh/config".into()), + workspace: None, + action: PermissionAction::Deny, + }], + )]); + + // An absolute rule must reach a call outside the workspace that the + // workspace-relative normalization cannot express. + let decision = engine + .check(ExecPolicyContext { + command: "", + cwd: "/workspace", + tool: Some("read_file"), + path: Some("/root/.ssh/config"), + ask_for_approval: AskForApproval::OnFailure, + sandbox_mode: Some("workspace-write"), + }) + .unwrap(); + assert_eq!(decision.matched_action, Some(PermissionAction::Deny)); + + // A different absolute path must not match. + let decision = engine + .check(ExecPolicyContext { + command: "", + cwd: "/workspace", + tool: Some("read_file"), + path: Some("/root/.ssh/known_hosts"), + ask_for_approval: AskForApproval::OnFailure, + sandbox_mode: Some("workspace-write"), + }) + .unwrap(); + assert_eq!(decision.matched_rule, None); + + // The fallback is exact: a traversal spelling of the same file is a + // different token string and must stay unmatchable (the documented + // "traversal is never matchable" stance, pinned through the rooted + // fallback too). + let decision = engine + .check(ExecPolicyContext { + command: "", + cwd: "/workspace", + tool: Some("read_file"), + path: Some("/root/../root/.ssh/config"), + ask_for_approval: AskForApproval::OnFailure, + sandbox_mode: Some("workspace-write"), + }) + .unwrap(); + assert_eq!(decision.matched_rule, None); + } + + #[test] + fn typed_ask_literal_tilde_rule_matches_unexpanded_call_spelling() { + let engine = + ExecPolicyEngine::with_rulesets(vec![Ruleset::user(vec![], vec![]).with_ask_rules( + vec![ToolAskRule { + tool: "read_file".into(), + command: None, + command_exact: false, + path: Some("~/.ssh/config".into()), + workspace: None, + action: PermissionAction::Deny, + }], + )]); + + let decision = engine + .check(ExecPolicyContext { + command: "", + cwd: "/workspace", + tool: Some("read_file"), + path: Some("~/.ssh/config"), + ask_for_approval: AskForApproval::OnFailure, + sandbox_mode: Some("workspace-write"), + }) + .unwrap(); + assert_eq!(decision.matched_action, Some(PermissionAction::Deny)); + + // The tilde-rooted channel is exact as well: a traversal spelling of + // the same file must not match (never-matchable-traversal stance). + let decision = engine + .check(ExecPolicyContext { + command: "", + cwd: "/workspace", + tool: Some("read_file"), + path: Some("~/.ssh/../ssh/config"), + ask_for_approval: AskForApproval::OnFailure, + sandbox_mode: Some("workspace-write"), + }) + .unwrap(); + assert_eq!(decision.matched_rule, None); + } + + #[test] + fn typed_ask_relative_path_rule_still_rejects_absolute_call() { + // The absolute fallback is rooted-rule-only: a relative rule keeps + // its workspace-relative semantics and must not reach an absolute + // call path through it. + let engine = ExecPolicyEngine::with_rulesets(vec![ + Ruleset::user(vec![], vec![]) + .with_ask_rules(vec![ToolAskRule::file_path("edit_file", "src/a.rs")]), + ]); + + let decision = engine + .check(ExecPolicyContext { + command: "", + cwd: "/workspace", + tool: Some("edit_file"), + path: Some("/src/a.rs"), + ask_for_approval: AskForApproval::OnFailure, + sandbox_mode: Some("workspace-write"), + }) + .unwrap(); + assert_eq!(decision.matched_rule, None); + } + + #[test] + fn typed_ask_absolute_path_rule_folds_separators_and_case_on_windows() { + let engine = + ExecPolicyEngine::with_rulesets(vec![Ruleset::user(vec![], vec![]).with_ask_rules( + vec![ToolAskRule { + tool: "read_file".into(), + command: None, + command_exact: false, + path: Some("C:/Users/u/.aws/credentials".into()), + workspace: None, + action: PermissionAction::Deny, + }], + )]); + + let decision = engine + .check(ExecPolicyContext { + command: "", + cwd: r"C:\workspace", + tool: Some("read_file"), + path: Some(r"C:\Users\U\.AWS\credentials"), + ask_for_approval: AskForApproval::OnFailure, + sandbox_mode: Some("workspace-write"), + }) + .unwrap(); + // The rule folds `C:/Users/u/...` and the call folds `C:\Users\U\...` + // to the same form on a case-insensitive platform; on a + // case-sensitive one the case difference is a different file. + if platform_paths_are_case_insensitive() { + assert_eq!(decision.matched_action, Some(PermissionAction::Deny)); + } else { + assert_eq!(decision.matched_rule, None); + } + } + // ── deny / allow action tests ────────────────────────────────────────── #[test] diff --git a/crates/tui/src/lib.rs b/crates/tui/src/lib.rs index 7c68a9cbd1..bb47c5b261 100644 --- a/crates/tui/src/lib.rs +++ b/crates/tui/src/lib.rs @@ -12226,6 +12226,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(), @@ -12234,6 +12235,7 @@ async fn build_direct_workflow_tool( Some(event_tx), manager.clone(), ) + .with_fleet_governor(fleet_governor) .with_locale_tag( codewhale_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..9b0ef10042 --- /dev/null +++ b/crates/tui/src/tools/subagent/governor.rs @@ -0,0 +1,889 @@ +//! 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 + }; + // Defensive: a failed receive requires the queued sender to be + // dropped without sending — which requires the gate itself to be + // dropped, impossible while this future holds an `Arc` to it. + // Loop anyway so a future refactor that breaks that invariant + // degrades to re-queueing instead of unwrapping. + 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" + ); + }); + } + + /// The mirror case of the redispatch test: a waiter cancelled *before* + /// its grant was dispatched leaves a stale queue entry with a dead + /// receiver. The granter must skip that entry — disarming the already + /// built permit instead of dropping it, which would re-enter the gate + /// lock held by `wake_locked` — and the slot must stay usable. + #[test] + fn forkguard_dynamic_gate_skips_stale_queued_waiter_without_leaking_slot() { + 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 behind the holder" + ); + + // Cancel while the gate is full: no grant was ever dispatched, + // so the stale entry stays queued with a dead receiver. + waiter.abort(); + tokio::time::sleep(ms(20)).await; + + // Releasing the holder runs the granter over the stale entry. + drop(holder); + assert_eq!( + gate.available_permits(), + 1, + "cancelled queued waiter must neither swallow nor leak the slot" + ); + let permit = gate + .try_acquire() + .expect("slot usable after the stale entry is skipped"); + drop(permit); + }); + } + + /// Stress: concurrent acquire/release with aborts and capacity + /// oscillation through 0 (a pause). Whatever the interleaving, every + /// slot must come home — a lost wakeup or a leaked (never released) + /// permit leaves the gate short of full capacity at the end of a round, + /// and an over-granted permit keeps a slot alive after all owners are + /// gone. Both fail the drain assertion. + #[test] + fn forkguard_dynamic_gate_stress_drains_to_full_capacity_despite_aborts() { + let (_governor, gate) = RateLimitGovernor::new(4); + let rt = tokio::runtime::Builder::new_multi_thread() + .worker_threads(2) + .enable_time() + .build() + .expect("test runtime"); + rt.block_on(async move { + for round in 0..24usize { + // Start every round with live headroom, then briefly drop to + // 0 mid-round on every third round: the pause case keeps a + // full queue parked while nothing holds a permit. Capacity 0 + // is never left in place while joining — with nobody holding + // a permit a permanent 0 would deadlock the round by design, + // so the restore below is part of the scenario. + gate.set_capacity(1 + (round % 2)); + let mut handles = Vec::new(); + for i in 0..16u32 { + let g = std::sync::Arc::clone(&gate); + handles.push(tokio::spawn(async move { + let _permit = g.acquire().await; + tokio::time::sleep(ms(u64::from(i % 4))).await; + })); + } + // Abort every third task: some while still queued (stale + // queue entries), some already holding a permit (the + // drop-releases-and-rewakes path). + for handle in handles.iter().step_by(3) { + handle.abort(); + } + if round % 3 == 0 { + gate.set_capacity(0); + tokio::time::sleep(ms(2)).await; + } + gate.set_capacity(4); + for handle in handles { + // A task that cannot finish inside the budget means a + // lost wakeup, a leaked permit, or a slot swallowed by a + // stale entry — fail the round instead of hanging. + tokio::time::timeout(ms(2000), handle) + .await + .expect("task must finish: stuck rounds mean lost wakeups or leaked slots") + .ok(); + } + // Let straggler permit drops (cancelled waiter re-release) + // run before asserting the drain. + tokio::time::sleep(ms(5)).await; + assert_eq!( + gate.available_permits(), + 4, + "round {round}: gate must drain to full capacity despite aborts and pauses" + ); + } + }); + } +} diff --git a/crates/tui/src/tools/subagent/mod.rs b/crates/tui/src/tools/subagent/mod.rs index 5727a1d4d0..f7d31745c6 100644 --- a/crates/tui/src/tools/subagent/mod.rs +++ b/crates/tui/src/tools/subagent/mod.rs @@ -19,7 +19,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; @@ -87,6 +87,7 @@ use coord::{ pub mod advisor; pub mod coord; +mod governor; pub mod mailbox; mod naming; mod worktree; @@ -358,6 +359,15 @@ const SUBAGENT_SESSION_CLOSED_REASON: &str = "Interrupted: parent session closed #[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"; +/// While queued, probe the governor for a drained rate-limit window at this +/// period (see `acquire_queued_launch_permit`). +const LAUNCH_RECOVERY_PROBE_PERIOD: Duration = Duration::from_secs(5); +/// Placeholder probe period when the runtime has no governor: the probe +/// branch no-ops, and a queued child's wall-time deadline always fires first. +const LAUNCH_RECOVERY_PROBE_PERIOD_WITHOUT_GOVERNOR: Duration = Duration::from_secs(3600); /// #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 @@ -2612,6 +2622,14 @@ pub struct SubAgentRuntime { /// Durable approval evidence inherited from the parent session. Legacy /// runtimes that do not install a store cannot open child approval prompts. approval_receipt_store: Option>, + /// 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 { @@ -2673,6 +2691,12 @@ impl SubAgentRuntime { ), parent_can_prompt: false, approval_receipt_store: None, + // 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, } } @@ -2726,6 +2750,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 { @@ -2969,6 +3006,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, approval_mode: self.approval_mode, auto_review_policy: Arc::clone(&self.auto_review_policy), @@ -3382,7 +3422,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 @@ -3490,6 +3543,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(), @@ -3514,7 +3570,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, @@ -3530,12 +3587,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] @@ -4266,9 +4336,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, @@ -4286,13 +4358,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. @@ -6597,6 +6667,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()?; @@ -10325,7 +10401,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>, /// Releases the parent turn's settlement barrier after this turn-owned /// child or descendant has completed its terminal fan-in. _foreground_child_registration: Option, @@ -10415,9 +10491,9 @@ async fn run_subagent_task_inner(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)), @@ -10428,12 +10504,6 @@ async fn run_subagent_task_inner(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 - )); - } } } @@ -10523,28 +10593,69 @@ async fn run_subagent_task_inner(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() => { - 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(if task.runtime.governor.is_some() { + LAUNCH_RECOVERY_PROBE_PERIOD + } else { + LAUNCH_RECOVERY_PROBE_PERIOD_WITHOUT_GOVERNOR + }); + // Hold the acquire future across select iterations. Re-creating it on + // every probe tick would leave one stale queue entry per tick per queued + // child inside the gate (purged only by the next grant wave), which adds + // up over a long pause with a full swarm queue. + let mut acquire_permit = std::pin::pin!(gate.acquire()); + loop { + tokio::select! { + biased; + () = task.runtime.cancel_token.cancelled() => { + // No Cancelled progress event here: turn-end parking cancels + // queued children through the same token and must project as + // Interrupted-with-checkpoint, and genuine cancellation is + // reported by `run_subagent`'s first-step cancel check. + 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; the pinned `acquire_permit` + // below observes the grant on the next poll. + } + permit = &mut acquire_permit => { + 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, ); @@ -10553,16 +10664,13 @@ async fn record_queued_launch_progress(task: &SubAgentTask) { task.runtime.event_tx.as_ref(), &task.runtime.context.state_namespace, &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)); } } @@ -11173,8 +11281,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)), }); } @@ -11250,14 +11361,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 bfa5641157..c96cfb6605 100644 --- a/crates/tui/src/tools/subagent/tests.rs +++ b/crates/tui/src/tools/subagent/tests.rs @@ -13681,6 +13681,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, } } @@ -16004,7 +16007,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"); @@ -16021,12 +16023,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(), @@ -16154,7 +16155,6 @@ async fn launch_gate_queues_extra_direct_children() { #[tokio::test] async fn queued_turn_owned_child_parks_without_a_false_start_transition() { - use tokio::sync::Semaphore; use tokio_util::sync::CancellationToken; let tmp = tempdir().expect("tempdir"); @@ -16192,10 +16192,9 @@ async fn queued_turn_owned_child_parks_without_a_false_start_transition() { let registration = foreground_children .register(&agent_id, runtime.cancel_token.clone()) .expect("turn-owned queued child registers before settlement"); - 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 only launch permit"); let task = SubAgentTask { manager_handle: Arc::clone(&manager), @@ -16311,7 +16310,6 @@ async fn queued_turn_owned_child_parks_without_a_false_start_transition() { #[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); @@ -16343,10 +16341,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), diff --git a/crates/tui/src/tools/subagent/tests/roster_routes.rs b/crates/tui/src/tools/subagent/tests/roster_routes.rs index 6875ec06ef..6d65b9c239 100644 --- a/crates/tui/src/tools/subagent/tests/roster_routes.rs +++ b/crates/tui/src/tools/subagent/tests/roster_routes.rs @@ -853,7 +853,7 @@ async fn fleet_editor_save_reload_reaches_type_only_admission_without_a_model_re let client = DeepSeekClient::new(&reloaded).unwrap(); let manager = new_shared_subagent_manager(root.path().to_path_buf(), 1); let gate = manager.read().await.launch_gate.clone(); - let held_permit = gate.acquire_owned().await.unwrap(); + let held_permit = gate.acquire().await; let (mailbox, mut mailbox_rx) = Mailbox::new(CancellationToken::new()); let context = ToolContext::new(root.path()).with_state_namespace("fleet-editor-restarted"); let mut runtime = SubAgentRuntime::new( @@ -968,7 +968,7 @@ model = "deepseek-v4-pro" let client = DeepSeekClient::new(&config).unwrap(); let manager = new_shared_subagent_manager(root.path().to_path_buf(), 1); let gate = manager.read().await.launch_gate.clone(); - let held_permit = gate.acquire_owned().await.unwrap(); + let held_permit = gate.acquire().await; let (mailbox, mut mailbox_rx) = Mailbox::new(CancellationToken::new()); let context = ToolContext::new(root.path()).with_state_namespace("manual-role-restarted"); let mut runtime = SubAgentRuntime::new( @@ -1152,7 +1152,7 @@ model = "deepseek/deepseek-v4-flash" let client = DeepSeekClient::new(&config).unwrap(); let manager = new_shared_subagent_manager(root.path().to_path_buf(), 1); let gate = manager.read().await.launch_gate.clone(); - let held_permit = gate.acquire_owned().await.unwrap(); + let held_permit = gate.acquire().await; let (mailbox, mut mailbox_rx) = Mailbox::new(CancellationToken::new()); let context = ToolContext::new(root.path()).with_state_namespace(name); let mut runtime = SubAgentRuntime::new(