475 lines
16 KiB
Rust
475 lines
16 KiB
Rust
//! A scripted stand-in for `llama-server`, for tests. Do not edit.
|
|
//!
|
|
//! It listens on a Unix socket in a temporary directory. Each path has a list of replies that are
|
|
//! served in order; the last one repeats. A reply is raw bytes, normally a response recorded from
|
|
//! the real server (`tests/fixtures/http/*.http`), and can be delayed, sent in small pieces, cut
|
|
//! short, or left hanging. A path can instead compute each reply from its request, as the fake
|
|
//! tokenizer does. Every request is recorded.
|
|
|
|
#![allow(dead_code)] // each test file uses a different part of this module
|
|
|
|
use std::collections::VecDeque;
|
|
use std::io::{Read, Write};
|
|
use std::os::unix::net::{UnixListener, UnixStream};
|
|
use std::path::{Path, PathBuf};
|
|
use std::sync::atomic::{AtomicU32, Ordering};
|
|
use std::sync::{Arc, Mutex};
|
|
use std::thread;
|
|
use std::time::Duration;
|
|
|
|
static NEXT: AtomicU32 = AtomicU32::new(0);
|
|
|
|
pub fn fixture_path(kind: &str, name: &str) -> PathBuf {
|
|
Path::new(env!("CARGO_MANIFEST_DIR"))
|
|
.join("tests/fixtures")
|
|
.join(kind)
|
|
.join(name)
|
|
}
|
|
|
|
pub fn fixture_bytes(kind: &str, name: &str) -> Vec<u8> {
|
|
let path = fixture_path(kind, name);
|
|
std::fs::read(&path).unwrap_or_else(|e| panic!("{}: {e}", path.display()))
|
|
}
|
|
|
|
/// A config that points at `socket`, with every limit short enough for a test.
|
|
pub fn test_config(socket: &Path) -> loopd::config::Config {
|
|
let text = format!(
|
|
r#"
|
|
[infer]
|
|
socket = "{}"
|
|
model = "test-model"
|
|
[slots]
|
|
main = 0
|
|
background = 1
|
|
[expect]
|
|
template_sha256 = "f55f52930aa8bf44ab5cb85f99370fcc3c56e9a85640b812086d5330bce5d86b"
|
|
n_ctx = 131072
|
|
slots = 2
|
|
[limits]
|
|
poll_ms = 40
|
|
busy_wait_ms = 400
|
|
load_wait_ms = 300
|
|
idle_grace_ms = 200
|
|
liveness_ms = 150
|
|
retry_backoff_ms = [10, 20, 30]
|
|
retry_window_ms = 5000
|
|
"#,
|
|
socket.display()
|
|
);
|
|
loopd::config::Config::parse(&text).unwrap()
|
|
}
|
|
|
|
/// A temporary `BOXMAKER_HOME` with a `system.md` beside a `config.toml`, for session tests.
|
|
pub struct Home {
|
|
pub dir: PathBuf,
|
|
}
|
|
|
|
impl Home {
|
|
pub fn new() -> Home {
|
|
let n = NEXT.fetch_add(1, Ordering::SeqCst);
|
|
let dir = std::env::temp_dir().join(format!("loopd-home-{}-{n}", std::process::id()));
|
|
let _ = std::fs::remove_dir_all(&dir);
|
|
std::fs::create_dir_all(&dir).unwrap();
|
|
std::fs::write(dir.join("system.md"), "You are Boxmaker, a test agent.\n").unwrap();
|
|
Home { dir }
|
|
}
|
|
|
|
/// A config for `socket` whose home, system prompt and channel socket are all under this
|
|
/// directory, with the fast test limits.
|
|
pub fn config(&self, socket: &Path) -> loopd::config::Config {
|
|
let mut cfg = test_config(socket);
|
|
cfg.paths.home = self.dir.clone();
|
|
cfg.baseline.system = self.dir.join("system.md");
|
|
cfg.channel.socket = self.dir.join("loop.sock");
|
|
cfg
|
|
}
|
|
|
|
pub fn write(&self, relative: &str, text: &str) {
|
|
let path = self.dir.join(relative);
|
|
std::fs::create_dir_all(path.parent().unwrap()).unwrap();
|
|
std::fs::write(path, text).unwrap();
|
|
}
|
|
|
|
pub fn read(&self, relative: &str) -> String {
|
|
std::fs::read_to_string(self.dir.join(relative)).unwrap()
|
|
}
|
|
|
|
/// The records of a session's log, epoch 0.
|
|
pub fn records(&self, session: &str) -> Vec<bxcore::LogRecord> {
|
|
let text = self.read(&format!("sessions/{session}/0.jsonl"));
|
|
text.lines()
|
|
.map(|l| serde_json::from_str(l).unwrap())
|
|
.collect()
|
|
}
|
|
}
|
|
|
|
/// A tool port that answers from a script and records what it was asked.
|
|
pub struct ScriptedPort {
|
|
replies: Mutex<VecDeque<bxcore::ToolResponse>>,
|
|
calls: Mutex<Vec<bxcore::ToolRequest>>,
|
|
}
|
|
|
|
impl ScriptedPort {
|
|
/// Replies are given in order; when they run out, every call gets `fallback`.
|
|
pub fn new(replies: Vec<bxcore::ToolResponse>) -> ScriptedPort {
|
|
ScriptedPort {
|
|
replies: Mutex::new(replies.into()),
|
|
calls: Mutex::new(Vec::new()),
|
|
}
|
|
}
|
|
|
|
pub fn calls(&self) -> Vec<bxcore::ToolRequest> {
|
|
self.calls.lock().unwrap().clone()
|
|
}
|
|
}
|
|
|
|
pub fn ok_result(content: &str) -> bxcore::ToolResponse {
|
|
bxcore::ToolResponse::Result {
|
|
content: content.to_string(),
|
|
class: bxcore::DataClass::Private,
|
|
untrusted: true,
|
|
truncated: false,
|
|
}
|
|
}
|
|
|
|
pub fn pending(approval: u64, expires: &str) -> bxcore::ToolResponse {
|
|
bxcore::ToolResponse::PendingApproval {
|
|
approval,
|
|
expires: bxcore::Timestamp::parse(expires).unwrap(),
|
|
}
|
|
}
|
|
|
|
impl loopd::tools::ToolPort for ScriptedPort {
|
|
/// A scripted `PendingApproval` is reported through `on_pending`, as the real port does with
|
|
/// the broker's pending frame, and the reply after it is the answer. Use
|
|
/// `ReturnsPendingPort` for a port that breaks the rule and returns one.
|
|
fn call(
|
|
&self,
|
|
request: &bxcore::ToolRequest,
|
|
on_pending: &mut dyn FnMut(&loopd::tools::Pending),
|
|
) -> bxcore::ToolResponse {
|
|
self.calls.lock().unwrap().push(request.clone());
|
|
let mut replies = self.replies.lock().unwrap();
|
|
let mut reply = replies.pop_front();
|
|
if let Some(bxcore::ToolResponse::PendingApproval { approval, expires }) = reply {
|
|
on_pending(&loopd::tools::Pending { approval, expires });
|
|
reply = replies.pop_front();
|
|
}
|
|
reply.unwrap_or_else(|| ok_result("scripted"))
|
|
}
|
|
}
|
|
|
|
/// A port that breaks the rule: it returns a pending frame as its final answer.
|
|
pub struct ReturnsPendingPort;
|
|
|
|
impl loopd::tools::ToolPort for ReturnsPendingPort {
|
|
fn call(
|
|
&self,
|
|
_request: &bxcore::ToolRequest,
|
|
_on_pending: &mut dyn FnMut(&loopd::tools::Pending),
|
|
) -> bxcore::ToolResponse {
|
|
pending(1, "2026-09-18T12:00:00.000Z")
|
|
}
|
|
}
|
|
|
|
/// The fake `/tokenize`, for `FakeServer::respond`: the text of the request's `content` is one
|
|
/// token when it is one of the model's tags (`bxcore::sanitize::TAGS`), the tag's id, and
|
|
/// otherwise one token per byte, the byte's value. A body without a string `content` gets a 400.
|
|
pub fn tokenizer(request: &Recorded) -> Reply {
|
|
let body: serde_json::Value = serde_json::from_slice(&request.body).unwrap_or_default();
|
|
let Some(content) = body.get("content").and_then(|c| c.as_str()) else {
|
|
return Reply::json(400, r#"{"error":"no content"}"#);
|
|
};
|
|
let ids: Vec<u32> = match bxcore::sanitize::TAGS.iter().find(|t| t.text == content) {
|
|
Some(tag) => vec![tag.id],
|
|
None => content.bytes().map(u32::from).collect(),
|
|
};
|
|
tokens(&ids)
|
|
}
|
|
|
|
/// A `/tokenize` answer as llama-server gives one without `with_pieces`: `{"tokens": [ids]}`.
|
|
pub fn tokens(ids: &[u32]) -> Reply {
|
|
Reply::json(200, &serde_json::json!({ "tokens": ids }).to_string())
|
|
}
|
|
|
|
/// What an independent script derived from a recorded stream: `tests/fixtures/expected/`.
|
|
pub fn expected(name: &str) -> serde_json::Value {
|
|
serde_json::from_slice(&fixture_bytes("expected", &format!("{name}.json"))).unwrap()
|
|
}
|
|
|
|
#[derive(Clone)]
|
|
pub struct Reply {
|
|
bytes: Vec<u8>,
|
|
head_delay_ms: u64,
|
|
piece: usize,
|
|
piece_delay_ms: u64,
|
|
stop_after: Option<usize>,
|
|
hang_ms: u64,
|
|
}
|
|
|
|
impl Reply {
|
|
/// Exactly these bytes, then close.
|
|
pub fn raw(bytes: impl Into<Vec<u8>>) -> Reply {
|
|
Reply {
|
|
bytes: bytes.into(),
|
|
head_delay_ms: 0,
|
|
piece: usize::MAX,
|
|
piece_delay_ms: 0,
|
|
stop_after: None,
|
|
hang_ms: 0,
|
|
}
|
|
}
|
|
|
|
/// A response recorded from the real server: `tests/fixtures/http/<name>.http`.
|
|
pub fn fixture(name: &str) -> Reply {
|
|
Reply::raw(fixture_bytes("http", &format!("{name}.http")))
|
|
}
|
|
|
|
/// A small JSON response with a content length.
|
|
pub fn json(status: u16, body: &str) -> Reply {
|
|
Reply::raw(format!(
|
|
"HTTP/1.1 {status} X\r\nContent-Type: application/json\r\nContent-Length: {}\r\n\r\n{body}",
|
|
body.len()
|
|
))
|
|
}
|
|
|
|
/// Wait this long before sending the first byte, as a queued request does.
|
|
pub fn head_delay(mut self, ms: u64) -> Reply {
|
|
self.head_delay_ms = ms;
|
|
self
|
|
}
|
|
|
|
/// Send `piece` bytes at a time, waiting `delay_ms` before each piece after the first.
|
|
pub fn trickle(mut self, piece: usize, delay_ms: u64) -> Reply {
|
|
self.piece = piece.max(1);
|
|
self.piece_delay_ms = delay_ms;
|
|
self
|
|
}
|
|
|
|
/// Send only the first `bytes` bytes, then close, as a server that dies does.
|
|
pub fn cut_after(mut self, bytes: usize) -> Reply {
|
|
self.stop_after = Some(bytes);
|
|
self
|
|
}
|
|
|
|
/// Send only the first `bytes` bytes, then stay silent for `ms` before closing.
|
|
pub fn hang_after(mut self, bytes: usize, ms: u64) -> Reply {
|
|
self.stop_after = Some(bytes);
|
|
self.hang_ms = ms;
|
|
self
|
|
}
|
|
|
|
/// The offset just after the `n`th `data:` line of the body, for use with `cut_after`.
|
|
pub fn offset_after_events(&self, n: usize) -> usize {
|
|
let mut seen = 0;
|
|
let mut at = 0;
|
|
while let Some(found) = find(&self.bytes[at..], b"\n\n") {
|
|
at += found + 2;
|
|
seen += 1;
|
|
if seen == n {
|
|
return at;
|
|
}
|
|
}
|
|
panic!("the reply has only {seen} events");
|
|
}
|
|
}
|
|
|
|
#[derive(Debug, Clone)]
|
|
pub struct Recorded {
|
|
pub method: String,
|
|
/// Path and query as sent.
|
|
pub target: String,
|
|
/// Header lines as sent, without the request line.
|
|
pub headers: Vec<String>,
|
|
pub body: Vec<u8>,
|
|
}
|
|
|
|
impl Recorded {
|
|
pub fn path(&self) -> &str {
|
|
self.target.split('?').next().unwrap_or("")
|
|
}
|
|
|
|
pub fn json(&self) -> serde_json::Value {
|
|
serde_json::from_slice(&self.body).expect("request body is JSON")
|
|
}
|
|
}
|
|
|
|
/// A reply computed from its request.
|
|
type Responder = Arc<dyn Fn(&Recorded) -> Reply + Send + Sync>;
|
|
|
|
struct State {
|
|
routes: Mutex<Vec<(String, VecDeque<Reply>)>>,
|
|
responders: Mutex<Vec<(String, Responder)>>,
|
|
requests: Mutex<Vec<Recorded>>,
|
|
}
|
|
|
|
pub struct FakeServer {
|
|
pub socket: PathBuf,
|
|
state: Arc<State>,
|
|
dir: PathBuf,
|
|
}
|
|
|
|
impl FakeServer {
|
|
pub fn start() -> FakeServer {
|
|
let n = NEXT.fetch_add(1, Ordering::SeqCst);
|
|
let dir = std::env::temp_dir().join(format!("loopd-fake-{}-{n}", std::process::id()));
|
|
let _ = std::fs::remove_dir_all(&dir);
|
|
std::fs::create_dir_all(&dir).unwrap();
|
|
let socket = dir.join("infer.sock");
|
|
let _ = std::fs::remove_file(&socket);
|
|
let listener = UnixListener::bind(&socket).unwrap();
|
|
let state = Arc::new(State {
|
|
routes: Mutex::new(Vec::new()),
|
|
responders: Mutex::new(Vec::new()),
|
|
requests: Mutex::new(Vec::new()),
|
|
});
|
|
let accept_state = Arc::clone(&state);
|
|
thread::spawn(move || {
|
|
for stream in listener.incoming() {
|
|
let Ok(stream) = stream else { break };
|
|
let state = Arc::clone(&accept_state);
|
|
thread::spawn(move || serve(stream, &state));
|
|
}
|
|
});
|
|
FakeServer { socket, state, dir }
|
|
}
|
|
|
|
/// Replies for `path` (the query is ignored), served in order. The last one repeats. They
|
|
/// replace a responder for `path`.
|
|
pub fn route(&self, path: &str, replies: Vec<Reply>) {
|
|
assert!(!replies.is_empty());
|
|
self.state
|
|
.responders
|
|
.lock()
|
|
.unwrap()
|
|
.retain(|(p, _)| p != path);
|
|
let mut routes = self.state.routes.lock().unwrap();
|
|
routes.retain(|(p, _)| p != path);
|
|
routes.push((path.to_string(), replies.into()));
|
|
}
|
|
|
|
/// For each request to `path` (the query is ignored), the reply `reply` computes from it. It
|
|
/// replaces the replies `route` gave `path`.
|
|
pub fn respond(&self, path: &str, reply: impl Fn(&Recorded) -> Reply + Send + Sync + 'static) {
|
|
self.state.routes.lock().unwrap().retain(|(p, _)| p != path);
|
|
let mut responders = self.state.responders.lock().unwrap();
|
|
responders.retain(|(p, _)| p != path);
|
|
responders.push((path.to_string(), Arc::new(reply)));
|
|
}
|
|
|
|
pub fn requests(&self) -> Vec<Recorded> {
|
|
self.state.requests.lock().unwrap().clone()
|
|
}
|
|
|
|
pub fn requests_to(&self, path: &str) -> Vec<Recorded> {
|
|
self.requests()
|
|
.into_iter()
|
|
.filter(|r| r.path() == path)
|
|
.collect()
|
|
}
|
|
}
|
|
|
|
fn find(haystack: &[u8], needle: &[u8]) -> Option<usize> {
|
|
haystack.windows(needle.len()).position(|w| w == needle)
|
|
}
|
|
|
|
fn read_request(stream: &mut UnixStream) -> Option<Recorded> {
|
|
let mut buf = Vec::new();
|
|
let mut chunk = [0u8; 4096];
|
|
let head_end = loop {
|
|
if let Some(end) = find(&buf, b"\r\n\r\n") {
|
|
break end;
|
|
}
|
|
match stream.read(&mut chunk) {
|
|
Ok(0) | Err(_) => return None,
|
|
Ok(n) => buf.extend_from_slice(&chunk[..n]),
|
|
}
|
|
};
|
|
let head = String::from_utf8_lossy(&buf[..head_end]).to_string();
|
|
let mut lines = head.split("\r\n");
|
|
let mut request_line = lines.next()?.split(' ');
|
|
let method = request_line.next()?.to_string();
|
|
let target = request_line.next()?.to_string();
|
|
let headers: Vec<String> = lines.map(str::to_string).collect();
|
|
let length = headers
|
|
.iter()
|
|
.filter_map(|h| h.split_once(':'))
|
|
.find(|(name, _)| name.eq_ignore_ascii_case("content-length"))
|
|
.and_then(|(_, v)| v.trim().parse::<usize>().ok())
|
|
.unwrap_or(0);
|
|
let mut body = buf[head_end + 4..].to_vec();
|
|
while body.len() < length {
|
|
match stream.read(&mut chunk) {
|
|
Ok(0) | Err(_) => return None,
|
|
Ok(n) => body.extend_from_slice(&chunk[..n]),
|
|
}
|
|
}
|
|
Some(Recorded {
|
|
method,
|
|
target,
|
|
headers,
|
|
body,
|
|
})
|
|
}
|
|
|
|
fn serve(mut stream: UnixStream, state: &State) {
|
|
let Some(request) = read_request(&mut stream) else {
|
|
return;
|
|
};
|
|
let path = request.path().to_string();
|
|
let responder = state
|
|
.responders
|
|
.lock()
|
|
.unwrap()
|
|
.iter()
|
|
.find(|(p, _)| *p == path)
|
|
.map(|(_, reply)| Arc::clone(reply));
|
|
let computed = responder.map(|reply| reply(&request));
|
|
state.requests.lock().unwrap().push(request);
|
|
let reply = computed.or_else(|| {
|
|
let mut routes = state.routes.lock().unwrap();
|
|
match routes.iter_mut().find(|(p, _)| *p == path) {
|
|
Some((_, replies)) if replies.len() > 1 => replies.pop_front(),
|
|
Some((_, replies)) => replies.front().cloned(),
|
|
None => None,
|
|
}
|
|
});
|
|
let reply =
|
|
reply.unwrap_or_else(|| Reply::json(404, r#"{"error":"no route in the fake server"}"#));
|
|
thread::sleep(Duration::from_millis(reply.head_delay_ms));
|
|
let end = reply
|
|
.stop_after
|
|
.unwrap_or(reply.bytes.len())
|
|
.min(reply.bytes.len());
|
|
for (i, piece) in reply.bytes[..end].chunks(reply.piece).enumerate() {
|
|
if i > 0 {
|
|
thread::sleep(Duration::from_millis(reply.piece_delay_ms));
|
|
}
|
|
if stream.write_all(piece).is_err() {
|
|
return; // the client went away, which some tests do on purpose
|
|
}
|
|
let _ = stream.flush();
|
|
}
|
|
thread::sleep(Duration::from_millis(reply.hang_ms));
|
|
}
|
|
|
|
/// Removes a test's directory, first putting back the permissions a test may have taken away,
|
|
/// or the removal fails (tip T44).
|
|
fn remove_test_dir(dir: &Path) {
|
|
use std::os::unix::fs::PermissionsExt;
|
|
let _ = std::fs::set_permissions(dir, std::fs::Permissions::from_mode(0o700));
|
|
let _ = std::fs::remove_dir_all(dir);
|
|
}
|
|
|
|
impl Drop for Home {
|
|
fn drop(&mut self) {
|
|
remove_test_dir(&self.dir);
|
|
}
|
|
}
|
|
|
|
impl Drop for FakeServer {
|
|
fn drop(&mut self) {
|
|
remove_test_dir(&self.dir);
|
|
}
|
|
}
|