sausage/cli/tests/agent.rs

//! 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_message, 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<String, String>,
    body: Value,
}

enum ReplyPlan {
    Send(Value),
    Hold(Value, oneshot::Receiver<()>),
}

struct FakeServer {
    base_url: String,
    plans: mpsc::Sender<ReplyPlan>,
    requests: mpsc::Receiver<HttpRequest>,
    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::<usize>().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<u8>,
    mut offset: usize,
) -> Vec<u8> {
    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<PathBuf> = 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<Envelope>) {
    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:?}"
    );
    let subscribed = request(
        &host,
        &mut mailbox,
        json!({"type": "subscribe", "actor": "agent"}),
    )
    .await;
    assert_eq!(subscribed.body["type"], "subscribed");
    (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_published_text(mailbox: &mut mpsc::Receiver<Envelope>, text: &str) {
    let message = receive_message(mailbox, Duration::from_secs(20)).await;
    assert_eq!(message.from, "agent");
    assert_eq!(message.to, "console");
    assert_eq!(message.reply_to, None);
    assert_eq!(message.body, json!({"type": "text", "text": text}));
}

fn history(world: &World) -> Vec<Value> {
    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::<Value>(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;
        send_text(&host, "say hello");
        api.check_request(&fake.next_request().await);
        assert_published_text(&mut mailbox, "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;
        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_published_text(&mut mailbox, "the actor list is ready").await;

        host.stop().await;
        let history = history(&world);
        assert_eq!(
            history
                .iter()
                .map(|entry| entry["kind"].as_str().unwrap())
                .collect::<Vec<_>>(),
            ["user", "tool_call", "tool_result", "assistant"]
        );
        assert_eq!(history[1]["id"], "list-1");
        assert_eq!(history[2]["id"], "list-1");
    }
}

#[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::<Value>(
                &std::fs::read(world.data_dir("agent").join("turn.json")).unwrap()
            )
            .unwrap(),
            json!(true)
        );
        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 started = receive_reply(&mut mailbox, &start_id, Duration::from_secs(20)).await;
        assert_eq!(started.body["type"], "started");
        assert_published_text(&mut mailbox, "resumed answer").await;

        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;
    }
}