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

agent_code_lib/state/
mod.rs
