agtrace_engine/domain/
model.rs1use 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 pub fn total_tokens(&self) -> crate::TokenCount {
54 self.current_usage.total_tokens()
55 }
56
57 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}