agentis_pay/mcp/

handler.rs

use std::collections::HashMap;
use std::sync::{Arc, Mutex as StdMutex};
use std::time::{Duration, Instant};

use rmcp::handler::server::ServerHandler;
use rmcp::handler::server::tool::ToolRouter;
use rmcp::handler::server::wrapper::Parameters;
use rmcp::model::{
    CallToolResult, Content, Implementation, InitializeRequestParams, InitializeResult,
    ServerCapabilities,
};
use rmcp::service::RequestContext;
use rmcp::{ErrorData as McpError, RoleServer, tool, tool_handler, tool_router};
use serde_json::json;
use tokio::sync::Mutex;

use crate::CommandName;
use crate::display::OutputEnvelope;
use crate::mcp::backend::{
    BrcodeDecodeInput, HistoryInput, HistoryQuery, HistoryRecord, PayInput, PixBackend,
};
use crate::pix_key::normalize_pix_key;

const MAX_BRCODE_LENGTH: usize = 4_096;
const PAY_RATE_LIMIT: RateLimitRule = RateLimitRule {
    limit: 5,
    window: Duration::from_secs(60),
};
const BRCODE_RATE_LIMIT: RateLimitRule = RateLimitRule {
    limit: 20,
    window: Duration::from_secs(60),
};

#[derive(Clone)]
pub struct AgentisPayServer {
    backend: Arc<Mutex<Box<dyn PixBackend>>>,
    rate_limiter: RateLimiter,
    tool_router: ToolRouter<Self>,
}

#[derive(Clone, Default)]
struct RateLimiter {
    buckets: Arc<StdMutex<HashMap<&'static str, FixedWindow>>>,
}

#[derive(Clone, Copy)]
struct RateLimitRule {
    limit: u32,
    window: Duration,
}

struct FixedWindow {
    started_at: Instant,
    used: u32,
}

impl RateLimiter {
    fn check(&self, bucket: &'static str, rule: RateLimitRule) -> Result<(), McpError> {
        let now = Instant::now();
        let mut buckets = self
            .buckets
            .lock()
            .map_err(|_| McpError::internal_error("rate limiter unavailable", None))?;
        let entry = buckets.entry(bucket).or_insert(FixedWindow {
            started_at: now,
            used: 0,
        });

        let elapsed = now.saturating_duration_since(entry.started_at);
        if elapsed >= rule.window {
            entry.started_at = now;
            entry.used = 0;
        }

        if entry.used >= rule.limit {
            let retry_after_seconds = rule
                .window
                .checked_sub(elapsed)
                .unwrap_or_default()
                .as_secs()
                .max(1);
            return Err(McpError::invalid_request(
                format!("rate limit exceeded for {bucket}; try again later"),
                Some(json!({
                    "bucket": bucket,
                    "retry_after_seconds": retry_after_seconds,
                })),
            ));
        }

        entry.used += 1;
        Ok(())
    }
}

fn tool_result(command: CommandName, data: serde_json::Value) -> CallToolResult {
    let envelope = OutputEnvelope { command, data };
    let text =
        serde_json::to_string(&envelope).unwrap_or_else(|e| format!("{{\"error\":\"{e}\"}}"));
    CallToolResult::success(vec![Content::text(text)])
}

fn backend_error(operation: &'static str, error: anyhow::Error) -> McpError {
    let raw = error.to_string();
    if let Some(message) = safe_backend_message(raw.as_str()) {
        return McpError::invalid_request(message, None);
    }

    McpError::internal_error(format!("{operation} failed; please retry later"), None)
}

fn safe_backend_message(raw: &str) -> Option<&'static str> {
    match raw {
        "transaction not found" => Some("transaction not found"),
        "recipient key was not found" => Some("recipient key was not found"),
        "recipient key lookup failed" => Some("recipient key lookup failed"),
        "recipient key lookup was rate-limited" => Some("recipient key lookup was rate-limited"),
        "recipient key is flagged as fraudulent" => Some("recipient key is flagged as fraudulent"),
        "recipient key response was empty" => Some("recipient key response was empty"),
        "No active session. Run `agentis-pay login` first." => {
            Some("no active Agentis Pay session; run `agentis-pay login` again")
        }
        _ if raw.starts_with("session expired") => {
            Some("stored Agentis Pay session expired; run `agentis-pay login` again")
        }
        _ => None,
    }
}

fn invalid_params(error: anyhow::Error) -> McpError {
    McpError::invalid_params(error.to_string(), None)
}

fn validate_brcode_input(raw: &str) -> Result<&str, McpError> {
    let value = raw.trim();
    if value.is_empty() {
        return Err(McpError::invalid_params(
            "brcode payload cannot be empty",
            None,
        ));
    }
    if value.len() > MAX_BRCODE_LENGTH {
        return Err(McpError::invalid_params(
            format!("brcode payload cannot exceed {MAX_BRCODE_LENGTH} characters"),
            None,
        ));
    }

    Ok(value)
}

#[tool_router]
impl AgentisPayServer {
    pub fn new(backend: Box<dyn PixBackend>) -> Self {
        Self {
            backend: Arc::new(Mutex::new(backend)),
            rate_limiter: RateLimiter::default(),
            tool_router: Self::tool_router(),
        }
    }

    #[tool(
        name = "agentispay_whoami",
        description = "Show current session status."
    )]
    async fn whoami(&self) -> Result<CallToolResult, McpError> {
        let mut backend = self.backend.lock().await;
        let data = backend
            .whoami()
            .await
            .map_err(|error| backend_error("session lookup", error))?;
        Ok(tool_result(CommandName::Whoami, json!(data)))
    }

    #[tool(
        name = "agentispay_account",
        description = "Read account profile and metadata."
    )]
    async fn account(&self) -> Result<CallToolResult, McpError> {
        let mut backend = self.backend.lock().await;
        let data = backend
            .account()
            .await
            .map_err(|error| backend_error("account lookup", error))?;
        Ok(tool_result(CommandName::Account, json!(data)))
    }

    #[tool(
        name = "agentispay_balance",
        description = "Read available balance in cents."
    )]
    async fn balance(&self) -> Result<CallToolResult, McpError> {
        let mut backend = self.backend.lock().await;
        let data = backend
            .balance()
            .await
            .map_err(|error| backend_error("balance lookup", error))?;
        Ok(tool_result(CommandName::Balance, json!(data)))
    }

    #[tool(
        name = "agentispay_history",
        description = "Read transaction history, optionally by id, with optional list limit."
    )]
    async fn history(&self, params: Parameters<HistoryInput>) -> Result<CallToolResult, McpError> {
        let mut backend = self.backend.lock().await;
        let query = if let Some(id) = params.0.id {
            HistoryQuery::Detail(id)
        } else {
            let limit = params.0.limit.unwrap_or(20);
            if limit == 0 {
                return Err(McpError::invalid_params(
                    "history limit must be greater than zero",
                    None,
                ));
            }
            if limit > 50 {
                return Err(McpError::invalid_params(
                    "history limit cannot exceed 50",
                    None,
                ));
            }
            HistoryQuery::List { limit }
        };
        let payload = match backend
            .history(query)
            .await
            .map_err(|error| backend_error("history lookup", error))?
        {
            HistoryRecord::List(list) => json!(list),
            HistoryRecord::Detail(detail) => json!(detail),
        };
        Ok(tool_result(CommandName::History, payload))
    }

    #[tool(
        name = "agentispay_brcode_decode",
        description = "Decode a PIX BR Code payload."
    )]
    async fn brcode_decode(
        &self,
        params: Parameters<BrcodeDecodeInput>,
    ) -> Result<CallToolResult, McpError> {
        self.rate_limiter
            .check("agentispay_brcode_decode", BRCODE_RATE_LIMIT)?;
        let code = validate_brcode_input(&params.0.code)?;
        let data = crate::commands::brcode::decode_payload(code).map_err(invalid_params)?;
        Ok(tool_result(CommandName::BrcodeDecode, json!(data)))
    }

    #[tool(name = "agentispay_deposit", description = "List deposit PIX keys.")]
    async fn deposit(&self) -> Result<CallToolResult, McpError> {
        let mut backend = self.backend.lock().await;
        let data = backend
            .deposit()
            .await
            .map_err(|error| backend_error("deposit key lookup", error))?;
        Ok(tool_result(CommandName::Deposit, json!(data)))
    }

    #[tool(
        name = "agentispay_pix_keys",
        description = "List configured PIX keys (CLI parity alias)."
    )]
    async fn pix_keys(&self) -> Result<CallToolResult, McpError> {
        let mut backend = self.backend.lock().await;
        let data = backend
            .pix_list()
            .await
            .map_err(|error| backend_error("pix key lookup", error))?;
        Ok(tool_result(CommandName::PixList, json!(data)))
    }

    #[tool(name = "agentispay_limits", description = "Read transfer risk limits.")]
    async fn limits(&self) -> Result<CallToolResult, McpError> {
        let mut backend = self.backend.lock().await;
        let data = backend
            .limits()
            .await
            .map_err(|error| backend_error("limits lookup", error))?;
        Ok(tool_result(CommandName::Limits, json!(data)))
    }

    #[tool(name = "agentispay_pay", description = "Create a PIX transfer.")]
    async fn pay(&self, params: Parameters<PayInput>) -> Result<CallToolResult, McpError> {
        self.pay_impl(params.0).await
    }

    #[tool(
        name = "agentispay_pix_send",
        description = "Alias for agentispay_pay."
    )]
    async fn pix_send(&self, params: Parameters<PayInput>) -> Result<CallToolResult, McpError> {
        self.pay_impl(params.0).await
    }
}

impl AgentisPayServer {
    async fn pay_impl(&self, mut input: PayInput) -> Result<CallToolResult, McpError> {
        if input.amount_cents <= 0 {
            return Err(McpError::invalid_params(
                "amount_cents must be greater than zero",
                None,
            ));
        }
        if input.request_id.trim().is_empty() {
            return Err(McpError::invalid_params("request_id is required", None));
        }
        if input.agent_message.trim().is_empty() {
            return Err(McpError::invalid_params("agent_message is required", None));
        }
        if input
            .idempotency_key
            .as_deref()
            .is_some_and(|token| token.trim().is_empty())
        {
            input.idempotency_key = None;
        }

        input.key = normalize_pix_key(&input.key).map_err(invalid_params)?;
        self.rate_limiter.check("agentispay_pay", PAY_RATE_LIMIT)?;

        let mut backend = self.backend.lock().await;
        let data = backend
            .pay(input)
            .await
            .map_err(|error| backend_error("payment request", error))?;
        Ok(tool_result(CommandName::PixSend, json!(data)))
    }
}

#[tool_handler]
impl ServerHandler for AgentisPayServer {
    async fn initialize(
        &self,
        request: InitializeRequestParams,
        context: RequestContext<RoleServer>,
    ) -> Result<InitializeResult, McpError> {
        if context.peer.peer_info().is_none() {
            context.peer.set_peer_info(request.clone());
        }

        let client_name = request.client_info.name.clone();
        if !client_name.is_empty() {
            let mut backend = self.backend.lock().await;
            backend.set_mcp_client_name(client_name);
        }

        let cli_version =
            std::env::var("AGENTIS_PAY_CLI_VERSION").unwrap_or_else(|_| "0.1.0".to_string());

        Ok(
            InitializeResult::new(ServerCapabilities::builder().enable_tools().build())
                .with_server_info(Implementation::new("agentis-pay", cli_version)),
        )
    }
}