a3s_code_core/session/
mod.rs

1//! Session management
2//!
3//! Provides session-based conversation management:
4//! - Multiple independent sessions per agent
5//! - Conversation history tracking
6//! - Context usage monitoring
7//! - Per-session LLM client configuration
8//! - Session state management (Active, Paused, Completed, Error)
9//! - Per-session command queue with lane-based priority
10//! - Human-in-the-Loop (HITL) confirmation support
11//! - Session persistence (JSONL file storage)
12
13pub(crate) mod compaction;
14pub mod manager;
15
16pub use manager::SessionManager;
17
18#[cfg(test)]
19#[path = "tests.rs"]
20mod tests_file;
21
22use crate::agent::AgentEvent;
23use crate::hitl::{ConfirmationManager, ConfirmationPolicy, ConfirmationProvider};
24use crate::llm::{LlmClient, Message, TokenUsage, ToolDefinition};
25use crate::permissions::{PermissionChecker, PermissionDecision, PermissionPolicy};
26use crate::planning::Task;
27use crate::prompts::PlanningMode;
28use crate::queue::{ExternalTaskResult, LaneHandlerConfig, SessionQueueConfig};
29use crate::session_lane_queue::SessionLaneQueue;
30use crate::store::{LlmConfigData, SessionData};
31use anyhow::Result;
32use serde::{Deserialize, Serialize};
33use std::sync::Arc;
34use tokio::sync::{broadcast, RwLock};
35
36/// Session state
37#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
38pub enum SessionState {
39    #[default]
40    Unknown = 0,
41    Active = 1,
42    Paused = 2,
43    Completed = 3,
44    Error = 4,
45}
46
47/// Context usage statistics
48#[derive(Debug, Clone, Serialize, Deserialize)]
49pub struct ContextUsage {
50    pub used_tokens: usize,
51    pub max_tokens: usize,
52    pub percent: f32,
53    pub turns: usize,
54}
55
56impl Default for ContextUsage {
57    fn default() -> Self {
58        Self {
59            used_tokens: 0,
60            max_tokens: 200_000,
61            percent: 0.0,
62            turns: 0,
63        }
64    }
65}
66
67/// Default auto-compact threshold (80% of context window)
68pub const DEFAULT_AUTO_COMPACT_THRESHOLD: f32 = 0.80;
69
70/// Serde default function for auto_compact_threshold
71fn default_auto_compact_threshold() -> f32 {
72    DEFAULT_AUTO_COMPACT_THRESHOLD
73}
74
75/// Session configuration
76#[derive(Debug, Clone, Serialize, Deserialize)]
77pub struct SessionConfig {
78    pub name: String,
79    pub workspace: String,
80    pub system_prompt: Option<String>,
81    pub max_context_length: u32,
82    pub auto_compact: bool,
83    /// Context usage percentage threshold to trigger auto-compaction (0.0 - 1.0).
84    /// Only used when `auto_compact` is true. Default: 0.80 (80%).
85    #[serde(default = "default_auto_compact_threshold")]
86    pub auto_compact_threshold: f32,
87    /// Storage type for this session
88    #[serde(default)]
89    pub storage_type: crate::config::StorageBackend,
90    /// Queue configuration (optional, uses defaults if None)
91    #[serde(skip_serializing_if = "Option::is_none")]
92    pub queue_config: Option<SessionQueueConfig>,
93    /// Confirmation policy (optional, uses defaults if None)
94    #[serde(skip_serializing_if = "Option::is_none")]
95    pub confirmation_policy: Option<ConfirmationPolicy>,
96    /// Permission policy (optional, uses defaults if None)
97    #[serde(skip_serializing_if = "Option::is_none")]
98    pub permission_policy: Option<PermissionPolicy>,
99    /// Parent session ID (for subagent sessions)
100    #[serde(skip_serializing_if = "Option::is_none")]
101    pub parent_id: Option<String>,
102    /// Security configuration (optional, enables security features)
103    #[serde(skip_serializing_if = "Option::is_none")]
104    pub security_config: Option<crate::security::SecurityConfig>,
105    /// Shared hook engine for lifecycle events
106    #[serde(skip)]
107    pub hook_engine: Option<std::sync::Arc<dyn crate::hooks::HookExecutor>>,
108    /// Enable planning phase before execution
109    #[serde(default)]
110    pub planning_mode: PlanningMode,
111    /// Enable goal tracking
112    #[serde(default)]
113    pub goal_tracking: bool,
114}
115
116impl Default for SessionConfig {
117    fn default() -> Self {
118        Self {
119            name: String::new(),
120            workspace: String::new(),
121            system_prompt: None,
122            max_context_length: 0,
123            auto_compact: false,
124            auto_compact_threshold: DEFAULT_AUTO_COMPACT_THRESHOLD,
125            storage_type: crate::config::StorageBackend::default(),
126            queue_config: None,
127            confirmation_policy: None,
128            permission_policy: None,
129            parent_id: None,
130            security_config: None,
131            hook_engine: None,
132            planning_mode: PlanningMode::default(),
133            goal_tracking: false,
134        }
135    }
136}
137
138pub struct Session {
139    pub id: String,
140    pub config: SessionConfig,
141    pub state: SessionState,
142    pub messages: Vec<Message>,
143    pub context_usage: ContextUsage,
144    pub total_usage: TokenUsage,
145    /// Cumulative dollar cost for this session
146    pub total_cost: f64,
147    /// Model name for cost calculation
148    pub model_name: Option<String>,
149    pub tools: Vec<ToolDefinition>,
150    pub thinking_enabled: bool,
151    pub thinking_budget: Option<usize>,
152    /// Per-session LLM client (overrides default if set)
153    pub llm_client: Option<Arc<dyn LlmClient>>,
154    /// Creation timestamp (Unix epoch seconds)
155    pub created_at: i64,
156    /// Last update timestamp (Unix epoch seconds)
157    pub updated_at: i64,
158    /// Per-session command queue (a3s-lane backed)
159    pub command_queue: SessionLaneQueue,
160    /// HITL confirmation manager
161    pub confirmation_manager: Arc<dyn ConfirmationProvider>,
162    /// Permission checker for tool execution
163    pub permission_checker: Arc<dyn PermissionChecker>,
164    /// Event broadcaster for this session
165    event_tx: broadcast::Sender<AgentEvent>,
166    /// Context providers for augmenting prompts with external context
167    pub context_providers: Vec<Arc<dyn crate::context::ContextProvider>>,
168    /// Task list for tracking
169    pub tasks: Vec<Task>,
170    /// Parent session ID (for subagent sessions)
171    pub parent_id: Option<String>,
172    /// Agent memory system for this session (externally injected)
173    pub memory: Option<Arc<RwLock<crate::memory::AgentMemory>>>,
174    /// Current execution plan (if any)
175    pub current_plan: Arc<RwLock<Option<crate::planning::ExecutionPlan>>>,
176    /// Security guard (if enabled)
177    pub security_provider: Option<Arc<dyn crate::security::SecurityProvider>>,
178    /// Per-session tool execution metrics
179    pub tool_metrics: Arc<RwLock<crate::telemetry::ToolMetrics>>,
180    /// Per-call LLM cost records for cross-session aggregation
181    pub cost_records: Vec<crate::telemetry::LlmCostRecord>,
182}
183
184/// Validate that an identifier is safe for use in file paths.
185/// Rejects path traversal attempts and non-alphanumeric characters.
186fn validate_path_safe_id(id: &str, label: &str) -> Result<()> {
187    if id.is_empty() {
188        anyhow::bail!("{label} must not be empty");
189    }
190    // Allow alphanumeric, hyphens, underscores, and dots (but not leading dots)
191    let is_safe = id
192        .chars()
193        .all(|c| c.is_ascii_alphanumeric() || c == '-' || c == '_' || c == '.')
194        && !id.starts_with('.')
195        && !id.contains("..");
196    if !is_safe {
197        anyhow::bail!("{label} contains unsafe characters: {id:?}");
198    }
199    Ok(())
200}
201
202impl Session {
203    /// Create a new session (async due to SessionLaneQueue initialization)
204    pub async fn new(
205        id: String,
206        config: SessionConfig,
207        tools: Vec<ToolDefinition>,
208    ) -> Result<Self> {
209        // Validate session ID to prevent path traversal
210        validate_path_safe_id(&id, "Session ID")?;
211
212        let now = std::time::SystemTime::now()
213            .duration_since(std::time::UNIX_EPOCH)
214            .map(|d| d.as_secs() as i64)
215            .unwrap_or(0);
216
217        // Create event broadcaster
218        let (event_tx, _) = broadcast::channel(100);
219
220        // Create command queue with config or defaults
221        let queue_config = config.queue_config.clone().unwrap_or_default();
222        let command_queue = SessionLaneQueue::new(&id, queue_config, event_tx.clone()).await?;
223
224        // Create confirmation manager with policy or secure default (HITL enabled)
225        let confirmation_policy = config
226            .confirmation_policy
227            .clone()
228            .unwrap_or_else(ConfirmationPolicy::enabled);
229        let confirmation_manager = Arc::new(ConfirmationManager::new(
230            confirmation_policy,
231            event_tx.clone(),
232        ));
233
234        // Create permission checker with config or defaults
235        let permission_checker: Arc<dyn PermissionChecker> =
236            Arc::new(config.permission_policy.clone().unwrap_or_default());
237
238        // Extract parent_id from config
239        let parent_id = config.parent_id.clone();
240
241        // Memory is externally injected; default to None
242        let memory = None;
243
244        let context_providers: Vec<Arc<dyn crate::context::ContextProvider>> = vec![];
245
246        // Initialize empty plan
247        let current_plan = Arc::new(RwLock::new(None));
248
249        // Create security provider if security config is enabled
250        let security_provider: Option<Arc<dyn crate::security::SecurityProvider>> =
251            config.security_config.as_ref().and_then(|sc| {
252                if sc.enabled {
253                    Some(Arc::new(crate::security::NoOpSecurityProvider)
254                        as Arc<dyn crate::security::SecurityProvider>)
255                } else {
256                    None
257                }
258            });
259
260        Ok(Self {
261            id,
262            config,
263            state: SessionState::Active,
264            messages: Vec::new(),
265            context_usage: ContextUsage::default(),
266            total_usage: TokenUsage::default(),
267            total_cost: 0.0,
268            model_name: None,
269            tools,
270            thinking_enabled: false,
271            thinking_budget: None,
272            llm_client: None,
273            created_at: now,
274            updated_at: now,
275            command_queue,
276            confirmation_manager,
277            permission_checker,
278            event_tx,
279            context_providers,
280            tasks: Vec::new(),
281            parent_id,
282            memory,
283            current_plan,
284
285            security_provider,
286            tool_metrics: Arc::new(RwLock::new(crate::telemetry::ToolMetrics::new())),
287            cost_records: Vec::new(),
288        })
289    }
290
291    /// Check if this is a child session (has a parent)
292    pub fn is_child_session(&self) -> bool {
293        self.parent_id.is_some()
294    }
295
296    /// Get the parent session ID if this is a child session
297    pub fn parent_session_id(&self) -> Option<&str> {
298        self.parent_id.as_deref()
299    }
300
301    /// Get a receiver for session events
302    pub fn subscribe_events(&self) -> broadcast::Receiver<AgentEvent> {
303        self.event_tx.subscribe()
304    }
305
306    /// Get the event broadcaster
307    pub fn event_tx(&self) -> broadcast::Sender<AgentEvent> {
308        self.event_tx.clone()
309    }
310
311    /// Update the confirmation policy
312    pub async fn set_confirmation_policy(&self, policy: ConfirmationPolicy) {
313        self.confirmation_manager.set_policy(policy).await;
314    }
315
316    /// Update the permission policy used for tool execution checks.
317    pub fn set_permission_policy(&mut self, policy: PermissionPolicy) {
318        self.permission_checker = Arc::new(policy.clone());
319        self.config.permission_policy = Some(policy);
320    }
321
322    /// Get the current confirmation policy
323    pub async fn confirmation_policy(&self) -> ConfirmationPolicy {
324        self.confirmation_manager.policy().await
325    }
326
327    /// Check permission for a tool invocation
328    pub fn check_permission(
329        &self,
330        tool_name: &str,
331        args: &serde_json::Value,
332    ) -> PermissionDecision {
333        self.permission_checker.check(tool_name, args)
334    }
335
336    /// Add a context provider to the session
337    pub fn add_context_provider(&mut self, provider: Arc<dyn crate::context::ContextProvider>) {
338        self.context_providers.push(provider);
339    }
340
341    /// Remove a context provider by name
342    ///
343    /// Returns true if a provider was removed, false otherwise.
344    pub fn remove_context_provider(&mut self, name: &str) -> bool {
345        let initial_len = self.context_providers.len();
346        self.context_providers.retain(|p| p.name() != name);
347        self.context_providers.len() < initial_len
348    }
349
350    /// Get the names of all registered context providers
351    pub fn context_provider_names(&self) -> Vec<String> {
352        self.context_providers
353            .iter()
354            .map(|p| p.name().to_string())
355            .collect()
356    }
357
358    // ========================================================================
359    // Task Management
360    // ========================================================================
361
362    /// Get the current task list
363    pub fn get_tasks(&self) -> &[Task] {
364        &self.tasks
365    }
366
367    /// Set the task list (replaces entire list)
368    ///
369    /// Broadcasts a TaskUpdated event after updating.
370    pub fn set_tasks(&mut self, tasks: Vec<Task>) {
371        self.tasks = tasks.clone();
372        self.touch();
373
374        // Broadcast event
375        let _ = self.event_tx.send(AgentEvent::TaskUpdated {
376            session_id: self.id.clone(),
377            tasks,
378        });
379    }
380
381    /// Get count of active (non-completed, non-cancelled) tasks
382    pub fn active_task_count(&self) -> usize {
383        self.tasks.iter().filter(|t| t.is_active()).count()
384    }
385
386    /// Set handler mode for a lane
387    pub async fn set_lane_handler(
388        &self,
389        lane: crate::hitl::SessionLane,
390        config: LaneHandlerConfig,
391    ) {
392        self.command_queue.set_lane_handler(lane, config).await;
393    }
394
395    /// Get handler config for a lane
396    pub async fn get_lane_handler(&self, lane: crate::hitl::SessionLane) -> LaneHandlerConfig {
397        self.command_queue.get_lane_handler(lane).await
398    }
399
400    /// Complete an external task
401    pub async fn complete_external_task(&self, task_id: &str, result: ExternalTaskResult) -> bool {
402        self.command_queue
403            .complete_external_task(task_id, result)
404            .await
405    }
406
407    /// Get pending external tasks
408    pub async fn pending_external_tasks(&self) -> Vec<crate::queue::ExternalTask> {
409        self.command_queue.pending_external_tasks().await
410    }
411
412    /// Get dead letters from the queue's DLQ
413    pub async fn dead_letters(&self) -> Vec<a3s_lane::DeadLetter> {
414        self.command_queue.dead_letters().await
415    }
416
417    /// Get queue metrics snapshot
418    pub async fn queue_metrics(&self) -> Option<a3s_lane::MetricsSnapshot> {
419        self.command_queue.metrics_snapshot().await
420    }
421
422    /// Get queue statistics
423    pub async fn queue_stats(&self) -> crate::queue::SessionQueueStats {
424        self.command_queue.stats().await
425    }
426
427    /// Start the command queue scheduler
428    pub async fn start_queue(&self) -> Result<()> {
429        self.command_queue.start().await
430    }
431
432    /// Stop the command queue scheduler
433    pub async fn stop_queue(&self) {
434        self.command_queue.stop().await;
435    }
436
437    /// Get the system prompt from config
438    pub fn system(&self) -> Option<&str> {
439        self.config.system_prompt.as_deref()
440    }
441
442    /// Get conversation history
443    pub fn history(&self) -> &[Message] {
444        &self.messages
445    }
446
447    /// Add a message to history
448    pub fn add_message(&mut self, message: Message) {
449        self.messages.push(message);
450        self.context_usage.turns = self.messages.len();
451        self.touch();
452    }
453
454    /// Update context usage after a response
455    pub fn update_usage(&mut self, usage: &TokenUsage) {
456        self.total_usage.prompt_tokens += usage.prompt_tokens;
457        self.total_usage.completion_tokens += usage.completion_tokens;
458        self.total_usage.total_tokens += usage.total_tokens;
459
460        // Calculate cost if model pricing is available
461        let cost_usd = if let Some(ref model) = self.model_name {
462            let pricing_map = crate::telemetry::default_model_pricing();
463            if let Some(pricing) = pricing_map.get(model) {
464                let cost = pricing.calculate_cost(usage.prompt_tokens, usage.completion_tokens);
465                self.total_cost += cost;
466                Some(cost)
467            } else {
468                None
469            }
470        } else {
471            None
472        };
473
474        // Record per-call cost for aggregation
475        let model_str = self.model_name.clone().unwrap_or_default();
476        self.cost_records.push(crate::telemetry::LlmCostRecord {
477            model: model_str.clone(),
478            provider: String::new(),
479            prompt_tokens: usage.prompt_tokens,
480            completion_tokens: usage.completion_tokens,
481            total_tokens: usage.total_tokens,
482            cost_usd,
483            timestamp: chrono::Utc::now(),
484            session_id: Some(self.id.clone()),
485        });
486
487        // Record OTLP metrics (counters only; duration recorded in agent loop)
488        crate::telemetry::record_llm_metrics(
489            if model_str.is_empty() {
490                "unknown"
491            } else {
492                &model_str
493            },
494            usage.prompt_tokens,
495            usage.completion_tokens,
496            cost_usd.unwrap_or(0.0),
497            0.0, // Duration not available here; recorded via spans
498        );
499
500        // Estimate context usage (rough approximation)
501        self.context_usage.used_tokens = usage.prompt_tokens;
502        self.context_usage.percent =
503            self.context_usage.used_tokens as f32 / self.context_usage.max_tokens as f32;
504        self.touch();
505    }
506
507    /// Clear conversation history
508    pub fn clear(&mut self) {
509        self.messages.clear();
510        self.context_usage = ContextUsage::default();
511        self.touch();
512    }
513
514    /// Compact context by summarizing old messages
515    pub async fn compact(&mut self, llm_client: &Arc<dyn LlmClient>) -> Result<()> {
516        if let Some(new_messages) =
517            compaction::compact_messages(&self.id, &self.messages, llm_client).await?
518        {
519            self.messages = new_messages;
520            self.touch();
521        }
522        Ok(())
523    }
524
525    /// Pause the session
526    pub fn pause(&mut self) -> bool {
527        if self.state == SessionState::Active {
528            self.state = SessionState::Paused;
529            self.touch();
530            true
531        } else {
532            false
533        }
534    }
535
536    /// Resume the session
537    pub fn resume(&mut self) -> bool {
538        if self.state == SessionState::Paused {
539            self.state = SessionState::Active;
540            self.touch();
541            true
542        } else {
543            false
544        }
545    }
546
547    /// Set session state to error
548    pub fn set_error(&mut self) {
549        self.state = SessionState::Error;
550        self.touch();
551    }
552
553    /// Set session state to completed
554    pub fn set_completed(&mut self) {
555        self.state = SessionState::Completed;
556        self.touch();
557    }
558
559    /// Update the updated_at timestamp
560    fn touch(&mut self) {
561        self.updated_at = std::time::SystemTime::now()
562            .duration_since(std::time::UNIX_EPOCH)
563            .map(|d| d.as_secs() as i64)
564            .unwrap_or(0);
565    }
566
567    /// Convert to serializable SessionData for persistence
568    pub fn to_session_data(&self, llm_config: Option<LlmConfigData>) -> SessionData {
569        SessionData {
570            id: self.id.clone(),
571            config: self.config.clone(),
572            state: self.state,
573            messages: self.messages.clone(),
574            context_usage: self.context_usage.clone(),
575            total_usage: self.total_usage.clone(),
576            total_cost: self.total_cost,
577            model_name: self.model_name.clone(),
578            cost_records: self.cost_records.clone(),
579            tool_names: SessionData::tool_names_from_definitions(&self.tools),
580            thinking_enabled: self.thinking_enabled,
581            thinking_budget: self.thinking_budget,
582            created_at: self.created_at,
583            updated_at: self.updated_at,
584            llm_config,
585            tasks: self.tasks.clone(),
586            parent_id: self.parent_id.clone(),
587        }
588    }
589
590    /// Restore session state from SessionData
591    ///
592    /// Note: This only restores serializable fields. Non-serializable fields
593    /// (event_tx, command_queue, confirmation_manager) are already initialized
594    /// in Session::new().
595    pub fn restore_from_data(&mut self, data: &SessionData) {
596        self.state = data.state;
597        self.messages = data.messages.clone();
598        self.context_usage = data.context_usage.clone();
599        self.total_usage = data.total_usage.clone();
600        self.total_cost = data.total_cost;
601        self.model_name = data.model_name.clone();
602        self.cost_records = data.cost_records.clone();
603        self.thinking_enabled = data.thinking_enabled;
604        self.thinking_budget = data.thinking_budget;
605        self.created_at = data.created_at;
606        self.updated_at = data.updated_at;
607        self.tasks = data.tasks.clone();
608        self.parent_id = data.parent_id.clone();
609    }
610}
a3s_code_core/session/mod.rs

a3s_code_core/session/
mod.rs
