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

vtcode_core/tools/registry/
pty.rs
