tsafe-mcp 2.0.1

//! `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, Tool,
};
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;
use crate::tools::schema;

/// MCP protocol version this build targets. This is the single wire-version
/// constant for the crate; rmcp's `ProtocolVersion::V_2025_06_18` is asserted
/// to agree with it in the unit tests below.
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 {
    let data = e.data.clone();
    match e.kind {
        McpErrorKind::ParseError | McpErrorKind::InvalidRequest | McpErrorKind::InvalidParams => {
            RmcpError::invalid_params(e.message, data)
        }
        McpErrorKind::MethodNotFound => {
            RmcpError::invalid_params(format!("method not found: {}", e.message), data)
        }
        _ => {
            tracing::warn!(
                error_code = e.code,
                error = %e.message,
                "mcp: tool error mapped to internal_error"
            );
            RmcpError::internal_error(format!("[{}] {}", e.code, e.message), data)
        }
    }
}

#[tool_router]
impl TsafeMcpServer {
    #[tool(
        name = "show_exec_plan",
        description = "Show the bounded `tsafe exec --contract ... --plan` invocation for this server's fixed profile, contract, and workdir."
    )]
    async fn show_exec_plan(
        &self,
        Parameters(p): Parameters<ToolArgs>,
    ) -> Result<Json<ToolResultBody>, RmcpError> {
        self.dispatch("show_exec_plan", p)
    }

    #[tool(
        name = "run_contract_command",
        description = "Run one command through the server's fixed `tsafe exec --contract` authority."
    )]
    async fn run_contract_command(
        &self,
        Parameters(p): Parameters<ToolArgs>,
    ) -> Result<Json<ToolResultBody>, RmcpError> {
        self.dispatch("run_contract_command", p)
    }

    #[tool(
        name = "tsafe_mcp_status",
        description = "Return safe status for the bound MCP server: profile, contract, workdir, agent/lock state, and compiled capabilities."
    )]
    async fn tsafe_mcp_status(
        &self,
        Parameters(p): Parameters<ToolArgs>,
    ) -> Result<Json<ToolResultBody>, RmcpError> {
        self.dispatch("tsafe_mcp_status", p)
    }

    #[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_suggest_keys",
        description = "Suggest a missing secret slot for the current repo. Writes metadata to .tsafe/tooling/keys.ini by default and never writes secret values to the vault."
    )]
    async fn tsafe_suggest_keys(
        &self,
        Parameters(p): Parameters<ToolArgs>,
    ) -> Result<Json<ToolResultBody>, RmcpError> {
        self.dispatch("tsafe_suggest_keys", p)
    }

    #[tool(
        name = "tsafe_inventory_check",
        description = "Validate the repo-local .tsafe/tooling/keys.ini secret-slot inventory without reading or returning secret values."
    )]
    async fn tsafe_inventory_check(
        &self,
        Parameters(p): Parameters<ToolArgs>,
    ) -> Result<Json<ToolResultBody>, RmcpError> {
        self.dispatch("tsafe_inventory_check", 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 = if self.session.is_bound_contract_mode() {
            "tsafe-mcp: bound contract command authority. This server is fixed to one \
             profile, one contract, and one workdir. Use show_exec_plan before \
             run_contract_command; use tsafe_mcp_status for safe operational metadata. \
             Secret values, vault browsing, profile switching, and request-time contract \
             or workdir switching are not available in bound mode."
        } else {
            "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> {
        if self.session.is_bound_contract_mode() {
            return Ok(ListToolsResult {
                tools: bound_contract_tools(self.session.as_ref())?,
                meta: None,
                next_cursor: None,
            });
        }

        let mut tools = Self::tool_router().list_all();
        tools.retain(|t| !BOUND_CONTRACT_TOOL_NAMES.contains(&t.name.as_ref()));
        if !self.session.allow_reveal {
            tools.retain(|t| t.name != "tsafe_reveal");
        }
        for tool in &mut tools {
            apply_real_schemas(tool);
        }
        Ok(ListToolsResult {
            tools,
            meta: None,
            next_cursor: None,
        })
    }
}

/// Replace the rmcp-derived generic `Map<String, Value>` passthrough schema on
/// a default-surface tool with the real per-tool schema derived from the
/// `schemars` parameter structs in [`crate::tools::schema`], and attach an
/// `outputSchema` for responses whose shape is well-defined.
///
/// Every default tool gets a real input schema here: the action tools from the
/// dedicated [`crate::tools::schema`] structs, and `tsafe_inventory_check` /
/// `tsafe_suggest_keys` from their own `serde` arg structs via the
/// `*_input_schema()` builders. The dispatch path still validates the raw
/// payload via each tool module, so these schemas are descriptive metadata for
/// `tools/list` only — they do not change runtime validation or its error
/// kinds.
fn apply_real_schemas(tool: &mut Tool) {
    let (input, output): (Map<String, Value>, Option<Map<String, Value>>) = match tool.name.as_ref()
    {
        "tsafe_run" => (
            schema::input_schema::<schema::RunParams>(),
            Some(schema::output_schema::<schema::RunResult>()),
        ),
        "tsafe_list_keys" => (
            schema::input_schema::<schema::ListKeysParams>(),
            Some(schema::string_array_result_schema(
                "Scope-filtered vault key names. Values are never returned.",
            )),
        ),
        "tsafe_search_keys" => (
            schema::input_schema::<schema::SearchKeysParams>(),
            Some(schema::string_array_result_schema(
                "Matching scope-filtered vault key names. Values are never returned.",
            )),
        ),
        "tsafe_has_key" => (
            schema::input_schema::<schema::HasKeyParams>(),
            Some(schema::output_schema::<schema::HasKeyResult>()),
        ),
        "tsafe_audit_tail" => (
            schema::input_schema::<schema::AuditTailParams>(),
            Some(schema::array_result_schema::<schema::AuditRow>(
                "Redacted audit rows; secret values are never present.",
            )),
        ),
        "tsafe_status" => (
            schema::input_schema::<schema::StatusParams>(),
            Some(schema::output_schema::<schema::StatusResult>()),
        ),
        "tsafe_inventory_check" => (
            schema_object(tools::tooling_inventory::inventory_check_input_schema()),
            None,
        ),
        "tsafe_suggest_keys" => (
            schema_object(tools::tooling_inventory::suggest_keys_input_schema()),
            None,
        ),
        "tsafe_reveal" => (schema::input_schema::<schema::HasKeyParams>(), None),
        // Any future tool keeps the rmcp-derived schema until added here.
        _ => return,
    };

    tool.input_schema = Arc::new(input);
    tool.output_schema = output.map(Arc::new);
}

/// Unwrap a `serde_json::Value` known to be an object into its map. Used for
/// the tooling-inventory schema builders that return `Value`.
fn schema_object(value: Value) -> Map<String, Value> {
    match value {
        Value::Object(map) => map,
        other => {
            tracing::error!(?other, "tool input schema builder did not return an object");
            Map::new()
        }
    }
}

/// 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] = &[
    "show_exec_plan",
    "run_contract_command",
    "tsafe_mcp_status",
    "tsafe_run",
    "tsafe_list_keys",
    "tsafe_search_keys",
    "tsafe_has_key",
    "tsafe_audit_tail",
    "tsafe_status",
    "tsafe_suggest_keys",
    "tsafe_inventory_check",
    "tsafe_reveal",
];

#[allow(dead_code)]
pub const BOUND_CONTRACT_TOOL_NAMES: &[&str] =
    &["show_exec_plan", "run_contract_command", "tsafe_mcp_status"];

fn bound_contract_tools(session: &Session) -> Result<Vec<Tool>, RmcpError> {
    let catalog = tools::list_tools(session);
    let tool_values = catalog
        .get("tools")
        .and_then(Value::as_array)
        .ok_or_else(|| RmcpError::internal_error("bound tools catalog is malformed", None))?;

    tool_values
        .iter()
        .cloned()
        .map(|tool| {
            serde_json::from_value::<Tool>(tool).map_err(|err| {
                RmcpError::internal_error(format!("bound tool schema is malformed: {err}"), None)
            })
        })
        .collect()
}

#[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(),
            12,
            "tsafe-mcp registers 3 bound tools plus 9 default tools"
        );
    }
}