scryer-mcp 0.2.1

//! One-line JSON handshake exchanged before MCP traffic on a daemon connection.

use std::path::PathBuf;

use serde::{Deserialize, Serialize, de::DeserializeOwned};
use tokio::io::{AsyncBufRead, AsyncBufReadExt, AsyncReadExt, AsyncWrite, AsyncWriteExt};

/// Handshake protocol version. Bump only on incompatible changes to these types.
pub const PROTOCOL_VERSION: u32 = 1;

/// Upper bound on a handshake line, so a bad peer cannot make us buffer forever.
const MAX_LINE_BYTES: u64 = 64 * 1024;

/// First line a client sends after connecting.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct Hello {
    /// Handshake protocol version ([`PROTOCOL_VERSION`]).
    pub scryer: u32,
    /// Client's Scryer version.
    pub version: String,
    /// Client's build ID ([`crate::build_info::build_id`]). Absent from older clients.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub build: Option<String>,
    #[serde(flatten)]
    pub mode: HelloMode,
}

/// What the client wants from the connection.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(tag = "mode", rename_all = "snake_case")]
pub enum HelloMode {
    /// Open an MCP session; raw JSON-RPC follows the handshake.
    Mcp {
        /// Client working directory, used for relative paths and `--watch`.
        cwd: Option<PathBuf>,
        /// Watch the project containing `cwd` for the life of the session.
        #[serde(default)]
        watch: bool,
    },
    /// Run a single control command; the connection closes after the reply.
    Control { cmd: ControlCommand },
}

/// Daemon control commands.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum ControlCommand {
    Status,
    Stop,
    /// Stop only if no session is open, atomically with refusing new sessions.
    ///
    /// Replies with no `status` when the daemon is stopping, or with its `status` when it
    /// has sessions and keeps running. Daemons that predate this command reject it.
    StopIfIdle,
}

/// Daemon reply to a [`Hello`].
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct HelloReply {
    pub ok: bool,
    /// Daemon's Scryer version.
    pub version: String,
    /// Daemon's build ID. Absent from older daemons.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub build: Option<String>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub error: Option<String>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub status: Option<DaemonStatus>,
}

/// Snapshot returned by [`ControlCommand::Status`].
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct DaemonStatus {
    pub pid: u32,
    pub version: String,
    /// Daemon's build ID. Absent from older daemons.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub build: Option<String>,
    pub db_path: PathBuf,
    pub sessions: usize,
    pub uptime_secs: u64,
    pub watched: Vec<WatchedProject>,
}

/// A project the daemon is currently watching.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct WatchedProject {
    pub project_id: u64,
    pub root: PathBuf,
    /// Number of directories watched for this project (ignored directories are not watched).
    #[serde(default)]
    pub directories: usize,
}

impl Hello {
    /// Build a handshake for this build's version.
    pub fn new(mode: HelloMode) -> Self {
        Self {
            scryer: PROTOCOL_VERSION,
            version: crate::mcp_version().to_string(),
            build: Some(crate::build_info::build_id().to_string()),
            mode,
        }
    }
}

impl HelloReply {
    /// Successful reply carrying this build's version.
    pub fn ok() -> Self {
        Self {
            ok: true,
            version: crate::mcp_version().to_string(),
            build: Some(crate::build_info::build_id().to_string()),
            error: None,
            status: None,
        }
    }

    /// Failed reply carrying this build's version.
    pub fn error(message: impl Into<String>) -> Self {
        Self {
            ok: false,
            error: Some(message.into()),
            ..Self::ok()
        }
    }
}

/// Write `value` as a single JSON line and flush.
pub async fn write_json_line<W, T>(writer: &mut W, value: &T) -> anyhow::Result<()>
where
    W: AsyncWrite + Unpin,
    T: Serialize,
{
    let mut line = serde_json::to_vec(value)?;
    line.push(b'\n');
    writer.write_all(&line).await?;
    writer.flush().await?;
    Ok(())
}

/// Read one JSON line. Bytes after the newline stay buffered in `reader`.
pub async fn read_json_line<R, T>(reader: &mut R) -> anyhow::Result<T>
where
    R: AsyncBufRead + Unpin,
    T: DeserializeOwned,
{
    let mut line = Vec::new();
    let read = (&mut *reader)
        .take(MAX_LINE_BYTES)
        .read_until(b'\n', &mut line)
        .await?;
    if read == 0 {
        anyhow::bail!("connection closed before handshake");
    }
    if line.last() != Some(&b'\n') {
        anyhow::bail!("handshake line is truncated or exceeds {MAX_LINE_BYTES} bytes");
    }
    Ok(serde_json::from_slice(&line)?)
}

#[cfg(test)]
mod tests {
    use super::*;
    use tokio::io::BufReader;

    #[tokio::test]
    async fn handshake_preserves_trailing_bytes() {
        let hello = Hello::new(HelloMode::Mcp {
            cwd: Some(PathBuf::from("/repo")),
            watch: true,
        });
        let mut wire = Vec::new();
        write_json_line(&mut wire, &hello).await.unwrap();
        wire.extend_from_slice(b"{\"jsonrpc\":\"2.0\"}\n");

        let mut reader = BufReader::new(wire.as_slice());
        let decoded: Hello = read_json_line(&mut reader).await.unwrap();
        assert_eq!(decoded, hello);

        let mut rest = String::new();
        reader.read_to_string(&mut rest).await.unwrap();
        assert_eq!(rest, "{\"jsonrpc\":\"2.0\"}\n");
    }

    #[test]
    fn hello_wire_format_is_flat() {
        let json = serde_json::to_value(Hello::new(HelloMode::Control {
            cmd: ControlCommand::Stop,
        }))
        .unwrap();
        assert_eq!(json["mode"], "control");
        assert_eq!(json["cmd"], "stop");
        assert_eq!(json["scryer"], PROTOCOL_VERSION);
    }
}