diff --git a/cli/tests/agent.rs b/cli/tests/agent.rs new file mode 100644 index 0000000..a3d77ab --- /dev/null +++ b/cli/tests/agent.rs @@ -0,0 +1,647 @@ +//! Integration tests for the agent's inference and restart behavior. + +#[allow(dead_code)] +mod common; + +use std::collections::HashMap; +use std::path::PathBuf; +use std::process::Command; +use std::sync::OnceLock; +use std::time::Duration; + +use common::{create_world, receive_replies, receive_reply, request, source_with_test, start}; +use sausage::{Envelope, Host, World}; +use serde_json::{Value, json}; +use tokio::io::{AsyncReadExt, AsyncWriteExt}; +use tokio::net::{TcpListener, TcpStream}; +use tokio::sync::{mpsc, oneshot}; +use tokio::task::{JoinHandle, JoinSet}; + +const APIS: [Api; 2] = [Api::OpenAi, Api::Anthropic]; +const TEST_TOKEN: &str = "fake-token"; + +#[derive(Clone, Copy)] +enum Api { + OpenAi, + Anthropic, +} + +impl Api { + fn name(self) -> &'static str { + match self { + Self::OpenAi => "openai-responses", + Self::Anthropic => "anthropic-messages", + } + } + + fn path(self) -> &'static str { + match self { + Self::OpenAi => "/v1/responses", + Self::Anthropic => "/v1/messages", + } + } + + fn text_response(self, text: &str) -> Value { + match self { + Self::OpenAi => { + json!({"output": [{"type": "message", "content": [{"type": "output_text", "text": text}]}]}) + } + Self::Anthropic => json!({"content": [{"type": "text", "text": text}]}), + } + } + + fn tool_response(self, id: &str, name: &str, arguments: &Value) -> Value { + match self { + Self::OpenAi => { + json!({"output": [{"type": "function_call", "call_id": id, "name": name, "arguments": arguments.to_string()}]}) + } + Self::Anthropic => { + json!({"content": [{"type": "tool_use", "id": id, "name": name, "input": arguments}]}) + } + } + } + + fn check_request(self, request: &HttpRequest) { + assert_eq!(request.method, "POST"); + assert_eq!(request.path, self.path()); + assert_eq!(request.headers["content-type"], "application/json"); + assert_eq!(request.body["model"], "fake-model"); + match self { + Self::OpenAi => assert_eq!(request.headers["authorization"], "Bearer fake-token"), + Self::Anthropic => { + assert_eq!(request.headers["x-api-key"], TEST_TOKEN); + assert_eq!(request.headers["anthropic-version"], "2023-06-01"); + } + } + } +} + +struct HttpRequest { + method: String, + path: String, + headers: HashMap, + body: Value, +} + +enum ReplyPlan { + Send(Value), + Hold(Value, oneshot::Receiver<()>), +} + +struct FakeServer { + base_url: String, + plans: mpsc::Sender, + requests: mpsc::Receiver, + task: JoinHandle<()>, +} + +impl FakeServer { + async fn start() -> Self { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let (plans, mut plan_rx) = mpsc::channel(16); + let (request_tx, requests) = mpsc::channel(16); + let task = tokio::spawn(async move { + let mut handlers = JoinSet::new(); + loop { + tokio::select! { + accepted = listener.accept() => { + let (mut stream, _) = accepted.expect("fake listener accept failed"); + let request = read_http_request(&mut stream).await; + let Some(plan) = plan_rx.recv().await else { break }; + if request_tx.send(request).await.is_err() { + break; + } + handlers.spawn(async move { + let (body, release) = match plan { + ReplyPlan::Send(body) => (body, None), + ReplyPlan::Hold(body, release) => (body, Some(release)), + }; + if let Some(release) = release { + let _ = release.await; + } + write_http_response(stream, body).await; + }); + } + Some(_) = handlers.join_next(), if !handlers.is_empty() => {} + } + } + }); + Self { + base_url: format!("http://127.0.0.1:{}", address.port()), + plans, + requests, + task, + } + } + + async fn answer(&self, body: Value) { + self.plans.send(ReplyPlan::Send(body)).await.unwrap(); + } + + async fn hold(&self, body: Value) -> oneshot::Sender<()> { + let (release, wait) = oneshot::channel(); + self.plans.send(ReplyPlan::Hold(body, wait)).await.unwrap(); + release + } + + async fn next_request(&mut self) -> HttpRequest { + tokio::time::timeout(Duration::from_secs(20), self.requests.recv()) + .await + .expect("timed out waiting for fake inference request") + .expect("fake server request channel closed") + } +} + +impl Drop for FakeServer { + fn drop(&mut self) { + self.task.abort(); + } +} + +async fn read_http_request(stream: &mut TcpStream) -> HttpRequest { + let mut received = Vec::new(); + let mut chunk = [0; 4096]; + let header_end = loop { + let count = stream.read(&mut chunk).await.expect("read HTTP request"); + assert_ne!(count, 0, "client closed before HTTP headers arrived"); + received.extend_from_slice(&chunk[..count]); + if let Some(end) = received.windows(4).position(|part| part == b"\r\n\r\n") { + break end; + } + }; + + let header_text = String::from_utf8_lossy(&received[..header_end]); + let mut lines = header_text.split("\r\n"); + let mut request_line = lines.next().unwrap().split_whitespace(); + let method = request_line.next().unwrap().to_string(); + let path = request_line.next().unwrap().to_string(); + let mut headers = HashMap::new(); + for line in lines { + if let Some((name, value)) = line.split_once(':') { + headers.insert(name.trim().to_ascii_lowercase(), value.trim().to_string()); + } + } + + let body_start = header_end + 4; + let request_body = if let Some(length) = headers.get("content-length") { + let body_length = length.parse::().unwrap(); + while received.len() < body_start + body_length { + let count = stream.read(&mut chunk).await.expect("read HTTP body"); + assert_ne!(count, 0, "client closed before HTTP body arrived"); + received.extend_from_slice(&chunk[..count]); + } + received[body_start..body_start + body_length].to_vec() + } else { + assert!( + headers + .get("transfer-encoding") + .is_some_and(|value| value.eq_ignore_ascii_case("chunked")), + "request has neither Content-Length nor chunked transfer encoding: {headers:?}" + ); + read_chunked_body(stream, &mut received, body_start).await + }; + let body = serde_json::from_slice(&request_body).unwrap(); + HttpRequest { + method, + path, + headers, + body, + } +} + +async fn read_chunked_body( + stream: &mut TcpStream, + received: &mut Vec, + mut offset: usize, +) -> Vec { + let mut body = Vec::new(); + loop { + let line_end = loop { + if let Some(end) = received[offset..] + .windows(2) + .position(|part| part == b"\r\n") + { + break offset + end; + } + let mut chunk = [0; 4096]; + let count = stream.read(&mut chunk).await.expect("read chunk size"); + assert_ne!(count, 0, "client closed before chunk size arrived"); + received.extend_from_slice(&chunk[..count]); + }; + let size = String::from_utf8_lossy(&received[offset..line_end]); + let size = size.split(';').next().unwrap().trim(); + let size = usize::from_str_radix(size, 16).unwrap(); + offset = line_end + 2; + if size == 0 { + return body; + } + while received.len() < offset + size + 2 { + let mut chunk = [0; 4096]; + let count = stream.read(&mut chunk).await.expect("read chunk data"); + assert_ne!(count, 0, "client closed before chunk data arrived"); + received.extend_from_slice(&chunk[..count]); + } + body.extend_from_slice(&received[offset..offset + size]); + assert_eq!(&received[offset + size..offset + size + 2], b"\r\n"); + offset += size + 2; + } +} + +async fn write_http_response(mut stream: TcpStream, body: Value) { + let body = serde_json::to_vec(&body).unwrap(); + let head = format!( + "HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n", + body.len() + ); + let mut response = head.into_bytes(); + response.extend_from_slice(&body); + let _ = stream.write_all(&response).await; + let _ = stream.shutdown().await; +} + +fn agent_wasm() -> &'static PathBuf { + static WASM: OnceLock = OnceLock::new(); + WASM.get_or_init(|| { + let root = common::repo_root(); + let output = Command::new("cargo") + .args([ + "build", + "--release", + "-p", + "agent", + "--target", + "wasm32-wasip2", + "--target-dir", + "target/test-actors", + ]) + .current_dir(&root) + .output() + .expect("run cargo to build the agent actor"); + assert!( + output.status.success(), + "building agent failed\nstdout:\n{}\nstderr:\n{}", + String::from_utf8_lossy(&output.stdout), + String::from_utf8_lossy(&output.stderr) + ); + let wasm = root.join("target/test-actors/wasm32-wasip2/release/agent.wasm"); + assert!( + wasm.is_file(), + "agent build did not create {}", + wasm.display() + ); + wasm + }) +} + +async fn start_agent( + api: Api, + fake: &FakeServer, + label: &str, +) -> (World, Host, mpsc::Receiver) { + let world = create_world(label); + let data = world.data_dir("agent"); + std::fs::create_dir_all(&data).unwrap(); + let config = json!({"api": api.name(), "base_url": fake.base_url, "model": "fake-model", "token": TEST_TOKEN}); + std::fs::write( + data.join("config.json"), + serde_json::to_vec(&config).unwrap(), + ) + .unwrap(); + let source = source_with_test( + &world, + "agent-candidate", + json!({"steps": [{"send": {"type": "text", "text": "candidate"}, "expect": {"type": "text"}}]}), + ); + let (host, mut mailbox) = start(&world).await; + let spawned = common::spawn_actor(&host, &mut mailbox, "agent", agent_wasm(), &source).await; + assert_eq!( + spawned.body["type"], "spawned", + "agent spawn failed: {spawned:?}" + ); + (world, host, mailbox) +} + +fn send_text(host: &Host, text: &str) -> String { + host.broker.send( + "console", + Some("agent"), + None, + None, + json!({"type": "text", "text": text}), + ) +} + +async fn assert_answer(mailbox: &mut mpsc::Receiver, id: &str, text: &str) { + let reply = receive_reply(mailbox, id, Duration::from_secs(20)).await; + assert_eq!(reply.to, "console"); + assert_eq!(reply.reply_to.as_deref(), Some(id)); + assert_eq!(reply.body, json!({"type": "text", "text": text})); +} + +fn history(world: &World) -> Vec { + std::fs::read_to_string(world.data_dir("agent").join("history.jsonl")) + .unwrap() + .lines() + .map(|line| serde_json::from_str(line).unwrap()) + .collect() +} + +fn assert_history_tools_have_results(history: &[Value]) { + let calls: Vec<_> = history + .iter() + .filter(|entry| entry["kind"] == "tool_call") + .collect(); + let results: Vec<_> = history + .iter() + .filter(|entry| entry["kind"] == "tool_result") + .collect(); + assert_eq!(calls.len(), results.len()); + for call in calls { + assert!( + results.iter().any(|result| result["id"] == call["id"]), + "tool call has no result: {call}" + ); + } +} + +fn assert_tool_pair( + request: &HttpRequest, + api: Api, + id: &str, + name: &str, + arguments: &Value, +) -> String { + match api { + Api::OpenAi => { + let input = request.body["input"].as_array().unwrap(); + let call_index = input + .iter() + .position(|item| item["type"] == "function_call") + .unwrap(); + let result_index = input + .iter() + .position(|item| item["type"] == "function_call_output") + .unwrap(); + assert!(call_index < result_index); + let call = &input[call_index]; + assert_eq!(call["call_id"], id); + assert_eq!(call["name"], name); + assert_eq!( + serde_json::from_str::(call["arguments"].as_str().unwrap()).unwrap(), + *arguments + ); + let result = &input[result_index]; + assert_eq!(result["call_id"], id); + result["output"].as_str().unwrap().to_string() + } + Api::Anthropic => { + let messages = request.body["messages"].as_array().unwrap(); + let mut call_position = None; + let mut result_position = None; + let mut call = None; + let mut result = None; + for (index, message) in messages.iter().enumerate() { + let Some(blocks) = message["content"].as_array() else { + continue; + }; + for block in blocks { + match block["type"].as_str() { + Some("tool_use") => { + call_position = Some(index); + call = Some(block); + } + Some("tool_result") => { + result_position = Some(index); + result = Some(block); + } + _ => {} + } + } + } + assert!(call_position.unwrap() < result_position.unwrap()); + let call = call.unwrap(); + assert_eq!(call["id"], id); + assert_eq!(call["name"], name); + assert_eq!(call["input"], *arguments); + let result = result.unwrap(); + assert_eq!(result["tool_use_id"], id); + result["content"].as_str().unwrap().to_string() + } + } +} + +#[tokio::test] +async fn text_answers_use_both_inference_apis() { + for api in APIS { + let mut fake = FakeServer::start().await; + fake.answer(api.text_response("hello from the fake model")) + .await; + let (_, host, mut mailbox) = + start_agent(api, &fake, &format!("agent-text-{}", api.name())).await; + let id = send_text(&host, "say hello"); + api.check_request(&fake.next_request().await); + assert_answer(&mut mailbox, &id, "hello from the fake model").await; + host.stop().await; + } +} + +#[tokio::test] +async fn tool_calls_pair_with_results_in_both_inference_apis() { + let args = json!({"path": "/actor"}); + for api in APIS { + let mut fake = FakeServer::start().await; + fake.answer(api.tool_response("list-1", "list_dir", &args)) + .await; + let (world, host, mut mailbox) = + start_agent(api, &fake, &format!("agent-tools-{}", api.name())).await; + let id = send_text(&host, "list actors"); + api.check_request(&fake.next_request().await); + fake.answer(api.text_response("the actor list is ready")) + .await; + let second = fake.next_request().await; + api.check_request(&second); + let output = assert_tool_pair(&second, api, "list-1", "list_dir", &args); + let result: Value = serde_json::from_str(&output).unwrap(); + assert!( + result["entries"] + .as_array() + .unwrap() + .iter() + .any(|entry| entry["name"] == "agent") + ); + assert_answer(&mut mailbox, &id, "the actor list is ready").await; + + let history = history(&world); + assert_eq!( + history + .iter() + .map(|entry| entry["kind"].as_str().unwrap()) + .collect::>(), + ["user", "tool_call", "tool_result", "assistant"] + ); + assert_eq!(history[1]["id"], "list-1"); + assert_eq!(history[2]["id"], "list-1"); + host.stop().await; + } +} + +#[tokio::test] +async fn restart_resumes_a_turn_with_completed_tool_history() { + let args = json!({"path": "/actor"}); + for api in [Api::OpenAi] { + let mut fake = FakeServer::start().await; + fake.answer(api.tool_response("list-1", "list_dir", &args)) + .await; + let (world, host, mut mailbox) = + start_agent(api, &fake, &format!("agent-resume-{}", api.name())).await; + let id = send_text(&host, "list actors before the restart"); + api.check_request(&fake.next_request().await); + let release = fake.hold(api.text_response("discarded response")).await; + let second = fake.next_request().await; + api.check_request(&second); + assert_tool_pair(&second, api, "list-1", "list_dir", &args); + + let stopped = request( + &host, + &mut mailbox, + json!({"type": "stop", "actor": "agent"}), + ) + .await; + assert_eq!(stopped.body["type"], "stopped"); + assert_eq!( + serde_json::from_slice::( + &std::fs::read(world.data_dir("agent").join("turn.json")).unwrap() + ) + .unwrap()["active"]["id"], + id + ); + let _ = release.send(()); + fake.answer(api.text_response("resumed answer")).await; + let start_id = host.broker.send( + "console", + Some("broker"), + None, + None, + json!({"type": "start", "actor": "agent"}), + ); + let resumed = fake.next_request().await; + api.check_request(&resumed); + assert_tool_pair(&resumed, api, "list-1", "list_dir", &args); + let replies = receive_replies( + &mut mailbox, + &[start_id.clone(), id.clone()], + Duration::from_secs(20), + ) + .await; + let started = replies + .iter() + .find(|reply| reply.reply_to.as_deref() == Some(start_id.as_str())) + .unwrap(); + assert_eq!(started.body["type"], "started"); + let answer = replies + .iter() + .find(|reply| reply.reply_to.as_deref() == Some(id.as_str())) + .unwrap(); + assert_eq!(answer.body["text"], "resumed answer"); + + let history = history(&world); + assert_history_tools_have_results(&history); + assert_eq!( + history + .iter() + .filter(|entry| entry["kind"] == "user" + && entry["text"].as_str().unwrap_or_default().contains(&id)) + .count(), + 1 + ); + host.stop().await; + } +} + +#[tokio::test] +async fn queued_text_resumes_after_the_active_turn_across_restart() { + let wait_args = json!({"seconds": 10}); + for api in [Api::Anthropic] { + let mut fake = FakeServer::start().await; + fake.answer(api.tool_response("wait-1", "wait_for_messages", &wait_args)) + .await; + let (world, host, mut mailbox) = + start_agent(api, &fake, &format!("agent-queue-{}", api.name())).await; + let first_id = send_text(&host, "first turn"); + api.check_request(&fake.next_request().await); + let release = fake.hold(api.text_response("discarded response")).await; + let second_id = send_text(&host, "second turn"); + let second_request = fake.next_request().await; + api.check_request(&second_request); + let output = assert_tool_pair( + &second_request, + api, + "wait-1", + "wait_for_messages", + &wait_args, + ); + let queued: Value = serde_json::from_str(&output).unwrap(); + assert!( + queued["messages"] + .as_array() + .unwrap() + .iter() + .any(|message| message["message"]["id"] == second_id) + ); + let turn: Value = serde_json::from_slice( + &std::fs::read(world.data_dir("agent").join("turn.json")).unwrap(), + ) + .unwrap(); + assert_eq!(turn["queue"][0]["id"], second_id); + + let stopped = request( + &host, + &mut mailbox, + json!({"type": "stop", "actor": "agent"}), + ) + .await; + assert_eq!(stopped.body["type"], "stopped"); + let _ = release.send(()); + fake.answer(api.text_response("first answer")).await; + fake.answer(api.text_response("second answer")).await; + let start_id = host.broker.send( + "console", + Some("broker"), + None, + None, + json!({"type": "start", "actor": "agent"}), + ); + let resumed_first = fake.next_request().await; + api.check_request(&resumed_first); + assert!(resumed_first.body.to_string().contains("first turn")); + let next_turn = fake.next_request().await; + api.check_request(&next_turn); + assert!(next_turn.body.to_string().contains("second turn")); + + let replies = receive_replies( + &mut mailbox, + &[start_id.clone(), first_id.clone(), second_id.clone()], + Duration::from_secs(20), + ) + .await; + let started = replies + .iter() + .find(|reply| reply.reply_to.as_deref() == Some(start_id.as_str())) + .unwrap(); + assert_eq!(started.body["type"], "started"); + let first_position = replies + .iter() + .position(|reply| reply.reply_to.as_deref() == Some(first_id.as_str())) + .unwrap(); + let second_position = replies + .iter() + .position(|reply| reply.reply_to.as_deref() == Some(second_id.as_str())) + .unwrap(); + assert!(first_position < second_position); + assert_eq!(replies[first_position].body["text"], "first answer"); + assert_eq!(replies[second_position].body["text"], "second answer"); + assert_history_tools_have_results(&history(&world)); + host.stop().await; + } +}