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

357 lines
14 KiB
Rust

//! `hive-bash-mcp` binary — stdio MCP server claude spawns per turn.
//! Thin protocol bridge: every tool call → connect to the daemon's
//! unix socket → write a JSON request line → read the JSON response →
//! return the result to claude.
//!
//! No subprocess management at this entrypoint — the daemon owns that.
//! Cold-starts in milliseconds.
use anyhow::{Context, Result};
use rmcp::{
ServerHandler, ServiceExt,
handler::server::wrapper::Parameters,
schemars::{self, JsonSchema},
tool, tool_handler, tool_router,
transport::stdio,
};
use serde::Deserialize;
use std::fmt::Write as _;
use tokio::io::{AsyncBufReadExt as _, AsyncWriteExt as _, BufReader};
use tokio::net::UnixStream;
use hive_bash_mcp::paths;
use hive_bash_mcp::protocol::{DaemonRequest, DaemonResponse};
use hive_bash_mcp::runner::SUMMARY_BYTES;
/// Send `req` to the daemon and read back the response. Each call is a
/// fresh unix-socket connection — short-lived (single round-trip) so
/// connection pooling is unnecessary.
async fn round_trip(req: DaemonRequest) -> Result<DaemonResponse> {
let socket = paths::daemon_socket();
let stream = UnixStream::connect(&socket)
.await
.with_context(|| format!("connect bash daemon socket {}", socket.display()))?;
let (reader, mut writer) = stream.into_split();
let mut line = serde_json::to_string(&req)?;
line.push('\n');
writer
.write_all(line.as_bytes())
.await
.context("write request to bash daemon socket")?;
writer.shutdown().await.ok();
let mut buf = String::new();
BufReader::new(reader)
.read_line(&mut buf)
.await
.context("read response from bash daemon socket")?;
serde_json::from_str(&buf).context("parse bash daemon response")
}
/// Format a `TaskFile` JSON value as a human-readable status string.
/// Mirrors `format_bash_status` in the old hive-ag3nt, adapted to work
/// from the daemon's JSON payload.
fn format_task(id: &str, task: &serde_json::Value) -> String {
let status = task["status"].as_str().unwrap_or("unknown");
let mut out = format!("task `{id}`: status={status}");
if let Some(code) = task["exit_code"].as_i64() {
let _ = write!(out, ", exit={code}");
}
if let (Some(started), None) = (
task["started_at"].as_i64(),
task["completed_at"].as_i64().map(|_| ()),
) {
// Running — show age.
let now = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_default()
.as_secs()
.cast_signed();
let _ = write!(out, ", running for {}s", now - started);
}
if let (Some(completed), Some(started)) =
(task["completed_at"].as_i64(), task["started_at"].as_i64())
{
let _ = write!(out, ", took {}s", completed - started);
}
let out_file = paths::task_out(id);
let err_file = paths::task_err(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"].as_str() {
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"].as_str() {
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
}
/// Turn a `DaemonResponse` from a `BashRun` call into the string
/// claude sees as the tool result.
fn render_bash_run(id: &str, resp: Result<DaemonResponse>) -> String {
match resp {
Ok(DaemonResponse::Ok { payload }) => {
let finished = payload["finished"].as_bool().unwrap_or(false);
if finished && let Some(task) = payload.get("task") {
return format_task(id, task);
}
format!("task started: id={id}")
}
Ok(DaemonResponse::Error { message }) => format!("bash_run error: {message}"),
Err(e) => format!("bash bridge error: {e:#}"),
}
}
/// 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 \
a wake fires when it completes) rather than immediately re-waiting.";
/// Turn a `DaemonResponse` from a `BashStatus` call into the string
/// claude sees as the tool result. When `waited` is set (the call
/// parked on a `wait_seconds` poll) and the task is still non-terminal,
/// [`BASH_IDLE_WAIT_HINT`] is appended.
fn render_bash_status(id: &str, resp: Result<DaemonResponse>, waited: bool) -> String {
match resp {
Ok(DaemonResponse::Ok { payload }) => {
let mut out = format_task(id, &payload);
if waited && matches!(payload["status"].as_str(), Some("pending" | "running")) {
out.push_str(BASH_IDLE_WAIT_HINT);
}
out
}
Ok(DaemonResponse::Error { message }) => message,
Err(e) => format!("bash bridge error: {e:#}"),
}
}
/// Turn a `DaemonResponse` from a `BashKill` call into the string claude sees.
fn render_bash_kill(resp: Result<DaemonResponse>) -> String {
match resp {
Ok(DaemonResponse::Ok { payload }) => {
let id = payload["id"].as_str().unwrap_or("unknown");
if payload["was_running"].as_bool().unwrap_or(false) {
let sig = payload["signal"].as_str().unwrap_or("SIGINT");
format!(
"task `{id}`: {sig} sent to its process group; it transitions to `killed` \
once the process exits and the usual completion wake fires."
)
} else {
format!("task `{id}` was pending — cancelled before it started.")
}
}
Ok(DaemonResponse::Error { message }) => format!("kill error: {message}"),
Err(e) => format!("bash bridge error: {e:#}"),
}
}
// ---------------------------------------------------------------------------
// 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 and no wake is fired;
/// 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 the completion wake (`from: "bash-task-<name>"`), in `status`
/// lookups, and in the 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: ASCII letters, digits, `.`,
/// `_`, `-` (max 64). 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, the harness fires a wake with \
`from: \"bash-task-<id>\"` and the exit code + last stdout lines in the body; \
handle it on a future turn. Use `status` to poll the task status within \
the same turn if needed. `timeout_secs` defaults to `None` (no timeout) — \
task runs until natural exit; pass an explicit value to kill after N seconds. \
Pass `wait_seconds` (capped at 30) to wait inline for fast commands: when the \
task finishes within the window the full status is returned immediately and no \
wake is fired; when the timeout expires the task keeps running and the normal \
`task started: id=<id>` response is returned. `wait_seconds` defaults to 3; \
pass `wait_seconds: 0` to disable inline waiting and always get the immediate \
response. Pass `name` to label the task with a memorable id (used in the wake \
`from`, `status` lookups, and the loose-ends list) instead of an opaque hex id; \
a name is reusable once its prior task has finished, and rejected while one is \
still running."
)]
async fn run(&self, Parameters(args): Parameters<BashRunArgs>) -> String {
let req = DaemonRequest::BashRun {
cmd: args.cmd,
timeout_secs: args.timeout_secs,
wait_seconds: args.wait_seconds,
name: args.name,
};
let resp = round_trip(req).await;
// Extract the id from the response to format the result.
match &resp {
Ok(DaemonResponse::Ok { payload }) => {
let id = payload["id"].as_str().unwrap_or("unknown").to_owned();
render_bash_run(&id, resp)
}
_ => render_bash_run("unknown", resp),
}
}
#[tool(
description = "Check the status of a background bash task by its ID (from `run`). \
Returns the current status (pending/running/done/timed_out/interrupted/killed), exit \
code if finished, and a tail of stdout/stderr. Full output lives in \
`harness/bash-tasks/<id>.out` / `.err`. \
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 id = args.id.clone();
let waited = args.wait_seconds.is_some_and(|w| w > 0);
let req = DaemonRequest::BashStatus {
id: args.id,
wait_seconds: args.wait_seconds,
};
render_bash_status(&id, round_trip(req).await, waited)
}
#[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 fires the usual \
completion wake."
)]
async fn kill(&self, Parameters(args): Parameters<BashKillArgs>) -> String {
let req = DaemonRequest::BashKill {
id: args.id,
force: args.force,
};
render_bash_kill(round_trip(req).await)
}
}
#[tool_handler]
impl ServerHandler for BashMcp {}
#[tokio::main]
async fn main() -> Result<()> {
tracing_subscriber::fmt()
.with_env_filter(
tracing_subscriber::EnvFilter::try_from_env("RUST_LOG")
.unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("warn")),
)
.with_writer(std::io::stderr)
.init();
let service = BashMcp.serve(stdio()).await?;
service.waiting().await?;
Ok(())
}
#[cfg(test)]
mod status_hint_tests {
use super::{BASH_IDLE_WAIT_HINT, render_bash_status};
use hive_bash_mcp::protocol::DaemonResponse;
fn status_resp(status: &str) -> DaemonResponse {
DaemonResponse::Ok {
payload: serde_json::json!({ "status": status, "started_at": 1 }),
}
}
#[test]
fn running_task_after_wait_appends_idle_hint() {
let out = render_bash_status("t1", Ok(status_resp("running")), true);
assert!(out.contains(BASH_IDLE_WAIT_HINT));
}
#[test]
fn running_task_without_wait_has_no_hint() {
let out = render_bash_status("t1", Ok(status_resp("running")), false);
assert!(!out.contains(BASH_IDLE_WAIT_HINT));
}
#[test]
fn finished_task_after_wait_has_no_hint() {
let out = render_bash_status("t1", Ok(status_resp("done")), true);
assert!(!out.contains(BASH_IDLE_WAIT_HINT));
}
}