use serde::{Deserialize, Serialize};
pub const AGENT_PROTOCOL_VERSION: u32 = 1;
pub const WORK_PACKET_KIND: &str = "agent_work";
pub const GRAPH_PACKET_KIND: &str = "agent_graph";
pub const AGENT_RECIPIENT_PREFIX: &str = "agent:";
pub const DEFAULT_HEARTBEAT_INTERVAL_SECONDS: u32 = 15;
pub fn host_target_triple() -> String {
format!("{}-{}", std::env::consts::ARCH, std::env::consts::OS)
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AgentRegisterRequest {
pub protocol_version: u32,
#[serde(default)]
pub agent_id: Option<String>,
pub max_concurrency: u32,
pub target_triple: String,
#[serde(default)]
pub capabilities: Vec<String>,
#[serde(default)]
pub ephemeral_public_key: Option<String>,
#[serde(default)]
pub ephemeral_key_pool: Vec<EphemeralKeyEntry>,
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct EphemeralKeyEntry {
pub key_id: String,
pub public_key_b64: String,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AgentRegisterResponse {
pub protocol_version: u32,
pub agent_id: String,
pub heartbeat_interval_seconds: u32,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AgentHeartbeatRequest {
pub protocol_version: u32,
pub agent_id: String,
pub in_flight: u32,
pub available_capacity: u32,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AgentHeartbeatResponse {
pub protocol_version: u32,
#[serde(default)]
pub replenish_keys: u32,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AgentKeyReplenishRequest {
pub protocol_version: u32,
pub agent_id: String,
pub keys: Vec<EphemeralKeyEntry>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AgentKeyReplenishResponse {
pub protocol_version: u32,
pub accepted: u32,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct WorkPacket {
pub protocol_version: u32,
pub task_execution_id: String,
pub workflow_execution_id: String,
pub task_name: String,
pub attempt: i32,
pub context: serde_json::Value,
pub artifact: ArtifactRef,
pub timeout_seconds: u32,
pub tenant_id: Option<String>,
#[serde(default)]
pub language: Option<String>,
#[serde(default)]
pub wrapped_secrets: Vec<WrappedSecret>,
#[serde(default)]
pub secret_key_id: Option<String>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct WrappedSecret {
pub name: String,
pub enc_b64: String,
pub ciphertext_b64: String,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct GraphWorkPacket {
pub protocol_version: u32,
pub firing_id: String,
pub graph_name: String,
pub cache: std::collections::HashMap<String, String>,
pub artifact: ArtifactRef,
pub timeout_seconds: u32,
pub tenant_id: Option<String>,
#[serde(default)]
pub language: Option<String>,
#[serde(default)]
pub wrapped_secrets: Vec<WrappedSecret>,
#[serde(default)]
pub secret_key_id: Option<String>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ArtifactRef {
pub digest: String,
pub fetch_url: String,
pub build_target_triple: String,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AgentResultRequest {
pub protocol_version: u32,
pub agent_id: String,
pub task_execution_id: String,
pub attempt: i32,
pub duration_ms: u64,
pub outcome: AgentOutcome,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AgentResultResponse {
pub protocol_version: u32,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "kind", rename_all = "snake_case")]
pub enum AgentOutcome {
Success { context: serde_json::Value },
Failure {
message: String,
classification: FailureClassification,
},
Refused {
reason: RefusalReason,
message: String,
},
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum FailureClassification {
TaskError,
Transient,
Validation,
Timeout,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum RefusalReason {
TargetTripleMismatch,
ArtifactFetchFailed,
RuntimeLoadFailed,
Shutdown,
TenantMismatch,
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn work_packet_round_trips_as_json() {
let p = WorkPacket {
protocol_version: AGENT_PROTOCOL_VERSION,
task_execution_id: "t1".into(),
workflow_execution_id: "w1".into(),
task_name: "ns::task".into(),
attempt: 1,
context: serde_json::json!({"k": 42}),
artifact: ArtifactRef {
digest: "deadbeef".into(),
fetch_url: "/v1/agent/artifact/deadbeef".into(),
build_target_triple: "aarch64-apple-darwin".into(),
},
timeout_seconds: 60,
tenant_id: Some("t1".into()),
language: Some("rust".into()),
wrapped_secrets: Vec::new(),
secret_key_id: None,
};
let json = serde_json::to_string(&p).unwrap();
let back: WorkPacket = serde_json::from_str(&json).unwrap();
assert_eq!(back.task_execution_id, "t1");
assert_eq!(back.artifact.build_target_triple, "aarch64-apple-darwin");
assert_eq!(back.context, serde_json::json!({"k": 42}));
}
#[test]
fn register_request_advertises_a_key_pool() {
let pool = vec![
EphemeralKeyEntry {
key_id: "k1".into(),
public_key_b64: "AAAA".into(),
},
EphemeralKeyEntry {
key_id: "k2".into(),
public_key_b64: "BBBB".into(),
},
];
let req = AgentRegisterRequest {
protocol_version: AGENT_PROTOCOL_VERSION,
agent_id: Some("a1".into()),
max_concurrency: 4,
target_triple: "aarch64-apple-darwin".into(),
capabilities: vec![],
ephemeral_public_key: None,
ephemeral_key_pool: pool.clone(),
};
let json = serde_json::to_string(&req).unwrap();
let back: AgentRegisterRequest = serde_json::from_str(&json).unwrap();
assert_eq!(back.ephemeral_key_pool, pool);
let legacy = r#"{"protocol_version":1,"max_concurrency":4,"target_triple":"x"}"#;
let back: AgentRegisterRequest = serde_json::from_str(legacy).unwrap();
assert!(back.ephemeral_key_pool.is_empty());
}
#[test]
fn work_packet_secret_key_id_round_trips() {
let p = WorkPacket {
protocol_version: AGENT_PROTOCOL_VERSION,
task_execution_id: "t1".into(),
workflow_execution_id: "w1".into(),
task_name: "ns::task".into(),
attempt: 1,
context: serde_json::json!({}),
artifact: ArtifactRef {
digest: "d".into(),
fetch_url: "/x".into(),
build_target_triple: "aarch64-apple-darwin".into(),
},
timeout_seconds: 60,
tenant_id: None,
language: None,
wrapped_secrets: Vec::new(),
secret_key_id: Some("key-42".into()),
};
let json = serde_json::to_string(&p).unwrap();
let back: WorkPacket = serde_json::from_str(&json).unwrap();
assert_eq!(back.secret_key_id.as_deref(), Some("key-42"));
}
#[test]
fn outcome_variants_round_trip_with_snake_case_tags() {
let success = AgentOutcome::Success {
context: serde_json::json!({}),
};
let json = serde_json::to_string(&success).unwrap();
assert!(json.contains("\"kind\":\"success\""));
let back: AgentOutcome = serde_json::from_str(&json).unwrap();
assert!(matches!(back, AgentOutcome::Success { .. }));
let failure = AgentOutcome::Failure {
message: "oops".into(),
classification: FailureClassification::Transient,
};
let json = serde_json::to_string(&failure).unwrap();
assert!(json.contains("\"kind\":\"failure\""));
assert!(json.contains("\"classification\":\"transient\""));
let back: AgentOutcome = serde_json::from_str(&json).unwrap();
assert!(matches!(
back,
AgentOutcome::Failure {
classification: FailureClassification::Transient,
..
}
));
let refused = AgentOutcome::Refused {
reason: RefusalReason::TargetTripleMismatch,
message: "expected x86_64, got aarch64".into(),
};
let json = serde_json::to_string(&refused).unwrap();
assert!(json.contains("\"reason\":\"target_triple_mismatch\""));
let back: AgentOutcome = serde_json::from_str(&json).unwrap();
assert!(matches!(
back,
AgentOutcome::Refused {
reason: RefusalReason::TargetTripleMismatch,
..
}
));
}
#[test]
fn agent_recipient_prefix_is_stable() {
assert_eq!(AGENT_RECIPIENT_PREFIX, "agent:");
let recipient = format!("{}{}", AGENT_RECIPIENT_PREFIX, "abc-123");
assert_eq!(recipient, "agent:abc-123");
}
#[test]
fn register_request_agent_id_defaults_to_none() {
let json = r#"{"protocol_version":1,"max_concurrency":4,"target_triple":"x86_64-unknown-linux-gnu"}"#;
let req: AgentRegisterRequest = serde_json::from_str(json).unwrap();
assert!(req.agent_id.is_none());
assert!(req.capabilities.is_empty());
}
}