tsafe-mcp 0.1.0

//! `TsafeMcpServer` — the rmcp 1.7 `ServerHandler` that exposes tsafe's tool
//! surface over a spec-compliant MCP 2025-06-18 stdio transport.
//!
//! This module replaces the hand-rolled JSON-RPC dispatcher that lived in the
//! former `serve_with` loop. rmcp handles the `initialize` / `tools/list` /
//! `tools/call` handshake; tool bodies continue to live in [`crate::tools`]
//! and are reached via a thin passthrough that preserves all existing
//! validation and audit logic unchanged.
//!
//! ## Why a passthrough instead of typed Params per tool
//!
//! Each tool module already validates its own raw `serde_json::Value` payload
//! and returns a typed `McpError` on failure. Re-deriving seven typed params
//! structs with `#[derive(schemars::JsonSchema)]` would duplicate that
//! validation surface and risk drift between the rmcp schema and the runtime
//! validator. The passthrough preserves the single source of truth at the
//! cost of presenting "any JSON object" as the schema in `tools/list`.
//!
//! ## schemars 1.x object schema requirement
//!
//! rmcp 1.7's `Parameters<T>` enforces that `T: JsonSchema` derives a schema
//! whose root has `"type": "object"` (MCP 2025-06-18 §6). Naked
//! `serde_json::Value` derives the "any" schema (no `type` field) and rmcp
//! panics at `tools/list` registration time. `serde_json::Map<String, Value>`
//! produces `{ "type": "object", "additionalProperties": true }` in schemars
//! 1.x, which satisfies the constraint while preserving the "any JSON object"
//! wire shape every tool module already validates.

use std::sync::Arc;

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

use crate::errors::{McpError, McpErrorKind};
use crate::session::Session;
use crate::tools;

/// MCP protocol version this build targets. Matches the wire constant used
/// previously by `rpc::MCP_PROTOCOL_VERSION`.
pub const MCP_PROTOCOL_VERSION: &str = "2025-06-18";

/// JSON object accepted as the argument map for every tsafe MCP tool.
///
/// See module-level doc for why `Map<String, Value>` is used instead of naked
/// `Value`.
type ToolArgs = Map<String, Value>;

/// JSON object returned as the response body for every tsafe MCP tool.
type ToolResultBody = Map<String, Value>;

/// tsafe MCP server. Holds an `Arc<Session>` so the struct is cheap to clone;
/// rmcp requires `ServerHandler: Clone` for its router.
#[derive(Clone)]
pub struct TsafeMcpServer {
    session: Arc<Session>,
}

impl TsafeMcpServer {
    pub fn new(session: Session) -> Self {
        Self {
            session: Arc::new(session),
        }
    }

    /// Route a tool call through [`crate::tools::dispatch`], which handles
    /// scope-widening rejection and per-tool dispatch.
    ///
    /// The raw `Value` from each tool module is wrapped in a JSON object so
    /// rmcp's output schema requirement (root must be `"type":"object"`) is
    /// satisfied. rmcp's `Json<T>` places the value into `structured_content`
    /// and also serialises it as a text `content` block — this double-serves
    /// both the structured and the text wire representations that MCP clients
    /// expect.
    fn dispatch(&self, name: &'static str, params: ToolArgs) -> Result<Json<ToolResultBody>, RmcpError> {
        let raw = Value::Object(params);
        match tools::dispatch(&self.session, name, raw) {
            Ok(payload) => {
                // Normalise to an object. Most tools already return an object;
                // tsafe_list_keys and tsafe_search_keys return arrays, so we
                // wrap them as {"result": [...]}.
                let body = match payload {
                    Value::Object(obj) => obj,
                    other => {
                        let mut map = serde_json::Map::new();
                        map.insert("result".to_string(), other);
                        map
                    }
                };
                Ok(Json(body))
            }
            Err(e) => Err(mcp_error_to_rmcp(e)),
        }
    }
}

fn mcp_error_to_rmcp(e: McpError) -> RmcpError {
    match e.kind {
        McpErrorKind::ParseError
        | McpErrorKind::InvalidRequest
        | McpErrorKind::InvalidParams => RmcpError::invalid_params(e.message, None),
        McpErrorKind::MethodNotFound => RmcpError::invalid_params(
            format!("method not found: {}", e.message),
            None,
        ),
        _ => {
            tracing::warn!(
                error_code = e.code,
                error = %e.message,
                "mcp: tool error mapped to internal_error"
            );
            RmcpError::internal_error(format!("[{}] {}", e.code, e.message), None)
        }
    }
}

#[tool_router]
impl TsafeMcpServer {
    #[tool(
        name = "tsafe_run",
        description = "Execute a command with explicitly-allowed vault keys injected as environment variables. Returns stdout/stderr/exit_code/duration_ms and the names of keys injected — never the secret values."
    )]
    async fn tsafe_run(
        &self,
        Parameters(p): Parameters<ToolArgs>,
    ) -> Result<Json<ToolResultBody>, RmcpError> {
        self.dispatch("tsafe_run", p)
    }

    #[tool(
        name = "tsafe_list_keys",
        description = "List vault key names visible to this server, filtered by scope. Optionally narrow by namespace prefix. Values are never returned."
    )]
    async fn tsafe_list_keys(
        &self,
        Parameters(p): Parameters<ToolArgs>,
    ) -> Result<Json<ToolResultBody>, RmcpError> {
        self.dispatch("tsafe_list_keys", p)
    }

    #[tool(
        name = "tsafe_search_keys",
        description = "Case-insensitive substring search across scope-filtered vault key names. Returns key names only."
    )]
    async fn tsafe_search_keys(
        &self,
        Parameters(p): Parameters<ToolArgs>,
    ) -> Result<Json<ToolResultBody>, RmcpError> {
        self.dispatch("tsafe_search_keys", p)
    }

    #[tool(
        name = "tsafe_has_key",
        description = "Check whether a vault key exists within this server's scope. Out-of-scope keys always return present=false regardless of vault contents."
    )]
    async fn tsafe_has_key(
        &self,
        Parameters(p): Parameters<ToolArgs>,
    ) -> Result<Json<ToolResultBody>, RmcpError> {
        self.dispatch("tsafe_has_key", p)
    }

    #[tool(
        name = "tsafe_audit_tail",
        description = "Return the most recent audit entries for the bound profile. Values are redacted; only id, timestamp, operation, key, status, and source are surfaced."
    )]
    async fn tsafe_audit_tail(
        &self,
        Parameters(p): Parameters<ToolArgs>,
    ) -> Result<Json<ToolResultBody>, RmcpError> {
        self.dispatch("tsafe_audit_tail", p)
    }

    #[tool(
        name = "tsafe_status",
        description = "Return agent/vault/profile status plus this server's configured scope. Matches ADR-029 schema version 1."
    )]
    async fn tsafe_status(
        &self,
        Parameters(p): Parameters<ToolArgs>,
    ) -> Result<Json<ToolResultBody>, RmcpError> {
        // tsafe_status is infallible on the tool side (returns Value directly).
        // Wrap it through dispatch to keep uniform scope-widening rejection.
        self.dispatch("tsafe_status", p)
    }

    #[tool(
        name = "tsafe_reveal",
        description = "Return the plaintext value of a single in-scope vault key. Gated by --allow-reveal; audited; biometric re-prompt when configured. The explicit escape hatch — every call appears in the profile audit log."
    )]
    async fn tsafe_reveal(
        &self,
        Parameters(p): Parameters<ToolArgs>,
    ) -> Result<Json<ToolResultBody>, RmcpError> {
        self.dispatch("tsafe_reveal", p)
    }
}

#[tool_handler]
impl ServerHandler for TsafeMcpServer {
    fn get_info(&self) -> rmcp::model::ServerInfo {
        let capabilities = ServerCapabilities::builder().enable_tools().build();
        let server_info = Implementation::new(
            "tsafe".to_string(),
            env!("CARGO_PKG_VERSION").to_string(),
        );
        let instructions = "tsafe-mcp: action-shaped secrets runtime. No secret values reach \
                            the LLM context by default. Use tsafe_run to execute commands with \
                            injected env vars. tsafe_reveal is only available when the server \
                            was started with --allow-reveal; every reveal call is audited.";
        InitializeResult::new(capabilities)
            .with_protocol_version(ProtocolVersion::V_2025_06_18)
            .with_server_info(server_info)
            .with_instructions(instructions)
    }

    /// Override list_tools to respect the session's allow_reveal flag.
    ///
    /// When `--allow-reveal` was NOT passed at startup, `tsafe_reveal` is
    /// excluded from the tool list per design §4.3. The dispatch path still
    /// rejects reveal calls via `tools::dispatch` when `allow_reveal` is false;
    /// this list override keeps the advertised surface consistent with the
    /// dispatch behavior.
    async fn list_tools(
        &self,
        _request: Option<PaginatedRequestParams>,
        _context: RequestContext<RoleServer>,
    ) -> Result<ListToolsResult, RmcpError> {
        let mut tools = Self::tool_router().list_all();
        if !self.session.allow_reveal {
            tools.retain(|t| t.name != "tsafe_reveal");
        }
        Ok(ListToolsResult {
            tools,
            meta: None,
            next_cursor: None,
        })
    }
}

/// Drive the stdio MCP loop until the peer closes stdin (EOF) or an OS
/// signal terminates the process.
///
/// All tracing/log output must be routed to stderr by the caller; stdout is
/// reserved for JSON-RPC frames.
///
/// # 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(session: Session) -> anyhow::Result<()> {
    tracing::info!("tsafe-mcp: rmcp 1.7 stdio server starting");
    let server = TsafeMcpServer::new(session);
    let service = server
        .serve(stdio())
        .await
        .map_err(|e| anyhow::anyhow!("serve_stdio init failed: {e}"))?;
    service
        .waiting()
        .await
        .map_err(|e| anyhow::anyhow!("serve_stdio loop failed: {e}"))?;
    tracing::info!("tsafe-mcp: rmcp stdio server shutdown (EOF)");
    Ok(())
}

/// Stable list of tool names this server registers. Tests pin against this.
#[allow(dead_code)]
pub const TOOL_NAMES: &[&str] = &[
    "tsafe_run",
    "tsafe_list_keys",
    "tsafe_search_keys",
    "tsafe_has_key",
    "tsafe_audit_tail",
    "tsafe_status",
    "tsafe_reveal",
];

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn tool_names_are_unique() {
        let mut seen = std::collections::HashSet::new();
        for name in TOOL_NAMES {
            assert!(seen.insert(name), "duplicate tool name: {name}");
        }
    }

    #[test]
    fn tool_names_count() {
        assert_eq!(TOOL_NAMES.len(), 7, "tsafe-mcp v1 ships 7 tools (6 default + reveal)");
    }
}