nika_engine/runtime/
limit_tracker.rs

1//! LimitTracker - Runtime limit tracking for agent execution
2//!
3//! Tracks resource consumption during agent execution and checks against
4//! configured limits to enable cost control and graceful partial completion.
5//!
6//! ## Features
7//!
8//! - **Turn tracking**: Counts agentic loop iterations
9//! - **Token tracking**: Accumulates input + output tokens
10//! - **Cost tracking**: Calculates USD cost based on provider pricing
11//! - **Duration tracking**: Wall-clock time from start
12//!
13//! ## Example
14//!
15//! ```rust,ignore
16//! use nika::runtime::LimitTracker;
17//! use nika::ast::LimitsConfig;
18//!
19//! let config = LimitsConfig {
20//!     max_turns: 20,
21//!     max_tokens: 50000,
22//!     max_cost_usd: 2.00,
23//!     max_duration_secs: 300,
24//!     ..Default::default()
25//! };
26//!
27//! let mut tracker = LimitTracker::new(config);
28//!
29//! // After each turn
30//! tracker.add_turn();
31//! tracker.add_tokens(500, 150);  // input, output
32//! tracker.add_cost(0.015);
33//!
34//! // Check limits
35//! if let Some(exceeded) = tracker.check_limits() {
36//!     println!("Limit exceeded: {:?}", exceeded);
37//! }
38//! ```
39
40use std::time::Instant;
41
42use crate::ast::limits::{LimitStatus, LimitType, LimitsConfig};
43
44// ═══════════════════════════════════════════════════════════════════════════
45// LimitTracker
46// ═══════════════════════════════════════════════════════════════════════════
47
48/// Runtime tracker for agent execution limits.
49///
50/// Accumulates resource consumption and checks against configured limits.
51/// Used by RigAgentLoop to enforce cost control and enable partial completion.
52#[derive(Debug, Clone)]
53pub struct LimitTracker {
54    /// Configuration with limit thresholds
55    config: LimitsConfig,
56
57    /// Current turn count
58    current_turns: u32,
59
60    /// Accumulated input tokens
61    input_tokens: u64,
62
63    /// Accumulated output tokens
64    output_tokens: u64,
65
66    /// Accumulated cost in USD
67    cost_usd: f64,
68
69    /// Start time for duration tracking
70    start_time: Instant,
71}
72
73impl LimitTracker {
74    /// Create a new tracker with the given configuration.
75    pub fn new(config: LimitsConfig) -> Self {
76        Self {
77            config,
78            current_turns: 0,
79            input_tokens: 0,
80            output_tokens: 0,
81            cost_usd: 0.0,
82            start_time: Instant::now(),
83        }
84    }
85
86    /// Create a tracker with no limits (unlimited execution).
87    pub fn unlimited() -> Self {
88        Self::new(LimitsConfig::default())
89    }
90
91    // ═══════════════════════════════════════════════════════════════════════
92    // Increment methods
93    // ═══════════════════════════════════════════════════════════════════════
94
95    /// Increment turn counter.
96    pub fn add_turn(&mut self) {
97        self.current_turns += 1;
98    }
99
100    /// Add tokens from a turn.
101    pub fn add_tokens(&mut self, input: u64, output: u64) {
102        self.input_tokens += input;
103        self.output_tokens += output;
104    }
105
106    /// Add cost from a turn.
107    pub fn add_cost(&mut self, cost: f64) {
108        self.cost_usd += cost;
109    }
110
111    /// Record a complete turn with all metrics.
112    pub fn record_turn(&mut self, input_tokens: u64, output_tokens: u64, cost: f64) {
113        self.add_turn();
114        self.add_tokens(input_tokens, output_tokens);
115        self.add_cost(cost);
116    }
117
118    // ═══════════════════════════════════════════════════════════════════════
119    // Getters
120    // ═══════════════════════════════════════════════════════════════════════
121
122    /// Get current turn count.
123    pub fn turns(&self) -> u32 {
124        self.current_turns
125    }
126
127    /// Get total tokens (input + output).
128    pub fn total_tokens(&self) -> u64 {
129        self.input_tokens + self.output_tokens
130    }
131
132    /// Get input tokens.
133    pub fn input_tokens(&self) -> u64 {
134        self.input_tokens
135    }
136
137    /// Get output tokens.
138    pub fn output_tokens(&self) -> u64 {
139        self.output_tokens
140    }
141
142    /// Get accumulated cost in USD.
143    pub fn cost_usd(&self) -> f64 {
144        self.cost_usd
145    }
146
147    /// Get elapsed duration in seconds.
148    pub fn duration_secs(&self) -> u64 {
149        self.start_time.elapsed().as_secs()
150    }
151
152    /// Get the configuration.
153    pub fn config(&self) -> &LimitsConfig {
154        &self.config
155    }
156
157    // ═══════════════════════════════════════════════════════════════════════
158    // Limit checking
159    // ═══════════════════════════════════════════════════════════════════════
160
161    /// Check all limits and return the first exceeded limit, if any.
162    ///
163    /// Returns `Some(LimitStatus)` if a limit is exceeded, `None` otherwise.
164    pub fn check_limits(&self) -> Option<LimitStatus> {
165        // Check in order of typical impact: turns, tokens, cost, duration
166        if let Some(status) = self.check_turns() {
167            if status.exceeded {
168                return Some(status);
169            }
170        }
171
172        if let Some(status) = self.check_tokens() {
173            if status.exceeded {
174                return Some(status);
175            }
176        }
177
178        if let Some(status) = self.check_cost() {
179            if status.exceeded {
180                return Some(status);
181            }
182        }
183
184        if let Some(status) = self.check_duration() {
185            if status.exceeded {
186                return Some(status);
187            }
188        }
189
190        None
191    }
192
193    /// Check turns limit.
194    pub fn check_turns(&self) -> Option<LimitStatus> {
195        if self.config.has_turns_limit() {
196            Some(LimitStatus::new(
197                LimitType::Turns,
198                self.current_turns as f64,
199                self.config.max_turns as f64,
200            ))
201        } else {
202            None
203        }
204    }
205
206    /// Check tokens limit.
207    pub fn check_tokens(&self) -> Option<LimitStatus> {
208        if self.config.has_tokens_limit() {
209            Some(LimitStatus::new(
210                LimitType::Tokens,
211                self.total_tokens() as f64,
212                self.config.max_tokens as f64,
213            ))
214        } else {
215            None
216        }
217    }
218
219    /// Check cost limit.
220    pub fn check_cost(&self) -> Option<LimitStatus> {
221        if self.config.has_cost_limit() {
222            Some(LimitStatus::new(
223                LimitType::Cost,
224                self.cost_usd,
225                self.config.max_cost_usd,
226            ))
227        } else {
228            None
229        }
230    }
231
232    /// Check duration limit.
233    pub fn check_duration(&self) -> Option<LimitStatus> {
234        if self.config.has_duration_limit() {
235            Some(LimitStatus::new(
236                LimitType::Duration,
237                self.duration_secs() as f64,
238                self.config.max_duration_secs as f64,
239            ))
240        } else {
241            None
242        }
243    }
244
245    /// Get status for all configured limits.
246    pub fn all_statuses(&self) -> Vec<LimitStatus> {
247        let mut statuses = Vec::new();
248
249        if let Some(s) = self.check_turns() {
250            statuses.push(s);
251        }
252        if let Some(s) = self.check_tokens() {
253            statuses.push(s);
254        }
255        if let Some(s) = self.check_cost() {
256            statuses.push(s);
257        }
258        if let Some(s) = self.check_duration() {
259            statuses.push(s);
260        }
261
262        statuses
263    }
264
265    /// Calculate progress as a percentage (0.0 - 1.0).
266    ///
267    /// Returns the maximum progress across all configured limits.
268    /// If no limits are configured, returns 0.0.
269    pub fn progress(&self) -> f64 {
270        self.all_statuses()
271            .iter()
272            .map(|s| s.usage_pct)
273            .fold(0.0, f64::max)
274    }
275
276    /// Check if any limit has been exceeded.
277    pub fn any_exceeded(&self) -> bool {
278        self.check_limits().is_some()
279    }
280
281    /// Check if we're approaching a limit (>80% usage).
282    pub fn approaching_limit(&self) -> bool {
283        self.all_statuses().iter().any(|s| s.usage_pct > 0.8)
284    }
285
286    // ═══════════════════════════════════════════════════════════════════════
287    // Reset
288    // ═══════════════════════════════════════════════════════════════════════
289
290    /// Reset all counters (keeps configuration).
291    pub fn reset(&mut self) {
292        self.current_turns = 0;
293        self.input_tokens = 0;
294        self.output_tokens = 0;
295        self.cost_usd = 0.0;
296        self.start_time = Instant::now();
297    }
298}
299
300impl Default for LimitTracker {
301    fn default() -> Self {
302        Self::unlimited()
303    }
304}
305
306// ═══════════════════════════════════════════════════════════════════════════
307// Tests
308// ═══════════════════════════════════════════════════════════════════════════
309
310#[cfg(test)]
311mod tests {
312    use super::*;
313    use crate::ast::limits::OnLimitReachedConfig;
314
315    fn config_with_turns(max: u32) -> LimitsConfig {
316        LimitsConfig {
317            max_turns: max,
318            ..Default::default()
319        }
320    }
321
322    fn config_with_tokens(max: u64) -> LimitsConfig {
323        LimitsConfig {
324            max_tokens: max,
325            ..Default::default()
326        }
327    }
328
329    fn config_with_cost(max: f64) -> LimitsConfig {
330        LimitsConfig {
331            max_cost_usd: max,
332            ..Default::default()
333        }
334    }
335
336    fn full_config() -> LimitsConfig {
337        LimitsConfig {
338            max_turns: 10,
339            max_tokens: 5000,
340            max_cost_usd: 1.00,
341            max_duration_secs: 60,
342            on_limit_reached: OnLimitReachedConfig::default(),
343        }
344    }
345
346    // ════════════════════════════════════════════════════════════════════
347    // Construction tests
348    // ════════════════════════════════════════════════════════════════════
349
350    #[test]
351    fn new_tracker_starts_at_zero() {
352        let tracker = LimitTracker::new(full_config());
353
354        assert_eq!(tracker.turns(), 0);
355        assert_eq!(tracker.total_tokens(), 0);
356        assert_eq!(tracker.input_tokens(), 0);
357        assert_eq!(tracker.output_tokens(), 0);
358        assert!((tracker.cost_usd() - 0.0).abs() < f64::EPSILON);
359    }
360
361    #[test]
362    fn unlimited_tracker_has_no_limits() {
363        let tracker = LimitTracker::unlimited();
364
365        assert!(!tracker.config().has_limits());
366        assert!(tracker.check_limits().is_none());
367    }
368
369    #[test]
370    fn default_is_unlimited() {
371        let tracker = LimitTracker::default();
372        assert!(!tracker.config().has_limits());
373    }
374
375    // ════════════════════════════════════════════════════════════════════
376    // Increment tests
377    // ════════════════════════════════════════════════════════════════════
378
379    #[test]
380    fn add_turn_increments() {
381        let mut tracker = LimitTracker::new(full_config());
382
383        tracker.add_turn();
384        assert_eq!(tracker.turns(), 1);
385
386        tracker.add_turn();
387        tracker.add_turn();
388        assert_eq!(tracker.turns(), 3);
389    }
390
391    #[test]
392    fn add_tokens_accumulates() {
393        let mut tracker = LimitTracker::new(full_config());
394
395        tracker.add_tokens(100, 50);
396        assert_eq!(tracker.input_tokens(), 100);
397        assert_eq!(tracker.output_tokens(), 50);
398        assert_eq!(tracker.total_tokens(), 150);
399
400        tracker.add_tokens(200, 100);
401        assert_eq!(tracker.input_tokens(), 300);
402        assert_eq!(tracker.output_tokens(), 150);
403        assert_eq!(tracker.total_tokens(), 450);
404    }
405
406    #[test]
407    fn add_cost_accumulates() {
408        let mut tracker = LimitTracker::new(full_config());
409
410        tracker.add_cost(0.10);
411        assert!((tracker.cost_usd() - 0.10).abs() < f64::EPSILON);
412
413        tracker.add_cost(0.25);
414        assert!((tracker.cost_usd() - 0.35).abs() < f64::EPSILON);
415    }
416
417    #[test]
418    fn record_turn_does_all() {
419        let mut tracker = LimitTracker::new(full_config());
420
421        tracker.record_turn(100, 50, 0.015);
422
423        assert_eq!(tracker.turns(), 1);
424        assert_eq!(tracker.input_tokens(), 100);
425        assert_eq!(tracker.output_tokens(), 50);
426        assert!((tracker.cost_usd() - 0.015).abs() < f64::EPSILON);
427    }
428
429    // ════════════════════════════════════════════════════════════════════
430    // Limit checking tests
431    // ════════════════════════════════════════════════════════════════════
432
433    #[test]
434    fn check_turns_not_exceeded() {
435        let mut tracker = LimitTracker::new(config_with_turns(10));
436
437        tracker.add_turn();
438        tracker.add_turn();
439        tracker.add_turn();
440
441        let status = tracker.check_turns().unwrap();
442        assert!(!status.exceeded);
443        assert!((status.current - 3.0).abs() < f64::EPSILON);
444        assert!((status.maximum - 10.0).abs() < f64::EPSILON);
445        assert!((status.usage_pct - 0.3).abs() < f64::EPSILON);
446    }
447
448    #[test]
449    fn check_turns_exceeded() {
450        let mut tracker = LimitTracker::new(config_with_turns(5));
451
452        for _ in 0..5 {
453            tracker.add_turn();
454        }
455
456        let status = tracker.check_turns().unwrap();
457        assert!(status.exceeded);
458        assert!((status.current - 5.0).abs() < f64::EPSILON);
459        assert!((status.usage_pct - 1.0).abs() < f64::EPSILON);
460    }
461
462    #[test]
463    fn check_tokens_not_exceeded() {
464        let mut tracker = LimitTracker::new(config_with_tokens(10000));
465
466        tracker.add_tokens(2000, 1000);
467
468        let status = tracker.check_tokens().unwrap();
469        assert!(!status.exceeded);
470        assert!((status.current - 3000.0).abs() < f64::EPSILON);
471        assert!((status.usage_pct - 0.3).abs() < f64::EPSILON);
472    }
473
474    #[test]
475    fn check_tokens_exceeded() {
476        let mut tracker = LimitTracker::new(config_with_tokens(5000));
477
478        tracker.add_tokens(3000, 2500); // 5500 total
479
480        let status = tracker.check_tokens().unwrap();
481        assert!(status.exceeded);
482    }
483
484    #[test]
485    fn check_cost_not_exceeded() {
486        let mut tracker = LimitTracker::new(config_with_cost(2.00));
487
488        tracker.add_cost(0.50);
489        tracker.add_cost(0.75);
490
491        let status = tracker.check_cost().unwrap();
492        assert!(!status.exceeded);
493        assert!((status.current - 1.25).abs() < f64::EPSILON);
494    }
495
496    #[test]
497    fn check_cost_exceeded() {
498        let mut tracker = LimitTracker::new(config_with_cost(1.00));
499
500        tracker.add_cost(0.80);
501        tracker.add_cost(0.30);
502
503        let status = tracker.check_cost().unwrap();
504        assert!(status.exceeded);
505    }
506
507    #[test]
508    fn check_limits_returns_first_exceeded() {
509        let config = LimitsConfig {
510            max_turns: 5,
511            max_tokens: 10000,
512            ..Default::default()
513        };
514        let mut tracker = LimitTracker::new(config);
515
516        // Exceed turns first
517        for _ in 0..6 {
518            tracker.add_turn();
519        }
520
521        let exceeded = tracker.check_limits().unwrap();
522        assert_eq!(exceeded.limit_type, LimitType::Turns);
523    }
524
525    #[test]
526    fn check_limits_none_when_ok() {
527        let mut tracker = LimitTracker::new(full_config());
528
529        tracker.add_turn();
530        tracker.add_tokens(100, 50);
531        tracker.add_cost(0.01);
532
533        assert!(tracker.check_limits().is_none());
534    }
535
536    // ════════════════════════════════════════════════════════════════════
537    // Progress and status tests
538    // ════════════════════════════════════════════════════════════════════
539
540    #[test]
541    fn progress_returns_max_usage() {
542        let config = LimitsConfig {
543            max_turns: 10,
544            max_tokens: 1000,
545            ..Default::default()
546        };
547        let mut tracker = LimitTracker::new(config);
548
549        tracker.add_turn(); // 10% turns
550        tracker.add_tokens(500, 0); // 50% tokens
551
552        let progress = tracker.progress();
553        assert!((progress - 0.5).abs() < f64::EPSILON);
554    }
555
556    #[test]
557    fn progress_zero_when_unlimited() {
558        let tracker = LimitTracker::unlimited();
559        assert!((tracker.progress() - 0.0).abs() < f64::EPSILON);
560    }
561
562    #[test]
563    fn all_statuses_returns_configured() {
564        let tracker = LimitTracker::new(full_config());
565        let statuses = tracker.all_statuses();
566
567        assert_eq!(statuses.len(), 4); // turns, tokens, cost, duration
568    }
569
570    #[test]
571    fn any_exceeded_true_when_over() {
572        let mut tracker = LimitTracker::new(config_with_turns(3));
573
574        assert!(!tracker.any_exceeded());
575
576        tracker.add_turn();
577        tracker.add_turn();
578        tracker.add_turn();
579
580        assert!(tracker.any_exceeded());
581    }
582
583    #[test]
584    fn approaching_limit_at_80_percent() {
585        let mut tracker = LimitTracker::new(config_with_turns(10));
586
587        for _ in 0..7 {
588            tracker.add_turn();
589        }
590        assert!(!tracker.approaching_limit()); // 70%
591
592        tracker.add_turn();
593        tracker.add_turn();
594        assert!(tracker.approaching_limit()); // 90%
595    }
596
597    // ════════════════════════════════════════════════════════════════════
598    // Reset tests
599    // ════════════════════════════════════════════════════════════════════
600
601    #[test]
602    fn reset_clears_counters() {
603        let mut tracker = LimitTracker::new(full_config());
604
605        tracker.add_turn();
606        tracker.add_tokens(100, 50);
607        tracker.add_cost(0.10);
608
609        tracker.reset();
610
611        assert_eq!(tracker.turns(), 0);
612        assert_eq!(tracker.total_tokens(), 0);
613        assert!((tracker.cost_usd() - 0.0).abs() < f64::EPSILON);
614    }
615
616    #[test]
617    fn reset_keeps_config() {
618        let config = config_with_turns(10);
619        let mut tracker = LimitTracker::new(config);
620
621        tracker.reset();
622
623        assert!(tracker.config().has_turns_limit());
624        assert_eq!(tracker.config().max_turns, 10);
625    }
626
627    // ════════════════════════════════════════════════════════════════════
628    // Edge-case: exactly at limit
629    // ════════════════════════════════════════════════════════════════════
630
631    #[test]
632    fn turns_exactly_at_limit_is_exceeded() {
633        let mut tracker = LimitTracker::new(config_with_turns(3));
634        tracker.add_turn();
635        tracker.add_turn();
636        tracker.add_turn();
637
638        let status = tracker.check_turns().unwrap();
639        assert!(status.exceeded, "exactly at max_turns should be exceeded");
640        assert!((status.current - 3.0).abs() < f64::EPSILON);
641        assert!((status.usage_pct - 1.0).abs() < f64::EPSILON);
642    }
643
644    #[test]
645    fn tokens_exactly_at_limit_is_exceeded() {
646        let mut tracker = LimitTracker::new(config_with_tokens(1000));
647        tracker.add_tokens(600, 400); // total = 1000
648
649        let status = tracker.check_tokens().unwrap();
650        assert!(status.exceeded, "exactly at max_tokens should be exceeded");
651        assert!((status.current - 1000.0).abs() < f64::EPSILON);
652        assert!((status.usage_pct - 1.0).abs() < f64::EPSILON);
653    }
654
655    #[test]
656    fn cost_exactly_at_limit_is_exceeded() {
657        let mut tracker = LimitTracker::new(config_with_cost(2.00));
658        tracker.add_cost(1.50);
659        tracker.add_cost(0.50);
660
661        let status = tracker.check_cost().unwrap();
662        assert!(
663            status.exceeded,
664            "exactly at max_cost_usd should be exceeded"
665        );
666        assert!((status.current - 2.0).abs() < f64::EPSILON);
667        assert!((status.usage_pct - 1.0).abs() < f64::EPSILON);
668    }
669
670    // ════════════════════════════════════════════════════════════════════
671    // Edge-case: one over limit
672    // ════════════════════════════════════════════════════════════════════
673
674    #[test]
675    fn turns_one_over_limit_is_exceeded() {
676        let mut tracker = LimitTracker::new(config_with_turns(5));
677        for _ in 0..6 {
678            tracker.add_turn();
679        }
680
681        let status = tracker.check_turns().unwrap();
682        assert!(status.exceeded);
683        assert!((status.current - 6.0).abs() < f64::EPSILON);
684        // usage_pct is capped at 1.0 by LimitStatus::new
685        assert!((status.usage_pct - 1.0).abs() < f64::EPSILON);
686    }
687
688    #[test]
689    fn tokens_one_over_limit_is_exceeded() {
690        let mut tracker = LimitTracker::new(config_with_tokens(5000));
691        tracker.add_tokens(3000, 2001); // total = 5001
692
693        let status = tracker.check_tokens().unwrap();
694        assert!(status.exceeded);
695        assert!((status.current - 5001.0).abs() < f64::EPSILON);
696    }
697
698    #[test]
699    fn cost_one_over_limit_is_exceeded() {
700        let mut tracker = LimitTracker::new(config_with_cost(1.00));
701        tracker.add_cost(1.01);
702
703        let status = tracker.check_cost().unwrap();
704        assert!(status.exceeded);
705    }
706
707    // ════════════════════════════════════════════════════════════════════
708    // Unlimited tracker never exceeds even under heavy usage
709    // ════════════════════════════════════════════════════════════════════
710
711    #[test]
712    fn unlimited_never_exceeds_under_heavy_usage() {
713        let mut tracker = LimitTracker::unlimited();
714
715        for _ in 0..1000 {
716            tracker.add_turn();
717        }
718        tracker.add_tokens(1_000_000, 500_000);
719        tracker.add_cost(999.99);
720
721        assert!(!tracker.any_exceeded());
722        assert!(tracker.check_turns().is_none());
723        assert!(tracker.check_tokens().is_none());
724        assert!(tracker.check_cost().is_none());
725        assert!(tracker.check_duration().is_none());
726        assert!(tracker.check_limits().is_none());
727        assert!((tracker.progress() - 0.0).abs() < f64::EPSILON);
728    }
729
730    // ════════════════════════════════════════════════════════════════════
731    // Unconfigured individual checks return None
732    // ════════════════════════════════════════════════════════════════════
733
734    #[test]
735    fn check_turns_none_when_not_configured() {
736        let tracker = LimitTracker::new(config_with_tokens(5000));
737        assert!(tracker.check_turns().is_none());
738    }
739
740    #[test]
741    fn check_tokens_none_when_not_configured() {
742        let tracker = LimitTracker::new(config_with_turns(10));
743        assert!(tracker.check_tokens().is_none());
744    }
745
746    #[test]
747    fn check_cost_none_when_not_configured() {
748        let tracker = LimitTracker::new(config_with_turns(10));
749        assert!(tracker.check_cost().is_none());
750    }
751
752    #[test]
753    fn check_duration_none_when_not_configured() {
754        let tracker = LimitTracker::new(config_with_turns(10));
755        assert!(tracker.check_duration().is_none());
756    }
757
758    // ════════════════════════════════════════════════════════════════════
759    // check_limits returns correct LimitType for non-turn exceeded
760    // ════════════════════════════════════════════════════════════════════
761
762    #[test]
763    fn check_limits_returns_tokens_when_tokens_exceeded() {
764        let mut tracker = LimitTracker::new(config_with_tokens(1000));
765        tracker.add_tokens(800, 300); // 1100 > 1000
766
767        let exceeded = tracker.check_limits().unwrap();
768        assert_eq!(exceeded.limit_type, LimitType::Tokens);
769    }
770
771    #[test]
772    fn check_limits_returns_cost_when_cost_exceeded() {
773        let mut tracker = LimitTracker::new(config_with_cost(0.50));
774        tracker.add_cost(0.75);
775
776        let exceeded = tracker.check_limits().unwrap();
777        assert_eq!(exceeded.limit_type, LimitType::Cost);
778    }
779
780    // ════════════════════════════════════════════════════════════════════
781    // check_limits priority: turns > tokens > cost > duration
782    // ════════════════════════════════════════════════════════════════════
783
784    #[test]
785    fn check_limits_turns_takes_priority_over_tokens() {
786        let config = LimitsConfig {
787            max_turns: 2,
788            max_tokens: 100,
789            ..Default::default()
790        };
791        let mut tracker = LimitTracker::new(config);
792        tracker.add_turn();
793        tracker.add_turn(); // turns exceeded
794        tracker.add_tokens(200, 0); // tokens also exceeded
795
796        let exceeded = tracker.check_limits().unwrap();
797        assert_eq!(exceeded.limit_type, LimitType::Turns);
798    }
799
800    #[test]
801    fn check_limits_tokens_takes_priority_over_cost() {
802        let config = LimitsConfig {
803            max_tokens: 100,
804            max_cost_usd: 0.01,
805            ..Default::default()
806        };
807        let mut tracker = LimitTracker::new(config);
808        tracker.add_tokens(200, 0); // tokens exceeded
809        tracker.add_cost(1.00); // cost also exceeded
810
811        let exceeded = tracker.check_limits().unwrap();
812        assert_eq!(exceeded.limit_type, LimitType::Tokens);
813    }
814
815    // ════════════════════════════════════════════════════════════════════
816    // approaching_limit edge case at exactly 80%
817    // ════════════════════════════════════════════════════════════════════
818
819    #[test]
820    fn approaching_limit_false_at_exactly_80_percent() {
821        let mut tracker = LimitTracker::new(config_with_turns(10));
822        for _ in 0..8 {
823            tracker.add_turn();
824        }
825        // 80% is NOT > 0.8, so this should be false
826        assert!(!tracker.approaching_limit());
827    }
828
829    #[test]
830    fn approaching_limit_true_at_81_percent() {
831        let mut tracker = LimitTracker::new(config_with_tokens(100));
832        tracker.add_tokens(81, 0); // 81%
833        assert!(tracker.approaching_limit());
834    }
835
836    // ════════════════════════════════════════════════════════════════════
837    // record_turn composes properly across multiple calls
838    // ════════════════════════════════════════════════════════════════════
839
840    #[test]
841    fn record_turn_accumulates_across_calls() {
842        let mut tracker = LimitTracker::new(full_config());
843
844        tracker.record_turn(100, 50, 0.01);
845        tracker.record_turn(200, 75, 0.02);
846        tracker.record_turn(150, 60, 0.015);
847
848        assert_eq!(tracker.turns(), 3);
849        assert_eq!(tracker.input_tokens(), 450);
850        assert_eq!(tracker.output_tokens(), 185);
851        assert_eq!(tracker.total_tokens(), 635);
852        assert!((tracker.cost_usd() - 0.045).abs() < f64::EPSILON);
853    }
854
855    // ════════════════════════════════════════════════════════════════════
856    // all_statuses only includes configured limits
857    // ════════════════════════════════════════════════════════════════════
858
859    #[test]
860    fn all_statuses_empty_for_unlimited() {
861        let tracker = LimitTracker::unlimited();
862        assert!(tracker.all_statuses().is_empty());
863    }
864
865    #[test]
866    fn all_statuses_single_when_one_configured() {
867        let tracker = LimitTracker::new(config_with_turns(10));
868        let statuses = tracker.all_statuses();
869        assert_eq!(statuses.len(), 1);
870        assert_eq!(statuses[0].limit_type, LimitType::Turns);
871    }
872
873    // ════════════════════════════════════════════════════════════════════
874    // Duration tests (basic - timing is hard to test)
875    // ════════════════════════════════════════════════════════════════════
876
877    #[test]
878    fn duration_starts_at_zero() {
879        let tracker = LimitTracker::new(full_config());
880        // Duration should be 0 or very small (< 1 sec)
881        assert!(tracker.duration_secs() < 1);
882    }
883
884    #[test]
885    fn duration_check_returns_status_when_configured() {
886        let config = LimitsConfig {
887            max_duration_secs: 300,
888            ..Default::default()
889        };
890        let tracker = LimitTracker::new(config);
891
892        let status = tracker.check_duration().unwrap();
893        assert!(!status.exceeded);
894        assert_eq!(status.limit_type, LimitType::Duration);
895        assert!((status.maximum - 300.0).abs() < f64::EPSILON);
896        // current should be near 0 since we just created it
897        assert!(status.current < 1.0);
898    }
899
900    #[test]
901    fn duration_not_exceeded_for_fresh_tracker() {
902        let config = LimitsConfig {
903            max_duration_secs: 1, // 1 second -- still won't be exceeded instantly
904            ..Default::default()
905        };
906        let tracker = LimitTracker::new(config);
907
908        // Freshly created tracker should not have exceeded even a 1s limit
909        let status = tracker.check_duration().unwrap();
910        assert!(!status.exceeded);
911    }
912
913    // ════════════════════════════════════════════════════════════════════
914    // Reset restarts duration timer
915    // ════════════════════════════════════════════════════════════════════
916
917    #[test]
918    fn reset_restarts_duration_timer() {
919        let mut tracker = LimitTracker::new(full_config());
920
921        tracker.add_turn();
922        tracker.add_tokens(100, 50);
923        tracker.add_cost(0.10);
924        tracker.reset();
925
926        // After reset, duration should be near 0 again
927        assert!(tracker.duration_secs() < 1);
928    }
929}
nika_engine/runtime/limit_tracker.rs

nika_engine/runtime/
limit_tracker.rs
