sausage/cli/tests/agent.rs

//! Run the agent in the real core against scripted OpenAI and Anthropic HTTP servers.
use bytes::Bytes;
use http_body_util::{BodyExt, Full};
use hyper::{Response, server::conn::http1, service::service_fn};
use hyper_util::rt::TokioIo;
use sausage::{Host, world::World};
use sausage_vfs::Layer;
use serde_json::{Value, json};
use std::{convert::Infallible, path::PathBuf, sync::OnceLock, time::Duration};
use tokio::{
    net::TcpListener,
    sync::{mpsc, oneshot},
};

/// The two HTTP contracts supported by the agent.
const APIS: [&str; 2] = ["openai-responses", "anthropic-messages"];
/// A captured request and the script's delayed HTTP response.
type Exchange = (String, Value, oneshot::Sender<(u16, Value)>);
/// The script controls each response, including concurrent requests.
struct Llm {
    url: String,
    api: &'static str,
    requests: mpsc::UnboundedReceiver<Exchange>,
    task: tokio::task::JoinHandle<()>,
}
impl Llm {
    async fn new(api: &'static str) -> Self {
        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
        let url = format!("http://{}", listener.local_addr().unwrap());
        let (tx, requests) = mpsc::unbounded_channel::<Exchange>();
        let task = tokio::spawn(async move {
            loop {
                let (stream, _) = listener.accept().await.unwrap();
                let tx = tx.clone();
                tokio::spawn(async move {
                    let service =
                        service_fn(move |request: hyper::Request<hyper::body::Incoming>| {
                            let tx = tx.clone();
                            async move {
                                let path = request.uri().path().to_string();
                                let bytes = request.into_body().collect().await.unwrap().to_bytes();
                                let body = serde_json::from_slice(&bytes).unwrap();
                                let (reply, answer) = oneshot::channel();
                                tx.send((path, body, reply)).unwrap();
                                let (status, body) = answer.await.unwrap();
                                Ok::<_, Infallible>(
                                    Response::builder()
                                        .status(status)
                                        .header("content-type", "application/json")
                                        .body(Full::new(Bytes::from(body.to_string())))
                                        .unwrap(),
                                )
                            }
                        });
                    http1::Builder::new()
                        .serve_connection(TokioIo::new(stream), service)
                        .await
                        .unwrap();
                });
            }
        });
        Self {
            url,
            api,
            requests,
            task,
        }
    }
    async fn next(&mut self) -> (Value, oneshot::Sender<(u16, Value)>) {
        let (path, body, reply) =
            tokio::time::timeout(Duration::from_secs(60), self.requests.recv())
                .await
                .unwrap()
                .unwrap();
        assert_eq!(
            path,
            if self.api == APIS[0] {
                "/v1/responses"
            } else {
                "/v1/messages"
            }
        );
        assert_eq!(body["model"], "scripted-test");
        (body, reply)
    }
    fn answer(&self, text: &str, calls: &[(&str, Value)]) -> Value {
        if self.api == APIS[0] {
            let mut output =
                vec![json!({"type":"message","content":[{"type":"output_text","text":text}]})];
            for (n, (name, args)) in calls.iter().enumerate() {
                output.push(json!({"type":"function_call","call_id":format!("tool-{n}"),"name":name,"arguments":args.to_string()}));
            }
            json!({"status":"completed","output":output})
        } else {
            let mut content = vec![json!({"type":"text","text":text})];
            for (n, (name, args)) in calls.iter().enumerate() {
                content.push(
                    json!({"type":"tool_use","id":format!("tool-{n}"),"name":name,"input":args}),
                );
            }
            json!({"content":content,"stop_reason":if calls.is_empty(){"end_turn"}else{"tool_use"}})
        }
    }
}
impl Drop for Llm {
    fn drop(&mut self) {
        self.task.abort();
    }
}

fn fixtures() -> &'static PathBuf {
    static TARGET: OnceLock<PathBuf> = OnceLock::new();
    TARGET.get_or_init(|| {
        let root = PathBuf::from(env!("CARGO_MANIFEST_DIR"))
            .parent()
            .unwrap()
            .to_owned();
        let target = PathBuf::from(concat!(env!("CARGO_TARGET_TMPDIR"), "/guests"));
        for args in [
            vec![
                "build",
                "--offline",
                "-p",
                "agent",
                "--target",
                "wasm32-wasip2",
            ],
            vec![
                "build",
                "--offline",
                "--manifest-path",
                "sdk/tests/plugins/probe/Cargo.toml",
                "--target",
                "wasm32-wasip2",
            ],
        ] {
            assert!(
                std::process::Command::new("cargo")
                    .args(args)
                    .current_dir(&root)
                    .env("CARGO_TARGET_DIR", &target)
                    .status()
                    .unwrap()
                    .success()
            );
        }
        target.join("wasm32-wasip2/debug")
    })
}
struct TestWorld(PathBuf);
impl TestWorld {
    fn new(server: &Llm) -> Self {
        let path = std::env::temp_dir().join(format!("sausage-agent-{}", uuid::Uuid::now_v7()));
        let world = World::create(&path).unwrap();
        for (plugin, binary) in [("agent", "agent"), ("probe", "sdk-probe")] {
            std::fs::create_dir_all(world.plugin_dir(plugin)).unwrap();
            std::fs::copy(
                fixtures().join(format!("{binary}.wasm")),
                world.plugin_dir(plugin).join("plugin.wasm"),
            )
            .unwrap();
        }
        std::fs::write(path.join("config/agent.json"),json!({"api":server.api,"base_url":server.url,"model":"scripted-test","token":"fake-token"}).to_string()).unwrap();
        Self(path)
    }
    async fn start(&self) -> Host {
        Host::start_with_cache(
            &self.0,
            std::path::Path::new(concat!(env!("CARGO_TARGET_TMPDIR"), "/cache")),
        )
        .await
        .unwrap()
    }
    fn history(&self) -> String {
        std::fs::read_to_string(self.0.join("state/agent/history.jsonl")).unwrap()
    }
}
impl Drop for TestWorld {
    fn drop(&mut self) {
        let _ = std::fs::remove_dir_all(&self.0);
    }
}
fn chat(host: &Host, text: &str) -> tokio::task::JoinHandle<Value> {
    let host = host.clone();
    let text = text.to_owned();
    tokio::spawn(async move {
        tokio::time::timeout(
            Duration::from_secs(90),
            host.call("agent.chat", json!({"text":text})),
        )
        .await
        .unwrap()
        .unwrap()
    })
}
fn input(body: &Value, api: &str) -> Value {
    body[if api == APIS[0] { "input" } else { "messages" }].clone()
}
fn outputs(body: &Value, api: &str) -> Vec<Value> {
    if api == APIS[0] {
        body["input"]
            .as_array()
            .unwrap()
            .iter()
            .filter(|e| e["type"] == "function_call_output")
            .map(|e| serde_json::from_str(e["output"].as_str().unwrap()).unwrap())
            .collect()
    } else {
        body["messages"]
            .as_array()
            .unwrap()
            .iter()
            .flat_map(|e| e["content"].as_array().unwrap())
            .filter(|e| e["type"] == "tool_result")
            .map(|e| serde_json::from_str(e["content"].as_str().unwrap()).unwrap())
            .collect()
    }
}

/// Both API shapes return final text and retain the complete history after a core restart.
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn final_answers_and_history_survive_restart() {
    for api in APIS {
        let mut server = Llm::new(api).await;
        let world = TestWorld::new(&server);
        let host = world.start().await;
        let turn = chat(&host, "first request alpha");
        let (body, reply) = server.next().await;
        assert!(
            input(&body, api)
                .to_string()
                .contains("first request alpha")
        );
        assert_eq!(body["tools"].as_array().unwrap().len(), 5);
        reply
            .send((200, server.answer("first answer alpha", &[])))
            .unwrap();
        assert_eq!(turn.await.unwrap(), "first answer alpha");
        host.stop().await;
        let prefix = world.history();
        drop(host);
        let host = world.start().await;
        let turn = chat(&host, "second request beta");
        let (body, reply) = server.next().await;
        let text = input(&body, api).to_string();
        for expected in [
            "first request alpha",
            "first answer alpha",
            "second request beta",
        ] {
            assert!(text.contains(expected));
        }
        reply
            .send((200, server.answer("second answer beta", &[])))
            .unwrap();
        assert_eq!(turn.await.unwrap(), "second answer beta");
        assert!(world.history().starts_with(&prefix));
        let entries: Vec<Value> = world
            .history()
            .lines()
            .map(|line| serde_json::from_str(line).unwrap())
            .collect();
        assert_eq!(
            entries
                .iter()
                .map(|e| e["kind"].as_str().unwrap())
                .collect::<Vec<_>>(),
            ["user", "assistant", "user", "assistant"]
        );
        let doc = host.call("agent.rpc.discover", json!({})).await.unwrap();
        assert_eq!(doc["x-sausage-grace"], 600);
        host.stop().await;
    }
}

/// Both API shapes execute every tool, return tool errors, and deliver interim text.
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn tools_and_interim_notifications() {
    for api in APIS {
        let mut server = Llm::new(api).await;
        let world = TestWorld::new(&server);
        let host = world.start().await;
        let mut events = host.subscribe();
        let turn = chat(&host, "exercise the tools");
        let (_, reply) = server.next().await;
        let calls = [
            (
                "call",
                json!({"method":"probe.add","params":{"a":4,"b":5},"detach":true}),
            ),
            ("read_file", json!({"path":"/proc/core/PROTOCOL.md"})),
            (
                "write_file",
                json!({"path":"/state/work/source/note.txt","content":"saved"}),
            ),
            ("list_dir", json!({"path":"/state/work/source"})),
            (
                "copy_dir",
                json!({"from":"/state/work/source","to":"/state/work/copied"}),
            ),
            ("read_file", json!({"path":"/state/work/missing"})),
            ("call", json!({"method":"absent.run","params":{}})),
            ("call", json!({"method":"probe.fail","params":{}})),
        ];
        reply
            .send((200, server.answer("I will inspect the world.", &calls)))
            .unwrap();
        let (body, reply) = server.next().await;
        let results = outputs(&body, api);
        assert_eq!(results.len(), 8);
        assert_eq!(results[0], 9);
        assert!(results[1]["content"].as_str().unwrap().contains("JSON-RPC"));
        assert_eq!(results[2]["written"], "/state/work/source/note.txt");
        assert_eq!(results[3]["entries"][0]["name"], "note.txt");
        assert_eq!(results[4]["copied_files"], 1);
        assert!(results[5]["error"]["message"].is_string());
        assert!(results[6]["error"]["code"].is_number());
        assert_eq!(
            results[7]["error"],
            json!({"code":-32010,"message":"deliberate","data":{"diagnostic":"probe error details"}})
        );
        assert_eq!(
            std::fs::read_to_string(world.0.join("state/work/copied/note.txt")).unwrap(),
            "saved"
        );
        let event = tokio::time::timeout(Duration::from_secs(5), events.recv())
            .await
            .unwrap()
            .unwrap();
        assert_eq!(event["params"]["name"], "text");
        assert_eq!(
            event["params"]["params"]["text"],
            "I will inspect the world."
        );
        reply
            .send((200, server.answer("The tools finished.", &[])))
            .unwrap();
        assert_eq!(turn.await.unwrap(), "The tools finished.");
        host.stop().await;
    }
}

/// The history lock holds the second chat until the complete first turn ends.
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn concurrent_chats_serialize_complete_turns() {
    for api in APIS {
        let mut server = Llm::new(api).await;
        let world = TestWorld::new(&server);
        let host = world.start().await;
        let first = chat(&host, "first concurrent request");
        let (_, reply) = server.next().await;
        let second = chat(&host, "second concurrent request");
        let view = host.filesystem().await.unwrap();
        tokio::time::timeout(Duration::from_secs(10), async {
            loop {
                let file = view
                    .open("/proc/core/locks.json", sausage_vfs::Open::default())
                    .await
                    .unwrap();
                let locks: Value =
                    serde_json::from_slice(&file.read(0, 65536).await.unwrap()).unwrap();
                if locks["/state/agent/history.jsonl"]["waiting"]
                    .as_array()
                    .is_some_and(|waiting| waiting.len() == 1)
                {
                    break;
                }
                tokio::task::yield_now().await;
            }
        })
        .await
        .unwrap();
        assert!(server.requests.try_recv().is_err());
        reply
            .send((200, server.answer("first concurrent answer", &[])))
            .unwrap();
        assert_eq!(first.await.unwrap(), "first concurrent answer");
        let (body, reply) = server.next().await;
        let text = input(&body, api).to_string();
        assert!(text.contains("first concurrent answer"));
        assert!(text.contains("second concurrent request"));
        reply
            .send((200, server.answer("second concurrent answer", &[])))
            .unwrap();
        assert_eq!(second.await.unwrap(), "second concurrent answer");
        host.stop().await;
    }
}

/// Context-size failures discard old complete turns without changing the append-only history.
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn context_window_advances_by_complete_turns() {
    for api in APIS {
        let mut server = Llm::new(api).await;
        let world = TestWorld::new(&server);
        let host = world.start().await;
        let first = chat(&host, "old turn");
        let (_, reply) = server.next().await;
        reply.send((200, server.answer("old answer", &[]))).unwrap();
        first.await.unwrap();
        let second = chat(&host, "current turn");
        let (_, reply) = server.next().await;
        reply
            .send((400, json!({"error":"context_length_exceeded"})))
            .unwrap();
        let (body, reply) = server.next().await;
        let text = input(&body, api).to_string();
        assert!(!text.contains("old turn"));
        assert!(!text.contains("old answer"));
        assert!(text.contains("current turn"));
        reply
            .send((200, server.answer("current answer", &[])))
            .unwrap();
        assert_eq!(second.await.unwrap(), "current answer");
        assert_eq!(world.history().lines().count(), 4);
        host.stop().await;
    }
}

/// Oversized tool outputs are bounded, and failed retries do not commit a new history cursor.
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn oversized_current_turn_preserves_the_previous_cursor() {
    for api in APIS {
        let mut server = Llm::new(api).await;
        let world = TestWorld::new(&server);
        std::fs::write(world.0.join("large.txt"), "x".repeat(200_000)).unwrap();
        let host = world.start().await;
        let first = chat(&host, "old retained turn");
        let (_, reply) = server.next().await;
        reply
            .send((200, server.answer("old retained answer", &[])))
            .unwrap();
        first.await.unwrap();
        let skip = world.0.join("state/agent/history.skip");
        std::fs::write(&skip, "0").unwrap();
        let caller = host.clone();
        let turn = tokio::spawn(async move {
            caller
                .call("agent.chat", json!({"text":"read a large tool result"}))
                .await
        });
        let (_, reply) = server.next().await;
        reply
            .send((
                200,
                server.answer("", &[("read_file", json!({"path":"/large.txt"}))]),
            ))
            .unwrap();
        let (body, reply) = server.next().await;
        let serialized = input(&body, api).to_string();
        assert!(serialized.contains("[tool output truncated"));
        assert!(serialized.len() < 70_000);
        reply
            .send((413, json!({"type":"request_too_large"})))
            .unwrap();
        let (body, reply) = server.next().await;
        assert!(!input(&body, api).to_string().contains("old retained turn"));
        reply
            .send((413, json!({"type":"request_too_large"})))
            .unwrap();
        assert!(turn.await.unwrap().is_err());
        assert_eq!(std::fs::read_to_string(&skip).unwrap(), "0");
        host.stop().await;
    }
}

/// Provider-private items are replayed only for their original API and model.
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn private_items_are_scoped_to_the_model() {
    for api in APIS {
        let mut server = Llm::new(api).await;
        let world = TestWorld::new(&server);
        let config_path = world.0.join("config/agent.json");
        let mut config: Value =
            serde_json::from_slice(&std::fs::read(&config_path).unwrap()).unwrap();
        config["thinking"] = json!("high");
        std::fs::write(config_path, config.to_string()).unwrap();
        std::fs::create_dir_all(world.0.join("state/agent")).unwrap();
        let mut entries = vec![json!({"kind":"user","text":"old turn"})];
        for (model, marker) in [
            ("another-model", "old-private-marker"),
            ("scripted-test", "current-private-marker"),
        ] {
            let item = if api == APIS[0] {
                json!({"type":"reasoning","encrypted_content":marker})
            } else {
                json!({"type":"thinking","thinking":"private","signature":marker})
            };
            entries.push(json!({"kind":"opaque","api":api,"model":model,"item":item}));
        }
        entries.push(json!({"kind":"assistant","text":"old answer"}));
        std::fs::write(
            world.0.join("state/agent/history.jsonl"),
            entries
                .iter()
                .map(|entry| format!("{entry}\n"))
                .collect::<String>(),
        )
        .unwrap();
        let host = world.start().await;
        let turn = chat(&host, "next turn");
        let (body, reply) = server.next().await;
        let text = input(&body, api).to_string();
        assert!(!text.contains("old-private-marker"));
        assert!(text.contains("current-private-marker"));
        reply.send((200, server.answer("done", &[]))).unwrap();
        assert_eq!(turn.await.unwrap(), "done");
        host.stop().await;
    }
}

/// Separate assistant messages retain a separator in the final answer.
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn final_texts_have_a_separator() {
    for api in APIS {
        let mut server = Llm::new(api).await;
        let world = TestWorld::new(&server);
        let host = world.start().await;
        let turn = chat(&host, "two paragraphs");
        let (_, reply) = server.next().await;
        let mut answer = server.answer("first", &[]);
        if api == APIS[0] {
            answer["output"]
                .as_array_mut()
                .unwrap()
                .push(json!({"type":"message","content":[{"type":"output_text","text":"second"}]}));
        } else {
            answer["content"]
                .as_array_mut()
                .unwrap()
                .push(json!({"type":"text","text":"second"}));
        }
        reply.send((200, answer)).unwrap();
        assert_eq!(turn.await.unwrap(), "first\n\nsecond");
        host.stop().await;
    }
}

/// An incomplete response containing reasoning but no text is an error.
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn incomplete_without_text_is_an_error() {
    let mut server = Llm::new(APIS[0]).await;
    let world = TestWorld::new(&server);
    let host = world.start().await;
    let caller = host.clone();
    let turn =
        tokio::spawn(async move { caller.call("agent.chat", json!({"text":"answer"})).await });
    let (_, reply) = server.next().await;
    reply.send((200, json!({"status":"incomplete","output":[{"type":"reasoning","encrypted_content":"private"}]}))).unwrap();
    assert!(
        turn.await.unwrap().unwrap_err()["message"]
            .as_str()
            .unwrap()
            .contains("incomplete response")
    );
    host.stop().await;
}