use std::sync::Arc;
use anyhow::Result;
use clap::{Parser, ValueEnum};
use rmcp::{
ServiceExt,
transport::{
sse_server::{SseServer, SseServerConfig},
stdio,
},
};
use terraphim_automata::builder::{Logseq, ThesaurusBuilder};
use terraphim_config::{Config, ConfigBuilder, ConfigState};
use terraphim_file_search::kg_scorer::KgPathScorer;
use terraphim_mcp_server::McpService;
use terraphim_types::RoleName;
use tracing::{Level, info, warn};
#[derive(Parser, Debug)]
#[command(name = "terraphim_mcp_server")]
#[command(about = "Terraphim MCP server with configurable profile")]
#[command(version)]
struct Args {
#[arg(short, long, value_enum, default_value_t = ConfigProfile::Desktop)]
profile: ConfigProfile,
#[arg(short, long)]
verbose: bool,
#[arg(long, default_value_t = false)]
sse: bool,
#[arg(long, default_value = "127.0.0.1:8000")]
bind: String,
}
#[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord, ValueEnum, Debug)]
enum ConfigProfile {
Desktop,
Server,
}
async fn shutdown_signal() {
tokio::signal::ctrl_c()
.await
.expect("Failed to listen for ctrl+c");
}
fn build_profile_config(profile: &ConfigProfile) -> terraphim_config::Config {
match profile {
ConfigProfile::Desktop => {
info!("Using desktop configuration (Terraphim Engineer role with local KG)");
ConfigBuilder::new()
.build_default_desktop()
.build()
.expect("Failed to build default desktop configuration")
}
ConfigProfile::Server => {
info!("Using server configuration (Default role without KG)");
ConfigBuilder::new()
.build_default_server()
.build()
.expect("Failed to build default server configuration")
}
}
}
fn merge_project_into_base(
mut config: Config,
project_config: &terraphim_config::project::ProjectConfig,
) -> Config {
if let Some(ref shortcut) = project_config.global_shortcut {
config.global_shortcut = Some(shortcut.clone());
}
for (name, role) in &project_config.roles {
config.roles.insert(RoleName::new(name), role.clone());
}
if let Ok(Some(role_name)) = project_config.resolve_role_name(None) {
let role_name = RoleName::new(&role_name);
if config.roles.contains_key(&role_name) {
config.selected_role = role_name.clone();
config.default_role = role_name;
}
}
repair_selected_roles(&mut config);
config
}
fn repair_selected_roles(config: &mut Config) {
if config.roles.contains_key(&config.selected_role)
&& config.roles.contains_key(&config.default_role)
{
return;
}
if let Some(first_role) = config.roles.keys().next().cloned() {
if !config.roles.contains_key(&config.selected_role) {
config.selected_role = first_role.clone();
}
if !config.roles.contains_key(&config.default_role) {
config.default_role = first_role;
}
}
}
#[tokio::main]
async fn main() -> Result<()> {
let args = Args::parse();
let level = if args.verbose {
Level::DEBUG
} else {
Level::INFO
};
let subscriber = tracing_subscriber::fmt()
.with_max_level(level)
.with_env_filter(
tracing_subscriber::EnvFilter::from_default_env().add_directive(level.into()),
);
if args.sse {
subscriber.init();
} else {
subscriber
.without_time()
.with_writer(std::io::stderr)
.init();
}
info!("Starting Terraphim MCP Server...");
info!("Args: {:?}", args);
let config = if let Ok(Some(project_dir)) = terraphim_config::project::discover(None) {
match terraphim_config::project::ProjectConfig::load_from_dir(&project_dir) {
Ok(project_config) if !project_config.is_empty() => {
info!(
"Using project configuration from '{}' ({} role(s))",
project_dir.display(),
project_config.roles.len()
);
let base = build_profile_config(&args.profile);
merge_project_into_base(base, &project_config)
}
_ => {
info!(
"No project roles found in '{}', using profile",
project_dir.display()
);
build_profile_config(&args.profile)
}
}
} else {
build_profile_config(&args.profile)
};
let mut temp_config = config.clone();
let config_state = ConfigState::new(&mut temp_config)
.await
.expect("Failed to create config state from config");
let config_state = Arc::new(config_state);
let (selected_role, kg_path) = {
let cfg = config_state.config.lock().await;
let selected = cfg.selected_role.clone();
let path = cfg
.roles
.get(&selected)
.and_then(|r| r.kg.as_ref())
.and_then(|kg| kg.knowledge_graph_local.as_ref())
.map(|kgl| kgl.path.clone());
(selected, path)
};
let mut service = McpService::new(Arc::clone(&config_state));
if let Some(kg_path) = kg_path {
info!("Building KgPathScorer from KG path: {:?}", kg_path);
let builder = Logseq::default();
let role_name = selected_role.as_lowercase().to_string();
match builder.build(role_name, kg_path).await {
Ok(thesaurus) => {
let term_count = thesaurus.len();
let scorer = Arc::new(KgPathScorer::new(thesaurus));
service = service.with_kg_scorer(scorer);
info!(
"KgPathScorer wired with {} terms for role '{}'",
term_count, selected_role
);
}
Err(e) => {
warn!(
"Failed to build thesaurus for KgPathScorer (role '{}'): {}",
selected_role, e
);
}
}
}
if args.sse {
info!("Starting SSE server on {}", args.bind);
let config = SseServerConfig {
bind: args.bind.parse().expect("Invalid bind address"),
sse_path: "/sse".to_string(),
post_path: "/message".to_string(),
ct: tokio_util::sync::CancellationToken::new(),
sse_keep_alive: None,
};
let (sse_server, router) = SseServer::new(config);
let listener = tokio::net::TcpListener::bind(sse_server.config.bind).await?;
let ct = sse_server.config.ct.child_token();
let server = axum::serve(listener, router).with_graceful_shutdown(async move {
ct.cancelled().await;
info!("SSE server cancelled");
});
tokio::spawn(async move {
if let Err(e) = server.await {
tracing::error!(error = %e, "SSE server shutdown with error");
}
});
let _ct = sse_server.with_service(move || service.clone());
shutdown_signal().await;
} else {
info!("Starting stdio server");
service.init_autocomplete_default().await;
info!("Initialized Terraphim MCP service");
let mcp_service = service.serve(stdio()).await?;
mcp_service.waiting().await?;
}
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
use std::collections::HashMap;
#[test]
fn repair_selected_roles_sets_missing_selected_to_existing_role() {
let mut config = Config::default();
config.roles.clear();
config.roles.insert(
RoleName::new("devops"),
terraphim_config::Role::new("DevOps"),
);
repair_selected_roles(&mut config);
assert_eq!(config.selected_role, RoleName::new("devops"));
assert_eq!(config.default_role, RoleName::new("devops"));
}
#[test]
fn merge_project_into_base_selects_single_project_role() {
let base = build_profile_config(&ConfigProfile::Desktop);
let mut roles = HashMap::new();
roles.insert("devops".to_string(), terraphim_config::Role::new("DevOps"));
let project_config = terraphim_config::project::ProjectConfig {
roles,
..Default::default()
};
let config = merge_project_into_base(base, &project_config);
assert!(config.roles.contains_key(&RoleName::new("devops")));
assert_eq!(config.selected_role, RoleName::new("devops"));
assert_eq!(config.default_role, RoleName::new("devops"));
}
}