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 no host for it to appear to honour.
//!
//! The quoting, the forced ssh options and the remote script are
//! `crate::ssh`'s: this module decides *which* transport, that one knows what a
//! correct ssh invocation is.
//!
//! A transport is therefore two operations rather than one: **run this** and
//! **reach this port**. The second is what a managed `llama-server` needs -- it
//! is spawned as a process and then spoken to over HTTP -- and it is a no-op
//! locally, where the port a program binds is already one this machine can
//! dial. Over ssh it is an `-L` tunnel on the same connection that runs the
//! command, so the model server binds loopback on the far machine and is never
//! exposed to its network. See [`Transport::reserve_port`].

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

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

use crate::config::SshConfig;
pub use crate::ssh::Forward;

/// What a driver needs run in order to exist as a process. Deliberately just
/// what every transport can carry -- the command, where it runs, and a port the
/// caller needs to reach; anything a particular machine needs 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>,
    /// A port this program will listen on, and the port that reaches it
    /// from here -- see [`Transport::reserve_port`], which is the only
    /// thing that should produce one.
    ///
    /// On the launch rather than in [`Transport::spawn`]'s signature
    /// because it is part of what is being run: a caller that needs to
    /// reach the process it is starting says so once, where it says
    /// everything else about it, and every transport reads it the same
    /// way.
    pub forward: Option<Forward>,
}

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),
            forward: None,
        }
    }

    /// Says that this program serves `forward.there`, and that the caller
    /// will reach it at `forward.here`.
    pub fn reaching(mut self, forward: Forward) -> Self {
        self.forward = Some(forward);
        self
    }
}

/// 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 answers
/// within one call, so pipes this server drains are 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.
pub enum Streams {
    /// Pipes owned by this server; the child is killed when they drop.
    Piped,
    /// The same, except that stdin is already open on something this server
    /// holds -- the file being copied to another machine. Bytes this process has
    /// in memory do not need this: [`Streams::Piped`] gives a pipe to write them
    /// into as the child reads.
    PipedFrom(Stdio),
    /// 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.
    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 not on the remote
    /// PATH -- so each says its own thing.
    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(),
            launch.forward,
        ));
        match streams {
            Streams::Piped => {
                command
                    .stdin(Stdio::piped())
                    .stdout(Stdio::piped())
                    .stderr(Stdio::piped())
                    .kill_on_drop(true);
            }
            Streams::PipedFrom(stdin) => {
                command
                    .stdin(stdin)
                    .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
                // does not travel to a session 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 already on a blocking task that
    /// would otherwise need a runtime to ask a machine a question. Both build the
    /// invocation the same way.
    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(),
            launch.forward,
        )
        .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())
    }

    /// Runs `launch` with `input` on its stdin and reports everything it
    /// produced -- stdout as bytes, stderr as text, and the exit status.
    ///
    /// The one description of "run this there, with this on stdin", so that
    /// shipping an attachment and writing a file through the explorer are the
    /// same operation rather than two. It is also the only capture that hands
    /// back the **status**: a script can answer with an exit code the caller
    /// distinguishes (the explorer's write says `exit 3` for "this file is not
    /// the one you read"), which [`Transport::capture`] cannot express.
    ///
    /// Bytes rather than a `String`, because a file's contents are not text
    /// until something has checked, and lossy decoding would replace the
    /// evidence that they are not.
    ///
    /// `Err` means the process could not be started at all; a process that ran
    /// and failed is a [`Captured`] with a status saying so.
    pub async fn capture_with_input(&self, launch: &Launch, input: Input) -> Result<Captured> {
        let (streams, to_write) = match input {
            Input::None => (Streams::Piped, None),
            Input::Bytes(bytes) => (Streams::Piped, Some(bytes)),
            Input::File(file) => (Streams::PipedFrom(file.into()), None),
        };
        let mut child = self.spawn(launch, streams)?;
        if let Some(bytes) = to_write {
            // Written from a task rather than before the wait, because the child
            // may not read all of it -- the write script exits without reading
            // when the file has changed underneath -- and a caller blocked on
            // filling a pipe nobody is draining would deadlock instead of getting
            // that answer. The broken pipe is the expected end of this write.
            let mut stdin = child.stdin.take().context("the child has no stdin")?;
            tokio::spawn(async move {
                use tokio::io::AsyncWriteExt;
                let _ = stdin.write_all(&bytes).await;
                let _ = stdin.shutdown().await;
            });
        }
        let output = child
            .wait_with_output()
            .await
            .with_context(|| format!("waiting for \"{}\" {}", launch.program, self.describe()))?;
        Ok(Captured {
            status: output.status,
            stdout: output.stdout,
            stderr: String::from_utf8_lossy(&output.stderr).trim().to_string(),
        })
    }

    /// Picks a port for a launched program to serve on, and the port that
    /// reaches it from here.
    ///
    /// The "reach this port" half of what a transport is. Locally there is
    /// one port and the OS chooses it, by binding and letting go -- racy
    /// in principle, and nothing on this machine is hunting for ports.
    ///
    /// Over ssh the near end is chosen the same way and the far end is a
    /// guess, because there is no portable way to ask a machine for a free
    /// port that does not race with binding it anyway. It is taken from
    /// [`FAR_PORTS`], below the range Linux hands out to outgoing
    /// connections, so a collision means something else deliberately
    /// listening there. That is not silent: the program fails to bind and
    /// exits, and `session::llama` reports what its log said rather than
    /// waiting out its readiness timeout.
    pub fn reserve_port(&self) -> Result<Forward> {
        let listener = std::net::TcpListener::bind("127.0.0.1:0")
            .context("asking this machine for a free port")?;
        let here = listener.local_addr()?.port();
        Ok(match self {
            Self::Here => Forward { there: here, here },
            Self::Ssh { .. } => Forward {
                there: rand::random_range(FAR_PORTS),
                here,
            },
        })
    }

    /// 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),
        }
    }
}

/// Where a port on another machine is guessed from: high enough to be out
/// of the way of services, and below the 32768-60999 Linux hands out to
/// outgoing connections, which is where a guess would most often collide.
const FAR_PORTS: std::ops::Range<u16> = 20000..30000;

/// What a command is given on its standard input.
///
/// Three cases rather than an `Option<Stdio>` because they are three genuinely
/// different arrangements and only this knows which: nothing to say, bytes this
/// process is holding, or a file it has open. The last is how a
/// several-hundred-megabyte attachment reaches another machine without passing
/// through this server's memory.
pub enum Input {
    None,
    Bytes(Vec<u8>),
    File(std::fs::File),
}

/// Everything a finished command produced, including the status.
pub struct Captured {
    pub status: std::process::ExitStatus,
    pub stdout: Vec<u8>,
    /// Trimmed, and what a failure is reported as: ssh's own refusals and a
    /// tool's own message about the file it could not open are both the useful
    /// half of why something did not work.
    pub stderr: String,
}

impl Captured {
    /// The stdout of a command that succeeded, or the machine's own words.
    pub fn ok(self) -> Result<Vec<u8>> {
        if self.status.success() {
            return Ok(self.stdout);
        }
        anyhow::bail!(if self.stderr.is_empty() {
            format!("it failed with no explanation ({})", self.status)
        } else {
            self.stderr
        })
    }
}