use std::fmt;
use async_trait::async_trait;
use uuid::Uuid;
use khive_runtime::Namespace;
use khive_runtime::{BackendId, NoteSearchHit, SearchHit};
use khive_storage::{Edge, EdgeRelation};
pub struct CoordLinkResult {
pub edge: Edge,
pub cross_backend: bool,
pub target_backend_id: Option<BackendId>,
}
pub enum CoordError {
UnknownNode { id: Uuid },
EdgeRuleViolation(String),
Backend(String),
}
impl fmt::Display for CoordError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
CoordError::UnknownNode { id } => write!(f, "node {id} not found on any backend"),
CoordError::EdgeRuleViolation(msg) => write!(f, "edge rule violation: {msg}"),
CoordError::Backend(msg) => write!(f, "backend error: {msg}"),
}
}
}
impl fmt::Debug for CoordError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
fmt::Display::fmt(self, f)
}
}
impl From<CoordError> for khive_runtime::RuntimeError {
fn from(e: CoordError) -> Self {
match e {
CoordError::UnknownNode { id } => {
khive_runtime::RuntimeError::NotFound(format!("node {id} not found on any backend"))
}
CoordError::EdgeRuleViolation(msg) => khive_runtime::RuntimeError::InvalidInput(msg),
CoordError::Backend(msg) => khive_runtime::RuntimeError::Internal(msg),
}
}
}
pub struct BackendSearchResult {
pub backend_id: BackendId,
pub entity_hits: Vec<SearchHit>,
pub note_hits: Vec<NoteSearchHit>,
pub error: Option<String>,
}
pub struct CoordSearchResult {
pub entity_hits: Vec<SearchHit>,
pub note_hits: Vec<NoteSearchHit>,
pub per_backend: Vec<BackendSearchResult>,
pub partial: bool,
pub entity_kinds: std::collections::HashMap<uuid::Uuid, String>,
pub note_kinds: std::collections::HashMap<uuid::Uuid, String>,
}
#[async_trait]
pub trait CoordinatorService: Send + Sync {
async fn locate(&self, id: Uuid) -> Option<BackendId>;
fn record_created(&self, id: Uuid, backend_id: BackendId);
fn primary_backend_id(&self) -> Option<BackendId>;
async fn link(
&self,
namespace: &Namespace,
source_id: Uuid,
target_id: Uuid,
relation: EdgeRelation,
weight: f64,
metadata: Option<serde_json::Value>,
) -> Result<CoordLinkResult, CoordError>;
#[allow(clippy::too_many_arguments)]
async fn fan_out_search(
&self,
kind: &str,
query: &str,
namespace: &Namespace,
limit: u32,
kind_filter: Option<&str>,
props_filter: Option<&serde_json::Value>,
tags: &[String],
) -> CoordSearchResult;
fn is_single_backend(&self) -> bool;
}
#[cfg(test)]
pub(crate) mod tests {
use super::*;
use std::sync::Arc;
pub struct MockCoordinator {
pub link_called: std::sync::atomic::AtomicBool,
pub search_called: std::sync::atomic::AtomicBool,
pub single_backend: bool,
pub last_limit: std::sync::atomic::AtomicU32,
}
impl MockCoordinator {
pub fn multi_backend() -> Arc<Self> {
Arc::new(Self {
link_called: std::sync::atomic::AtomicBool::new(false),
search_called: std::sync::atomic::AtomicBool::new(false),
single_backend: false,
last_limit: std::sync::atomic::AtomicU32::new(0),
})
}
pub fn single_backend_instance() -> Arc<Self> {
Arc::new(Self {
link_called: std::sync::atomic::AtomicBool::new(false),
search_called: std::sync::atomic::AtomicBool::new(false),
single_backend: true,
last_limit: std::sync::atomic::AtomicU32::new(0),
})
}
}
#[async_trait]
impl CoordinatorService for MockCoordinator {
async fn locate(&self, _id: Uuid) -> Option<BackendId> {
Some(BackendId::main())
}
fn record_created(&self, _id: Uuid, _backend_id: BackendId) {}
fn primary_backend_id(&self) -> Option<BackendId> {
Some(BackendId::main())
}
async fn link(
&self,
_namespace: &Namespace,
_source_id: Uuid,
_target_id: Uuid,
_relation: EdgeRelation,
_weight: f64,
_metadata: Option<serde_json::Value>,
) -> Result<CoordLinkResult, CoordError> {
self.link_called
.store(true, std::sync::atomic::Ordering::SeqCst);
Err(CoordError::UnknownNode { id: Uuid::new_v4() })
}
async fn fan_out_search(
&self,
_kind: &str,
_query: &str,
_namespace: &Namespace,
limit: u32,
_kind_filter: Option<&str>,
_props_filter: Option<&serde_json::Value>,
_tags: &[String],
) -> CoordSearchResult {
self.search_called
.store(true, std::sync::atomic::Ordering::SeqCst);
self.last_limit
.store(limit, std::sync::atomic::Ordering::SeqCst);
CoordSearchResult {
entity_hits: vec![],
note_hits: vec![],
per_backend: vec![],
partial: false,
entity_kinds: std::collections::HashMap::new(),
note_kinds: std::collections::HashMap::new(),
}
}
fn is_single_backend(&self) -> bool {
self.single_backend
}
}
use crate::server::KhiveMcpServer;
use crate::tools::request::RequestParams;
use khive_runtime::{KhiveRuntime, Namespace as RuntimeNamespace, RuntimeConfig};
fn make_registry() -> (khive_runtime::VerbRegistry, khive_runtime::KhiveRuntime) {
let config = RuntimeConfig {
db_path: None,
default_namespace: RuntimeNamespace::parse("local").unwrap(),
embedding_model: None,
additional_embedding_models: vec![],
packs: vec!["kg".to_string()],
..RuntimeConfig::default()
};
let runtime = KhiveRuntime::new(config).expect("in-memory runtime");
let gate = runtime.config().gate.clone();
let default_ns = runtime.config().default_namespace.clone();
let actor_id = runtime.config().actor_id.clone();
let mut builder = khive_runtime::VerbRegistryBuilder::new();
builder.with_gate(gate);
builder.with_default_namespace(default_ns.as_str());
builder.with_actor_id(actor_id);
khive_runtime::PackRegistry::register_packs(
&["kg".to_string()],
runtime.clone(),
&mut builder,
)
.expect("register kg");
let registry = builder.build().expect("build registry");
runtime.install_edge_rules(registry.all_edge_rules());
(registry, runtime)
}
#[tokio::test]
async fn t6a_multi_backend_server_routes_link_through_coordinator() {
let (registry, _runtime) = make_registry();
let coord = MockCoordinator::multi_backend();
let server = KhiveMcpServer::from_registry_with_meta(registry, "local", "test-cfg")
.with_coordinator(Arc::clone(&coord) as Arc<dyn CoordinatorService>);
let src_id = Uuid::new_v4();
let tgt_id = Uuid::new_v4();
let ops = format!(
r#"link(source_id="{}", target_id="{}", relation="implements")"#,
src_id, tgt_id
);
let _result = server
.dispatch_request_local(RequestParams {
ops,
presentation: None,
presentation_per_op: None,
save_to: None,
format: None,
format_per_op: None,
})
.await;
assert!(
coord
.link_called
.load(std::sync::atomic::Ordering::SeqCst),
"T6a: coordinator.link must be called when a link op is dispatched through a multi-backend server"
);
}
#[tokio::test]
async fn t6b_multi_backend_server_routes_search_through_coordinator() {
let (registry, _runtime) = make_registry();
let coord = MockCoordinator::multi_backend();
let server = KhiveMcpServer::from_registry_with_meta(registry, "local", "test-cfg")
.with_coordinator(Arc::clone(&coord) as Arc<dyn CoordinatorService>);
let _result = server
.dispatch_request_local(RequestParams {
ops: r#"search(kind="entity", query="anything")"#.to_string(),
presentation: None,
presentation_per_op: None,
save_to: None,
format: None,
format_per_op: None,
})
.await;
assert!(
coord
.search_called
.load(std::sync::atomic::Ordering::SeqCst),
"T6b: coordinator.fan_out_search must be called when a search op is dispatched through a multi-backend server"
);
}
#[tokio::test]
async fn t6d_malformed_tags_return_per_op_error_in_multi_backend() {
let (registry, _runtime) = make_registry();
let coord = MockCoordinator::multi_backend();
let server = KhiveMcpServer::from_registry_with_meta(registry, "local", "test-cfg")
.with_coordinator(Arc::clone(&coord) as Arc<dyn CoordinatorService>);
let raw = server
.dispatch_request_local(RequestParams {
ops: r#"search(kind="entity", query="anything", tags=[42])"#.to_string(),
presentation: None,
presentation_per_op: None,
save_to: None,
format: None,
format_per_op: None,
})
.await
.expect("T6d: dispatch must not return an MCP-level error");
let result_val: serde_json::Value =
serde_json::from_str(&raw).expect("T6d: response must be valid JSON");
let first = result_val
.get("results")
.and_then(|r| r.as_array())
.and_then(|a| a.first())
.expect("T6d: results array must be non-empty");
assert_eq!(
first.get("ok").and_then(serde_json::Value::as_bool),
Some(false),
"T6d: malformed tags must produce ok=false; got {:?}",
first
);
assert!(
!coord
.search_called
.load(std::sync::atomic::Ordering::SeqCst),
"T6d: coordinator must not be reached when tags validation fails"
);
}
#[tokio::test]
async fn t6e_multi_backend_search_malformed_namespace_fails_closed() {
let cases: [(&str, &str); 5] = [
("null", "null"),
("number", "42"),
("boolean", "true"),
("array", r#"["local"]"#),
("object", r#"{"ns":"local"}"#),
];
for (label, ns_literal) in cases {
let (registry, _runtime) = make_registry();
let coord = MockCoordinator::multi_backend();
let server = KhiveMcpServer::from_registry_with_meta(registry, "local", "test-cfg")
.with_coordinator(Arc::clone(&coord) as Arc<dyn CoordinatorService>);
let ops = format!(r#"search(kind="entity", query="anything", namespace={ns_literal})"#);
let raw = server
.dispatch_request_local(RequestParams {
ops,
presentation: None,
presentation_per_op: None,
save_to: None,
format: None,
format_per_op: None,
})
.await
.unwrap_or_else(|e| panic!("T6e case {label}: dispatch must not MCP-error: {e}"));
let result_val: serde_json::Value =
serde_json::from_str(&raw).expect("T6e: response must be valid JSON");
let first = result_val
.get("results")
.and_then(|r| r.as_array())
.and_then(|a| a.first())
.unwrap_or_else(|| panic!("T6e case {label}: results array must be non-empty"));
assert_eq!(
first.get("ok").and_then(serde_json::Value::as_bool),
Some(false),
"T6e case {label}: malformed namespace must fail closed; got {first:?}"
);
let err_text = first.get("error").map(|e| e.to_string().to_lowercase());
assert!(
err_text.as_deref().is_some_and(|e| e.contains("namespace")),
"T6e case {label}: error must name the namespace; got {first:?}"
);
assert!(
!coord
.search_called
.load(std::sync::atomic::Ordering::SeqCst),
"T6e case {label}: coordinator.fan_out_search must NOT be called for a malformed namespace"
);
}
}
#[tokio::test]
async fn t6f_multi_backend_link_malformed_namespace_fails_closed() {
let cases: [(&str, &str); 5] = [
("null", "null"),
("number", "42"),
("boolean", "true"),
("array", r#"["local"]"#),
("object", r#"{"ns":"local"}"#),
];
for (label, ns_literal) in cases {
let (registry, _runtime) = make_registry();
let coord = MockCoordinator::multi_backend();
let server = KhiveMcpServer::from_registry_with_meta(registry, "local", "test-cfg")
.with_coordinator(Arc::clone(&coord) as Arc<dyn CoordinatorService>);
let src_id = Uuid::new_v4();
let tgt_id = Uuid::new_v4();
let ops = format!(
r#"link(source_id="{src_id}", target_id="{tgt_id}", relation="implements", namespace={ns_literal})"#
);
let raw = server
.dispatch_request_local(RequestParams {
ops,
presentation: None,
presentation_per_op: None,
save_to: None,
format: None,
format_per_op: None,
})
.await
.unwrap_or_else(|e| panic!("T6f case {label}: dispatch must not MCP-error: {e}"));
let result_val: serde_json::Value =
serde_json::from_str(&raw).expect("T6f: response must be valid JSON");
let first = result_val
.get("results")
.and_then(|r| r.as_array())
.and_then(|a| a.first())
.unwrap_or_else(|| panic!("T6f case {label}: results array must be non-empty"));
assert_eq!(
first.get("ok").and_then(serde_json::Value::as_bool),
Some(false),
"T6f case {label}: malformed namespace must fail closed; got {first:?}"
);
let err_text = first.get("error").map(|e| e.to_string().to_lowercase());
assert!(
err_text.as_deref().is_some_and(|e| e.contains("namespace")),
"T6f case {label}: error must name the namespace; got {first:?}"
);
assert!(
!coord.link_called.load(std::sync::atomic::Ordering::SeqCst),
"T6f case {label}: coordinator.link must NOT be called for a malformed namespace"
);
}
}
#[tokio::test]
async fn t6c_single_backend_server_bypasses_coordinator() {
let (registry, runtime) = make_registry();
let coord = MockCoordinator::single_backend_instance();
let server = KhiveMcpServer::from_registry_with_meta(registry, "local", "test-cfg")
.with_coordinator(Arc::clone(&coord) as Arc<dyn CoordinatorService>);
let ns = RuntimeNamespace::local();
let token = runtime.authorize(ns).expect("authorize");
let entity = runtime
.create_entity(&token, "concept", None, "T6cEntity", None, None, vec![])
.await
.expect("create entity");
let _ = entity;
let _result = server
.dispatch_request_local(RequestParams {
ops: r#"search(kind="entity", query="T6cEntity")"#.to_string(),
presentation: None,
presentation_per_op: None,
save_to: None,
format: None,
format_per_op: None,
})
.await;
assert!(
!coord
.search_called
.load(std::sync::atomic::Ordering::SeqCst),
"T6c: coordinator.fan_out_search must NOT be called for a single-backend server"
);
assert!(
!coord.link_called.load(std::sync::atomic::Ordering::SeqCst),
"T6c: coordinator.link must NOT be called for a single-backend server"
);
}
#[tokio::test]
async fn t6e_multi_backend_search_limit_matches_single_backend_u32_contract() {
let (registry, _runtime) = make_registry();
let coord = MockCoordinator::multi_backend();
let server = KhiveMcpServer::from_registry_with_meta(registry, "local", "test-cfg")
.with_coordinator(Arc::clone(&coord) as Arc<dyn CoordinatorService>);
let too_large: u64 = u64::from(u32::MAX) + 2;
let raw = server
.dispatch_request_local(RequestParams {
ops: format!(r#"search(kind="entity", query="anything", limit={too_large})"#),
presentation: None,
presentation_per_op: None,
save_to: None,
format: None,
format_per_op: None,
})
.await
.expect("T6e: dispatch must not return an MCP-level error");
let result_val: serde_json::Value =
serde_json::from_str(&raw).expect("T6e: response must be valid JSON");
let first = result_val
.get("results")
.and_then(|r| r.as_array())
.and_then(|a| a.first())
.expect("T6e: results array must be non-empty");
assert_eq!(
first.get("ok").and_then(serde_json::Value::as_bool),
Some(false),
"T6e: an out-of-range limit must produce ok=false; got {:?}",
first
);
assert!(
!coord
.search_called
.load(std::sync::atomic::Ordering::SeqCst),
"T6e: coordinator must not be called with an out-of-range limit \
(it must not silently wrap to a small value); recorded last_limit={}",
coord.last_limit.load(std::sync::atomic::Ordering::SeqCst)
);
}
#[tokio::test]
async fn t6e_multi_backend_search_limit_u32_max_is_capped_at_100() {
let (registry, _runtime) = make_registry();
let coord = MockCoordinator::multi_backend();
let server = KhiveMcpServer::from_registry_with_meta(registry, "local", "test-cfg")
.with_coordinator(Arc::clone(&coord) as Arc<dyn CoordinatorService>);
let raw = server
.dispatch_request_local(RequestParams {
ops: format!(
r#"search(kind="entity", query="anything", limit={})"#,
u32::MAX
),
presentation: None,
presentation_per_op: None,
save_to: None,
format: None,
format_per_op: None,
})
.await
.expect("T6e: dispatch must not return an MCP-level error");
let _ = raw;
assert!(
coord
.search_called
.load(std::sync::atomic::Ordering::SeqCst),
"T6e: coordinator.fan_out_search must be called for a valid in-range limit"
);
assert_eq!(
coord.last_limit.load(std::sync::atomic::Ordering::SeqCst),
100,
"T6e: u32::MAX must be capped at 100 before reaching the coordinator"
);
}
}