agtrace_engine/domain/
model.rs

1use crate::ContextWindowUsage;
2use chrono::{DateTime, Utc};
3use std::path::PathBuf;
4
5#[derive(Debug, Clone)]
6pub struct SessionState {
7    pub session_id: String,
8    pub project_root: Option<PathBuf>,
9    pub log_path: Option<PathBuf>,
10    pub start_time: DateTime<Utc>,
11    pub last_activity: DateTime<Utc>,
12    pub model: Option<String>,
13    pub context_window_limit: Option<u64>,
14    pub current_usage: ContextWindowUsage,
15    pub current_reasoning_tokens: i32,
16    pub error_count: u32,
17    pub event_count: usize,
18    pub turn_count: usize,
19}
20
21impl SessionState {
22    pub fn new(
23        session_id: String,
24        project_root: Option<PathBuf>,
25        log_path: Option<PathBuf>,
26        start_time: DateTime<Utc>,
27    ) -> Self {
28        Self {
29            session_id,
30            project_root,
31            log_path,
32            start_time,
33            last_activity: start_time,
34            model: None,
35            context_window_limit: None,
36            current_usage: ContextWindowUsage::default(),
37            current_reasoning_tokens: 0,
38            error_count: 0,
39            event_count: 0,
40            turn_count: 0,
41        }
42    }
43
44    pub fn total_input_side_tokens(&self) -> i32 {
45        self.current_usage.input_tokens()
46    }
47
48    pub fn total_output_side_tokens(&self) -> i32 {
49        self.current_usage.output_tokens()
50    }
51
52    /// Get total tokens as type-safe TokenCount
53    pub fn total_tokens(&self) -> crate::TokenCount {
54        self.current_usage.total_tokens()
55    }
56
57    /// Get context limit as type-safe ContextLimit
58    pub fn context_limit(&self) -> Option<crate::ContextLimit> {
59        self.context_window_limit.map(crate::ContextLimit::new)
60    }
61}
62
63#[cfg(test)]
64mod tests {
65    use super::*;
66
67    #[test]
68    fn test_session_state_initialization() {
69        let state = SessionState::new("test-id".to_string(), None, None, Utc::now());
70
71        assert_eq!(state.session_id, "test-id");
72        assert!(state.current_usage.is_empty());
73        assert_eq!(state.current_reasoning_tokens, 0);
74        assert_eq!(state.error_count, 0);
75        assert_eq!(state.event_count, 0);
76        assert_eq!(state.turn_count, 0);
77    }
78
79    #[test]
80    fn test_session_state_token_snapshot() {
81        let mut state = SessionState::new("test-id".to_string(), None, None, Utc::now());
82
83        state.current_usage = ContextWindowUsage::from_raw(100, 0, 0, 50);
84        assert_eq!(state.total_input_side_tokens(), 100);
85        assert_eq!(state.total_output_side_tokens(), 50);
86        assert_eq!(state.total_tokens(), crate::TokenCount::new(150));
87
88        state.current_usage = ContextWindowUsage::from_raw(10, 0, 1000, 60);
89        assert_eq!(state.total_input_side_tokens(), 1010);
90        assert_eq!(state.total_output_side_tokens(), 60);
91        assert_eq!(state.total_tokens(), crate::TokenCount::new(1070));
92    }
93}