tuitbot_server/routes/

mcp.rs

//! MCP governance and telemetry endpoints.

use std::sync::Arc;
use std::time::{SystemTime, UNIX_EPOCH};

use axum::extract::{Path, Query, State};
use axum::Json;
use serde::Deserialize;
use serde_json::{json, Value};
use tuitbot_core::config::Config;
use tuitbot_core::mcp_policy::templates;
use tuitbot_core::mcp_policy::types::PolicyTemplateName;
use tuitbot_core::storage::{mcp_telemetry, rate_limits};

use crate::error::ApiError;
use crate::state::AppState;

// ---------------------------------------------------------------------------
// Query types
// ---------------------------------------------------------------------------

#[derive(Deserialize)]
pub struct TimeWindowQuery {
    /// Lookback window in hours (default: 24).
    #[serde(default = "default_hours")]
    pub hours: u32,
}

fn default_hours() -> u32 {
    24
}

#[derive(Deserialize)]
pub struct RecentQuery {
    /// Number of recent entries to return (default: 50).
    #[serde(default = "default_limit")]
    pub limit: u32,
}

fn default_limit() -> u32 {
    50
}

// ---------------------------------------------------------------------------
// Policy endpoints
// ---------------------------------------------------------------------------

/// `GET /api/mcp/policy` — current MCP policy config + rate limit usage + v2 fields.
pub async fn get_policy(State(state): State<Arc<AppState>>) -> Result<Json<Value>, ApiError> {
    let config = read_config(&state)?;

    let rate_limit_info = match rate_limits::get_all_rate_limits(&state.db).await {
        Ok(limits) => {
            let mcp = limits.iter().find(|l| l.action_type == "mcp_mutation");
            match mcp {
                Some(rl) => json!({
                    "used": rl.request_count,
                    "max": rl.max_requests,
                    "period_seconds": rl.period_seconds,
                    "period_start": rl.period_start,
                }),
                None => json!({ "used": 0, "max": config.mcp_policy.max_mutations_per_hour }),
            }
        }
        Err(_) => json!({ "used": 0, "max": config.mcp_policy.max_mutations_per_hour }),
    };

    Ok(Json(json!({
        "enforce_for_mutations": config.mcp_policy.enforce_for_mutations,
        "require_approval_for": config.mcp_policy.require_approval_for,
        "blocked_tools": config.mcp_policy.blocked_tools,
        "dry_run_mutations": config.mcp_policy.dry_run_mutations,
        "max_mutations_per_hour": config.mcp_policy.max_mutations_per_hour,
        "mode": format!("{}", config.mode),
        "rate_limit": rate_limit_info,
        "template": config.mcp_policy.template,
        "rules": config.mcp_policy.rules,
        "rate_limits": config.mcp_policy.rate_limits,
    })))
}

/// `PATCH /api/mcp/policy` — update MCP policy config fields.
///
/// Accepts partial JSON with `mcp_policy` fields and merges into config.
pub async fn patch_policy(
    State(state): State<Arc<AppState>>,
    Json(patch): Json<Value>,
) -> Result<Json<Value>, ApiError> {
    if !patch.is_object() {
        return Err(ApiError::BadRequest(
            "request body must be a JSON object".to_string(),
        ));
    }

    // Wrap the patch under `mcp_policy` key for the settings merge.
    let wrapped = json!({ "mcp_policy": patch });

    let contents = std::fs::read_to_string(&state.config_path).map_err(|e| {
        ApiError::BadRequest(format!(
            "could not read config file {}: {e}",
            state.config_path.display()
        ))
    })?;

    let mut toml_value: toml::Value = contents.parse().map_err(|e: toml::de::Error| {
        ApiError::BadRequest(format!("failed to parse existing config: {e}"))
    })?;

    let patch_toml = json_to_toml(&wrapped)
        .map_err(|e| ApiError::BadRequest(format!("patch contains invalid values: {e}")))?;

    merge_toml(&mut toml_value, &patch_toml);

    let merged_str = toml::to_string_pretty(&toml_value)
        .map_err(|e| ApiError::BadRequest(format!("failed to serialize merged config: {e}")))?;

    let config: Config = toml::from_str(&merged_str)
        .map_err(|e| ApiError::BadRequest(format!("merged config is invalid: {e}")))?;

    std::fs::write(&state.config_path, &merged_str).map_err(|e| {
        ApiError::BadRequest(format!(
            "could not write config file {}: {e}",
            state.config_path.display()
        ))
    })?;

    Ok(Json(json!({
        "enforce_for_mutations": config.mcp_policy.enforce_for_mutations,
        "require_approval_for": config.mcp_policy.require_approval_for,
        "blocked_tools": config.mcp_policy.blocked_tools,
        "dry_run_mutations": config.mcp_policy.dry_run_mutations,
        "max_mutations_per_hour": config.mcp_policy.max_mutations_per_hour,
        "template": config.mcp_policy.template,
        "rules": config.mcp_policy.rules,
        "rate_limits": config.mcp_policy.rate_limits,
    })))
}

// ---------------------------------------------------------------------------
// Template endpoints
// ---------------------------------------------------------------------------

/// `GET /api/mcp/policy/templates` — list available policy templates.
pub async fn list_templates() -> Json<Value> {
    let templates = templates::list_templates();
    Json(json!(templates))
}

/// `POST /api/mcp/policy/templates/{name}` — apply a template.
pub async fn apply_template(
    State(state): State<Arc<AppState>>,
    Path(name): Path<String>,
) -> Result<Json<Value>, ApiError> {
    let template_name: PolicyTemplateName =
        name.parse().map_err(|e: String| ApiError::BadRequest(e))?;

    let template = templates::get_template(&template_name);

    // Build a patch that sets the template and its rules/rate_limits
    let patch = json!({
        "template": template_name,
        "rules": template.rules,
        "rate_limits": template.rate_limits,
    });

    // Wrap under mcp_policy and merge into config
    let wrapped = json!({ "mcp_policy": patch });

    let contents = std::fs::read_to_string(&state.config_path).map_err(|e| {
        ApiError::BadRequest(format!(
            "could not read config file {}: {e}",
            state.config_path.display()
        ))
    })?;

    let mut toml_value: toml::Value = contents.parse().map_err(|e: toml::de::Error| {
        ApiError::BadRequest(format!("failed to parse existing config: {e}"))
    })?;

    let patch_toml = json_to_toml(&wrapped)
        .map_err(|e| ApiError::BadRequest(format!("patch contains invalid values: {e}")))?;

    merge_toml(&mut toml_value, &patch_toml);

    let merged_str = toml::to_string_pretty(&toml_value)
        .map_err(|e| ApiError::BadRequest(format!("failed to serialize merged config: {e}")))?;

    let config: Config = toml::from_str(&merged_str)
        .map_err(|e| ApiError::BadRequest(format!("merged config is invalid: {e}")))?;

    std::fs::write(&state.config_path, &merged_str).map_err(|e| {
        ApiError::BadRequest(format!(
            "could not write config file {}: {e}",
            state.config_path.display()
        ))
    })?;

    // Initialize rate limit rows for the new template limits
    if let Err(e) =
        rate_limits::init_policy_rate_limits(&state.db, &config.mcp_policy.rate_limits).await
    {
        tracing::warn!("Failed to initialize policy rate limits: {e}");
    }

    Ok(Json(json!({
        "applied_template": template_name,
        "description": template.description,
        "rules_count": config.mcp_policy.rules.len(),
        "rate_limits_count": config.mcp_policy.rate_limits.len(),
    })))
}

// ---------------------------------------------------------------------------
// Telemetry endpoints
// ---------------------------------------------------------------------------

/// `GET /api/mcp/telemetry/summary` — aggregate stats over a time window.
pub async fn telemetry_summary(
    State(state): State<Arc<AppState>>,
    Query(params): Query<TimeWindowQuery>,
) -> Result<Json<Value>, ApiError> {
    let since = since_timestamp(params.hours);
    let summary = mcp_telemetry::get_summary(&state.db, &since).await?;
    // `serde_json::to_value` can only fail if T contains a non-string map key or
    // an f64 NaN/infinity; the telemetry summary struct contains neither.
    let value = serde_json::to_value(summary)
        .map_err(|e| ApiError::Internal(format!("serialization error: {e}")))?;
    Ok(Json(value))
}

/// `GET /api/mcp/telemetry/metrics` — per-tool metrics over a time window.
pub async fn telemetry_metrics(
    State(state): State<Arc<AppState>>,
    Query(params): Query<TimeWindowQuery>,
) -> Result<Json<Value>, ApiError> {
    let since = since_timestamp(params.hours);
    let metrics = mcp_telemetry::get_metrics_since(&state.db, &since).await?;
    Ok(Json(json!(metrics)))
}

/// `GET /api/mcp/telemetry/errors` — error breakdown over a time window.
pub async fn telemetry_errors(
    State(state): State<Arc<AppState>>,
    Query(params): Query<TimeWindowQuery>,
) -> Result<Json<Value>, ApiError> {
    let since = since_timestamp(params.hours);
    let errors = mcp_telemetry::get_error_breakdown(&state.db, &since).await?;
    Ok(Json(json!(errors)))
}

/// `GET /api/mcp/telemetry/recent` — recent tool executions.
pub async fn telemetry_recent(
    State(state): State<Arc<AppState>>,
    Query(params): Query<RecentQuery>,
) -> Result<Json<Value>, ApiError> {
    let entries = mcp_telemetry::get_recent_entries(&state.db, params.limit).await?;
    Ok(Json(json!(entries)))
}

// ---------------------------------------------------------------------------
// Helpers
// ---------------------------------------------------------------------------

fn read_config(state: &AppState) -> Result<Config, ApiError> {
    let contents = std::fs::read_to_string(&state.config_path).map_err(|e| {
        ApiError::BadRequest(format!(
            "could not read config file {}: {e}",
            state.config_path.display()
        ))
    })?;
    let config: Config = toml::from_str(&contents)
        .map_err(|e| ApiError::BadRequest(format!("failed to parse config: {e}")))?;
    Ok(config)
}

fn since_timestamp(hours: u32) -> String {
    let now = SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .unwrap_or_default()
        .as_secs();
    let since_epoch = now.saturating_sub(u64::from(hours) * 3600);

    // Convert epoch seconds to ISO-8601 UTC (YYYY-MM-DDTHH:MM:SSZ).
    let secs = since_epoch as i64;
    let days = secs.div_euclid(86400);
    let day_secs = secs.rem_euclid(86400);
    let h = day_secs / 3600;
    let m = (day_secs % 3600) / 60;
    let s = day_secs % 60;

    // Days since epoch → date using the civil-from-days algorithm.
    let z = days + 719468;
    let era = z.div_euclid(146097);
    let doe = z.rem_euclid(146097);
    let yoe = (doe - doe / 1460 + doe / 36524 - doe / 146096) / 365;
    let y = yoe + era * 400;
    let doy = doe - (365 * yoe + yoe / 4 - yoe / 100);
    let mp = (5 * doy + 2) / 153;
    let d = doy - (153 * mp + 2) / 5 + 1;
    let month = if mp < 10 { mp + 3 } else { mp - 9 };
    let year = if month <= 2 { y + 1 } else { y };

    format!("{year:04}-{month:02}-{d:02}T{h:02}:{m:02}:{s:02}Z")
}

/// Recursively merge `patch` into `base`.
fn merge_toml(base: &mut toml::Value, patch: &toml::Value) {
    match (base, patch) {
        (toml::Value::Table(base_table), toml::Value::Table(patch_table)) => {
            for (key, patch_val) in patch_table {
                if let Some(base_val) = base_table.get_mut(key) {
                    merge_toml(base_val, patch_val);
                } else {
                    base_table.insert(key.clone(), patch_val.clone());
                }
            }
        }
        (base, _) => {
            *base = patch.clone();
        }
    }
}

/// Convert JSON to TOML, skipping nulls in objects.
fn json_to_toml(json: &serde_json::Value) -> Result<toml::Value, String> {
    match json {
        serde_json::Value::Object(map) => {
            let mut table = toml::map::Map::new();
            for (key, val) in map {
                if val.is_null() {
                    continue;
                }
                table.insert(key.clone(), json_to_toml(val)?);
            }
            Ok(toml::Value::Table(table))
        }
        serde_json::Value::Array(arr) => {
            let values: Result<Vec<_>, _> = arr.iter().map(json_to_toml).collect();
            Ok(toml::Value::Array(values?))
        }
        serde_json::Value::String(s) => Ok(toml::Value::String(s.clone())),
        serde_json::Value::Number(n) => {
            if let Some(i) = n.as_i64() {
                Ok(toml::Value::Integer(i))
            } else if let Some(f) = n.as_f64() {
                Ok(toml::Value::Float(f))
            } else {
                Err(format!("unsupported number: {n}"))
            }
        }
        serde_json::Value::Bool(b) => Ok(toml::Value::Boolean(*b)),
        serde_json::Value::Null => Err("null values are not supported in TOML arrays".to_string()),
    }
}

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

    #[test]
    fn since_timestamp_is_valid_utc() {
        let ts = since_timestamp(24);
        assert!(ts.ends_with('Z'));
        assert!(ts.contains('T'));
    }

    #[test]
    fn since_timestamp_zero_hours_is_now() {
        let ts = since_timestamp(0);
        assert!(ts.ends_with('Z'));
        // Should be close to current time
        assert!(ts.starts_with("20"));
    }

    #[test]
    fn since_timestamp_large_hours() {
        let ts = since_timestamp(8760); // 1 year
        assert!(ts.ends_with('Z'));
        assert!(ts.contains('T'));
    }

    #[test]
    fn since_timestamp_format_correct() {
        let ts = since_timestamp(48);
        // Should match YYYY-MM-DDTHH:MM:SSZ
        assert_eq!(ts.len(), 20);
        assert_eq!(&ts[4..5], "-");
        assert_eq!(&ts[7..8], "-");
        assert_eq!(&ts[10..11], "T");
        assert_eq!(&ts[13..14], ":");
        assert_eq!(&ts[16..17], ":");
    }

    // --- merge_toml ---

    #[test]
    fn merge_toml_adds_new_keys() {
        let mut base: toml::Value = "key1 = \"value1\"".parse().unwrap();
        let patch: toml::Value = "key2 = \"value2\"".parse().unwrap();
        merge_toml(&mut base, &patch);
        assert_eq!(
            base.as_table().unwrap().get("key2").unwrap().as_str(),
            Some("value2")
        );
    }

    #[test]
    fn merge_toml_overwrites_existing() {
        let mut base: toml::Value = "key = \"old\"".parse().unwrap();
        let patch: toml::Value = "key = \"new\"".parse().unwrap();
        merge_toml(&mut base, &patch);
        assert_eq!(
            base.as_table().unwrap().get("key").unwrap().as_str(),
            Some("new")
        );
    }

    #[test]
    fn merge_toml_deep_merge() {
        let mut base: toml::Value = "[section]\na = 1".parse().unwrap();
        let patch: toml::Value = "[section]\nb = 2".parse().unwrap();
        merge_toml(&mut base, &patch);
        let section = base.as_table().unwrap().get("section").unwrap();
        assert_eq!(section.get("a").unwrap().as_integer(), Some(1));
        assert_eq!(section.get("b").unwrap().as_integer(), Some(2));
    }

    // --- json_to_toml ---

    #[test]
    fn json_to_toml_string() {
        let json = serde_json::json!("hello");
        let toml = json_to_toml(&json).unwrap();
        assert_eq!(toml.as_str(), Some("hello"));
    }

    #[test]
    fn json_to_toml_integer() {
        let json = serde_json::json!(42);
        let toml = json_to_toml(&json).unwrap();
        assert_eq!(toml.as_integer(), Some(42));
    }

    #[test]
    fn json_to_toml_float() {
        let json = serde_json::json!(3.14);
        let toml = json_to_toml(&json).unwrap();
        assert!((toml.as_float().unwrap() - 3.14).abs() < 0.001);
    }

    #[test]
    fn json_to_toml_boolean() {
        let json = serde_json::json!(true);
        let toml = json_to_toml(&json).unwrap();
        assert_eq!(toml.as_bool(), Some(true));
    }

    #[test]
    fn json_to_toml_array() {
        let json = serde_json::json!([1, 2, 3]);
        let toml = json_to_toml(&json).unwrap();
        assert_eq!(toml.as_array().unwrap().len(), 3);
    }

    #[test]
    fn json_to_toml_object() {
        let json = serde_json::json!({"key": "value"});
        let toml = json_to_toml(&json).unwrap();
        assert!(toml.as_table().is_some());
    }

    #[test]
    fn json_to_toml_skips_null_in_objects() {
        let json = serde_json::json!({"key": "value", "null_key": null});
        let toml = json_to_toml(&json).unwrap();
        let table = toml.as_table().unwrap();
        assert!(table.contains_key("key"));
        assert!(!table.contains_key("null_key"));
    }

    #[test]
    fn json_to_toml_null_in_array_errors() {
        let json = serde_json::json!([null]);
        assert!(json_to_toml(&json).is_err());
    }

    // --- TimeWindowQuery defaults ---

    #[test]
    fn time_window_query_defaults() {
        let json = "{}";
        let q: TimeWindowQuery = serde_json::from_str(json).expect("deser");
        assert_eq!(q.hours, 24);
    }

    // --- RecentQuery defaults ---

    #[test]
    fn recent_query_defaults() {
        let json = "{}";
        let q: RecentQuery = serde_json::from_str(json).expect("deser");
        assert_eq!(q.limit, 50);
    }
}