agtrace_engine/domain/
model.rs

1use crate::ContextWindowUsage;
2use chrono::{DateTime, Utc};
3use std::path::PathBuf;
4
5/// Real-time session state for live monitoring and watch operations.
6///
7/// Tracks cumulative metrics, token usage, and context window state
8/// as events stream in. Used by the watch service to provide live updates.
9#[derive(Debug, Clone)]
10pub struct SessionState {
11    /// Session UUID.
12    pub session_id: String,
13    /// Project root directory, if known.
14    pub project_root: Option<PathBuf>,
15    /// Path to the session's log file.
16    pub log_path: Option<PathBuf>,
17    /// Session start timestamp.
18    pub start_time: DateTime<Utc>,
19    /// Timestamp of the most recent event.
20    pub last_activity: DateTime<Utc>,
21    /// Model name/ID being used in this session.
22    pub model: Option<String>,
23    /// Maximum context window size for the model.
24    pub context_window_limit: Option<u64>,
25    /// Current cumulative token usage.
26    pub current_usage: ContextWindowUsage,
27    /// Current cumulative reasoning tokens (o1-style extended thinking).
28    pub current_reasoning_tokens: i32,
29    /// Count of errors encountered so far.
30    pub error_count: u32,
31    /// Total number of events processed.
32    pub event_count: usize,
33    /// Total number of user turns.
34    pub turn_count: usize,
35}
36
37impl SessionState {
38    pub fn new(
39        session_id: String,
40        project_root: Option<PathBuf>,
41        log_path: Option<PathBuf>,
42        start_time: DateTime<Utc>,
43    ) -> Self {
44        Self {
45            session_id,
46            project_root,
47            log_path,
48            start_time,
49            last_activity: start_time,
50            model: None,
51            context_window_limit: None,
52            current_usage: ContextWindowUsage::default(),
53            current_reasoning_tokens: 0,
54            error_count: 0,
55            event_count: 0,
56            turn_count: 0,
57        }
58    }
59
60    pub fn total_input_side_tokens(&self) -> i32 {
61        self.current_usage.input_tokens()
62    }
63
64    pub fn total_output_side_tokens(&self) -> i32 {
65        self.current_usage.output_tokens()
66    }
67
68    /// Get total tokens as type-safe TokenCount
69    pub fn total_tokens(&self) -> crate::TokenCount {
70        self.current_usage.total_tokens()
71    }
72
73    /// Get context limit as type-safe ContextLimit
74    pub fn context_limit(&self) -> Option<crate::ContextLimit> {
75        self.context_window_limit.map(crate::ContextLimit::new)
76    }
77}
78
79#[cfg(test)]
80mod tests {
81    use super::*;
82
83    #[test]
84    fn test_session_state_initialization() {
85        let state = SessionState::new("test-id".to_string(), None, None, Utc::now());
86
87        assert_eq!(state.session_id, "test-id");
88        assert!(state.current_usage.is_empty());
89        assert_eq!(state.current_reasoning_tokens, 0);
90        assert_eq!(state.error_count, 0);
91        assert_eq!(state.event_count, 0);
92        assert_eq!(state.turn_count, 0);
93    }
94
95    #[test]
96    fn test_session_state_token_snapshot() {
97        let mut state = SessionState::new("test-id".to_string(), None, None, Utc::now());
98
99        state.current_usage = ContextWindowUsage::from_raw(100, 0, 0, 50);
100        assert_eq!(state.total_input_side_tokens(), 100);
101        assert_eq!(state.total_output_side_tokens(), 50);
102        assert_eq!(state.total_tokens(), crate::TokenCount::new(150));
103
104        state.current_usage = ContextWindowUsage::from_raw(10, 0, 1000, 60);
105        assert_eq!(state.total_input_side_tokens(), 1010);
106        assert_eq!(state.total_output_side_tokens(), 60);
107        assert_eq!(state.total_tokens(), crate::TokenCount::new(1070));
108    }
109}