compute/src/mcp.rs

//! MCP server: the Governor's surface, exposed to agents as tools.
//!
//! Agents already drive Compute by shelling out to the CLI and scraping the
//! column-aligned output. That works until a column moves. Speaking MCP means
//! the agent gets the same numbers the Governor reasons about, as data, with a
//! schema attached — and it means the *shape* of what an agent may do becomes
//! something this file decides rather than something `bash` decides.
//!
//! Two deliberate constraints shape everything below.
//!
//! **Read-first.** Every tool here except `scene_activate` is a pure read.
//! Activating a Scene stops every registered model that is not in it, which on
//! this box means taking down whatever is currently serving — a live voice
//! pipeline included. There is no undo an agent can reach for: if the target
//! Scene then fails to come up, recovery depends on rollback that may itself
//! fail. That asymmetry is why `scene adopt`, `scene resume`, and
//! `scene deactivate` are absent entirely; they are operator verbs whose
//! correctness depends on knowing what the box was doing five minutes ago.
//! `scene_activate` is exposed because planning a switch is genuinely the
//! useful thing an agent wants, and it defaults to planning only.
//!
//! **The safety boundary is the operator's, not the agent's.** A `dry_run`
//! argument defaulting to `true` documents intent but guards nothing: the agent
//! writes the arguments, so it can write `false`. The only boundary an agent
//! cannot cross is one set before it connects, so a real transition also
//! requires `lumbridge-compute mcp --allow-activate`, chosen by the human who
//! launched the server. Without that flag `dry_run: false` is refused, and the
//! refusal says so rather than silently planning instead.

use anyhow::{Context, Result};
use rmcp::handler::server::router::tool::ToolRouter;
use rmcp::handler::server::wrapper::Parameters;
use rmcp::model::{Implementation, ServerCapabilities, ServerInfo};
use rmcp::{tool, tool_handler, tool_router, ErrorData, Json, ServerHandler, ServiceExt};
use schemars::JsonSchema;
use serde::{Deserialize, Serialize};
use std::os::fd::FromRawFd;
use std::path::{Path, PathBuf};
use std::sync::Arc;

use crate::config::{find_scene, load_scenes, Registry};
use crate::proc::State;
use crate::{eval, governor, lifecycle, mem};

pub fn run(root: &Path, allow_activate: bool) -> Result<()> {
    let transport_stdout = hand_over_stdout()?;
    // The banner has to go to stderr for the same reason everything else does;
    // it doubles as confirmation to the operator that the mutating tool is off.
    eprintln!(
        "Lumbridge Compute MCP server on stdio · root {} · scene activation {}",
        root.display(),
        if allow_activate {
            "ENABLED (--allow-activate)"
        } else {
            "disabled; plan only"
        }
    );

    let server = ComputeMcp::new(root.to_path_buf(), allow_activate);
    let runtime = tokio::runtime::Builder::new_current_thread()
        .enable_all()
        .build()
        .context("building the MCP runtime")?;
    runtime.block_on(async move {
        let transport = (
            tokio::io::stdin(),
            tokio::fs::File::from_std(transport_stdout),
        );
        let service = server
            .serve(transport)
            .await
            .context("negotiating the MCP stdio session")?;
        service.waiting().await.context("serving MCP over stdio")?;
        Ok(())
    })
}

/// Hand the real stdout to the transport and point fd 1 at stderr.
///
/// On stdio transport fd 1 *is* the JSON-RPC framing, and this crate reports
/// progress with `println!` throughout — `lifecycle::activate` narrates every
/// stop and start. One such line interleaved into the framing desynchronises
/// the client mid-transition, which is the worst possible moment to lose it.
/// Rather than audit every print (and every future one), move the file
/// descriptor: library output lands on stderr, where operators already read
/// this server's logs, and the protocol gets a channel nothing else can write.
fn hand_over_stdout() -> Result<std::fs::File> {
    let saved = unsafe { libc::dup(libc::STDOUT_FILENO) };
    if saved < 0 {
        return Err(std::io::Error::last_os_error())
            .context("duplicating stdout for the MCP transport");
    }
    if unsafe { libc::dup2(libc::STDERR_FILENO, libc::STDOUT_FILENO) } < 0 {
        let error = std::io::Error::last_os_error();
        unsafe { libc::close(saved) };
        return Err(error).context("redirecting stdout to stderr");
    }
    // SAFETY: `saved` is a fresh descriptor from `dup` that nothing else owns.
    Ok(unsafe { std::fs::File::from_raw_fd(saved) })
}

#[derive(Clone)]
struct ComputeMcp {
    root: Arc<PathBuf>,
    allow_activate: bool,
    tool_router: ToolRouter<ComputeMcp>,
}

#[tool_router]
impl ComputeMcp {
    fn new(root: PathBuf, allow_activate: bool) -> Self {
        Self {
            root: Arc::new(root),
            allow_activate,
            tool_router: Self::tool_router(),
        }
    }

    /// Governor status for this node: unified-memory totals, the hard budget
    /// and its margins, how much is committed by models that are actually
    /// serving, how much is still admittable, the persisted Scene state, and
    /// the last transition or watchdog error. Read this before reasoning about
    /// whether anything else will fit.
    #[tool(
        name = "governor_status",
        annotations(
            title = "Governor status",
            read_only_hint = true,
            open_world_hint = false
        )
    )]
    async fn governor_status(&self) -> Result<Json<StatusReport>, ErrorData> {
        let root = self.root.clone();
        offload(move || collect_status(&root)).await.map(Json)
    }

    /// List every model in the node's registry with its declared worst-case
    /// footprint, serving port, and whether it is currently serving. Footprints
    /// are what admission control is decided against, not observed usage.
    #[tool(
        name = "model_list",
        annotations(
            title = "List registered models",
            read_only_hint = true,
            open_world_hint = false
        )
    )]
    async fn model_list(&self) -> Result<Json<ModelList>, ErrorData> {
        let root = self.root.clone();
        offload(move || collect_models(&root)).await.map(Json)
    }

    /// List the Scenes this node can activate, with each one's total footprint
    /// and member models. A Scene is a named set of models brought up as a
    /// single unit.
    #[tool(
        name = "scene_list",
        annotations(title = "List scenes", read_only_hint = true, open_world_hint = false)
    )]
    async fn scene_list(&self) -> Result<Json<SceneList>, ErrorData> {
        let root = self.root.clone();
        offload(move || collect_scenes(&root)).await.map(Json)
    }

    /// Show one Scene: its models, their footprints and live state, and the
    /// Governor's admission verdict — whether the Scene fits its budget once
    /// the safety margin is added, and by how much it fits or misses.
    #[tool(
        name = "scene_show",
        annotations(
            title = "Show a scene and its admission verdict",
            read_only_hint = true,
            open_world_hint = false
        )
    )]
    async fn scene_show(
        &self,
        Parameters(params): Parameters<SceneNameParams>,
    ) -> Result<Json<SceneDetail>, ErrorData> {
        let root = self.root.clone();
        offload(move || collect_scene(&root, &params.name))
            .await
            .map(Json)
    }

    /// List the evaluation suites available on this node, with their versions
    /// and case counts. Running a suite is deliberately not exposed: it drives
    /// real load against a serving model for an unbounded time.
    #[tool(
        name = "eval_list",
        annotations(
            title = "List eval suites",
            read_only_hint = true,
            open_world_hint = false
        )
    )]
    async fn eval_list(&self) -> Result<Json<EvalList>, ErrorData> {
        let root = self.root.clone();
        offload(move || collect_evals(&root)).await.map(Json)
    }

    /// Plan — or, if the operator allowed it, perform — a Scene transition.
    ///
    /// With `dry_run` true (the default) this runs every check that guards a
    /// real transition and returns the models it would stop and start, changing
    /// nothing. With `dry_run` false it actually switches the node: models
    /// outside the Scene are stopped, which will interrupt anything they serve,
    /// and the call blocks until every Scene model reports healthy or the
    /// transition rolls back. A real transition additionally requires the
    /// server to have been started with `--allow-activate`; without it,
    /// `dry_run: false` is refused. Ask a human before setting it.
    #[tool(
        name = "scene_activate",
        annotations(
            title = "Activate a scene (plans by default)",
            read_only_hint = false,
            destructive_hint = true,
            idempotent_hint = false,
            open_world_hint = false
        )
    )]
    async fn scene_activate(
        &self,
        Parameters(params): Parameters<SceneActivateParams>,
    ) -> Result<Json<ActivationReport>, ErrorData> {
        if !params.dry_run && !self.allow_activate {
            return Err(ErrorData::invalid_params(
                "this MCP server is running plan-only: a real Scene transition needs an \
                 operator to restart it as `lumbridge-compute mcp --allow-activate`. \
                 Re-run with dry_run: true to see the plan.",
                None,
            ));
        }
        let root = self.root.clone();
        offload(move || activate_scene(&root, &params.name, params.dry_run))
            .await
            .map(Json)
    }
}

#[tool_handler(router = self.tool_router)]
impl ServerHandler for ComputeMcp {
    fn get_info(&self) -> ServerInfo {
        ServerInfo::new(ServerCapabilities::builder().enable_tools().build())
            .with_server_info(Implementation::new(
                "lumbridge-compute",
                env!("CARGO_PKG_VERSION"),
            ))
            .with_instructions(
                "Lumbridge Compute runs many AI models on one unified-memory box. There is a \
                 single memory pool shared by CPU and GPU, so over-committing does not fail \
                 gracefully — the machine thrashes and wedges before the OOM killer acts. The \
                 Governor prevents that by admitting a model only if its declared footprint \
                 plus a safety margin still fits a hard budget.\n\n\
                 Start from `governor_status` for what is committed and what is admittable, \
                 and `scene_show` for whether a named Scene would be admitted. Everything is \
                 read-only except `scene_activate`, which plans by default. Activating a Scene \
                 stops every registered model outside it, so treat it as a production change \
                 and get a human's agreement before asking for a non-dry run.",
            )
    }
}

// ---- tool parameters -------------------------------------------------------

#[derive(Debug, Deserialize, JsonSchema)]
struct SceneNameParams {
    /// Scene name, as reported by `scene_list`.
    name: String,
}

#[derive(Debug, Deserialize, JsonSchema)]
struct SceneActivateParams {
    /// Scene name, as reported by `scene_list`.
    name: String,
    /// Validate and report the plan without changing anything. Defaults to
    /// true; set it to false only with a human's explicit agreement.
    #[serde(default = "yes")]
    dry_run: bool,
}

fn yes() -> bool {
    true
}

// ---- tool results ----------------------------------------------------------
//
// Lists are wrapped in objects rather than returned bare: MCP structured
// content must be a JSON object, and a named field leaves room to report
// alongside the list later without breaking a client's schema.

#[derive(Debug, Serialize, JsonSchema)]
struct StatusReport {
    memory_total_gb: f64,
    memory_available_gb: f64,
    budget_gb: f64,
    safety_margin_gb: f64,
    watchdog_floor_gb: f64,
    /// Sum of the footprints of every registered model currently serving.
    committed_gb: f64,
    /// What a new model could still claim without breaching the safety margin.
    headroom_gb: f64,
    /// Scene whose exact model health was last verified.
    active_scene: Option<String>,
    /// Scene the node should return to after a restart.
    desired_scene: Option<String>,
    /// Previous proven Scene, used if the desired one cannot resume.
    last_known_good_scene: Option<String>,
    last_error: Option<String>,
    running: Vec<RunningModel>,
}

#[derive(Debug, Serialize, JsonSchema)]
struct RunningModel {
    id: String,
    name: String,
    footprint_gb: f64,
    port: Option<u16>,
    /// True when Compute owns this process identity and may signal it. False
    /// means the model is serving but was started outside Compute, so a Scene
    /// transition will refuse to touch it until it is adopted.
    compute_owned: bool,
}

#[derive(Debug, Serialize, JsonSchema)]
struct ModelList {
    models: Vec<ModelEntry>,
}

#[derive(Debug, Serialize, JsonSchema)]
struct ModelEntry {
    id: String,
    name: String,
    /// Declared worst-case unified memory once serving, not observed usage.
    footprint_gb: f64,
    port: Option<u16>,
    running: bool,
}

#[derive(Debug, Serialize, JsonSchema)]
struct SceneList {
    scenes: Vec<SceneEntry>,
}

#[derive(Debug, Serialize, JsonSchema)]
struct SceneEntry {
    name: String,
    version: u32,
    description: String,
    /// Total footprint of the Scene's models that exist in the registry.
    footprint_gb: f64,
    models: Vec<String>,
}

#[derive(Debug, Serialize, JsonSchema)]
struct SceneDetail {
    name: String,
    version: u32,
    description: String,
    budget_gb: f64,
    models: Vec<SceneModel>,
    /// Model ids the Scene references that the registry does not define.
    missing_models: Vec<String>,
    footprint_gb: f64,
    /// Footprint plus the Governor's safety margin — the figure compared
    /// against the budget.
    required_gb: f64,
    /// The Governor's verdict: would this Scene be admitted?
    admits: bool,
    verdict: String,
}

#[derive(Debug, Serialize, JsonSchema)]
struct SceneModel {
    id: String,
    /// Absent when the id is missing from the registry.
    name: Option<String>,
    footprint_gb: Option<f64>,
    running: bool,
}

#[derive(Debug, Serialize, JsonSchema)]
struct EvalList {
    suites: Vec<EvalSuite>,
}

#[derive(Debug, Serialize, JsonSchema)]
struct EvalSuite {
    name: String,
    version: u32,
    cases: usize,
    description: String,
}

#[derive(Debug, Serialize, JsonSchema)]
struct ActivationReport {
    scene: String,
    budget_gb: f64,
    /// Serving models outside the Scene, which the transition stops first.
    stop: Vec<String>,
    /// Scene models not yet serving, in the order they would be admitted.
    start: Vec<String>,
    /// False when this was a plan and the node was left untouched.
    applied: bool,
}

// ---- collectors ------------------------------------------------------------
//
// These compose the same primitives the CLI prints from — `governor`, `mem`,
// `config`, `proc::State` — rather than parsing the CLI's output. The numbers
// an agent sees are therefore the numbers admission control uses, by
// construction.

fn collect_status(root: &Path) -> Result<StatusReport> {
    let registry = Registry::load(root)?;
    let memory = mem::read()?;
    let committed = governor::committed_gb(&registry);
    let managed = State::load_checked(root)?;
    let budget = governor::DEFAULT_BUDGET_GB;

    let running = governor::running_ids(&registry)
        .into_iter()
        .map(|id| {
            let model = &registry.models[&id];
            RunningModel {
                name: model.name.clone(),
                footprint_gb: model.footprint_gb,
                port: model.serve.port,
                compute_owned: managed
                    .procs
                    .get(&id)
                    .is_some_and(|process| process.owned_alive()),
                id,
            }
        })
        .collect();

    Ok(StatusReport {
        memory_total_gb: memory.total_gb,
        memory_available_gb: memory.available_gb,
        budget_gb: budget,
        safety_margin_gb: governor::SAFETY_MARGIN_GB,
        watchdog_floor_gb: governor::WATCHDOG_FLOOR_GB,
        committed_gb: committed,
        headroom_gb: (budget - committed - governor::SAFETY_MARGIN_GB).max(0.0),
        active_scene: managed.active_scene,
        desired_scene: managed.desired_scene,
        last_known_good_scene: managed.last_known_good_scene,
        last_error: managed.last_error,
        running,
    })
}

fn collect_models(root: &Path) -> Result<ModelList> {
    let registry = Registry::load(root)?;
    Ok(ModelList {
        models: registry
            .models
            .iter()
            .map(|(id, model)| ModelEntry {
                id: id.clone(),
                name: model.name.clone(),
                footprint_gb: model.footprint_gb,
                port: model.serve.port,
                running: governor::is_running(model),
            })
            .collect(),
    })
}

fn collect_scenes(root: &Path) -> Result<SceneList> {
    let registry = Registry::load(root)?;
    Ok(SceneList {
        scenes: load_scenes(root)?
            .into_iter()
            .map(|scene| SceneEntry {
                name: scene.metadata.name,
                version: scene.metadata.version,
                description: scene.metadata.description,
                footprint_gb: scene
                    .models
                    .iter()
                    .filter_map(|id| registry.models.get(id))
                    .map(|model| model.footprint_gb)
                    .sum(),
                models: scene.models,
            })
            .collect(),
    })
}

fn collect_scene(root: &Path, name: &str) -> Result<SceneDetail> {
    let registry = Registry::load(root)?;
    let scene = find_scene(root, name)?;
    let budget = scene.budget_gb.unwrap_or(governor::DEFAULT_BUDGET_GB);

    let mut footprint = 0.0;
    let mut missing = Vec::new();
    let mut models = Vec::new();
    for id in &scene.models {
        match registry.models.get(id) {
            Some(model) => {
                footprint += model.footprint_gb;
                models.push(SceneModel {
                    id: id.clone(),
                    name: Some(model.name.clone()),
                    footprint_gb: Some(model.footprint_gb),
                    running: governor::is_running(model),
                });
            }
            None => {
                missing.push(id.clone());
                models.push(SceneModel {
                    id: id.clone(),
                    name: None,
                    footprint_gb: None,
                    running: false,
                });
            }
        }
    }

    // Same three-way verdict the CLI prints: a missing id is fatal regardless
    // of arithmetic, because the Scene cannot be resolved at all.
    let required = footprint + governor::SAFETY_MARGIN_GB;
    let (admits, verdict) = if !missing.is_empty() {
        (
            false,
            format!("{} model(s) missing from the registry", missing.len()),
        )
    } else if required <= budget {
        (
            true,
            format!("fits with {:.1} GB to spare", budget - required),
        )
    } else {
        (
            false,
            format!("exceeds the budget by {:.1} GB", required - budget),
        )
    };

    Ok(SceneDetail {
        name: scene.metadata.name,
        version: scene.metadata.version,
        description: scene.metadata.description,
        budget_gb: budget,
        models,
        missing_models: missing,
        footprint_gb: footprint,
        required_gb: required,
        admits,
        verdict,
    })
}

fn collect_evals(root: &Path) -> Result<EvalList> {
    Ok(EvalList {
        suites: eval::catalogue(root)?
            .into_iter()
            .map(|suite| EvalSuite {
                name: suite.name,
                version: suite.version,
                cases: suite.cases,
                description: suite.description,
            })
            .collect(),
    })
}

fn activate_scene(root: &Path, name: &str, dry_run: bool) -> Result<ActivationReport> {
    // Plan first either way. On a real transition this is the plan we report;
    // `activate` re-derives and re-validates its own under the transition lock,
    // so the report describes intent and the lock still owns the truth.
    let plan = lifecycle::plan(root, name)?;
    if !dry_run {
        lifecycle::activate(root, name, false)?;
    }
    Ok(ActivationReport {
        scene: plan.scene,
        budget_gb: plan.budget_gb,
        stop: plan.stop,
        start: plan.start,
        applied: !dry_run,
    })
}

/// Run a synchronous Governor call off the transport's thread.
///
/// Every read here touches `/proc` and probes serving ports with blocking
/// socket timeouts, and a real transition can sit for minutes waiting on model
/// health. Doing that inline would stall the JSON-RPC reader for the duration,
/// so the synchronous core stays on the blocking pool where it belongs.
async fn offload<T, F>(work: F) -> Result<T, ErrorData>
where
    F: FnOnce() -> Result<T> + Send + 'static,
    T: Send + 'static,
{
    match tokio::task::spawn_blocking(work).await {
        Ok(Ok(value)) => Ok(value),
        // `{:#}` keeps anyhow's context chain, which is where the actionable
        // half of a Compute error lives ("adopt the active Scene first").
        Ok(Err(error)) => Err(ErrorData::internal_error(format!("{error:#}"), None)),
        Err(join) => Err(ErrorData::internal_error(
            format!("Compute worker task failed: {join}"),
            None,
        )),
    }
}