1use std::collections::HashMap;
2use std::fs;
3use std::io::{Read, Write};
4use std::path::{Component, Path, PathBuf};
5use std::sync::{Arc, Mutex};
6use std::thread::{self, JoinHandle};
7use std::time::{Duration, Instant};
8
9use anyhow::{Context, Result, anyhow};
10use portable_pty::{Child, CommandBuilder, MasterPty, PtySize, native_pty_system};
11use rexpect::{
12 process::wait::WaitStatus,
13 session::{Options, spawn_with_options},
14};
15use tracing::{debug, warn};
16use tui_term::vt100::Parser;
17
18use crate::config::PtyConfig;
19use crate::tools::types::VTCodePtySession;
20
21#[derive(Clone)]
22pub struct PtyManager {
23 workspace_root: PathBuf,
24 config: PtyConfig,
25 inner: Arc<PtyState>,
26}
27
28#[derive(Default)]
29struct PtyState {
30 sessions: Mutex<HashMap<String, PtySessionHandle>>,
31}
32
33struct PtySessionHandle {
34 master: Box<dyn MasterPty + Send>,
35 child: Mutex<Box<dyn Child + Send>>,
36 writer: Mutex<Option<Box<dyn Write + Send>>>,
37 parser: Arc<Mutex<Parser>>,
38 reader_thread: Mutex<Option<JoinHandle<()>>>,
39 metadata: VTCodePtySession,
40}
41
42impl PtySessionHandle {
43 fn snapshot_metadata(&self) -> VTCodePtySession {
44 let mut metadata = self.metadata.clone();
45 if let Ok(size) = self.master.get_size() {
46 metadata.rows = size.rows;
47 metadata.cols = size.cols;
48 }
49 if let Ok(parser) = self.parser.lock() {
50 metadata.screen_contents = Some(parser.screen().contents());
51 }
52 metadata
53 }
54}
55
56pub struct PtyCommandRequest {
57 pub command: Vec<String>,
58 pub working_dir: PathBuf,
59 pub timeout: Duration,
60 pub size: PtySize,
61}
62
63pub struct PtyCommandResult {
64 pub exit_code: i32,
65 pub output: String,
66 pub duration: Duration,
67 pub size: PtySize,
68}
69
70impl PtyManager {
71 pub fn new(workspace_root: PathBuf, config: PtyConfig) -> Self {
72 let resolved_root = workspace_root
73 .canonicalize()
74 .unwrap_or(workspace_root.clone());
75
76 Self {
77 workspace_root: resolved_root,
78 config,
79 inner: Arc::new(PtyState::default()),
80 }
81 }
82
83 pub fn config(&self) -> &PtyConfig {
84 &self.config
85 }
86
87 pub fn describe_working_dir(&self, path: &Path) -> String {
88 self.format_working_dir(path)
89 }
90
91 pub async fn run_command(&self, request: PtyCommandRequest) -> Result<PtyCommandResult> {
92 if request.command.is_empty() {
93 return Err(anyhow!("PTY command cannot be empty"));
94 }
95
96 let mut command = request.command.clone();
97 let program = command.remove(0);
98 let args = command;
99 let timeout = clamp_timeout(request.timeout);
100 let work_dir = request.working_dir.clone();
101 let size = request.size;
102 let start = Instant::now();
103
104 let result = tokio::task::spawn_blocking(move || -> Result<PtyCommandResult> {
105 let mut cmd = std::process::Command::new(&program);
106 cmd.args(&args);
107 cmd.current_dir(&work_dir);
108 cmd.env("TERM", "xterm-256color");
109 cmd.env("COLUMNS", size.cols.to_string());
110 cmd.env("LINES", size.rows.to_string());
111
112 let options = Options {
113 timeout_ms: Some(timeout),
114 strip_ansi_escape_codes: false,
115 };
116
117 let mut session = spawn_with_options(cmd, options)
118 .with_context(|| format!("failed to spawn PTY command '{program}'"))?;
119
120 let mut output = String::new();
121 let collected = session
122 .exp_eof()
123 .context("failed to read PTY command output")?;
124 output.push_str(&collected);
125
126 let status = session
127 .process
128 .wait()
129 .context("failed to wait for PTY command to exit")?;
130 let exit_code = wait_status_code(status);
131
132 Ok(PtyCommandResult {
133 exit_code,
134 output,
135 duration: start.elapsed(),
136 size,
137 })
138 })
139 .await
140 .context("failed to join PTY command task")??;
141
142 Ok(result)
143 }
144
145 pub fn resolve_working_dir(&self, requested: Option<&str>) -> Result<PathBuf> {
146 let requested = match requested {
147 Some(dir) if !dir.trim().is_empty() => dir,
148 _ => return Ok(self.workspace_root.clone()),
149 };
150
151 let candidate = self.workspace_root.join(requested);
152 let normalized = normalize_path(&candidate);
153 if !normalized.starts_with(&self.workspace_root) {
154 return Err(anyhow!(
155 "Working directory '{}' escapes the workspace root",
156 candidate.display()
157 ));
158 }
159 let metadata = fs::metadata(&normalized).with_context(|| {
160 format!(
161 "Working directory '{}' does not exist",
162 normalized.display()
163 )
164 })?;
165 if !metadata.is_dir() {
166 return Err(anyhow!(
167 "Working directory '{}' is not a directory",
168 normalized.display()
169 ));
170 }
171 Ok(normalized)
172 }
173
174 pub fn create_session(
175 &self,
176 session_id: String,
177 command: Vec<String>,
178 working_dir: PathBuf,
179 size: PtySize,
180 ) -> Result<VTCodePtySession> {
181 if command.is_empty() {
182 return Err(anyhow!("PTY session command cannot be empty"));
183 }
184
185 let mut sessions = self
186 .inner
187 .sessions
188 .lock()
189 .expect("PTY session mutex poisoned");
190 if sessions.contains_key(&session_id) {
191 return Err(anyhow!("PTY session '{}' already exists", session_id));
192 }
193
194 let mut command_parts = command.clone();
195 let program = command_parts.remove(0);
196 let args = command_parts;
197
198 let pty_system = native_pty_system();
199 let pair = pty_system
200 .openpty(size)
201 .context("failed to allocate PTY pair")?;
202
203 let mut builder = CommandBuilder::new(program.clone());
204 for arg in &args {
205 builder.arg(arg);
206 }
207 builder.cwd(&working_dir);
208 builder.env("TERM", "xterm-256color");
209 builder.env("COLUMNS", size.cols.to_string());
210 builder.env("LINES", size.rows.to_string());
211
212 let child = pair
213 .slave
214 .spawn_command(builder)
215 .context("failed to spawn PTY session command")?;
216 drop(pair.slave);
217
218 let master = pair.master;
219 let mut reader = master
220 .try_clone_reader()
221 .context("failed to clone PTY reader")?;
222 let writer = master.take_writer().context("failed to take PTY writer")?;
223
224 let parser = Arc::new(Mutex::new(Parser::new(size.rows, size.cols, 0)));
225 let parser_clone = Arc::clone(&parser);
226 let session_name = session_id.clone();
227 let reader_thread = thread::Builder::new()
228 .name(format!("vtcode-pty-reader-{session_name}"))
229 .spawn(move || {
230 let mut buffer = [0u8; 4096];
231 loop {
232 match reader.read(&mut buffer) {
233 Ok(0) => {
234 debug!("PTY session '{}' reader reached EOF", session_name);
235 break;
236 }
237 Ok(bytes_read) => {
238 if let Ok(mut parser) = parser_clone.lock() {
239 parser.process(&buffer[..bytes_read]);
240 }
241 }
242 Err(error) => {
243 warn!("PTY session '{}' reader error: {}", session_name, error);
244 break;
245 }
246 }
247 }
248 })
249 .context("failed to spawn PTY reader thread")?;
250
251 let metadata = VTCodePtySession {
252 id: session_id.clone(),
253 command: program,
254 args,
255 working_dir: Some(self.format_working_dir(&working_dir)),
256 rows: size.rows,
257 cols: size.cols,
258 screen_contents: None,
259 };
260
261 sessions.insert(
262 session_id.clone(),
263 PtySessionHandle {
264 master,
265 child: Mutex::new(child),
266 writer: Mutex::new(Some(writer)),
267 parser,
268 reader_thread: Mutex::new(Some(reader_thread)),
269 metadata: metadata.clone(),
270 },
271 );
272
273 Ok(metadata)
274 }
275
276 pub fn list_sessions(&self) -> Vec<VTCodePtySession> {
277 let sessions = self
278 .inner
279 .sessions
280 .lock()
281 .expect("PTY session mutex poisoned");
282 sessions
283 .values()
284 .map(PtySessionHandle::snapshot_metadata)
285 .collect()
286 }
287
288 pub fn close_session(&self, session_id: &str) -> Result<VTCodePtySession> {
289 let handle = {
290 let mut sessions = self
291 .inner
292 .sessions
293 .lock()
294 .expect("PTY session mutex poisoned");
295 sessions
296 .remove(session_id)
297 .ok_or_else(|| anyhow!("PTY session '{}' not found", session_id))?
298 };
299
300 if let Ok(mut writer_guard) = handle.writer.lock() {
301 if let Some(mut writer) = writer_guard.take() {
302 let _ = writer.write_all(b"exit\n");
303 let _ = writer.flush();
304 }
305 }
306
307 let mut child = handle.child.lock().expect("PTY child mutex poisoned");
308 if child
309 .try_wait()
310 .context("failed to poll PTY session status")?
311 .is_none()
312 {
313 child.kill().context("failed to terminate PTY session")?;
314 let _ = child.wait();
315 }
316
317 if let Ok(mut thread_guard) = handle.reader_thread.lock() {
318 if let Some(reader_thread) = thread_guard.take() {
319 if let Err(panic) = reader_thread.join() {
320 warn!(
321 "PTY session '{}' reader thread panicked: {:?}",
322 session_id, panic
323 );
324 }
325 }
326 }
327
328 Ok(handle.snapshot_metadata())
329 }
330
331 fn format_working_dir(&self, path: &Path) -> String {
332 match path.strip_prefix(&self.workspace_root) {
333 Ok(relative) if relative.as_os_str().is_empty() => ".".to_string(),
334 Ok(relative) => relative.to_string_lossy().replace("\\", "/"),
335 Err(_) => path.to_string_lossy().to_string(),
336 }
337 }
338}
339
340fn clamp_timeout(duration: Duration) -> u64 {
341 duration.as_millis().min(u64::MAX as u128) as u64
342}
343
344fn wait_status_code(status: WaitStatus) -> i32 {
345 match status {
346 WaitStatus::Exited(_, code) => code,
347 WaitStatus::Signaled(_, signal, _) | WaitStatus::Stopped(_, signal) => -(signal as i32),
348 WaitStatus::StillAlive | WaitStatus::Continued(_) => 0,
349 #[cfg(any(target_os = "linux", target_os = "android"))]
350 WaitStatus::PtraceEvent(_, _, _) | WaitStatus::PtraceSyscall(_) => 0,
351 }
352}
353
354fn normalize_path(path: &Path) -> PathBuf {
355 let mut normalized = PathBuf::new();
356 for component in path.components() {
357 match component {
358 Component::ParentDir => {
359 normalized.pop();
360 }
361 Component::CurDir => {}
362 Component::Prefix(prefix) => normalized.push(prefix.as_os_str()),
363 Component::RootDir => normalized.push(component.as_os_str()),
364 Component::Normal(part) => normalized.push(part),
365 }
366 }
367 normalized
368}