zeph_subagent/budget.rs
1// SPDX-FileCopyrightText: 2026 Andrei G <bug-ops>
2// SPDX-License-Identifier: MIT OR Apache-2.0
3
4//! Session-wide cumulative subagent spawn budget (issue #6545).
5
6use std::sync::atomic::{AtomicUsize, Ordering};
7
8use crate::error::SubAgentError;
9
10/// Session-wide cumulative counter of subagent spawns.
11///
12/// Bounds the *total* number of subagents spawned over a session's lifetime, independent of
13/// [`SubAgentManager::spawn`](crate::manager::SubAgentManager::spawn)'s existing
14/// `max_concurrent` (in-flight) and `max_spawn_depth` (recursion) guardrails — a shallow,
15/// low-concurrency but high-frequency sequential delegation loop trips neither of those.
16///
17/// # Ownership, not a shared handle
18///
19/// Deliberately a plain `AtomicUsize` newtype — no `Arc`, no `Clone`. Nothing in this crate
20/// needs a shared, cloned handle to a budget: `SubAgentManager` owns one instance as the origin
21/// of truth, and `zeph-core`'s `OrchestrationState` owns an independent fallback instance used
22/// only when no manager is wired. Both the manager-side spawn path and the ACP `/subagent
23/// spawn` chokepoint (which never touches `SubAgentManager` at all) reach whichever instance
24/// applies through an accessor (`Agent::session_budget` in `zeph-core`) that hands out a plain
25/// `&SessionSpawnBudget` reference, never a copy.
26///
27/// This also makes a would-be TOCTOU hazard structurally impossible rather than merely
28/// documented: such a hazard could only arise if `SubAgentManager` were ever shared (e.g.
29/// behind an `Arc`) across concurrent tasks, and a future refactor down that path would have to
30/// deliberately reintroduce `Clone`/`Arc` here — a change reviewable on its own, rather than one
31/// hiding behind an innocuous `.clone()` call at some unrelated call site.
32///
33/// # Concurrency
34///
35/// `AtomicUsize` provides interior mutability so [`check`](Self::check) and
36/// [`record_spawn`](Self::record_spawn) work through a shared `&self` reference (every access
37/// goes through `&self` accessors, never `&mut self`), satisfying NFR-001's atomic-counter
38/// requirement. Every call path that reaches either method — `SubAgentManager::spawn`/`resume`,
39/// the orchestration scheduler, and the ACP chokepoint — is serialized behind `&mut Agent` on
40/// the single agent task, so `Relaxed` ordering suffices. The check/consume split (budget is
41/// checked at the spawn guard but only consumed at the true commit point, so a rejected or
42/// transiently retried spawn never burns budget it never used) introduces no reachable race
43/// under that serialization.
44///
45/// # Examples
46///
47/// ```rust
48/// use zeph_subagent::SessionSpawnBudget;
49///
50/// let budget = SessionSpawnBudget::default();
51/// assert_eq!(budget.spawned(), 0);
52///
53/// budget.check(1).expect("budget not yet exhausted");
54/// budget.record_spawn();
55/// assert_eq!(budget.spawned(), 1);
56/// assert!(budget.check(1).is_err(), "cap of 1 must now be exhausted");
57///
58/// // `0` is the unlimited sentinel: check() always succeeds regardless of count.
59/// assert!(budget.check(0).is_ok());
60/// ```
61#[derive(Default)]
62pub struct SessionSpawnBudget(AtomicUsize);
63
64impl SessionSpawnBudget {
65 /// Check the budget without consuming it.
66 ///
67 /// `max == 0` is the unlimited sentinel and always succeeds, so callers never need to
68 /// duplicate the sentinel check themselves (mirrors
69 /// [`DelegationMode::permits_explicit`](zeph_config::DelegationMode::permits_explicit)'s
70 /// anti-drift rationale for a check shared across multiple chokepoints).
71 ///
72 /// # Errors
73 ///
74 /// Returns [`SubAgentError::SessionSpawnLimit`] when the cumulative spawn count has
75 /// already reached `max`.
76 pub fn check(&self, max: usize) -> Result<(), SubAgentError> {
77 if max == 0 {
78 return Ok(());
79 }
80 let spawned = self.0.load(Ordering::Relaxed);
81 if spawned >= max {
82 return Err(SubAgentError::SessionSpawnLimit { spawned, max });
83 }
84 Ok(())
85 }
86
87 /// Record a successful spawn, incrementing the cumulative count by one.
88 ///
89 /// Must be called only at a spawn's true commit point — see the check/consume split
90 /// described in the type-level concurrency note.
91 pub fn record_spawn(&self) {
92 self.0.fetch_add(1, Ordering::Relaxed);
93 }
94
95 /// Current cumulative spawn count.
96 #[must_use]
97 pub fn spawned(&self) -> usize {
98 self.0.load(Ordering::Relaxed)
99 }
100}
101
102impl std::fmt::Debug for SessionSpawnBudget {
103 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
104 f.debug_tuple("SessionSpawnBudget")
105 .field(&self.spawned())
106 .finish()
107 }
108}
109
110#[cfg(test)]
111mod tests {
112 use super::*;
113
114 #[test]
115 fn default_mints_independent_counters() {
116 let a = SessionSpawnBudget::default();
117 let b = SessionSpawnBudget::default();
118 a.record_spawn();
119 assert_eq!(a.spawned(), 1);
120 assert_eq!(
121 b.spawned(),
122 0,
123 "default() must not share state across instances"
124 );
125 }
126
127 #[test]
128 fn zero_is_unlimited_sentinel() {
129 let budget = SessionSpawnBudget::default();
130 for _ in 0..1000 {
131 budget.record_spawn();
132 }
133 assert!(budget.check(0).is_ok());
134 }
135
136 #[test]
137 fn check_does_not_consume() {
138 let budget = SessionSpawnBudget::default();
139 budget.check(5).unwrap();
140 budget.check(5).unwrap();
141 assert_eq!(budget.spawned(), 0, "check() must be read-only");
142 }
143
144 #[test]
145 fn cap_reached_returns_session_spawn_limit() {
146 let budget = SessionSpawnBudget::default();
147 budget.record_spawn();
148 let err = budget.check(1).unwrap_err();
149 assert!(matches!(
150 err,
151 SubAgentError::SessionSpawnLimit { spawned: 1, max: 1 }
152 ));
153 }
154
155 #[test]
156 fn debug_prints_count() {
157 let budget = SessionSpawnBudget::default();
158 budget.record_spawn();
159 let debug = format!("{budget:?}");
160 assert!(
161 debug.contains('1'),
162 "Debug output must surface the count: {debug}"
163 );
164 }
165}