ai-app/server/src/session/transport.rs

//! Where a session's process runs, and the only place that knows how.
//!
//! A driver says *what* to run -- a [`Launch`] -- and hands it here.
//! Whether that becomes a child of this process or an `ssh host …`
//! invocation is settled in this module, so a driver carries no transport
//! knowledge and a second one cannot forget to handle the remote case. It
//! also means the wrapping is honest about drivers that run nothing at
//! all: `EchoDriver` builds no [`Launch`], so there is nothing to wrap and
//! no host for it to appear to honour.
//!
//! The quoting, the forced ssh options and the remote script are
//! `crate::ssh`'s, which this dispatches to. That split is deliberate:
//! this module decides *which* transport, that one knows what a correct
//! ssh invocation is.
//!
//! Known second operation, not built because nothing needs it yet: a
//! managed `llama-server` is spawned as a process but then spoken to over
//! HTTP, so a remote one needs a forwarded port (`ssh -L`) as well. A
//! transport is eventually "run this" plus "reach this port", where the
//! second is a no-op locally. See PLAN.md's SSH section.

use std::path::{Path, PathBuf};
use std::process::Stdio;

use anyhow::{Context, Result};
use tokio::process::Child;

use crate::config::SshConfig;

/// What a driver needs run in order to exist as a process.
///
/// Deliberately just the three things every transport can carry. Anything
/// a particular machine needs -- a port, a key, extra ssh options -- is
/// the transport's own configuration, not something a driver states.
pub struct Launch {
    pub program: String,
    pub args: Vec<String>,
    pub cwd: Option<PathBuf>,
}

impl Launch {
    pub fn new(program: impl Into<String>, args: Vec<String>, cwd: Option<&Path>) -> Self {
        Self {
            program: program.into(),
            args,
            cwd: cwd.map(Path::to_path_buf),
        }
    }
}

/// How a launched process's standard streams are connected.
///
/// The choice is not the transport's and not the driver's dialect: it is
/// whether the process is expected to outlive this server. A probe is
/// asked a question and answers within one call, so pipes this server
/// drains are right and dying with it is right. A session is a
/// conversation somebody is having, so its streams live in the session
/// directory where a later run of this server can pick them up again --
/// see `session::process`.
pub enum Streams {
    /// Pipes owned by this server; the child is killed when they drop.
    Piped,
    /// Files -- and, for stdin, a fifo the child itself holds open so it
    /// never reads EOF -- that outlast this process.
    Detached {
        stdin: Stdio,
        stdout: Stdio,
        stderr: Stdio,
    },
}

/// The machine a session's process runs on.
pub enum Transport {
    /// The machine this server is running on.
    Here,
    /// Reached with the system `ssh` client. Owns its entry rather than
    /// borrowing it, so a session keeps working against the config it was
    /// spawned with even if the setup is edited afterwards. Carries the
    /// setup's name only to say where things are running.
    Ssh { name: String, ssh: SshConfig },
}

impl Transport {
    /// The transport a setup describes; a setup with no `ssh` is here.
    pub fn for_setup(setup: &crate::config::SetupConfig) -> Self {
        match &setup.ssh {
            Some(ssh) => Self::Ssh {
                name: setup.name.clone(),
                ssh: ssh.clone(),
            },
            None => Self::Here,
        }
    }

    /// Starts `launch` with its streams connected as `streams` says.
    ///
    /// The failure names what to check, and the two transports fail for
    /// genuinely different reasons -- a missing ssh client here versus a
    /// program that is not on the remote PATH -- so each says its own
    /// thing rather than one message hedging between them.
    pub fn spawn(&self, launch: &Launch, streams: Streams) -> Result<Child> {
        let host = match self {
            Self::Here => None,
            Self::Ssh { ssh, .. } => Some(ssh),
        };
        let mut command = tokio::process::Command::from(crate::ssh::command(
            host,
            &launch.program,
            &launch.args,
            launch.cwd.as_deref(),
        ));
        match streams {
            Streams::Piped => {
                command
                    .stdin(Stdio::piped())
                    .stdout(Stdio::piped())
                    .stderr(Stdio::piped())
                    .kill_on_drop(true);
            }
            Streams::Detached {
                stdin,
                stdout,
                stderr,
            } => {
                command.stdin(stdin).stdout(stdout).stderr(stderr);
                // No `kill_on_drop`: outliving this server is the point.
                // Its own process group as well, so a signal sent to the
                // server's group -- which is how a terminal or a
                // supervisor stops it -- does not travel to a session that
                // is meant to survive being stopped.
                command.process_group(0);
            }
        }
        command.spawn().with_context(|| match self {
            Self::Ssh { name, .. } => format!(
                "couldn't start ssh to run \"{}\" on {name} -- is the ssh client installed \
                     here?",
                launch.program,
            ),
            Self::Here => format!(
                "couldn't run \"{}\" on this machine -- is it installed and on PATH? If it \
                     lives on another machine, give the session a host to run on.",
                launch.program,
            ),
        })
    }

    /// Runs `launch` to completion and returns its stdout, blocking.
    ///
    /// The synchronous twin of `capture`, for callers that are already on a
    /// blocking task and would otherwise need a runtime to ask a machine a
    /// question. Both build the invocation the same way -- see
    /// `crate::ssh::command` -- so there is still only one description of
    /// what running something on another machine means.
    pub fn capture_blocking(&self, launch: &Launch) -> Result<String> {
        let host = match self {
            Self::Here => None,
            Self::Ssh { ssh, .. } => Some(ssh),
        };
        let output =
            crate::ssh::command(host, &launch.program, &launch.args, launch.cwd.as_deref())
                .output()
                .with_context(|| {
                    format!("couldn't run \"{}\" {}", launch.program, self.describe())
                })?;
        if !output.status.success() {
            let stderr = String::from_utf8_lossy(&output.stderr).trim().to_string();
            anyhow::bail!(if stderr.is_empty() {
                format!("couldn't reach it ({})", output.status)
            } else {
                stderr
            });
        }
        Ok(String::from_utf8_lossy(&output.stdout).into_owned())
    }

    /// How to say where this runs, for a log line a person reads.
    pub fn describe(&self) -> String {
        match self {
            Self::Here => "on this machine".to_string(),
            Self::Ssh { name, ssh } => format!("on {name} ({})", ssh.address),
        }
    }
}