hyperhive/hive-sock-client/src/lib.rs

//! JSON-line-over-unix-socket client, shared by every daemon that talks to
//! a hyperhive socket.
//!
//! The wire protocol is identical everywhere — connect, write one line of
//! JSON, read one line of JSON back — so this crate is generic over the
//! request and response types and knows nothing about either protocol. The
//! host-served control socket and the harness's in-agent socket both use
//! it with their own wire-type crates.
//!
//! Two axes of behaviour are real and stay caller-selectable; everything
//! else is shared:
//!
//! - **retry**: [`Retry::RideOutRestart`] for callers with no natural
//!   retry of their own, [`Retry::None`] for callers already inside a poll
//!   loop whose interval *is* the retry.
//! - **response**: [`request`] decodes it, [`notify`] drains and discards
//!   it.
//!
//! Whether a failure propagates or is logged and swallowed is the caller's
//! choice and stays at the call site — it is not a property of the
//! transport.

use std::path::Path;
use std::time::Duration;

use anyhow::{Context, Result, anyhow};
use serde::Serialize;
use serde::de::DeserializeOwned;
use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
use tokio::net::UnixStream;

/// Backoff schedule for [`Retry::RideOutRestart`]. Five entries → up to 5
/// retries on top of the initial attempt; total wall-clock cap =
/// 2+4+8+16+30 = 60s. Sized to ride out a service restart (systemd usually
/// has the unix socket back inside ~5s) without the caller having to
/// handle the transient itself.
const RIDE_OUT_RESTART_BACKOFFS_MS: &[u64] = &[2_000, 4_000, 8_000, 16_000, 30_000];

/// What to do when a connect or I/O attempt fails.
///
/// Deliberately two named policies rather than a configurable schedule:
/// only two behaviours exist in the tree, and naming them keeps the
/// *reason* for each choice readable at the call site.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum Retry {
    /// Fail on the first failure. For callers inside a poll loop, where
    /// the poll interval already is the retry — a second, in-request
    /// backoff would stack sleeps on top of it and delay the rest of the
    /// batch.
    None,
    /// Back off on [`RIDE_OUT_RESTART_BACKOFFS_MS`] (~60s total). For
    /// callers with no natural retry of their own, where a surfaced
    /// transient costs more than the wait.
    RideOutRestart,
}

impl Retry {
    /// Sleep schedule between attempts; its length is the retry budget.
    fn backoffs(self) -> &'static [u64] {
        match self {
            Self::None => &[],
            Self::RideOutRestart => RIDE_OUT_RESTART_BACKOFFS_MS,
        }
    }
}

/// Send `req` over `socket` and decode the single-line JSON response.
///
/// # Errors
///
/// Returns an error if `req` cannot be serialised, if the socket is
/// unreachable after the retry budget is spent, if the server closes
/// without responding, or if the response does not deserialise into
/// `Resp`.
pub async fn request<Req, Resp>(socket: &Path, req: &Req, retry: Retry) -> Result<Resp>
where
    Req: Serialize + ?Sized,
    Resp: DeserializeOwned,
{
    request_retried(socket, req, retry)
        .await
        .map(|(resp, _)| resp)
}

/// Same as [`request`], but also reports how many retries it took past the
/// initial attempt (0 = succeeded first try).
///
/// MCP tool handlers use this to append a one-line hint to the tool result
/// when retries happened, so claude reads the earlier socket flake as a
/// transient rather than a content error worth an LLM-level retry.
///
/// # Errors
///
/// Same as [`request`].
pub async fn request_retried<Req, Resp>(
    socket: &Path,
    req: &Req,
    retry: Retry,
) -> Result<(Resp, u32)>
where
    Req: Serialize + ?Sized,
    Resp: DeserializeOwned,
{
    let payload = encode(req)?;
    let (line, retries) = with_retry(socket, &payload, Mode::Decode, retry).await?;
    let resp = serde_json::from_str(line.trim())
        .with_context(|| format!("decode response from {}", socket.display()))?;
    Ok((resp, retries))
}

/// Send `req` over `socket`, half-close, and drain the response line
/// without decoding it.
///
/// For fire-and-forget ops whose reply carries nothing the caller acts on.
/// The drain is not optional: without it the server's write-back lands on
/// a closed socket.
///
/// # Errors
///
/// Returns an error if `req` cannot be serialised, or if the socket is
/// unreachable / the write fails after the retry budget is spent. A
/// failure to read the discarded response is not an error — the request
/// was already delivered.
pub async fn notify<Req>(socket: &Path, req: &Req, retry: Retry) -> Result<()>
where
    Req: Serialize + ?Sized,
{
    let payload = encode(req)?;
    with_retry(socket, &payload, Mode::Drain, retry).await?;
    Ok(())
}

/// What to do with the server's response line.
#[derive(Clone, Copy)]
enum Mode {
    /// Flush the write half, read the response, hand it back to be parsed.
    Decode,
    /// Half-close the write half, read-and-discard the response.
    Drain,
}

/// Serialise `req` into one newline-terminated JSON line.
///
/// Kept outside the retry loop on purpose: a serialisation failure is
/// deterministic, so retrying it would only reproduce it.
fn encode<Req>(req: &Req) -> Result<Vec<u8>>
where
    Req: Serialize + ?Sized,
{
    let mut payload = serde_json::to_string(req).context("serialise request")?;
    payload.push('\n');
    Ok(payload.into_bytes())
}

/// Run [`try_once`] until it succeeds or the retry budget is spent,
/// returning the response line and the number of retries it took.
async fn with_retry(
    socket: &Path,
    payload: &[u8],
    mode: Mode,
    retry: Retry,
) -> Result<(String, u32)> {
    let backoffs = retry.backoffs();
    let mut attempt: usize = 0;
    loop {
        match try_once(socket, payload, mode).await {
            Ok(line) => return Ok((line, u32::try_from(attempt).unwrap_or(u32::MAX))),
            Err(e) => {
                let Some(&sleep_ms) = backoffs.get(attempt) else {
                    return Err(exhausted(e, backoffs));
                };
                tracing::warn!(
                    attempt = attempt + 1,
                    sleep_ms,
                    socket = %socket.display(),
                    error = %e,
                    "hive socket attempt failed; retrying"
                );
                tokio::time::sleep(Duration::from_millis(sleep_ms)).await;
                attempt += 1;
            }
        }
    }
}

/// Annotate the final failure with how long the retry schedule tried, so a
/// surfaced error says whether it was one shot or a full ride-out.
fn exhausted(err: anyhow::Error, backoffs: &[u64]) -> anyhow::Error {
    if backoffs.is_empty() {
        return err;
    }
    let total_s = backoffs.iter().sum::<u64>() / 1_000;
    err.context(format!(
        "gave up after {} retries over ~{total_s}s",
        backoffs.len()
    ))
}

/// One connect / write / read cycle. Every error it returns is retryable
/// by construction — the deterministic failures (serialise, deserialise)
/// happen outside the retry loop.
async fn try_once(socket: &Path, payload: &[u8], mode: Mode) -> Result<String> {
    let stream = UnixStream::connect(socket).await.map_err(|e| {
        // A refused or missing socket usually means the listener is
        // mid-restart (operator redeploy, harness restart) — it is
        // recreated on boot and a retrying caller rides it out. When the
        // error does surface (budget spent, or a fail-fast caller) the
        // hint marks it as a likely transient rather than a hard failure
        // worth escalating. The path stays in the error either way: a
        // permission problem that reads as "is the daemon running?" sends
        // the operator to fix the wrong thing.
        let restarting = matches!(
            e.kind(),
            std::io::ErrorKind::ConnectionRefused | std::io::ErrorKind::NotFound
        );
        let err = anyhow::Error::new(e).context(format!("connect to {}", socket.display()));
        if restarting {
            err.context("the listener may be restarting (e.g. an operator redeploy)")
        } else {
            err
        }
    })?;
    let (read, mut write) = stream.into_split();

    write
        .write_all(payload)
        .await
        .with_context(|| format!("write to {}", socket.display()))?;
    match mode {
        Mode::Decode => write
            .flush()
            .await
            .with_context(|| format!("flush {}", socket.display()))?,
        Mode::Drain => write
            .shutdown()
            .await
            .with_context(|| format!("shutdown write to {}", socket.display()))?,
    }

    let mut reader = BufReader::new(read);
    let mut line = String::new();
    match mode {
        Mode::Decode => {
            let read_bytes = reader
                .read_line(&mut line)
                .await
                .with_context(|| format!("read from {}", socket.display()))?;
            if read_bytes == 0 || line.is_empty() {
                return Err(anyhow!(
                    "{} closed the connection without responding",
                    socket.display()
                ));
            }
        }
        Mode::Drain => {
            // Discarded, including any error reading it: the request is
            // already delivered and the caller acts on nothing here.
            let _ = reader.read_line(&mut line).await;
            line.clear();
        }
    }
    Ok(line)
}

#[cfg(test)]
mod tests {
    use super::{Retry, notify, request};

    /// A connect to a non-existent socket path (ENOENT → `NotFound`) is
    /// annotated with both the socket path and the "may be restarting"
    /// hint, so a surfaced tool error reads as the expected transient.
    #[tokio::test]
    async fn missing_socket_names_the_path_and_hints_restart() {
        let bogus = std::path::Path::new("/nonexistent/hive/mcp.sock");
        let err = request::<(), serde_json::Value>(bogus, &(), Retry::None)
            .await
            .expect_err("connect to a non-existent socket must fail");
        let msg = format!("{err:#}");
        assert!(msg.contains("restarting"), "missing restart hint: {msg}");
        assert!(msg.contains("connect to"), "missing connect context: {msg}");
    }

    /// `Retry::None` returns promptly rather than sleeping the ride-out
    /// schedule — its callers are poll ticks that need the rest of the
    /// batch to still run this cycle.
    #[tokio::test]
    async fn no_retry_fails_fast() {
        let bogus = std::path::Path::new("/nonexistent/hive/agent.sock");
        let started = std::time::Instant::now();
        notify(bogus, &(), Retry::None)
            .await
            .expect_err("connect to a non-existent socket must fail");
        assert!(
            started.elapsed() < std::time::Duration::from_secs(1),
            "Retry::None slept: {:?}",
            started.elapsed()
        );
    }

    /// The ride-out schedule is a full minute of patience — the value the
    /// harness's tool callers depend on to not surface a restart.
    #[test]
    fn ride_out_restart_budget_is_a_minute() {
        let total_ms: u64 = Retry::RideOutRestart.backoffs().iter().sum();
        assert_eq!(total_ms, 60_000);
        assert!(Retry::None.backoffs().is_empty());
    }
}