use std::path::PathBuf;
use std::sync::Arc;
use crate::core::{
db::Database,
types::{Drawer, SourceType, Triple},
utils::{build_drawer_id, build_triple_id, current_timestamp, source_file_or_synthetic},
};
use crate::embed::EmbedderFactory;
use crate::search::{resolve_route, search_with_vector};
use anyhow::Context;
use rmcp::{
ErrorData, Json, ServerHandler, ServiceExt,
handler::server::{router::tool::ToolRouter, wrapper::Parameters},
model::{ServerCapabilities, ServerInfo},
tool, tool_handler, tool_router,
};
use super::tools::{
DeleteRequest, DeleteResponse, DuplicateWarning, IngestRequest, IngestResponse, KgRequest,
KgResponse, KgStatsDto, ScopeCount, SearchRequest, SearchResponse, SearchResultDto,
StatusResponse, TaxonomyEntryDto, TaxonomyRequest, TaxonomyResponse, TripleDto, TunnelDto,
TunnelsResponse,
};
#[derive(Clone)]
pub struct MempalMcpServer {
db_path: PathBuf,
embedder_factory: Arc<dyn EmbedderFactory>,
tool_router: ToolRouter<Self>,
}
impl MempalMcpServer {
pub fn new(db_path: PathBuf, config: crate::core::config::Config) -> Self {
Self::new_with_factory(
db_path,
Arc::new(crate::embed::ConfiguredEmbedderFactory::new(config)),
)
}
pub fn new_with_factory(db_path: PathBuf, embedder_factory: Arc<dyn EmbedderFactory>) -> Self {
Self {
db_path,
embedder_factory,
tool_router: Self::tool_router(),
}
}
pub async fn serve_stdio(
self,
) -> anyhow::Result<rmcp::service::RunningService<rmcp::RoleServer, Self>> {
self.serve(rmcp::transport::stdio())
.await
.context("failed to initialize MCP stdio transport")
}
fn open_db(&self) -> std::result::Result<Database, ErrorData> {
Database::open(&self.db_path).map_err(|error| {
ErrorData::internal_error(format!("failed to open database: {error}"), None)
})
}
}
#[tool_router(router = tool_router)]
impl MempalMcpServer {
#[tool(
name = "mempal_status",
description = "Return schema version, drawer counts, taxonomy counts, database size, scope breakdown, the AAAK format spec, and the memory protocol. Call once at session start if you haven't seen the protocol yet."
)]
async fn mempal_status(&self) -> std::result::Result<Json<StatusResponse>, ErrorData> {
let db = self.open_db()?;
let schema_version = db.schema_version().map_err(db_error)?;
let drawer_count = db.drawer_count().map_err(db_error)?;
let taxonomy_count = db.taxonomy_count().map_err(db_error)?;
let db_size_bytes = db.database_size_bytes().map_err(db_error)?;
let scopes = db
.scope_counts()
.map_err(db_error)?
.into_iter()
.map(|(wing, room, drawer_count)| ScopeCount {
wing,
room,
drawer_count,
})
.collect();
Ok(Json(StatusResponse {
schema_version,
drawer_count,
taxonomy_count,
db_size_bytes,
scopes,
aaak_spec: crate::aaak::generate_spec(),
memory_protocol: crate::core::protocol::MEMORY_PROTOCOL.to_string(),
}))
}
#[tool(
name = "mempal_search",
description = "Search persistent project memory via vector embedding with optional wing/room filters. PREFER THIS over grepping files or guessing from general knowledge when answering ANY project-specific question — past decisions, design rationale, implementation details, bug history, how a component works, why something was built a certain way, or any other project knowledge. Every result includes drawer_id and source_file for citation."
)]
async fn mempal_search(
&self,
Parameters(request): Parameters<SearchRequest>,
) -> std::result::Result<Json<SearchResponse>, ErrorData> {
let embedder = self.embedder_factory.build().await.map_err(|error| {
ErrorData::internal_error(format!("failed to build embedder: {error}"), None)
})?;
let query_vector = embedder
.embed(&[request.query.as_str()])
.await
.map_err(|error| ErrorData::internal_error(format!("embedding failed: {error}"), None))?
.into_iter()
.next()
.ok_or_else(|| ErrorData::internal_error("embedder returned no query vector", None))?;
let db = self.open_db()?;
let route = resolve_route(
&db,
&request.query,
request.wing.as_deref(),
request.room.as_deref(),
)
.map_err(|error| ErrorData::internal_error(format!("routing failed: {error}"), None))?;
let results = search_with_vector(
&db,
&request.query,
&query_vector,
route,
request.top_k.unwrap_or(10),
)
.map_err(|error| ErrorData::internal_error(format!("search failed: {error}"), None))?;
Ok(Json(SearchResponse {
results: results.into_iter().map(SearchResultDto::from).collect(),
}))
}
#[tool(
name = "mempal_ingest",
description = "Persist a decision, bug fix, or design insight to project memory. Call this when a decision is reached in conversation — include the rationale, not just the outcome. Wing is required; let mempal auto-route the room. Set dry_run=true to preview the drawer_id without writing."
)]
async fn mempal_ingest(
&self,
Parameters(request): Parameters<IngestRequest>,
) -> std::result::Result<Json<IngestResponse>, ErrorData> {
let room = request.room.as_deref();
let drawer_id = build_drawer_id(&request.wing, room, &request.content);
if request.dry_run.unwrap_or(false) {
return Ok(Json(IngestResponse {
drawer_id,
duplicate_warning: None,
}));
}
let embedder = self.embedder_factory.build().await.map_err(|error| {
ErrorData::internal_error(format!("failed to build embedder: {error}"), None)
})?;
let vector = embedder
.embed(&[request.content.as_str()])
.await
.map_err(|error| ErrorData::internal_error(format!("embedding failed: {error}"), None))?
.into_iter()
.next()
.ok_or_else(|| ErrorData::internal_error("embedder returned no vector", None))?;
let db = self.open_db()?;
let duplicate_warning = check_semantic_duplicate(&db, &vector, &request.content);
if !db.drawer_exists(&drawer_id).map_err(db_error)? {
let source_file = source_file_or_synthetic(&drawer_id, request.source.as_deref());
db.insert_drawer(&Drawer {
id: drawer_id.clone(),
content: request.content,
wing: request.wing,
room: request.room,
source_file: Some(source_file),
source_type: SourceType::Manual,
added_at: current_timestamp(),
chunk_index: Some(0),
importance: request.importance.unwrap_or(0),
})
.map_err(db_error)?;
db.insert_vector(&drawer_id, &vector).map_err(db_error)?;
}
Ok(Json(IngestResponse {
drawer_id,
duplicate_warning,
}))
}
#[tool(
name = "mempal_delete",
description = "Soft-delete a drawer by ID. The drawer is marked with a deleted_at timestamp and excluded from search results, but not physically removed. Use the CLI `mempal purge` to permanently remove soft-deleted drawers. Returns the drawer_id and whether it was found."
)]
async fn mempal_delete(
&self,
Parameters(request): Parameters<DeleteRequest>,
) -> std::result::Result<Json<DeleteResponse>, ErrorData> {
let db = self.open_db()?;
let deleted = db
.soft_delete_drawer(&request.drawer_id)
.map_err(db_error)?;
let message = if deleted {
format!("drawer {} soft-deleted", request.drawer_id)
} else {
format!("drawer {} not found or already deleted", request.drawer_id)
};
Ok(Json(DeleteResponse {
drawer_id: request.drawer_id,
deleted,
message,
}))
}
#[tool(
name = "mempal_taxonomy",
description = "List or edit wing/room taxonomy entries that drive query routing keywords."
)]
async fn mempal_taxonomy(
&self,
Parameters(request): Parameters<TaxonomyRequest>,
) -> std::result::Result<Json<TaxonomyResponse>, ErrorData> {
let db = self.open_db()?;
match request.action.as_str() {
"list" => {
let entries = db
.taxonomy_entries()
.map_err(db_error)?
.into_iter()
.map(TaxonomyEntryDto::from)
.collect();
Ok(Json(TaxonomyResponse {
action: "list".to_string(),
entries,
}))
}
"edit" => {
let wing = request
.wing
.ok_or_else(|| ErrorData::invalid_params("missing wing", None))?;
let room = request
.room
.ok_or_else(|| ErrorData::invalid_params("missing room", None))?;
let keywords = request
.keywords
.ok_or_else(|| ErrorData::invalid_params("missing keywords", None))?;
let entry = crate::core::types::TaxonomyEntry {
wing,
room,
display_name: None,
keywords,
};
db.upsert_taxonomy_entry(&entry).map_err(db_error)?;
Ok(Json(TaxonomyResponse {
action: "edit".to_string(),
entries: vec![TaxonomyEntryDto::from(entry)],
}))
}
action => Err(ErrorData::invalid_params(
format!("unsupported taxonomy action: {action}"),
None,
)),
}
}
#[tool(
name = "mempal_kg",
description = "Knowledge graph: add, query, or invalidate triples (subject-predicate-object). Use 'add' to record structured relationships between entities. Use 'query' to find relationships by subject, predicate, or object. Use 'invalidate' to mark a triple as no longer valid."
)]
async fn mempal_kg(
&self,
Parameters(request): Parameters<KgRequest>,
) -> std::result::Result<Json<KgResponse>, ErrorData> {
let db = self.open_db()?;
match request.action.as_str() {
"add" => {
let subject = request
.subject
.ok_or_else(|| ErrorData::invalid_params("missing subject", None))?;
let predicate = request
.predicate
.ok_or_else(|| ErrorData::invalid_params("missing predicate", None))?;
let object = request
.object
.ok_or_else(|| ErrorData::invalid_params("missing object", None))?;
let id = build_triple_id(&subject, &predicate, &object);
let triple = Triple {
id: id.clone(),
subject,
predicate,
object,
valid_from: Some(current_timestamp()),
valid_to: None,
confidence: 1.0,
source_drawer: request.source_drawer,
};
db.insert_triple(&triple).map_err(db_error)?;
Ok(Json(KgResponse {
action: "add".to_string(),
triples: vec![triple_to_dto(&triple)],
stats: None,
}))
}
"query" => {
let active_only = request.active_only.unwrap_or(true);
let triples = db
.query_triples(
request.subject.as_deref(),
request.predicate.as_deref(),
request.object.as_deref(),
active_only,
)
.map_err(db_error)?;
Ok(Json(KgResponse {
action: "query".to_string(),
triples: triples.iter().map(triple_to_dto).collect(),
stats: None,
}))
}
"invalidate" => {
let triple_id = request
.triple_id
.ok_or_else(|| ErrorData::invalid_params("missing triple_id", None))?;
let invalidated = db.invalidate_triple(&triple_id).map_err(db_error)?;
let message = if invalidated {
format!("triple {triple_id} invalidated")
} else {
format!("triple {triple_id} not found or already invalidated")
};
Ok(Json(KgResponse {
action: message,
triples: vec![],
stats: None,
}))
}
"timeline" => {
let entity = request.subject.ok_or_else(|| {
ErrorData::invalid_params("missing subject for timeline", None)
})?;
let triples = db.timeline_for_entity(&entity).map_err(db_error)?;
Ok(Json(KgResponse {
action: format!("timeline for {entity}"),
triples: triples.iter().map(triple_to_dto).collect(),
stats: None,
}))
}
"stats" => {
let stats = db.triple_stats().map_err(db_error)?;
Ok(Json(KgResponse {
action: "stats".to_string(),
triples: vec![],
stats: Some(KgStatsDto {
total: stats.total,
active: stats.active,
expired: stats.expired,
entities: stats.entities,
top_predicates: stats.top_predicates,
}),
}))
}
action => Err(ErrorData::invalid_params(
format!("unsupported kg action: {action}"),
None,
)),
}
}
#[tool(
name = "mempal_tunnels",
description = "Discover cross-wing tunnels: rooms that appear in multiple wings, enabling cross-domain knowledge discovery. Returns an empty list if only one wing exists."
)]
async fn mempal_tunnels(&self) -> std::result::Result<Json<TunnelsResponse>, ErrorData> {
let db = self.open_db()?;
let tunnels = db
.find_tunnels()
.map_err(db_error)?
.into_iter()
.map(|(room, wings)| TunnelDto { room, wings })
.collect();
Ok(Json(TunnelsResponse { tunnels }))
}
}
#[tool_handler(router = self.tool_router)]
impl ServerHandler for MempalMcpServer {
fn get_info(&self) -> ServerInfo {
ServerInfo::new(ServerCapabilities::builder().enable_tools().build())
.with_instructions(crate::core::protocol::MEMORY_PROTOCOL)
}
}
fn db_error(error: impl std::fmt::Display) -> ErrorData {
ErrorData::internal_error(format!("{error}"), None)
}
const DEDUP_THRESHOLD: f32 = 0.85;
fn check_semantic_duplicate(
db: &Database,
vector: &[f32],
_content: &str,
) -> Option<DuplicateWarning> {
use crate::core::types::RouteDecision;
let route = RouteDecision {
wing: None,
room: None,
confidence: 0.0,
reason: "dedup check".to_string(),
};
let results = crate::search::search_by_vector(db, vector, route, 1).ok()?;
let top = results.first()?;
if top.similarity >= DEDUP_THRESHOLD {
Some(DuplicateWarning {
similar_drawer_id: top.drawer_id.clone(),
similarity: top.similarity,
preview: top.content.chars().take(100).collect(),
})
} else {
None
}
}
fn triple_to_dto(triple: &Triple) -> TripleDto {
TripleDto {
id: triple.id.clone(),
subject: triple.subject.clone(),
predicate: triple.predicate.clone(),
object: triple.object.clone(),
valid_from: triple.valid_from.clone(),
valid_to: triple.valid_to.clone(),
confidence: triple.confidence,
source_drawer: triple.source_drawer.clone(),
}
}