use std::path::PathBuf;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, RwLock};
use anyhow::{Context, Result};
use clap::Parser;
use mcp_methods::server::manifest::{
find_sibling_manifest, find_workspace_manifest, ManifestError,
};
use mcp_methods::server::{
init_tracing, load_env_for_mode, maybe_watch, resolve_source_roots, serve_prompts, watch,
workspace, BundledSkill, Manifest, McpServer, PredicateClause, ServerOptions,
SkillPredicateEvaluator, SkillRegistry, WorkspaceKind,
};
use rmcp::transport::stdio;
use rmcp::ServiceExt;
mod code_source;
mod csv_http;
mod cypher_tools;
mod explore;
mod tools;
use crate::tools::GraphState;
#[derive(Parser, Debug)]
#[command(
name = "kglite-mcp-server",
about = "MCP server for KGLite knowledge graphs (Rust-native)"
)]
struct Cli {
#[arg(long, conflicts_with_all = ["workspace", "watch", "source_root"])]
graph: Option<PathBuf>,
#[arg(long = "source-root", conflicts_with_all = ["graph", "workspace", "watch"])]
source_root: Option<PathBuf>,
#[arg(long, conflicts_with_all = ["graph", "source_root", "watch"])]
workspace: Option<PathBuf>,
#[arg(long, conflicts_with_all = ["graph", "source_root", "workspace"])]
watch: Option<PathBuf>,
#[arg(long = "mcp-config")]
mcp_config: Option<PathBuf>,
#[arg(long)]
name: Option<String>,
#[arg(long = "trust-tools")]
#[allow(dead_code)]
trust_tools: bool,
#[arg(long = "stale-after-days", default_value_t = 7)]
stale_after_days: u32,
}
#[derive(Debug, Clone)]
enum Mode {
Graph {
path: PathBuf,
},
SourceRoot {
dir: PathBuf,
},
Workspace {
dir: PathBuf,
},
LocalWorkspace {
root: PathBuf,
watch: bool,
},
Watch {
dir: PathBuf,
},
Bare,
}
fn pick_mode(cli: &Cli) -> Mode {
if let Some(p) = &cli.graph {
Mode::Graph { path: p.clone() }
} else if let Some(d) = &cli.source_root {
Mode::SourceRoot { dir: d.clone() }
} else if let Some(d) = &cli.workspace {
Mode::Workspace { dir: d.clone() }
} else if let Some(d) = &cli.watch {
Mode::Watch { dir: d.clone() }
} else {
Mode::Bare
}
}
fn fallback_name(mode: &Mode) -> &'static str {
match mode {
Mode::Graph { .. } => "KGLite (single-graph)",
Mode::SourceRoot { .. } => "KGLite (source-root)",
Mode::Workspace { .. } => "KGLite (workspace)",
Mode::LocalWorkspace { .. } => "KGLite (local-workspace)",
Mode::Watch { .. } => "KGLite (watch)",
Mode::Bare => "KGLite",
}
}
fn default_manifest_path(mode: &Mode) -> Option<PathBuf> {
match mode {
Mode::Graph { path } => find_sibling_manifest(path),
Mode::Workspace { dir } | Mode::Watch { dir } => find_workspace_manifest(dir),
Mode::LocalWorkspace { root, .. } => find_workspace_manifest(root),
Mode::SourceRoot { .. } | Mode::Bare => None,
}
}
fn load_manifest(cli: &Cli, mode: &Mode) -> Result<Option<Manifest>, ManifestError> {
let path = match &cli.mcp_config {
Some(p) if !p.is_file() => {
return Err(ManifestError::bare(format!(
"--mcp-config path does not exist: {}",
p.display()
)))
}
Some(p) => Some(p.clone()),
None => default_manifest_path(mode),
};
match path {
Some(p) => Ok(Some(mcp_methods::server::load_manifest(&p)?)),
None => Ok(None),
}
}
#[tokio::main]
async fn main() -> Result<()> {
init_tracing();
let cli = Cli::parse();
let mut mode = pick_mode(&cli);
if let Mode::Graph { path } = &mode {
if !path.is_file() {
anyhow::bail!("--graph path does not exist: {}", path.display());
}
}
if let Mode::SourceRoot { dir } | Mode::Watch { dir } = &mode {
if !dir.is_dir() {
anyhow::bail!(
"path does not exist or is not a directory: {}",
dir.display()
);
}
}
let manifest = load_manifest(&cli, &mode).context("manifest load failed")?;
if let Some(m) = manifest.as_ref() {
if let Some(wcfg) = m.workspace.as_ref() {
if wcfg.kind == WorkspaceKind::Local {
let raw_root = wcfg.root.as_ref().ok_or_else(|| {
anyhow::anyhow!("manifest.workspace.kind=local is missing required `root`")
})?;
let base = m
.yaml_path
.parent()
.map(|p| p.to_path_buf())
.unwrap_or_else(|| PathBuf::from("."));
let resolved = base.join(raw_root).canonicalize().with_context(|| {
format!("workspace.root {raw_root:?} resolves to a path that does not exist")
})?;
mode = Mode::LocalWorkspace {
root: resolved,
watch: wcfg.watch,
};
}
}
}
let env_start_dir: PathBuf = match &mode {
Mode::Graph { path } => path
.canonicalize()
.ok()
.and_then(|p| p.parent().map(PathBuf::from))
.unwrap_or_else(|| std::env::current_dir().unwrap_or_else(|_| PathBuf::from("."))),
Mode::SourceRoot { dir } | Mode::Workspace { dir } | Mode::Watch { dir } => dir.clone(),
Mode::LocalWorkspace { root, .. } => root.clone(),
Mode::Bare => std::env::current_dir().unwrap_or_else(|_| PathBuf::from(".")),
};
let env_file_loaded = load_env_for_mode(manifest.as_ref(), &env_start_dir)
.context("manifest env_file load failed")?;
let mut options = ServerOptions::from_manifest(manifest.as_ref(), fallback_name(&mode));
if cli.name.is_some() {
options.name = cli.name.clone();
}
let graph_state = GraphState::new();
let local_active_root: Arc<RwLock<Option<PathBuf>>> = Arc::new(RwLock::new(None));
match &mode {
Mode::Graph { path } => {
let canon = path.canonicalize()?;
graph_state.load_kgl(&canon).context("kglite.load failed")?;
let manifest_roots = manifest
.as_ref()
.filter(|m| !m.source_roots.is_empty())
.map(resolve_source_roots)
.transpose()
.context("manifest source_root resolution failed")?;
let roots = if let Some(rs) = manifest_roots {
rs
} else if let Some(parent) = canon.parent() {
vec![parent.to_string_lossy().into_owned()]
} else {
Vec::new()
};
if !roots.is_empty() {
options = options.with_static_source_roots(roots);
}
}
Mode::SourceRoot { dir } | Mode::Watch { dir } => {
let canon = dir.canonicalize()?;
options = options.with_static_source_roots(vec![canon.to_string_lossy().into_owned()]);
if matches!(mode, Mode::Watch { .. }) {
graph_state
.build_code_tree(&canon)
.context("initial code_tree build failed")?;
}
}
Mode::Workspace { dir } => {
let canon = dir.canonicalize().unwrap_or_else(|_| dir.clone());
let gs = graph_state.clone();
let hook: workspace::PostActivateHook = Arc::new(move |path, name| {
tracing::info!(repo = name, "code_tree::build on activate");
gs.build_code_tree(path)
});
let ws = workspace::Workspace::open(canon, cli.stale_after_days, Some(hook))
.context("workspace init failed")?;
options = options.with_workspace(ws);
}
Mode::LocalWorkspace { root, .. } => {
let gs = graph_state.clone();
let initial_activate_seen = Arc::new(AtomicBool::new(false));
let initial_flag = initial_activate_seen.clone();
let active_root_for_hook = local_active_root.clone();
let hook: workspace::PostActivateHook = Arc::new(move |path, name| {
if !initial_flag.swap(true, Ordering::SeqCst) {
tracing::info!(
root = %path.display(),
"deferring local-workspace code_tree build until first set_root_dir"
);
return Ok(());
}
if let Ok(mut guard) = active_root_for_hook.write() {
*guard = Some(path.to_path_buf());
}
tracing::info!(repo = name, "code_tree::build on local-workspace activate");
gs.build_code_tree(path)
});
let ws = workspace::Workspace::open_local(root.clone(), Some(hook))
.context("local-workspace init failed")?;
options = options.with_workspace(ws);
}
Mode::Bare => {
if let Some(m) = manifest.as_ref() {
if !m.source_roots.is_empty() {
let resolved =
resolve_source_roots(m).context("source root resolution failed")?;
options = options.with_static_source_roots(resolved);
}
}
}
}
let source_roots_provider = options.source_roots.clone();
let manifest_base: PathBuf = manifest
.as_ref()
.and_then(|m| m.yaml_path.parent().map(|p| p.to_path_buf()))
.unwrap_or_else(|| std::env::current_dir().unwrap_or_else(|_| PathBuf::from(".")));
let csv_http_cfg = match manifest.as_ref() {
Some(m) => match m.extensions.get("csv_http_server") {
Some(raw) => csv_http::CsvHttpConfig::from_manifest_value(raw, &manifest_base)
.context("extensions.csv_http_server parse failed")?,
None => None,
},
None => None,
};
if let Some(cfg) = csv_http_cfg.as_ref() {
csv_http::spawn(cfg.clone())
.await
.context("csv_http_server failed to bind")?;
}
let builtins = tools::Builtins {
save_graph: manifest
.as_ref()
.map(|m| m.builtins.save_graph)
.unwrap_or(false),
temp_cleanup_on_overview: manifest
.as_ref()
.map(|m| {
matches!(
m.builtins.temp_cleanup,
mcp_methods::server::manifest::TempCleanup::OnOverview
)
})
.unwrap_or(false),
temp_dir: Some(
csv_http_cfg
.as_ref()
.map(|c| c.dir.clone())
.unwrap_or_else(|| manifest_base.join("temp")),
),
};
let csv_http_arc = csv_http_cfg.map(Arc::new);
let mut server = McpServer::new(options);
tools::register(
&mut server,
graph_state.clone(),
builtins,
csv_http_arc.clone(),
);
code_source::register(
&mut server,
graph_state.clone(),
source_roots_provider.clone(),
)
.context("read_code_source registration failed")?;
explore::register(
&mut server,
graph_state.clone(),
source_roots_provider.clone(),
)
.context("explore registration failed")?;
if let Some(m) = manifest.as_ref() {
if let Some(embedder) = build_embedder_from_manifest(m)? {
graph_state
.bind_embedder(embedder)
.context("graph.set_embedder_native failed")?;
}
}
if let Some(m) = manifest.as_ref() {
let runner = cypher_tools::make_runner(graph_state.clone(), csv_http_arc.clone());
let registered = cypher_tools::register_cypher_tools(&mut server, m, runner)
.context("YAML cypher tool registration failed")?;
if registered > 0 {
tracing::info!(count = registered, "manifest cypher tools registered");
}
}
let watch_handle = match &mode {
Mode::Watch { dir } => {
let canon = dir.canonicalize()?;
let gs = graph_state.clone();
let cb: watch::ChangeHandler = Arc::new(move |paths| {
let any_code = paths
.iter()
.any(|p| kglite::api::language_for_path(p).is_some());
if !any_code {
return;
}
gs.tag_code_tree_dirty(canon.clone());
});
maybe_watch(Some(dir), Some(cb))?
}
Mode::LocalWorkspace { root, watch: true } => {
let gs = graph_state.clone();
let active_root_for_watch = local_active_root.clone();
let cb: watch::ChangeHandler = Arc::new(move |paths| {
let active = match active_root_for_watch.read() {
Ok(g) => g.clone(),
Err(_) => return,
};
let Some(active) = active else {
return;
};
let any_under_active_and_code = paths
.iter()
.any(|p| p.starts_with(&active) && kglite::api::language_for_path(p).is_some());
if !any_under_active_and_code {
return;
}
gs.tag_code_tree_dirty(active.clone());
});
maybe_watch(Some(root), Some(cb))?
}
_ => None,
};
let _watch_handle = watch_handle;
if let Some(m) = manifest.as_ref() {
let registry_result = SkillRegistry::new()
.add_bundled(BundledSkill {
name: "cypher_query",
body: include_str!("../skills/cypher_query.md"),
})
.add_bundled(BundledSkill {
name: "graph_overview",
body: include_str!("../skills/graph_overview.md"),
})
.add_bundled(BundledSkill {
name: "save_graph",
body: include_str!("../skills/save_graph.md"),
})
.add_bundled(BundledSkill {
name: "read_code_source",
body: include_str!("../skills/read_code_source.md"),
})
.add_bundled(BundledSkill {
name: "explore",
body: include_str!("../skills/explore.md"),
})
.merge_framework_defaults()
.auto_detect_project_layer(&m.yaml_path)
.layer_dirs(&m.skills, &m.yaml_path)
.and_then(|r| {
r.with_predicate_evaluator(KglitePredicateEvaluator {
state: graph_state.clone(),
})
.finalise()
});
match registry_result {
Ok(registry) => serve_prompts(®istry, &mut server),
Err(e) => {
tracing::warn!(error = %e, "skills registry build failed; skills disabled for this session");
}
}
}
print_boot_summary(
&mode,
manifest.as_ref(),
&graph_state,
env_file_loaded.as_deref(),
);
let service = server
.serve(stdio())
.await
.context("failed to start MCP service over stdio")?;
service.waiting().await?;
Ok(())
}
struct KglitePredicateEvaluator {
state: GraphState,
}
impl SkillPredicateEvaluator for KglitePredicateEvaluator {
fn evaluate(&self, clause: &PredicateClause<'_>) -> Option<bool> {
match clause {
PredicateClause::GraphHasNodeType(types) => {
Some(types.iter().any(|t| self.state.has_node_type(t)))
}
PredicateClause::GraphHasProperty {
node_type,
prop_name,
} => Some(self.state.has_property(node_type, prop_name)),
_ => None,
}
}
}
fn build_embedder_from_manifest(
manifest: &Manifest,
) -> Result<Option<Arc<dyn kglite::api::Embedder>>> {
let Some(raw) = manifest.extensions.get("embedder") else {
return Ok(None);
};
let obj = raw
.as_object()
.ok_or_else(|| anyhow::anyhow!("extensions.embedder must be a mapping (got: {raw:?})"))?;
let backend = obj
.get("backend")
.and_then(|v| v.as_str())
.unwrap_or("fastembed");
match backend {
#[cfg(feature = "fastembed")]
"fastembed" => {
let model = obj.get("model").and_then(|v| v.as_str()).ok_or_else(|| {
anyhow::anyhow!("extensions.embedder.model is required for the fastembed backend")
})?;
let adapter = kglite::api::FastEmbedAdapter::new(model)
.map_err(|e| anyhow::anyhow!("fastembed init failed: {e}"))?;
tracing::info!(model, backend, "registered Rust-native embedder");
Ok(Some(Arc::new(adapter)))
}
#[cfg(not(feature = "fastembed"))]
"fastembed" => anyhow::bail!(
"extensions.embedder.backend = \"fastembed\" requires this binary \
to be built with the `fastembed` feature enabled. Rebuild with: \
`cargo install kglite-mcp-server --features fastembed`. The \
default build excludes fastembed because its ort-sys dependency \
has a flaky upstream binary download — opt in only when you need \
text_score() semantic search."
),
other => anyhow::bail!(
"extensions.embedder.backend = {other:?} is not supported. \
Known: fastembed (requires --features fastembed at install time)."
),
}
}
fn print_boot_summary(
mode: &Mode,
manifest: Option<&Manifest>,
graph_state: &GraphState,
env_file_loaded: Option<&std::path::Path>,
) {
let label = match mode {
Mode::Graph { path } => format!("graph [{}]", path.display()),
Mode::SourceRoot { dir } => format!("source-root [{}]", dir.display()),
Mode::Workspace { dir } => format!("workspace [{}]", dir.display()),
Mode::LocalWorkspace { root, watch } => format!(
"local-workspace [{}{}]",
root.display(),
if *watch { " +watch" } else { "" }
),
Mode::Watch { dir } => format!("watch [{}]", dir.display()),
Mode::Bare => "bare".to_string(),
};
let mut parts = vec![format!("mode: {label}")];
if let Some(p) = env_file_loaded {
parts.push(format!("env: {}", p.display()));
} else {
parts.push("env: (no .env found)".to_string());
}
if let Some(m) = manifest {
parts.push(format!("manifest: {}", m.yaml_path.display()));
}
if let Some((nodes, edges)) = graph_state.schema() {
parts.push(format!("graph: {nodes} nodes, {edges} edges"));
}
eprintln!("kglite-mcp-server: {}", parts.join("; "));
}