use std::path::{Path, PathBuf};
use serde::{Deserialize, Serialize};
use tokio::io::{AsyncBufReadExt, AsyncRead, AsyncWrite, AsyncWriteExt, BufReader};
use tokio::sync::mpsc::UnboundedSender;
use tokio::sync::{broadcast, oneshot};
use crate::components::AgentStatus;
use crate::host::{ControlOp, SpawnArgs, WorldEvent};
use leviath_core::interaction::{InteractionRequest, InteractionResponse};
#[cfg(unix)]
mod unix;
#[cfg(unix)]
pub use unix::{
ClientStream, ControlId, ControlListener, ServerStream, bind_control_listener, connect,
control_id, control_id_from_str, is_daemon_running,
};
#[cfg(windows)]
mod windows;
#[cfg(windows)]
pub use windows::{
ClientStream, ControlId, ControlListener, ServerStream, bind_control_listener, connect,
control_id, control_id_from_str, is_daemon_running,
};
const AUTH_REQUIRED: &str = "authentication";
#[derive(Clone)]
pub struct ControlToken(String);
impl std::fmt::Debug for ControlToken {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.write_str("ControlToken(<redacted>)")
}
}
impl ControlToken {
pub fn path(dir: &Path) -> PathBuf {
dir.join("control.token")
}
pub fn pid_path(dir: &Path) -> PathBuf {
dir.join("daemon.pid")
}
pub fn write_pid(dir: &Path) -> std::io::Result<()> {
leviath_sys::write_private(
&Self::pid_path(dir),
std::process::id().to_string().as_bytes(),
)
}
pub fn read_pid(dir: &Path) -> Option<u32> {
std::fs::read_to_string(Self::pid_path(dir))
.ok()?
.trim()
.parse()
.ok()
}
pub fn create(dir: &Path) -> std::io::Result<Self> {
use rand::RngExt as _;
let bytes: [u8; 32] = rand::rng().random();
let token: String = bytes.iter().map(|b| format!("{b:02x}")).collect();
std::fs::create_dir_all(dir)?;
let _ = leviath_sys::secure_dir_perms(dir);
leviath_sys::write_private(&Self::path(dir), token.as_bytes())?;
Ok(Self(token))
}
pub fn load(dir: &Path) -> std::io::Result<Self> {
let token = std::fs::read_to_string(Self::path(dir))?;
Ok(Self(token.trim().to_string()))
}
pub fn matches(&self, presented: &str) -> bool {
leviath_core::constant_time_eq(&self.0, presented)
}
pub fn expose(&self) -> &str {
&self.0
}
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
#[serde(tag = "op", rename_all = "snake_case")]
pub enum ControlRequest {
Authenticate {
token: String,
},
Spawn {
args: Box<SpawnArgs>,
},
Status {
run_id: String,
},
Pause {
run_id: String,
},
Resume {
run_id: String,
},
Cancel {
run_id: String,
},
List,
Message {
agent_id: String,
content: String,
#[serde(default)]
target_region: Option<String>,
},
ListInteractions,
AnswerInteraction {
response: InteractionResponse,
},
CancelInteraction {
request_id: String,
},
Shutdown,
Subscribe,
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
#[serde(tag = "result", rename_all = "snake_case")]
pub enum ControlResponse {
Spawned {
run_id: String,
},
Status {
status: Option<AgentStatus>,
},
Ok {
ok: bool,
},
List {
runs: Vec<(String, AgentStatus)>,
},
Interactions {
interactions: Vec<(String, InteractionRequest)>,
},
Error {
message: String,
},
}
async fn dispatch(req: ControlRequest, op_tx: &UnboundedSender<ControlOp>) -> ControlResponse {
match req {
ControlRequest::Authenticate { .. } => ControlResponse::Ok { ok: true },
ControlRequest::Spawn { args } => {
let (reply, rx) = oneshot::channel();
let _ = op_tx.send(ControlOp::Spawn { args, reply });
match rx.await {
Ok(Ok(run_id)) => ControlResponse::Spawned { run_id },
Ok(Err(message)) => ControlResponse::Error { message },
Err(_) => ControlResponse::Error {
message: "daemon is shutting down".to_string(),
},
}
}
ControlRequest::Status { run_id } => {
let (reply, rx) = oneshot::channel();
let _ = op_tx.send(ControlOp::Status { run_id, reply });
ControlResponse::Status {
status: rx.await.unwrap_or(None),
}
}
ControlRequest::Pause { run_id } => {
let (reply, rx) = oneshot::channel();
let _ = op_tx.send(ControlOp::Pause { run_id, reply });
ControlResponse::Ok {
ok: rx.await.unwrap_or(false),
}
}
ControlRequest::Resume { run_id } => {
let (reply, rx) = oneshot::channel();
let _ = op_tx.send(ControlOp::Resume { run_id, reply });
ControlResponse::Ok {
ok: rx.await.unwrap_or(false),
}
}
ControlRequest::Cancel { run_id } => {
let (reply, rx) = oneshot::channel();
let _ = op_tx.send(ControlOp::Cancel { run_id, reply });
ControlResponse::Ok {
ok: rx.await.unwrap_or(false),
}
}
ControlRequest::List => {
let (reply, rx) = oneshot::channel();
let _ = op_tx.send(ControlOp::List { reply });
ControlResponse::List {
runs: rx.await.unwrap_or_default(),
}
}
ControlRequest::Message {
agent_id,
content,
target_region,
} => {
let (reply, rx) = oneshot::channel();
let _ = op_tx.send(ControlOp::Message {
agent_id,
content,
target_region,
reply,
});
ControlResponse::Ok {
ok: rx.await.unwrap_or(false),
}
}
ControlRequest::ListInteractions => {
let (reply, rx) = oneshot::channel();
let _ = op_tx.send(ControlOp::ListInteractions { reply });
ControlResponse::Interactions {
interactions: rx.await.unwrap_or_default(),
}
}
ControlRequest::AnswerInteraction { response } => {
let (reply, rx) = oneshot::channel();
let _ = op_tx.send(ControlOp::AnswerInteraction { response, reply });
ControlResponse::Ok {
ok: rx.await.unwrap_or(false),
}
}
ControlRequest::CancelInteraction { request_id } => {
let (reply, rx) = oneshot::channel();
let _ = op_tx.send(ControlOp::CancelInteraction { request_id, reply });
ControlResponse::Ok {
ok: rx.await.unwrap_or(false),
}
}
ControlRequest::Shutdown => {
let (reply, rx) = oneshot::channel();
let _ = op_tx.send(ControlOp::Shutdown { reply });
ControlResponse::Ok {
ok: rx.await.unwrap_or(false),
}
}
ControlRequest::Subscribe => ControlResponse::Error {
message: "subscribe is a streaming request, not a single-reply op".to_string(),
},
}
}
async fn stream_events<W>(
write: &mut W,
mut rx: broadcast::Receiver<WorldEvent>,
) -> std::io::Result<()>
where
W: AsyncWrite + Unpin,
{
loop {
match rx.recv().await {
Ok(event) => {
let mut line = serde_json::to_string(&event).expect("WorldEvent serializes");
line.push('\n');
if write.write_all(line.as_bytes()).await.is_err() {
return Ok(()); }
}
Err(broadcast::error::RecvError::Lagged(_)) => continue,
Err(broadcast::error::RecvError::Closed) => return Ok(()),
}
}
}
pub async fn handle_connection<S>(
stream: S,
op_tx: UnboundedSender<ControlOp>,
events: broadcast::Sender<WorldEvent>,
token: Option<ControlToken>,
) -> std::io::Result<()>
where
S: AsyncRead + AsyncWrite + Unpin,
{
let (read_half, mut write_half) = tokio::io::split(stream);
let mut lines = BufReader::new(read_half).lines();
let mut authenticated = token.is_none();
while let Some(line) = lines.next_line().await? {
if line.trim().is_empty() {
continue;
}
if !authenticated {
let refused = match serde_json::from_str::<ControlRequest>(&line) {
Ok(ControlRequest::Authenticate { token: presented }) => {
match token.as_ref().is_some_and(|t| t.matches(&presented)) {
true => {
authenticated = true;
write_line(&mut write_half, &ControlResponse::Ok { ok: true }).await;
continue;
}
false => "authentication failed",
}
}
_ => "authentication required: send an `authenticate` request first",
};
write_line(
&mut write_half,
&ControlResponse::Error {
message: refused.to_string(),
},
)
.await;
return Ok(());
}
let response = match serde_json::from_str::<ControlRequest>(&line) {
Ok(ControlRequest::Subscribe) => {
let rx = events.subscribe();
drop(events);
return stream_events(&mut write_half, rx).await;
}
Ok(ControlRequest::Authenticate { .. }) => ControlResponse::Ok { ok: true },
Ok(req) => dispatch(req, &op_tx).await,
Err(e) => ControlResponse::Error {
message: format!("invalid request: {e}"),
},
};
write_line(&mut write_half, &response).await;
}
Ok(())
}
async fn write_line<W>(write_half: &mut W, response: &ControlResponse)
where
W: AsyncWrite + Unpin,
{
let mut out = serde_json::to_string(response).expect("ControlResponse serializes");
out.push('\n');
let _ = write_half.write_all(out.as_bytes()).await;
}
pub const DEFAULT_CONTROL_TIMEOUT_SECS: u64 = 30;
pub const SPAWN_CONTROL_TIMEOUT_SECS: u64 = 300;
pub fn request_timeout() -> std::time::Duration {
let secs = std::env::var("LEVIATH_CONTROL_TIMEOUT_SECS")
.ok()
.and_then(|v| v.trim().parse::<u64>().ok())
.unwrap_or(DEFAULT_CONTROL_TIMEOUT_SECS);
match secs {
0 => std::time::Duration::MAX,
secs => std::time::Duration::from_secs(secs),
}
}
fn timeout_for(req: &ControlRequest) -> std::time::Duration {
let base = request_timeout();
match req {
ControlRequest::Spawn { .. } if base != std::time::Duration::MAX => {
base.max(std::time::Duration::from_secs(SPAWN_CONTROL_TIMEOUT_SECS))
}
_ => base,
}
}
#[derive(Clone)]
pub struct ControlClient {
id: ControlId,
token: Option<ControlToken>,
token_dir: Option<PathBuf>,
}
impl ControlClient {
pub fn new(id: impl Into<ControlId>) -> Self {
Self {
id: id.into(),
token: None,
token_dir: None,
}
}
pub fn with_token(mut self, token: ControlToken) -> Self {
self.token = Some(token);
self
}
pub fn for_home(id: impl Into<ControlId>, dir: &Path) -> Self {
Self {
id: id.into(),
token: ControlToken::load(dir).ok(),
token_dir: Some(dir.to_path_buf()),
}
}
fn refused(&self) -> std::io::Error {
let detail = match (&self.token, &self.token_dir) {
(None, Some(dir)) => format!(
"no control token was found at {}. If a daemon is running, it was \
started by a different user or before this file existed - restart \
it with `lev daemon restart`.",
ControlToken::path(dir).display()
),
_ => "the daemon refused this client's control token. Restart it with \
`lev daemon restart` to issue a fresh one."
.to_string(),
};
std::io::Error::new(std::io::ErrorKind::PermissionDenied, detail)
}
pub async fn request(&self, req: &ControlRequest) -> std::io::Result<ControlResponse> {
let deadline = timeout_for(req);
tokio::time::timeout(deadline, self.request_uncapped(req))
.await
.unwrap_or_else(|_| {
Err(std::io::Error::new(
std::io::ErrorKind::TimedOut,
format!("the daemon did not respond within {}s", deadline.as_secs()),
))
})
}
async fn request_uncapped(&self, req: &ControlRequest) -> std::io::Result<ControlResponse> {
let stream = connect(&self.id).await?;
let (read_half, mut write_half) = tokio::io::split(stream);
let mut lines = BufReader::new(read_half).lines();
if let Some(token) = &self.token {
let hello = ControlRequest::Authenticate {
token: token.expose().to_string(),
};
let mut line = serde_json::to_string(&hello).expect("ControlRequest serializes");
line.push('\n');
let _ = write_half.write_all(line.as_bytes()).await;
match lines.next_line().await.ok().flatten() {
Some(resp) => match serde_json::from_str::<ControlResponse>(&resp) {
Ok(ControlResponse::Ok { ok: true }) => {}
_ => return Err(self.refused()),
},
None => {
return Err(std::io::Error::new(
std::io::ErrorKind::UnexpectedEof,
"control connection closed during authentication",
));
}
}
}
let mut line = serde_json::to_string(req).expect("ControlRequest serializes");
line.push('\n');
let _ = write_half.write_all(line.as_bytes()).await;
match lines.next_line().await? {
Some(resp_line) => {
let parsed: ControlResponse = serde_json::from_str(&resp_line)
.map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))?;
match &parsed {
ControlResponse::Error { message } if message.starts_with(AUTH_REQUIRED) => {
Err(self.refused())
}
_ => Ok(parsed),
}
}
None => Err(std::io::Error::new(
std::io::ErrorKind::UnexpectedEof,
"control connection closed before a response",
)),
}
}
pub async fn spawn(&self, args: SpawnArgs) -> std::io::Result<ControlResponse> {
self.request(&ControlRequest::Spawn {
args: Box::new(args),
})
.await
}
pub async fn status(&self, run_id: &str) -> std::io::Result<ControlResponse> {
self.request(&ControlRequest::Status {
run_id: run_id.to_string(),
})
.await
}
pub async fn list(&self) -> std::io::Result<ControlResponse> {
self.request(&ControlRequest::List).await
}
pub async fn shutdown(&self) -> std::io::Result<ControlResponse> {
self.request(&ControlRequest::Shutdown).await
}
pub async fn subscribe(&self) -> std::io::Result<WorldEventStream> {
let stream = connect(&self.id).await?;
let (read_half, mut write_half) = tokio::io::split(stream);
let mut line =
serde_json::to_string(&ControlRequest::Subscribe).expect("ControlRequest serializes");
line.push('\n');
let _ = write_half.write_all(line.as_bytes()).await;
Ok(WorldEventStream {
lines: BufReader::new(read_half).lines(),
_write: write_half,
})
}
}
pub struct WorldEventStream {
lines: tokio::io::Lines<BufReader<tokio::io::ReadHalf<ClientStream>>>,
_write: tokio::io::WriteHalf<ClientStream>,
}
impl WorldEventStream {
pub async fn next(&mut self) -> Option<WorldEvent> {
let line = self.lines.next_line().await.ok().flatten()?;
serde_json::from_str(&line).ok()
}
}
#[cfg(test)]
mod tests {
#[test]
fn the_daemon_pid_round_trips_and_a_missing_or_junk_file_is_no_pid() {
let dir = tempfile::tempdir().unwrap();
assert_eq!(
ControlToken::read_pid(dir.path()),
None,
"no file yet is not a pid"
);
ControlToken::write_pid(dir.path()).unwrap();
assert_eq!(
ControlToken::read_pid(dir.path()),
Some(std::process::id()),
"what was written is what comes back"
);
std::fs::write(ControlToken::pid_path(dir.path()), " 4242\n").unwrap();
assert_eq!(ControlToken::read_pid(dir.path()), Some(4242));
std::fs::write(ControlToken::pid_path(dir.path()), "not-a-pid").unwrap();
assert_eq!(ControlToken::read_pid(dir.path()), None);
}
#[test]
fn a_token_round_trips_through_its_file_and_is_owner_only() {
let dir = tempfile::tempdir().unwrap();
let created = ControlToken::create(dir.path()).unwrap();
let loaded = ControlToken::load(dir.path()).unwrap();
assert!(
created.matches(loaded.expose()),
"the same secret comes back"
);
assert_eq!(created.expose().len(), 64, "256 bits, hex encoded");
let rendered = created.expose().to_string();
assert!(
rendered.chars().all(|c| c.is_ascii_hexdigit()),
"one printable line: {rendered}"
);
#[cfg(unix)]
{
use std::os::unix::fs::PermissionsExt;
let mode = std::fs::metadata(ControlToken::path(dir.path()))
.unwrap()
.permissions()
.mode();
assert_eq!(mode & 0o777, 0o600, "the token must be owner-only");
}
}
#[test]
fn each_token_is_different() {
let a = tempfile::tempdir().unwrap();
let b = tempfile::tempdir().unwrap();
let first = ControlToken::create(a.path()).unwrap();
let second = ControlToken::create(b.path()).unwrap();
assert!(!first.matches(second.expose()), "tokens must not repeat");
}
#[test]
fn a_wrong_or_truncated_token_does_not_match() {
let dir = tempfile::tempdir().unwrap();
let token = ControlToken::create(dir.path()).unwrap();
assert!(!token.matches(""));
assert!(!token.matches("deadbeef"));
let half: String = token.expose().chars().take(32).collect();
assert!(!token.matches(&half));
assert!(token.matches(token.expose()));
}
#[test]
fn the_token_is_redacted_in_debug_output() {
let dir = tempfile::tempdir().unwrap();
let token = ControlToken::create(dir.path()).unwrap();
let rendered = format!("{token:?}");
assert!(!rendered.contains(token.expose()), "{rendered}");
assert!(rendered.contains("redacted"), "{rendered}");
}
#[test]
fn a_missing_token_still_builds_a_client() {
let dir = tempfile::tempdir().unwrap();
let client = ControlClient::for_home(control_id(dir.path()), dir.path());
let err = client.refused().to_string();
assert!(err.contains("no control token was found"), "{err}");
assert!(err.contains("lev daemon restart"), "{err}");
}
#[test]
fn a_rejected_token_reads_differently_from_a_missing_one() {
let dir = tempfile::tempdir().unwrap();
let _token = ControlToken::create(dir.path()).unwrap();
let client = ControlClient::for_home(control_id(dir.path()), dir.path());
let err = client.refused().to_string();
assert!(err.contains("refused this client's control token"), "{err}");
assert!(!err.contains("no control token was found"), "{err}");
}
use super::*;
use tokio::sync::mpsc;
fn no_events() -> broadcast::Sender<WorldEvent> {
broadcast::channel(16).0
}
fn spawn_fake_host(mut rx: mpsc::UnboundedReceiver<ControlOp>) {
tokio::spawn(async move {
while let Some(op) = rx.recv().await {
match op {
ControlOp::Spawn { args, reply } => {
let result = if args.run_id == "FAIL" {
Err("bad blueprint".to_string())
} else {
Ok(args.run_id)
};
let _ = reply.send(result);
}
ControlOp::Status { reply, .. } => {
let _ = reply.send(Some(AgentStatus::Active));
}
ControlOp::Pause { reply, .. }
| ControlOp::Resume { reply, .. }
| ControlOp::Cancel { reply, .. } => {
let _ = reply.send(true);
}
ControlOp::Message { reply, .. }
| ControlOp::AnswerInteraction { reply, .. }
| ControlOp::CancelInteraction { reply, .. }
| ControlOp::Shutdown { reply } => {
let _ = reply.send(true);
}
ControlOp::List { reply } => {
let _ = reply.send(vec![("run-a".to_string(), AgentStatus::Active)]);
}
ControlOp::ListInteractions { reply } => {
let _ = reply.send(vec![]);
}
}
}
});
}
fn test_listener() -> (ControlListener, ControlId, tempfile::TempDir) {
let dir = tempfile::tempdir().unwrap();
let id = control_id(dir.path());
let listener = bind_control_listener(&id).unwrap();
(listener, id, dir)
}
async fn round_trip(req: &ControlRequest) -> ControlResponse {
let (op_tx, op_rx) = mpsc::unbounded_channel();
spawn_fake_host(op_rx);
let (mut listener, id, _dir) = test_listener();
tokio::spawn(async move {
let stream = listener
.accept()
.await
.expect("accept succeeds")
.expect("our own connection is admitted");
let _ = handle_connection(stream, op_tx, no_events(), None).await;
});
let stream = connect(&id).await.unwrap();
let (read_half, mut write_half) = tokio::io::split(stream);
let mut line = serde_json::to_string(req).unwrap();
line.push('\n');
write_half.write_all(line.as_bytes()).await.unwrap();
let mut lines = BufReader::new(read_half).lines();
let resp_line = lines.next_line().await.unwrap().unwrap();
serde_json::from_str(&resp_line).unwrap()
}
#[tokio::test]
async fn status_request_round_trips() {
let resp = round_trip(&ControlRequest::Status {
run_id: "run-a".to_string(),
})
.await;
assert_eq!(
resp,
ControlResponse::Status {
status: Some(AgentStatus::Active)
}
);
}
#[tokio::test]
async fn control_ops_round_trip() {
for req in [
ControlRequest::Pause {
run_id: "r".to_string(),
},
ControlRequest::Resume {
run_id: "r".to_string(),
},
ControlRequest::Cancel {
run_id: "r".to_string(),
},
ControlRequest::Message {
agent_id: "a".to_string(),
content: "hi".to_string(),
target_region: None,
},
ControlRequest::AnswerInteraction {
response: InteractionResponse::text("q1", "yes"),
},
ControlRequest::CancelInteraction {
request_id: "q1".to_string(),
},
] {
assert_eq!(round_trip(&req).await, ControlResponse::Ok { ok: true });
}
}
#[tokio::test]
async fn spawn_request_round_trips() {
let resp = round_trip(&ControlRequest::Spawn {
args: Box::new(SpawnArgs {
run_id: "run-9".to_string(),
blueprint_path: "/agents/x".to_string(),
task: "do it".to_string(),
regions: Default::default(),
model: None,
workdir: "/w".to_string(),
metadata: Default::default(),
callback_url: None,
callback_secret: None,
yolo: false,
no_seed_commands: false,
allow: Vec::new(),
max_depth: None,
parent_run_id: None,
}),
})
.await;
assert_eq!(
resp,
ControlResponse::Spawned {
run_id: "run-9".to_string()
}
);
}
#[tokio::test]
async fn spawn_error_from_host_becomes_error_response() {
let resp = round_trip(&ControlRequest::Spawn {
args: Box::new(SpawnArgs {
run_id: "FAIL".to_string(),
..Default::default()
}),
})
.await;
assert_eq!(
std::mem::discriminant(&resp),
std::mem::discriminant(&ControlResponse::Error {
message: String::new()
})
);
}
#[tokio::test]
async fn list_interactions_round_trips() {
let resp = round_trip(&ControlRequest::ListInteractions).await;
assert_eq!(
resp,
ControlResponse::Interactions {
interactions: vec![]
}
);
}
#[tokio::test]
async fn list_request_round_trips() {
let resp = round_trip(&ControlRequest::List).await;
assert_eq!(
resp,
ControlResponse::List {
runs: vec![("run-a".to_string(), AgentStatus::Active)]
}
);
}
#[tokio::test]
async fn shutdown_request_round_trips() {
assert_eq!(
round_trip(&ControlRequest::Shutdown).await,
ControlResponse::Ok { ok: true }
);
}
fn completed(run_id: &str) -> WorldEvent {
WorldEvent::Completed {
run_id: run_id.to_string(),
agent_id: "a".to_string(),
status: "complete".to_string(),
}
}
#[tokio::test]
async fn dispatch_rejects_subscribe_as_a_single_reply_op() {
let (op_tx, _rx) = mpsc::unbounded_channel();
let resp = dispatch(ControlRequest::Subscribe, &op_tx).await;
assert_eq!(
std::mem::discriminant(&resp),
std::mem::discriminant(&ControlResponse::Error {
message: String::new()
})
);
}
#[tokio::test]
async fn stream_events_skips_lagged_writes_ok_and_stops_on_closed() {
use tokio::io::AsyncReadExt;
let (tx, rx) = broadcast::channel::<WorldEvent>(1);
tx.send(completed("first")).unwrap();
tx.send(completed("second")).unwrap();
tx.send(completed("third")).unwrap();
drop(tx);
let (mut w, mut r) = tokio::io::duplex(4096);
let server = tokio::spawn(async move { stream_events(&mut w, rx).await });
let mut buf = String::new();
r.read_to_string(&mut buf).await.unwrap();
server.await.unwrap().unwrap();
assert!(buf.contains("third"));
assert!(!buf.contains("first"));
}
#[tokio::test]
async fn stream_events_returns_when_the_client_hangs_up() {
let (tx, rx) = broadcast::channel::<WorldEvent>(4);
tx.send(completed("x")).unwrap();
let (mut w, r) = tokio::io::duplex(64);
drop(r); stream_events(&mut w, rx).await.unwrap();
drop(tx);
}
#[test]
fn creating_a_token_in_an_unwritable_place_is_an_error() {
let dir = tempfile::tempdir().unwrap();
let blocker = dir.path().join("blocker");
std::fs::write(&blocker, b"x").unwrap();
assert!(
ControlToken::create(&blocker.join("nested")).is_err(),
"a directory that cannot be created is an error"
);
let occupied = dir.path().join("occupied");
std::fs::create_dir_all(ControlToken::path(&occupied)).unwrap();
assert!(
ControlToken::create(&occupied).is_err(),
"a token file that cannot be written is an error"
);
}
#[tokio::test]
async fn dispatching_authenticate_is_inert() {
let (op_tx, _op_rx) = mpsc::unbounded_channel();
let response = dispatch(
ControlRequest::Authenticate {
token: "irrelevant".to_string(),
},
&op_tx,
)
.await;
assert!(matches!(response, ControlResponse::Ok { ok: true }));
}
#[tokio::test]
async fn re_authenticating_on_an_open_connection_is_accepted() {
let (events, _r) = broadcast::channel::<WorldEvent>(16);
let (op_tx, op_rx) = mpsc::unbounded_channel();
spawn_fake_host(op_rx);
let (mut listener, id, dir) = test_listener();
let token = ControlToken::create(dir.path()).unwrap();
let server_token = token.clone();
let server = tokio::spawn(async move {
let stream = listener.accept().await.unwrap().unwrap();
let _ = handle_connection(stream, op_tx, events, Some(server_token)).await;
});
let stream = connect(&id).await.unwrap();
let (read_half, mut write_half) = tokio::io::split(stream);
let mut lines = BufReader::new(read_half).lines();
let hello = serde_json::to_string(&ControlRequest::Authenticate {
token: token.expose().to_string(),
})
.unwrap();
for _ in 0..2 {
write_half
.write_all(format!("{hello}\n").as_bytes())
.await
.unwrap();
let resp: ControlResponse =
serde_json::from_str(&lines.next_line().await.unwrap().unwrap()).unwrap();
let rendered = format!("{resp:?}");
assert!(rendered.starts_with("Ok"), "{rendered}");
}
drop(write_half);
drop(lines);
server.await.unwrap();
}
#[tokio::test]
async fn a_connection_closed_during_authentication_is_reported() {
let (mut listener, id, dir) = test_listener();
let token = ControlToken::create(dir.path()).unwrap();
tokio::spawn(async move {
let _ = listener.accept().await;
});
let err = ControlClient::new(id)
.with_token(token)
.list()
.await
.expect_err("a hang-up during authentication is an error");
assert_eq!(err.kind(), std::io::ErrorKind::UnexpectedEof);
assert!(err.to_string().contains("during authentication"), "{err}");
}
#[tokio::test]
async fn an_unauthenticated_caller_is_refused_and_disconnected() {
let (events, _r) = broadcast::channel::<WorldEvent>(16);
let (op_tx, op_rx) = mpsc::unbounded_channel();
spawn_fake_host(op_rx);
let (mut listener, id, dir) = test_listener();
let token = ControlToken::create(dir.path()).unwrap();
let server_token = token.clone();
let server = tokio::spawn(async move {
let stream = listener.accept().await.unwrap().unwrap();
let _ = handle_connection(stream, op_tx, events, Some(server_token)).await;
});
let err = ControlClient::new(id)
.list()
.await
.expect_err("an unauthenticated List must not be served");
assert_eq!(err.kind(), std::io::ErrorKind::PermissionDenied);
assert!(err.to_string().contains("token"), "{err}");
server.await.unwrap();
}
#[tokio::test]
async fn a_client_presenting_the_wrong_token_is_refused() {
let (events, _r) = broadcast::channel::<WorldEvent>(16);
let (op_tx, op_rx) = mpsc::unbounded_channel();
spawn_fake_host(op_rx);
let (mut listener, id, dir) = test_listener();
let real = ControlToken::create(dir.path()).unwrap();
tokio::spawn(async move {
let stream = listener.accept().await.unwrap().unwrap();
let _ = handle_connection(stream, op_tx, events, Some(real)).await;
});
let other_dir = tempfile::tempdir().unwrap();
let wrong = ControlToken::create(other_dir.path()).unwrap();
let err = ControlClient::new(id)
.with_token(wrong)
.list()
.await
.expect_err("a wrong token is refused");
assert_eq!(err.kind(), std::io::ErrorKind::PermissionDenied);
assert!(err.to_string().contains("refused"), "{err}");
}
#[tokio::test]
async fn a_client_presenting_the_right_token_is_served() {
let (events, _r) = broadcast::channel::<WorldEvent>(16);
let (op_tx, op_rx) = mpsc::unbounded_channel();
spawn_fake_host(op_rx);
let (mut listener, id, dir) = test_listener();
let token = ControlToken::create(dir.path()).unwrap();
let server_token = token.clone();
let server = tokio::spawn(async move {
let stream = listener.accept().await.unwrap().unwrap();
let _ = handle_connection(stream, op_tx, events, Some(server_token)).await;
});
let client = ControlClient::for_home(id, dir.path());
let response = client
.list()
.await
.expect("an authenticated List is served");
let rendered = format!("{response:?}");
assert!(
rendered.starts_with("List"),
"expected a run list: {rendered}"
);
server.await.unwrap();
}
#[tokio::test]
async fn subscribe_streams_events_to_the_client() {
let (events, _r) = broadcast::channel::<WorldEvent>(16);
let (op_tx, op_rx) = mpsc::unbounded_channel();
spawn_fake_host(op_rx);
let (mut listener, id, _dir) = test_listener();
let server_events = events.clone();
let server = tokio::spawn(async move {
let stream = listener
.accept()
.await
.expect("accept succeeds")
.expect("our own connection is admitted");
let _ = handle_connection(stream, op_tx, server_events, None).await;
});
let mut stream = ControlClient::new(id).subscribe().await.unwrap();
let received = loop {
events.send(completed("run-1")).unwrap();
tokio::select! {
e = stream.next() => break e,
_ = tokio::time::sleep(std::time::Duration::from_millis(5)) => {}
}
};
let received = received.expect("an event should have streamed to the client");
assert_eq!(
std::mem::discriminant(&received),
std::mem::discriminant(&completed("x"))
);
drop(events);
server.await.unwrap();
}
#[tokio::test]
async fn subscribe_errors_when_daemon_absent() {
let dir = tempfile::tempdir().unwrap();
let client = ControlClient::new(control_id(&dir.path().join("no-daemon")));
assert!(client.subscribe().await.is_err());
}
#[tokio::test]
async fn subscribe_stream_ends_when_the_daemon_closes() {
let (events, _r) = broadcast::channel::<WorldEvent>(16);
let (op_tx, op_rx) = mpsc::unbounded_channel();
spawn_fake_host(op_rx);
let (mut listener, id, _dir) = test_listener();
let server_events = events.clone();
tokio::spawn(async move {
let stream = listener
.accept()
.await
.expect("accept succeeds")
.expect("our own connection is admitted");
let _ = handle_connection(stream, op_tx, server_events, None).await;
});
let mut stream = ControlClient::new(id).subscribe().await.unwrap();
tokio::time::sleep(std::time::Duration::from_millis(50)).await;
drop(events);
assert!(stream.next().await.is_none());
}
async fn connected_pair() -> (ClientStream, ServerStream, tempfile::TempDir) {
let (mut listener, id, dir) = test_listener();
let (client, server) = tokio::join!(connect(&id), listener.accept());
let server = server
.expect("accept succeeds")
.expect("our own connection is admitted");
(client.unwrap(), server, dir)
}
#[tokio::test]
async fn malformed_request_gets_error_and_connection_continues() {
let (op_tx, op_rx) = mpsc::unbounded_channel();
spawn_fake_host(op_rx);
let (client, server, _dir) = connected_pair().await;
let handle =
tokio::spawn(async move { handle_connection(server, op_tx, no_events(), None).await });
let (read_half, mut write_half) = tokio::io::split(client);
write_half.write_all(b"\nnot json\n").await.unwrap();
let mut lines = BufReader::new(read_half).lines();
let err_line = lines.next_line().await.unwrap().unwrap();
let resp: ControlResponse = serde_json::from_str(&err_line).unwrap();
assert_eq!(
std::mem::discriminant(&resp),
std::mem::discriminant(&ControlResponse::Error {
message: String::new()
})
);
let mut valid = serde_json::to_string(&ControlRequest::List).unwrap();
valid.push('\n');
write_half.write_all(valid.as_bytes()).await.unwrap();
let ok_line = lines.next_line().await.unwrap().unwrap();
let ok: ControlResponse = serde_json::from_str(&ok_line).unwrap();
assert_eq!(
std::mem::discriminant(&ok),
std::mem::discriminant(&ControlResponse::List { runs: vec![] })
);
drop(write_half);
drop(lines);
handle.await.unwrap().unwrap();
}
#[tokio::test]
async fn invalid_utf8_line_ends_connection_with_error() {
let (op_tx, op_rx) = mpsc::unbounded_channel();
spawn_fake_host(op_rx);
let (client, server, _dir) = connected_pair().await;
let handle =
tokio::spawn(async move { handle_connection(server, op_tx, no_events(), None).await });
let (_read_half, mut write_half) = tokio::io::split(client);
write_half.write_all(&[0xff, 0xfe, b'\n']).await.unwrap();
let result = handle.await.unwrap();
assert!(result.is_err());
}
#[tokio::test]
async fn client_round_trips_status_and_list() {
let (op_tx, op_rx) = mpsc::unbounded_channel();
spawn_fake_host(op_rx);
let (mut listener, id, _dir) = test_listener();
tokio::spawn(async move {
for _ in 0..4 {
let stream = listener
.accept()
.await
.expect("accept succeeds")
.expect("our own connection is admitted");
let op_tx = op_tx.clone();
tokio::spawn(async move {
let _ = handle_connection(stream, op_tx, no_events(), None).await;
});
}
});
let client = ControlClient::new(id);
let spawned = client
.spawn(SpawnArgs {
run_id: "r-c".to_string(),
..Default::default()
})
.await
.unwrap();
assert_eq!(
spawned,
ControlResponse::Spawned {
run_id: "r-c".to_string()
}
);
let status = client.status("run-a").await.unwrap();
assert_eq!(
status,
ControlResponse::Status {
status: Some(AgentStatus::Active)
}
);
let list = client.list().await.unwrap();
assert_eq!(
std::mem::discriminant(&list),
std::mem::discriminant(&ControlResponse::List { runs: vec![] })
);
assert_eq!(
client.shutdown().await.unwrap(),
ControlResponse::Ok { ok: true }
);
}
#[tokio::test]
async fn client_errors_when_daemon_absent() {
let dir = tempfile::tempdir().unwrap();
let id = control_id(&dir.path().join("no-daemon-here"));
assert!(ControlClient::new(id).list().await.is_err());
}
async fn raw_server(bytes: &'static [u8]) -> (ControlId, tempfile::TempDir) {
let (mut listener, id, dir) = test_listener();
tokio::spawn(async move {
let stream = listener
.accept()
.await
.expect("accept succeeds")
.expect("our own connection is admitted");
let (_r, mut w) = tokio::io::split(stream);
let _ = w.write_all(bytes).await;
});
(id, dir)
}
#[tokio::test]
async fn client_errors_on_unparseable_response() {
let (id, _dir) = raw_server(b"not json\n").await;
let err = ControlClient::new(id).list().await.unwrap_err();
assert_eq!(err.kind(), std::io::ErrorKind::InvalidData);
}
#[tokio::test]
async fn client_errors_on_invalid_utf8_response() {
let (id, _dir) = raw_server(&[0xff, 0xfe, b'\n']).await;
assert!(ControlClient::new(id).list().await.is_err());
}
#[tokio::test]
async fn client_errors_on_closed_connection_without_reply() {
let (mut listener, id, _dir) = test_listener();
tokio::spawn(async move {
let stream = listener
.accept()
.await
.expect("accept succeeds")
.expect("our own connection is admitted");
let (read_half, _write_half) = tokio::io::split(stream);
let mut lines = BufReader::new(read_half).lines();
let _ = lines.next_line().await;
});
let err = ControlClient::new(id).list().await.unwrap_err();
assert_eq!(err.kind(), std::io::ErrorKind::UnexpectedEof);
}
#[tokio::test]
async fn client_times_out_on_a_daemon_that_never_answers() {
let (mut listener, id, _dir) = test_listener();
let (tx, rx) = oneshot::channel();
tokio::spawn(async move {
let stream = listener
.accept()
.await
.expect("accept succeeds")
.expect("our own connection is admitted");
let (read_half, write_half) = tokio::io::split(stream);
let mut lines = BufReader::new(read_half).lines();
let _ = lines.next_line().await;
let _ = tx.send((lines, write_half));
});
let err = temp_env::async_with_vars([("LEVIATH_CONTROL_TIMEOUT_SECS", Some("1"))], async {
ControlClient::new(id).list().await.unwrap_err()
})
.await;
assert_eq!(err.kind(), std::io::ErrorKind::TimedOut);
assert!(err.to_string().contains("did not respond"), "got: {err}");
drop(rx);
}
#[test]
fn request_timeout_honors_the_override_and_falls_back() {
temp_env::with_var("LEVIATH_CONTROL_TIMEOUT_SECS", Some("7"), || {
assert_eq!(request_timeout(), std::time::Duration::from_secs(7));
});
temp_env::with_var("LEVIATH_CONTROL_TIMEOUT_SECS", Some("0"), || {
assert_eq!(request_timeout(), std::time::Duration::MAX);
});
temp_env::with_var("LEVIATH_CONTROL_TIMEOUT_SECS", Some("soon"), || {
assert_eq!(
request_timeout(),
std::time::Duration::from_secs(DEFAULT_CONTROL_TIMEOUT_SECS)
);
});
temp_env::with_var_unset("LEVIATH_CONTROL_TIMEOUT_SECS", || {
assert_eq!(
request_timeout(),
std::time::Duration::from_secs(DEFAULT_CONTROL_TIMEOUT_SECS)
);
});
}
#[test]
fn spawn_gets_a_longer_deadline_than_other_ops() {
let spawn = ControlRequest::Spawn {
args: Box::new(SpawnArgs::default()),
};
let cancel = ControlRequest::Cancel {
run_id: "r".to_string(),
};
temp_env::with_var_unset("LEVIATH_CONTROL_TIMEOUT_SECS", || {
assert_eq!(
timeout_for(&spawn),
std::time::Duration::from_secs(SPAWN_CONTROL_TIMEOUT_SECS)
);
assert_eq!(
timeout_for(&cancel),
std::time::Duration::from_secs(DEFAULT_CONTROL_TIMEOUT_SECS)
);
});
temp_env::with_var("LEVIATH_CONTROL_TIMEOUT_SECS", Some("900"), || {
assert_eq!(timeout_for(&spawn), std::time::Duration::from_secs(900));
assert_eq!(timeout_for(&cancel), std::time::Duration::from_secs(900));
});
temp_env::with_var("LEVIATH_CONTROL_TIMEOUT_SECS", Some("0"), || {
assert_eq!(timeout_for(&spawn), std::time::Duration::MAX);
});
}
#[tokio::test]
async fn bind_rejects_when_daemon_already_running() {
let (_live, id, _dir) = test_listener(); let err = bind_control_listener(&id).unwrap_err();
assert_eq!(err.kind(), std::io::ErrorKind::AddrInUse);
}
#[tokio::test]
async fn is_daemon_running_reflects_a_live_listener() {
let dir = tempfile::tempdir().unwrap();
let id = control_id(dir.path());
assert!(!is_daemon_running(&id)); let _live = bind_control_listener(&id).unwrap();
assert!(is_daemon_running(&id)); }
#[tokio::test]
async fn dispatch_returns_neutral_when_host_gone() {
let (op_tx, op_rx) = mpsc::unbounded_channel();
drop(op_rx);
assert_eq!(
dispatch(
ControlRequest::Status {
run_id: "r".to_string()
},
&op_tx
)
.await,
ControlResponse::Status { status: None }
);
assert_eq!(
dispatch(
ControlRequest::Cancel {
run_id: "r".to_string()
},
&op_tx
)
.await,
ControlResponse::Ok { ok: false }
);
assert_eq!(
dispatch(ControlRequest::List, &op_tx).await,
ControlResponse::List { runs: vec![] }
);
assert_eq!(
dispatch(ControlRequest::ListInteractions, &op_tx).await,
ControlResponse::Interactions {
interactions: vec![]
}
);
assert_eq!(
std::mem::discriminant(
&dispatch(
ControlRequest::Spawn {
args: Box::new(SpawnArgs::default())
},
&op_tx
)
.await
),
std::mem::discriminant(&ControlResponse::Error {
message: String::new()
})
);
}
}