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

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);
+    }
+}