Genarrative/rust/crates/agent-cli/tests/app_server.rs

//! `agent app-server --stdio` 的黑盒协议测试。
//!
//! 测试只通过子进程的 stdin/stdout 交互，不读取宿主环境中的凭据，也不依赖
//! AGC 的实现细节；这样可以把协议回归和 CLI 的装配/生命周期一起验收。

use std::fs;
use std::io::{BufRead, BufReader, Write};
use std::path::{Path, PathBuf};
use std::process::{Child, ChildStdin, Command, Stdio};
use std::sync::mpsc::{self, Receiver};
use std::thread;
use std::time::Duration;

use serde_json::{Value, json};

const IO_TIMEOUT: Duration = Duration::from_secs(10);

struct TempDir {
    path: PathBuf,
}

impl TempDir {
    fn new() -> Self {
        let base = std::env::var_os("TMPDIR")
            .map(PathBuf::from)
            .unwrap_or_else(|| PathBuf::from("/home/dsk/data/tmp"));
        fs::create_dir_all(&base).expect("create test temp base");
        for n in 0..1000u32 {
            let path = base.join(format!("agent-app-server-{n}-{}", std::process::id()));
            if fs::create_dir(&path).is_ok() {
                return Self { path };
            }
        }
        panic!("unable to allocate test temp directory");
    }

    fn path(&self) -> &Path {
        &self.path
    }
}

impl Drop for TempDir {
    fn drop(&mut self) {
        // 测试只删除自己生成的精确目录；目录内是该测试进程独占的数据库和配置。
        let _ = fs::remove_dir_all(&self.path);
    }
}

struct AppServer {
    child: Child,
    stdin: Option<ChildStdin>,
    lines: Receiver<String>,
}

impl AppServer {
    fn start(temp: &TempDir) -> Self {
        let config = temp.path().join("empty.toml");
        fs::write(&config, "").expect("write empty config");
        let db = temp.path().join("agent.db");
        let home = temp.path().join("home");
        fs::create_dir_all(&home).expect("create isolated home");
        let mut child = Command::new(env!("CARGO_BIN_EXE_agent"))
            .args(["app-server", "--stdio"])
            .current_dir(temp.path())
            .env_clear()
            // `env_clear` prevents credential leakage, while PATH keeps the
            // child runtime's normal process environment usable.
            .env("PATH", std::env::var_os("PATH").unwrap_or_default())
            .env("AGENT_PROVIDER", "fake")
            .env("AGENT_DB", &db)
            .env("AGENT_CONFIG", &config)
            .env("TMPDIR", temp.path())
            .env("HOME", &home)
            .stdin(Stdio::piped())
            .stdout(Stdio::piped())
            .stderr(Stdio::null())
            .spawn()
            .expect("spawn agent app-server");
        let stdin = child.stdin.take().expect("app-server stdin");
        let stdout = child.stdout.take().expect("app-server stdout");
        let (tx, rx) = mpsc::channel();
        thread::spawn(move || {
            let reader = BufReader::new(stdout);
            for line in reader.lines() {
                match line {
                    Ok(line) => {
                        if tx.send(line).is_err() {
                            break;
                        }
                    }
                    Err(_) => break,
                }
            }
        });
        Self {
            child,
            stdin: Some(stdin),
            lines: rx,
        }
    }

    fn send(&mut self, value: Value) {
        let stdin = self.stdin.as_mut().expect("app-server stdin is open");
        writeln!(stdin, "{value}").expect("write app-server request");
        stdin.flush().expect("flush app-server request");
    }

    fn close_stdin(&mut self) {
        // Drop the pipe explicitly to exercise EOF cooperative shutdown.
        self.stdin.take();
    }

    fn next(&self) -> Value {
        let line = self
            .lines
            .recv_timeout(IO_TIMEOUT)
            .expect("timed out waiting for app-server output");
        serde_json::from_str(&line)
            .unwrap_or_else(|error| panic!("invalid JSONL output: {line}: {error}"))
    }

    fn response(&self, id: i64) -> Value {
        loop {
            let message = self.next();
            if message.get("id") == Some(&json!(id)) {
                return message;
            }
        }
    }

    fn initialize(&mut self) {
        self.send(
            json!({"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":1}}),
        );
        let response = self.response(1);
        assert_eq!(response["result"]["protocolVersion"], 1);
    }

    fn run_to_completion(&mut self, id: i64, task: &str, stream: bool) -> String {
        self.send(json!({"jsonrpc":"2.0","id":id,"method":"run/start","params":{"task":task,"stream":stream}}));
        let accepted = self.response(id);
        let run_id = accepted["result"]["runId"]
            .as_str()
            .expect("runId")
            .to_owned();
        for _ in 0..64 {
            let message = self.next();
            if message["method"] == "run/completed" && message["params"]["runId"] == run_id {
                assert!(message["params"]["result"]["output"].is_object());
                return run_id;
            }
        }
        panic!("run/completed notification not received");
    }
}

impl Drop for AppServer {
    fn drop(&mut self) {
        if let Some(stdin) = self.stdin.as_mut() {
            let _ = stdin.flush();
        }
        if self.child.try_wait().ok().flatten().is_none() {
            let _ = self.child.kill();
        }
        let _ = self.child.wait();
    }
}

#[test]
fn protocol_requires_initialize_and_rejects_invalid_requests() {
    let temp = TempDir::new();
    let mut server = AppServer::start(&temp);
    server.send(json!({"jsonrpc":"2.0","id":1,"method":"run/get","params":{"runId":"missing"}}));
    let response = server.response(1);
    assert_eq!(response["error"]["code"], -32002);

    server.send(
        json!({"jsonrpc":"2.0","id":2,"method":"initialize","params":{"protocolVersion":99}}),
    );
    assert_eq!(server.response(2)["error"]["code"], -32602);
    server.initialize();

    server.send(json!({"jsonrpc":"2.0","id":6,"method":"run/get","params":{"runId":"missing"}}));
    assert_eq!(server.response(6)["error"]["code"], -32004);

    server
        .send(json!({"jsonrpc":"2.0","id":3,"method":"initialize","params":{"protocolVersion":1}}));
    assert_eq!(server.response(3)["error"]["code"], -32600);
    server.send(json!({"jsonrpc":"2.0","id":4,"method":"unknown","params":{}}));
    assert_eq!(server.response(4)["error"]["code"], -32601);
    server.send(json!({"jsonrpc":"2.0","id":5,"method":"run/start","params":{"task":""}}));
    assert_eq!(server.response(5)["error"]["code"], -32602);
    server.send(json!({"jsonrpc":"2.0","id":7,"method":"run/start","params":{"task":42}}));
    assert_eq!(server.response(7)["error"]["code"], -32602);
}

#[test]
fn run_stream_tool_and_persisted_queries_form_one_closed_loop() {
    let temp = TempDir::new();
    let mut server = AppServer::start(&temp);
    server.initialize();
    server.send(json!({"jsonrpc":"2.0","id":2,"method":"run/start","params":{"task":"hello","stream":true}}));
    let accepted = server.response(2);
    let run_id = accepted["result"]["runId"]
        .as_str()
        .expect("runId")
        .to_owned();
    assert!(accepted["result"]["sessionId"].is_string());
    assert!(accepted["result"]["runtimeId"].is_string());

    let mut completed = None;
    let mut saw_stream = false;
    let mut saw_tool = false;
    for _ in 0..32 {
        let message = server.next();
        match message["method"].as_str() {
            Some("run/stream") => {
                saw_stream = true;
            }
            Some("run/event") => {
                // Engine 的实时流只转发 Provider 事件；工具闭环属于 durable
                // run/event，事件标签按协议保持 snake_case。
                saw_tool |= message["params"]["event"]["type"]
                    .as_str()
                    .is_some_and(|kind| kind.contains("tool"));
            }
            Some("run/completed") => {
                if message["params"]["runId"] == run_id {
                    completed = Some(message);
                    break;
                }
            }
            Some("run/failed") | Some("run/cancelled") | Some("run/paused")
                if message["params"]["runId"] == run_id =>
            {
                panic!("run did not complete: {message}");
            }
            _ => {}
        }
    }
    let completed = completed.expect("run/completed notification");
    assert!(
        saw_stream,
        "stream=true must produce run/stream notifications"
    );
    assert!(saw_tool, "fake provider tool loop must be observable");
    assert!(completed["params"]["result"]["output"].is_object());

    // A long-lived app-server connection must be reusable. The implementation
    // reassembles the deterministic fake provider for each accepted run.
    let second_run = server.run_to_completion(6, "second run", false);
    assert_ne!(second_run, run_id);

    server.send(json!({"jsonrpc":"2.0","id":3,"method":"run/get","params":{"runId":run_id}}));
    let record = server.response(3)["result"].clone();
    assert!(
        record["status"]
            .as_str()
            .is_some_and(|status| status == "completed" || status == "succeeded")
    );
    server.send(json!({"jsonrpc":"2.0","id":4,"method":"run/events","params":{"runId":run_id,"afterRevision":0}}));
    assert!(server.response(4)["result"].is_array());

    server.send(json!({"jsonrpc":"2.0","id":5,"method":"shutdown","params":{}}));
    assert_eq!(server.response(5)["result"]["shutdown"], true);

    // Durable records remain queryable after a process restart.
    let mut restarted = AppServer::start(&temp);
    restarted.initialize();
    restarted.send(json!({"jsonrpc":"2.0","id":2,"method":"run/get","params":{"runId":run_id}}));
    assert!(restarted.response(2)["result"]["id"].is_string());
    restarted.send(json!({"jsonrpc":"2.0","id":3,"method":"run/events","params":{"runId":run_id,"afterRevision":0}}));
    assert!(
        restarted.response(3)["result"]
            .as_array()
            .is_some_and(|events| !events.is_empty())
    );
    restarted.send(json!({"jsonrpc":"2.0","id":4,"method":"shutdown","params":{}}));
    assert_eq!(restarted.response(4)["result"]["shutdown"], true);
}

#[test]
fn shutdown_is_acknowledged_and_process_exits_without_orphan() {
    let temp = TempDir::new();
    let mut server = AppServer::start(&temp);
    server.initialize();
    server.send(json!({"jsonrpc":"2.0","id":2,"method":"shutdown","params":{}}));
    assert_eq!(server.response(2)["result"]["shutdown"], true);
    assert!(
        server.child.wait_timeout(IO_TIMEOUT).is_ok(),
        "shutdown must terminate process"
    );
}

#[test]
fn stdin_eof_terminates_server_without_orphan_process() {
    let temp = TempDir::new();
    let mut server = AppServer::start(&temp);
    server.initialize();
    server.close_stdin();
    assert!(
        server.child.wait_timeout(IO_TIMEOUT).is_ok(),
        "EOF must terminate process"
    );
}

trait ChildWaitTimeout {
    fn wait_timeout(&mut self, timeout: Duration) -> std::io::Result<std::process::ExitStatus>;
}

impl ChildWaitTimeout for Child {
    fn wait_timeout(&mut self, timeout: Duration) -> std::io::Result<std::process::ExitStatus> {
        let start = std::time::Instant::now();
        loop {
            if let Some(status) = self.try_wait()? {
                return Ok(status);
            }
            if start.elapsed() >= timeout {
                return Err(std::io::Error::new(
                    std::io::ErrorKind::TimedOut,
                    "child still running",
                ));
            }
            thread::sleep(Duration::from_millis(20));
        }
    }
}