102 lines
2.7 KiB
Rust
102 lines
2.7 KiB
Rust
//! A fake `loopd` on a Unix socket: it records each turn and answers with the frames the test's
|
|
//! function gives for it. Do not edit.
|
|
|
|
#![allow(dead_code)] // each test file uses a different part of this module
|
|
|
|
use std::io::Write;
|
|
use std::os::unix::net::UnixListener;
|
|
use std::path::Path;
|
|
use std::sync::mpsc;
|
|
|
|
use proto::{
|
|
Envelope, ErrorCode, Message, PROTOCOL_VERSION, Turn, TurnDone, TurnEvent, Usage, WireError,
|
|
read_frame, write_frame,
|
|
};
|
|
|
|
const USAGE: Usage = Usage {
|
|
cache_n: 0,
|
|
prompt_n: 1,
|
|
predicted_n: 1,
|
|
reasoning_tokens: 0,
|
|
thinking_capped: false,
|
|
};
|
|
|
|
/// What the fake sends back for a turn.
|
|
pub enum Reply {
|
|
Frame(Envelope),
|
|
/// Raw bytes, for broken frames.
|
|
Bytes(Vec<u8>),
|
|
/// Stop answering this connection (the caller sees it close).
|
|
Close,
|
|
}
|
|
|
|
pub fn event(e: TurnEvent) -> Reply {
|
|
Reply::Frame(Envelope {
|
|
v: PROTOCOL_VERSION,
|
|
id: 1,
|
|
r#final: false,
|
|
msg: Message::TurnEvent(e),
|
|
})
|
|
}
|
|
|
|
pub fn done(content: &str) -> Reply {
|
|
let msg = Message::TurnDone(TurnDone {
|
|
content: content.to_string(),
|
|
usage: USAGE,
|
|
});
|
|
Reply::Frame(Envelope {
|
|
v: PROTOCOL_VERSION,
|
|
id: 1,
|
|
r#final: true,
|
|
msg,
|
|
})
|
|
}
|
|
|
|
pub fn error(code: ErrorCode, detail: &str) -> Reply {
|
|
let msg = Message::Error(WireError {
|
|
code,
|
|
detail: detail.to_string(),
|
|
});
|
|
Reply::Frame(Envelope {
|
|
v: PROTOCOL_VERSION,
|
|
id: 1,
|
|
r#final: true,
|
|
msg,
|
|
})
|
|
}
|
|
|
|
/// Listen on `socket`; `script(n, &turn)` gives the replies to the n-th turn (from 0).
|
|
pub fn serve_loop<F>(socket: &Path, script: F) -> mpsc::Receiver<Turn>
|
|
where
|
|
F: Fn(usize, &Turn) -> Vec<Reply> + Send + 'static,
|
|
{
|
|
let listener = UnixListener::bind(socket).unwrap();
|
|
let (tx, rx) = mpsc::channel();
|
|
std::thread::spawn(move || {
|
|
for (n, stream) in listener.incoming().enumerate() {
|
|
let Ok(mut stream) = stream else { continue };
|
|
let Ok(Envelope {
|
|
msg: Message::Turn(turn),
|
|
..
|
|
}) = read_frame(&mut stream)
|
|
else {
|
|
continue;
|
|
};
|
|
// Reported before the script runs: a script may wait for the test.
|
|
let _ = tx.send(turn.clone());
|
|
let replies = script(n, &turn);
|
|
for reply in replies {
|
|
let ok = match reply {
|
|
Reply::Frame(env) => write_frame(&mut stream, &env).is_ok(),
|
|
Reply::Bytes(b) => stream.write_all(&b).is_ok(),
|
|
Reply::Close => false,
|
|
};
|
|
if !ok {
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
});
|
|
rx
|
|
}
|