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::pricing::calculate_cost_optional_split;
7use crate::SharedMessage;
8use crate::{Runtime, Session};
9
10/// Conversation state tracked by the engine.
11pub struct ConversationState {
12    pub session: Session,
13    pub api_messages: Vec<SharedMessage>,
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,
119            input_tokens,
120            output_tokens,
121            cache_read,
122            cache_creation,
123            cache_creation_5m,
124            cache_creation_1h,
125        );
126    }
127}
agent_engine/engine/session.rs

agent_engine/engine/
session.rs
