use serde::{de, ser::SerializeStruct, Deserialize, Deserializer, Serialize, Serializer};
use serde_json::Value;
use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
pub const MAX_FRAME_BYTES: u32 = 16 * 1024 * 1024;
#[derive(Debug, thiserror::Error)]
pub enum ProtocolError {
#[error("io: {0}")]
Io(#[from] std::io::Error),
#[error("frame exceeds max size: {0} > {MAX_FRAME_BYTES}")]
FrameTooLarge(u32),
#[error("malformed json frame: {0}")]
Json(#[from] serde_json::Error),
#[error("connection closed before a full frame was read")]
UnexpectedEof,
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct Request {
pub id: u64,
pub method: String,
#[serde(default)]
pub params: Value,
}
impl Request {
pub fn new(id: u64, method: impl Into<String>, params: Value) -> Self {
Request {
id,
method: method.into(),
params,
}
}
}
#[derive(Debug, Clone, PartialEq)]
pub enum ResponsePayload {
Ok(Value),
Err(RpcError),
}
#[derive(Debug, Clone, PartialEq)]
pub struct Response {
pub id: u64,
pub payload: ResponsePayload,
}
impl Response {
pub fn ok(id: u64, result: Value) -> Self {
Response {
id,
payload: ResponsePayload::Ok(result),
}
}
pub fn err(id: u64, code: ErrorCode, message: impl Into<String>) -> Self {
Response {
id,
payload: ResponsePayload::Err(RpcError {
code,
message: message.into(),
}),
}
}
pub fn is_err(&self) -> bool {
matches!(self.payload, ResponsePayload::Err(_))
}
pub fn result(&self) -> Option<&Value> {
match &self.payload {
ResponsePayload::Ok(v) => Some(v),
ResponsePayload::Err(_) => None,
}
}
pub fn error(&self) -> Option<&RpcError> {
match &self.payload {
ResponsePayload::Err(e) => Some(e),
ResponsePayload::Ok(_) => None,
}
}
}
impl Serialize for Response {
fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
where
S: Serializer,
{
let mut st = serializer.serialize_struct("Response", 2)?;
st.serialize_field("id", &self.id)?;
match &self.payload {
ResponsePayload::Ok(v) => st.serialize_field("result", v)?,
ResponsePayload::Err(e) => st.serialize_field("error", e)?,
}
st.end()
}
}
impl<'de> Deserialize<'de> for Response {
fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
where
D: Deserializer<'de>,
{
fn present_value<'de, D>(deserializer: D) -> Result<Option<Value>, D::Error>
where
D: Deserializer<'de>,
{
Value::deserialize(deserializer).map(Some)
}
#[derive(Deserialize)]
struct Wire {
id: u64,
#[serde(default, deserialize_with = "present_value")]
result: Option<Value>,
#[serde(default)]
error: Option<RpcError>,
}
let w = Wire::deserialize(deserializer)?;
let payload = match (w.result, w.error) {
(Some(_), Some(_)) => {
return Err(de::Error::custom(
"response carries both `result` and `error`",
))
}
(Some(r), None) => ResponsePayload::Ok(r),
(None, Some(e)) => ResponsePayload::Err(e),
(None, None) => {
return Err(de::Error::custom(
"response carries neither `result` nor `error`",
))
}
};
Ok(Response { id: w.id, payload })
}
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct RpcError {
pub code: ErrorCode,
pub message: String,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum ErrorCode {
MalformedFrame,
UnknownMethod,
InvalidParams,
AgentNotFound,
AgentExists,
InvalidStatus,
Busy,
LockTimeout,
SpawnFailed,
ChannelUnknown,
Internal,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Namespace {
Agent,
Channel,
Unknown,
}
impl Namespace {
pub fn of(method: &str) -> Namespace {
match method.split_once('.') {
Some(("agent", _)) => Namespace::Agent,
Some(("channel", _)) => Namespace::Channel,
_ => Namespace::Unknown,
}
}
pub fn verb(method: &str) -> Option<&str> {
method.split_once('.').map(|(_, v)| v)
}
}
pub async fn read_frame<R: AsyncRead + Unpin>(reader: &mut R) -> Result<Vec<u8>, ProtocolError> {
let mut len_buf = [0u8; 4];
match reader.read_exact(&mut len_buf).await {
Ok(_) => {}
Err(e) if e.kind() == std::io::ErrorKind::UnexpectedEof => {
return Err(ProtocolError::UnexpectedEof)
}
Err(e) => return Err(e.into()),
}
let len = u32::from_le_bytes(len_buf);
if len > MAX_FRAME_BYTES {
return Err(ProtocolError::FrameTooLarge(len));
}
let mut body = vec![0u8; len as usize];
reader
.read_exact(&mut body)
.await
.map_err(|e| match e.kind() {
std::io::ErrorKind::UnexpectedEof => ProtocolError::UnexpectedEof,
_ => ProtocolError::Io(e),
})?;
Ok(body)
}
pub async fn write_frame<W: AsyncWrite + Unpin>(
writer: &mut W,
body: &[u8],
) -> Result<(), ProtocolError> {
let len: u32 = body
.len()
.try_into()
.map_err(|_| ProtocolError::FrameTooLarge(u32::MAX))?;
if len > MAX_FRAME_BYTES {
return Err(ProtocolError::FrameTooLarge(len));
}
writer.write_all(&len.to_le_bytes()).await?;
writer.write_all(body).await?;
writer.flush().await?;
Ok(())
}
pub async fn read_request<R: AsyncRead + Unpin>(reader: &mut R) -> Result<Request, ProtocolError> {
let body = read_frame(reader).await?;
Ok(serde_json::from_slice(&body)?)
}
pub async fn write_request<W: AsyncWrite + Unpin>(
writer: &mut W,
req: &Request,
) -> Result<(), ProtocolError> {
let body = serde_json::to_vec(req)?;
write_frame(writer, &body).await
}
pub async fn read_response<R: AsyncRead + Unpin>(
reader: &mut R,
) -> Result<Response, ProtocolError> {
let body = read_frame(reader).await?;
Ok(serde_json::from_slice(&body)?)
}
pub async fn write_response<W: AsyncWrite + Unpin>(
writer: &mut W,
resp: &Response,
) -> Result<(), ProtocolError> {
let body = serde_json::to_vec(resp)?;
write_frame(writer, &body).await
}
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
#[test]
fn namespace_classification() {
assert_eq!(Namespace::of("agent.spawn"), Namespace::Agent);
assert_eq!(
Namespace::of("channel.register_channel"),
Namespace::Channel
);
assert_eq!(Namespace::of("bogus.method"), Namespace::Unknown);
assert_eq!(Namespace::of("noseparator"), Namespace::Unknown);
assert_eq!(Namespace::verb("agent.spawn"), Some("spawn"));
assert_eq!(Namespace::verb("nope"), None);
}
#[tokio::test]
async fn request_roundtrips_over_duplex() {
let (mut a, mut b) = tokio::io::duplex(4096);
let req = Request::new(7, "agent.spawn", json!({"name": "worker-A"}));
write_request(&mut a, &req).await.unwrap();
let got = read_request(&mut b).await.unwrap();
assert_eq!(got, req);
}
#[tokio::test]
async fn response_ok_and_err_roundtrip() {
let (mut a, mut b) = tokio::io::duplex(4096);
let ok = Response::ok(7, json!({"status": "live"}));
write_response(&mut a, &ok).await.unwrap();
let got = read_response(&mut b).await.unwrap();
assert_eq!(got, ok);
assert!(!got.is_err());
let err = Response::err(8, ErrorCode::AgentNotFound, "no such agent");
write_response(&mut a, &err).await.unwrap();
let got = read_response(&mut b).await.unwrap();
assert!(got.is_err());
assert_eq!(got.error().unwrap().code, ErrorCode::AgentNotFound);
}
#[tokio::test]
async fn frame_too_large_is_rejected_not_allocated() {
let (mut a, mut b) = tokio::io::duplex(64);
let writer = tokio::spawn(async move {
let bogus_len = (MAX_FRAME_BYTES + 1).to_le_bytes();
a.write_all(&bogus_len).await.unwrap();
a.flush().await.unwrap();
a
});
let err = read_frame(&mut b).await.unwrap_err();
assert!(matches!(err, ProtocolError::FrameTooLarge(_)));
let _a = writer.await.unwrap();
}
#[tokio::test]
async fn clean_eof_is_distinguished_from_io_error() {
let (a, mut b) = tokio::io::duplex(64);
drop(a); let err = read_frame(&mut b).await.unwrap_err();
assert!(matches!(err, ProtocolError::UnexpectedEof));
}
#[test]
fn response_wire_shape_is_flat() {
let ok = Response::ok(7, json!({"status": "live"}));
assert_eq!(
serde_json::to_value(&ok).unwrap(),
json!({"id": 7, "result": {"status": "live"}})
);
let err = Response::err(8, ErrorCode::AgentNotFound, "no such agent");
assert_eq!(
serde_json::to_value(&err).unwrap(),
json!({"id": 8, "error": {"code": "agent_not_found", "message": "no such agent"}})
);
}
#[test]
fn response_rejects_both_or_neither_payload() {
let both = json!({"id": 1, "result": {}, "error": {"code": "internal", "message": "x"}});
assert!(serde_json::from_value::<Response>(both).is_err());
let neither = json!({"id": 1});
assert!(serde_json::from_value::<Response>(neither).is_err());
}
#[test]
fn response_accepts_explicit_null_result() {
let parsed: Response =
serde_json::from_value(json!({"id": 7, "result": null})).expect("null result parses");
assert!(!parsed.is_err());
assert_eq!(parsed.result(), Some(&Value::Null));
assert_eq!(
serde_json::to_value(&parsed).unwrap(),
json!({"id": 7, "result": null})
);
}
#[tokio::test]
async fn malformed_json_body_surfaces_json_error() {
let (mut a, mut b) = tokio::io::duplex(4096);
write_frame(&mut a, b"{not json").await.unwrap();
let err = read_request(&mut b).await.unwrap_err();
assert!(matches!(err, ProtocolError::Json(_)));
}
}