Skip to main content

vtcode_core/pods/
transport.rs

1use anyhow::{Context, Result, anyhow};
2use async_trait::async_trait;
3use std::process::Stdio;
4use tokio::io::AsyncWriteExt;
5use tokio::process::Command;
6
7/// Captured command output from a pod transport.
8#[derive(Debug, Clone, Default)]
9pub struct CommandOutput {
10    pub success: bool,
11    pub stdout: String,
12    pub stderr: String,
13}
14
15#[async_trait]
16pub trait PodTransport: Send + Sync {
17    async fn exec_capture(&self, ssh_target: &str, command: &str) -> Result<CommandOutput>;
18    async fn write_file(&self, ssh_target: &str, remote_path: &str, contents: &str) -> Result<()>;
19    async fn exec_stream(&self, ssh_target: &str, command: &str) -> Result<()>;
20}
21
22/// SSH-backed transport used by the real CLI.
23#[derive(Debug, Clone, Default)]
24pub struct SshTransport;
25
26#[async_trait]
27impl PodTransport for SshTransport {
28    async fn exec_capture(&self, ssh_target: &str, command: &str) -> Result<CommandOutput> {
29        let mut ssh = build_ssh_command(ssh_target, command)?;
30        ssh.stdout(Stdio::piped()).stderr(Stdio::piped());
31
32        let output = ssh
33            .output()
34            .await
35            .with_context(|| format!("failed to execute SSH command: {command}"))?;
36
37        Ok(CommandOutput {
38            success: output.status.success(),
39            stdout: String::from_utf8_lossy(&output.stdout).into_owned(),
40            stderr: String::from_utf8_lossy(&output.stderr).into_owned(),
41        })
42    }
43
44    async fn write_file(&self, ssh_target: &str, remote_path: &str, contents: &str) -> Result<()> {
45        let remote_command = format!("cat > {remote_path}");
46        let mut ssh = build_ssh_command(ssh_target, &remote_command)?;
47        ssh.stdin(Stdio::piped()).stdout(Stdio::null()).stderr(Stdio::piped());
48
49        let mut child = ssh
50            .spawn()
51            .with_context(|| format!("failed to spawn SSH writer for {remote_path}"))?;
52
53        if let Some(mut stdin) = child.stdin.take() {
54            stdin
55                .write_all(contents.as_bytes())
56                .await
57                .with_context(|| format!("failed to write remote file {remote_path}"))?;
58        } else {
59            return Err(anyhow!("SSH writer did not provide stdin"));
60        }
61
62        let output = child
63            .wait_with_output()
64            .await
65            .with_context(|| format!("failed to finish SSH writer for {remote_path}"))?;
66
67        if output.status.success() {
68            Ok(())
69        } else {
70            Err(anyhow!("remote file write failed for {remote_path}: {}", String::from_utf8_lossy(&output.stderr)))
71        }
72    }
73
74    async fn exec_stream(&self, ssh_target: &str, command: &str) -> Result<()> {
75        let mut ssh = build_ssh_command(ssh_target, command)?;
76        ssh.stdin(Stdio::null()).stdout(Stdio::inherit()).stderr(Stdio::inherit());
77
78        let status = ssh
79            .status()
80            .await
81            .with_context(|| format!("failed to stream SSH command: {command}"))?;
82
83        if status.success() {
84            Ok(())
85        } else {
86            Err(anyhow!("SSH stream command failed with status {status}"))
87        }
88    }
89}
90
91fn build_ssh_command(ssh_target: &str, remote_command: &str) -> Result<Command> {
92    let parts = shell_words::split(ssh_target).with_context(|| format!("failed to parse SSH target: {ssh_target}"))?;
93
94    let Some((program, args)) = parts.split_first() else {
95        return Err(anyhow!("SSH target is empty"));
96    };
97
98    let mut command = Command::new(program);
99    command.args(args);
100    command.arg(remote_command);
101    Ok(command)
102}
103
104#[cfg(test)]
105mod tests {}