Archived
267 lines
7.8 KiB
Rust
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);
|
|
}
|
|
}
|