Marcos Paulo
224 lines
8.3 KiB
Rust
224 lines
8.3 KiB
Rust
//! Test-only ACP agent stub for the integration tests (std + `serde_json`).
|
|
//!
|
|
//! One reader loop, no nested stdin locks: an agent request that expects a
|
|
//! client answer (the permission request) is *remembered* in `State` and
|
|
//! resolved when the matching response arrives.
|
|
//!
|
|
//! Modes via `STUB_AGENT_MODE`:
|
|
//! - unset/`normal`: streams two message chunks, a tool call
|
|
//! (pending -> completed) and a permission request, then ends the turn once
|
|
//! the client answers.
|
|
//! - `cancel`: stays silent until `session/cancel`, then reports
|
|
//! `stopReason: cancelled`.
|
|
//! - `garbage`: injects malformed lines before the normal turn.
|
|
//! - `giant`: writes one >4 MiB line before the normal turn.
|
|
//! - `exit`: exits immediately with code 0.
|
|
//! - `silent`: never answers anything (handshake timeout).
|
|
//! - `crash`: streams one chunk, then exits with code 3 mid-prompt.
|
|
//! - `hold`: prints `PID <pid>` on stderr and then blocks on stdin.
|
|
//! - `unsupported`: probes the client with `fs/read_text_file` and reports
|
|
//! the JSON-RPC error code the client answered with.
|
|
|
|
use std::io::{BufRead, Write};
|
|
use std::process::exit;
|
|
|
|
use serde_json::{Value, json};
|
|
|
|
/// What the stub is waiting for between client messages.
|
|
#[derive(Default)]
|
|
struct State {
|
|
/// Id of an agent request the client still has to answer (the permission
|
|
/// request, or the `fs/*` probe of the `unsupported` mode).
|
|
awaiting: Option<Value>,
|
|
/// Id of the `session/prompt` request being served.
|
|
prompt: Option<Value>,
|
|
}
|
|
|
|
fn main() {
|
|
let mode = std::env::var("STUB_AGENT_MODE").unwrap_or_else(|_| "normal".to_string());
|
|
if mode == "exit" {
|
|
exit(0);
|
|
}
|
|
let stdin = std::io::stdin();
|
|
let stdout = std::io::stdout();
|
|
let mut out = stdout.lock();
|
|
if mode == "hold" {
|
|
let _ = writeln!(std::io::stderr(), "PID {}", std::process::id());
|
|
}
|
|
let mut state = State::default();
|
|
let mut line = String::new();
|
|
loop {
|
|
line.clear();
|
|
match stdin.lock().read_line(&mut line) {
|
|
Ok(0) | Err(_) => return,
|
|
Ok(_) => {}
|
|
}
|
|
let Some(object) = parse_object(&line) else { continue };
|
|
if resolve_awaited(&mut state, &mut out, &mode, &object) {
|
|
continue;
|
|
}
|
|
handle_request(&mut state, &mut out, &mode, &object);
|
|
}
|
|
}
|
|
|
|
/// Answers the pending agent request when `object` is the client response.
|
|
fn resolve_awaited(state: &mut State, out: &mut std::io::StdoutLock<'_>, mode: &str, object: &serde_json::Map<String, Value>) -> bool {
|
|
let Some(pending) = state.awaiting.as_ref() else {
|
|
return false;
|
|
};
|
|
if object.get("id") != Some(pending) {
|
|
return false;
|
|
}
|
|
let option_id = object
|
|
.get("result")
|
|
.and_then(|result| result.get("outcome"))
|
|
.and_then(|outcome| outcome.get("optionId"))
|
|
.and_then(Value::as_str)
|
|
.unwrap_or("cancelled")
|
|
.to_string();
|
|
let code = object.get("error").and_then(|error| error.get("code")).and_then(Value::as_i64).unwrap_or(0);
|
|
state.awaiting = None;
|
|
send_json(
|
|
out,
|
|
mode,
|
|
&json!({
|
|
"jsonrpc": "2.0", "method": "session/update",
|
|
"params": {"sessionId": "sess-1", "update": {
|
|
"sessionUpdate": "ack", "optionId": option_id, "code": code
|
|
}}
|
|
}),
|
|
);
|
|
if let Some(prompt_id) = state.prompt.take() {
|
|
send_json(out, mode, &response(&prompt_id, &json!({"stopReason": "end_turn"})));
|
|
}
|
|
true
|
|
}
|
|
|
|
/// Dispatches one client request or notification.
|
|
fn handle_request(state: &mut State, out: &mut std::io::StdoutLock<'_>, mode: &str, object: &serde_json::Map<String, Value>) {
|
|
let id = object.get("id").cloned();
|
|
match object.get("method").and_then(Value::as_str) {
|
|
Some("initialize") if id.is_some() => {
|
|
send_json(out, mode, &response(id.as_ref().unwrap_or(&Value::Null), &json!({"protocolVersion": 1})));
|
|
}
|
|
Some("session/new") if id.is_some() => {
|
|
send_json(out, mode, &response(id.as_ref().unwrap_or(&Value::Null), &json!({"sessionId": "sess-1"})));
|
|
}
|
|
Some("session/prompt") if id.is_some() => {
|
|
state.prompt = id;
|
|
start_prompt(out, state, mode);
|
|
}
|
|
Some("session/cancel") => {
|
|
state.awaiting = None;
|
|
send_json(out, mode, &update_json("agent_message_chunk", "cancelled"));
|
|
if let Some(prompt_id) = state.prompt.take() {
|
|
send_json(out, mode, &response(&prompt_id, &json!({"stopReason": "cancelled"})));
|
|
}
|
|
}
|
|
_ => {}
|
|
}
|
|
}
|
|
|
|
/// Starts one turn: emits the mode-specific prelude and, in the default
|
|
/// modes, the streaming updates plus the permission request.
|
|
fn start_prompt(out: &mut std::io::StdoutLock<'_>, state: &mut State, mode: &str) {
|
|
match mode {
|
|
"silent" | "cancel" | "hold" => return,
|
|
"unsupported" => {
|
|
send_json(out, mode, &update_json("agent_message_chunk", "probing"));
|
|
state.awaiting = Some(Value::String("fs-1".to_string()));
|
|
send_json(
|
|
out,
|
|
mode,
|
|
&json!({
|
|
"jsonrpc": "2.0", "id": "fs-1", "method": "fs/read_text_file",
|
|
"params": {"path": "notes/a.md"}
|
|
}),
|
|
);
|
|
return;
|
|
}
|
|
"garbage" => {
|
|
let _ = writeln!(out, "not json {{{{");
|
|
let _ = writeln!(out, "[1, 2, 3]");
|
|
let _ = out.flush();
|
|
}
|
|
"giant" => {
|
|
let giant = "x".repeat(5 * 1024 * 1024);
|
|
let _ = writeln!(out, "{giant}");
|
|
let _ = out.flush();
|
|
}
|
|
"crash" => {
|
|
send_json(out, mode, &update_json("agent_message_chunk", "dying"));
|
|
exit(3);
|
|
}
|
|
_ => {}
|
|
}
|
|
send_json(out, mode, &update_json("agent_message_chunk", "Hello "));
|
|
send_json(out, mode, &update_json("agent_message_chunk", "world"));
|
|
send_json(
|
|
out,
|
|
mode,
|
|
&json!({
|
|
"jsonrpc": "2.0", "method": "session/update",
|
|
"params": {"sessionId": "sess-1", "update": {
|
|
"sessionUpdate": "tool_call", "toolCallId": "t1",
|
|
"title": "Read file", "kind": "read", "status": "pending",
|
|
"locations": [{"path": "notes/a.md"}],
|
|
"rawInput": {"path": "notes/a.md"}
|
|
}}
|
|
}),
|
|
);
|
|
send_json(
|
|
out,
|
|
mode,
|
|
&json!({
|
|
"jsonrpc": "2.0", "method": "session/update",
|
|
"params": {"sessionId": "sess-1", "update": {
|
|
"sessionUpdate": "tool_call_update", "toolCallId": "t1",
|
|
"status": "completed", "title": "Read file",
|
|
"content": [{"type": "text", "text": "Read 1 file"}]
|
|
}}
|
|
}),
|
|
);
|
|
let permission = json!({
|
|
"jsonrpc": "2.0", "id": "perm-1", "method": "session/request_permission",
|
|
"params": {"sessionId": "sess-1",
|
|
"toolCall": {"toolCallId": "t1", "title": "Write file", "kind": "edit", "status": "pending"},
|
|
"options": [{"id": "allow", "name": "Allow", "kind": "allow_once"},
|
|
{"id": "deny", "name": "Deny", "kind": "reject_once"}]
|
|
}
|
|
});
|
|
state.awaiting = Some(Value::String("perm-1".to_string()));
|
|
send_json(out, mode, &permission);
|
|
}
|
|
|
|
/// Parses one client line into a JSON object, ignoring anything else.
|
|
fn parse_object(line: &str) -> Option<serde_json::Map<String, Value>> {
|
|
serde_json::from_str::<Value>(line.trim()).ok()?.as_object().cloned()
|
|
}
|
|
|
|
/// Builds a JSON-RPC response object.
|
|
fn response(id: &Value, result: &Value) -> Value {
|
|
json!({"jsonrpc": "2.0", "id": id.clone(), "result": result.clone()})
|
|
}
|
|
|
|
/// Builds a `session/update` notification carrying one content chunk.
|
|
fn update_json(kind: &str, text: &str) -> Value {
|
|
json!({
|
|
"jsonrpc": "2.0", "method": "session/update",
|
|
"params": {"sessionId": "sess-1", "update": {
|
|
"sessionUpdate": kind, "content": {"type": "text", "text": text}
|
|
}}
|
|
})
|
|
}
|
|
|
|
/// Serialises and sends one JSON line (silent mode never answers).
|
|
fn send_json(out: &mut std::io::StdoutLock<'_>, mode: &str, value: &Value) {
|
|
if mode == "silent" {
|
|
return;
|
|
}
|
|
let Ok(line) = serde_json::to_string(value) else { return };
|
|
let _ = writeln!(out, "{line}");
|
|
let _ = out.flush();
|
|
}
|