| 1 | //! The raw JSON-RPC test client and its environment. |
| 2 | |
| 3 | use std::collections::VecDeque; |
| 4 | use std::path::{Path, PathBuf}; |
| 5 | use std::process::Stdio; |
| 6 | use std::time::Duration; |
| 7 | |
| 8 | use serde_json::{Value, json}; |
| 9 | use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader}; |
| 10 | use tokio::process::{Child, ChildStdin, ChildStdout, Command}; |
| 11 | |
| 12 | /// How long to wait for any one message from the agent. |
| 13 | pub const TIMEOUT: Duration = Duration::from_secs(30); |
| 14 | |
| 15 | /// What the fake client does when the agent asks for permission. |
| 16 | #[derive(Debug, Clone, Copy)] |
| 17 | pub enum Perm { |
| 18 | Allow, |
| 19 | Reject, |
| 20 | RejectAlways, |
| 21 | } |
| 22 | |
| 23 | /// An ACP client speaking raw JSON-RPC to the agent's stdio, so tests see exactly what is on |
| 24 | /// the wire. It answers `session/request_permission` per [`Raw::perm`] and fails every other |
| 25 | /// agent-to-client request. |
| 26 | pub struct Raw { |
| 27 | pub child: Child, |
| 28 | stdin: Option<ChildStdin>, |
| 29 | lines: tokio::io::Lines<BufReader<ChildStdout>>, |
| 30 | pub next_id: u64, |
| 31 | /// Notifications (`session/update` etc.), in arrival order. |
| 32 | pub notes: Vec<Value>, |
| 33 | /// Permission requests seen, as the full JSON-RPC `params`. |
| 34 | pub permissions: Vec<Value>, |
| 35 | pub perm: Perm, |
| 36 | queued: VecDeque<Value>, |
| 37 | } |
| 38 | |
| 39 | /// The agent's command and environment for one test. The process gets nothing else from the |
| 40 | /// test's environment. |
| 41 | #[derive(Debug, Clone)] |
| 42 | pub struct Env { |
| 43 | pub bin: PathBuf, |
| 44 | pub prefix: &'static str, |
| 45 | pub vars: Vec<(String, String)>, |
| 46 | } |
| 47 | |
| 48 | impl Env { |
| 49 | /// `<PREFIX>_<suffix>`. |
| 50 | pub fn var(&self, suffix: &str) -> String { |
| 51 | format!("{}_{suffix}", self.prefix) |
| 52 | } |
| 53 | pub fn set(&mut self, k: &str, v: impl Into<String>) -> &mut Self { |
| 54 | self.vars.retain(|(n, _)| n != k); |
| 55 | self.vars.push((k.to_string(), v.into())); |
| 56 | self |
| 57 | } |
| 58 | pub fn unset(&mut self, k: &str) -> &mut Self { |
| 59 | self.vars.retain(|(n, _)| n != k); |
| 60 | self |
| 61 | } |
| 62 | } |
| 63 | |
| 64 | impl Raw { |
| 65 | /// Start the agent with `--acp`. |
| 66 | pub fn spawn(env: &Env) -> Raw { |
| 67 | let mut cmd = Command::new(&env.bin); |
| 68 | cmd.arg("--acp") |
| 69 | .env_clear() |
| 70 | .stdin(Stdio::piped()) |
| 71 | .stdout(Stdio::piped()) |
| 72 | .stderr(Stdio::inherit()) |
| 73 | .kill_on_drop(true); |
| 74 | for (k, v) in &env.vars { |
| 75 | cmd.env(k, v); |
| 76 | } |
| 77 | let mut child = cmd.spawn().expect("spawn agent"); |
| 78 | let stdin = child.stdin.take(); |
| 79 | let stdout = child.stdout.take().unwrap(); |
| 80 | Raw { |
| 81 | child, |
| 82 | stdin, |
| 83 | lines: BufReader::new(stdout).lines(), |
| 84 | next_id: 1, |
| 85 | notes: Vec::new(), |
| 86 | permissions: Vec::new(), |
| 87 | perm: Perm::Allow, |
| 88 | queued: VecDeque::new(), |
| 89 | } |
| 90 | } |
| 91 | |
| 92 | pub async fn send(&mut self, msg: Value) { |
| 93 | let stdin = self.stdin.as_mut().expect("stdin closed"); |
| 94 | stdin |
| 95 | .write_all(format!("{msg}\n").as_bytes()) |
| 96 | .await |
| 97 | .unwrap(); |
| 98 | stdin.flush().await.unwrap(); |
| 99 | } |
| 100 | |
| 101 | pub async fn notify(&mut self, method: &str, params: Value) { |
| 102 | self.send(json!({"jsonrpc": "2.0", "method": method, "params": params})) |
| 103 | .await; |
| 104 | } |
| 105 | |
| 106 | pub async fn next_message(&mut self) -> Value { |
| 107 | if let Some(m) = self.queued.pop_front() { |
| 108 | return m; |
| 109 | } |
| 110 | let line = tokio::time::timeout(TIMEOUT, self.lines.next_line()) |
| 111 | .await |
| 112 | .expect("timed out waiting for the agent") |
| 113 | .unwrap() |
| 114 | .expect("agent closed stdout"); |
| 115 | serde_json::from_str(&line).unwrap_or_else(|e| panic!("non-JSON on stdout ({e}): {line}")) |
| 116 | } |
| 117 | |
| 118 | /// Send a request and wait for its response, serving agent→client requests meanwhile. |
| 119 | pub async fn call(&mut self, method: &str, params: Value) -> Result<Value, Value> { |
| 120 | let id = self.next_id; |
| 121 | self.next_id += 1; |
| 122 | self.send(json!({"jsonrpc": "2.0", "id": id, "method": method, "params": params})) |
| 123 | .await; |
| 124 | loop { |
| 125 | let msg = self.next_message().await; |
| 126 | if msg.get("method").is_some() { |
| 127 | if msg.get("id").is_some() { |
| 128 | self.serve(&msg).await; |
| 129 | } else { |
| 130 | self.notes.push(msg); |
| 131 | } |
| 132 | continue; |
| 133 | } |
| 134 | if msg["id"] == json!(id) { |
| 135 | return match msg.get("error") { |
| 136 | Some(e) => Err(e.clone()), |
| 137 | None => Ok(msg["result"].clone()), |
| 138 | }; |
| 139 | } |
| 140 | } |
| 141 | } |
| 142 | |
| 143 | pub async fn serve(&mut self, req: &Value) { |
| 144 | let id = req["id"].clone(); |
| 145 | let reply = match req["method"].as_str().unwrap() { |
| 146 | "session/request_permission" => { |
| 147 | self.permissions.push(req["params"].clone()); |
| 148 | match self.perm { |
| 149 | Perm::Allow => { |
| 150 | json!({"outcome": {"outcome": "selected", "optionId": "allow_once"}}) |
| 151 | } |
| 152 | Perm::Reject => { |
| 153 | json!({"outcome": {"outcome": "selected", "optionId": "reject_once"}}) |
| 154 | } |
| 155 | Perm::RejectAlways => { |
| 156 | json!({"outcome": {"outcome": "selected", "optionId": "reject_always"}}) |
| 157 | } |
| 158 | } |
| 159 | } |
| 160 | other => { |
| 161 | self.send(json!({"jsonrpc": "2.0", "id": id, "error": {"code": -32601, "message": format!("unsupported {other}")}})).await; |
| 162 | return; |
| 163 | } |
| 164 | }; |
| 165 | self.send(json!({"jsonrpc": "2.0", "id": id, "result": reply})) |
| 166 | .await; |
| 167 | } |
| 168 | |
| 169 | /// Read whatever the agent sends right after a response (e.g. the commands update that |
| 170 | /// follows `session/new`) until it has been quiet for a moment. |
| 171 | pub async fn settle(&mut self) { |
| 172 | while let Ok(line) = |
| 173 | tokio::time::timeout(Duration::from_millis(300), self.lines.next_line()).await |
| 174 | { |
| 175 | let Ok(Some(line)) = line else { return }; |
| 176 | let msg: Value = serde_json::from_str(&line).unwrap(); |
| 177 | if msg.get("id").is_some() && msg.get("method").is_some() { |
| 178 | self.serve(&msg).await; |
| 179 | } else { |
| 180 | self.notes.push(msg); |
| 181 | } |
| 182 | } |
| 183 | } |
| 184 | |
| 185 | pub async fn ok(&mut self, method: &str, params: Value) -> Value { |
| 186 | self.call(method, params) |
| 187 | .await |
| 188 | .unwrap_or_else(|e| panic!("{method} failed: {e}")) |
| 189 | } |
| 190 | |
| 191 | pub async fn initialize(&mut self) -> Value { |
| 192 | self.ok( |
| 193 | "initialize", |
| 194 | json!({"protocolVersion": 1, "clientCapabilities": {}}), |
| 195 | ) |
| 196 | .await |
| 197 | } |
| 198 | |
| 199 | pub async fn new_session(&mut self, cwd: &Path, mcp: Value) -> String { |
| 200 | let r = self |
| 201 | .ok("session/new", json!({"cwd": cwd, "mcpServers": mcp})) |
| 202 | .await; |
| 203 | self.settle().await; |
| 204 | r["sessionId"].as_str().unwrap().to_string() |
| 205 | } |
| 206 | |
| 207 | pub async fn prompt(&mut self, sid: &str, blocks: Value) -> Value { |
| 208 | self.ok( |
| 209 | "session/prompt", |
| 210 | json!({"sessionId": sid, "prompt": blocks}), |
| 211 | ) |
| 212 | .await |
| 213 | } |
| 214 | |
| 215 | pub async fn text_prompt(&mut self, sid: &str, text: &str) -> Value { |
| 216 | self.prompt(sid, json!([{"type": "text", "text": text}])) |
| 217 | .await |
| 218 | } |
| 219 | |
| 220 | /// Notifications for a session update kind, e.g. `tool_call_update`. |
| 221 | pub fn updates(&self, kind: &str) -> Vec<Value> { |
| 222 | self.notes |
| 223 | .iter() |
| 224 | .filter(|n| { |
| 225 | n["method"] == "session/update" && n["params"]["update"]["sessionUpdate"] == kind |
| 226 | }) |
| 227 | .map(|n| n["params"]["update"].clone()) |
| 228 | .collect() |
| 229 | } |
| 230 | |
| 231 | /// Close stdin and wait for the agent to exit on its own. |
| 232 | pub async fn finish(mut self) { |
| 233 | drop(self.stdin.take()); |
| 234 | let _ = tokio::time::timeout(TIMEOUT, self.child.wait()) |
| 235 | .await |
| 236 | .expect("agent did not exit"); |
| 237 | } |
| 238 | } |
| 239 | |
| 240 | /// The JSON-RPC error code. |
| 241 | pub fn error_code(e: &Value) -> i64 { |
| 242 | e["code"].as_i64().unwrap() |
| 243 | } |
| 244 | |
| 245 | /// The text content of a tool call or update. |
| 246 | pub fn tool_text(update: &Value) -> String { |
| 247 | update["content"] |
| 248 | .as_array() |
| 249 | .into_iter() |
| 250 | .flatten() |
| 251 | .filter_map(|c| c.pointer("/content/text").and_then(Value::as_str)) |
| 252 | .collect::<Vec<_>>() |
| 253 | .join("\n") |
| 254 | } |
| 255 | |
| 256 | pub fn alive(pid: u32) -> bool { |
| 257 | std::process::Command::new("kill") |
| 258 | .args(["-0", &pid.to_string()]) |
| 259 | .stderr(Stdio::null()) |
| 260 | .status() |
| 261 | .is_ok_and(|s| s.success()) |
| 262 | } |
| 263 | |
| 264 | /// Wait up to five seconds for `pid` to exit. |
| 265 | pub async fn wait_dead(pid: u32) -> bool { |
| 266 | for _ in 0..100 { |
| 267 | if !alive(pid) { |
| 268 | return true; |
| 269 | } |
| 270 | tokio::time::sleep(Duration::from_millis(50)).await; |
| 271 | } |
| 272 | false |
| 273 | } |
| 274 | |
| 275 | pub fn read_pid(file: &Path) -> u32 { |
| 276 | std::fs::read_to_string(file) |
| 277 | .expect("stub pid file") |
| 278 | .trim() |
| 279 | .parse() |
| 280 | .unwrap() |
| 281 | } |
| 282 | |
| 283 | /// An `mcpServers` entry for the test MCP server under `name`. |
| 284 | pub fn stub_server(stub: &Path, name: &str, pid_file: Option<&Path>) -> Value { |
| 285 | let env: Vec<Value> = pid_file |
| 286 | .map(|p| vec![json!({"name": "STUB_PID_FILE", "value": p})]) |
| 287 | .unwrap_or_default(); |
| 288 | json!({"name": name, "command": stub, "args": [], "env": env}) |
| 289 | } |