use serde::{Deserialize, Serialize};
use std::time::{SystemTime, UNIX_EPOCH};
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RecordedCommand {
pub timestamp: u64,
pub command: String,
pub payload: String,
}
impl RecordedCommand {
#[must_use]
pub fn new(command: impl Into<String>, payload: impl Into<String>) -> Self {
#[allow(clippy::cast_possible_truncation)]
let timestamp = SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap_or_default()
.as_millis() as u64;
Self {
timestamp,
command: command.into(),
payload: payload.into(),
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
#[derive(Default)]
pub enum DumpFormat {
#[default]
Text,
Jsonl,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum Request {
Spawn {
cmd: Vec<String>,
#[serde(default = "default_rows")]
rows: u16,
#[serde(default = "default_cols")]
cols: u16,
#[serde(default)]
name: Option<String>,
#[serde(default)]
labels: Vec<String>,
#[serde(default)]
timeout: Option<u64>,
#[serde(default)]
max_output: Option<u64>,
#[serde(default)]
env: Vec<String>,
#[serde(default)]
cwd: Option<String>,
#[serde(default)]
no_resize: bool,
#[serde(default)]
record: bool,
#[serde(default)]
memory_limit: Option<String>,
},
List {
#[serde(default)]
labels: Vec<String>,
},
Kill {
#[serde(default)]
id: Option<String>,
#[serde(default)]
labels: Vec<String>,
#[serde(default)]
all: bool,
#[serde(default = "default_signal")]
signal: i32,
#[serde(default)]
proc_filter: Option<String>,
},
Send {
#[serde(default)]
id: Option<String>,
#[serde(default)]
labels: Vec<String>,
#[serde(default)]
all: bool,
#[serde(default)]
proc_filter: Option<String>,
data: String,
#[serde(default)]
newline: bool,
#[serde(default)]
enter: bool,
#[serde(default)]
submit_delay_ms: Option<u64>,
#[serde(default)]
paste: bool,
},
SendBytes {
#[serde(default)]
id: Option<String>,
#[serde(default)]
labels: Vec<String>,
#[serde(default)]
all: bool,
#[serde(default)]
proc_filter: Option<String>,
#[serde(with = "base64_bytes")]
data: Vec<u8>,
},
Tail {
id: String,
#[serde(default = "default_tail_lines")]
lines: usize,
#[serde(default)]
follow: bool,
},
Dump {
id: String,
#[serde(default)]
since: Option<u64>,
#[serde(default)]
format: DumpFormat,
},
Snapshot {
id: String,
#[serde(default = "default_true")]
strip_colors: bool,
},
Attach {
id: String,
#[serde(default)]
readonly: bool,
},
Shutdown,
Ping,
Events {
#[serde(default)]
filter: Vec<String>,
#[serde(default)]
include_output: bool,
},
Resize {
id: String,
rows: u16,
cols: u16,
#[serde(default)]
clear_transcript: bool,
},
GetRecording {
id: String,
},
GetEnv {
id: String,
},
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AgentInfo {
pub id: String,
pub pid: u32,
pub state: AgentState,
pub command: Vec<String>,
#[serde(default)]
pub labels: Vec<String>,
pub size: (u16, u16),
pub started_at: u64,
pub exit_code: Option<i32>,
#[serde(default)]
pub exit_reason: Option<ExitReason>,
#[serde(default)]
pub limits: Option<ResourceLimits>,
#[serde(default)]
pub no_resize: bool,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub rss_bytes: Option<u64>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum ExitReason {
Normal,
Timeout,
Killed,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub struct ResourceLimits {
#[serde(default, skip_serializing_if = "Option::is_none")]
pub timeout: Option<u64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub max_output: Option<u64>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum AgentState {
Running,
Exited,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TranscriptEntry {
pub timestamp: u64,
#[serde(with = "base64_bytes")]
pub data: Vec<u8>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum Response {
Ok,
Pong,
Spawned {
id: String,
pid: u32,
},
Agents {
agents: Vec<AgentInfo>,
},
SendResults {
results: Vec<SendOutcome>,
},
Output {
#[serde(with = "base64_bytes")]
data: Vec<u8>,
#[serde(default)]
exited: bool,
},
Transcript {
entries: Vec<TranscriptEntry>,
},
Snapshot {
content: String,
cursor: (u16, u16),
size: (u16, u16),
},
Error {
message: String,
},
AgentExited {
id: String,
exit_code: Option<i32>,
},
AttachStarted {
id: String,
size: (u16, u16),
},
AttachEnded {
reason: AttachEndReason,
},
Event(Event),
Recording {
agent_id: String,
commands: Vec<RecordedCommand>,
},
AgentEnv {
id: String,
env: Vec<(String, String)>,
},
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum AttachEndReason {
Detached,
AgentExited { exit_code: Option<i32> },
Error { message: String },
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "event", rename_all = "snake_case")]
pub enum Event {
AgentSpawned {
id: String,
pid: u32,
command: Vec<String>,
#[serde(default)]
labels: Vec<String>,
},
AgentOutput {
id: String,
#[serde(with = "base64_bytes")]
data: Vec<u8>,
},
AgentExited {
id: String,
exit_code: Option<i32>,
},
}
impl Response {
pub fn error(message: impl Into<String>) -> Self {
Self::Error {
message: message.into(),
}
}
}
const fn default_rows() -> u16 {
24
}
const fn default_cols() -> u16 {
80
}
const fn default_signal() -> i32 {
15 }
const fn default_tail_lines() -> usize {
10
}
const fn default_true() -> bool {
true
}
pub const DEFAULT_SUBMIT_DELAY_MS: u64 = 50;
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct SendOutcome {
pub id: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub error: Option<String>,
}
impl SendOutcome {
#[must_use]
pub const fn delivered(id: String) -> Self {
Self { id, error: None }
}
#[must_use]
pub const fn failed(id: String, error: String) -> Self {
Self {
id,
error: Some(error),
}
}
#[must_use]
pub const fn is_ok(&self) -> bool {
self.error.is_none()
}
}
pub const PASTE_START: &[u8] = b"\x1b[200~";
pub const PASTE_END: &[u8] = b"\x1b[201~";
const MAX_SEND_BYTES_BASE64_LEN: usize = 1024 * 1024;
mod base64_bytes {
use super::MAX_SEND_BYTES_BASE64_LEN;
use serde::{Deserialize, Deserializer, Serialize, Serializer};
pub fn serialize<S>(bytes: &[u8], serializer: S) -> Result<S::Ok, S::Error>
where
S: Serializer,
{
use base64::Engine;
let encoded = base64::engine::general_purpose::STANDARD.encode(bytes);
encoded.serialize(serializer)
}
pub fn deserialize<'de, D>(deserializer: D) -> Result<Vec<u8>, D::Error>
where
D: Deserializer<'de>,
{
use base64::Engine;
let s = String::deserialize(deserializer)?;
if s.len() > MAX_SEND_BYTES_BASE64_LEN {
return Err(serde::de::Error::custom(format!(
"send_bytes payload too large: {} encoded bytes exceeds limit of {MAX_SEND_BYTES_BASE64_LEN}",
s.len()
)));
}
base64::engine::general_purpose::STANDARD
.decode(&s)
.map_err(serde::de::Error::custom)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_request_serialization_roundtrip() {
let requests = vec![
Request::Spawn {
cmd: vec!["bash".into(), "-c".into(), "echo hello".into()],
rows: 24,
cols: 80,
name: None,
labels: vec!["worker".into()],
timeout: Some(60),
max_output: Some(1024 * 1024),
env: vec![],
cwd: None,
no_resize: false,
record: false,
memory_limit: Some("4G".into()),
},
Request::List { labels: vec![] },
Request::Kill {
id: Some("test-agent".into()),
labels: vec![],
all: false,
signal: 9,
proc_filter: None,
},
Request::Send {
id: Some("test-agent".into()),
labels: Vec::new(),
all: false,
proc_filter: None,
data: "hello\n".into(),
newline: false,
enter: false,
submit_delay_ms: None,
paste: false,
},
Request::SendBytes {
id: Some("test-agent".into()),
labels: Vec::new(),
all: false,
proc_filter: None,
data: vec![0x1b, 0x5b, 0x41], },
Request::Tail {
id: "test-agent".into(),
lines: 20,
follow: true,
},
Request::Snapshot {
id: "test-agent".into(),
strip_colors: true,
},
Request::Ping,
Request::Shutdown,
Request::Events {
filter: vec!["agent-1".into()],
include_output: true,
},
Request::Resize {
id: "test-agent".into(),
rows: 40,
cols: 120,
clear_transcript: false,
},
Request::GetRecording {
id: "test-agent".into(),
},
Request::GetEnv {
id: "test-agent".into(),
},
];
for req in requests {
let json = serde_json::to_string(&req).expect("serialize");
let parsed: Request = serde_json::from_str(&json).expect("deserialize");
let json2 = serde_json::to_string(&parsed).expect("re-serialize");
assert_eq!(json, json2, "roundtrip failed for {:?}", req);
}
}
#[test]
fn test_send_without_submit_delay_defaults_to_none() {
let json = r#"{"type":"send","id":"a","data":"hi","newline":true}"#;
let parsed: Request = serde_json::from_str(json).expect("deserialize");
match parsed {
Request::Send {
submit_delay_ms,
newline,
enter,
..
} => {
assert_eq!(submit_delay_ms, None);
assert!(newline);
assert!(!enter);
}
other => panic!("expected Send, got {other:?}"),
}
}
#[test]
fn test_send_bytes_oversized_payload_rejected() {
let oversized = "A".repeat(MAX_SEND_BYTES_BASE64_LEN + 4);
let json = format!(r#"{{"type":"send_bytes","id":"a","data":"{oversized}"}}"#);
let err = serde_json::from_str::<Request>(&json)
.expect_err("oversized send_bytes payload should be rejected");
assert!(
err.to_string().contains("send_bytes payload too large"),
"unexpected error: {err}"
);
}
#[test]
fn test_send_bytes_at_limit_accepted() {
let at_limit = "A".repeat(MAX_SEND_BYTES_BASE64_LEN);
let json = format!(r#"{{"type":"send_bytes","id":"a","data":"{at_limit}"}}"#);
let parsed: Request =
serde_json::from_str(&json).expect("payload at the limit should be accepted");
match parsed {
Request::SendBytes { data, .. } => {
assert_eq!(data.len(), MAX_SEND_BYTES_BASE64_LEN / 4 * 3);
}
other => panic!("expected SendBytes, got {other:?}"),
}
}
#[test]
fn test_response_serialization_roundtrip() {
let responses = vec![
Response::Ok,
Response::Pong,
Response::Spawned {
id: "rusty-nail".into(),
pid: 12345,
},
Response::Agents {
agents: vec![AgentInfo {
id: "rusty-nail".into(),
pid: 12345,
state: AgentState::Running,
command: vec!["bash".into()],
labels: vec!["worker".into()],
size: (24, 80),
started_at: 1706140800000,
exit_code: None,
exit_reason: None,
limits: Some(ResourceLimits {
timeout: Some(60),
max_output: None,
}),
no_resize: false,
rss_bytes: Some(142_000_000),
}],
},
Response::Output {
data: b"hello world\n".to_vec(),
exited: false,
},
Response::Snapshot {
content: "$ echo hello\nhello\n$ ".into(),
cursor: (2, 2),
size: (24, 80),
},
Response::error("agent not found"),
Response::Event(Event::AgentSpawned {
id: "test-agent".into(),
pid: 12345,
command: vec!["bash".into()],
labels: vec![],
}),
Response::Event(Event::AgentOutput {
id: "test-agent".into(),
data: b"hello".to_vec(),
}),
Response::Event(Event::AgentExited {
id: "test-agent".into(),
exit_code: Some(0),
}),
Response::Recording {
agent_id: "test-agent".into(),
commands: vec![RecordedCommand {
timestamp: 1706140800000,
command: "send".into(),
payload: "hello\n".into(),
}],
},
Response::AgentEnv {
id: "test-agent".into(),
env: vec![
("CARGO_BUILD_JOBS".into(), "2".into()),
("PATH".into(), "/usr/bin".into()),
],
},
];
for resp in responses {
let json = serde_json::to_string(&resp).expect("serialize");
let parsed: Response = serde_json::from_str(&json).expect("deserialize");
let json2 = serde_json::to_string(&parsed).expect("re-serialize");
assert_eq!(json, json2, "roundtrip failed for {:?}", resp);
}
}
#[test]
fn test_base64_bytes_encoding() {
let req = Request::SendBytes {
id: Some("test".into()),
labels: Vec::new(),
all: false,
proc_filter: None,
data: vec![0x1b, 0x5b, 0x41],
};
let json = serde_json::to_string(&req).unwrap();
assert!(json.contains("G1tB")); }
}