|
| 1 | +//! Shared subprocess test harness for the buzz-agent ACP integration suites. |
| 2 | +//! |
| 3 | +//! Every integration test file is its own crate, so this module is included |
| 4 | +//! with `mod common;` in each and compiles once per binary — each binary uses |
| 5 | +//! only the subset it needs, hence the module-wide `dead_code` allow. |
| 6 | +//! |
| 7 | +//! It drives a real `buzz-agent` child over the ACP wire against a fake LLM |
| 8 | +//! (`CapturingLlm`, which records each request body) and answers the |
| 9 | +//! `session/request_permission` surface (`approve_permission`, selecting the |
| 10 | +//! offered `allow_once` option by `kind`, never a hardcoded `optionId`). |
| 11 | +
|
| 12 | +#![allow(dead_code)] |
| 13 | + |
| 14 | +use std::collections::VecDeque; |
| 15 | +use std::process::Stdio; |
| 16 | +use std::sync::{Arc, Mutex as StdMutex}; |
| 17 | +use std::time::Duration; |
| 18 | + |
| 19 | +use serde_json::{json, Value}; |
| 20 | +use tokio::io::{AsyncBufReadExt, AsyncReadExt, AsyncWriteExt, BufReader}; |
| 21 | +use tokio::net::TcpListener; |
| 22 | +use tokio::sync::Mutex; |
| 23 | + |
| 24 | +pub struct CapturingLlm { |
| 25 | + pub url: String, |
| 26 | + pub captured: Arc<Mutex<Vec<Value>>>, |
| 27 | +} |
| 28 | + |
| 29 | +pub async fn spawn_capturing_llm(responses: Vec<Value>) -> CapturingLlm { |
| 30 | + spawn_capturing_llm_with_status(responses.into_iter().map(|v| (200u16, v)).collect()).await |
| 31 | +} |
| 32 | + |
| 33 | +/// Like `spawn_capturing_llm` but each canned response carries its own HTTP |
| 34 | +/// status, so a test can serve a real provider rejection (e.g. a context-window |
| 35 | +/// 400) instead of only success bodies. |
| 36 | +pub async fn spawn_capturing_llm_with_status(responses: Vec<(u16, Value)>) -> CapturingLlm { |
| 37 | + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); |
| 38 | + let url = format!("http://{}", listener.local_addr().unwrap()); |
| 39 | + let queue = Arc::new(Mutex::new(VecDeque::from(responses))); |
| 40 | + let captured: Arc<Mutex<Vec<Value>>> = Arc::new(Mutex::new(Vec::new())); |
| 41 | + let cap2 = captured.clone(); |
| 42 | + tokio::spawn(async move { |
| 43 | + loop { |
| 44 | + let (mut sock, _) = match listener.accept().await { |
| 45 | + Ok(p) => p, |
| 46 | + Err(_) => return, |
| 47 | + }; |
| 48 | + let queue = queue.clone(); |
| 49 | + let captured = cap2.clone(); |
| 50 | + tokio::spawn(async move { |
| 51 | + let mut buf = Vec::new(); |
| 52 | + let mut tmp = [0u8; 8192]; |
| 53 | + // Read until headers complete. |
| 54 | + while !buf.windows(4).any(|w| w == b"\r\n\r\n") { |
| 55 | + match sock.read(&mut tmp).await { |
| 56 | + Ok(0) | Err(_) => return, |
| 57 | + Ok(n) => buf.extend_from_slice(&tmp[..n]), |
| 58 | + } |
| 59 | + if buf.len() > 4_000_000 { |
| 60 | + return; |
| 61 | + } |
| 62 | + } |
| 63 | + // Parse Content-Length and read body. |
| 64 | + let header_end = buf.windows(4).position(|w| w == b"\r\n\r\n").unwrap() + 4; |
| 65 | + let headers = &buf[..header_end]; |
| 66 | + let mut body_len = 0usize; |
| 67 | + for line in headers.split(|b| *b == b'\n') { |
| 68 | + let line = std::str::from_utf8(line).unwrap_or(""); |
| 69 | + if let Some(rest) = line.to_ascii_lowercase().strip_prefix("content-length:") { |
| 70 | + body_len = rest.trim().trim_end_matches('\r').parse().unwrap_or(0); |
| 71 | + } |
| 72 | + } |
| 73 | + while buf.len() < header_end + body_len { |
| 74 | + match sock.read(&mut tmp).await { |
| 75 | + Ok(0) | Err(_) => return, |
| 76 | + Ok(n) => buf.extend_from_slice(&tmp[..n]), |
| 77 | + } |
| 78 | + } |
| 79 | + if let Ok(req) = serde_json::from_slice::<Value>(&buf[header_end..]) { |
| 80 | + captured.lock().await.push(req); |
| 81 | + } |
| 82 | + let (status, body) = queue |
| 83 | + .lock() |
| 84 | + .await |
| 85 | + .pop_front() |
| 86 | + .unwrap_or_else(|| (200, json!({ "error": "no canned response" }))); |
| 87 | + let body_s = serde_json::to_string(&body).unwrap(); |
| 88 | + let reason = match status { |
| 89 | + 200 => "OK", |
| 90 | + 400 => "Bad Request", |
| 91 | + _ => "Error", |
| 92 | + }; |
| 93 | + let resp = format!( |
| 94 | + "HTTP/1.1 {status} {reason}\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}", |
| 95 | + body_s.len(), body_s, |
| 96 | + ); |
| 97 | + let _ = sock.write_all(resp.as_bytes()).await; |
| 98 | + let _ = sock.shutdown().await; |
| 99 | + }); |
| 100 | + } |
| 101 | + }); |
| 102 | + CapturingLlm { url, captured } |
| 103 | +} |
| 104 | + |
| 105 | +pub struct Harness { |
| 106 | + child: tokio::process::Child, |
| 107 | + stdin: tokio::process::ChildStdin, |
| 108 | + stdout: BufReader<tokio::process::ChildStdout>, |
| 109 | + stderr: Arc<StdMutex<String>>, |
| 110 | + next_id: i64, |
| 111 | +} |
| 112 | + |
| 113 | +impl Harness { |
| 114 | + pub async fn spawn_with_env(base_url: &str, extra: &[(&str, &str)]) -> Self { |
| 115 | + let bin = env!("CARGO_BIN_EXE_buzz-agent"); |
| 116 | + let mut cmd = tokio::process::Command::new(bin); |
| 117 | + cmd.env("BUZZ_AGENT_PROVIDER", "openai") |
| 118 | + .env("OPENAI_COMPAT_API_KEY", "test") |
| 119 | + .env("OPENAI_COMPAT_MODEL", "fake-model") |
| 120 | + .env("OPENAI_COMPAT_BASE_URL", base_url) |
| 121 | + .env("BUZZ_AGENT_LLM_TIMEOUT_SECS", "5") |
| 122 | + .env("BUZZ_AGENT_TOOL_TIMEOUT_SECS", "5") |
| 123 | + .env("BUZZ_AGENT_MAX_ROUNDS", "8") |
| 124 | + .env("BUZZ_AGENT_MCP_INIT_TIMEOUT_SECS", "2"); |
| 125 | + for (k, v) in extra { |
| 126 | + cmd.env(k, v); |
| 127 | + } |
| 128 | + cmd.stdin(Stdio::piped()) |
| 129 | + .stdout(Stdio::piped()) |
| 130 | + .stderr(Stdio::piped()) |
| 131 | + .kill_on_drop(true); |
| 132 | + let mut child = cmd.spawn().expect("spawn buzz-agent"); |
| 133 | + let stdin = child.stdin.take().unwrap(); |
| 134 | + let stdout = BufReader::new(child.stdout.take().unwrap()); |
| 135 | + let stderr = child.stderr.take().unwrap(); |
| 136 | + let stderr_buf = Arc::new(StdMutex::new(String::new())); |
| 137 | + let stderr_out = Arc::clone(&stderr_buf); |
| 138 | + tokio::spawn(async move { |
| 139 | + let mut reader = BufReader::new(stderr); |
| 140 | + let mut line = String::new(); |
| 141 | + loop { |
| 142 | + line.clear(); |
| 143 | + let n = match reader.read_line(&mut line).await { |
| 144 | + Ok(n) => n, |
| 145 | + Err(_) => break, |
| 146 | + }; |
| 147 | + if n == 0 { |
| 148 | + break; |
| 149 | + } |
| 150 | + if let Ok(mut out) = stderr_out.lock() { |
| 151 | + out.push_str(&line); |
| 152 | + } |
| 153 | + } |
| 154 | + }); |
| 155 | + Self { |
| 156 | + child, |
| 157 | + stdin, |
| 158 | + stdout, |
| 159 | + stderr: stderr_buf, |
| 160 | + next_id: 1, |
| 161 | + } |
| 162 | + } |
| 163 | + |
| 164 | + pub async fn spawn(base_url: &str) -> Self { |
| 165 | + Self::spawn_with_env(base_url, &[]).await |
| 166 | + } |
| 167 | + |
| 168 | + pub async fn send(&mut self, method: &str, params: Value) -> i64 { |
| 169 | + let id = self.next_id; |
| 170 | + self.next_id += 1; |
| 171 | + self.write(json!({ "jsonrpc": "2.0", "id": id, "method": method, "params": params })) |
| 172 | + .await; |
| 173 | + id |
| 174 | + } |
| 175 | + |
| 176 | + pub async fn notify(&mut self, method: &str, params: Value) { |
| 177 | + self.write(json!({ "jsonrpc": "2.0", "method": method, "params": params })) |
| 178 | + .await; |
| 179 | + } |
| 180 | + |
| 181 | + pub async fn write(&mut self, msg: Value) { |
| 182 | + let mut s = serde_json::to_string(&msg).unwrap(); |
| 183 | + s.push('\n'); |
| 184 | + self.stdin.write_all(s.as_bytes()).await.unwrap(); |
| 185 | + self.stdin.flush().await.unwrap(); |
| 186 | + } |
| 187 | + |
| 188 | + pub async fn recv(&mut self) -> Value { |
| 189 | + let mut line = String::new(); |
| 190 | + let n = tokio::time::timeout(Duration::from_secs(15), self.stdout.read_line(&mut line)) |
| 191 | + .await |
| 192 | + .expect("recv timeout") |
| 193 | + .expect("read line"); |
| 194 | + assert!(n > 0, "agent EOF; stderr={}", self.stderr_text()); |
| 195 | + serde_json::from_str(&line).expect("non-JSON line") |
| 196 | + } |
| 197 | + |
| 198 | + pub async fn recv_until<F: FnMut(&Value) -> bool>(&mut self, mut pred: F) -> Value { |
| 199 | + loop { |
| 200 | + let v = self.recv().await; |
| 201 | + if pred(&v) { |
| 202 | + return v; |
| 203 | + } |
| 204 | + } |
| 205 | + } |
| 206 | + |
| 207 | + /// Like `recv_until`, but auto-approves any `session/request_permission` |
| 208 | + /// seen while waiting. Tests that exercise tool execution, not the |
| 209 | + /// permission boundary (that lives in `permission_boundary.rs`), must |
| 210 | + /// approve a model-issued tool call so it reaches the server. |
| 211 | + pub async fn recv_until_approving<F: FnMut(&Value) -> bool>(&mut self, mut pred: F) -> Value { |
| 212 | + loop { |
| 213 | + let v = self.recv().await; |
| 214 | + if v.get("method") == Some(&json!("session/request_permission")) { |
| 215 | + let resp = approve_permission(&v); |
| 216 | + self.write(resp).await; |
| 217 | + continue; |
| 218 | + } |
| 219 | + if pred(&v) { |
| 220 | + return v; |
| 221 | + } |
| 222 | + } |
| 223 | + } |
| 224 | + |
| 225 | + pub async fn shutdown(mut self) { |
| 226 | + drop(self.stdin); |
| 227 | + let _ = tokio::time::timeout(Duration::from_secs(2), self.child.wait()).await; |
| 228 | + let _ = self.child.start_kill(); |
| 229 | + } |
| 230 | + |
| 231 | + pub fn stderr_text(&self) -> String { |
| 232 | + self.stderr.lock().map(|s| s.clone()).unwrap_or_default() |
| 233 | + } |
| 234 | +} |
| 235 | + |
| 236 | +pub fn openai_text(content: &str) -> Value { |
| 237 | + json!({ |
| 238 | + "id": "cc-1", "object": "chat.completion", "model": "fake-model", |
| 239 | + "choices": [{ |
| 240 | + "index": 0, |
| 241 | + "message": { "role": "assistant", "content": content }, |
| 242 | + "finish_reason": "stop", |
| 243 | + }], |
| 244 | + }) |
| 245 | +} |
| 246 | + |
| 247 | +pub fn openai_tool_call(id: &str, name: &str, args: Value) -> Value { |
| 248 | + json!({ |
| 249 | + "id": "cc-2", "object": "chat.completion", "model": "fake-model", |
| 250 | + "choices": [{ |
| 251 | + "index": 0, |
| 252 | + "message": { |
| 253 | + "role": "assistant", "content": null, |
| 254 | + "tool_calls": [{ |
| 255 | + "id": id, "type": "function", |
| 256 | + "function": { "name": name, "arguments": args.to_string() }, |
| 257 | + }], |
| 258 | + }, |
| 259 | + "finish_reason": "tool_calls", |
| 260 | + }], |
| 261 | + }) |
| 262 | +} |
| 263 | + |
| 264 | +/// Select the offered option whose `kind == "allow_once"` and return the |
| 265 | +/// `session/request_permission` response. Mirrors buzz-acp's answering side, |
| 266 | +/// which selects by `kind`, never by a hardcoded `optionId`. Centralizing this |
| 267 | +/// means a future option-id rename can't silently turn allow into a denial. |
| 268 | +pub fn approve_permission(request: &Value) -> Value { |
| 269 | + let option_id = request["params"]["options"] |
| 270 | + .as_array() |
| 271 | + .and_then(|opts| opts.iter().find(|o| o["kind"] == "allow_once")) |
| 272 | + .and_then(|o| o["optionId"].as_str()) |
| 273 | + .expect("request must offer an allow_once option"); |
| 274 | + json!({ |
| 275 | + "jsonrpc": "2.0", |
| 276 | + "id": request["id"], |
| 277 | + "result": { "outcome": { "outcome": "selected", "optionId": option_id } }, |
| 278 | + }) |
| 279 | +} |
0 commit comments