agent_code_lib/state/
mod.rs

1//! Application state management.
2//!
3//! Centralized state store for the session. Tracks conversation
4//! messages, active queries, costs, token usage, and UI state.
5
6use std::collections::HashMap;
7
8use crate::config::Config;
9use crate::llm::message::{Message, Usage};
10
11/// Global application state for the session.
12pub struct AppState {
13    /// Configuration snapshot.
14    pub config: Config,
15    /// Full conversation history.
16    pub messages: Vec<Message>,
17    /// Whether a query is currently in progress.
18    pub is_query_active: bool,
19    /// Accumulated token usage across all turns.
20    pub total_usage: Usage,
21    /// Total estimated cost in USD.
22    pub total_cost_usd: f64,
23    /// Number of agent turns completed.
24    pub turn_count: usize,
25    /// Current working directory.
26    pub cwd: String,
27    /// Per-model token usage.
28    pub model_usage: HashMap<String, Usage>,
29    /// Whether plan mode is active (read-only tools only).
30    pub plan_mode: bool,
31    /// Shared background task manager.
32    pub task_manager: std::sync::Arc<crate::services::background::TaskManager>,
33    /// Session ID for persistence.
34    pub session_id: String,
35}
36
37impl AppState {
38    pub fn new(config: Config) -> Self {
39        let cwd = std::env::current_dir()
40            .map(|p| p.display().to_string())
41            .unwrap_or_else(|_| ".".into());
42
43        Self {
44            config,
45            messages: Vec::new(),
46            is_query_active: false,
47            total_usage: Usage::default(),
48            total_cost_usd: 0.0,
49            turn_count: 0,
50            cwd,
51            model_usage: HashMap::new(),
52            plan_mode: false,
53            task_manager: std::sync::Arc::new(crate::services::background::TaskManager::new()),
54            session_id: crate::services::session::new_session_id(),
55        }
56    }
57
58    /// Record usage from a completed API call.
59    pub fn record_usage(&mut self, usage: &Usage, model: &str) {
60        self.total_usage.merge(usage);
61        self.model_usage
62            .entry(model.to_string())
63            .or_default()
64            .merge(usage);
65        self.total_cost_usd += estimate_cost(usage, model);
66    }
67
68    /// Push a message into the conversation history.
69    pub fn push_message(&mut self, msg: Message) {
70        self.messages.push(msg);
71    }
72
73    /// Get the conversation history.
74    pub fn history(&self) -> &[Message] {
75        &self.messages
76    }
77}
78
79/// Cost estimation using the per-model pricing database.
80fn estimate_cost(usage: &Usage, model: &str) -> f64 {
81    crate::services::pricing::calculate_cost(
82        model,
83        usage.input_tokens,
84        usage.output_tokens,
85        usage.cache_read_input_tokens,
86        usage.cache_creation_input_tokens,
87    )
88}
89
90#[cfg(test)]
91mod tests {
92    use super::*;
93
94    #[test]
95    fn test_new_state() {
96        let state = AppState::new(crate::config::Config::default());
97        assert!(!state.cwd.is_empty());
98        assert_eq!(state.turn_count, 0);
99        assert_eq!(state.total_cost_usd, 0.0);
100        assert!(state.messages.is_empty());
101    }
102
103    #[test]
104    fn test_push_message() {
105        let mut state = AppState::new(crate::config::Config::default());
106        state.push_message(crate::llm::message::user_message("hello"));
107        assert_eq!(state.messages.len(), 1);
108        assert_eq!(state.history().len(), 1);
109    }
110
111    #[test]
112    fn test_record_usage() {
113        let mut state = AppState::new(crate::config::Config::default());
114        let usage = Usage {
115            input_tokens: 1000,
116            output_tokens: 500,
117            ..Default::default()
118        };
119        state.record_usage(&usage, "claude-sonnet-4");
120        assert_eq!(state.total_usage.input_tokens, 1000);
121        assert_eq!(state.total_usage.output_tokens, 500);
122        assert!(state.total_cost_usd > 0.0);
123    }
124
125    #[test]
126    fn test_record_usage_accumulates() {
127        let mut state = AppState::new(crate::config::Config::default());
128        let u1 = Usage {
129            input_tokens: 100,
130            output_tokens: 50,
131            ..Default::default()
132        };
133        let u2 = Usage {
134            input_tokens: 200,
135            output_tokens: 30,
136            ..Default::default()
137        };
138        state.record_usage(&u1, "claude-sonnet-4");
139        state.record_usage(&u2, "claude-sonnet-4");
140        assert_eq!(state.total_usage.output_tokens, 80); // 50 + 30.
141    }
142
143    #[test]
144    fn test_model_usage_tracking() {
145        let mut state = AppState::new(crate::config::Config::default());
146        let u1 = Usage {
147            input_tokens: 100,
148            output_tokens: 50,
149            ..Default::default()
150        };
151        state.record_usage(&u1, "model-a");
152        state.record_usage(&u1, "model-b");
153        assert!(state.model_usage.contains_key("model-a"));
154        assert!(state.model_usage.contains_key("model-b"));
155    }
156}
agent_code_lib/state/mod.rs

agent_code_lib/state/
mod.rs
