This repository has been archived on 2026-08-20. You can view files and clone it. You cannot open issues or pull requests or push a commit.
Files
neuron-rs/crates/neuron-protocol/src/axon.rs
T

267 lines
7.8 KiB
Rust

use serde::{Deserialize, Serialize};
use serde_json::Value;
// ------------------------------------------------------------------
// AxonMethod
// ------------------------------------------------------------------
/// Typed method discriminant for Axon protocol messages.
///
/// Replaces the raw `method: String` field — callers that need a string
/// representation can use [`AxonMethod::as_str`] or the `Display` impl.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum AxonMethod {
/// Invoke a named tool.
ToolCall,
/// Return value from a tool invocation.
ToolResult,
/// Deliver a typed Axon event envelope.
Event,
/// Liveness check.
Ping,
/// Liveness reply.
Pong,
/// Subscribe to an event stream.
Subscribe,
/// Cancel a subscription.
Unsubscribe,
}
impl AxonMethod {
/// Return the canonical lowercase string for this method.
pub fn as_str(self) -> &'static str {
match self {
Self::ToolCall => "tool_call",
Self::ToolResult => "tool_result",
Self::Event => "event",
Self::Ping => "ping",
Self::Pong => "pong",
Self::Subscribe => "subscribe",
Self::Unsubscribe => "unsubscribe",
}
}
/// Parse a string into an `AxonMethod`, returning `None` on mismatch.
pub fn from_str(s: &str) -> Option<Self> {
match s {
"tool_call" => Some(Self::ToolCall),
"tool_result" => Some(Self::ToolResult),
"event" => Some(Self::Event),
"ping" => Some(Self::Ping),
"pong" => Some(Self::Pong),
"subscribe" => Some(Self::Subscribe),
"unsubscribe" => Some(Self::Unsubscribe),
_ => None,
}
}
}
impl std::fmt::Display for AxonMethod {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.write_str(self.as_str())
}
}
// ------------------------------------------------------------------
// AxonMessage
// ------------------------------------------------------------------
/// An Axon protocol message — wraps a tool call, event delivery, or
/// other method invocation.
///
/// The `method` field is now typed as [`AxonMethod`]. For backward
/// compatibility the raw string representation is available via
/// [`AxonMessage::method_str`].
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AxonMessage {
/// Unique message ID.
pub id: String,
/// Typed method discriminant.
pub method: AxonMethod,
/// Method-specific parameters.
pub params: Value,
}
impl AxonMessage {
/// Return the canonical string representation of the method (e.g.
/// `"tool_call"`, `"ping"`). Useful when serialising to a transport
/// that expects a raw string rather than a tagged enum.
pub fn method_str(&self) -> &'static str {
self.method.as_str()
}
/// Construct a ping message.
pub fn ping(id: impl Into<String>) -> Self {
Self {
id: id.into(),
method: AxonMethod::Ping,
params: Value::Object(Default::default()),
}
}
}
// ------------------------------------------------------------------
// AxonResponse
// ------------------------------------------------------------------
/// An Axon protocol response.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AxonResponse {
/// Echo of the request ID.
pub id: String,
/// Whether the call succeeded.
pub success: bool,
/// The result payload (on success) or error message (on failure).
pub result: Value,
}
impl AxonResponse {
pub fn ok(id: impl Into<String>, result: Value) -> Self {
Self {
id: id.into(),
success: true,
result,
}
}
pub fn err(id: impl Into<String>, message: impl Into<String>) -> Self {
Self {
id: id.into(),
success: false,
result: serde_json::json!({ "error": message.into() }),
}
}
}
// ------------------------------------------------------------------
// AxonEvent (legacy SSE event — kept for backward compat)
// ------------------------------------------------------------------
/// An Axon SSE event — sent over the streaming channel.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AxonEvent {
pub event_type: String,
pub payload: Value,
pub timestamp: i64,
}
impl AxonEvent {
pub fn new(event_type: impl Into<String>, payload: Value) -> Self {
Self {
event_type: event_type.into(),
payload,
timestamp: neuron_domain::now_ms(),
}
}
}
// ------------------------------------------------------------------
// Tests
// ------------------------------------------------------------------
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn axon_response_ok() {
let r = AxonResponse::ok("req-1", serde_json::json!({"id": "abc"}));
assert!(r.success);
assert_eq!(r.id, "req-1");
}
#[test]
fn axon_response_err() {
let r = AxonResponse::err("req-2", "not found");
assert!(!r.success);
assert_eq!(r.result["error"], "not found");
}
#[test]
fn axon_message_roundtrip_with_method_enum() {
let msg = AxonMessage {
id: "m1".into(),
method: AxonMethod::ToolCall,
params: serde_json::json!({ "tool": "remember", "content": "hello" }),
};
let json = serde_json::to_string(&msg).unwrap();
let back: AxonMessage = serde_json::from_str(&json).unwrap();
assert_eq!(back.method, AxonMethod::ToolCall);
assert_eq!(back.method_str(), "tool_call");
}
#[test]
fn axon_method_as_str_all_variants() {
let cases = [
(AxonMethod::ToolCall, "tool_call"),
(AxonMethod::ToolResult, "tool_result"),
(AxonMethod::Event, "event"),
(AxonMethod::Ping, "ping"),
(AxonMethod::Pong, "pong"),
(AxonMethod::Subscribe, "subscribe"),
(AxonMethod::Unsubscribe, "unsubscribe"),
];
for (method, expected) in cases {
assert_eq!(method.as_str(), expected);
assert_eq!(method.to_string(), expected);
}
}
#[test]
fn axon_method_from_str_roundtrip() {
let variants = [
AxonMethod::ToolCall,
AxonMethod::ToolResult,
AxonMethod::Event,
AxonMethod::Ping,
AxonMethod::Pong,
AxonMethod::Subscribe,
AxonMethod::Unsubscribe,
];
for v in variants {
let s = v.as_str();
assert_eq!(AxonMethod::from_str(s), Some(v));
}
}
#[test]
fn axon_method_from_str_unknown() {
assert_eq!(AxonMethod::from_str("mystery"), None);
assert_eq!(AxonMethod::from_str(""), None);
}
#[test]
fn axon_method_serde_roundtrip() {
let m = AxonMethod::Subscribe;
let json = serde_json::to_string(&m).unwrap();
let back: AxonMethod = serde_json::from_str(&json).unwrap();
assert_eq!(back, AxonMethod::Subscribe);
}
#[test]
fn axon_message_ping_helper() {
let msg = AxonMessage::ping("p-1");
assert_eq!(msg.method, AxonMethod::Ping);
assert_eq!(msg.id, "p-1");
}
#[test]
fn axon_method_all_variants_have_unique_str() {
use std::collections::HashSet;
let strs: HashSet<&str> = [
AxonMethod::ToolCall,
AxonMethod::ToolResult,
AxonMethod::Event,
AxonMethod::Ping,
AxonMethod::Pong,
AxonMethod::Subscribe,
AxonMethod::Unsubscribe,
]
.iter()
.map(|m| m.as_str())
.collect();
assert_eq!(strs.len(), 7);
}
}