perspt_sdk/search/
budget.rs

1//! Search limits, monotone usage, and the reservation discipline
2//! (PSP-10 Definition 5, system 20, Gate AC).
3//!
4//! Every branch action reserves resources before it starts and always
5//! consumes one nonrefundable action unit. Unused resource reservations
6//! may be released; consumed usage never decreases — there is no public
7//! decrement. A closed forest cannot reopen: closing consumes the budget
8//! by type.
9
10use serde::{Deserialize, Serialize};
11
12use crate::error::{Result, SdkError};
13
14/// The forest's finite limits (Definition 5). One set per forest.
15#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
16pub struct SearchLimits {
17    pub actions: u32,
18    pub model_turns: u32,
19    pub tool_calls: u32,
20    pub mutations: u32,
21    pub verifier_runs: u32,
22    pub tokens: u64,
23    pub elapsed_secs: u64,
24    pub result_bytes: u64,
25    /// Cumulative eager-copy file reservations across every fork.
26    pub workspace_files: u64,
27    /// Cumulative eager-copy byte reservations across every fork.
28    pub workspace_bytes: u64,
29}
30
31impl SearchLimits {
32    /// A finite release-default limit set (system 20's starting posture).
33    pub fn release_default() -> Self {
34        Self {
35            actions: 64,
36            model_turns: 24,
37            tool_calls: 192,
38            mutations: 96,
39            verifier_runs: 48,
40            tokens: 400_000,
41            elapsed_secs: 1_800,
42            result_bytes: 16 * 1024 * 1024,
43            workspace_files: 300_000,
44            workspace_bytes: 4 * 1024 * 1024 * 1024,
45        }
46    }
47}
48
49/// One action's resource request, reserved before the action starts.
50#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
51pub struct ReservationRequest {
52    pub model_turns: u32,
53    pub tool_calls: u32,
54    pub mutations: u32,
55    pub verifier_runs: u32,
56    pub tokens: u64,
57    pub result_bytes: u64,
58    pub workspace_files: u64,
59    pub workspace_bytes: u64,
60}
61
62/// Proof one reservation was taken; releasing unused amounts requires it.
63#[derive(Debug, Clone, PartialEq, Eq)]
64pub struct ReservationTicket {
65    request: ReservationRequest,
66}
67
68impl ReservationTicket {
69    pub fn request(&self) -> &ReservationRequest {
70        &self.request
71    }
72}
73
74/// Monotone usage. All fields are reserved-or-consumed totals; nothing
75/// public decrements except [`SearchUsage::release_unused`], which may only
76/// return part of a held ticket — never below what was actually consumed.
77#[derive(Debug, Clone, PartialEq, Eq, Default, Serialize, Deserialize)]
78pub struct SearchUsage {
79    pub actions: u32,
80    pub model_turns: u32,
81    pub tool_calls: u32,
82    pub mutations: u32,
83    pub verifier_runs: u32,
84    pub tokens: u64,
85    /// Wall-clock seconds charged against `SearchLimits::elapsed_secs`
86    /// (Definition 5's time dimension; monotone like every other field).
87    #[serde(default)]
88    pub elapsed_secs: u64,
89    pub result_bytes: u64,
90    pub workspace_files: u64,
91    pub workspace_bytes: u64,
92    closed: bool,
93}
94
95impl SearchUsage {
96    /// Reserve one action's resources. Consumes one action unit — action
97    /// units are never returned — and fails without reserving anything if
98    /// any limit would be exceeded or the forest is closed.
99    pub fn reserve(
100        &mut self,
101        limits: &SearchLimits,
102        request: ReservationRequest,
103    ) -> Result<ReservationTicket> {
104        if self.closed {
105            return Err(SdkError::Domain(
106                "search_closed: a closed forest cannot act".into(),
107            ));
108        }
109        let over = self.actions + 1 > limits.actions
110            || self.model_turns + request.model_turns > limits.model_turns
111            || self.tool_calls + request.tool_calls > limits.tool_calls
112            || self.mutations + request.mutations > limits.mutations
113            || self.verifier_runs + request.verifier_runs > limits.verifier_runs
114            || self.tokens + request.tokens > limits.tokens
115            || self.result_bytes + request.result_bytes > limits.result_bytes
116            || self.workspace_files + request.workspace_files > limits.workspace_files
117            || self.workspace_bytes + request.workspace_bytes > limits.workspace_bytes;
118        if over {
119            return Err(SdkError::Domain(
120                "search budget exceeded; the forest must close".into(),
121            ));
122        }
123        self.actions += 1;
124        self.model_turns += request.model_turns;
125        self.tool_calls += request.tool_calls;
126        self.mutations += request.mutations;
127        self.verifier_runs += request.verifier_runs;
128        self.tokens += request.tokens;
129        self.result_bytes += request.result_bytes;
130        self.workspace_files += request.workspace_files;
131        self.workspace_bytes += request.workspace_bytes;
132        Ok(ReservationTicket { request })
133    }
134
135    /// Charge consumed actuals a reservation could not know in advance
136    /// (observed turns, calls, mutations, verifier runs, tokens, bytes).
137    /// Monotone; crossing any limit closes the forest by type — the error
138    /// is terminal for further reservations.
139    pub fn charge(&mut self, limits: &SearchLimits, consumed: ReservationRequest) -> Result<()> {
140        self.model_turns += consumed.model_turns;
141        self.tool_calls += consumed.tool_calls;
142        self.mutations += consumed.mutations;
143        self.verifier_runs += consumed.verifier_runs;
144        self.tokens += consumed.tokens;
145        self.result_bytes += consumed.result_bytes;
146        self.workspace_files += consumed.workspace_files;
147        self.workspace_bytes += consumed.workspace_bytes;
148        self.ensure_within(limits)
149    }
150
151    /// Charge one wall-clock interval (Definition 5's time dimension).
152    pub fn charge_elapsed(&mut self, limits: &SearchLimits, seconds: u64) -> Result<()> {
153        self.elapsed_secs += seconds;
154        self.ensure_within(limits)
155    }
156
157    fn ensure_within(&mut self, limits: &SearchLimits) -> Result<()> {
158        let over = self.model_turns > limits.model_turns
159            || self.tool_calls > limits.tool_calls
160            || self.mutations > limits.mutations
161            || self.verifier_runs > limits.verifier_runs
162            || self.tokens > limits.tokens
163            || self.elapsed_secs > limits.elapsed_secs
164            || self.result_bytes > limits.result_bytes
165            || self.workspace_files > limits.workspace_files
166            || self.workspace_bytes > limits.workspace_bytes;
167        if over {
168            self.close();
169            return Err(SdkError::Domain(
170                "search budget exhausted by consumed actuals; the forest closed".into(),
171            ));
172        }
173        Ok(())
174    }
175
176    /// Release the unused part of a held reservation. The action unit is
177    /// never returned; each dimension may release at most what the ticket
178    /// reserved.
179    pub fn release_unused(&mut self, ticket: ReservationTicket, consumed: ReservationRequest) {
180        let held = ticket.request;
181        self.model_turns -= held
182            .model_turns
183            .saturating_sub(consumed.model_turns.min(held.model_turns));
184        self.tool_calls -= held
185            .tool_calls
186            .saturating_sub(consumed.tool_calls.min(held.tool_calls));
187        self.mutations -= held
188            .mutations
189            .saturating_sub(consumed.mutations.min(held.mutations));
190        self.verifier_runs -= held
191            .verifier_runs
192            .saturating_sub(consumed.verifier_runs.min(held.verifier_runs));
193        self.tokens -= held.tokens.saturating_sub(consumed.tokens.min(held.tokens));
194        self.result_bytes -= held
195            .result_bytes
196            .saturating_sub(consumed.result_bytes.min(held.result_bytes));
197        self.workspace_files -= held
198            .workspace_files
199            .saturating_sub(consumed.workspace_files.min(held.workspace_files));
200        self.workspace_bytes -= held
201            .workspace_bytes
202            .saturating_sub(consumed.workspace_bytes.min(held.workspace_bytes));
203    }
204
205    /// Close the forest. Terminal: no reservation succeeds afterwards.
206    pub fn close(&mut self) {
207        self.closed = true;
208    }
209
210    pub fn is_closed(&self) -> bool {
211        self.closed
212    }
213}
214
215/// Add with a ceiling; reports whether the ceiling truncated the add.
216fn add_capped_u32(current: &mut u32, extra: u32, cap: u32) -> bool {
217    let target = current.saturating_add(extra);
218    *current = target.min(cap);
219    target > cap
220}
221
222fn add_capped_u64(current: &mut u64, extra: u64, cap: u64) -> bool {
223    let target = current.saturating_add(extra);
224    *current = target.min(cap);
225    target > cap
226}
227
228/// One forest's usage behind a shared handle, so the branch tool loops can
229/// reserve before every action while the forest holds the same ledgered
230/// totals (Gate AC: reservations precede actions; usage never decreases;
231/// usage never persists above a limit — overflow clamps at the limit and
232/// closes the forest).
233#[derive(Debug, Clone)]
234pub struct SharedSearchBudget {
235    inner: std::sync::Arc<std::sync::Mutex<SearchUsage>>,
236    limits: SearchLimits,
237    denied: std::sync::Arc<std::sync::atomic::AtomicBool>,
238    /// The owning forest's `(forest_id, epoch)`, so every holder of the
239    /// handle can ledger a usage snapshot under the right identity —
240    /// crash-durable accounting needs snapshots at action granularity,
241    /// not only at branch boundaries.
242    identity: std::sync::Arc<std::sync::Mutex<(String, u64)>>,
243}
244
245impl SharedSearchBudget {
246    pub fn new(limits: SearchLimits) -> Self {
247        Self::with_usage(limits, SearchUsage::default())
248    }
249
250    /// Open with prior usage (resume: an interrupted forest's consumption
251    /// is not silently forgotten).
252    pub fn with_usage(limits: SearchLimits, usage: SearchUsage) -> Self {
253        Self {
254            inner: std::sync::Arc::new(std::sync::Mutex::new(usage)),
255            limits,
256            denied: std::sync::Arc::new(std::sync::atomic::AtomicBool::new(false)),
257            identity: std::sync::Arc::new(std::sync::Mutex::new((String::new(), 0))),
258        }
259    }
260
261    /// Bind the owning forest's identity (set once at forest opening).
262    pub fn bind_forest(&self, forest_id: &str) {
263        self.identity.lock().expect("identity lock").0 = forest_id.to_string();
264    }
265
266    /// Advance the identity's epoch (set at each frontier epoch opening).
267    pub fn set_epoch(&self, epoch: u64) {
268        self.identity.lock().expect("identity lock").1 = epoch;
269    }
270
271    /// The bound `(forest_id, epoch)`; the forest id is empty outside a
272    /// forest.
273    pub fn identity(&self) -> (String, u64) {
274        self.identity.lock().expect("identity lock").clone()
275    }
276
277    pub fn limits(&self) -> &SearchLimits {
278        &self.limits
279    }
280
281    /// Reserve one action's resources before it starts. A refusal marks
282    /// the handle denied so the forest can abandon the branch gracefully.
283    pub fn reserve(&self, request: ReservationRequest) -> Result<ReservationTicket> {
284        let mut usage = self.inner.lock().expect("budget lock");
285        usage.reserve(&self.limits, request).inspect_err(|_| {
286            self.denied.store(true, std::sync::atomic::Ordering::SeqCst);
287        })
288    }
289
290    /// Settle one action: release the unused part of the ticket; charge any
291    /// overshoot clamped at the limits. Overshoot closes the forest and
292    /// errs — the caller ledgers the overflow as an observation.
293    pub fn settle(&self, ticket: ReservationTicket, actual: ReservationRequest) -> Result<()> {
294        let mut usage = self.inner.lock().expect("budget lock");
295        let held = *ticket.request();
296        usage.release_unused(ticket, actual);
297        let overshoot = ReservationRequest {
298            model_turns: actual.model_turns.saturating_sub(held.model_turns),
299            tool_calls: actual.tool_calls.saturating_sub(held.tool_calls),
300            mutations: actual.mutations.saturating_sub(held.mutations),
301            verifier_runs: actual.verifier_runs.saturating_sub(held.verifier_runs),
302            tokens: actual.tokens.saturating_sub(held.tokens),
303            result_bytes: actual.result_bytes.saturating_sub(held.result_bytes),
304            workspace_files: actual.workspace_files.saturating_sub(held.workspace_files),
305            workspace_bytes: actual.workspace_bytes.saturating_sub(held.workspace_bytes),
306        };
307        if overshoot == ReservationRequest::default() {
308            return Ok(());
309        }
310        self.charge_clamped(&mut usage, overshoot)
311    }
312
313    fn charge_clamped(&self, usage: &mut SearchUsage, extra: ReservationRequest) -> Result<()> {
314        let limits = &self.limits;
315        let mut over = false;
316        over |= add_capped_u32(
317            &mut usage.model_turns,
318            extra.model_turns,
319            limits.model_turns,
320        );
321        over |= add_capped_u32(&mut usage.tool_calls, extra.tool_calls, limits.tool_calls);
322        over |= add_capped_u32(&mut usage.mutations, extra.mutations, limits.mutations);
323        over |= add_capped_u32(
324            &mut usage.verifier_runs,
325            extra.verifier_runs,
326            limits.verifier_runs,
327        );
328        over |= add_capped_u64(&mut usage.tokens, extra.tokens, limits.tokens);
329        over |= add_capped_u64(
330            &mut usage.result_bytes,
331            extra.result_bytes,
332            limits.result_bytes,
333        );
334        over |= add_capped_u64(
335            &mut usage.workspace_files,
336            extra.workspace_files,
337            limits.workspace_files,
338        );
339        over |= add_capped_u64(
340            &mut usage.workspace_bytes,
341            extra.workspace_bytes,
342            limits.workspace_bytes,
343        );
344        if over {
345            usage.close();
346            return Err(SdkError::Domain(
347                "consumed actuals crossed a search limit; charged to the limit and closed".into(),
348            ));
349        }
350        Ok(())
351    }
352
353    /// Charge one wall-clock interval, clamped at the limit; crossing it
354    /// closes the forest.
355    pub fn charge_elapsed(&self, seconds: u64) -> Result<()> {
356        let mut usage = self.inner.lock().expect("budget lock");
357        let over = add_capped_u64(&mut usage.elapsed_secs, seconds, self.limits.elapsed_secs);
358        if over {
359            usage.close();
360            return Err(SdkError::Domain(
361                "elapsed time crossed the search limit; the forest closed".into(),
362            ));
363        }
364        Ok(())
365    }
366
367    /// Non-consuming check: would this reservation be admitted right now?
368    pub fn headroom(&self, request: &ReservationRequest) -> bool {
369        let usage = self.inner.lock().expect("budget lock");
370        let limits = &self.limits;
371        !usage.is_closed()
372            && usage.actions < limits.actions
373            && usage.model_turns + request.model_turns <= limits.model_turns
374            && usage.tool_calls + request.tool_calls <= limits.tool_calls
375            && usage.verifier_runs + request.verifier_runs <= limits.verifier_runs
376            && usage.tokens + request.tokens <= limits.tokens
377    }
378
379    pub fn snapshot(&self) -> SearchUsage {
380        self.inner.lock().expect("budget lock").clone()
381    }
382
383    /// Terminal snapshot for `SearchClosed`; the forest cannot reopen.
384    pub fn close(&self) -> SearchUsage {
385        let mut usage = self.inner.lock().expect("budget lock");
386        usage.close();
387        usage.clone()
388    }
389
390    pub fn is_closed(&self) -> bool {
391        self.inner.lock().expect("budget lock").is_closed()
392    }
393
394    /// Whether a reservation was refused since the last call (one-shot).
395    pub fn take_denied(&self) -> bool {
396        self.denied.swap(false, std::sync::atomic::Ordering::SeqCst)
397    }
398}
399
400#[cfg(test)]
401mod tests {
402    use super::*;
403
404    fn limits() -> SearchLimits {
405        SearchLimits {
406            actions: 2,
407            model_turns: 2,
408            tool_calls: 4,
409            mutations: 4,
410            verifier_runs: 2,
411            tokens: 1_000,
412            elapsed_secs: 60,
413            result_bytes: 1_000,
414            workspace_files: 10,
415            workspace_bytes: 1_000,
416        }
417    }
418
419    #[test]
420    fn reservation_precedes_action_and_units_are_nonrefundable() {
421        let mut usage = SearchUsage::default();
422        let ticket = usage
423            .reserve(
424                &limits(),
425                ReservationRequest {
426                    tokens: 600,
427                    ..Default::default()
428                },
429            )
430            .unwrap();
431        assert_eq!(usage.actions, 1);
432        assert_eq!(usage.tokens, 600);
433        // Release everything unused: tokens return, the action unit stays.
434        usage.release_unused(ticket, ReservationRequest::default());
435        assert_eq!(usage.tokens, 0);
436        assert_eq!(usage.actions, 1, "action units are never returned");
437    }
438
439    #[test]
440    fn limits_refuse_without_partial_reservation() {
441        let mut usage = SearchUsage::default();
442        usage
443            .reserve(
444                &limits(),
445                ReservationRequest {
446                    tokens: 900,
447                    ..Default::default()
448                },
449            )
450            .unwrap();
451        let before = usage.clone();
452        assert!(usage
453            .reserve(
454                &limits(),
455                ReservationRequest {
456                    tokens: 200,
457                    ..Default::default()
458                },
459            )
460            .is_err());
461        assert_eq!(usage, before, "a refused reservation reserves nothing");
462    }
463
464    #[test]
465    fn consumed_actuals_and_elapsed_time_close_the_forest_at_the_limit() {
466        let mut usage = SearchUsage::default();
467        usage
468            .charge(
469                &limits(),
470                ReservationRequest {
471                    model_turns: 1,
472                    tool_calls: 2,
473                    tokens: 500,
474                    ..Default::default()
475                },
476            )
477            .unwrap();
478        usage.charge_elapsed(&limits(), 59).unwrap();
479        // The next interval crosses elapsed_secs = 60: terminal.
480        assert!(usage.charge_elapsed(&limits(), 2).is_err());
481        assert!(usage.is_closed(), "crossing a limit closes the forest");
482        assert!(usage
483            .reserve(&limits(), ReservationRequest::default())
484            .is_err());
485    }
486
487    #[test]
488    fn a_closed_forest_cannot_reopen() {
489        let mut usage = SearchUsage::default();
490        usage.close();
491        assert!(usage.is_closed());
492        assert!(usage
493            .reserve(&limits(), ReservationRequest::default())
494            .is_err());
495    }
496
497    #[test]
498    fn shared_budget_reserves_settles_and_never_persists_over_limit() {
499        let shared = SharedSearchBudget::new(limits());
500        let ticket = shared
501            .reserve(ReservationRequest {
502                model_turns: 1,
503                tokens: 400,
504                ..Default::default()
505            })
506            .unwrap();
507        // Actuals overshoot the ticket's tokens past the 1_000 limit:
508        // charged to the limit, closed, and reported.
509        let result = shared.settle(
510            ticket,
511            ReservationRequest {
512                model_turns: 1,
513                tokens: 1_500,
514                ..Default::default()
515            },
516        );
517        assert!(result.is_err());
518        let usage = shared.snapshot();
519        assert!(
520            usage.tokens <= limits().tokens,
521            "usage never sits above a limit"
522        );
523        assert!(shared.is_closed());
524        assert!(shared.reserve(ReservationRequest::default()).is_err());
525    }
526
527    #[test]
528    fn shared_budget_denial_is_observable_once() {
529        let shared = SharedSearchBudget::new(limits());
530        assert!(!shared.take_denied());
531        let _ = shared.reserve(ReservationRequest {
532            tokens: 5_000,
533            ..Default::default()
534        });
535        assert!(shared.take_denied(), "the refusal is observable");
536        assert!(!shared.take_denied(), "and one-shot");
537    }
538
539    #[test]
540    fn shared_budget_headroom_is_non_consuming() {
541        let shared = SharedSearchBudget::new(limits());
542        let probe = ReservationRequest {
543            model_turns: 1,
544            ..Default::default()
545        };
546        assert!(shared.headroom(&probe));
547        assert_eq!(shared.snapshot(), SearchUsage::default());
548    }
549
550    #[test]
551    fn the_action_limit_ends_search() {
552        let mut usage = SearchUsage::default();
553        for _ in 0..2 {
554            usage
555                .reserve(&limits(), ReservationRequest::default())
556                .unwrap();
557        }
558        assert!(usage
559            .reserve(&limits(), ReservationRequest::default())
560            .is_err());
561    }
562}
perspt_sdk/search/budget.rs

perspt_sdk/search/
budget.rs
