mnote/rust/crates/mnote-web/src/acp_client.rs

/// ACP (Agent Client Protocol) JSON-RPC 2.0 client.
///
/// Walks an agent runtime subprocess (e.g. `hermes acp` or `node reasonix-acp-wrapper.mjs`)
/// over NDJSON stdio: one JSON object per line, newline-delimited.
///
/// Reference implementations:
/// - `reference-code/hermes-vscode-main/src/acpClient.ts` (primary reference)
/// - `reference-code/DeepSeek-Reasonix-main/src/acp/protocol.ts`
///
/// Wire format:
///   Request:     { jsonrpc: "2.0", id: number, method: string, params?: object }
///   Response:    { jsonrpc: "2.0", id: number, result?: any, error?: { code, message } }
///   Notification:{ jsonrpc: "2.0", method: string, params?: object }           (no id)
///   Incoming:    { jsonrpc: "2.0", id: number, method: string, params?: object } (from agent)
use serde::de::DeserializeOwned;
use serde::Serialize;
use serde_json::{json, Value};
use std::collections::HashMap;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Arc;
use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader, BufWriter};
use tokio::process::{Child, ChildStdin, ChildStdout, Command};
use tokio::sync::{oneshot, Mutex};
use tokio::time::{timeout, Duration};
use tracing::{debug, info, warn};

// ── Error types ──────────────────────────────────────

#[derive(Debug)]
pub enum AcpError {
    Spawn(std::io::Error),
    JsonParse(serde_json::Error),
    JsonRpc { code: i64, message: String },
    Timeout(u64),
    Closed,
    Internal(String),
}

impl std::fmt::Display for AcpError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            AcpError::Spawn(e) => write!(f, "ACP spawn failed: {e}"),
            AcpError::JsonParse(e) => write!(f, "ACP JSON parse error: {e}"),
            AcpError::JsonRpc { code, message } => {
                write!(f, "ACP JSON-RPC error [{code}]: {message}")
            }
            AcpError::Timeout(secs) => write!(f, "ACP request timed out after {secs}s"),
            AcpError::Closed => write!(f, "ACP connection closed"),
            AcpError::Internal(msg) => write!(f, "ACP internal: {msg}"),
        }
    }
}

impl std::error::Error for AcpError {
    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
        match self {
            AcpError::Spawn(e) => Some(e),
            AcpError::JsonParse(e) => Some(e),
            _ => None,
        }
    }
}

impl From<std::io::Error> for AcpError {
    fn from(e: std::io::Error) -> Self {
        AcpError::Spawn(e)
    }
}

impl From<serde_json::Error> for AcpError {
    fn from(e: serde_json::Error) -> Self {
        AcpError::JsonParse(e)
    }
}

// ── Notification handler type ────────────────────────

type NotificationHandler = Box<dyn Fn(String, Value) + Send + 'static>;

/// Thread-safe mutex for notification handler (std mutex — lightweight, never held across awaits).
type NotificationHandlerMutex = std::sync::Mutex<Option<NotificationHandler>>;

// ── Pending request entry ────────────────────────────

type PendingEntry = oneshot::Sender<Result<Value, AcpError>>;

// ── AcpClient ────────────────────────────────────────

/// ACP JSON-RPC 2.0 client over stdio.
///
/// Create via [`AcpClient::spawn`], then use [`request`](Self::request) for RPC
/// calls and [`notification`](Self::notification) for fire-and-forget messages.
/// Register a handler with [`on_notification`](Self::on_notification) to receive
/// agent push events (e.g. `session/update`).
// Manual Debug impl: Child doesn't impl Debug, so we skip it
impl std::fmt::Debug for AcpClient {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("AcpClient")
            .field("next_id", &self.next_id)
            .field("pending_count", &self.pending.blocking_lock().len())
            .finish_non_exhaustive()
    }
}

pub struct AcpClient {
    child: Option<Child>,
    writer: Arc<Mutex<BufWriter<ChildStdin>>>,
    pending: Arc<Mutex<HashMap<u64, PendingEntry>>>,
    next_id: AtomicU64,
    notification_handler: Arc<NotificationHandlerMutex>,
}

impl AcpClient {
    /// Spawn an ACP subprocess and establish the JSON-RPC connection.
    ///
    /// After spawn, sends an `initialize` handshake (as Hermes does in acpClient.ts
    /// `start()` → `call('initialize', { protocolVersion: 1 })`).
    /// Launches a background tokio task that reads NDJSON lines from the child's stdout.
    ///
    /// Reference: `hermes-vscode-main/src/acpClient.ts` L50-80 (spawn + stdio setup)
    pub async fn spawn(bin: &str, args: &[&str]) -> Result<Self, AcpError> {
        Self::spawn_with_env(bin, args, None).await
    }

    /// Spawn an ACP subprocess with extra environment variables.
    pub async fn spawn_with_env(
        bin: &str,
        args: &[&str],
        env_overrides: Option<&HashMap<String, String>>,
    ) -> Result<Self, AcpError> {
        let mut command = Command::new(bin);
        command
            .args(args)
            .stdin(std::process::Stdio::piped())
            .stdout(std::process::Stdio::piped())
            .stderr(std::process::Stdio::inherit())
            .kill_on_drop(true);
        if let Some(env) = env_overrides {
            if let Some(workspace_root) = env.get("MNOTE_AI_WORKSPACE_ROOT") {
                let workspace_root = std::path::Path::new(workspace_root);
                if workspace_root.is_dir() {
                    // 本地 workspace run 以授权根目录作为进程工作目录，贴近 VSCode agent 行为。
                    command.current_dir(workspace_root);
                }
            }
            command.envs(env);
        }
        let mut child = command.spawn().map_err(AcpError::Spawn)?;

        let stdin = child
            .stdin
            .take()
            .ok_or_else(|| AcpError::Internal("failed to take child stdin".into()))?;
        let stdout = child
            .stdout
            .take()
            .ok_or_else(|| AcpError::Internal("failed to take child stdout".into()))?;

        let writer = Arc::new(Mutex::new(BufWriter::new(stdin)));
        let reader = BufReader::new(stdout);

        let pending: Arc<Mutex<HashMap<u64, PendingEntry>>> = Arc::new(Mutex::new(HashMap::new()));
        let notification_handler: Arc<NotificationHandlerMutex> =
            Arc::new(std::sync::Mutex::new(None));

        // Start background reader task
        let pending_clone = pending.clone();
        let handler_clone = notification_handler.clone();
        let writer_clone = writer.clone();
        let child_pid = child.id().unwrap_or(0);
        tokio::spawn(async move {
            Self::reader_loop(reader, writer_clone, pending_clone, handler_clone).await;
            info!("ACP reader loop ended (pid={})", child_pid);
        });

        let client = Self {
            child: Some(child),
            writer,
            pending,
            next_id: AtomicU64::new(1),
            notification_handler,
        };

        // Handshake: initialize (reference: acpClient.ts L108 → call('initialize', {protocolVersion: 1}))
        let init_result: Value = client
            .request("initialize", json!({ "protocolVersion": 1 }))
            .await?;
        debug!(?init_result, "ACP initialize OK");

        Ok(client)
    }

    /// Send a JSON-RPC request and await the response.
    ///
    /// Returns `Result<R>` where `R` is the deserialized `result` field.
    /// On JSON-RPC error, returns [`AcpError::JsonRpc`].
    /// Default timeout: 300 seconds.
    ///
    /// Reference: `acpClient.ts` L95-110 (`call()` method)
    pub async fn request<P: Serialize, R: DeserializeOwned>(
        &self,
        method: &str,
        params: P,
    ) -> Result<R, AcpError> {
        self.request_with_timeout(method, params, Duration::from_secs(300))
            .await
    }

    /// Same as [`request`] but with a configurable timeout.
    pub async fn request_with_timeout<P: Serialize, R: DeserializeOwned>(
        &self,
        method: &str,
        params: P,
        dur: Duration,
    ) -> Result<R, AcpError> {
        let id = self.next_id.fetch_add(1, Ordering::Relaxed);
        let (tx, rx) = oneshot::channel();

        self.pending.lock().await.insert(id, tx);

        let req = json!({
            "jsonrpc": "2.0",
            "id": id,
            "method": method,
            "params": params,
        });
        let line = serde_json::to_string(&req)?;
        debug!("ACP --> {} #{} ({} bytes)", method, id, line.len());

        let mut writer = self.writer.lock().await;
        writer.write_all(line.as_bytes()).await?;
        writer.write_all(b"\n").await?;
        writer.flush().await?;

        match timeout(dur, rx).await {
            Ok(Ok(Ok(value))) => {
                let result: R = serde_json::from_value(value)?;
                Ok(result)
            }
            Ok(Ok(Err(err))) => Err(err),
            Ok(Err(_recv_err)) => Err(AcpError::Closed),
            Err(_elapsed) => Err(AcpError::Timeout(dur.as_secs())),
        }
    }

    /// Send a fire-and-forget notification (no id, no response expected).
    ///
    /// Reference: `acpClient.ts` L115-118 (`notify()`)
    pub async fn notification(&self, method: &str, params: Value) -> Result<(), AcpError> {
        let msg = json!({
            "jsonrpc": "2.0",
            "method": method,
            "params": params,
        });
        let line = serde_json::to_string(&msg)?;
        debug!("ACP ~~> {} ({} bytes)", method, line.len());

        let mut writer = self.writer.lock().await;
        writer.write_all(line.as_bytes()).await?;
        writer.write_all(b"\n").await?;
        writer.flush().await?;
        Ok(())
    }

    /// Register a handler for incoming notifications (messages with `method` but no `id`).
    /// Only one handler at a time — subsequent calls replace the previous.
    pub fn on_notification<F>(&self, handler: F)
    where
        F: Fn(String, Value) + Send + 'static,
    {
        let mut guard = self.notification_handler.lock().unwrap();
        *guard = Some(Box::new(handler));
    }

    /// Gracefully close the ACP connection and kill the subprocess.
    pub async fn close(&mut self) {
        if let Some(mut child) = self.child.take() {
            let _ = child.start_kill();
            let _ = child.wait().await;
        }
        // Resolve all pending with Closed error
        let mut pending = self.pending.lock().await;
        for (_, tx) in pending.drain() {
            let _ = tx.send(Err(AcpError::Closed));
        }
    }

    // ── Background reader ────────────────────────────

    /// Background loop: reads NDJSON lines from the child's stdout,
    /// routes responses to pending requests and notifications to the handler.
    ///
    /// Reference: `acpClient.ts` L120-180 (onData + dispatch)
    async fn reader_loop(
        mut reader: BufReader<ChildStdout>,
        writer: Arc<Mutex<BufWriter<ChildStdin>>>,
        pending: Arc<Mutex<HashMap<u64, PendingEntry>>>,
        notification_handler: Arc<NotificationHandlerMutex>,
    ) {
        let mut line_buf = String::new();
        loop {
            line_buf.clear();
            match reader.read_line(&mut line_buf).await {
                Ok(0) => {
                    info!("ACP stdout closed (EOF)");
                    break;
                }
                Ok(_n) => {}
                Err(e) => {
                    warn!("ACP read error: {e}");
                    break;
                }
            }

            let trimmed = line_buf.trim();
            if trimmed.is_empty() {
                continue;
            }

            let msg: Value = match serde_json::from_str(trimmed) {
                Ok(v) => v,
                Err(e) => {
                    warn!(
                        "ACP parse error: {e} (line: {})",
                        &trimmed[..trimmed.len().min(80)]
                    );
                    continue;
                }
            };

            Self::dispatch_message(msg, &writer, &pending, &notification_handler).await;
        }

        // Process died or EOF — resolve all pending
        let mut pending_guard = pending.lock().await;
        for (_, tx) in pending_guard.drain() {
            let _ = tx.send(Err(AcpError::Closed));
        }
    }

    /// Route a single JSON message to pending request, notification handler, or incoming request.
    ///
    /// Reference: `acpClient.ts` L160-200 (dispatch)
    async fn dispatch_message(
        msg: Value,
        writer: &Arc<Mutex<BufWriter<ChildStdin>>>,
        pending: &Arc<Mutex<HashMap<u64, PendingEntry>>>,
        notification_handler: &Arc<NotificationHandlerMutex>,
    ) {
        let has_id = msg.get("id").is_some();
        let has_method = msg
            .get("method")
            .and_then(|v| v.as_str())
            .map(|s| !s.is_empty())
            .unwrap_or(false);

        if has_id && has_method {
            // Incoming request from agent (e.g. session/request_permission)
            let method = msg["method"].as_str().unwrap_or("unknown").to_string();
            let params = msg.get("params").cloned().unwrap_or(Value::Null);
            let id = msg.get("id").cloned().unwrap_or(Value::Null);
            // 当前还没有权限确认 UI，必须明确拒绝，避免 agent 等待到超时。
            warn!("ACP incoming request not handled: {method} (params={params:?})");
            let response = json!({
                "jsonrpc": "2.0",
                "id": id,
                "error": {
                    "code": -32601,
                    "message": format!("ACP incoming request not supported: {method}")
                }
            });
            if let Err(error) = Self::write_jsonrpc_message(writer, &response).await {
                warn!("ACP incoming request response write failed: {error}");
            }
            if method == "session/request_permission" {
                let mut event_params = params;
                if let Some(object) = event_params.as_object_mut() {
                    object.insert("decision".into(), Value::String("denied".into()));
                    object.insert("method".into(), Value::String(method.clone()));
                    object.insert("jsonrpcId".into(), id);
                }
                let handler_guard = notification_handler.lock().unwrap();
                if let Some(ref handler) = *handler_guard {
                    handler(method, event_params);
                }
            }
        } else if has_id {
            // Response to one of our requests
            if let Some(id) = msg["id"].as_u64() {
                let mut pending_guard = pending.lock().await;
                if let Some(tx) = pending_guard.remove(&id) {
                    if let Some(error) = msg.get("error") {
                        let code = error["code"].as_i64().unwrap_or(-1);
                        let message = error["message"]
                            .as_str()
                            .unwrap_or("unknown error")
                            .to_string();
                        let _ = tx.send(Err(AcpError::JsonRpc { code, message }));
                    } else if let Some(result) = msg.get("result") {
                        let _ = tx.send(Ok(result.clone()));
                    } else {
                        let _ = tx.send(Err(AcpError::Internal(
                            "response without result or error".into(),
                        )));
                    }
                } else {
                    debug!("ACP response for unknown request id={id}");
                }
            }
        } else if has_method {
            // Notification (no id)
            let method = msg["method"].as_str().unwrap_or("unknown").to_string();
            let params = msg.get("params").cloned().unwrap_or(Value::Null);
            let handler_guard = notification_handler.lock().unwrap();
            if let Some(ref handler) = *handler_guard {
                handler(method, params);
            } else {
                debug!("ACP notification unhandled: {method}");
            }
        }
    }

    async fn write_jsonrpc_message(
        writer: &Arc<Mutex<BufWriter<ChildStdin>>>,
        msg: &Value,
    ) -> Result<(), AcpError> {
        let line = serde_json::to_string(msg)?;
        let mut writer = writer.lock().await;
        writer.write_all(line.as_bytes()).await?;
        writer.write_all(b"\n").await?;
        writer.flush().await?;
        Ok(())
    }
}

impl Drop for AcpClient {
    fn drop(&mut self) {
        if let Some(mut child) = self.child.take() {
            let _ = child.start_kill();
        }
    }
}

// ── Tests ────────────────────────────────────────────

#[cfg(test)]
mod tests {
    use super::*;
    /// Helper: create a mock subprocess that echoes back requests as responses.
    /// Simulates a minimal ACP server for testing.
    async fn spawn_mock_acp_server() -> AcpClient {
        // We spawn a small node script that reads NDJSON and echoes back
        let script = r#"
import * as readline from 'node:readline';
import { stdin as input, stdout as output } from 'node:process';
const rl = readline.createInterface({ input, output, terminal: false });
rl.on('line', (line) => {
    const msg = JSON.parse(line);
    if (msg.id !== undefined && msg.method) {
        if (msg.method === 'initialize') {
            process.stdout.write(JSON.stringify({
                jsonrpc: '2.0',
                id: msg.id,
                result: { protocolVersion: 1, agentCapabilities: {}, agentInfo: { name: 'mock', version: '1.0' }, authMethods: [] }
            }) + '\n');
        } else {
            process.stdout.write(JSON.stringify({
                jsonrpc: '2.0',
                id: msg.id,
                result: { ok: true, echo: msg.params }
            }) + '\n');
        }
    } else if (msg.method && msg.id === undefined) {
        // Notification → ignore
    }
});
        "#;

        // Write script to temp file
        let dir = std::env::temp_dir();
        let script_path = dir.join("acp_test_mock.mjs");
        std::fs::write(&script_path, script).expect("write mock script");

        AcpClient::spawn("node", &[script_path.to_str().unwrap()])
            .await
            .expect("spawn mock ACP")
    }

    async fn spawn_permission_request_mock_server() -> AcpClient {
        let script = r#"
import * as readline from 'node:readline';
import { stdin as input, stdout as output } from 'node:process';
let permissionResponse = null;
const rl = readline.createInterface({ input, output, terminal: false });
function send(msg) { process.stdout.write(JSON.stringify(msg) + '\n'); }
rl.on('line', (line) => {
    const msg = JSON.parse(line);
    if (msg.id !== undefined && msg.method === 'initialize') {
        send({
            jsonrpc: '2.0',
            id: msg.id,
            result: { protocolVersion: 1, agentCapabilities: {}, agentInfo: { name: 'mock', version: '1.0' }, authMethods: [] }
        });
        setTimeout(() => send({
            jsonrpc: '2.0',
            id: 77,
            method: 'session/request_permission',
            params: { reason: 'test permission' }
        }), 10);
    } else if (msg.id === 77 && msg.method === undefined) {
        permissionResponse = msg;
    } else if (msg.id !== undefined && msg.method === 'get_permission_response') {
        send({
            jsonrpc: '2.0',
            id: msg.id,
            result: { permissionResponse }
        });
    }
});
        "#;

        let dir = std::env::temp_dir();
        let script_path = dir.join("acp_permission_request_mock.mjs");
        std::fs::write(&script_path, script).expect("write permission mock script");

        AcpClient::spawn("node", &[script_path.to_str().unwrap()])
            .await
            .expect("spawn permission mock ACP")
    }

    #[tokio::test]
    async fn test_request_response() {
        let client = spawn_mock_acp_server().await;
        let result: Value = client
            .request("test_method", json!({ "hello": "world" }))
            .await
            .expect("request should succeed");

        assert_eq!(result["ok"], true);
        assert_eq!(result["echo"]["hello"], "world");
    }

    #[tokio::test]
    async fn test_notification() {
        let client = spawn_mock_acp_server().await;
        // Notifications are fire-and-forget, no response expected
        client
            .notification("test_notify", json!({ "foo": "bar" }))
            .await
            .expect("notification should succeed");
    }

    #[tokio::test]
    async fn test_on_notification_received() {
        use std::sync::atomic::AtomicBool;
        let client = spawn_mock_acp_server().await;
        let received = Arc::new(AtomicBool::new(false));
        let received_clone = received.clone();

        client.on_notification(move |method, _params| {
            if method == "test_push" {
                received_clone.store(true, Ordering::SeqCst);
            }
        });

        // Send a notification that the mock server will echo back as...
        // Actually the mock doesn't send unsolicited notifications.
        // This test just validates the handler registration doesn't crash.
        client
            .notification("test_push", json!({}))
            .await
            .expect("notification");

        // Give background task time to process
        tokio::time::sleep(Duration::from_millis(100)).await;
        // In this mock, no notification will be received; that's OK
    }

    #[tokio::test]
    async fn test_close() {
        let mut client = spawn_mock_acp_server().await;
        client.close().await;
        // Second close should be no-op
        client.close().await;
    }

    #[tokio::test]
    async fn test_initialize_handshake() {
        // spawn already calls initialize; if it fails, the test fails
        let _client = spawn_mock_acp_server().await;
    }

    #[tokio::test]
    async fn test_incoming_permission_request_gets_response() {
        let client = spawn_permission_request_mock_server().await;
        tokio::time::sleep(Duration::from_millis(100)).await;
        let result: Value = client
            .request("get_permission_response", json!({}))
            .await
            .expect("permission response probe");
        let response = &result["permissionResponse"];
        assert_eq!(response["jsonrpc"], "2.0");
        assert_eq!(response["id"], 77);
        assert!(response.get("result").is_some() || response.get("error").is_some());
    }
}