pushkin_daemon/

server.rs

//! The daemon server: accepts UDS connections under `.pushkin/`, answers
//! protocol requests with the shared pipeline (spec §4.3). One tokio
//! runtime per `serve` call; connections are tracked in a `JoinSet` so no
//! task is fire-and-forget (AGENTS.md tokio rule). The daemon is a
//! transport for the cold pipeline, never a second brain: `Check` calls
//! `pushkin_core::pipeline::check_write` — the same function, the same
//! envelope. Read-only daemons (spec §8.4: the offer made to non-canonical
//! binaries) serve on private sockets and refuse wire mutations.

use crate::protocol::{DaemonInfo, Request, Response, PROTOCOL_VERSION};
use crate::warm::WarmState;
use pushkin_core::manifest::Manifest;
use pushkin_core::pipeline::WriteRequest;
use std::path::{Path, PathBuf};
use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
use tokio::net::{UnixListener, UnixStream};
use tokio::task::JoinSet;

/// Socket location relative to the repo root — inside `.pushkin/` so the
/// gate surface owns its own transport.
pub const SOCKET_FILE: &str = ".pushkin/daemon.sock";

/// One request line must fit comfortably; a whole-file write payload can
/// be large, so the cap is generous but bounded (no unbounded reads from
/// an untrusted local peer).
const MAX_LINE_BYTES: usize = 16 * 1024 * 1024;

#[derive(Debug, thiserror::Error)]
pub enum ServerError {
    #[error("daemon io failure: {0}")]
    Io(#[from] std::io::Error),
    #[error("daemon not running")]
    NotRunning,
    #[error("protocol failure: {0}")]
    Protocol(String),
}

#[must_use]
pub fn socket_path(repo_root: &Path) -> PathBuf {
    repo_root.join(SOCKET_FILE)
}

/// Serve requests on the canonical socket until a `Shutdown` request
/// arrives. Blocking: owns a current-thread tokio runtime for its
/// lifetime. The socket file is created on bind and removed before
/// returning.
///
/// # Errors
/// `Io` when the socket cannot be created/bound; `Protocol` when the
/// runtime cannot be built.
pub fn serve(repo_root: &Path, manifest: Manifest) -> Result<(), ServerError> {
    serve_at(&socket_path(repo_root), manifest, false)
}

/// Serve on an explicit socket path. `read_only` daemons answer checks
/// and pings (flagged in `DaemonInfo`) but refuse `Shutdown` over the
/// wire — their lifecycle belongs to the OS session that spawned them
/// (spec §8.4), so they run until killed.
///
/// # Errors
/// `Io` when the socket cannot be created/bound; `Protocol` when the
/// runtime cannot be built.
pub fn serve_at(socket: &Path, manifest: Manifest, read_only: bool) -> Result<(), ServerError> {
    let warm = WarmState::new(manifest);
    // Watch the repo root (the socket's grandparent via .pushkin/, or cwd
    // for private sockets elsewhere): manifest edits reload, file edits
    // invalidate. A watcher failure degrades to compute-every-time — the
    // gate stays correct, only slower.
    let watch_root = socket
        .parent()
        .and_then(Path::parent)
        .filter(|root| !root.as_os_str().is_empty())
        .unwrap_or_else(|| Path::new("."));
    let _watch_guard = warm.watch(watch_root).ok();
    if let Some(parent) = socket.parent() {
        std::fs::create_dir_all(parent)?;
    }
    // A previous unclean exit leaves a stale socket; bind() would fail.
    if socket.exists() {
        std::fs::remove_file(socket)?;
    }

    let runtime = tokio::runtime::Builder::new_current_thread()
        .enable_io()
        .enable_time()
        .build()
        .map_err(|e| ServerError::Protocol(format!("runtime build failed: {e}")))?;

    let result = runtime.block_on(serve_inner(socket, &warm, read_only));
    // Socket removal is part of shutdown's contract regardless of outcome.
    let _ = std::fs::remove_file(socket);
    result
}

async fn serve_inner(socket: &Path, warm: &WarmState, read_only: bool) -> Result<(), ServerError> {
    let listener = UnixListener::bind(socket)?;
    let (shutdown_tx, mut shutdown_rx) = tokio::sync::mpsc::channel::<()>(1);
    let mut connections: JoinSet<()> = JoinSet::new();

    loop {
        tokio::select! {
            accepted = listener.accept() => {
                let Ok((stream, _addr)) = accepted else { continue };
                let warm = warm.share();
                let shutdown_tx = shutdown_tx.clone();
                connections.spawn(async move {
                    // Per-connection failures are that connection's problem,
                    // never the accept loop's.
                    let _ = handle_connection(stream, &warm, read_only, &shutdown_tx).await;
                });
            }
            _ = shutdown_rx.recv() => break,
            // Reap finished connection tasks so the set doesn't grow.
            Some(_) = connections.join_next(), if !connections.is_empty() => {}
        }
    }
    // Drain in-flight connections before tearing the socket down.
    while connections.join_next().await.is_some() {}
    Ok(())
}

async fn handle_connection(
    stream: UnixStream,
    warm: &WarmState,
    read_only: bool,
    shutdown_tx: &tokio::sync::mpsc::Sender<()>,
) -> Result<(), ServerError> {
    let (read_half, mut write_half) = stream.into_split();
    let mut lines = BufReader::with_capacity(64 * 1024, read_half).lines();

    while let Ok(Some(line)) = lines.next_line().await {
        if line.len() > MAX_LINE_BYTES {
            let response = Response::Error {
                message: "request exceeds size cap".to_owned(),
            };
            write_response(&mut write_half, &response).await?;
            continue;
        }
        let response = match serde_json::from_str::<Request>(&line) {
            Ok(request) => {
                let response = respond(&request, warm, read_only);
                let stop_serving = !read_only && matches!(request, Request::Shutdown { .. });
                write_response(&mut write_half, &response).await?;
                if stop_serving {
                    let _ = shutdown_tx.send(()).await;
                    return Ok(());
                }
                continue;
            }
            Err(error) => Response::Error {
                message: format!("unrecognized request: {error}"),
            },
        };
        write_response(&mut write_half, &response).await?;
    }
    Ok(())
}

async fn write_response(
    write_half: &mut tokio::net::unix::OwnedWriteHalf,
    response: &Response,
) -> Result<(), ServerError> {
    let mut payload = serde_json::to_string(response)
        .map_err(|e| ServerError::Protocol(format!("response encode failed: {e}")))?;
    payload.push('\n');
    write_half.write_all(payload.as_bytes()).await?;
    write_half.flush().await?;
    Ok(())
}

/// Version check then dispatch. Wrong-version requests get a typed error
/// (a stale daemon and a newer shim must fail loudly, not weirdly).
fn respond(request: &Request, warm: &WarmState, read_only: bool) -> Response {
    let v = match request {
        Request::Check { v, .. } | Request::Ping { v } | Request::Shutdown { v } => *v,
    };
    if v != PROTOCOL_VERSION {
        return Response::Error {
            message: format!("protocol version {v} unsupported (daemon speaks {PROTOCOL_VERSION})"),
        };
    }
    match request {
        Request::Check {
            file_path, content, ..
        } => Response::Check {
            result: warm.check(&WriteRequest {
                file_path: file_path.clone(),
                content: content.clone(),
            }),
        },
        Request::Ping { .. } => Response::Pong {
            info: DaemonInfo {
                pid: std::process::id(),
                version: env!("CARGO_PKG_VERSION").to_owned(),
                read_only,
            },
        },
        Request::Shutdown { .. } => {
            if read_only {
                Response::Error {
                    message: "read-only daemon: lifecycle mutations are refused over the wire \
                              (kill the process from the session that spawned it)"
                        .to_owned(),
                }
            } else {
                Response::ShuttingDown
            }
        }
    }
}

/// One round trip against the canonical socket of `repo_root`.
///
/// # Errors
/// `NotRunning` when the socket is absent or refuses connection (the
/// shim's fall-back-to-cold signal); `Io`/`Protocol` for transport and
/// encoding failures.
pub fn request(repo_root: &Path, request: &Request) -> Result<Response, ServerError> {
    request_at(&socket_path(repo_root), request)
}

/// One round trip against an explicit socket: connect, send one request
/// line, read one response line. Synchronous std transport — the shim
/// side has no runtime, and one blocking round trip is exactly its job.
///
/// # Errors
/// `NotRunning` when the socket is absent or refuses connection;
/// `Io`/`Protocol` for transport and encoding failures.
pub fn request_at(socket: &Path, request: &Request) -> Result<Response, ServerError> {
    use std::io::{BufRead, BufReader as StdBufReader, Write};

    let mut stream = match std::os::unix::net::UnixStream::connect(socket) {
        Ok(stream) => stream,
        Err(error) => {
            return Err(match error.kind() {
                std::io::ErrorKind::NotFound | std::io::ErrorKind::ConnectionRefused => {
                    ServerError::NotRunning
                }
                _ => ServerError::Io(error),
            })
        }
    };
    let mut payload = serde_json::to_string(request)
        .map_err(|e| ServerError::Protocol(format!("request encode failed: {e}")))?;
    payload.push('\n');
    stream.write_all(payload.as_bytes())?;
    stream.flush()?;

    let mut line = String::new();
    StdBufReader::new(&mut stream).read_line(&mut line)?;
    if line.is_empty() {
        return Err(ServerError::Protocol(
            "daemon closed the connection without responding".to_owned(),
        ));
    }
    serde_json::from_str(&line)
        .map_err(|e| ServerError::Protocol(format!("unrecognized response: {e}")))
}