vtcode_core/tools/registry/
pty.rs

1use std::path::PathBuf;
2use std::sync::Arc;
3use std::sync::atomic::{AtomicUsize, Ordering};
4
5use anyhow::{Result, anyhow};
6
7use crate::config::PtyConfig;
8
9use super::PtyManager;
10
11/// RAII guard to automatically decrement session count when dropped
12#[derive(Debug)]
13pub struct PtySessionGuard {
14    active_sessions: Arc<AtomicUsize>,
15}
16
17impl Drop for PtySessionGuard {
18    fn drop(&mut self) {
19        decrement_active_sessions(&self.active_sessions);
20    }
21}
22
23fn decrement_active_sessions(active_sessions: &AtomicUsize) {
24    let _ = active_sessions.fetch_update(Ordering::Relaxed, Ordering::Relaxed, |current| current.checked_sub(1));
25}
26
27#[derive(Clone)]
28pub struct PtySessionManager {
29    config: PtyConfig,
30    manager: PtyManager,
31    active_sessions: Arc<AtomicUsize>,
32}
33
34impl PtySessionManager {
35    pub fn new(workspace_root: PathBuf, config: PtyConfig) -> Self {
36        let manager = PtyManager::new(workspace_root, config.clone());
37
38        Self {
39            config,
40            manager,
41            active_sessions: Arc::new(AtomicUsize::new(0)),
42        }
43    }
44
45    pub fn config(&self) -> &PtyConfig {
46        &self.config
47    }
48
49    pub fn manager(&self) -> &PtyManager {
50        &self.manager
51    }
52
53    pub fn can_start_session(&self) -> bool {
54        if !self.config.enabled {
55            return false;
56        }
57
58        self.active_sessions.load(Ordering::Relaxed) < self.config.max_sessions
59    }
60
61    /// Start a PTY session and return an RAII guard that will automatically decrement
62    /// the session count when dropped, even if an error occurs during execution.
63    pub fn start_session(&self) -> Result<PtySessionGuard> {
64        if !self.config.enabled {
65            return Err(anyhow!(
66                "Maximum PTY sessions ({}) exceeded. Current active sessions: {}",
67                self.config.max_sessions,
68                self.active_sessions.load(Ordering::Relaxed)
69            ));
70        }
71
72        loop {
73            let current = self.active_sessions.load(Ordering::Relaxed);
74            if current >= self.config.max_sessions {
75                return Err(anyhow!(
76                    "Maximum PTY sessions ({}) exceeded. Current active sessions: {}",
77                    self.config.max_sessions,
78                    current
79                ));
80            }
81
82            if self
83                .active_sessions
84                .compare_exchange_weak(current, current + 1, Ordering::Relaxed, Ordering::Relaxed)
85                .is_ok()
86            {
87                return Ok(PtySessionGuard { active_sessions: Arc::clone(&self.active_sessions) });
88            }
89        }
90    }
91
92    pub fn end_session(&self) {
93        decrement_active_sessions(&self.active_sessions);
94    }
95
96    pub fn active_sessions(&self) -> usize {
97        self.active_sessions.load(Ordering::Relaxed)
98    }
99
100    pub fn terminate_all(&self) {
101        self.manager.terminate_all_sessions();
102        self.active_sessions.store(0, Ordering::Relaxed);
103    }
104
105    pub async fn terminate_all_async(&self) -> Result<()> {
106        let session_manager = self.clone();
107        tokio::task::spawn_blocking(move || session_manager.terminate_all())
108            .await
109            .map_err(|join_err| anyhow!("terminate_all_pty_sessions task failed to join: {join_err}"))?;
110        Ok(())
111    }
112}
vtcode_core/tools/registry/pty.rs

vtcode_core/tools/registry/
pty.rs
