cortex_mcp/

server.rs

//! `CortexServer` — the `rmcp` 1.7 `ServerHandler` that exposes Cortex's tool
//! surface over a spec-compliant MCP 2025-06-18 stdio transport.
//!
//! This module replaces the hand-rolled JSON-RPC dispatcher in
//! [`crate::serve`]. The legacy dispatcher treated the request `method`
//! field as a tool name (e.g. `"cortex_memory_note"`) and therefore failed
//! Claude Code's MCP client, which speaks the canonical
//! `initialize` / `tools/list` / `tools/call` handshake. The handshake itself
//! is now handled by `rmcp`; tool dispatch routes through the existing
//! [`ToolRegistry`] so every [`crate::ToolHandler`] implementation — and
//! every per-tool unit test — keeps working unchanged.
//!
//! ## Wire mapping
//!
//! Every `#[tool]` method below accepts `Parameters<serde_json::Value>` and
//! delegates to [`CortexServer::dispatch`], which looks up the static tool
//! name in the registry and returns either `Json<Value>` (success) or an
//! [`McpError`]. The mapping from [`ToolError`] to [`McpError`] is:
//!
//! | [`ToolError`]                    | [`McpError`] constructor      |
//! |----------------------------------|-------------------------------|
//! | `InvalidParams(msg)`             | `invalid_params(msg, None)`   |
//! | `PolicyRejected(msg)`            | `invalid_params(msg, None)`   |
//! | `SizeLimitExceeded(msg)`         | `invalid_params(msg, None)`   |
//! | `Internal(msg)`                  | `internal_error(msg, None)`   |
//!
//! `PolicyRejected` maps to `invalid_params` because every policy rejection
//! today is "the caller's input was refused by an authority gate" — there is
//! no separate rmcp variant for that case and the JSON-RPC `-32602` semantics
//! ("the parameters are invalid for this method as configured") are correct.
//!
//! ## Why a registry passthrough instead of typed Params per tool
//!
//! Cordance's `rmcp` integration uses one `JsonSchema`-derived params struct
//! per tool, which gives Claude Code's UI a rich schema for argument
//! prompting. Cortex's 18 tools all currently validate their own raw
//! `serde_json::Value` payload in `ToolHandler::call`. Re-deriving 18 typed
//! params structs would duplicate that validation surface and risk drift
//! between the rmcp schema and the runtime validator. The passthrough
//! preserves the existing single source of truth at the cost of presenting
//! "any JSON object" as the schema in `tools/list`. The handshake itself
//! still works, which is the bug this module was written to fix.

use std::sync::Arc;

use rmcp::handler::server::wrapper::{Json, Parameters};
use rmcp::model::{Implementation, InitializeResult, ProtocolVersion, ServerCapabilities};
use rmcp::transport::stdio;
use rmcp::{ErrorData as McpError, ServerHandler, ServiceExt, tool, tool_handler, tool_router};
use serde_json::{Map, Value};

use crate::tool_handler::ToolError;
use crate::tool_registry::ToolRegistry;

/// JSON object accepted as the argument map for every Cortex MCP tool.
///
/// rmcp 1.7's `Parameters<T>` enforces that `T: JsonSchema` derives a schema
/// whose root has `"type": "object"` (MCP 2025-06-18 §6 — tool argument schemas
/// MUST be objects). Naked `serde_json::Value` derives the "any" schema (no
/// `type` field) and rmcp panics at `tools/list` registration time. schemars
/// 1.x's impl of `JsonSchema for serde_json::Map<String, Value>` produces a
/// proper `{ "type": "object", "additionalProperties": true }` schema, which
/// satisfies the rmcp/MCP contract while preserving the "any JSON object"
/// wire shape every existing [`crate::ToolHandler`] already validates.
type ToolArgs = Map<String, Value>;

/// JSON object returned as the response body for every Cortex MCP tool.
///
/// Mirrors [`ToolArgs`] on the response side — the rmcp output-schema check
/// requires `"type": "object"` at the root, which `Map<String, Value>`
/// produces directly. Every existing Cortex tool already wraps its response
/// in `serde_json::json!({...})`, so the invariant holds at the call site;
/// [`CortexServer::dispatch`] errors out with a clear message if a future
/// tool ever returns a non-object value.
type ToolResultBody = Map<String, Value>;

/// Cortex MCP server.
///
/// Holds an `Arc<ToolRegistry>` so the struct is cheap to clone. `rmcp`
/// requires `ServerHandler: Clone` for its router.
#[derive(Clone)]
pub struct CortexServer {
    registry: Arc<ToolRegistry>,
}

impl std::fmt::Debug for CortexServer {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("CortexServer")
            .field("registry", &self.registry)
            .finish()
    }
}

impl CortexServer {
    /// Construct over a fully-wired registry. The registry must already
    /// contain every tool the server will expose; rmcp's `tools/list` reply
    /// is generated from the static `#[tool(...)]` attributes on this impl
    /// and not from the registry contents, so a tool advertised here but not
    /// registered will return `internal_error` at call time.
    #[must_use]
    pub fn new(registry: Arc<ToolRegistry>) -> Self {
        Self { registry }
    }

    /// Route a tool call through the underlying [`ToolRegistry`].
    ///
    /// `name` is a compile-time tool name supplied by each `#[tool]` shim
    /// below; `params` is the peer-supplied arguments object. The registry
    /// returns `None` only if the name was never registered (a wiring bug
    /// that surfaces as `internal_error`). A successful tool response is
    /// asserted to be a JSON object — every current Cortex tool returns
    /// `serde_json::json!({...})`; a non-object value indicates a programming
    /// error in the tool and is reported as `internal_error` rather than
    /// silently coerced.
    fn dispatch(
        &self,
        name: &'static str,
        params: ToolArgs,
    ) -> Result<Json<ToolResultBody>, McpError> {
        let raw_params = Value::Object(params);
        match self.registry.dispatch(name, raw_params) {
            Some(Ok(Value::Object(body))) => Ok(Json(body)),
            Some(Ok(other)) => {
                tracing::error!(
                    tool = name,
                    value_kind = value_kind(&other),
                    "mcp: tool returned non-object value (violates output schema invariant)"
                );
                Err(McpError::internal_error(
                    format!(
                        "tool '{name}' returned non-object value (kind: {})",
                        value_kind(&other)
                    ),
                    None,
                ))
            }
            Some(Err(err)) => Err(tool_error_to_mcp(name, err)),
            None => {
                tracing::error!(
                    tool = name,
                    "mcp: tool advertised by #[tool] but not registered in ToolRegistry"
                );
                Err(McpError::internal_error(
                    format!("tool '{name}' not registered"),
                    None,
                ))
            }
        }
    }
}

fn value_kind(v: &Value) -> &'static str {
    match v {
        Value::Null => "null",
        Value::Bool(_) => "bool",
        Value::Number(_) => "number",
        Value::String(_) => "string",
        Value::Array(_) => "array",
        Value::Object(_) => "object",
    }
}

fn tool_error_to_mcp(tool_name: &str, err: ToolError) -> McpError {
    match err {
        ToolError::InvalidParams(msg) => McpError::invalid_params(msg, None),
        ToolError::PolicyRejected(msg) => McpError::invalid_params(msg, None),
        ToolError::SizeLimitExceeded(msg) => McpError::invalid_params(msg, None),
        ToolError::Internal(msg) => {
            tracing::warn!(tool = tool_name, error = %msg, "mcp: tool internal error");
            McpError::internal_error(msg, None)
        }
    }
}

#[tool_router]
impl CortexServer {
    // ── Session tier (read-only, no confirmation token) ─────────────────

    #[tool(
        name = "cortex_search",
        description = "Find active memories matching the query. FTS5 by default; \
                       set `semantic: true` for Ollama-embedding similarity search. \
                       Returns top-K with relevance scores."
    )]
    async fn cortex_search(
        &self,
        Parameters(p): Parameters<ToolArgs>,
    ) -> Result<Json<ToolResultBody>, McpError> {
        self.dispatch("cortex_search", p)
    }

    #[tool(
        name = "cortex_context",
        description = "Build a context pack for the current session. Optionally include \
                       doctrine snippets and filter by tag, domain, or query."
    )]
    async fn cortex_context(
        &self,
        Parameters(p): Parameters<ToolArgs>,
    ) -> Result<Json<ToolResultBody>, McpError> {
        self.dispatch("cortex_context", p)
    }

    #[tool(
        name = "cortex_memory_health",
        description = "Return aggregate counts for active and quarantined memories: \
                       total, stale (>30 days old), unvalidated, and quarantined."
    )]
    async fn cortex_memory_health(
        &self,
        Parameters(p): Parameters<ToolArgs>,
    ) -> Result<Json<ToolResultBody>, McpError> {
        self.dispatch("cortex_memory_health", p)
    }

    #[tool(
        name = "cortex_config",
        description = "Return the active LLM and embedding backend configuration \
                       (Ollama / OpenAI-compat / Claude HTTP) loaded from cortex.toml."
    )]
    async fn cortex_config(
        &self,
        Parameters(p): Parameters<ToolArgs>,
    ) -> Result<Json<ToolResultBody>, McpError> {
        self.dispatch("cortex_config", p)
    }

    #[tool(
        name = "cortex_suggest",
        description = "Server-initiated memory suggestions for the current focus. \
                       Ranks by FTS5 match + salience; never mutates state."
    )]
    async fn cortex_suggest(
        &self,
        Parameters(p): Parameters<ToolArgs>,
    ) -> Result<Json<ToolResultBody>, McpError> {
        self.dispatch("cortex_suggest", p)
    }

    // ── Supervised tier (executes + logs; no confirmation token) ────────

    #[tool(
        name = "cortex_memory_list",
        description = "Browse active memories with optional tag/domain/status filters \
                       and a paging cursor. Read-only."
    )]
    async fn cortex_memory_list(
        &self,
        Parameters(p): Parameters<ToolArgs>,
    ) -> Result<Json<ToolResultBody>, McpError> {
        self.dispatch("cortex_memory_list", p)
    }

    #[tool(
        name = "cortex_memory_outcome",
        description = "Mark a specific memory as `helpful` or `not_helpful` for outcome \
                       tracking. Logs a structured outcome record (ADR 0020 §6)."
    )]
    async fn cortex_memory_outcome(
        &self,
        Parameters(p): Parameters<ToolArgs>,
    ) -> Result<Json<ToolResultBody>, McpError> {
        self.dispatch("cortex_memory_outcome", p)
    }

    #[tool(
        name = "cortex_decay_status",
        description = "Inspect the decay job queue: pending evictions and the next \
                       scheduled decay window."
    )]
    async fn cortex_decay_status(
        &self,
        Parameters(p): Parameters<ToolArgs>,
    ) -> Result<Json<ToolResultBody>, McpError> {
        self.dispatch("cortex_decay_status", p)
    }

    #[tool(
        name = "cortex_doctor",
        description = "Run health checks on the store, event log, and configured \
                       backends. Stores the result for `cortex doctor` to read back."
    )]
    async fn cortex_doctor(
        &self,
        Parameters(p): Parameters<ToolArgs>,
    ) -> Result<Json<ToolResultBody>, McpError> {
        self.dispatch("cortex_doctor", p)
    }

    #[tool(
        name = "cortex_audit_verify",
        description = "Verify the JSONL audit log's hash chain end-to-end. Returns \
                       pass/fail and the first divergence offset on failure."
    )]
    async fn cortex_audit_verify(
        &self,
        Parameters(p): Parameters<ToolArgs>,
    ) -> Result<Json<ToolResultBody>, McpError> {
        self.dispatch("cortex_audit_verify", p)
    }

    #[tool(
        name = "cortex_reflect",
        description = "Run a reflection pass over a session trace (optionally with a \
                       live LLM via `live_reflect: true`). Returns memory candidates."
    )]
    async fn cortex_reflect(
        &self,
        Parameters(p): Parameters<ToolArgs>,
    ) -> Result<Json<ToolResultBody>, McpError> {
        self.dispatch("cortex_reflect", p)
    }

    #[tool(
        name = "cortex_models_list",
        description = "List models available to the configured backends: pulled Ollama \
                       tags and the compile-time Claude allowlist."
    )]
    async fn cortex_models_list(
        &self,
        Parameters(p): Parameters<ToolArgs>,
    ) -> Result<Json<ToolResultBody>, McpError> {
        self.dispatch("cortex_models_list", p)
    }

    #[tool(
        name = "cortex_memory_embed",
        description = "Enrich pending memory rows with Ollama embeddings. Idempotent; \
                       `preview: true` reports the candidate set without writing."
    )]
    async fn cortex_memory_embed(
        &self,
        Parameters(p): Parameters<ToolArgs>,
    ) -> Result<Json<ToolResultBody>, McpError> {
        self.dispatch("cortex_memory_embed", p)
    }

    #[tool(
        name = "cortex_memory_note",
        description = "Store an operator-attested fact directly as an active memory. \
                       Bypasses the reflection pipeline. Required: `claim` (non-empty)."
    )]
    async fn cortex_memory_note(
        &self,
        Parameters(p): Parameters<ToolArgs>,
    ) -> Result<Json<ToolResultBody>, McpError> {
        self.dispatch("cortex_memory_note", p)
    }

    #[tool(
        name = "cortex_session_close",
        description = "Index the current session's events into pending memories. \
                       Use `live_reflect: true` for an LLM pass; otherwise heuristic only."
    )]
    async fn cortex_session_close(
        &self,
        Parameters(p): Parameters<ToolArgs>,
    ) -> Result<Json<ToolResultBody>, McpError> {
        self.dispatch("cortex_session_close", p)
    }

    // ── Confirmed tier (operator token from stderr) ─────────────────────

    #[tool(
        name = "cortex_memory_accept",
        description = "Promote a specific pending memory candidate to active. Requires \
                       the operator confirmation token printed to stderr (ADR 0047)."
    )]
    async fn cortex_memory_accept(
        &self,
        Parameters(p): Parameters<ToolArgs>,
    ) -> Result<Json<ToolResultBody>, McpError> {
        self.dispatch("cortex_memory_accept", p)
    }

    #[tool(
        name = "cortex_admit_axiom",
        description = "Admit a pinned-authority axiom into the ledger. Requires the \
                       operator confirmation token (ADR 0026 §4)."
    )]
    async fn cortex_admit_axiom(
        &self,
        Parameters(p): Parameters<ToolArgs>,
    ) -> Result<Json<ToolResultBody>, McpError> {
        self.dispatch("cortex_admit_axiom", p)
    }

    #[tool(
        name = "cortex_session_commit",
        description = "Activate the current session's pending_mcp_commit memories. \
                       Requires the operator confirmation token printed to stderr \
                       at server startup (ADR 0047 §3)."
    )]
    async fn cortex_session_commit(
        &self,
        Parameters(p): Parameters<ToolArgs>,
    ) -> Result<Json<ToolResultBody>, McpError> {
        self.dispatch("cortex_session_commit", p)
    }
}

#[tool_handler]
impl ServerHandler for CortexServer {
    fn get_info(&self) -> rmcp::model::ServerInfo {
        let capabilities = ServerCapabilities::builder().enable_tools().build();
        let server_info = Implementation::new(
            "cortex".to_string(),
            env!("CARGO_PKG_VERSION").to_string(),
        );
        let instructions = "Cortex MCP server. Memory mutations and session commits \
                            require an operator-issued confirmation token printed to \
                            stderr at startup (ADR 0047). Paste the token when prompted \
                            for `cortex_session_commit` or `cortex_memory_accept`. \
                            Sensitivity-gated context can be requested via \
                            `cortex_context` and `cortex_search`.";
        InitializeResult::new(capabilities)
            .with_protocol_version(ProtocolVersion::V_2025_06_18)
            .with_server_info(server_info)
            .with_instructions(instructions)
    }
}

/// Drive the stdio MCP loop until the peer closes stdin (EOF) or an OS
/// signal terminates the process.
///
/// This is the rmcp-based replacement for [`crate::serve::run_stdio_server`].
/// The caller is responsible for building a `tokio` runtime (rmcp's
/// `transport-io` feature uses the async stdin/stdout primitives).
///
/// # Errors
///
/// Returns an error when rmcp fails to bind the stdio transport or when the
/// loop exits abnormally. Clean EOF returns `Ok(())`.
pub async fn serve_stdio(server: CortexServer) -> Result<(), McpError> {
    tracing::info!("cortex mcp: rmcp 1.7 stdio server starting");
    let service = server.serve(stdio()).await.map_err(|e| {
        McpError::internal_error(format!("serve_stdio init failed: {e}"), None)
    })?;
    service.waiting().await.map_err(|e| {
        McpError::internal_error(format!("serve_stdio loop failed: {e}"), None)
    })?;
    tracing::info!("cortex mcp: rmcp stdio server shutdown (EOF)");
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::tool_handler::{GateId, ToolHandler};

    struct EchoTool;
    impl ToolHandler for EchoTool {
        fn name(&self) -> &'static str {
            "cortex_search"
        }
        fn gate_set(&self) -> &'static [GateId] {
            &[GateId::FtsRead]
        }
        fn call(&self, params: Value) -> Result<Value, ToolError> {
            Ok(params)
        }
    }

    fn server_with_echo() -> CortexServer {
        let mut registry = ToolRegistry::new();
        registry.register(Box::new(EchoTool));
        CortexServer::new(Arc::new(registry))
    }

    fn args(v: Value) -> ToolArgs {
        match v {
            Value::Object(fields) => fields,
            _ => panic!("test fixture must pass an object"),
        }
    }

    #[test]
    fn dispatch_returns_registry_value_on_success() {
        let server = server_with_echo();
        let result = server
            .dispatch("cortex_search", args(serde_json::json!({"q": "hello"})))
            .expect("registered tool dispatches");
        let Json(body) = result;
        assert_eq!(body.get("q"), Some(&Value::String("hello".into())));
    }

    #[test]
    fn dispatch_unregistered_tool_returns_internal_error() {
        // `rmcp::Json` deliberately does not implement `Debug` so the
        // `.expect_err(_)` shortcut isn't available; match the result by hand.
        let server = server_with_echo();
        let result = server.dispatch("cortex_missing", ToolArgs::new());
        match result {
            Ok(_) => panic!("unregistered tool must error"),
            Err(err) => assert!(
                err.message.contains("cortex_missing"),
                "error must name the missing tool: {}",
                err.message
            ),
        }
    }

    #[test]
    fn dispatch_non_object_tool_value_returns_internal_error() {
        // A tool that returns a non-object Value is a programming error;
        // `dispatch` must surface that as `internal_error` rather than
        // silently coercing — the output JSON Schema is `type: object`.
        struct NonObjectTool;
        impl ToolHandler for NonObjectTool {
            fn name(&self) -> &'static str {
                "cortex_search"
            }
            fn gate_set(&self) -> &'static [GateId] {
                &[GateId::FtsRead]
            }
            fn call(&self, _params: Value) -> Result<Value, ToolError> {
                Ok(Value::String("not an object".into()))
            }
        }
        let mut registry = ToolRegistry::new();
        registry.register(Box::new(NonObjectTool));
        let server = CortexServer::new(Arc::new(registry));
        let result = server.dispatch("cortex_search", ToolArgs::new());
        match result {
            Ok(_) => panic!("non-object tool value must error"),
            Err(err) => {
                assert_eq!(err.code, rmcp::model::ErrorCode::INTERNAL_ERROR);
                assert!(
                    err.message.contains("non-object"),
                    "error must say non-object: {}",
                    err.message
                );
            }
        }
    }

    #[test]
    fn tool_error_invalid_params_maps_to_invalid_params() {
        let err = tool_error_to_mcp("t", ToolError::InvalidParams("bad".into()));
        assert_eq!(err.code, rmcp::model::ErrorCode::INVALID_PARAMS);
    }

    #[test]
    fn tool_error_policy_rejected_maps_to_invalid_params() {
        let err = tool_error_to_mcp("t", ToolError::PolicyRejected("nope".into()));
        assert_eq!(err.code, rmcp::model::ErrorCode::INVALID_PARAMS);
    }

    #[test]
    fn tool_error_internal_maps_to_internal_error() {
        let err = tool_error_to_mcp("t", ToolError::Internal("kaboom".into()));
        assert_eq!(err.code, rmcp::model::ErrorCode::INTERNAL_ERROR);
    }
}