haematite-cli 0.21.0

The haematite command: serve a store, its door and Studio, answer MCP, query with jq, and host the store's assistant seat
use super::config::Seat;
use super::wire::{Frame, Pipe};
use agent_client_protocol_schema::{ProtocolVersion, v1, v2};
use serde_json::{Value, json};

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(super) enum Version {
    One,
    Two,
}

impl Version {
    pub const fn number(self) -> u16 {
        match self {
            Self::One => 1,
            Self::Two => 2,
        }
    }
}

pub(super) struct Session {
    pub id: String,
    pub version: Version,
    pub resumable: bool,
    pub capability: &'static str,
    pub updates: Vec<(bool, Value)>,
    pub config: Value,
    pub info: Value,
}

pub(super) fn capabilities(answer: Value) -> anyhow::Result<(Version, bool, &'static str)> {
    let number: ProtocolVersion = serde_json::from_value(answer["protocolVersion"].clone())
        .map_err(|_| anyhow::anyhow!("acp_initialize_invalid: protocolVersion"))?;
    match number.as_u16() {
        1 => {
            let response: v1::InitializeResponse = decode(answer, "initialize")?;
            Ok((
                Version::One,
                response.agent_capabilities.load_session,
                "agentCapabilities.loadSession",
            ))
        }
        2 => {
            let response: v2::InitializeResponse = decode(answer, "initialize")?;
            anyhow::ensure!(
                response.capabilities.session.is_some(),
                "ACP 2 agent did not advertise capabilities.session"
            );
            // The pinned session capability includes resume among its baseline methods.
            Ok((
                Version::Two,
                true,
                "capabilities.session (schema 1.5.0 includes resume)",
            ))
        }
        version => {
            anyhow::bail!("ACP protocol version {version} refused: offered 2, supported 2 and 1")
        }
    }
}

pub(super) async fn open(pipe: &mut Pipe, seat: &Seat) -> anyhow::Result<Session> {
    let mut initialize = serde_json::to_value(v2::InitializeRequest::new(
        ProtocolVersion::V2,
        v2::Implementation::new("haematite", env!("CARGO_PKG_VERSION")),
    ))?;
    initialize["clientInfo"] = initialize["info"].clone();
    initialize["clientCapabilities"] =
        json!({"fs":{"readTextFile":false,"writeTextFile":false},"terminal":false});
    let mut updates = Vec::new();
    let answer = tokio::time::timeout(
        std::time::Duration::from_millis(seat.harness.initialize_timeout_ms),
        call(
            pipe,
            1,
            "initialize",
            initialize,
            &mut updates,
            seat.harness.stream_queue,
        ),
    )
    .await
    .map_err(|_| anyhow::anyhow!("acp_initialize_timeout"))??;
    let info = answer
        .get("info")
        .or_else(|| answer.get("agentInfo"))
        .cloned()
        .unwrap_or(Value::Null);
    let (version, resumable, capability) = capabilities(answer)?;
    anyhow::ensure!(
        seat.harness.session.is_none() || resumable,
        "session_resume_not_offered"
    );
    Ok(Session {
        id: seat.harness.session.clone().unwrap_or_default(),
        version,
        resumable,
        capability,
        updates: updates.into_iter().map(|params| (false, params)).collect(),
        config: json!([]),
        info,
    })
}

pub(super) fn session_request(
    seat: &Seat,
    version: Version,
    prior: Option<&str>,
    agent: &str,
) -> anyhow::Result<(&'static str, Value)> {
    let method = match (version, prior) {
        (_, None) => "session/new",
        (Version::One, Some(_)) => "session/load",
        (Version::Two, Some(_)) => "session/resume",
    };
    let mut params = json!({"cwd":seat.cwd()?,"mcpServers":[super::tools::server(seat,agent)?]});
    if let Some(id) = prior {
        params["sessionId"] = json!(id);
    }
    Ok((method, params))
}

async fn call(
    pipe: &mut Pipe,
    id: u64,
    method: &str,
    params: Value,
    deferred: &mut Vec<Value>,
    bound: usize,
) -> anyhow::Result<Value> {
    pipe.output.request(id, method, params)?;
    loop {
        match pipe
            .input
            .recv()
            .await
            .ok_or_else(|| anyhow::anyhow!("ACP reader ended during handshake"))??
            .ok_or_else(|| anyhow::anyhow!("ACP stdout closed during handshake"))?
        {
            Frame::Response {
                id: received,
                result,
            } => {
                anyhow::ensure!(
                    received == json!(id),
                    "ACP handshake response id was never requested"
                );
                return result.map_err(|error| {
                    let code = error.code;
                    if code == -32000 {
                        anyhow::anyhow!("ACP authentication required before {method} (code {code})")
                    } else {
                        anyhow::anyhow!("ACP {method} refused (code {code})")
                    }
                });
            }
            Frame::Request { id, method, .. } if method == "session/request_permission" => {
                pipe.output.cancelled(id)?;
            }
            Frame::Request { id, .. } => pipe.output.refuse(
                id,
                -32601,
                "no filesystem, terminal or credential callbacks",
            )?,
            Frame::Notification { method, .. } if method.starts_with('_') => {}
            Frame::Notification { method, params } => {
                anyhow::ensure!(
                    method == "session/update",
                    "ACP notification is not advertised during handshake"
                );
                anyhow::ensure!(
                    deferred.len() < bound,
                    "handshake updates exceed stream_queue {bound}"
                );
                deferred.push(params);
            }
        }
    }
}

pub(super) fn decode<T: serde::de::DeserializeOwned>(
    value: Value,
    method: &str,
) -> anyhow::Result<T> {
    serde_json::from_value(value).map_err(|_| anyhow::anyhow!("acp_{method}_invalid"))
}

pub(super) fn session_answer(
    version: Version,
    prior: Option<&str>,
    answer: Value,
) -> anyhow::Result<(String, Value)> {
    match (version, prior) {
        (Version::One, None) => {
            let answer: v1::NewSessionResponse = decode(answer, "session_new")?;
            Ok((
                answer.session_id.to_string(),
                serde_json::to_value(answer.config_options.unwrap_or_default())?,
            ))
        }
        (Version::Two, None) => {
            let answer: v2::NewSessionResponse = decode(answer, "session_new")?;
            Ok((
                answer.session_id.to_string(),
                serde_json::to_value(answer.config_options)?,
            ))
        }
        (Version::One, Some(id)) => {
            let answer: v1::LoadSessionResponse = decode(answer, "session_load")?;
            Ok((
                id.to_owned(),
                serde_json::to_value(answer.config_options.unwrap_or_default())?,
            ))
        }
        (Version::Two, Some(id)) => {
            let answer: v2::ResumeSessionResponse = decode(answer, "session_resume")?;
            Ok((id.to_owned(), serde_json::to_value(answer.config_options)?))
        }
    }
}