agent_engine/engine/
session.rs

1//! Engine-level session management — save, load, resume, clear.
2//!
3//! Owns the conversation state that both TUI and headless modes need:
4//! messages, token counts, cost, abort context.
5
6use crate::{Session, Runtime};
7use crate::pricing::calculate_cost_optional_split;
8use serde_json::Value;
9
10/// Conversation state tracked by the engine.
11pub struct ConversationState {
12    pub session: Session,
13    pub api_messages: Vec<Value>,
14    pub total_input_tokens: u64,
15    pub total_output_tokens: u64,
16    pub total_cache_read_tokens: u64,
17    pub total_cache_creation_tokens: u64,
18    pub session_cost: f64,
19    pub abort_context: Option<String>,
20    /// Message queued to send after current stream completes.
21    pub queued_message: Option<String>,
22    /// Events buffered during streaming — drained on stream completion.
23    pub pending_events: Vec<String>,
24}
25
26impl ConversationState {
27    /// Create a new conversation with a fresh session.
28    pub fn new(session: Session) -> Self {
29        Self {
30            session,
31            api_messages: Vec::new(),
32            total_input_tokens: 0,
33            total_output_tokens: 0,
34            total_cache_read_tokens: 0,
35            total_cache_creation_tokens: 0,
36            session_cost: 0.0,
37            abort_context: None,
38            queued_message: None,
39            pending_events: Vec::new(),
40        }
41    }
42
43    /// Create from a resumed session.
44    pub fn from_resumed(session: Session) -> Self {
45        Self {
46            api_messages: session.api_messages.clone(),
47            total_input_tokens: session.total_input_tokens,
48            total_output_tokens: session.total_output_tokens,
49            total_cache_read_tokens: 0,
50            total_cache_creation_tokens: 0,
51            session_cost: session.session_cost,
52            abort_context: session.abort_context.clone(),
53            queued_message: None,
54            pending_events: Vec::new(),
55            session,
56        }
57    }
58
59    /// Save the current conversation state to disk.
60    pub async fn save(&mut self) {
61        if self.api_messages.is_empty() {
62            return;
63        }
64        self.session.api_messages = self.api_messages.clone();
65        self.session.total_input_tokens = self.total_input_tokens;
66        self.session.total_output_tokens = self.total_output_tokens;
67        self.session.session_cost = self.session_cost;
68        self.session.abort_context = self.abort_context.clone();
69        self.session.updated_at = chrono::Utc::now();
70        self.session.auto_title();
71        if let Err(e) = self.session.save().await {
72            tracing::error!("Failed to save session: {}", e);
73        }
74    }
75
76    /// Clear the current session and start fresh.
77    pub async fn clear(&mut self, runtime: &Runtime) {
78        self.save().await;
79        self.api_messages.clear();
80        self.total_input_tokens = 0;
81        self.total_output_tokens = 0;
82        self.total_cache_read_tokens = 0;
83        self.total_cache_creation_tokens = 0;
84        self.session_cost = 0.0;
85        self.abort_context = None;
86        self.queued_message = None;
87        self.pending_events.clear();
88        self.session = Session::new(
89            runtime.model(),
90            runtime.thinking_level(),
91            runtime.system_prompt(),
92        );
93    }
94
95    /// Add usage from a model turn.
96    ///
97    /// `cache_creation_5m` / `cache_creation_1h` are the cache-write TTL
98    /// split. When either is present the cost uses the split rates
99    /// (5m: 1.25×, 1h: 2.0×); when both are `None` the aggregate
100    /// `cache_creation` is billed at the 5m rate (fail-cheap fallback).
101    #[allow(clippy::too_many_arguments)]
102    pub fn add_usage(
103        &mut self,
104        input_tokens: u64,
105        output_tokens: u64,
106        cache_read: u64,
107        cache_creation: u64,
108        cache_creation_5m: Option<u64>,
109        cache_creation_1h: Option<u64>,
110        model: &str,
111    ) {
112        self.total_input_tokens += input_tokens;
113        self.total_output_tokens += output_tokens;
114        self.total_cache_read_tokens += cache_read;
115        self.total_cache_creation_tokens += cache_creation;
116
117        self.session_cost += calculate_cost_optional_split(
118            model, input_tokens, output_tokens, cache_read,
119            cache_creation, cache_creation_5m, cache_creation_1h,
120        );
121    }
122
123    /// Estimate current token count (for compaction decisions).
124    pub fn estimate_tokens(&self) -> usize {
125        let mut total_chars = 0usize;
126        for msg in &self.api_messages {
127            if let Some(s) = msg["content"].as_str() {
128                total_chars += s.len();
129            } else if let Some(arr) = msg["content"].as_array() {
130                for block in arr {
131                    if let Some(s) = block["text"].as_str() {
132                        total_chars += s.len();
133                    }
134                    if let Some(s) = block["thinking"].as_str() {
135                        total_chars += s.len();
136                    }
137                    if let Some(s) = block["content"].as_str() {
138                        total_chars += s.len();
139                    } else if let Some(content_arr) = block["content"].as_array() {
140                        for inner in content_arr {
141                            if let Some(s) = inner["text"].as_str() {
142                                total_chars += s.len();
143                            }
144                        }
145                    }
146                    if let Some(input) = block.get("input") {
147                        total_chars += input.to_string().len();
148                    }
149                }
150            }
151        }
152        total_chars / 4
153    }
154}
agent_engine/engine/session.rs

agent_engine/engine/
session.rs
