vtcode_core/tools/registry/
pty_facade.rs

1//! PTY-related ToolRegistry accessors.
2
3use std::collections::HashSet;
4use std::future::Future;
5use std::sync::Arc;
6
7use anyhow::Result;
8
9use crate::config::PtyConfig;
10use crate::tools::exec_session::ExecSessionManager;
11use crate::tools::pty::PtyManager;
12
13use super::ToolRegistry;
14use super::pty;
15
16impl ToolRegistry {
17    pub fn pty_manager(&self) -> &PtyManager {
18        self.pty_sessions.manager()
19    }
20
21    pub fn pty_config(&self) -> &PtyConfig {
22        self.pty_sessions.config()
23    }
24
25    pub fn can_start_pty_session(&self) -> bool {
26        self.pty_sessions.can_start_session()
27    }
28
29    pub fn start_pty_session(&self) -> Result<pty::PtySessionGuard> {
30        self.pty_sessions.start_session()
31    }
32
33    pub fn end_pty_session(&self) {
34        self.pty_sessions.end_session();
35    }
36
37    pub fn active_pty_sessions(&self) -> usize {
38        self.pty_sessions.active_sessions()
39    }
40
41    pub fn terminate_all_pty_sessions(&self) {
42        self.pty_sessions.terminate_all();
43    }
44
45    /// Inline-delegating wrapper over [`pty::PtySessionManager::terminate_all_async`].
46    /// Returns the inner future directly to avoid an extra coroutine state machine
47    /// (per the async-state-machine bloat audit, section 16).
48    pub fn terminate_all_pty_sessions_async(&self) -> impl Future<Output = Result<()>> + '_ {
49        self.pty_sessions.terminate_all_async()
50    }
51
52    /// Terminates every exec session, then retires verification bookkeeping
53    /// for the sessions that are actually gone (see
54    /// [`Self::terminate_exec_sessions_and_observe_cleanup`]).
55    pub async fn terminate_all_exec_sessions_async(&self) -> Result<()> {
56        self.terminate_exec_sessions_and_observe_cleanup(false, false).await
57    }
58
59    /// Exit-path variant: terminates every exec session with `Immediate` PTY
60    /// semantics (group SIGKILL, no SIGTERM grace window) so user-initiated
61    /// exit is not parked ~600 ms per live PTY session.
62    pub async fn terminate_all_exec_sessions_for_exit_async(&self) -> Result<()> {
63        self.terminate_exec_sessions_and_observe_cleanup(false, true).await
64    }
65
66    /// Terminates only foreground exec sessions after a turn interruption,
67    /// preserving user-owned background sessions, then retires verification
68    /// bookkeeping for the sessions that are actually gone.
69    pub async fn terminate_active_exec_sessions_async(&self) -> Result<()> {
70        self.terminate_exec_sessions_and_observe_cleanup(true, false).await
71    }
72
73    /// Exit-path variant of [`Self::terminate_active_exec_sessions_async`]
74    /// using `Immediate` PTY termination.
75    pub async fn terminate_active_exec_sessions_for_exit_async(&self) -> Result<()> {
76        self.terminate_exec_sessions_and_observe_cleanup(true, true).await
77    }
78
79    /// Shared bulk-termination path: snapshots the in-scope session ids,
80    /// terminates, then retires verification bookkeeping for the sessions
81    /// that are actually gone. Only gone ids are observed so a live mutator
82    /// (e.g. when termination partially fails) keeps its pending boundary.
83    /// Verification debt (`generation`) is preserved; only the
84    /// terminated-session block is released so fresh checks can satisfy it.
85    async fn terminate_exec_sessions_and_observe_cleanup(&self, active_only: bool, immediate: bool) -> Result<()> {
86        let before = self.scoped_exec_session_ids(active_only).await;
87        let result = if immediate {
88            if active_only {
89                self.exec_sessions.terminate_active_sessions_for_exit_async().await
90            } else {
91                self.exec_sessions.terminate_all_sessions_for_exit_async().await
92            }
93        } else if active_only {
94            self.exec_sessions.terminate_active_sessions_async().await
95        } else {
96            self.exec_sessions.terminate_all_sessions_async().await
97        };
98        self.observe_terminated_exec_sessions(before).await;
99        result
100    }
101
102    pub fn exec_session_manager(&self) -> ExecSessionManager {
103        self.exec_sessions.clone()
104    }
105
106    /// Force-stop every foreground exec session and close it (TUI
107    /// ForceCancel escape hatch), then retire verification bookkeeping for
108    /// the sessions that are actually gone so a force-cancelled mutator
109    /// cannot strand `pending_mutation_sessions` forever. Returns
110    /// `(stopped, closed, failed)` counts from the session manager.
111    /// Background sessions are user-owned and are left alone.
112    pub async fn force_cancel_foreground_sessions_and_observe_cleanup(&self) -> (usize, usize, usize) {
113        let before = self.scoped_exec_session_ids(true).await;
114        let counts = self.exec_sessions.force_cancel_foreground_sessions().await;
115        self.observe_terminated_exec_sessions(before).await;
116        counts
117    }
118
119    /// Snapshot the ids of in-scope exec sessions (foreground only when
120    /// `active_only`) for post-cleanup reconciliation.
121    async fn scoped_exec_session_ids(&self, active_only: bool) -> Vec<String> {
122        self.exec_sessions
123            .list_sessions()
124            .await
125            .into_iter()
126            .filter(|session| !active_only || !session.background)
127            .map(|session| session.id.as_str().to_owned())
128            .collect()
129    }
130
131    /// Retire verification bookkeeping for snapshotted ids that are gone
132    /// now. Only gone ids are observed so a live mutator (e.g. when
133    /// termination partially fails) keeps its pending boundary. Verification
134    /// debt is preserved; only the terminated-session block is released.
135    async fn observe_terminated_exec_sessions(&self, before: Vec<String>) {
136        if before.is_empty() {
137            return;
138        }
139        let remaining: HashSet<String> = self
140            .exec_sessions
141            .list_sessions()
142            .await
143            .into_iter()
144            .map(|session| session.id.as_str().to_owned())
145            .collect();
146        let terminated: Vec<String> = before.into_iter().filter(|id| !remaining.contains(id)).collect();
147        if !terminated.is_empty() {
148            self.observe_completion_exec_sessions_terminated(&terminated);
149        }
150    }
151
152    /// Set the active PTY sessions counter for tracking
153    pub fn set_active_pty_sessions(&self, counter: Arc<std::sync::atomic::AtomicUsize>) {
154        self.exec_sessions.set_foreground_pty_counter(Arc::clone(&counter));
155        if let Ok(mut guard) = self.active_pty_sessions.write() {
156            *guard = Some(counter);
157        }
158    }
159
160    /// Increment active PTY sessions count
161    pub fn increment_active_pty_sessions(&self) {
162        if let Some(counter) = self
163            .active_pty_sessions
164            .read()
165            .unwrap_or_else(|e| e.into_inner())
166            .as_ref()
167            .map(Arc::clone)
168        {
169            counter.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
170        }
171    }
172
173    /// Decrement active PTY sessions count
174    pub fn decrement_active_pty_sessions(&self) {
175        if let Some(counter) = self
176            .active_pty_sessions
177            .read()
178            .unwrap_or_else(|e| e.into_inner())
179            .as_ref()
180            .map(Arc::clone)
181        {
182            counter.fetch_sub(1, std::sync::atomic::Ordering::Relaxed);
183        }
184    }
185
186    /// Get the current active PTY sessions count
187    pub fn active_pty_sessions_count(&self) -> usize {
188        self.active_pty_sessions
189            .read()
190            .ok()
191            .and_then(|g| g.as_ref().map(|c| c.load(std::sync::atomic::Ordering::Relaxed)))
192            .unwrap_or(0)
193    }
194}
vtcode_core/tools/registry/pty_facade.rs

vtcode_core/tools/registry/
pty_facade.rs
