hyperhive/hive-runtime/src/acp/rpc.rs

//! JSON-RPC 2.0 over an ACP agent's stdio: newline-delimited messages, the
//! agent's requests answered here, its notifications queued for the turn.

use std::collections::HashMap;
use std::path::Path;
use std::process::Stdio;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, Mutex, PoisonError};

use serde_json::{Value, json};
use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
use tokio::process::{Child, ChildStdin, Command};
use tokio::sync::{mpsc, oneshot};

use super::provider_key::KeyVar;
use super::stream::mcp_server_of;
use super::{AcpError, PermissionAsk, PermissionPolicy};
use crate::spec::AcpCommand;

/// Something the agent sent that is not a response to one of our requests.
pub(super) enum Incoming {
    /// The `params` of a `session/update` notification.
    Update(Value),
    Stdout(String),
    Stderr(String),
    /// The agent closed its stdout: it has exited or is about to.
    Closed,
}

type Reply = Result<Value, AcpError>;
type Pending = Arc<Mutex<HashMap<u64, (&'static str, oneshot::Sender<Reply>)>>>;

/// The agent runs as the leader of its own process group; dropping the
/// connection sends that group SIGINT, so the agent can stop the commands
/// it started, and they get the signal too. Nothing escalates: an agent
/// that ignores SIGINT keeps running. A descendant that left the group
/// (`setsid`, `setpgid`) is not reached.
pub(super) struct Connection {
    child: Child,
    /// The agent's pid, which is its group's id. Read at spawn: a reaped
    /// child has no id, and the group can outlive its leader.
    pgid: Option<i32>,
    stdin: Arc<tokio::sync::Mutex<ChildStdin>>,
    next_id: AtomicU64,
    pending: Pending,
    pub(super) incoming: mpsc::UnboundedReceiver<Incoming>,
}

impl Connection {
    /// Spawn the agent in `cwd`, with `key` added to its environment, and
    /// start reading its output.
    pub(super) fn spawn(
        command: &AcpCommand,
        key: Option<&KeyVar>,
        cwd: &Path,
        permit: PermissionPolicy,
        servers: Vec<String>,
    ) -> Result<Self, AcpError> {
        let mut cmd = Command::new(&command.command);
        cmd.args(&command.args).envs(&command.env);
        if let Some(key) = key {
            cmd.env(&key.name, &key.value);
        }
        let mut child = cmd
            .current_dir(cwd)
            .stdin(Stdio::piped())
            .stdout(Stdio::piped())
            .stderr(Stdio::piped())
            .process_group(0)
            .spawn()
            .map_err(|source| AcpError::Spawn {
                program: command.command.clone(),
                source,
            })?;
        let pgid = group_id(child.id());
        let (Some(stdin), Some(stdout), Some(stderr)) =
            (child.stdin.take(), child.stdout.take(), child.stderr.take())
        else {
            unreachable!("all three stdio handles were requested as pipes");
        };
        let stdin = Arc::new(tokio::sync::Mutex::new(stdin));
        let pending: Pending = Arc::default();
        let (tx, incoming) = mpsc::unbounded_channel();

        let stderr_tx = tx.clone();
        tokio::spawn(async move {
            let mut lines = BufReader::new(stderr).lines();
            while let Ok(Some(line)) = lines.next_line().await {
                if stderr_tx.send(Incoming::Stderr(line)).is_err() {
                    break;
                }
            }
        });

        let reader = Reader {
            stdin: stdin.clone(),
            pending: pending.clone(),
            tx,
            permit,
            servers,
        };
        tokio::spawn(async move {
            let mut lines = BufReader::new(stdout).lines();
            while let Ok(Some(line)) = lines.next_line().await {
                reader.dispatch(&line).await;
            }
            // Fail every request still waiting: no reply can come now.
            reader
                .pending
                .lock()
                .unwrap_or_else(PoisonError::into_inner)
                .clear();
            let _ = reader.tx.send(Incoming::Closed);
        });

        Ok(Self {
            child,
            pgid,
            stdin,
            next_id: AtomicU64::new(1),
            pending,
            incoming,
        })
    }

    /// Whether the agent process has exited.
    pub(super) fn exited(&mut self) -> bool {
        !matches!(self.child.try_wait(), Ok(None))
    }

    /// Send a request and return the receiver its reply arrives on.
    pub(super) async fn send(
        &self,
        method: &'static str,
        params: Value,
    ) -> Result<oneshot::Receiver<Reply>, AcpError> {
        let id = self.next_id.fetch_add(1, Ordering::Relaxed);
        let (tx, rx) = oneshot::channel();
        self.pending
            .lock()
            .unwrap_or_else(PoisonError::into_inner)
            .insert(id, (method, tx));
        let message = json!({ "jsonrpc": "2.0", "id": id, "method": method, "params": params });
        if let Err(e) = write(&self.stdin, &message).await {
            self.pending
                .lock()
                .unwrap_or_else(PoisonError::into_inner)
                .remove(&id);
            return Err(e);
        }
        Ok(rx)
    }

    /// Send a notification: a message the agent does not reply to.
    pub(super) async fn notify(&self, method: &'static str, params: Value) -> Result<(), AcpError> {
        let message = json!({ "jsonrpc": "2.0", "method": method, "params": params });
        write(&self.stdin, &message).await
    }

    /// Send a request and wait for its reply, queueing anything else the agent
    /// sends meanwhile on [`Self::incoming`].
    pub(super) async fn request(
        &self,
        method: &'static str,
        params: Value,
    ) -> Result<Value, AcpError> {
        let rx = self.send(method, params).await?;
        rx.await.unwrap_or(Err(AcpError::Closed))
    }
}

/// The process group a leader with `pid` heads, as `kill(2)` takes it.
/// Never 0: `kill(-0, …)` signals the caller's own group.
fn group_id(pid: Option<u32>) -> Option<i32> {
    pid.and_then(|pid| i32::try_from(pid).ok())
        .filter(|&pgid| pgid > 0)
}

impl Drop for Connection {
    fn drop(&mut self) {
        if let Some(pgid) = self.pgid {
            // SAFETY: plain `kill(2)`; `-pgid` targets the agent's own
            // process group, which `process_group(0)` made it lead.
            unsafe {
                libc::kill(-pgid, libc::SIGINT);
            }
        }
    }
}

/// The stdout side: routes responses to their waiting request, answers the
/// agent's own requests, and queues its notifications.
struct Reader {
    stdin: Arc<tokio::sync::Mutex<ChildStdin>>,
    pending: Pending,
    tx: mpsc::UnboundedSender<Incoming>,
    permit: PermissionPolicy,
    servers: Vec<String>,
}

impl Reader {
    async fn dispatch(&self, line: &str) {
        let Ok(message) = serde_json::from_str::<Value>(line) else {
            if !line.trim().is_empty() {
                let _ = self.tx.send(Incoming::Stdout(line.to_owned()));
            }
            return;
        };
        let method = message.get("method").and_then(Value::as_str);
        let id = message.get("id");
        match (method, id) {
            (None, Some(id)) => self.resolve(id, &message),
            (Some(method), Some(id)) => {
                let reply = match method {
                    "session/request_permission" => json!({
                        "jsonrpc": "2.0", "id": id,
                        "result": permission_outcome(&message["params"], &self.permit, &self.servers),
                    }),
                    _ => json!({
                        "jsonrpc": "2.0", "id": id,
                        "error": { "code": -32601, "message": format!("method not found: {method}") },
                    }),
                };
                if let Err(e) = write(&self.stdin, &reply).await {
                    tracing::warn!(error = %e, method, "failed to answer ACP agent request");
                }
            }
            (Some("session/update"), None) => {
                let params = message.get("params").cloned().unwrap_or(Value::Null);
                let _ = self.tx.send(Incoming::Update(params));
            }
            _ => {}
        }
    }

    fn resolve(&self, id: &Value, message: &Value) {
        let Some(id) = id.as_u64() else { return };
        let Some((method, tx)) = self
            .pending
            .lock()
            .unwrap_or_else(PoisonError::into_inner)
            .remove(&id)
        else {
            return;
        };
        let reply = match message.get("error") {
            Some(error) => Err(AcpError::Rpc {
                method,
                code: error.get("code").and_then(Value::as_i64).unwrap_or(0),
                message: error
                    .get("message")
                    .and_then(Value::as_str)
                    .unwrap_or("")
                    .to_owned(),
                data: error.get("data").map(Value::to_string).unwrap_or_default(),
            }),
            None => Ok(message.get("result").cloned().unwrap_or(Value::Null)),
        };
        let _ = tx.send(reply);
    }
}

async fn write(stdin: &tokio::sync::Mutex<ChildStdin>, message: &Value) -> Result<(), AcpError> {
    let mut line = message.to_string();
    line.push('\n');
    let mut stdin = stdin.lock().await;
    stdin.write_all(line.as_bytes()).await?;
    stdin.flush().await?;
    Ok(())
}

/// The `session/request_permission` result: the agent's first option of the
/// matching kind, allow if `permit` accepts the request and reject otherwise.
/// With no such option the request is answered `cancelled`. `servers` are the
/// MCP servers handed to the session, matched against the tool call's title.
pub(super) fn permission_outcome(
    params: &Value,
    permit: &PermissionPolicy,
    servers: &[String],
) -> Value {
    let tool_call = params.get("toolCall");
    let field = |k: &str| tool_call.and_then(|t| t.get(k)).and_then(Value::as_str);
    let kind = field("kind").unwrap_or("other");
    let title = field("title");
    let ask = PermissionAsk {
        kind,
        // MCP tool calls are kind `other`; a title alone must not lift another
        // kind (`execute`, `fetch`) into one.
        mcp_server: (kind == "other")
            .then(|| title.and_then(|title| mcp_server_of(title, servers)))
            .flatten(),
    };
    let permitted = permit(&ask);
    let wanted = if permitted {
        ["allow_once", "allow_always"]
    } else {
        ["reject_once", "reject_always"]
    };
    let options = params
        .get("options")
        .and_then(Value::as_array)
        .map_or(&[][..], Vec::as_slice);
    let chosen = wanted.iter().find_map(|want| {
        options
            .iter()
            .find(|o| o.get("kind").and_then(Value::as_str) == Some(want))
            .and_then(|o| o.get("optionId"))
    });
    // Log what the agent was actually told, not just the policy's
    // pre-resolution answer: `permit` saying yes doesn't mean the agent
    // offered a matching `allow_*` option, and that silent-cancel case is
    // exactly the shape a reporter sees as "tool call refused, nothing
    // actionable in the logs" — the earlier version of this log line
    // reported `permitted` and would have called that case "allowed".
    // `title` is only logged for `kind == "other"` (it identifies the MCP
    // tool/server) — an `execute` title can carry a command line and must
    // not land in the journal at this level.
    let other_title = (kind == "other").then_some(title).flatten();
    match chosen {
        Some(option) if permitted => {
            tracing::debug!(kind, mcp_server = ?ask.mcp_server, option = %option, "ACP permission request allowed");
        }
        Some(option) => {
            tracing::info!(kind, mcp_server = ?ask.mcp_server, title = other_title, option = %option, "ACP permission request denied by policy");
        }
        None => {
            // An anomaly, not a policy decision: the agent offered no
            // option matching what the policy wanted, so the request is
            // cancelled regardless of `permitted`.
            tracing::warn!(kind, mcp_server = ?ask.mcp_server, title = other_title, permitted, "ACP permission request cancelled — agent offered no matching option");
        }
    }
    match chosen {
        Some(option) => json!({ "outcome": { "outcome": "selected", "optionId": option } }),
        None => json!({ "outcome": { "outcome": "cancelled" } }),
    }
}

#[cfg(test)]
mod tests {
    use super::{group_id, permission_outcome};
    use crate::acp::{PermissionAsk, PermissionPolicy};
    use serde_json::json;
    use std::sync::Arc;

    fn request(kind: &str, title: &str) -> serde_json::Value {
        json!({
            "sessionId": "s",
            "toolCall": { "toolCallId": "t", "kind": kind, "title": title },
            "options": [
                { "optionId": "once", "kind": "allow_once", "name": "Allow once" },
                { "optionId": "always", "kind": "allow_always", "name": "Always allow" },
                { "optionId": "reject", "kind": "reject_once", "name": "Reject" },
            ],
        })
    }

    fn servers() -> Vec<String> {
        vec!["hyperhive".to_owned()]
    }

    /// Allows `edit` and MCP tools, nothing else.
    fn edit_or_mcp() -> PermissionPolicy {
        Arc::new(|ask: &PermissionAsk<'_>| ask.kind == "edit" || ask.mcp_server.is_some())
    }

    #[test]
    fn a_permitted_request_is_allowed_once() {
        assert_eq!(
            permission_outcome(&request("edit", "edit"), &edit_or_mcp(), &servers()),
            json!({ "outcome": { "outcome": "selected", "optionId": "once" } })
        );
    }

    #[test]
    fn a_refused_request_is_rejected() {
        assert_eq!(
            permission_outcome(&request("execute", "bash"), &edit_or_mcp(), &servers()),
            json!({ "outcome": { "outcome": "selected", "optionId": "reject" } })
        );
    }

    #[test]
    fn the_policy_sees_which_mcp_server_a_tool_belongs_to() {
        let seen = Arc::new(std::sync::Mutex::new(Vec::new()));
        let record = seen.clone();
        let policy: PermissionPolicy = Arc::new(move |ask: &PermissionAsk<'_>| {
            record
                .lock()
                .unwrap()
                .push((ask.kind.to_owned(), ask.mcp_server.map(str::to_owned)));
            true
        });
        permission_outcome(&request("other", "hyperhive_send"), &policy, &servers());
        permission_outcome(&request("other", "skill"), &policy, &servers());
        let mut req = request("x", "read");
        req["toolCall"].as_object_mut().unwrap().remove("kind");
        permission_outcome(&req, &policy, &servers());
        assert_eq!(
            *seen.lock().unwrap(),
            vec![
                ("other".to_owned(), Some("hyperhive".to_owned())),
                ("other".to_owned(), None),
                ("other".to_owned(), None),
            ]
        );
    }

    /// A title shaped like `<server>_<tool>` marks an MCP call only on kind
    /// `other`, which is what an agent gives an MCP tool call.
    #[test]
    fn an_mcp_looking_title_on_execute_or_fetch_is_not_an_mcp_call() {
        let policy: PermissionPolicy = Arc::new(|ask: &PermissionAsk<'_>| ask.mcp_server.is_some());
        for kind in ["execute", "fetch"] {
            assert_eq!(
                permission_outcome(&request(kind, "hyperhive_send"), &policy, &servers())["outcome"]
                    ["optionId"],
                "reject",
                "{kind}"
            );
        }
        assert_eq!(
            permission_outcome(&request("other", "hyperhive_send"), &policy, &servers())["outcome"]
                ["optionId"],
            "once"
        );
    }

    #[test]
    fn no_matching_option_cancels() {
        let req = json!({ "toolCall": { "kind": "execute" },
                          "options": [{ "optionId": "once", "kind": "allow_once" }] });
        assert_eq!(
            permission_outcome(&req, &edit_or_mcp(), &servers()),
            json!({ "outcome": { "outcome": "cancelled" } })
        );
    }

    #[test]
    fn a_group_id_is_a_positive_pid() {
        assert_eq!(group_id(Some(42)), Some(42));
        assert_eq!(group_id(Some(0)), None);
        assert_eq!(group_id(Some(u32::MAX)), None);
        assert_eq!(group_id(None), None);
    }
}