use std::{
future::Future,
sync::{
atomic::{AtomicI64, Ordering},
Arc,
},
};
use futures::{stream::FuturesUnordered, StreamExt};
use rmcp::{
handler::server::wrapper::Parameters,
model::{Implementation, ServerCapabilities, ServerInfo},
tool, tool_handler, tool_router, ErrorData as McpError, ServerHandler, ServiceExt,
};
use serde_json::{json, Value};
use sha2::{Digest, Sha256};
use khive_db::ConnectionPool;
use khive_request::{parse_request, ArgValue, DslError, ExecutionMode, ParsedOp, PrevFailure};
use khive_runtime::{
present, render_format, KhiveRuntime, OutputFormat, PackLoadError, PackRegistry,
PresentationMode, RuntimeConfig, RuntimeError, VerbPresentationPolicy, VerbRegistry,
VerbRegistryBuilder,
};
use khive_storage::EdgeRelation;
use crate::coordinator::CoordinatorService;
use crate::tools::request::RequestParams;
const MAX_BATCH_CONCURRENCY: usize = 8;
const BATCH_RESPONSE_BUDGET_BYTES: usize = khive_runtime::daemon::MAX_FRAME_BYTES / 2;
struct BatchTask<F> {
index: usize,
tool: String,
future: F,
}
#[derive(Clone, Copy, PartialEq, Eq)]
pub(crate) enum DispatchOrigin {
Local,
Daemon,
}
#[derive(Clone, Copy)]
struct RunParsedContext<'a> {
enforce_response_budget: bool,
from_wire: bool,
identity: Option<&'a khive_runtime::RequestIdentity>,
}
pub fn compute_config_id(
config: &RuntimeConfig,
khive_cfg: Option<&khive_runtime::KhiveConfig>,
) -> String {
let mut packs = config.packs.clone();
packs.sort();
let db = config
.db_path
.as_deref()
.map(canonical_fingerprint_path)
.unwrap_or_else(|| ":memory:".to_string());
let primary = config
.embedding_model
.as_ref()
.map(|m| format!("{m:?}"))
.unwrap_or_else(|| "none".to_string());
let mut extra: Vec<String> = config
.additional_embedding_models
.iter()
.map(|m| format!("{m:?}"))
.collect();
extra.sort();
let mut outbound: Vec<String> = config
.allowed_outbound_namespaces
.iter()
.map(|ns| ns.as_str().to_owned())
.collect();
outbound.sort();
outbound.dedup();
let mut git_write_hasher = Sha256::new();
git_write_hasher.update(b"khive.git-write-policy.v1");
git_write_hasher.update((config.git_write.allowed.len() as u64).to_be_bytes());
for entry in &config.git_write.allowed {
git_write_hasher.update((entry.repo.len() as u64).to_be_bytes());
git_write_hasher.update(entry.repo.as_bytes());
git_write_hasher.update((entry.branches.len() as u64).to_be_bytes());
for branch in &entry.branches {
git_write_hasher.update((branch.len() as u64).to_be_bytes());
git_write_hasher.update(branch.as_bytes());
}
}
let git_write = format!("{:x}", git_write_hasher.finalize());
let base = format!(
"packs=[{}];db={};embed={};extra=[{}];backend={:?};outbound=[{}];git_write={}",
packs.join(","),
db,
primary,
extra.join(","),
config.backend_id,
outbound.join(","),
git_write,
);
let topology = khive_cfg
.filter(|cfg| !cfg.backends.is_empty())
.map(|cfg| {
let mut backend_entries: Vec<String> = cfg
.backends
.iter()
.map(|b| {
let path = b
.path
.as_deref()
.map(canonical_fingerprint_path)
.unwrap_or_else(|| ":memory:".to_string());
format!("{}:{:?}:{}", b.name, b.kind, path)
})
.collect();
backend_entries.sort();
let mut pack_entries: Vec<String> = cfg
.packs
.iter()
.map(|(pack, pc)| format!("{}={}", pack, pc.backend))
.collect();
pack_entries.sort();
format!(
";backends=[{}];pack_backends=[{}]",
backend_entries.join(","),
pack_entries.join(","),
)
})
.unwrap_or_default();
format!("{base}{topology}")
}
fn canonical_fingerprint_path(path: &std::path::Path) -> String {
crate::serve::canonical_path_no_side_effects(path)
.map(|p| p.display().to_string())
.unwrap_or_else(|_| path.display().to_string())
}
fn build_verb_catalog(verbs: impl IntoIterator<Item = (String, String, String)>) -> String {
let mut by_verb: std::collections::BTreeMap<String, Vec<(String, String)>> =
std::collections::BTreeMap::new();
for (pack_name, verb_name, description) in verbs {
by_verb
.entry(verb_name)
.or_default()
.push((pack_name, description));
}
let mut out = String::new();
for (name, pack_descs) in &by_verb {
if pack_descs.len() > 1 {
let packs: Vec<&str> = pack_descs.iter().map(|(p, _)| p.as_str()).collect();
tracing::warn!(
verb = %name,
packs = ?packs,
"verb registered by multiple packs; all descriptions included in catalog"
);
}
out.push_str(" ");
out.push_str(name);
out.push_str(" — ");
if pack_descs.len() == 1 {
out.push_str(&pack_descs[0].1);
} else {
for (i, (pack, desc)) in pack_descs.iter().enumerate() {
if i > 0 {
out.push_str("\n ");
}
out.push('[');
out.push_str(pack);
out.push_str("] ");
out.push_str(desc);
}
}
out.push('\n');
}
out
}
#[derive(Clone)]
pub struct KhiveMcpServer {
registry: VerbRegistry,
default_namespace: String,
config_id: String,
coordinator: Option<Arc<dyn CoordinatorService>>,
pool: Option<Arc<ConnectionPool>>,
secondary_pools: Vec<Arc<ConnectionPool>>,
default_output_format: OutputFormat,
schedule_ticker_last_tick_micros: Arc<AtomicI64>,
}
pub enum PackRegFailure {
UnknownPack(String),
MissingDependency { pack: String, dep: String },
Registry(khive_runtime::RuntimeError),
}
pub struct PackRegError {
pub failure: PackRegFailure,
pub runtime: KhiveRuntime,
}
impl std::fmt::Debug for PackRegError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
let mut dbg = f.debug_struct("PackRegError");
match &self.failure {
PackRegFailure::UnknownPack(unknown) => dbg.field("unknown", unknown),
PackRegFailure::MissingDependency { pack, dep } => {
dbg.field("pack", pack).field("missing_dep", dep)
}
PackRegFailure::Registry(source) => dbg.field("source", source),
}
.finish_non_exhaustive()
}
}
impl std::fmt::Display for PackRegError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match &self.failure {
PackRegFailure::UnknownPack(unknown) => write!(
f,
"unknown pack name {:?} — built-in packs: {}",
unknown,
builtin_pack_names().join(", ")
),
PackRegFailure::MissingDependency { pack, dep } => write!(
f,
"pack {pack:?} requires {dep:?}, which is not in the requested pack list; \
add --pack {dep} before --pack {pack}"
),
PackRegFailure::Registry(source) => write!(f, "pack registry build failed: {source}"),
}
}
}
impl std::error::Error for PackRegError {}
pub fn builtin_pack_names() -> Vec<&'static str> {
PackRegistry::discovered_names()
}
#[cfg(unix)]
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum StdioServeMode {
Handshake,
Resumed,
}
#[cfg(unix)]
fn stdio_serve_mode_for(resumed_generation: Option<u32>) -> StdioServeMode {
match resumed_generation {
Some(_) => StdioServeMode::Resumed,
None => StdioServeMode::Handshake,
}
}
impl KhiveMcpServer {
#[allow(clippy::result_large_err)]
pub fn new(runtime: KhiveRuntime) -> Result<Self, PackRegError> {
let packs: Vec<String> = runtime.config().packs.clone();
Self::with_packs(runtime, &packs)
}
#[allow(clippy::result_large_err)]
pub fn with_packs(runtime: KhiveRuntime, packs: &[String]) -> Result<Self, PackRegError> {
let gate = runtime.config().gate.clone();
let default_namespace = runtime.config().default_namespace.clone();
let config_id = compute_config_id(runtime.config(), None);
let visible_namespaces = runtime.config().visible_namespaces.clone();
let actor_id = runtime.config().actor_id.clone();
let mut builder = VerbRegistryBuilder::new();
builder.with_gate(gate);
builder.with_default_namespace(default_namespace.as_str());
builder.with_visible_namespaces(visible_namespaces);
builder.with_actor_id(actor_id);
if let Ok(tok) = runtime.authorize(khive_runtime::Namespace::local()) {
if let Ok(event_store) = runtime.events(&tok) {
builder.with_event_store(event_store);
}
}
if let Err(load_err) = PackRegistry::register_packs(packs, runtime.clone(), &mut builder) {
let failure = match load_err {
PackLoadError::UnknownPack(name) => PackRegFailure::UnknownPack(name),
PackLoadError::MissingDependency { pack, dep } => {
PackRegFailure::MissingDependency { pack, dep }
}
};
return Err(PackRegError { failure, runtime });
}
let registry = builder.build().map_err(|source| PackRegError {
failure: PackRegFailure::Registry(source),
runtime: runtime.clone(),
})?;
runtime.install_edge_rules(registry.all_edge_rules());
registry.call_register_embedders(&runtime);
registry.call_register_entity_type_validators(&runtime);
registry.call_register_note_mutation_hooks(&runtime);
registry.apply_schema_plans(runtime.backend());
let pool = if runtime.backend().is_file_backed() {
Some(runtime.backend().pool_arc())
} else {
None
};
Ok(Self {
registry,
default_namespace: default_namespace.as_str().to_string(),
config_id,
coordinator: None,
pool,
secondary_pools: Vec::new(),
default_output_format: OutputFormat::Json,
schedule_ticker_last_tick_micros: Arc::new(AtomicI64::new(0)),
})
}
#[doc(hidden)]
pub fn from_registry(registry: VerbRegistry) -> Self {
Self {
registry,
default_namespace: "local".to_string(),
config_id: "registry-only".to_string(),
coordinator: None,
pool: None,
secondary_pools: Vec::new(),
default_output_format: OutputFormat::Json,
schedule_ticker_last_tick_micros: Arc::new(AtomicI64::new(0)),
}
}
pub fn from_registry_with_meta(
registry: VerbRegistry,
default_namespace: &str,
config_id: &str,
) -> Self {
Self {
registry,
default_namespace: default_namespace.to_string(),
config_id: config_id.to_string(),
coordinator: None,
pool: None,
secondary_pools: Vec::new(),
default_output_format: OutputFormat::Json,
schedule_ticker_last_tick_micros: Arc::new(AtomicI64::new(0)),
}
}
pub fn with_default_output_format(mut self, fmt: OutputFormat) -> Self {
self.default_output_format = fmt;
self
}
pub fn with_coordinator(mut self, coordinator: Arc<dyn CoordinatorService>) -> Self {
self.coordinator = Some(coordinator);
self
}
pub fn with_pool(mut self, pool: Arc<ConnectionPool>) -> Self {
self.pool = Some(pool);
self
}
pub fn with_secondary_pools(mut self, pools: Vec<Arc<ConnectionPool>>) -> Self {
self.secondary_pools = pools;
self
}
#[cfg(any(feature = "channel-email", feature = "channel-telegram"))]
pub(crate) fn verb_registry_clone(&self) -> VerbRegistry {
self.registry.clone()
}
async fn dispatch_via_coordinator(
&self,
tool: &str,
args_value: &Value,
identity: Option<&khive_runtime::RequestIdentity>,
) -> Option<Result<Value, (String, Value)>> {
let coord = self.coordinator.as_ref()?;
if coord.is_single_backend() {
return None;
}
dispatch_via_coordinator_inner(coord.as_ref(), &self.registry, tool, args_value, identity)
.await
}
pub fn default_namespace(&self) -> &str {
&self.default_namespace
}
pub fn config_id(&self) -> &str {
&self.config_id
}
pub fn actor_id(&self) -> Option<&str> {
self.registry.actor_id()
}
pub fn visible_namespaces(&self) -> &[khive_runtime::Namespace] {
self.registry.visible_namespaces()
}
pub fn pool(&self) -> Option<Arc<ConnectionPool>> {
self.pool.clone()
}
pub fn secondary_pools(&self) -> Vec<Arc<ConnectionPool>> {
self.secondary_pools.clone()
}
pub fn event_store(&self) -> Option<Arc<dyn khive_storage::EventStore>> {
self.registry.event_store()
}
pub fn default_output_format(&self) -> OutputFormat {
self.default_output_format
}
pub(crate) fn record_schedule_ticker_tick(&self) {
self.schedule_ticker_last_tick_micros
.store(chrono::Utc::now().timestamp_micros(), Ordering::Release);
}
pub async fn warm_all(&self) {
self.registry.call_warm_all().await;
}
#[cfg(unix)]
pub async fn serve_stdio(self) -> anyhow::Result<()> {
use rmcp::transport::{async_rw::AsyncRwTransport, stdio};
let build_transport = || {
let (read, write) = stdio();
crate::daemon::SelfHealOnFlushTransport::new(AsyncRwTransport::new_server(read, write))
};
match stdio_serve_mode_for(crate::daemon::resumed_generation()) {
StdioServeMode::Resumed => {
let service = rmcp::service::serve_directly(self, build_transport(), None);
service.waiting().await?;
}
StdioServeMode::Handshake => {
let service = self.serve(build_transport()).await?;
service.waiting().await?;
}
}
Ok(())
}
#[cfg(not(unix))]
pub async fn serve_stdio(self) -> anyhow::Result<()> {
use rmcp::transport::stdio;
let service = self.serve(stdio()).await?;
service.waiting().await?;
Ok(())
}
fn verb_catalog(&self) -> String {
let verbs = self
.registry
.all_verbs_with_names()
.into_iter()
.map(|(pack, v)| (pack.to_owned(), v.name.to_owned(), v.description.to_owned()));
build_verb_catalog(verbs)
}
async fn dispatch_op(
&self,
op: ParsedOp,
prev_result: Option<&Value>,
from_wire: bool,
identity: Option<&khive_runtime::RequestIdentity>,
) -> Result<Value, (String, Value)> {
let ParsedOp { tool, args } = op;
let mut resolved: serde_json::Map<String, Value> = serde_json::Map::new();
for (name, arg_val) in args {
let needs_prev = !matches!(&arg_val, ArgValue::Value(_));
let value = if needs_prev {
let prev = prev_result.ok_or_else(|| {
(
tool.clone(),
json!({
"kind": "substitution_error",
"reason": "no_preceding_op",
"message": format!(
"argument {name:?}: $prev has no preceding op to resolve \
against — this is the first operation in the chain. $prev \
always resolves against the immediately preceding op's result \
only; it cannot reach further back or forward. Move this op \
after the one that produces the value, or pass a literal value \
here instead of $prev."
)
}),
)
})?;
let resolved_val = arg_val.resolve_all(prev).ok_or_else(|| {
(
tool.clone(),
substitution_error_payload(&name, &arg_val, prev),
)
})?;
if matches!(&arg_val, ArgValue::PrevRef { path } if path.is_empty()) {
match &resolved_val {
Value::Object(map) => {
let fields: Vec<&str> = map.keys().map(String::as_str).collect();
return Err((
tool.clone(),
json!({
"kind": "substitution_error",
"reason": "bare_ref_ambiguous",
"message": format!(
"argument {name:?}: $prev requires a dotted path \
(e.g. $prev.id) when the prior result is a map. \
Available top-level fields: [{}]",
fields.join(", ")
),
}),
));
}
Value::Array(_) => {
return Err((
tool.clone(),
json!({
"kind": "substitution_error",
"reason": "bare_ref_ambiguous",
"message": format!(
"argument {name:?}: $prev requires a dotted path \
(e.g. $prev.0) when the prior result is an array. \
Use $prev.N to select a specific element."
),
}),
));
}
_ => {}
}
}
resolved_val
} else {
match arg_val {
ArgValue::Value(v) => v,
_ => unreachable!(),
}
};
resolved.insert(name, value);
}
let args_value = Value::Object(resolved);
let is_help = args_value
.get("help")
.and_then(Value::as_bool)
.unwrap_or(false);
if from_wire && !is_help && self.registry.is_subhandler_verb(&tool) {
return Err((
tool.clone(),
json!(format!(
"permission denied for verb {tool:?}: verb '{tool}' is an internal \
subhandler and cannot be invoked via the MCP request surface"
)),
));
}
if let Some(coord_result) = self
.dispatch_via_coordinator(&tool, &args_value, identity)
.await
{
return coord_result.and_then(|result| chain_ok_envelope_or_depth_error(tool, result));
}
match self
.registry
.dispatch_with_identity(&tool, args_value, identity.cloned())
.await
{
Ok(result) => {
let result = decorate_schedule_agenda_with_ticker_health(
&tool,
is_help,
result,
self.schedule_ticker_last_tick_micros.as_ref(),
);
chain_ok_envelope_or_depth_error(tool, result)
}
Err(RuntimeError::Khive(k)) => {
let error_payload = serde_json::to_value(&k)
.unwrap_or_else(|_| json!({ "kind": "internal", "message": k.to_string() }));
Err((tool, error_payload))
}
Err(e) => Err((tool, json!(e.to_string()))),
}
}
async fn run_parsed(
&self,
ops: Vec<ParsedOp>,
mode: ExecutionMode,
presentation: PresentationMode,
presentation_per_op: Option<Vec<Option<PresentationMode>>>,
context: RunParsedContext<'_>,
) -> Value {
let RunParsedContext {
enforce_response_budget,
from_wire,
identity,
} = context;
let response_budget = if mode == ExecutionMode::Parallel && enforce_response_budget {
BATCH_RESPONSE_BUDGET_BYTES
} else {
usize::MAX
};
let now_unix = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.ok()
.and_then(|d| i64::try_from(d.as_secs()).ok())
.unwrap_or(0);
let mode_for_op = |i: usize| -> PresentationMode {
presentation_per_op
.as_ref()
.and_then(|v| v.get(i))
.and_then(|o| *o)
.unwrap_or(presentation)
};
match mode {
ExecutionMode::Single | ExecutionMode::Parallel => {
let conflict_indices: std::collections::HashSet<usize> = {
let mut seen: std::collections::HashMap<String, usize> =
std::collections::HashMap::new();
let mut bad: std::collections::HashSet<usize> =
std::collections::HashSet::new();
for (i, op) in ops.iter().enumerate() {
for key in khive_request::write_keys_for_op_pub(op) {
if let Some(&prior) = seen.get(&key) {
bad.insert(prior);
bad.insert(i);
} else {
seen.insert(key, i);
}
}
}
bad
};
let coordinator: Option<Arc<dyn CoordinatorService>> = self.coordinator.clone();
let schedule_ticker_last_tick_micros =
self.schedule_ticker_last_tick_micros.clone();
let identity_owned: Option<khive_runtime::RequestIdentity> = identity.cloned();
let futures = ops.into_iter().enumerate().map(|(i, op)| {
let conflict_with: Option<String> = if conflict_indices.contains(&i) {
Some(format!(
"conflict: writes overlap with another op in this batch (op #{})",
i
))
} else {
None
};
let registry = self.registry.clone();
let coord = coordinator.clone();
let schedule_ticker_last_tick_micros =
schedule_ticker_last_tick_micros.clone();
let op_identity = identity_owned.clone();
let op_mode = mode_for_op(i);
let task_tool = op.tool.clone();
BatchTask {
index: i,
tool: task_tool,
future: async move {
let usage_ctx = khive_runtime::usage::UsageContext::new();
let mut entry = khive_runtime::usage::scope(usage_ctx.clone(), async {
let tool = op.tool.clone();
if let Some(msg) = conflict_with {
return json!({ "ok": false, "tool": tool, "error": msg });
}
let effective_mode =
if registry.presentation_policy_for(&tool)
== VerbPresentationPolicy::AlwaysVerbose
{
PresentationMode::Verbose
} else {
op_mode
};
let mut resolved: serde_json::Map<String, Value> =
serde_json::Map::new();
let mut prev_error: Option<Value> = None;
for (name, arg_val) in &op.args {
if matches!(arg_val, ArgValue::Value(_)) {
if let ArgValue::Value(v) = arg_val {
resolved.insert(name.clone(), v.clone());
}
} else {
prev_error = Some(json!({
"ok": false,
"tool": tool,
"error": format!(
"argument {name:?}: $prev reference is only valid in chain (|) mode"
)
}));
break;
}
}
if let Some(err) = prev_error {
return err;
}
let args_value = Value::Object(resolved);
let is_help = args_value
.get("help")
.and_then(Value::as_bool)
.unwrap_or(false);
if from_wire && !is_help && registry.is_subhandler_verb(&tool) {
return json!({
"ok": false,
"tool": tool,
"error": format!(
"permission denied for verb {tool:?}: verb '{tool}' is an \
internal subhandler and cannot be invoked via the MCP \
request surface"
)
});
}
if let Some(active_coord) = coord.as_ref() {
if !active_coord.is_single_backend() {
if let Some(coord_result) = dispatch_via_coordinator_inner(
active_coord.as_ref(),
®istry,
&tool,
&args_value,
op_identity.as_ref(),
)
.await
{
return match coord_result {
Ok(result) => present_ok_envelope_or_depth_error(
tool,
result,
effective_mode,
now_unix,
),
Err((_, error_payload)) => {
json!({ "ok": false, "tool": tool, "error": error_payload })
}
};
}
}
}
match registry
.dispatch_with_identity(&tool, args_value, op_identity)
.await
{
Ok(result) => {
let result = decorate_schedule_agenda_with_ticker_health(
&tool,
is_help,
result,
schedule_ticker_last_tick_micros.as_ref(),
);
present_ok_envelope_or_depth_error(
tool,
result,
effective_mode,
now_unix,
)
}
Err(RuntimeError::Khive(k)) => {
let error_payload = serde_json::to_value(&k).unwrap_or_else(
|_| json!({ "kind": "internal", "message": k.to_string() }),
);
json!({ "ok": false, "tool": tool, "error": error_payload })
}
Err(e) => json!({ "ok": false, "tool": tool, "error": e.to_string() }),
}
})
.await;
stamp_usage(&mut entry, &usage_ctx);
entry
},
}
});
let results = execute_bounded_batch(futures, response_budget).await;
parallel_batch_envelope(results)
}
ExecutionMode::Chain => {
let total = ops.len();
let mut results: Vec<Value> = Vec::with_capacity(total);
let mut prev_result: Option<Value> = None;
let mut aborted_from: Option<usize> = None;
for (i, op) in ops.into_iter().enumerate() {
if let Some(failed_at) = aborted_from {
let failed_index = failed_at - 1;
let failed_tool = results
.get(failed_index)
.and_then(|r| r.get("tool"))
.and_then(Value::as_str)
.unwrap_or("<unknown>");
results.push(json!({
"ok": false,
"tool": op.tool,
"aborted": true,
"message": format!(
"not executed: op #{failed_index} ({failed_tool:?}) failed \
earlier in this chain, so the chain aborted before reaching \
this op. Fix op #{failed_index} — this op's own arguments, \
including any $prev reference, were never evaluated."
),
}));
continue;
}
let op_mode = mode_for_op(i);
let effective_mode = if self.registry.presentation_policy_for(&op.tool)
== VerbPresentationPolicy::AlwaysVerbose
{
PresentationMode::Verbose
} else {
op_mode
};
let usage_ctx = khive_runtime::usage::UsageContext::new();
match khive_runtime::usage::scope(
usage_ctx.clone(),
self.dispatch_op(op, prev_result.as_ref(), from_wire, identity),
)
.await
{
Ok(mut result_obj) => {
stamp_usage(&mut result_obj, &usage_ctx);
match chain_aggregation_depth_reject(result_obj) {
Err(error_entry) => {
results.push(error_entry);
prev_result = None;
aborted_from = Some(i + 1);
continue;
}
Ok(result_obj) => {
prev_result = result_obj.get("result").cloned();
let presented_obj = apply_presentation_to_result(
result_obj,
effective_mode,
now_unix,
);
results.push(presented_obj);
}
}
}
Err((tool, error_payload)) => {
let mut entry =
json!({ "ok": false, "tool": tool, "error": error_payload });
stamp_usage(&mut entry, &usage_ctx);
results.push(entry);
aborted_from = Some(i + 1);
}
}
}
let succeeded = results
.iter()
.filter(|r| r.get("ok").and_then(Value::as_bool) == Some(true))
.count();
let aborted = results
.iter()
.filter(|r| r.get("aborted").and_then(Value::as_bool) == Some(true))
.count();
let failed = total - succeeded - aborted;
json!({
"results": results,
"summary": { "total": total, "succeeded": succeeded, "failed": failed, "aborted": aborted },
"status": batch_status(failed, aborted),
})
}
}
}
}
async fn dispatch_via_coordinator_inner(
coord: &dyn CoordinatorService,
registry: &VerbRegistry,
tool: &str,
args_value: &Value,
identity: Option<&khive_runtime::RequestIdentity>,
) -> Option<Result<Value, (String, Value)>> {
if !matches!(tool, "link" | "search") {
return None;
}
match tool {
"link" => {
if args_value.get("links").is_some() {
return None;
}
let source_str = args_value.get("source_id")?.as_str()?;
let target_str = args_value.get("target_id")?.as_str()?;
let relation_str = args_value.get("relation")?.as_str()?;
let source_id: uuid::Uuid = source_str.parse().ok()?;
let target_id: uuid::Uuid = target_str.parse().ok()?;
let relation: EdgeRelation = relation_str.parse().ok()?;
let weight = args_value
.get("weight")
.and_then(Value::as_f64)
.unwrap_or(1.0);
let metadata = args_value.get("metadata").cloned();
let result = registry
.dispatch_intercepted_with_identity(
tool,
args_value,
identity,
|namespace| async move {
let coord_result = coord
.link(&namespace, source_id, target_id, relation, weight, metadata)
.await
.map_err(RuntimeError::from)?;
let mut raw = serde_json::to_value(&coord_result.edge)
.unwrap_or_else(|e| json!({"error": format!("serialize edge: {e}")}));
if relation.is_symmetric() {
if let Some(obj) = raw.as_object_mut() {
obj.insert("source_id".to_string(), json!(source_id.to_string()));
obj.insert("target_id".to_string(), json!(target_id.to_string()));
}
}
Ok(raw)
},
)
.await;
Some(result.map_err(|error| runtime_error_payload(tool, error)))
}
"search" => {
let kind = args_value.get("kind")?.as_str()?;
let query = args_value.get("query")?.as_str()?;
let result = registry
.dispatch_intercepted_with_identity(
tool,
args_value,
identity,
|namespace| async move {
let limit = match args_value.get("limit") {
None | Some(Value::Null) => 10,
Some(v) => match serde_json::from_value::<u32>(v.clone()) {
Ok(limit) => limit.min(100),
Err(_) => {
return Err(RuntimeError::InvalidInput(
"limit must be an unsigned 32-bit integer".to_string(),
));
}
},
};
let score_floor = args_value
.get("min_score")
.and_then(Value::as_f64)
.unwrap_or(0.0)
.max(0.0);
let kind_filter: Option<&str> = match kind {
"entity" | "note" => None,
other => Some(other),
};
let props_filter: Option<&serde_json::Value> =
args_value.get("properties").and_then(|v| {
if v.as_object().is_some_and(|m| !m.is_empty()) {
Some(v)
} else {
None
}
});
let tags_owned: Vec<String> = match args_value.get("tags") {
None | Some(Value::Null) => vec![],
Some(v) => match serde_json::from_value::<Vec<String>>(v.clone()) {
Ok(t) => t,
Err(_) => {
return Err(RuntimeError::InvalidInput(
"tags must be an array of strings".to_string(),
));
}
},
};
let coord_result = coord
.fan_out_search(
kind,
query,
&namespace,
limit,
kind_filter,
props_filter,
&tags_owned,
)
.await;
let result_val = if !coord_result.note_hits.is_empty()
|| (coord_result.entity_hits.is_empty()
&& coord_result.note_hits.is_empty())
{
let items: Vec<Value> = coord_result
.note_hits
.iter()
.filter(|h| h.score.to_f64() >= score_floor)
.map(|h| {
let note_kind = coord_result.note_kinds.get(&h.note_id);
json!({
"id": h.note_id.to_string(),
"note_kind": note_kind,
"score": h.score.to_f64(),
"source": h.source.as_str(),
"title": h.title,
"snippet": h.snippet,
})
})
.collect();
serde_json::to_value(items).unwrap_or_else(|_| json!([]))
} else {
let items: Vec<Value> = coord_result
.entity_hits
.iter()
.filter(|h| h.score.to_f64() >= score_floor)
.map(|h| {
let entity_kind = coord_result.entity_kinds.get(&h.entity_id);
json!({
"id": h.entity_id.to_string(),
"entity_kind": entity_kind,
"score": h.score.to_f64(),
"source": h.source.as_str(),
"title": h.title,
"snippet": h.snippet,
})
})
.collect();
serde_json::to_value(items).unwrap_or_else(|_| json!([]))
};
Ok(result_val)
},
)
.await;
Some(result.map_err(|error| runtime_error_payload(tool, error)))
}
_ => None,
}
}
fn runtime_error_payload(tool: &str, error: RuntimeError) -> (String, Value) {
match error {
RuntimeError::Khive(k) => {
let error_payload = serde_json::to_value(&k)
.unwrap_or_else(|_| json!({"kind": "internal", "message": k.to_string()}));
(tool.to_string(), error_payload)
}
other => (tool.to_string(), json!(other.to_string())),
}
}
fn result_within_depth_limit(result: &Value) -> bool {
khive_request::value_nesting_within_limit(result, khive_request::NESTING_DEPTH_LIMIT)
}
fn depth_error_payload(context: &str) -> Value {
json!({
"kind": "result_too_deep",
"message": format!(
"op result nesting depth exceeds max {}{context}",
khive_request::NESTING_DEPTH_LIMIT
),
})
}
fn ok_envelope(tool: String, result: Value) -> Value {
let mut map = serde_json::Map::with_capacity(3);
map.insert("ok".to_string(), Value::Bool(true));
map.insert("tool".to_string(), Value::String(tool));
map.insert("result".to_string(), result);
Value::Object(map)
}
fn drop_value_iteratively(value: Value) {
let mut stack = vec![value];
while let Some(v) = stack.pop() {
match v {
Value::Array(items) => stack.extend(items),
Value::Object(map) => stack.extend(map.into_values()),
_ => {}
}
}
}
fn substitution_error_payload(name: &str, arg_val: &ArgValue, prev: &Value) -> Value {
let Some(failure) = arg_val.find_prev_failure(prev) else {
let fields_hint = if let Value::Object(map) = prev {
let mut fields: Vec<&str> = map.keys().map(String::as_str).collect();
fields.sort_unstable();
format!(" Available top-level fields: [{}]", fields.join(", "))
} else {
String::new()
};
return json!({
"kind": "substitution_error",
"reason": "path_not_found",
"message": format!(
"argument {name:?}: one or more $prev paths not found in prior result.{fields_hint}"
),
});
};
let reason = match &failure {
PrevFailure::NotFound { .. } => "path_not_found",
PrevFailure::WrongType { .. } => "path_wrong_type",
PrevFailure::Unsupported { .. } => "path_unsupported",
};
json!({
"kind": "substitution_error",
"reason": reason,
"message": format!(
"argument {name:?}: {failure}. $prev resolves only against the immediately \
preceding op's result — a non-adjacent dependency cannot be expressed inside \
one chain; split into separate calls and carry the value across yourself."
),
})
}
fn stamp_usage(entry: &mut Value, ctx: &khive_runtime::usage::UsageContext) {
if let Value::Object(map) = entry {
map.insert("usage".to_string(), ctx.frozen_or_snapshot());
}
}
fn decorate_schedule_agenda_with_ticker_health(
tool: &str,
is_help: bool,
mut result: Value,
last_tick_micros: &AtomicI64,
) -> Value {
if tool != "schedule.agenda" || is_help {
return result;
}
let last_tick = last_tick_micros.load(Ordering::Acquire);
let last_tick_at = (last_tick > 0).then(|| khive_runtime::micros_to_iso(last_tick));
if let Some(result) = result.as_object_mut() {
result.insert(
"ticker".to_string(),
json!({ "last_tick_at": last_tick_at }),
);
}
result
}
fn chain_ok_envelope_or_depth_error(tool: String, result: Value) -> Result<Value, (String, Value)> {
if !result_within_depth_limit(&result) {
drop_value_iteratively(result);
return Err((
tool,
depth_error_payload("; cannot be used as $prev chain context"),
));
}
Ok(ok_envelope(tool, result))
}
fn present_ok_envelope_or_depth_error(
tool: String,
result: Value,
mode: PresentationMode,
now_unix: i64,
) -> Value {
if !result_within_depth_limit(&result) {
drop_value_iteratively(result);
return json!({ "ok": false, "tool": tool, "error": depth_error_payload("") });
}
let presented = present(result, mode, now_unix);
ok_envelope(tool, presented)
}
fn result_exceeds_depth_limit(result_obj: &Value) -> bool {
result_obj
.get("result")
.is_some_and(|v| !result_within_depth_limit(v))
}
fn chain_aggregation_depth_reject(result_obj: Value) -> Result<Value, Value> {
if result_exceeds_depth_limit(&result_obj) {
let tool_name = result_obj
.get("tool")
.and_then(Value::as_str)
.unwrap_or_default()
.to_string();
let error_entry = json!({
"ok": false,
"tool": tool_name,
"error": {
"kind": "result_too_deep",
"message": format!(
"op result nesting depth exceeds max {}; \
cannot be used as $prev chain context",
khive_request::NESTING_DEPTH_LIMIT
),
},
});
drop_value_iteratively(result_obj);
return Err(error_entry);
}
Ok(result_obj)
}
fn apply_presentation_to_result(
mut result_obj: Value,
mode: PresentationMode,
now_unix: i64,
) -> Value {
if result_obj.get("ok").and_then(Value::as_bool) == Some(true) {
if let Some(result_field) = result_obj.get("result").cloned() {
let presented = present(result_field, mode, now_unix);
if let Some(obj) = result_obj.as_object_mut() {
obj.insert("result".to_string(), presented);
}
}
}
result_obj
}
#[tool_router]
impl KhiveMcpServer {
#[tool(description = r#"Run one or more khive verbs in a single MCP call.
ops syntax:
Single op : verb(name=value, name=value)
Batch : [verb(...), verb(...)] — parallel, max 100
Chain : verb1(...) | verb2(id=$prev.id) — sequential, $prev
JSON form : [{"tool":"verb","args":{...}}, ...] — INDEPENDENT ops only
Argument values are JSON literals: strings (double-quoted), numbers, booleans,
null, arrays, objects. Strings may contain commas / parens; escape with \".
Chain-only: $prev resolves to the prior op's result. Path extraction syntax:
$prev — full result
$prev.field — nested object field
$prev.items[0].id — array index
$prev[2] — top-level array index
Quoted strings that contain $prev are promoted to substitutions (e.g. id="$prev.id"
is the same as id=$prev.id). To pass a literal "$prev", escape with backslash:
\"\\$prev\". JSON form is for independent ops only — any $prev string in JSON
form is rejected.
Response shape:
{
"results": [ {"ok": true, "tool": "verb", "result": {...}}, ... ],
"summary": { "total": N, "succeeded": N, "failed": N, "aborted": N },
"status": "success" | "partial"
}
Parallel: a failed op does NOT abort siblings. Chain: failure aborts remaining
ops (reported as {"ok": false, "aborted": true}). Committed ops are not rolled back.
`status` is "partial" whenever summary.failed or summary.aborted is non-zero — check
it (or summary) rather than relying on the absence of a top-level error.
Verb discovery: install the `kg` / `gtd` plugins for usage skills. The verbs
currently registered on this server (pack-derived) are listed below. Argument
schemas live in each pack's docs and SKILL.md files.
Tip: for one-shot calls, the single-op form is the densest. Use batch when
several independent ops can run together; use chain when each op needs the prior
result (e.g. create then link with the new entity's id)."#)]
async fn request(&self, Parameters(p): Parameters<RequestParams>) -> Result<String, McpError> {
#[cfg(unix)]
if p.save_to.is_none() {
let frame = self.wire_daemon_frame(&p);
if let Some(res) = crate::daemon::forward_or_spawn(&frame).await {
return match res {
Ok(s) => Ok(s),
Err(e) => match strict_fallback_reason(&e) {
Some(reason) => strict_fallback_envelope_response(&p, reason),
None => Err(e),
},
};
}
}
self.dispatch_request_wire(p).await
}
}
fn batch_status(failed: usize, aborted: usize) -> &'static str {
if failed == 0 && aborted == 0 {
"success"
} else {
"partial"
}
}
fn batch_budget_error(tool: &str, response_budget: usize) -> Value {
json!({
"ok": false,
"tool": tool,
"error": format!(
"batch response budget of {response_budget} serialized bytes exceeded"
),
})
}
async fn execute_bounded_batch<I, F>(tasks: I, response_budget: usize) -> Vec<Value>
where
I: IntoIterator<Item = BatchTask<F>>,
F: Future<Output = Value>,
{
let mut queued: std::collections::VecDeque<_> = tasks.into_iter().collect();
let total = queued.len();
let mut in_flight = FuturesUnordered::new();
let start = |task: BatchTask<F>| async move {
let entry = task.future.await;
(task.index, task.tool, entry)
};
for _ in 0..MAX_BATCH_CONCURRENCY {
if let Some(task) = queued.pop_front() {
in_flight.push(start(task));
}
}
let mut results: Vec<Option<Value>> = (0..total).map(|_| None).collect();
let mut accumulated_bytes = 0usize;
let mut budget_breached = false;
while let Some((index, _tool, entry)) = in_flight.next().await {
if budget_breached {
results[index] = Some(entry);
continue;
}
let serialized_bytes = serde_json::to_vec(&entry)
.expect("serde_json::Value is always serializable")
.len();
if serialized_bytes > response_budget.saturating_sub(accumulated_bytes) {
results[index] = Some(entry);
budget_breached = true;
continue;
}
accumulated_bytes += serialized_bytes;
results[index] = Some(entry);
if let Some(task) = queued.pop_front() {
in_flight.push(start(task));
}
}
for task in queued {
results[task.index] = Some(batch_budget_error(&task.tool, response_budget));
}
results
.into_iter()
.map(|entry| entry.expect("every started or queued batch task has a result"))
.collect()
}
fn parallel_batch_envelope(results: Vec<Value>) -> Value {
let total = results.len();
let succeeded = results
.iter()
.filter(|result| result.get("ok").and_then(Value::as_bool) == Some(true))
.count();
let failed = total - succeeded;
json!({
"results": results,
"summary": { "total": total, "succeeded": succeeded, "failed": failed, "aborted": 0 },
"status": batch_status(failed, 0),
})
}
#[cfg(unix)]
fn strict_fallback_reason(e: &McpError) -> Option<String> {
let data = e.data.as_ref()?;
if data.get(crate::daemon::STRICT_FALLBACK_MARKER)?.as_bool() != Some(true) {
return None;
}
data.get("reason")
.and_then(Value::as_str)
.map(str::to_string)
}
#[cfg(unix)]
fn strict_fallback_envelope_response(
p: &RequestParams,
reason: String,
) -> Result<String, McpError> {
let parsed = parse_request(&p.ops).map_err(dsl_err_to_mcp)?;
let total = parsed.ops.len();
let error_msg = format!(
"daemon fallback rejected under KHIVE_DAEMON_STRICT=1: reason={reason}; \
refusing to complete the request via local dispatch; \
rebuild with `make local` and retry"
);
let results: Vec<Value> = match parsed.mode {
ExecutionMode::Chain => parsed
.ops
.iter()
.enumerate()
.map(|(i, op)| {
if i == 0 {
json!({ "ok": false, "tool": op.tool, "error": error_msg })
} else {
json!({ "ok": false, "tool": op.tool, "aborted": true })
}
})
.collect(),
ExecutionMode::Single | ExecutionMode::Parallel => parsed
.ops
.iter()
.map(|op| json!({ "ok": false, "tool": op.tool, "error": error_msg }))
.collect(),
};
let aborted = if parsed.mode == ExecutionMode::Chain {
total.saturating_sub(1)
} else {
0
};
let failed = total - aborted;
Ok(serde_json::to_string(&json!({
"results": results,
"summary": { "total": total, "succeeded": 0, "failed": failed, "aborted": aborted },
"status": batch_status(failed, aborted),
}))
.expect("envelope of string/bool JSON values always serializes"))
}
impl KhiveMcpServer {
#[cfg(unix)]
pub(crate) fn wire_daemon_frame(&self, p: &RequestParams) -> khive_runtime::DaemonRequestFrame {
khive_runtime::DaemonRequestFrame {
ops: p.ops.clone(),
presentation: p.presentation.clone(),
presentation_per_op: p.presentation_per_op.clone(),
namespace: self.default_namespace.clone(),
actor_id: self.actor_id().map(str::to_string),
process_ref: khive_runtime::process_ref_from_env(),
visible_namespaces: self
.visible_namespaces()
.iter()
.map(|ns| ns.as_str().to_string())
.collect(),
config_id: self.config_id.clone(),
protocol_version: khive_runtime::daemon::PROTOCOL_VERSION,
probe_only: false,
metrics_only: false,
format: p.format.clone(),
format_per_op: p.format_per_op.clone(),
from_wire: true,
request_id: p.request_id,
}
}
pub async fn dispatch_request_local(&self, p: RequestParams) -> Result<String, McpError> {
self.dispatch_request_inner(p, false, None, DispatchOrigin::Local)
.await
}
pub(crate) async fn dispatch_request_replay_as(
&self,
p: RequestParams,
namespace: &str,
verified_actor: Option<khive_runtime::VerifiedActor>,
) -> Result<String, McpError> {
let identity = khive_runtime::RequestIdentity {
namespace: namespace.to_string(),
actor_id: verified_actor.map(|actor| actor.as_str().to_string()),
process_ref: khive_runtime::process_ref_from_env(),
visible_namespaces: Vec::new(),
request_id: None,
};
self.dispatch_request_inner(p, true, Some(identity), DispatchOrigin::Local)
.await
}
pub(crate) async fn dispatch_request_wire(&self, p: RequestParams) -> Result<String, McpError> {
self.dispatch_request_inner(p, true, None, DispatchOrigin::Local)
.await
}
pub(crate) async fn dispatch_request_inner(
&self,
p: RequestParams,
from_wire: bool,
identity: Option<khive_runtime::RequestIdentity>,
origin: DispatchOrigin,
) -> Result<String, McpError> {
let save_to = p.save_to.clone();
let identity = identity.or_else(|| {
p.request_id
.map(|request_id| khive_runtime::RequestIdentity {
namespace: self.default_namespace.clone(),
actor_id: self.actor_id().map(str::to_string),
process_ref: khive_runtime::process_ref_from_env(),
visible_namespaces: self
.visible_namespaces()
.iter()
.map(|ns| ns.as_str().to_string())
.collect(),
request_id: Some(request_id),
})
});
let parsed = parse_request(&p.ops).map_err(dsl_err_to_mcp)?;
let presentation = parse_presentation_mode(p.presentation.as_deref())
.map_err(|e| McpError::invalid_params(e, None))?;
let presentation_per_op: Option<Vec<Option<PresentationMode>>> =
if let Some(per_op_strs) = p.presentation_per_op {
let mut modes = Vec::with_capacity(per_op_strs.len());
for s in per_op_strs {
let mode = match s.as_deref() {
None => None,
Some(v) => Some(
parse_presentation_mode(Some(v))
.map_err(|e| McpError::invalid_params(e, None))?,
),
};
modes.push(mode);
}
Some(modes)
} else {
None
};
let batch_format = parse_output_format(p.format.as_deref())
.map_err(|e| McpError::invalid_params(e, None))?
.unwrap_or(self.default_output_format);
let format_per_op: Option<Vec<Option<OutputFormat>>> =
if let Some(per_op_strs) = p.format_per_op {
let mut fmts = Vec::with_capacity(per_op_strs.len());
for s in per_op_strs {
let fmt = match s.as_deref() {
None => None,
Some(v) => Some(
parse_output_format(Some(v))
.map_err(|e| McpError::invalid_params(e, None))?
.unwrap_or(batch_format),
),
};
fmts.push(fmt);
}
Some(fmts)
} else {
None
};
let result = self
.run_parsed(
parsed.ops,
parsed.mode,
presentation,
presentation_per_op.clone(),
RunParsedContext {
enforce_response_budget: save_to.is_none(),
from_wire,
identity: identity.as_ref(),
},
)
.await;
if let Some(path_str) = save_to {
let path = std::path::Path::new(&path_str);
let manifest = crate::save_sink::write_and_manifest(&result, path, from_wire)
.map_err(|e| McpError::internal_error(format!("save_to: {e}"), None))?;
return serde_json::to_string(&manifest)
.map_err(|e| McpError::internal_error(format!("serialize manifest: {e}"), None));
}
Ok(render_result(
result,
batch_format,
&format_per_op,
presentation,
&presentation_per_op,
&self.registry,
(origin == DispatchOrigin::Daemon).then_some(self.config_id.as_str()),
))
}
}
fn dsl_err_to_mcp(e: DslError) -> McpError {
McpError::invalid_params(e.to_string(), None)
}
fn parse_presentation_mode(s: Option<&str>) -> Result<PresentationMode, String> {
match s {
None | Some("agent") => Ok(PresentationMode::Agent),
Some("verbose") => Ok(PresentationMode::Verbose),
Some("human") => Ok(PresentationMode::Human),
Some(other) => Err(format!(
"unknown presentation mode {other:?}; valid values: \"agent\", \"verbose\", \"human\""
)),
}
}
fn parse_output_format(s: Option<&str>) -> Result<Option<OutputFormat>, String> {
match s {
None => Ok(None),
Some("json") => Ok(Some(OutputFormat::Json)),
Some("auto") => Ok(Some(OutputFormat::Auto)),
Some("table") => Ok(Some(OutputFormat::Table)),
Some(other) => Err(format!(
"unknown output format {other:?}; valid values: \"json\", \"auto\", \"table\""
)),
}
}
fn render_result(
value: serde_json::Value,
batch_format: OutputFormat,
format_per_op: &Option<Vec<Option<OutputFormat>>>,
presentation: PresentationMode,
presentation_per_op: &Option<Vec<Option<PresentationMode>>>,
registry: &VerbRegistry,
daemon_frame_config_id: Option<&str>,
) -> String {
if let serde_json::Value::Object(ref map) = value {
if let Some(serde_json::Value::Array(results)) = map.get("results") {
let out_results = results
.iter()
.enumerate()
.map(|(index, entry)| {
render_batch_entry(
index,
entry,
batch_format,
format_per_op,
presentation,
presentation_per_op,
registry,
)
})
.collect();
let out_map = match daemon_frame_config_id {
Some(config_id) => {
fit_rendered_batch_envelope(map, results, out_results, config_id)
}
None => {
let mut out_map = map.clone();
out_map.insert("results".to_string(), Value::Array(out_results));
out_map
}
};
return serialize_response_value(&serde_json::Value::Object(out_map));
}
}
let rendered = render_format(value.clone(), batch_format, presentation);
let Some(config_id) = daemon_frame_config_id else {
return rendered;
};
if rendered_response_fits_daemon_frame(&rendered, config_id) {
return rendered;
}
let compact = serialize_response_value(&value);
if rendered_response_fits_daemon_frame(&compact, config_id) {
return compact;
}
serde_json::to_string(&json!({
"ok": false,
"error": "response payload omitted because it exceeds the daemon frame budget",
}))
.expect("static frame-budget error is serializable")
}
fn render_batch_entry(
index: usize,
entry: &Value,
batch_format: OutputFormat,
format_per_op: &Option<Vec<Option<OutputFormat>>>,
presentation: PresentationMode,
presentation_per_op: &Option<Vec<Option<PresentationMode>>>,
registry: &VerbRegistry,
) -> Value {
let per_op_format = format_per_op
.as_ref()
.and_then(|formats| formats.get(index))
.and_then(|format| *format)
.unwrap_or(batch_format);
let is_ok = entry.get("ok").and_then(Value::as_bool).unwrap_or(false);
if !is_ok || per_op_format == OutputFormat::Json {
return entry.clone();
}
let base_presentation = presentation_per_op
.as_ref()
.and_then(|modes| modes.get(index))
.and_then(|mode| *mode)
.unwrap_or(presentation);
let effective_presentation = match entry.get("tool").and_then(Value::as_str) {
Some(tool)
if registry.presentation_policy_for(tool) == VerbPresentationPolicy::AlwaysVerbose =>
{
PresentationMode::Verbose
}
_ => base_presentation,
};
let Some(result) = entry.get("result") else {
return entry.clone();
};
let mut rendered_entry = entry.clone();
if let Value::Object(ref mut fields) = rendered_entry {
fields.insert(
"result".to_string(),
Value::String(render_format(
result.clone(),
per_op_format,
effective_presentation,
)),
);
}
rendered_entry
}
fn fit_rendered_batch_envelope(
map: &serde_json::Map<String, Value>,
compact_results: &[Value],
mut out_results: Vec<Value>,
served_config_id: &str,
) -> serde_json::Map<String, Value> {
let mut out_map = map.clone();
out_map.insert(
"results".to_string(),
serde_json::Value::Array(out_results.clone()),
);
if response_value_fits_daemon_frame(
&serde_json::Value::Object(out_map.clone()),
served_config_id,
) {
return out_map;
}
let rendered_frame_bytes = response_value_daemon_frame_len(
&serde_json::Value::Object(out_map.clone()),
served_config_id,
);
let mut compact_fallbacks: Vec<(usize, usize)> = compact_results
.iter()
.zip(&out_results)
.enumerate()
.filter_map(|(index, (compact, rendered))| {
if compact == rendered {
return None;
}
let mut candidate_results = out_results.clone();
candidate_results[index] = compact.clone();
let mut candidate_map = out_map.clone();
candidate_map.insert("results".to_string(), Value::Array(candidate_results));
let compact_frame_bytes =
response_value_daemon_frame_len(&Value::Object(candidate_map), served_config_id);
(compact_frame_bytes < rendered_frame_bytes)
.then_some((index, rendered_frame_bytes - compact_frame_bytes))
})
.collect();
compact_fallbacks.sort_unstable_by_key(|&(_, saved_bytes)| std::cmp::Reverse(saved_bytes));
for (index, _) in compact_fallbacks {
out_results[index] = compact_results[index].clone();
out_map.insert("results".to_string(), Value::Array(out_results.clone()));
if response_value_fits_daemon_frame(&Value::Object(out_map.clone()), served_config_id) {
return out_map;
}
}
let mut by_size: Vec<(usize, usize)> = out_results
.iter()
.enumerate()
.map(|(index, entry)| (index, serialized_response_len(entry)))
.collect();
by_size.sort_unstable_by_key(|&(_, bytes)| std::cmp::Reverse(bytes));
for (index, _) in by_size {
out_results[index] = frame_budget_omission(&compact_results[index]);
out_map.insert(
"results".to_string(),
serde_json::Value::Array(out_results.clone()),
);
if response_value_fits_daemon_frame(
&serde_json::Value::Object(out_map.clone()),
served_config_id,
) {
break;
}
}
out_map
}
fn frame_budget_omission(entry: &Value) -> Value {
let ok = entry.get("ok").and_then(Value::as_bool).unwrap_or(false);
let mut omitted = serde_json::Map::new();
for key in ["ok", "tool", "usage", "aborted"] {
if let Some(value) = entry.get(key) {
omitted.insert(key.to_string(), value.clone());
}
}
if ok {
omitted.insert(
"result_omitted".to_string(),
json!("operation succeeded; result omitted because the response frame budget was exceeded"),
);
} else {
omitted.insert(
"error".to_string(),
json!("operation failed; error details omitted because the response frame budget was exceeded"),
);
}
Value::Object(omitted)
}
fn serialized_response_len(value: &Value) -> usize {
serde_json::to_vec(value)
.expect("serde_json::Value is always serializable")
.len()
}
fn serialize_response_value(value: &Value) -> String {
serde_json::to_string(value).expect("serde_json::Value is always serializable")
}
fn response_value_fits_daemon_frame(value: &Value, served_config_id: &str) -> bool {
response_value_daemon_frame_len(value, served_config_id)
<= khive_runtime::daemon::MAX_FRAME_BYTES
}
fn response_value_daemon_frame_len(value: &Value, served_config_id: &str) -> usize {
rendered_response_daemon_frame_len(&serialize_response_value(value), served_config_id)
}
fn rendered_response_fits_daemon_frame(rendered: &str, served_config_id: &str) -> bool {
rendered_response_daemon_frame_len(rendered, served_config_id)
<= khive_runtime::daemon::MAX_FRAME_BYTES
}
fn rendered_response_daemon_frame_len(rendered: &str, served_config_id: &str) -> usize {
let frame = khive_runtime::DaemonResponseFrame {
ok: true,
result: Some(rendered.to_string()),
error: None,
namespace_mismatch: false,
config_mismatch: false,
served_config_id: Some(served_config_id.to_string()),
version_mismatch: false,
daemon_protocol_version: khive_runtime::PROTOCOL_VERSION,
metrics: None,
request_id: Some(u64::MAX),
};
serde_json::to_vec(&frame)
.expect("daemon response frame is always serializable")
.len()
}
fn build_instructions(catalog: &str, builtins: &str) -> String {
format!(
"khive — request-only MCP surface. One tool, `request`, \
dispatches verbs through the loaded pack registry. Configure packs via \
KHIVE_PACKS or --pack (built-ins: {builtins}). Verbs registered on this \
server:\n{catalog}\nFor detailed usage of each verb, see the corresponding \
plugin's SKILL.md files.\n\
Docs: https://ohdearquant.github.io/khive/ (hosted) or docs/*.md in the repo \
checkout. Treat the live verb catalog above and help=true as authoritative over \
cached/training knowledge. Config/backend issues: docs/configuration.md. Usage \
patterns: docs/guide/tips-and-tricks.md."
)
}
#[tool_handler]
impl ServerHandler for KhiveMcpServer {
fn get_info(&self) -> ServerInfo {
let catalog = self.verb_catalog();
let builtins = builtin_pack_names().join(", ");
let instructions = build_instructions(&catalog, &builtins);
ServerInfo::new(ServerCapabilities::builder().enable_tools().build())
.with_server_info(Implementation::new(
env!("CARGO_PKG_NAME"),
env!("CARGO_PKG_VERSION"),
))
.with_instructions(instructions)
}
async fn list_tools(
&self,
_request: Option<rmcp::model::PaginatedRequestParams>,
_context: rmcp::service::RequestContext<rmcp::RoleServer>,
) -> Result<rmcp::model::ListToolsResult, McpError> {
let mut tools = Self::tool_router().list_all();
let catalog = self.verb_catalog();
for t in &mut tools {
if t.name == "request" {
let base = t.description.as_deref().unwrap_or("");
t.description = Some(std::borrow::Cow::Owned(format!(
"{base}\n\nVerbs registered on this server:\n{catalog}"
)));
}
}
Ok(rmcp::model::ListToolsResult {
tools,
meta: None,
next_cursor: None,
})
}
}
#[cfg(test)]
mod tests {
use super::*;
use khive_runtime::Namespace;
use khive_storage::{EventFilter, PageRequest};
use serial_test::serial;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::time::Duration;
async fn observed_batch_entry(
_index: usize,
delay_ms: u64,
entry: Value,
in_flight: Arc<AtomicUsize>,
max_in_flight: Arc<AtomicUsize>,
) -> Value {
let current = in_flight.fetch_add(1, Ordering::SeqCst) + 1;
max_in_flight.fetch_max(current, Ordering::SeqCst);
tokio::time::sleep(Duration::from_millis(delay_ms)).await;
in_flight.fetch_sub(1, Ordering::SeqCst);
entry
}
fn batch_task<F>(index: usize, future: F) -> BatchTask<F>
where
F: Future<Output = Value>,
{
BatchTask {
index,
tool: "probe".to_string(),
future,
}
}
struct LargeResultPack;
impl khive_types::Pack for LargeResultPack {
const NAME: &'static str = "large-result-test";
const NOTE_KINDS: &'static [&'static str] = &[];
const ENTITY_KINDS: &'static [&'static str] = &[];
const HANDLERS: &'static [khive_runtime::HandlerDef] = &[khive_runtime::HandlerDef {
name: "large_result",
description: "returns a caller-sized test result",
visibility: khive_runtime::Visibility::Verb,
category: khive_runtime::VerbCategory::Assertive,
params: &[],
}];
}
#[async_trait::async_trait]
impl khive_runtime::PackRuntime for LargeResultPack {
fn name(&self) -> &str {
<Self as khive_types::Pack>::NAME
}
fn note_kinds(&self) -> &'static [&'static str] {
<Self as khive_types::Pack>::NOTE_KINDS
}
fn entity_kinds(&self) -> &'static [&'static str] {
<Self as khive_types::Pack>::ENTITY_KINDS
}
fn handlers(&self) -> &'static [khive_runtime::HandlerDef] {
<Self as khive_types::Pack>::HANDLERS
}
async fn dispatch(
&self,
_verb: &str,
params: Value,
_registry: &VerbRegistry,
_token: &khive_runtime::NamespaceToken,
) -> Result<Value, RuntimeError> {
if let Some(bytes) = params
.get("table_bytes")
.and_then(Value::as_u64)
.and_then(|n| usize::try_from(n).ok())
{
let payload = "x".repeat(bytes);
return Ok(json!([
{"payload": payload},
{"payload": payload},
]));
}
let bytes = params
.get("bytes")
.and_then(Value::as_u64)
.and_then(|n| usize::try_from(n).ok())
.expect("test supplies a valid byte count");
Ok(json!("x".repeat(bytes)))
}
}
async fn dispatch_large_result_through_daemon(
server: &KhiveMcpServer,
ops: String,
format: Option<String>,
) -> String {
khive_runtime::daemon::DaemonDispatch::dispatch(
server, ops, None, None, format, None, false, None,
)
.await
.expect("daemon dispatch")
}
fn large_result_test_server() -> KhiveMcpServer {
let mut builder = VerbRegistryBuilder::new();
builder.register(LargeResultPack);
KhiveMcpServer::from_registry(builder.build().expect("test registry"))
}
#[tokio::test]
async fn bounded_batch_preserves_input_order() {
let count = MAX_BATCH_CONCURRENCY + 3;
let in_flight = Arc::new(AtomicUsize::new(0));
let max_in_flight = Arc::new(AtomicUsize::new(0));
let futures = (0..count).map(|index| {
batch_task(
index,
observed_batch_entry(
index,
(count - index) as u64,
json!({"ok": true, "tool": "probe", "result": {"index": index}}),
in_flight.clone(),
max_in_flight.clone(),
),
)
});
let results = execute_bounded_batch(futures, usize::MAX).await;
let indices: Vec<u64> = results
.iter()
.map(|entry| entry["result"]["index"].as_u64().expect("result index"))
.collect();
assert_eq!(indices, (0..count as u64).collect::<Vec<_>>());
}
#[tokio::test]
async fn bounded_batch_enforces_aggregate_response_budget() {
assert_eq!(
BATCH_RESPONSE_BUDGET_BYTES,
khive_runtime::daemon::MAX_FRAME_BYTES / 2
);
let count = MAX_BATCH_CONCURRENCY * 2;
let budget = BATCH_RESPONSE_BUDGET_BYTES;
let small = json!({
"ok": true,
"tool": "probe",
"result": "x".repeat(budget / 4 - 128),
});
let in_flight = Arc::new(AtomicUsize::new(0));
let max_in_flight = Arc::new(AtomicUsize::new(0));
let futures = (0..count).map(|index| {
let (delay_ms, entry) = if index < 2 {
(index as u64, small.clone())
} else if index == 2 {
(
30,
json!({"ok": true, "tool": "probe", "result": "x".repeat(budget)}),
)
} else {
(
60,
json!({"ok": true, "tool": "probe", "result": {"index": index}}),
)
};
batch_task(
index,
observed_batch_entry(
index,
delay_ms,
entry,
in_flight.clone(),
max_in_flight.clone(),
),
)
});
let results = tokio::time::timeout(
Duration::from_secs(1),
execute_bounded_batch(futures, budget),
)
.await
.expect("started operations must settle promptly after a budget breach");
let response = parallel_batch_envelope(results);
assert_eq!(
response["summary"],
json!({"total": count, "succeeded": 10, "failed": count - 10, "aborted": 0})
);
assert_eq!(response["results"][0]["ok"], true);
assert_eq!(response["results"][1]["ok"], true);
assert_eq!(
response["results"][2]["result"]
.as_str()
.expect("breaching result remains truthful")
.len(),
budget
);
for index in 3..10 {
assert_eq!(response["results"][index]["ok"], true);
assert_eq!(response["results"][index]["result"]["index"], index);
}
for entry in response["results"]
.as_array()
.expect("results")
.iter()
.skip(10)
{
assert_eq!(entry["ok"], false);
let error = entry["error"].as_str().expect("budget error string");
assert!(error.contains("batch response budget"));
assert!(error.contains(&budget.to_string()));
}
assert_eq!(in_flight.load(Ordering::SeqCst), 0);
serde_json::to_vec(&response).expect("budgeted response must serialize");
}
#[tokio::test]
async fn save_to_writes_full_results_without_inline_response_budgeting() {
let server = large_result_test_server();
let dir = tempfile::tempdir().expect("tempdir");
let sink_path = dir.path().join("full-results.jsonl");
let result_bytes = BATCH_RESPONSE_BUDGET_BYTES * 3 / 4;
let response = server
.dispatch_request_inner(
RequestParams {
ops: format!(
"[large_result(bytes={result_bytes}), large_result(bytes={result_bytes})]"
),
presentation: None,
presentation_per_op: None,
save_to: Some(sink_path.to_string_lossy().into_owned()),
format: None,
format_per_op: None,
request_id: None,
},
false,
None,
DispatchOrigin::Local,
)
.await
.expect("save_to dispatch");
let manifest: Value = serde_json::from_str(&response).expect("manifest JSON");
assert_eq!(manifest["rows"], 2);
assert_eq!(manifest["summary"]["succeeded"], 2);
let rows: Vec<Value> = std::fs::read_to_string(&sink_path)
.expect("read JSONL")
.lines()
.map(|line| serde_json::from_str(line).expect("valid JSONL row"))
.collect();
assert_eq!(rows.len(), 2);
for row in rows {
assert_eq!(row["ok"], true);
assert_eq!(
row["result"].as_str().expect("full result string").len(),
result_bytes
);
}
}
#[tokio::test]
async fn local_dispatch_returns_result_larger_than_daemon_frame() {
let server = large_result_test_server();
let result_bytes = khive_runtime::daemon::MAX_FRAME_BYTES + 1_024;
let response = server
.dispatch_request_local(RequestParams {
ops: format!("large_result(bytes={result_bytes})"),
presentation: None,
presentation_per_op: None,
save_to: None,
format: None,
format_per_op: None,
request_id: None,
})
.await
.expect("local dispatch");
let envelope: Value = serde_json::from_str(&response).expect("response envelope");
assert_eq!(envelope["results"][0]["ok"], true);
assert_eq!(
envelope["results"][0]["result"]
.as_str()
.expect("full local result")
.len(),
result_bytes
);
assert!(envelope["results"][0].get("result_omitted").is_none());
}
#[tokio::test]
async fn daemon_dispatch_degrades_result_larger_than_frame() {
let server = large_result_test_server();
let result_bytes = khive_runtime::daemon::MAX_FRAME_BYTES + 1_024;
let response = dispatch_large_result_through_daemon(
&server,
format!("large_result(bytes={result_bytes})"),
None,
)
.await;
let envelope: Value = serde_json::from_str(&response).expect("response envelope");
assert_eq!(envelope["results"][0]["ok"], true);
assert!(envelope["results"][0].get("result").is_none());
assert!(envelope["results"][0].get("result_omitted").is_some());
assert!(rendered_response_fits_daemon_frame(
&response,
&server.config_id
));
}
#[tokio::test]
async fn daemon_batch_keeps_rendered_result_when_compact_result_exceeds_frame() {
let server = large_result_test_server();
let row_bytes = khive_runtime::daemon::MAX_FRAME_BYTES / 2;
let response = dispatch_large_result_through_daemon(
&server,
format!("large_result(table_bytes={row_bytes})"),
Some("auto".to_string()),
)
.await;
let envelope: Value = serde_json::from_str(&response).expect("response envelope");
let entry = &envelope["results"][0];
assert_eq!(entry["ok"], true);
assert!(entry.get("result_omitted").is_none());
let rendered = entry["result"].as_str().expect("rendered table result");
assert!(rendered.starts_with("| payload |"));
assert!(rendered.len() < row_bytes);
assert!(rendered_response_fits_daemon_frame(
&response,
&server.config_id
));
}
#[test]
fn auto_rendered_batch_stays_within_daemon_frame_cap() {
let mut leaves = serde_json::Map::new();
for index in 0..80_000 {
leaves.insert(format!("k{index}"), json!(0));
}
let result = nest_object(60, Value::Object(leaves));
let envelope = parallel_batch_envelope(vec![json!({
"ok": true,
"tool": "probe",
"result": result,
})]);
let compact_bytes = serde_json::to_vec(&envelope)
.expect("compact envelope")
.len();
assert!(compact_bytes < BATCH_RESPONSE_BUDGET_BYTES);
let rendered = render_result(
envelope,
OutputFormat::Auto,
&None,
PresentationMode::Agent,
&None,
&large_result_test_server().registry,
Some("test"),
);
let rendered_value: Value = serde_json::from_str(&rendered).expect("response envelope");
assert_eq!(rendered_value["status"], "success");
assert_eq!(rendered_value["results"][0]["ok"], true);
assert!(
rendered_value["results"][0]["result"].is_object(),
"oversized auto output must fall back to the truthful compact result"
);
let frame = khive_runtime::DaemonResponseFrame {
ok: true,
result: Some(rendered),
error: None,
namespace_mismatch: false,
config_mismatch: false,
served_config_id: Some("test".to_string()),
version_mismatch: false,
daemon_protocol_version: khive_runtime::PROTOCOL_VERSION,
metrics: None,
request_id: None,
};
let frame_bytes = serde_json::to_vec(&frame).expect("daemon frame").len();
assert!(
frame_bytes <= khive_runtime::daemon::MAX_FRAME_BYTES,
"rendered daemon frame was {frame_bytes} bytes"
);
}
#[tokio::test]
async fn bounded_batch_op_error_does_not_abort_siblings() {
let count = 5;
let in_flight = Arc::new(AtomicUsize::new(0));
let max_in_flight = Arc::new(AtomicUsize::new(0));
let futures = (0..count).map(|index| {
let entry = if index == 2 {
json!({"ok": false, "tool": "probe", "error": "expected failure"})
} else {
json!({"ok": true, "tool": "probe", "result": {"index": index}})
};
batch_task(
index,
observed_batch_entry(
index,
(count - index) as u64,
entry,
in_flight.clone(),
max_in_flight.clone(),
),
)
});
let response = parallel_batch_envelope(execute_bounded_batch(futures, usize::MAX).await);
assert_eq!(
response["summary"],
json!({"total": 5, "succeeded": 4, "failed": 1, "aborted": 0})
);
assert_eq!(response["results"][2]["error"], "expected failure");
assert!(response["results"][3]["ok"].as_bool().unwrap_or(false));
assert!(response["results"][4]["ok"].as_bool().unwrap_or(false));
}
#[tokio::test]
async fn bounded_batch_never_exceeds_concurrency_limit() {
let count = MAX_BATCH_CONCURRENCY * 3;
let in_flight = Arc::new(AtomicUsize::new(0));
let max_in_flight = Arc::new(AtomicUsize::new(0));
let futures = (0..count).map(|index| {
batch_task(
index,
observed_batch_entry(
index,
10,
json!({"ok": true, "tool": "probe", "result": index}),
in_flight.clone(),
max_in_flight.clone(),
),
)
});
let results = execute_bounded_batch(futures, usize::MAX).await;
assert_eq!(results.len(), count);
assert_eq!(max_in_flight.load(Ordering::SeqCst), MAX_BATCH_CONCURRENCY);
assert_eq!(in_flight.load(Ordering::SeqCst), 0);
}
fn t(pack: &str, verb: &str, desc: &str) -> (String, String, String) {
(pack.to_owned(), verb.to_owned(), desc.to_owned())
}
#[cfg(unix)]
#[test]
fn stdio_serve_mode_cold_start_uses_handshake() {
assert_eq!(stdio_serve_mode_for(None), StdioServeMode::Handshake);
}
#[cfg(unix)]
#[test]
fn stdio_serve_mode_resumed_generation_skips_handshake() {
assert_eq!(stdio_serve_mode_for(Some(1)), StdioServeMode::Resumed);
}
#[test]
fn single_pack_verbs_unchanged() {
let catalog = build_verb_catalog([
t("kg", "create", "Create an entity or note."),
t("kg", "list", "List entities."),
]);
assert_eq!(
catalog,
" create — Create an entity or note.\n list — List entities.\n"
);
}
#[test]
fn duplicate_verb_concatenates_descriptions_with_pack_attribution() {
let catalog = build_verb_catalog([
t("kg", "create", "Create an entity or note."),
t("gtd", "create", "Create a task."),
]);
assert!(catalog.contains("[kg] Create an entity or note."));
assert!(catalog.contains("[gtd] Create a task."));
assert_eq!(catalog.matches(" create — ").count(), 1);
}
#[test]
fn instructions_carry_docs_address_and_guidance_pointers() {
let instructions = build_instructions(" create — Create an entity or note.\n", "kg, gtd");
assert!(instructions.contains("https://ohdearquant.github.io/khive/"));
assert!(instructions.contains("docs/configuration.md"));
assert!(instructions.contains("docs/guide/tips-and-tricks.md"));
assert!(instructions.contains("help=true"));
}
#[test]
fn catalog_is_sorted_alphabetically() {
let catalog = build_verb_catalog([
t("kg", "search", "Search."),
t("kg", "assign", "Assign."),
t("kg", "list", "List."),
]);
let names: Vec<&str> = catalog
.lines()
.filter(|l| l.starts_with(" "))
.map(|l| l.trim_start().split(' ').next().unwrap())
.collect();
assert_eq!(names, vec!["assign", "list", "search"]);
}
#[tokio::test]
async fn brain_dispatch_hook_updates_state_visible_through_same_instance() {
let config = RuntimeConfig {
db_path: None,
default_namespace: Namespace::local(),
embedding_model: None,
additional_embedding_models: vec![],
packs: vec!["kg".to_string(), "brain".to_string()],
..RuntimeConfig::default()
};
let runtime = KhiveRuntime::new(config).expect("in-memory runtime");
let server = KhiveMcpServer::with_packs(runtime, &["kg".to_string(), "brain".to_string()])
.expect("server builds with kg + brain");
server
.registry
.dispatch("brain.state", serde_json::Value::Null)
.await
.expect("brain.state loads the default namespace into the active slot");
server
.registry
.dispatch("stats", serde_json::json!({}))
.await
.expect("kg.stats dispatch succeeds");
let state = server
.registry
.dispatch("brain.state", serde_json::Value::Null)
.await
.expect("brain.state dispatch");
let total_events = state["balanced_recall"]["total_events"]
.as_u64()
.unwrap_or(0);
assert!(
total_events > 0,
"dispatch hook must update the same BrainPack instance the registry \
dispatches brain.* verbs to; got snapshot {state:?}"
);
}
struct CwdGuard {
original: std::path::PathBuf,
}
impl CwdGuard {
fn enter(dir: &std::path::Path) -> Self {
let original = std::env::current_dir().expect("read cwd");
std::env::set_current_dir(dir).expect("chdir into test project root");
Self { original }
}
}
impl Drop for CwdGuard {
fn drop(&mut self) {
let _ = std::env::set_current_dir(&self.original);
}
}
#[test]
#[serial]
fn config_id_does_not_collide_across_projects_with_same_relative_backend_path() {
use khive_runtime::{BackendId, BackendKind, KhiveConfig, Namespace};
let project_a = tempfile::tempdir().expect("project a tempdir");
let project_b = tempfile::tempdir().expect("project b tempdir");
let base_rt = RuntimeConfig {
db_path: None,
default_namespace: Namespace::parse("local").unwrap(),
embedding_model: None,
packs: vec!["kg".to_string()],
backend_id: BackendId::main(),
..RuntimeConfig::default()
};
let relative_backend_cfg = || KhiveConfig {
backends: vec![khive_runtime::BackendConfig {
name: "main".to_string(),
kind: BackendKind::Sqlite,
path: Some(std::path::PathBuf::from("./data/main.db")),
cache_mb: None,
journal_mode: None,
read_only: false,
}],
..KhiveConfig::default()
};
let id_a = {
let _cwd = CwdGuard::enter(project_a.path());
compute_config_id(&base_rt, Some(&relative_backend_cfg()))
};
let id_b = {
let _cwd = CwdGuard::enter(project_b.path());
compute_config_id(&base_rt, Some(&relative_backend_cfg()))
};
assert_ne!(
id_a, id_b,
"two projects declaring the same relative backend path string from \
different working directories must not share a config_id; both \
produced: {id_a}"
);
}
#[test]
#[serial]
fn config_id_does_not_collide_across_projects_with_same_relative_db_override() {
use khive_runtime::Namespace;
let project_a = tempfile::tempdir().expect("project a tempdir");
let project_b = tempfile::tempdir().expect("project b tempdir");
let rt_with_db = |db_path: Option<std::path::PathBuf>| RuntimeConfig {
db_path,
default_namespace: Namespace::parse("local").unwrap(),
embedding_model: None,
packs: vec!["kg".to_string()],
..RuntimeConfig::default()
};
let relative_db = std::path::PathBuf::from("./data/main.db");
let id_a = {
let _cwd = CwdGuard::enter(project_a.path());
compute_config_id(&rt_with_db(Some(relative_db.clone())), None)
};
let id_b = {
let _cwd = CwdGuard::enter(project_b.path());
compute_config_id(&rt_with_db(Some(relative_db.clone())), None)
};
assert_ne!(
id_a, id_b,
"two projects overriding KHIVE_DB with the same relative path from \
different working directories must not share a config_id; both \
produced: {id_a}"
);
let id_a_again = {
let _cwd = CwdGuard::enter(project_a.path());
compute_config_id(&rt_with_db(Some(relative_db.clone())), None)
};
assert_eq!(
id_a, id_a_again,
"resolving the same project's KHIVE_DB override twice must produce \
the same config_id"
);
}
fn nest_object(depth: usize, leaf: Value) -> Value {
let mut v = leaf;
for _ in 0..depth {
let mut map = serde_json::Map::with_capacity(1);
map.insert("nested".to_string(), v);
v = Value::Object(map);
}
v
}
#[test]
fn deep_nested_result_over_limit_is_flagged() {
let deep = nest_object(
khive_request::NESTING_DEPTH_LIMIT + 5,
json!({"leaf": true}),
);
let result_obj = json!({ "ok": true, "tool": "traverse", "result": deep });
assert!(
result_exceeds_depth_limit(&result_obj),
"result nested past NESTING_DEPTH_LIMIT must be flagged"
);
}
#[test]
fn result_at_exactly_the_depth_limit_is_not_flagged() {
let at_limit = nest_object(khive_request::NESTING_DEPTH_LIMIT, json!(true));
let result_obj = json!({ "ok": true, "tool": "traverse", "result": at_limit });
assert!(
!result_exceeds_depth_limit(&result_obj),
"result nested exactly at the limit must still be usable as $prev context"
);
}
#[test]
fn shallow_result_is_not_flagged() {
let shallow = json!({"a": {"b": {"c": 1}}});
let result_obj = json!({ "ok": true, "tool": "get", "result": shallow });
assert!(!result_exceeds_depth_limit(&result_obj));
}
#[test]
fn result_missing_field_is_not_flagged() {
let result_obj = json!({ "ok": false, "tool": "get", "error": "not found" });
assert!(!result_exceeds_depth_limit(&result_obj));
}
#[test]
fn chain_aggregation_seam_rejects_over_limit_result_via_iterative_drop() {
let deep = nest_object(khive_request::NESTING_DEPTH_LIMIT + 50_000, json!(true));
let mut envelope = serde_json::Map::with_capacity(3);
envelope.insert("ok".to_string(), Value::Bool(true));
envelope.insert("tool".to_string(), Value::String("traverse".to_string()));
envelope.insert("result".to_string(), deep);
let result_obj = Value::Object(envelope);
let err = chain_aggregation_depth_reject(result_obj)
.expect_err("result nested past NESTING_DEPTH_LIMIT must be rejected");
assert_eq!(err["ok"], json!(false));
assert_eq!(err["tool"], json!("traverse"));
assert_eq!(err["error"]["kind"], json!("result_too_deep"));
assert!(err.get("result").is_none());
}
#[test]
fn chain_aggregation_seam_accepts_result_within_limit_unchanged() {
let shallow = json!({ "ok": true, "tool": "get", "result": {"a": {"b": 1}} });
let accepted = chain_aggregation_depth_reject(shallow.clone())
.expect("result within the limit must be passed through unchanged");
assert_eq!(accepted, shallow);
}
#[test]
fn chain_seam_rejects_over_limit_result_before_envelope_build() {
let pathological = nest_object(khive_request::NESTING_DEPTH_LIMIT + 50_000, json!(true));
let err = chain_ok_envelope_or_depth_error("traverse".to_string(), pathological)
.expect_err("over-limit result must be rejected, not enveloped");
assert_eq!(err.0, "traverse");
assert_eq!(err.1["kind"], json!("result_too_deep"));
assert!(err.1.get("result").is_none());
assert!(err.1.get("nested").is_none());
}
#[test]
fn chain_seam_accepts_at_limit_result_and_moves_value_without_reserializing() {
let at_limit = nest_object(khive_request::NESTING_DEPTH_LIMIT, json!("leaf"));
let envelope = chain_ok_envelope_or_depth_error("get".to_string(), at_limit.clone())
.expect("result at exactly the limit must be accepted");
assert_eq!(envelope["ok"], json!(true));
assert_eq!(envelope["tool"], json!("get"));
assert_eq!(envelope["result"], at_limit);
}
#[test]
fn parallel_seam_rejects_over_limit_result_before_present() {
let pathological = nest_object(khive_request::NESTING_DEPTH_LIMIT + 50_000, json!(true));
let envelope = present_ok_envelope_or_depth_error(
"context".to_string(),
pathological,
PresentationMode::Agent,
0,
);
assert_eq!(envelope["ok"], json!(false));
assert_eq!(envelope["tool"], json!("context"));
assert_eq!(envelope["error"]["kind"], json!("result_too_deep"));
assert!(envelope["error"].get("result").is_none());
}
#[test]
fn parallel_seam_accepts_shallow_result_and_applies_presentation() {
let shallow = json!({"id": "11111111-1111-1111-1111-111111111111"});
let envelope = present_ok_envelope_or_depth_error(
"get".to_string(),
shallow,
PresentationMode::Verbose,
0,
);
assert_eq!(envelope["ok"], json!(true));
assert_eq!(
envelope["result"]["id"],
json!("11111111-1111-1111-1111-111111111111")
);
}
#[tokio::test]
async fn chain_with_deep_accumulated_prev_result_errors_cleanly() {
let config = RuntimeConfig {
db_path: None,
default_namespace: Namespace::local(),
embedding_model: None,
additional_embedding_models: vec![],
packs: vec!["kg".to_string()],
..RuntimeConfig::default()
};
let runtime = KhiveRuntime::new(config).expect("in-memory runtime");
let server = KhiveMcpServer::new(runtime).expect("server builds with kg");
let steps = khive_request::NESTING_DEPTH_LIMIT + 6;
let mut dsl = String::from(
r#"create(kind="entity", entity_kind="concept", name="d0", properties={"n": 0})"#,
);
for i in 1..steps {
dsl.push_str(&format!(
r#" | create(kind="entity", entity_kind="concept", name="d{i}", properties={{"inner": $prev.properties}})"#
));
}
let parsed = parse_request(&dsl).expect("each op's own args stay shallow; DSL must parse");
assert_eq!(parsed.mode, ExecutionMode::Chain);
let response = server
.run_parsed(
parsed.ops,
parsed.mode,
PresentationMode::Verbose,
None,
RunParsedContext {
enforce_response_budget: true,
from_wire: false,
identity: None,
},
)
.await;
let results = response["results"]
.as_array()
.expect("results must be an array");
assert_eq!(results.len(), steps);
let failure_idx = results
.iter()
.position(|r| r["ok"] == json!(false))
.expect("accumulated nesting must trip the depth guard before the chain completes");
assert_eq!(
results[failure_idx]["error"]["kind"],
json!("result_too_deep"),
"unexpected failure shape at index {failure_idx}: {:?}",
results[failure_idx]
);
for r in &results[failure_idx + 1..] {
assert_eq!(
r["aborted"],
json!(true),
"expected abort after the depth guard trips: {r:?}"
);
}
}
#[tokio::test]
async fn request_boundary_raw_control_bytes_reach_handler() {
let wire = "{\"ops\":\"create(kind=\\\"entity\\\", entity_kind=\\\"concept\\\", name=\\\"line1\\nline2\\\")\"}";
let params: RequestParams = serde_json::from_str(wire).expect("wire JSON deserializes");
assert!(
params.ops.contains('\n'),
"deserialized ops must carry a raw LF, not the two-char escape: {:?}",
params.ops
);
let config = RuntimeConfig {
db_path: None,
default_namespace: Namespace::local(),
embedding_model: None,
additional_embedding_models: vec![],
packs: vec!["kg".to_string()],
..RuntimeConfig::default()
};
let runtime = KhiveRuntime::new(config).expect("in-memory runtime");
let server = KhiveMcpServer::new(runtime).expect("server builds with kg");
let parsed = parse_request(¶ms.ops).expect("literal newline inside quotes must parse");
let response = server
.run_parsed(
parsed.ops,
parsed.mode,
PresentationMode::Verbose,
None,
RunParsedContext {
enforce_response_budget: true,
from_wire: false,
identity: None,
},
)
.await;
let results = response["results"]
.as_array()
.expect("results must be an array");
assert_eq!(results.len(), 1);
assert_eq!(
results[0]["ok"],
json!(true),
"unexpected result: {response:?}"
);
assert_eq!(results[0]["result"]["name"], json!("line1\nline2"));
}
fn make_daemon_save_to_test_server() -> KhiveMcpServer {
let config = RuntimeConfig {
db_path: None,
default_namespace: Namespace::parse("test").unwrap(),
embedding_model: None,
additional_embedding_models: vec![],
packs: vec!["kg".to_string()],
..RuntimeConfig::default()
};
let runtime = KhiveRuntime::new(config).expect("in-memory runtime");
KhiveMcpServer::new(runtime).expect("server builds with kg")
}
fn clear_daemon_env() {
std::env::remove_var("KHIVE_SOCKET");
std::env::remove_var("KHIVE_PID");
std::env::remove_var("KHIVE_NO_DAEMON");
std::env::remove_var("KHIVE_LOCK");
std::env::remove_var("KHIVE_PROCESS_REF");
}
#[cfg(unix)]
#[test]
fn wire_daemon_frame_forwards_request_id() {
let server = make_daemon_save_to_test_server();
let with_id = RequestParams {
ops: "stats()".to_string(),
request_id: Some(123),
..Default::default()
};
let frame = server.wire_daemon_frame(&with_id);
assert_eq!(frame.request_id, Some(123));
let without_id = RequestParams {
ops: "stats()".to_string(),
..Default::default()
};
let frame = server.wire_daemon_frame(&without_id);
assert_eq!(frame.request_id, None);
}
async fn find_audit_event_with_request_id(
store: &Arc<dyn khive_storage::EventStore>,
id: u64,
) -> Option<khive_storage::Event> {
let page = store
.query_events(
EventFilter::default(),
PageRequest {
limit: 50,
offset: 0,
},
)
.await
.expect("query_events must succeed");
page.items
.into_iter()
.find(|ev| ev.payload["resource"]["request_id"] == json!(id))
}
#[tokio::test]
#[serial]
async fn request_no_daemon_fallback_preserves_request_id_in_audit_event() {
clear_daemon_env();
std::env::set_var("KHIVE_NO_DAEMON", "1");
let server = make_daemon_save_to_test_server();
server
.request(Parameters(RequestParams {
ops: "stats(namespace=\"local\")".to_string(),
request_id: Some(9001),
..Default::default()
}))
.await
.expect("request() must succeed via local dispatch under KHIVE_NO_DAEMON");
let store = server
.event_store()
.expect("in-memory runtime must configure an EventStore");
let matched = find_audit_event_with_request_id(&store, 9001).await;
assert!(
matched.is_some(),
"KHIVE_NO_DAEMON local dispatch must stamp request_id onto the persisted \
audit event"
);
clear_daemon_env();
}
#[tokio::test]
#[serial]
async fn request_save_to_bypass_preserves_request_id_in_audit_event() {
clear_daemon_env();
let dir = tempfile::tempdir().expect("tempdir");
std::env::set_var("KHIVE_SAVE_TO_ROOT", dir.path());
let server = make_daemon_save_to_test_server();
let sink_path = dir.path().join("out.jsonl");
server
.request(Parameters(RequestParams {
ops: "stats(namespace=\"local\")".to_string(),
save_to: Some(sink_path.to_string_lossy().to_string()),
request_id: Some(9002),
..Default::default()
}))
.await
.expect("request() with save_to must succeed");
let store = server
.event_store()
.expect("in-memory runtime must configure an EventStore");
let matched = find_audit_event_with_request_id(&store, 9002).await;
assert!(
matched.is_some(),
"save_to local-dispatch bypass must stamp request_id onto the persisted \
audit event"
);
clear_daemon_env();
std::env::remove_var("KHIVE_SAVE_TO_ROOT");
}
#[cfg(unix)]
async fn connect_when_daemon_ready(sock: &std::path::Path) {
let deadline = tokio::time::Instant::now() + std::time::Duration::from_secs(5);
loop {
if tokio::net::UnixStream::connect(sock).await.is_ok() {
return;
}
assert!(
tokio::time::Instant::now() < deadline,
"daemon never bound {sock:?} within 5s"
);
tokio::time::sleep(std::time::Duration::from_millis(20)).await;
}
}
#[cfg(unix)]
#[tokio::test]
#[serial]
async fn request_save_to_bypasses_daemon_forwarding_and_writes_manifest() {
clear_daemon_env();
let dir = tempfile::tempdir().expect("tempdir");
let sock = dir.path().join("khived.sock");
let pid = dir.path().join("khived.pid");
std::env::set_var("KHIVE_SOCKET", &sock);
std::env::set_var("KHIVE_PID", &pid);
std::env::remove_var("KHIVE_NO_DAEMON");
std::env::set_var("KHIVE_SAVE_TO_ROOT", dir.path());
let server = make_daemon_save_to_test_server();
let daemon_server = server.clone();
let handle = tokio::spawn(async move {
let _ = khive_runtime::daemon::run_daemon(daemon_server).await;
});
connect_when_daemon_ready(&sock).await;
let sink_path = dir.path().join("out.jsonl");
let resp = server
.request(Parameters(RequestParams {
ops: "stats()".to_string(),
presentation: None,
presentation_per_op: None,
save_to: Some(sink_path.to_string_lossy().to_string()),
format: None,
format_per_op: None,
request_id: None,
}))
.await
.expect("request with save_to must succeed even with a warm daemon reachable");
let manifest: serde_json::Value =
serde_json::from_str(&resp).expect("response must be the save_to manifest JSON");
assert!(
manifest.get("rows").is_some() && manifest.get("path").is_some(),
"response must be the save_to manifest, not an inline daemon result; got: {resp}"
);
assert!(
sink_path.exists(),
"save_to file must be written even when a daemon is reachable"
);
let contents = std::fs::read_to_string(&sink_path).expect("read sink file");
assert!(
!contents.trim().is_empty(),
"sink file must contain JSONL content"
);
handle.abort();
let _ = handle.await;
clear_daemon_env();
std::env::remove_var("KHIVE_SAVE_TO_ROOT");
}
#[cfg(unix)]
#[tokio::test]
#[serial]
async fn request_returns_ambiguous_forward_error_without_local_double_dispatch() {
clear_daemon_env();
let dir = tempfile::tempdir().expect("tempdir");
let sock = dir.path().join("khived.sock");
let pid = dir.path().join("khived.pid");
std::env::set_var("KHIVE_SOCKET", &sock);
std::env::set_var("KHIVE_PID", &pid);
std::env::remove_var("KHIVE_NO_DAEMON");
let config = RuntimeConfig {
db_path: None,
default_namespace: Namespace::parse("test").unwrap(),
embedding_model: None,
additional_embedding_models: vec![],
packs: vec!["kg".to_string(), "comm".to_string()],
..RuntimeConfig::default()
};
let runtime = KhiveRuntime::new(config).expect("in-memory runtime");
let server = KhiveMcpServer::new(runtime).expect("server builds with kg + comm");
let listener =
tokio::net::UnixListener::bind(&sock).expect("bind fake crash-daemon socket");
let fake_handle = tokio::spawn(async move {
if let Ok((mut stream, _)) = listener.accept().await {
let _ = khive_runtime::daemon::read_frame(&mut stream).await;
}
});
let baseline = server
.dispatch_request_local(RequestParams {
ops: "stats()".to_string(),
presentation: None,
presentation_per_op: None,
save_to: None,
format: None,
format_per_op: None,
request_id: None,
})
.await
.expect("baseline stats() must succeed");
let resp = server
.request(Parameters(RequestParams {
ops: "comm.send(to=\"bob\", content=\"double-forward-probe\")".to_string(),
presentation: None,
presentation_per_op: None,
save_to: None,
format: None,
format_per_op: None,
request_id: None,
}))
.await;
match resp {
Err(McpError { message, .. }) => {
assert!(
message.contains(
"not retrying or locally dispatching to avoid duplicate execution"
),
"request() must surface forward_or_spawn's ambiguous-forward error \
verbatim, not a local dispatch result; got: {message}"
);
}
Ok(v) => panic!(
"request() must return the ambiguous-forward error directly, not fall \
through to local dispatch; got Ok({v})"
),
}
let after = server
.dispatch_request_local(RequestParams {
ops: "stats()".to_string(),
presentation: None,
presentation_per_op: None,
save_to: None,
format: None,
format_per_op: None,
request_id: None,
})
.await
.expect("post-request stats() must succeed");
assert_eq!(
after, baseline,
"the comm.send op must NEVER have run locally after the ambiguous \
forward outcome — a double-dispatch would mutate local state here"
);
let _ = tokio::time::timeout(std::time::Duration::from_secs(2), fake_handle).await;
clear_daemon_env();
}
#[cfg(unix)]
#[tokio::test]
#[serial]
async fn request_strict_fallback_lands_as_failed_op_envelope_not_rpc_error() {
clear_daemon_env();
crate::daemon::reset_fallback_counters();
let dir = tempfile::tempdir().expect("tempdir");
std::env::set_var("KHIVE_SOCKET", dir.path().join("khived.sock"));
std::env::set_var("KHIVE_PID", dir.path().join("khived.pid"));
std::env::set_var("KHIVE_LOCK", dir.path().join("khived.recovery.lock"));
std::env::remove_var("KHIVE_NO_DAEMON");
std::env::set_var("KHIVE_DAEMON_STRICT", "1");
let config = RuntimeConfig {
db_path: None,
default_namespace: Namespace::parse("test").unwrap(),
embedding_model: None,
additional_embedding_models: vec![],
packs: vec!["kg".to_string(), "comm".to_string()],
..RuntimeConfig::default()
};
let runtime = KhiveRuntime::new(config).expect("in-memory runtime");
let server = KhiveMcpServer::new(runtime).expect("server builds with kg + comm");
let baseline = server
.dispatch_request_local(RequestParams {
ops: "stats()".to_string(),
presentation: None,
presentation_per_op: None,
save_to: None,
format: None,
format_per_op: None,
request_id: None,
})
.await
.expect("baseline stats() must succeed");
fn assert_fallback_error(entry: &Value, tool: &str) {
assert_eq!(entry["ok"], json!(false), "entry: {entry}");
assert_eq!(entry["tool"], json!(tool), "entry: {entry}");
let msg = entry["error"].as_str().expect("error must be a string");
assert!(
msg.contains("KHIVE_DAEMON_STRICT"),
"error must name the strict mode that rejected the fallback: {msg}"
);
assert!(
msg.contains("respawn_failed"),
"error must name the confirmed respawn failure: {msg}"
);
assert!(
msg.contains("make local"),
"error must include the safe respawn remediation: {msg}"
);
}
let single_resp = server
.request(Parameters(RequestParams {
ops: "comm.send(to=\"bob\", content=\"strict-single-probe\")".to_string(),
presentation: None,
presentation_per_op: None,
save_to: None,
format: None,
format_per_op: None,
request_id: None,
}))
.await
.expect("strict fallback must land as a normal Ok(envelope), not Err(McpError)");
let single: Value =
serde_json::from_str(&single_resp).expect("response must be the request envelope");
assert_eq!(
single["results"].as_array().expect("results array").len(),
1
);
assert_fallback_error(&single["results"][0], "comm.send");
assert_eq!(
single["summary"],
json!({ "total": 1, "succeeded": 0, "failed": 1, "aborted": 0 })
);
let batch_resp = server
.request(Parameters(RequestParams {
ops: "[comm.send(to=\"bob\", content=\"strict-batch-1\"), \
comm.send(to=\"bob\", content=\"strict-batch-2\")]"
.to_string(),
presentation: None,
presentation_per_op: None,
save_to: None,
format: None,
format_per_op: None,
request_id: None,
}))
.await
.expect("strict fallback must land as a normal Ok(envelope), not Err(McpError)");
let batch: Value =
serde_json::from_str(&batch_resp).expect("response must be the request envelope");
let batch_results = batch["results"].as_array().expect("results array");
assert_eq!(batch_results.len(), 2);
for entry in batch_results {
assert_fallback_error(entry, "comm.send");
}
assert_eq!(
batch["summary"],
json!({ "total": 2, "succeeded": 0, "failed": 2, "aborted": 0 })
);
let chain_resp = server
.request(Parameters(RequestParams {
ops: "comm.send(to=\"bob\", content=\"strict-chain-1\") | \
comm.send(to=\"bob\", content=\"strict-chain-2\")"
.to_string(),
presentation: None,
presentation_per_op: None,
save_to: None,
format: None,
format_per_op: None,
request_id: None,
}))
.await
.expect("strict fallback must land as a normal Ok(envelope), not Err(McpError)");
let chain: Value =
serde_json::from_str(&chain_resp).expect("response must be the request envelope");
let chain_results = chain["results"].as_array().expect("results array");
assert_eq!(chain_results.len(), 2);
assert_fallback_error(&chain_results[0], "comm.send");
assert_eq!(
chain_results[1],
json!({ "ok": false, "tool": "comm.send", "aborted": true })
);
assert_eq!(
chain["summary"],
json!({ "total": 2, "succeeded": 0, "failed": 1, "aborted": 1 })
);
let after = server
.dispatch_request_local(RequestParams {
ops: "stats()".to_string(),
presentation: None,
presentation_per_op: None,
save_to: None,
format: None,
format_per_op: None,
request_id: None,
})
.await
.expect("post-request stats() must succeed");
assert_eq!(
after, baseline,
"no comm.send op must ever have run locally under strict-mode fallback \
rejection — a local dispatch would mutate local state here"
);
crate::daemon::reset_fallback_counters();
clear_daemon_env();
}
fn in_memory_kg_server() -> KhiveMcpServer {
let config = RuntimeConfig {
db_path: None,
default_namespace: Namespace::local(),
embedding_model: None,
additional_embedding_models: vec![],
packs: vec!["kg".to_string()],
..RuntimeConfig::default()
};
let runtime = KhiveRuntime::new(config).expect("in-memory runtime");
KhiveMcpServer::new(runtime).expect("server builds with kg")
}
#[tokio::test]
async fn request_status_is_success_when_every_op_in_batch_succeeds() {
let server = in_memory_kg_server();
let resp = server
.dispatch_request_local(RequestParams {
ops: "[create(kind=\"entity\", entity_kind=\"concept\", name=\"status-ok-1\"), \
create(kind=\"entity\", entity_kind=\"concept\", name=\"status-ok-2\")]"
.to_string(),
presentation: None,
presentation_per_op: None,
save_to: None,
format: None,
format_per_op: None,
request_id: None,
})
.await
.expect("batch dispatch must succeed");
let parsed: Value = serde_json::from_str(&resp).expect("envelope must be JSON");
assert_eq!(parsed["summary"]["failed"], 0);
assert_eq!(
parsed["status"], "success",
"an all-succeeding batch must report status=success; got {parsed}"
);
}
#[tokio::test]
async fn request_status_is_partial_when_a_batch_op_fails() {
let server = in_memory_kg_server();
let resp = server
.dispatch_request_local(RequestParams {
ops:
"[create(kind=\"entity\", entity_kind=\"concept\", name=\"status-partial-1\"), \
search(kind=\"not_a_real_kind\", query=\"x\")]"
.to_string(),
presentation: None,
presentation_per_op: None,
save_to: None,
format: None,
format_per_op: None,
request_id: None,
})
.await
.expect("batch dispatch must succeed at the RPC level even with a failed op");
let parsed: Value = serde_json::from_str(&resp).expect("envelope must be JSON");
assert!(
parsed["summary"]["failed"].as_u64().unwrap_or(0) > 0,
"expected at least one failed op; got {parsed}"
);
assert_eq!(
parsed["status"], "partial",
"a batch with a failed op must report status=partial; got {parsed}"
);
}
#[tokio::test]
async fn request_status_is_partial_when_a_chain_op_is_aborted() {
let server = in_memory_kg_server();
let resp = server
.dispatch_request_local(RequestParams {
ops: "search(kind=\"not_a_real_kind\", query=\"x\") | \
create(kind=\"entity\", entity_kind=\"concept\", name=\"status-chain-aborted\")"
.to_string(),
presentation: None,
presentation_per_op: None,
save_to: None,
format: None,
format_per_op: None,
request_id: None,
})
.await
.expect("chain dispatch must succeed at the RPC level even with an aborted op");
let parsed: Value = serde_json::from_str(&resp).expect("envelope must be JSON");
assert!(
parsed["summary"]["aborted"].as_u64().unwrap_or(0) > 0,
"expected the second chain op to be aborted; got {parsed}"
);
assert_eq!(
parsed["status"], "partial",
"a chain with an aborted op must report status=partial; got {parsed}"
);
}
}