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

//! Unix socket server: the daemon listens here, the stdio MCP bridge
//! `connect()`s on every tool call. One JSON request line in, one
//! JSON response line out. Connections are short-lived (per tool call)
//! so the loop is just accept → dispatch → reply → close.

use std::path::Path;

use anyhow::{Context, Result};
use tokio::io::{AsyncBufReadExt as _, AsyncWriteExt as _, BufReader};
use tokio::net::{UnixListener, UnixStream};

use crate::protocol::{DaemonRequest, DaemonResponse, TaskStatus};
use crate::runner;

/// Start listening on `socket_path` and serve forever. Removes any
/// stale socket file first so a daemon restart doesn't hit EADDRINUSE.
pub async fn serve(socket_path: &Path) -> Result<()> {
    let _ = tokio::fs::remove_file(socket_path).await;
    if let Some(parent) = socket_path.parent() {
        tokio::fs::create_dir_all(parent)
            .await
            .with_context(|| format!("mkdir {}", parent.display()))?;
    }
    let listener = UnixListener::bind(socket_path)
        .with_context(|| format!("bind unix socket {}", socket_path.display()))?;
    tracing::info!(path = %socket_path.display(), "bash daemon socket up");
    loop {
        let (stream, _) = listener
            .accept()
            .await
            .context("accept on bash daemon socket")?;
        tokio::spawn(async move {
            if let Err(e) = handle_connection(stream).await {
                tracing::warn!(error = %e, "bash socket connection error");
            }
        });
    }
}

async fn handle_connection(stream: UnixStream) -> Result<()> {
    let (reader, mut writer) = stream.into_split();
    let mut lines = BufReader::new(reader).lines();
    while let Some(line) = lines.next_line().await? {
        let response = match serde_json::from_str::<DaemonRequest>(&line) {
            Ok(req) => dispatch(req).await,
            Err(e) => DaemonResponse::error(format!("parse request: {e}")),
        };
        let mut json = serde_json::to_string(&response)?;
        json.push('\n');
        writer.write_all(json.as_bytes()).await?;
        writer.flush().await?;
    }
    Ok(())
}

async fn dispatch(req: DaemonRequest) -> DaemonResponse {
    match req {
        DaemonRequest::Ping => DaemonResponse::ok(&serde_json::json!({"ok": true})),

        DaemonRequest::BashRun {
            cmd,
            timeout_secs,
            wait_seconds,
            name,
        } => {
            let id = match runner::submit_task(cmd, timeout_secs, name) {
                Ok(id) => id,
                Err(e) => return DaemonResponse::error(format!("submit_task: {e:#}")),
            };
            // Inline wait: if requested and the task finishes quickly,
            // return the full status instead of just the task ID.
            let wait = 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 DaemonResponse::ok(&serde_json::json!({
                    "id": id,
                    "finished": true,
                    "task": task,
                }));
            }
            DaemonResponse::ok(&serde_json::json!({ "id": id, "finished": false }))
        }

        DaemonRequest::BashStatus { id, wait_seconds } => {
            let wait = wait_seconds.unwrap_or(0);
            let task = if wait > 0 {
                runner::wait_for_task(&id, wait).await
            } else {
                runner::read_task(&id)
            };
            match task {
                Some(t) => DaemonResponse::ok(&t),
                None => DaemonResponse::error(format!("unknown task id `{id}`")),
            }
        }

        DaemonRequest::ActiveTasks => {
            let tasks = runner::active_tasks();
            DaemonResponse::ok(&tasks)
        }

        DaemonRequest::BashKill { id, force } => {
            let (killed, was_running) = runner::kill_task(&id, force);
            if !killed {
                return DaemonResponse::error(format!(
                    "no running or pending task with id `{id}` (already finished or unknown)"
                ));
            }
            let payload = if was_running {
                serde_json::json!({
                    "id": id,
                    "killed": true,
                    "was_running": true,
                    "signal": if force { "SIGKILL" } else { "SIGINT" },
                })
            } else {
                serde_json::json!({
                    "id": id,
                    "killed": true,
                    "was_running": false,
                    "note": "task was pending — cancelled before it started",
                })
            };
            DaemonResponse::ok(&payload)
        }
    }
}