// Copyright (c) 2025, 2026 Julius ML
// MIT License
//! Post-Cortex MCP Service — consolidated 9-tool API.
//!
//! Tools:
//! 1. session — create/list/load/search/update_metadata/delete sessions
//! 2. update_conversation_context — single or bulk updates
//! 3. semantic_search — unified search with scope (structured payload)
//! 4. get_structured_summary — summary with optional sections
//! 5. query_conversation_context — flexible queries
//! 6. manage_workspace — workspace operations
//! 7. assemble_context — graph-aware retrieval (entity match + traversal + impact)
//! 8. manage_entity — delete entities / context updates from a session
//! 9. admin — daemon health, vectorization, checkpoints
//!
//! The `#[tool_router]`-decorated impl block must stay in a single file so the
//! rmcp macro can see every `#[tool]` method at once. Tool method bodies are
//! delegated to per-tool submodules to keep this file readable.
use crate::daemon::coerce::CoercionError;
use post_cortex_mcp::MCPToolResult;
use post_cortex_memory::ConversationMemorySystem;
use rmcp::{
Peer, RoleServer, ServerHandler,
handler::server::router::tool::ToolRouter,
handler::server::wrapper::Parameters,
model::{
CallToolRequestParams, CallToolResponse, CallToolResult, ContentBlock,
ErrorData as McpError, ListToolsResult, PaginatedRequestParams, *,
},
service::RequestContext,
tool, tool_router,
};
use std::sync::Arc;
use tracing::debug;
mod admin;
mod assembly;
mod entity;
mod query;
mod requests;
mod search;
mod session;
mod summary;
mod update_context;
mod workspace;
// Request types are public so external code (e.g. integration tests) can
// build them, mirroring the original monolith's surface.
pub use requests::{
AdminRequest, AssembleContextRequest, ContextUpdateItem, GetStructuredSummaryRequest,
ManageEntityRequest, ManageWorkspaceRequest, QueryConversationContextRequest,
SemanticSearchRequest, SessionRequest, UpdateConversationContextRequest,
};
/// HTTP header naming the caller as the author of its writes, for clients
/// whose MCP `clientInfo` does not say who they are.
pub const AUTHOR_HEADER: &str = "x-pcx-author";
/// `clientInfo.name` values that name an MCP SDK rather than the
/// application using it.
const SDK_CLIENT_NAMES: &[&str] = &["rmcp", "mcp"];
/// Author of a write that names none: the [`AUTHOR_HEADER`] of the HTTP
/// request, else the MCP client's name unless that only names its SDK.
fn default_author(
parts: Option<&axum::http::request::Parts>,
client_name: Option<&str>,
) -> Option<String> {
let header = parts
.and_then(|p| p.headers.get(AUTHOR_HEADER))
.and_then(|v| v.to_str().ok())
.map(str::trim)
.filter(|v| !v.is_empty());
let client = client_name
.map(str::trim)
.filter(|n| !n.is_empty() && !SDK_CLIENT_NAMES.contains(&n.to_ascii_lowercase().as_str()));
header.or(client).map(String::from)
}
/// Post-Cortex MCP Service
#[derive(Clone)]
pub struct PostCortexService {
memory_system: Arc<ConversationMemorySystem>,
tool_router: ToolRouter<PostCortexService>,
}
// =============================================================================
// Helper Functions (shared across tool submodules)
// =============================================================================
/// JSON Schema for a tool's input, in the shape the daemon has always
/// published: draft 2020-12 with optional fields as `"nullable": true`.
///
/// rmcp 3's `schema_for_type` dropped schemars' `AddNullable` transform and
/// emits `"type": ["string", "null"]` instead; generating the schema here
/// keeps `tools/list` byte-for-byte compatible for existing MCP clients.
fn tool_input_schema<T: schemars::JsonSchema>() -> Arc<JsonObject> {
let mut settings = schemars::generate::SchemaSettings::draft2020_12();
settings.transforms = vec![Box::new(schemars::transform::AddNullable::default())];
let schema = settings.into_generator().into_root_schema_for::<T>();
match serde_json::to_value(schema) {
Ok(serde_json::Value::Object(object)) => Arc::new(object),
_ => Arc::new(JsonObject::new()),
}
}
/// Emit an MCPToolResult through CallToolResult, preserving structured data.
///
/// When `result.data` is present, the structured JSON payload is appended as a
/// second `ContentBlock::text` item so MCP clients keep access to scores, snippets
/// and entity tags instead of the human-readable summary alone.
pub(crate) fn mcp_result_to_call_result(result: MCPToolResult) -> CallToolResult {
let mut contents = vec![ContentBlock::text(result.message)];
if let Some(data) = result.data {
contents.push(ContentBlock::text(format!(
"\n\nResults:\n{}",
serde_json::to_string(&data).unwrap_or_default()
)));
}
CallToolResult::success(contents)
}
/// Check that a session exists in storage; used by dry_run paths so we don't
/// approve an update for a session that never existed.
pub(crate) async fn check_session_exists_for_dry_run(
memory_system: &Arc<ConversationMemorySystem>,
session_id: uuid::Uuid,
) -> Result<(), McpError> {
match memory_system.storage_actor.load_session(session_id).await {
Ok(Some(_)) => Ok(()),
Ok(None) => Err(CoercionError::new(
"Session not found",
std::io::Error::new(std::io::ErrorKind::NotFound, "session does not exist"),
Some(serde_json::Value::String(session_id.to_string())),
)
.with_parameter_path("session_id".to_string())
.with_hint("Create a session first using the 'session' tool with action='create'")
.to_mcp_error()),
Err(e) => Err(McpError::internal_error(
format!("Failed to check session existence: {}", e),
None,
)),
}
}
/// Parse `(date_from, date_to)` into a UTC range. Both must be supplied; partial
/// pairs return an MCP error. `(None, None)` returns `Ok(None)`.
pub(crate) fn parse_date_range(
date_from: Option<String>,
date_to: Option<String>,
) -> Result<Option<(chrono::DateTime<chrono::Utc>, chrono::DateTime<chrono::Utc>)>, McpError> {
match (date_from, date_to) {
(Some(from_str), Some(to_str)) => {
let from = chrono::DateTime::parse_from_rfc3339(&from_str)
.map_err(|_| {
McpError::invalid_params(
format!(
"Invalid date_from format: '{}'. Expected RFC3339 format (e.g., 2024-01-01T00:00:00Z)",
from_str
),
Some(serde_json::Value::String("date_from".to_string())),
)
})?
.with_timezone(&chrono::Utc);
let to = chrono::DateTime::parse_from_rfc3339(&to_str)
.map_err(|_| {
McpError::invalid_params(
format!(
"Invalid date_to format: '{}'. Expected RFC3339 format (e.g., 2024-12-31T23:59:59Z)",
to_str
),
Some(serde_json::Value::String("date_to".to_string())),
)
})?
.with_timezone(&chrono::Utc);
Ok(Some((from, to)))
}
(Some(_), None) | (None, Some(_)) => Err(McpError::invalid_params(
"Both date_from and date_to must be provided together".to_string(),
Some(serde_json::Value::String("date_from,date_to".to_string())),
)),
(None, None) => Ok(None),
}
}
// =============================================================================
// Tool dispatch — every method here is a thin call into a per-tool submodule.
// =============================================================================
#[tool_router]
impl PostCortexService {
/// Create a new MCP service backed by the given memory system.
pub fn new(memory_system: Arc<ConversationMemorySystem>) -> Self {
debug!("Initializing Post-Cortex MCP Service (9 consolidated tools)");
Self {
memory_system,
tool_router: Self::tool_router(),
}
}
#[tool(
description = "Manage sessions. Actions: create, list, load, search, update_metadata, delete.",
input_schema = tool_input_schema::<SessionRequest>()
)]
async fn session(
&self,
params: Parameters<serde_json::Value>,
) -> Result<CallToolResult, McpError> {
session::handle(&self.memory_system, params).await
}
#[tool(
description = "Add context updates to conversation. Supports single update or bulk mode with updates array.",
input_schema = tool_input_schema::<UpdateConversationContextRequest>()
)]
async fn update_conversation_context(
&self,
params: Parameters<serde_json::Value>,
peer: Peer<RoleServer>,
extensions: Extensions,
) -> Result<CallToolResult, McpError> {
let info = peer.peer_info();
let author = default_author(
extensions.get::<axum::http::request::Parts>(),
info.as_ref().map(|info| info.client_info.name.as_str()),
);
update_context::handle(&self.memory_system, params, author).await
}
#[tool(
description = "Universal semantic search. Scope: session (requires scope_id), workspace (requires scope_id), or global (default).",
input_schema = tool_input_schema::<SemanticSearchRequest>()
)]
async fn semantic_search(
&self,
params: Parameters<serde_json::Value>,
) -> Result<CallToolResult, McpError> {
search::handle(&self.memory_system, params).await
}
#[tool(
description = "Get session summary. Use 'include' to specify sections: decisions, insights, entities, or all (default).",
input_schema = tool_input_schema::<GetStructuredSummaryRequest>()
)]
async fn get_structured_summary(
&self,
params: Parameters<serde_json::Value>,
) -> Result<CallToolResult, McpError> {
summary::handle(&self.memory_system, params).await
}
#[tool(
description = "Query session data. Types include: find_related_entities, get_entity_context, search_updates, entity_importance, entity_network, find_related_content, key_decisions, key_insights, session_statistics, structured_summary, decisions, open_questions, assemble_context.",
input_schema = tool_input_schema::<QueryConversationContextRequest>()
)]
async fn query_conversation_context(
&self,
params: Parameters<serde_json::Value>,
) -> Result<CallToolResult, McpError> {
query::handle(&self.memory_system, params).await
}
#[tool(
description = "Manage workspaces and their cards. Actions: create, list, get, delete, add_session, remove_session; get_card / update_card for the project card (workspace_id) or the user card (card: \"user\"): short markdown notes put in front of every session",
input_schema = tool_input_schema::<ManageWorkspaceRequest>()
)]
async fn manage_workspace(
&self,
params: Parameters<serde_json::Value>,
) -> Result<CallToolResult, McpError> {
workspace::handle(&self.memory_system, params).await
}
#[tool(
description = "Graph-aware context assembly: finds graph entities named in the query, expands them one hop through the entity graph, runs impact analysis, and ranks recent session updates by entity matches and recency (no embeddings). With sources=messages or all, also adds captured conversation turns matching the query (semantic search) as excerpts. For embedding-based similarity over records use semantic_search. Returns ranked items, entity neighbourhood, dependency impact, and an LLM-ready formatted block. Scope: session or workspace.",
input_schema = tool_input_schema::<AssembleContextRequest>()
)]
async fn assemble_context(
&self,
params: Parameters<serde_json::Value>,
) -> Result<CallToolResult, McpError> {
assembly::handle(&self.memory_system, params).await
}
#[tool(
description = "Manage entities and context entries in a session. Actions: delete (remove entity + cascade typed edges; requires entity_name), delete_update (remove a single context entry by entry_id).",
input_schema = tool_input_schema::<ManageEntityRequest>()
)]
async fn manage_entity(
&self,
params: Parameters<serde_json::Value>,
) -> Result<CallToolResult, McpError> {
entity::handle(&self.memory_system, params).await
}
#[tool(
description = "Daemon administration. Actions: health, vectorize_session (session_id), vectorize_stats, create_checkpoint (session_id).",
input_schema = tool_input_schema::<AdminRequest>()
)]
async fn admin(
&self,
params: Parameters<serde_json::Value>,
) -> Result<CallToolResult, McpError> {
admin::handle(&self.memory_system, params).await
}
}
// NOTE: We implement ServerHandler manually instead of using #[tool_handler] macro
// to strip $schema from tool input schemas for broader MCP client compatibility.
// Many clients (Cursor, Windsurf, Continue.dev) don't support JSON Schema draft/2020-12.
impl ServerHandler for PostCortexService {
fn get_info(&self) -> ServerConfig {
ServerConfig::new(ServerCapabilities::builder().enable_tools().build())
.with_protocol_version(ProtocolVersion::V_2024_11_05)
.with_server_info(Implementation::from_build_env())
.with_instructions(
"Post-Cortex: Intelligent conversation memory system with 9 consolidated MCP tools. \
Tools: session, update_conversation_context, semantic_search, get_structured_summary, \
query_conversation_context, manage_workspace, assemble_context, manage_entity, admin. \
All use shared RocksDB for centralized knowledge management.",
)
}
async fn initialize(
&self,
request: InitializeRequestParams,
context: RequestContext<RoleServer>,
) -> Result<InitializeResult, McpError> {
context.peer.set_peer_info(request);
if let Some(parts) = context.extensions.get::<axum::http::request::Parts>() {
debug!("Client initialized from {}", parts.uri);
}
Ok(self.get_info())
}
// Custom list_tools that strips $schema from input schemas for client compatibility
async fn list_tools(
&self,
_request: Option<PaginatedRequestParams>,
_context: RequestContext<RoleServer>,
) -> Result<ListToolsResult, McpError> {
let tools = self.tool_router.list_all();
// Strip $schema from each tool's input_schema for broader client compatibility
// Clone and modify since input_schema is Arc<JsonObject>
let tools: Vec<Tool> = tools
.into_iter()
.map(|mut tool| {
let mut schema = (*tool.input_schema).clone();
schema.remove("$schema");
// Also strip from $defs if present
if let Some(defs) = schema.get_mut("$defs")
&& let Some(defs_obj) = defs.as_object_mut()
{
for (_, def) in defs_obj.iter_mut() {
if let Some(def_obj) = def.as_object_mut() {
def_obj.remove("$schema");
}
}
}
tool.input_schema = std::sync::Arc::new(schema);
tool
})
.collect();
Ok(ListToolsResult::with_all_items(tools))
}
// Route tool calls to the tool router
async fn call_tool(
&self,
request: CallToolRequestParams,
context: RequestContext<RoleServer>,
) -> Result<CallToolResponse, McpError> {
let tcc = rmcp::handler::server::tool::ToolCallContext::new(self, request, context);
self.tool_router.call(tcc).await
}
}
#[cfg(test)]
mod schema_tests {
use super::*;
#[test]
fn optional_fields_stay_nullable() {
let schema = serde_json::Value::Object((*tool_input_schema::<SessionRequest>()).clone());
let text = schema.to_string();
assert!(text.contains("\"nullable\":true"), "{text}");
assert!(
!text.contains("\"null\""),
"no type arrays with null: {text}"
);
assert_eq!(schema["type"], "object");
}
}
#[cfg(test)]
mod tests {
use super::default_author;
fn parts(author: Option<&str>) -> axum::http::request::Parts {
let mut request = axum::http::Request::builder().uri("/mcp");
if let Some(author) = author {
request = request.header("X-PCX-Author", author);
}
request.body(()).unwrap().into_parts().0
}
/// The header wins; a client name that only names its SDK is no author.
#[test]
fn default_author_prefers_the_header_and_skips_sdk_names() {
let with = parts(Some("claude-code"));
let without = parts(None);
assert_eq!(
default_author(Some(&with), Some("rmcp")).as_deref(),
Some("claude-code")
);
assert_eq!(
default_author(Some(&with), Some("Zed")).as_deref(),
Some("claude-code")
);
assert_eq!(
default_author(Some(&without), Some("Zed")).as_deref(),
Some("Zed")
);
assert_eq!(default_author(Some(&without), Some("rmcp")), None);
assert_eq!(default_author(None, Some("MCP")), None);
assert_eq!(default_author(Some(&parts(Some(" "))), None), None);
}
}