hyperhive/hive-bash-mcp/src/mcp.rs

//! MCP tool surface for `hive-bash-daemon`, served directly over
//! streamable-http — no stdio bridge, no round-trip socket. The daemon
//! already owns the subprocess runner in-process, so the tool handlers
//! below call straight into [`crate::runner`]. Mirrors
//! `hive-agent-mcp::mcp::serve_http`'s shape (persistent daemon, stable
//! URL claude reconnects to every turn instead of respawning a stdio
//! child).

use std::fmt::Write as _;

use rmcp::{
    ServerHandler,
    handler::server::wrapper::Parameters,
    schemars::{self, JsonSchema},
    tool, tool_handler, tool_router,
};
use serde::Deserialize;

use crate::protocol::{TaskFile, TaskStatus};
use crate::runner::{self, SUMMARY_BYTES};

fn status_str(status: &TaskStatus) -> &'static str {
    match status {
        TaskStatus::Pending => "pending",
        TaskStatus::Running => "running",
        TaskStatus::Done => "done",
        TaskStatus::TimedOut => "timed_out",
        TaskStatus::Interrupted => "interrupted",
        TaskStatus::Killed => "killed",
    }
}

/// What a state means and what the caller should do about it — carried on
/// every `format_task` answer instead of the tool description enumerating
/// all six up front, so a caller who only ever sees one state at a time
/// still learns what it means without reading anything else.
fn status_explanation(status: &TaskStatus) -> &'static str {
    match status {
        TaskStatus::Pending => "queued, not started yet",
        TaskStatus::Running => {
            "still going — check back later, or pass wait_seconds to block briefly"
        }
        TaskStatus::Done => "finished normally",
        TaskStatus::TimedOut => "exceeded timeout_secs and was killed",
        TaskStatus::Interrupted => {
            "the daemon restarted mid-task — the process is gone, no result to report"
        }
        TaskStatus::Killed => {
            "killed on request via `kill` — a signal was sent to the whole process group"
        }
    }
}

/// Format a `TaskFile` as the human-readable status string claude sees.
fn format_task(task: &TaskFile) -> String {
    let mut out = format!(
        "task `{}`: status={} ({})",
        task.id,
        status_str(&task.status),
        status_explanation(&task.status)
    );

    if let Some(code) = task.exit_code {
        let _ = write!(out, ", exit={code}");
    }
    if let (Some(started), None) = (task.started_at, task.completed_at) {
        let elapsed = (chrono::Utc::now() - started).num_seconds();
        let _ = write!(out, ", running for {elapsed}s");
    }
    if let (Some(completed), Some(started)) = (task.completed_at, task.started_at) {
        let took = (completed - started).num_seconds();
        let _ = write!(out, ", took {took}s");
    }

    let out_file = crate::paths::task_out(&task.id);
    let err_file = crate::paths::task_err(&task.id);
    let out_len = std::fs::metadata(&out_file).map_or(0, |m| m.len());
    let err_len = std::fs::metadata(&err_file).map_or(0, |m| m.len());

    if let Some(stdout) = &task.stdout_tail {
        let s = stdout.trim();
        if !s.is_empty() {
            let _ = write!(out, "\n\nstdout:\n```\n{s}\n```");
        }
    }
    if out_len > SUMMARY_BYTES as u64 {
        let _ = write!(out, "\n\nFull stdout lives in `{}`", out_file.display());
    }

    if let Some(stderr) = &task.stderr_tail {
        let s = stderr.trim();
        if !s.is_empty() {
            let _ = write!(out, "\n\nstderr:\n```\n{s}\n```");
        }
    }
    if err_len > SUMMARY_BYTES as u64 {
        let _ = write!(out, "\n\nFull stderr lives in `{}`", err_file.display());
    }

    out
}

/// Appended to a `status` result when the caller parked on a
/// `wait_seconds` poll that expired while the task was still running —
/// nudges the model to spend the idle time on other work instead of
/// immediately re-waiting on the same task.
const BASH_IDLE_WAIT_HINT: &str = "\n\nThe task is still running — your wait timed out before it \
finished. If you have other useful work, do that and check back later (the task keeps running, and \
it surfaces in your loose-ends when it completes) rather than immediately re-waiting.";

// ---------------------------------------------------------------------------
// MCP server
// ---------------------------------------------------------------------------

#[derive(Debug, Deserialize, JsonSchema)]
struct BashRunArgs {
    /// Shell command to run (passed to `bash -c`).
    cmd: String,
    /// Timeout in seconds. Defaults to `None` (no timeout) — task runs until
    /// natural exit. Pass an explicit value to kill the task after N seconds
    /// and mark it `timed_out`.
    #[serde(default)]
    timeout_secs: Option<u64>,
    /// Optional inline wait: `run` polls for up to `wait_seconds`
    /// (capped at 30) before returning. When the task finishes within the
    /// window the full status is returned immediately (nothing to handle
    /// later); when the timeout expires the task keeps running and the
    /// normal `task started: id=<id>` response is returned. Defaults to 3s.
    /// Pass `0` to disable inline waiting and always get the immediate response.
    #[serde(default = "default_wait")]
    wait_seconds: Option<u64>,
    /// Optional task name. When set it becomes the task id, so it appears
    /// in `status` lookups and your loose-ends list — handy for recognising
    /// a task later instead of an opaque hex id. A name can be reused once
    /// its previous task has finished; reusing a name whose task is still
    /// running is rejected. Allowed chars: `[a-z0-9-]` (a valid
    /// identifier — lowercase, digits, hyphen; max 63). Omit to get the
    /// auto-generated id.
    #[serde(default)]
    name: Option<String>,
}

#[allow(
    clippy::unnecessary_wraps,
    reason = "serde `#[serde(default = ...)]` requires the default fn's return type to match the field's `Option<u64>`, so the `Some` wrap is mandatory even though the value is constant"
)]
fn default_wait() -> Option<u64> {
    Some(3)
}

#[derive(Debug, Deserialize, JsonSchema)]
struct BashStatusArgs {
    /// Task ID returned by `run`.
    id: String,
    /// Optional inline wait: `status` polls for up to `wait_seconds`
    /// (capped at 30) before returning. Useful to avoid a separate
    /// round-trip when the task is expected to finish soon.
    #[serde(default)]
    wait_seconds: Option<u64>,
}

#[derive(Debug, Deserialize, JsonSchema)]
struct BashKillArgs {
    /// Task ID (from `run`) to kill.
    id: String,
    /// `false` (default): SIGINT to the task's process group — graceful, the
    /// process can clean up. `true`: SIGKILL. Either way the whole process
    /// group is signalled, so children of the shell (e.g. a `cargo`/`nix`
    /// invocation) are stopped too, not just the shell. Fire-and-forget: the
    /// call sends the signal and returns without waiting. If a SIGINT'd task
    /// doesn't exit, call `kill` again with `force: true` to send SIGKILL.
    #[serde(default)]
    force: bool,
}

#[derive(Clone)]
struct BashMcp;

#[tool_router]
impl BashMcp {
    #[tool(
        description = "Run a shell command in the background. Returns a task ID immediately — \
            do NOT wait inline. When the command finishes it surfaces as a todo in your \
            loose-ends (the exit status plus a `Read(<path>)` pointer to the captured \
            `.out`/`.err` files — read only what you need); handle it on a future turn. Use \
            `status` to poll within the same turn if needed. `timeout_secs` defaults to \
            `None` (no timeout) — task runs until natural exit; pass a value to kill after \
            N seconds. Pass `wait_seconds` (capped at 30) to wait inline for fast commands: \
            if the task finishes within the window you get the full status immediately \
            (nothing to handle later); otherwise it keeps running and you get \
            `task started: id=<id>`. Defaults to 3; pass `0` to disable inline waiting. \
            Pass `name` to label the task with a memorable id (shown in `status` lookups \
            and your loose-ends) instead of an opaque hex id; reusable once the prior task \
            has finished, rejected while one is still running."
    )]
    async fn run(&self, Parameters(args): Parameters<BashRunArgs>) -> String {
        let id = match runner::submit_task(args.cmd, args.timeout_secs, args.name) {
            Ok(id) => id,
            Err(e) => return format!("bash_run error: {e:#}"),
        };
        let wait = args.wait_seconds.unwrap_or(0);
        if wait > 0
            && let Some(task) = runner::wait_for_task(&id, wait).await
            && matches!(
                task.status,
                TaskStatus::Done
                    | TaskStatus::TimedOut
                    | TaskStatus::Interrupted
                    | TaskStatus::Killed
            )
        {
            return format_task(&task);
        }
        format!("task started: id={id}")
    }

    #[tool(
        description = "Check the status of a background bash task by its ID (from `run`). The \
            answer names the state it found and what to do about it, plus the exit code if \
            finished and a tail of stdout/stderr. \
            Pass `wait_seconds` (capped at 30) to wait inline for the task to finish: when the \
            task finishes within the window the full status is returned immediately. Useful to \
            avoid a separate round-trip when the task is expected to finish soon."
    )]
    async fn status(&self, Parameters(args): Parameters<BashStatusArgs>) -> String {
        let waited = args.wait_seconds.is_some_and(|w| w > 0);
        let wait = args.wait_seconds.unwrap_or(0);
        let task = if wait > 0 {
            runner::wait_for_task(&args.id, wait).await
        } else {
            runner::read_task(&args.id)
        };
        match task {
            Some(t) => {
                let mut out = format_task(&t);
                if waited && matches!(t.status, TaskStatus::Pending | TaskStatus::Running) {
                    out.push_str(BASH_IDLE_WAIT_HINT);
                }
                out
            }
            None => format!("unknown task id `{}`", args.id),
        }
    }

    #[tool(
        description = "Kill a background bash task you started (by its ID from `run`). \
            `force: false` (default) sends SIGINT — graceful, lets the process clean up; \
            `force: true` sends SIGKILL. Either way the task's whole process group is \
            signalled, so a runaway child (cargo/nix/etc.) is stopped too, not just the shell. \
            Fire-and-forget: sends the signal and returns without waiting. If a SIGINT'd task \
            doesn't exit, call kill again with `force: true` to SIGKILL. A still-pending task \
            is cancelled before it starts. The task ends as `killed` and surfaces in your \
            loose-ends like any completion."
    )]
    async fn kill(&self, Parameters(args): Parameters<BashKillArgs>) -> String {
        let (killed, was_running) = runner::kill_task(&args.id, args.force);
        if !killed {
            return format!(
                "no running or pending task with id `{}` could be killed (already finished, unknown, \
                 or its pending file could not be updated — see the daemon log)",
                args.id
            );
        }
        if was_running {
            let sig = if args.force { "SIGKILL" } else { "SIGINT" };
            format!(
                "task `{}`: {sig} sent to its process group; it transitions to `killed` \
                 once the process exits and then surfaces in your loose-ends.",
                args.id
            )
        } else {
            format!(
                "task `{}` was pending — cancelled before it started.",
                args.id
            )
        }
    }
}

#[tool_handler]
impl ServerHandler for BashMcp {}

/// Run the MCP server over HTTP (rmcp streamable-http transport) on `addr`.
///
/// Sole transport — there is no stdio mode. Long-lived so claude
/// reconnects to the stable URL each turn instead of respawning a stdio
/// child; since the daemon already owns [`crate::runner`] in-process,
/// tool calls need no round-trip to anywhere.
///
/// Binds loopback only in practice; the default `allowed_hosts`
/// (`localhost`/`127.0.0.1`/`::1`) rejects Host headers from anywhere else.
///
/// # Errors
///
/// Returns an error if the listener cannot bind `addr` or the HTTP server
/// exits with a fatal error.
pub async fn serve_http(addr: std::net::SocketAddr) -> anyhow::Result<()> {
    use rmcp::transport::streamable_http_server::{
        StreamableHttpServerConfig, StreamableHttpService, session::local::LocalSessionManager,
    };
    let mut session_manager = LocalSessionManager::default();
    // rmcp's default session idle timeout is 5 minutes, meant to reap zombie
    // connections lost to an HTTP/2 RST_STREAM — far too short for a claude
    // turn that goes several minutes between bash-MCP calls (a long Read/Edit
    // stretch, a background-task wait). That trips a client-visible "MCP
    // server ... session expired" mid-turn; self-recovering on the next call,
    // but confusing. Give it real headroom instead of turning it off outright.
    session_manager.session_config.keep_alive = Some(std::time::Duration::from_hours(24));
    let session_manager = std::sync::Arc::new(session_manager);
    let service = StreamableHttpService::new(
        || Ok(BashMcp),
        session_manager,
        StreamableHttpServerConfig::default(),
    );
    let app = axum::Router::new().nest_service("/mcp", service);
    let listener = tokio::net::TcpListener::bind(addr).await?;
    tracing::info!(%addr, "serving hive-bash MCP over streamable-http at /mcp");
    axum::serve(listener, app).await?;
    Ok(())
}

#[cfg(test)]
mod status_hint_tests {
    use super::{BASH_IDLE_WAIT_HINT, format_task};
    use crate::protocol::{TaskFile, TaskStatus};
    use crate::test_util::with_harness_dir;

    fn task(status: TaskStatus) -> TaskFile {
        TaskFile {
            id: "t1".to_owned(),
            cmd: "echo hi".to_owned(),
            timeout_secs: None,
            status,
            created_at: chrono::DateTime::from_timestamp(1, 0).unwrap_or_default(),
            started_at: Some(chrono::DateTime::from_timestamp(1, 0).unwrap_or_default()),
            completed_at: None,
            exit_code: None,
            stdout_tail: None,
            stderr_tail: None,
        }
    }

    #[test]
    fn running_task_after_wait_appends_idle_hint() {
        with_harness_dir(|| {
            let t = task(TaskStatus::Running);
            let mut out = format_task(&t);
            out.push_str(BASH_IDLE_WAIT_HINT);
            assert!(out.contains(BASH_IDLE_WAIT_HINT));
        });
    }

    #[test]
    fn done_task_formats_without_hint() {
        with_harness_dir(|| {
            let out = format_task(&task(TaskStatus::Done));
            assert!(!out.contains(BASH_IDLE_WAIT_HINT));
            assert!(out.contains("status=done"));
        });
    }

    #[test]
    fn every_status_answer_names_its_state_and_the_next_move() {
        // Pins the split this module's doc describes: the tool description
        // says what the tool is for, and each individual answer — not a
        // list the caller has to cross-reference — says what state it got
        // and what to do about it. One state's explanation text is
        // distinct enough from the others that a copy-paste regression
        // (e.g. `Killed` picking up `TimedOut`'s wording) would fail this.
        with_harness_dir(|| {
            for status in [
                TaskStatus::Pending,
                TaskStatus::Running,
                TaskStatus::Done,
                TaskStatus::TimedOut,
                TaskStatus::Interrupted,
                TaskStatus::Killed,
            ] {
                let explanation = super::status_explanation(&status);
                let out = format_task(&task(status.clone()));
                assert!(
                    out.contains(explanation),
                    "status={status:?} explanation missing from: {out}"
                );
            }
        });
    }
}