返回 CodeWhale
turn_budget.rs
根目录 / crates / tui / src / core / engine / turn_budget.rs
1 //! Turn budgets shared by interactive hosts and headless execution.
2 //!
3 //! Model steps are uncapped by default. Explicit positive limits still
4 //! apply; a fixed step counter is not a measure of useful progress.
5 //! Cumulative wall-clock and per-step stream budgets remain finite.
6 //!
7 //! `EngineConfig` retains its integer representation for embedders:
8 //! `u32::MAX` represents no model-step limit. `TurnContext::step_limit`
9 //! resolves that representation to `None` before checking a ceiling or
10 //! emitting diagnostics. It is not a very large finite fallback.
11 //!
12 //! ## Honesty at the limit
13 //!
14 //! Hitting a budget is never a clean success. The step ceiling already ends
15 //! the turn as `TurnOutcomeStatus::Failed` with the limit named (or, when
16 //! the model still owes work, grants exactly one bounded final-report turn
17 //! first). [`TurnWallClock`] follows the same contract in `run_turn`.
18
19 use std::time::{Duration, Instant};
20
21 /// No model-step limit unless the caller configures one. This is the
22 /// compatibility representation, not a ceiling checked at `u32::MAX`.
23 pub const DEFAULT_MAX_MODEL_STEPS: u32 = u32::MAX;
24 /// Smallest accepted model-step ceiling. One step still lets the model
25 /// answer once.
26 pub const MIN_MAX_MODEL_STEPS: u32 = 1;
27 /// Largest explicitly configured model-step ceiling.
28 pub const MAX_MAX_MODEL_STEPS: u32 = 100_000;
29
30 /// Default cumulative per-turn wall-clock budget, in seconds.
31 ///
32 /// Measured across every model step of one turn, not per request. Time the
33 /// turn spends blocked on a human approval decision is excluded (see
34 /// [`TurnWallClock::begin_human_wait`]) so an unanswered prompt cannot
35 /// consume the budget.
36 pub const DEFAULT_TURN_WALL_CLOCK_SECS: u64 = 3_600;
37 /// Smallest accepted per-turn wall-clock budget. Below this a single slow
38 /// reasoning request would trip the budget before it could finish.
39 pub const MIN_TURN_WALL_CLOCK_SECS: u64 = 30;
40 /// Largest accepted per-turn wall-clock budget (24 hours).
41 pub const MAX_TURN_WALL_CLOCK_SECS: u64 = 86_400;
42
43 /// Default per-step cap on accumulated streamed content, in bytes.
44 /// Preserves the pre-R1 hard-coded value; R1 only makes it overridable.
45 pub const DEFAULT_STREAM_MAX_CONTENT_BYTES: usize = super::streaming::STREAM_MAX_CONTENT_BYTES;
46 /// Smallest accepted per-step stream content cap (64 KiB).
47 pub const MIN_STREAM_MAX_CONTENT_BYTES: usize = 64 * 1024;
48 /// Largest accepted per-step stream content cap (512 MiB).
49 pub const MAX_STREAM_MAX_CONTENT_BYTES: usize = 512 * 1024 * 1024;
50
51 /// Default per-step cap on a single stream's wall-clock duration, in
52 /// seconds. Preserves the pre-R1 hard-coded value.
53 pub const DEFAULT_STREAM_MAX_DURATION_SECS: u64 = super::streaming::STREAM_MAX_DURATION_SECS;
54 /// Smallest accepted per-step stream duration cap.
55 pub const MIN_STREAM_MAX_DURATION_SECS: u64 = 10;
56 /// Largest accepted per-step stream duration cap (24 hours).
57 pub const MAX_STREAM_MAX_DURATION_SECS: u64 = 86_400;
58
59 /// Resolve a configured model-step ceiling.
60 ///
61 /// `None` and `0` select the uncapped default. Explicit positive values
62 /// clamp into `MIN_MAX_MODEL_STEPS..=MAX_MAX_MODEL_STEPS`. Resolve raw input
63 /// once: passing an already resolved default back as an explicit value
64 /// would incorrectly install the maximum configurable ceiling.
65 #[must_use]
66 pub fn resolve_max_model_steps(raw: Option<u32>) -> u32 {
67 match raw {
68 None | Some(0) => DEFAULT_MAX_MODEL_STEPS,
69 Some(value) => value.clamp(MIN_MAX_MODEL_STEPS, MAX_MAX_MODEL_STEPS),
70 }
71 }
72
73 /// Resolve a configured per-turn wall-clock budget, in seconds.
74 ///
75 /// `None` and `0` both resolve to [`DEFAULT_TURN_WALL_CLOCK_SECS`]; `0` is
76 /// invalid, not "unlimited".
77 #[must_use]
78 pub fn resolve_turn_wall_clock_secs(raw: Option<u64>) -> u64 {
79 match raw {
80 None | Some(0) => DEFAULT_TURN_WALL_CLOCK_SECS,
81 Some(value) => value.clamp(MIN_TURN_WALL_CLOCK_SECS, MAX_TURN_WALL_CLOCK_SECS),
82 }
83 }
84
85 /// Resolve a configured per-turn wall-clock budget as a [`Duration`].
86 #[must_use]
87 pub fn resolve_turn_wall_clock(raw: Option<u64>) -> Duration {
88 Duration::from_secs(resolve_turn_wall_clock_secs(raw))
89 }
90
91 /// Resolve a configured per-step stream content cap, given megabytes.
92 ///
93 /// `None` and `0` both resolve to [`DEFAULT_STREAM_MAX_CONTENT_BYTES`].
94 #[must_use]
95 pub fn resolve_stream_max_content_bytes(raw_mb: Option<u64>) -> usize {
96 match raw_mb {
97 None | Some(0) => DEFAULT_STREAM_MAX_CONTENT_BYTES,
98 Some(mb) => usize::try_from(mb.saturating_mul(1024 * 1024))
99 .unwrap_or(MAX_STREAM_MAX_CONTENT_BYTES)
100 .clamp(MIN_STREAM_MAX_CONTENT_BYTES, MAX_STREAM_MAX_CONTENT_BYTES),
101 }
102 }
103
104 /// Resolve a configured per-step stream duration cap, in seconds.
105 ///
106 /// `None` and `0` both resolve to [`DEFAULT_STREAM_MAX_DURATION_SECS`].
107 #[must_use]
108 pub fn resolve_stream_max_duration_secs(raw: Option<u64>) -> u64 {
109 match raw {
110 None | Some(0) => DEFAULT_STREAM_MAX_DURATION_SECS,
111 Some(value) => value.clamp(MIN_STREAM_MAX_DURATION_SECS, MAX_STREAM_MAX_DURATION_SECS),
112 }
113 }
114
115 /// Cumulative wall-clock budget for one turn.
116 ///
117 /// Started once at the top of `Engine::run_turn` and checked at the
118 /// provider-request boundary, so a turn that trips the budget stops before
119 /// authorizing another billable request and keeps every tool result already
120 /// in the transcript.
121 ///
122 /// Time blocked on a human approval decision is excluded: the budget bounds
123 /// what the agent spends on its own, not how long a person takes to answer.
124 /// Without that exclusion an approval prompt left open overnight would fail
125 /// the turn — and discard the work the user just approved — the moment they
126 /// came back.
127 #[derive(Debug)]
128 pub(crate) struct TurnWallClock {
129 budget: Duration,
130 started_at: Instant,
131 /// Total time already excluded because the turn was blocked on a human.
132 excluded: Duration,
133 /// Set while currently blocked on a human decision.
134 blocked_since: Option<Instant>,
135 }
136
137 impl TurnWallClock {
138 /// Start a fresh budget. A zero budget is legal here (and only here):
139 /// it is how tests assert the stop path without sleeping. Configuration
140 /// never produces one — [`resolve_turn_wall_clock`] rejects `0`.
141 pub(crate) fn start(budget: Duration) -> Self {
142 Self {
143 budget,
144 started_at: Instant::now(),
145 excluded: Duration::ZERO,
146 blocked_since: None,
147 }
148 }
149
150 /// The budget this clock was started with.
151 pub(crate) fn budget(&self) -> Duration {
152 self.budget
153 }
154
155 /// Wall-clock time this turn has spent on its own work, excluding time
156 /// blocked on a human decision.
157 pub(crate) fn spent(&self) -> Duration {
158 let blocked_now = self
159 .blocked_since
160 .map_or(Duration::ZERO, |since| since.elapsed());
161 self.started_at
162 .elapsed()
163 .saturating_sub(self.excluded)
164 .saturating_sub(blocked_now)
165 }
166
167 /// Whether the cumulative budget is spent.
168 pub(crate) fn exhausted(&self) -> bool {
169 self.spent() >= self.budget
170 }
171
172 /// Stop counting: the turn is now waiting on a human decision.
173 /// Idempotent — a second call while already blocked does nothing, so a
174 /// nested or re-entered approval cannot double-exclude.
175 pub(crate) fn begin_human_wait(&mut self) {
176 if self.blocked_since.is_none() {
177 self.blocked_since = Some(Instant::now());
178 }
179 }
180
181 /// Resume counting after a human decision, banking the blocked time.
182 pub(crate) fn end_human_wait(&mut self) {
183 if let Some(since) = self.blocked_since.take() {
184 self.excluded = self.excluded.saturating_add(since.elapsed());
185 }
186 }
187
188 /// Test-only: pretend `elapsed` more wall-clock time has passed, so the
189 /// exhaustion path can be exercised without sleeping. An in-progress
190 /// human wait is rewound too — otherwise the simulated time would count
191 /// as agent-owned work that never actually happened.
192 #[cfg(test)]
193 pub(crate) fn rewind_for_test(&mut self, elapsed: Duration) {
194 self.started_at = self
195 .started_at
196 .checked_sub(elapsed)
197 .unwrap_or(self.started_at);
198 if let Some(since) = self.blocked_since {
199 self.blocked_since = Some(since.checked_sub(elapsed).unwrap_or(since));
200 }
201 }
202 }
203
204 #[cfg(test)]
205 mod tests {
206 use super::*;
207
208 #[test]
209 fn model_step_defaults_do_not_install_a_ceiling() {
210 assert_eq!(resolve_max_model_steps(None), DEFAULT_MAX_MODEL_STEPS);
211 assert_eq!(resolve_max_model_steps(Some(0)), DEFAULT_MAX_MODEL_STEPS);
212 assert_eq!(DEFAULT_MAX_MODEL_STEPS, u32::MAX);
213 const { assert!(MAX_MAX_MODEL_STEPS < u32::MAX) };
214 }
215
216 #[test]
217 fn model_step_ceiling_is_overridable_and_clamped() {
218 assert_eq!(resolve_max_model_steps(Some(7)), 7);
219 assert_eq!(resolve_max_model_steps(Some(1)), 1);
220 assert_eq!(
221 resolve_max_model_steps(Some(u32::MAX)),
222 MAX_MAX_MODEL_STEPS,
223 "explicit positive overrides retain the configured ceiling"
224 );
225 }
226
227 #[test]
228 fn turn_wall_clock_defaults_are_finite_and_reject_the_zero_sentinel() {
229 assert_eq!(
230 resolve_turn_wall_clock_secs(None),
231 DEFAULT_TURN_WALL_CLOCK_SECS
232 );
233 assert_eq!(
234 resolve_turn_wall_clock_secs(Some(0)),
235 DEFAULT_TURN_WALL_CLOCK_SECS
236 );
237 assert_eq!(
238 resolve_turn_wall_clock(None),
239 Duration::from_secs(DEFAULT_TURN_WALL_CLOCK_SECS)
240 );
241 }
242
243 #[test]
244 fn turn_wall_clock_is_overridable_and_clamped() {
245 assert_eq!(resolve_turn_wall_clock_secs(Some(120)), 120);
246 assert_eq!(
247 resolve_turn_wall_clock_secs(Some(1)),
248 MIN_TURN_WALL_CLOCK_SECS
249 );
250 assert_eq!(
251 resolve_turn_wall_clock_secs(Some(u64::MAX)),
252 MAX_TURN_WALL_CLOCK_SECS
253 );
254 }
255
256 #[test]
257 fn stream_caps_default_reject_zero_and_are_overridable() {
258 assert_eq!(
259 resolve_stream_max_content_bytes(None),
260 DEFAULT_STREAM_MAX_CONTENT_BYTES
261 );
262 assert_eq!(
263 resolve_stream_max_content_bytes(Some(0)),
264 DEFAULT_STREAM_MAX_CONTENT_BYTES
265 );
266 assert_eq!(resolve_stream_max_content_bytes(Some(1)), 1024 * 1024);
267 assert_eq!(
268 resolve_stream_max_content_bytes(Some(u64::MAX)),
269 MAX_STREAM_MAX_CONTENT_BYTES
270 );
271
272 assert_eq!(
273 resolve_stream_max_duration_secs(None),
274 DEFAULT_STREAM_MAX_DURATION_SECS
275 );
276 assert_eq!(
277 resolve_stream_max_duration_secs(Some(0)),
278 DEFAULT_STREAM_MAX_DURATION_SECS
279 );
280 assert_eq!(resolve_stream_max_duration_secs(Some(60)), 60);
281 assert_eq!(
282 resolve_stream_max_duration_secs(Some(u64::MAX)),
283 MAX_STREAM_MAX_DURATION_SECS
284 );
285 }
286
287 #[test]
288 fn wall_clock_exhausts_once_the_budget_is_spent() {
289 let mut clock = TurnWallClock::start(Duration::from_secs(60));
290 assert!(!clock.exhausted());
291 clock.rewind_for_test(Duration::from_secs(61));
292 assert!(clock.exhausted());
293 assert!(clock.spent() >= clock.budget());
294 }
295
296 #[test]
297 fn wall_clock_excludes_time_blocked_on_a_human_decision() {
298 let mut clock = TurnWallClock::start(Duration::from_secs(60));
299 clock.begin_human_wait();
300 // The human took two minutes; the agent spent none of its budget.
301 clock.rewind_for_test(Duration::from_secs(120));
302 assert!(
303 !clock.exhausted(),
304 "an unanswered approval prompt must not burn the turn budget"
305 );
306 clock.end_human_wait();
307 assert!(!clock.exhausted());
308 // Agent-owned time after the decision still counts.
309 clock.rewind_for_test(Duration::from_secs(61));
310 assert!(clock.exhausted());
311 }
312
313 #[test]
314 fn nested_human_waits_cannot_double_exclude() {
315 let mut clock = TurnWallClock::start(Duration::from_secs(60));
316 clock.begin_human_wait();
317 clock.begin_human_wait();
318 clock.rewind_for_test(Duration::from_secs(30));
319 clock.end_human_wait();
320 // A second unmatched end is a no-op, not another exclusion.
321 clock.end_human_wait();
322 clock.rewind_for_test(Duration::from_secs(61));
323 assert!(clock.exhausted());
324 }
325 }
326
326 lines RUST