ai-app/server/src/machines.rs

//! Finding out what a machine can run, rather than being told.
//!
//! The phone adds a machine by giving connection details; this asks the machine
//! itself which of the known programs it has, and the answer becomes its
//! providers. That is a security property, not a convenience: **no route accepts
//! a command from the phone.** If it did, the enrolled token could introduce
//! arbitrary programs to run on every machine already configured here.
//!
//! It is also the better interface: nobody wants to type an absolute path on a
//! phone keyboard, and a machine that has moved its binaries answers correctly
//! on the next probe.
//!
//! The cost is that a program somewhere unusual is invisible. The escape hatch
//! is editing `config.ron` on the backend, which is exactly the authority the
//! phone is not being given.

use anyhow::{Context, Result};
use serde::Serialize;
use serde_json::{Value, json};

use crate::config::{DriverKind, ProviderConfig};
use crate::session::transport::{Launch, Transport};

/// What is looked for, and what finding it makes. Extending this is how a new
/// driver becomes discoverable -- one row, not a branch anywhere. The name is
/// what the provider gets called, so it is what the phone shows and what a
/// session stores.
const PROBES: &[(&str, &str, DriverKind)] = &[
    ("claude-cli", "claude", DriverKind::ClaudeCli),
    ("codex-cli", "codex", DriverKind::CodexCli),
    // Named for the program rather than for where it runs: it runs
    // wherever the machine is, and "local" was true only while a llama
    // session could not be spawned on another machine.
    ("llama-cpp", "llama-server", DriverKind::LlamaCpp),
];

/// Models offered for a discovered Claude CLI. A shortcut list for the spawn
/// screen, not a restriction -- the field stays free text.
const CLAUDE_MODELS: &[&str] = &["fable", "opus", "sonnet", "haiku"];

/// Asks `transport`'s machine which of [`PROBES`] it has.
///
/// One round trip rather than one per program: over ssh each would be a separate
/// connection and handshake. `command -v` is POSIX and a shell builtin, so it
/// works whatever is installed -- and `|| true` keeps a missing program from
/// ending the loop, since the caller wants the whole answer.
pub async fn discover(transport: &Transport) -> Result<Vec<ProviderConfig>> {
    let wanted: Vec<&str> = PROBES.iter().map(|(_, binary, _)| *binary).collect();
    let script = format!(
        "for p in {}; do command -v \"$p\" || true; done",
        wanted.join(" ")
    );
    let launch = Launch::new("sh", vec!["-c".to_string(), script], None);
    let found = transport.capture(&launch).await.map_err(explain)?;

    let mut providers = Vec::new();
    // Echo runs inside this server, so it exists exactly where this server does
    // and nowhere else. Offering it on a remote machine would be a choice that
    // changes nothing.
    if matches!(transport, Transport::Here) {
        providers.push(crate::config::Config::echo_provider());
    }
    for (name, binary, kind) in PROBES {
        let path = found
            .lines()
            .map(str::trim)
            .find(|line| line.rsplit('/').next() == Some(*binary));
        let Some(path) = path else {
            continue;
        };
        providers.push(ProviderConfig {
            name: (*name).to_string(),
            kind: *kind,
            // The resolved path rather than the bare name: PATH under a
            // non-interactive ssh session is not the one a person sees when they
            // log in, so "it is on my PATH" is not enough.
            command: Some(path.to_string()),
            models: match kind {
                DriverKind::ClaudeCli => CLAUDE_MODELS.iter().map(|m| (*m).to_string()).collect(),
                _ => Vec::new(),
            },
            mcp_servers: mcp_defaults(*kind),
            // What a probe cannot know: how this machine's models are loaded
            // is configured after the fact, and a re-probe keeps it -- see
            // `SessionManager::update_machine`.
            model_settings: Default::default(),
            max_loaded: None,
        });
    }
    Ok(providers)
}

/// One model a picker can offer, and what to call it there.
///
/// Two fields rather than one string because for one provider they differ:
/// a llama.cpp model is chosen by the path it lives at and read as the name
/// its own metadata gives it. Every other provider's id is already the name,
/// and says so by repeating it -- which is what keeps the picker free of a
/// branch on the session kind.
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct OfferedModel {
    /// What a spawn or a model change is given. Opaque to the phone.
    pub id: String,
    /// What a person reads on the chip.
    pub label: String,
}

impl OfferedModel {
    /// A model whose id is its own name, which is every provider but llama.
    fn plain(id: impl Into<String>) -> Self {
        let id = id.into();
        Self {
            label: id.clone(),
            id,
        }
    }
}

/// The MCP servers a newly discovered provider of this kind starts with.
///
/// A default rather than something to be typed in: a llama session with no web
/// search is the state somebody would then have to find out how to leave, and
/// Exa is what llama.cpp's own web UI offers under the same name. It is an
/// ordinary config entry once written, so removing it is deleting a line.
///
/// Only llama.cpp, because only a driver that runs its own agent loop can use
/// one -- the coding CLIs configure MCP themselves and a second answer here
/// would quietly disagree with theirs.
fn mcp_defaults(kind: DriverKind) -> Vec<crate::config::McpServerConfig> {
    match kind {
        DriverKind::LlamaCpp => vec![crate::config::McpServerConfig {
            name: "exa".to_string(),
            url: crate::session::llama::EXA_MCP_URL.to_string(),
        }],
        _ => Vec::new(),
    }
}

/// Models the selected provider currently offers on this machine.
///
/// Codex's catalog is account- and CLI-version-specific, so it is asked at the
/// moment the picker opens rather than copied into `config.ron`. A llama.cpp
/// provider offers the GGUFs on the machine it runs on, through this same call
/// -- there was a second route answering that alone, and it went when this one
/// learned to, because a picker offering a model the spawn screen does not, or
/// naming it differently, is two answers to one question. Other providers
/// retain the shortcut list discovery stored for them.
pub async fn provider_models(
    transport: &Transport,
    provider: &ProviderConfig,
    models_dir: &std::path::Path,
) -> Result<Vec<OfferedModel>> {
    if provider.kind == DriverKind::LlamaCpp {
        let dir = crate::models::dir_on(transport, models_dir);
        let mut found = crate::models::on_machine(transport, &dir).await?;
        // A vision model's projector is a file beside it rather than a model,
        // and the session that reads pictures is the one on the model: offered
        // here it is a chip that starts a server which cannot load it. It is
        // still in the machine's own model list, which is where a file on a
        // disk is managed and deleted.
        found.retain(|model| {
            !crate::session::llama::is_projector(model.key.rsplit('/').next().unwrap_or_default())
        });
        let labels = crate::models::labels(&found);
        return Ok(found
            .into_iter()
            .zip(labels)
            .map(|(model, label)| OfferedModel {
                id: model.key,
                label,
            })
            .collect());
    }
    if provider.kind != DriverKind::CodexCli {
        return Ok(provider.models.iter().map(OfferedModel::plain).collect());
    }
    let transport = transport.clone();
    let program = provider.program().to_string();
    tokio::task::spawn_blocking(move || {
        let launch = Launch::new(program, vec!["app-server".into(), "--stdio".into()], None);
        let initial = json!({
            "id": 1,
            "method": "initialize",
            "params": {"clientInfo": {"name": "ai-app", "title": "AI Sessions", "version": env!("CARGO_PKG_VERSION")}}
        });
        let requests = [
            json!({"method": "initialized"}),
            json!({"id": 2, "method": "model/list", "params": {"includeHidden": false, "limit": 100}}),
        ];
        let answer = transport.request_json_blocking(&launch, &initial, &requests, 2)?;
        Ok(parse_codex_models(&answer)?
            .into_iter()
            .map(OfferedModel::plain)
            .collect())
    })
    .await?
}

fn parse_codex_models(answer: &Value) -> Result<Vec<String>> {
    if let Some(message) = answer.pointer("/error/message").and_then(Value::as_str) {
        anyhow::bail!("Codex could not list models: {message}");
    }
    let entries = answer
        .pointer("/result/data")
        .and_then(Value::as_array)
        .context("Codex returned no model catalog")?;
    let mut models = entries
        .iter()
        .filter(|entry| {
            !entry
                .get("hidden")
                .and_then(Value::as_bool)
                .unwrap_or(false)
        })
        .filter_map(|entry| entry.get("model").and_then(Value::as_str))
        .map(str::to_string)
        .collect::<Vec<_>>();
    models.dedup();
    Ok(models)
}

/// Adds what to do to failures whose own wording does not say.
///
/// ssh's messages are written for someone at a terminal on the backend, which is
/// exactly who is not reading this one. Host key verification is the case that
/// matters: **every** machine fails it the first time, so without this, adding a
/// machine from the phone looks broken rather than unfinished.
///
/// Deliberately not fixed by relaxing the check. `StrictHostKeyChecking` stays
/// at its default, so a first connection is a decision somebody makes on the
/// backend with the key in front of them.
fn explain(err: anyhow::Error) -> anyhow::Error {
    let message = format!("{err:#}");
    if message.contains("Host key verification failed") {
        return anyhow::anyhow!(
            "{message} This machine has not been connected to before, so its key is not \
             trusted yet. Ssh to it once from the backend -- that is where the decision to \
             trust a key belongs -- and try again.",
        );
    }
    if message.contains("Permission denied") {
        return anyhow::anyhow!(
            "{message} The key named here has to be authorized on that machine, and the path \
             is read on the backend rather than on the phone.",
        );
    }
    err
}

/// A short, stable, filename-safe id derived from a label. Derived once when a
/// machine is added and then fixed, so the label stays editable. Collisions are
/// resolved by the caller, which is the only place that knows what exists.
pub fn id_from(label: &str) -> String {
    let slug: String = label
        .chars()
        .map(|c| {
            if c.is_ascii_alphanumeric() {
                c.to_ascii_lowercase()
            } else {
                '-'
            }
        })
        .collect();
    let slug = slug.trim_matches('-').replace("--", "-");
    if slug.is_empty() {
        crate::session::random_hex()
    } else {
        slug.chars().take(32).collect()
    }
}

/// Normalises what a phone keyboard produced: trims, drops blanks, and
/// expands a leading `~` the way a shell would.
pub fn tidy(value: &str) -> Option<String> {
    let value = value.trim();
    if value.is_empty() {
        return None;
    }
    Some(match value.strip_prefix("~/") {
        Some(rest) => match std::env::home_dir() {
            Some(home) => home.join(rest).to_string_lossy().into_owned(),
            None => value.to_string(),
        },
        None => value.to_string(),
    })
}

/// The inverse of [`tidy`]'s expansion: an absolute path under this machine's
/// home, written back as `~/…`, so that a working directory reads on a phone the
/// way it is written by hand.
///
/// Applied only to paths on **this** machine. `$HOME` here says nothing about
/// the home directory of a machine reached over ssh, so a remote path is stored
/// exactly as it was typed and the remote shell is what expands it.
pub fn shorten_home(path: &str) -> String {
    let Some(home) = std::env::home_dir() else {
        return path.to_string();
    };
    let home = home.to_string_lossy();
    // The separator has to be part of the match, or `/home/bobby` would be read
    // as a path inside `/home/bob`.
    match path.strip_prefix(home.as_ref()) {
        Some("") => "~".to_string(),
        Some(rest) if rest.starts_with('/') => format!("~{rest}"),
        _ => path.to_string(),
    }
}

/// Runs a launch to completion and returns its stdout as text.
///
/// The common case of [`Transport::capture_with_input`]: nothing on stdin, a
/// failure reported as the machine's own words (ssh's "Permission denied" is the
/// useful half of why a machine cannot be reached), and the output read as text
/// because every caller here is asking a question whose answer is words.
impl Transport {
    pub async fn capture(&self, launch: &Launch) -> Result<String> {
        let captured = self
            .capture_with_input(launch, super::session::transport::Input::None)
            .await?;
        Ok(String::from_utf8_lossy(&captured.ok()?).into_owned())
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    /// The two halves of a home-relative path, which have to be inverses: what
    /// is stored is what the phone draws, and what the phone sends back is what
    /// a process is started in.
    #[test]
    fn a_home_path_shortens_and_expands_back() {
        let Some(home) = std::env::home_dir() else {
            return;
        };
        let full = home.join("repos/ai-app-2");
        let full = full.to_string_lossy();
        assert_eq!(shorten_home(&full), "~/repos/ai-app-2");
        assert_eq!(shorten_home(&home.to_string_lossy()), "~");
        assert_eq!(tidy("~/repos/ai-app-2").as_deref(), Some(full.as_ref()));

        // Not a prefix match on the characters: a sibling directory whose name
        // merely starts with the home directory's is not inside it.
        let sibling = format!("{}-backup/notes", home.to_string_lossy());
        assert_eq!(shorten_home(&sibling), sibling);
        assert_eq!(shorten_home("/etc/hosts"), "/etc/hosts");
    }

    #[test]
    fn codex_models_are_the_selectable_non_hidden_catalog_entries() {
        let answer = json!({"result": {"data": [
            {"model": "gpt-small", "hidden": false},
            {"model": "gpt-hidden", "hidden": true},
            {"model": "gpt-large", "hidden": false}
        ]}});
        assert_eq!(
            parse_codex_models(&answer).unwrap(),
            vec!["gpt-small".to_string(), "gpt-large".to_string()]
        );
    }

    #[test]
    fn a_failed_codex_catalog_is_not_reported_as_an_empty_one() {
        let answer = json!({"error": {"message": "login required"}});
        assert_eq!(
            parse_codex_models(&answer).unwrap_err().to_string(),
            "Codex could not list models: login required"
        );
    }
}