feat: neuron-rs — Rust runtime, Engram-backed, Axon protocol

crates/neuron-protocol/src/axon.rs

@@ -0,0 +1,91 @@
+use serde::{Deserialize, Serialize};
+use serde_json::Value;
+
+/// An Axon protocol message — wraps a tool call or other method invocation.
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct AxonMessage {
+    /// Unique message ID
+    pub id: String,
+    /// Method name: "tool_call", "tool_result", "ping", etc.
+    pub method: String,
+    /// Method-specific parameters
+    pub params: Value,
+}
+
+/// 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() }),
+        }
+    }
+}
+
+/// 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(),
+        }
+    }
+}
+
+#[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() {
+        let msg = AxonMessage {
+            id: "m1".into(),
+            method: "tool_call".into(),
+            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, "tool_call");
+    }
+}

crates/neuron-protocol/src/handler.rs

@@ -0,0 +1,142 @@
+use crate::axon::{AxonMessage, AxonResponse};
+use serde_json::Value;
+use std::collections::HashMap;
+use std::sync::Arc;
+
+type ToolFn = Arc<dyn Fn(Value) -> Value + Send + Sync>;
+
+/// Routes Axon method names to handler functions.
+///
+/// Each handler takes the `params` field of an AxonMessage and returns
+/// a JSON Value to be wrapped in an AxonResponse.
+pub struct AxonHandler {
+    tools: HashMap<String, ToolFn>,
+}
+
+impl AxonHandler {
+    pub fn new() -> Self {
+        Self {
+            tools: HashMap::new(),
+        }
+    }
+
+    /// Register a tool handler by name.
+    pub fn register<F>(&mut self, name: impl Into<String>, f: F)
+    where
+        F: Fn(Value) -> Value + Send + Sync + 'static,
+    {
+        self.tools.insert(name.into(), Arc::new(f));
+    }
+
+    /// Dispatch an AxonMessage to the appropriate handler.
+    pub fn dispatch(&self, msg: AxonMessage) -> AxonResponse {
+        match msg.method.as_str() {
+            "tool_call" => {
+                let tool_name = msg
+                    .params
+                    .get("tool")
+                    .and_then(|v| v.as_str())
+                    .unwrap_or("");
+                let tool_params = msg
+                    .params
+                    .get("params")
+                    .cloned()
+                    .unwrap_or(Value::Object(Default::default()));
+
+                if let Some(f) = self.tools.get(tool_name) {
+                    let result = f(tool_params);
+                    AxonResponse::ok(msg.id, result)
+                } else {
+                    AxonResponse::err(msg.id, format!("unknown tool: {}", tool_name))
+                }
+            }
+            "ping" => AxonResponse::ok(msg.id, serde_json::json!({ "pong": true })),
+            other => AxonResponse::err(msg.id, format!("unknown method: {}", other)),
+        }
+    }
+
+    /// List all registered tool names.
+    pub fn tool_names(&self) -> Vec<String> {
+        let mut names: Vec<_> = self.tools.keys().cloned().collect();
+        names.sort();
+        names
+    }
+}
+
+impl Default for AxonHandler {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn dispatch_ping() {
+        let handler = AxonHandler::new();
+        let msg = AxonMessage {
+            id: "1".into(),
+            method: "ping".into(),
+            params: serde_json::json!({}),
+        };
+        let resp = handler.dispatch(msg);
+        assert!(resp.success);
+        assert_eq!(resp.result["pong"], true);
+    }
+
+    #[test]
+    fn dispatch_tool_call() {
+        let mut handler = AxonHandler::new();
+        handler.register("remember", |params| {
+            serde_json::json!({ "stored": params.get("content").cloned().unwrap_or_default() })
+        });
+
+        let msg = AxonMessage {
+            id: "2".into(),
+            method: "tool_call".into(),
+            params: serde_json::json!({ "tool": "remember", "params": { "content": "hello" } }),
+        };
+        let resp = handler.dispatch(msg);
+        assert!(resp.success);
+        assert_eq!(resp.result["stored"], "hello");
+    }
+
+    #[test]
+    fn dispatch_unknown_tool() {
+        let handler = AxonHandler::new();
+        let msg = AxonMessage {
+            id: "3".into(),
+            method: "tool_call".into(),
+            params: serde_json::json!({ "tool": "nonexistent", "params": {} }),
+        };
+        let resp = handler.dispatch(msg);
+        assert!(!resp.success);
+    }
+
+    #[test]
+    fn dispatch_unknown_method() {
+        let handler = AxonHandler::new();
+        let msg = AxonMessage {
+            id: "4".into(),
+            method: "mystery".into(),
+            params: serde_json::json!({}),
+        };
+        let resp = handler.dispatch(msg);
+        assert!(!resp.success);
+        assert!(resp.result["error"].as_str().unwrap().contains("unknown method"));
+    }
+
+    #[test]
+    fn tool_names_sorted() {
+        let mut handler = AxonHandler::new();
+        handler.register("remember", |_| Value::Null);
+        handler.register("recall", |_| Value::Null);
+        handler.register("plan_work", |_| Value::Null);
+        let names = handler.tool_names();
+        let mut expected = vec!["plan_work", "recall", "remember"];
+        expected.sort();
+        assert_eq!(names, expected);
+    }
+}

crates/neuron-protocol/src/lib.rs

@@ -0,0 +1,7 @@
+pub mod axon;
+pub mod handler;
+pub mod sse;
+pub mod tools;
+
+pub use axon::{AxonMessage, AxonResponse};
+pub use handler::AxonHandler;

crates/neuron-protocol/src/sse.rs

@@ -0,0 +1,47 @@
+use crate::axon::AxonEvent;
+use tokio::sync::broadcast;
+
+/// SSE channel — broadcasts AxonEvents to all connected SSE clients.
+#[derive(Clone)]
+pub struct SseChannel {
+    sender: broadcast::Sender<String>,
+}
+
+impl SseChannel {
+    pub fn new(capacity: usize) -> Self {
+        let (sender, _) = broadcast::channel(capacity);
+        Self { sender }
+    }
+
+    /// Send an event to all subscribers.
+    pub fn send(&self, event: &AxonEvent) -> Result<(), String> {
+        let json = serde_json::to_string(event).map_err(|e| e.to_string())?;
+        // Ignore errors when no receivers are connected
+        let _ = self.sender.send(json);
+        Ok(())
+    }
+
+    /// Subscribe to the SSE channel.
+    pub fn subscribe(&self) -> broadcast::Receiver<String> {
+        self.sender.subscribe()
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[tokio::test]
+    async fn sse_channel_broadcast() {
+        let ch = SseChannel::new(16);
+        let mut rx = ch.subscribe();
+
+        let event = AxonEvent::new("test", serde_json::json!({ "msg": "hello" }));
+        ch.send(&event).unwrap();
+
+        let received = rx.recv().await.unwrap();
+        let parsed: serde_json::Value = serde_json::from_str(&received).unwrap();
+        assert_eq!(parsed["event_type"], "test");
+        assert_eq!(parsed["payload"]["msg"], "hello");
+    }
+}

crates/neuron-protocol/src/tools.rs

@@ -0,0 +1,51 @@
+/// Tool name constants — must match the existing MCP tool names so Claude Code
+/// can switch over without configuration changes.
+pub const REMEMBER: &str = "remember";
+pub const RECALL: &str = "recall";
+pub const SEARCH_ENTITIES: &str = "search_entities";
+pub const CAPTURE_KNOWLEDGE: &str = "capture_knowledge";
+pub const SEARCH_KNOWLEDGE: &str = "search_knowledge";
+pub const RETRIEVE_KNOWLEDGE: &str = "retrieve_knowledge";
+pub const PLAN_WORK: &str = "plan_work";
+pub const REVIEW_BACKLOG: &str = "review_backlog";
+pub const TRACK_WORK: &str = "track_work";
+pub const BEGIN_WORK: &str = "begin_work";
+pub const PROGRESS_WORK: &str = "progress_work";
+pub const CHECK_WORK: &str = "check_work";
+pub const LOG_INTERNAL_STATE_EVENT: &str = "log_internal_state_event";
+pub const LIST_INTERNAL_STATE_EVENTS: &str = "list_internal_state_events";
+
+/// All registered tool names in canonical order.
+pub const ALL_TOOLS: &[&str] = &[
+    REMEMBER,
+    RECALL,
+    SEARCH_ENTITIES,
+    CAPTURE_KNOWLEDGE,
+    SEARCH_KNOWLEDGE,
+    RETRIEVE_KNOWLEDGE,
+    PLAN_WORK,
+    REVIEW_BACKLOG,
+    TRACK_WORK,
+    BEGIN_WORK,
+    PROGRESS_WORK,
+    CHECK_WORK,
+    LOG_INTERNAL_STATE_EVENT,
+    LIST_INTERNAL_STATE_EVENTS,
+];
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn all_tools_count() {
+        assert_eq!(ALL_TOOLS.len(), 14);
+    }
+
+    #[test]
+    fn tool_names_are_non_empty() {
+        for name in ALL_TOOLS {
+            assert!(!name.is_empty(), "tool name must not be empty");
+        }
+    }
+}