147 lines
4.3 KiB
Rust
147 lines
4.3 KiB
Rust
use crate::axon::{AxonMessage, AxonMethod, 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 {
|
|
AxonMethod::ToolCall => {
|
|
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))
|
|
}
|
|
}
|
|
AxonMethod::Ping => {
|
|
AxonResponse::ok(msg.id, serde_json::json!({ "pong": true }))
|
|
}
|
|
other => {
|
|
AxonResponse::err(msg.id, format!("unknown method: {}", other.as_str()))
|
|
}
|
|
}
|
|
}
|
|
|
|
/// 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: AxonMethod::Ping,
|
|
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: AxonMethod::ToolCall,
|
|
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: AxonMethod::ToolCall,
|
|
params: serde_json::json!({ "tool": "nonexistent", "params": {} }),
|
|
};
|
|
let resp = handler.dispatch(msg);
|
|
assert!(!resp.success);
|
|
}
|
|
|
|
#[test]
|
|
fn dispatch_unhandled_method() {
|
|
let handler = AxonHandler::new();
|
|
let msg = AxonMessage {
|
|
id: "4".into(),
|
|
method: AxonMethod::Subscribe,
|
|
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);
|
|
}
|
|
}
|