crank/crates/crank-community-mcp/src/rate_limit.rs

use std::sync::Arc;

use axum::{
    http::{HeaderMap, HeaderValue, StatusCode, header::RETRY_AFTER},
    response::{IntoResponse, Response},
};
use crank_runtime::RateLimitCheckError;
use serde_json::{Value, json};

use crate::{
    access::{bearer_token, hash_access_secret},
    app::{AgentRoutePath, AppState},
    jsonrpc::request_id,
    transport::{ResponseMode, session_id_from_headers, transport_response},
};

pub(super) async fn enforce_transport_rate_limit(
    state: &Arc<AppState>,
    path: &AgentRoutePath,
    headers: &HeaderMap,
) -> Result<(), RateLimitCheckError> {
    let key = rate_limit_key(path, headers);
    state.api_rate_limiter.check(&key).await
}

pub(super) fn rate_limited_jsonrpc_response(
    message: &Value,
    response_mode: ResponseMode,
    protocol_version: &str,
    error: RateLimitCheckError,
) -> Response {
    let RateLimitCheckError::Rejected(rejection) = error else {
        return transport_response(
            StatusCode::SERVICE_UNAVAILABLE,
            json!({
                "jsonrpc": "2.0",
                "id": request_id(message),
                "error": {
                    "code": -32603,
                    "message": "rate limit service unavailable",
                    "data": { "code": "rate_limit_unavailable" }
                }
            }),
            response_mode,
            None,
            Some(protocol_version),
        );
    };
    let payload = json!({
        "jsonrpc": "2.0",
        "id": request_id(message),
        "error": {
            "code": -32029,
            "message": "request rate limit exceeded",
            "data": {
                "code": "request_rate_limited",
                "retry_after_ms": rejection.retry_after_ms,
            }
        }
    });

    let mut response = transport_response(
        StatusCode::TOO_MANY_REQUESTS,
        payload,
        response_mode,
        None,
        Some(protocol_version),
    );
    attach_retry_after_header(&mut response, rejection.retry_after_ms);
    response
}

pub(super) fn rate_limited_status_response(error: RateLimitCheckError) -> Response {
    match error {
        RateLimitCheckError::Rejected(rejection) => {
            let mut response = StatusCode::TOO_MANY_REQUESTS.into_response();
            attach_retry_after_header(&mut response, rejection.retry_after_ms);
            response
        }
        RateLimitCheckError::StoreUnavailable => StatusCode::SERVICE_UNAVAILABLE.into_response(),
    }
}

fn rate_limit_key(path: &AgentRoutePath, headers: &HeaderMap) -> String {
    if let Ok(Some(session_id)) = session_id_from_headers(headers) {
        return format!(
            "session:{}:{}:{}",
            path.workspace_slug, path.agent_slug, session_id
        );
    }

    if let Some(secret) = bearer_token(headers) {
        return format!("api_key:{}", hash_access_secret(secret));
    }

    format!("workspace:{}:anonymous", path.workspace_slug)
}

fn attach_retry_after_header(response: &mut Response, retry_after_ms: u64) {
    let retry_after_seconds = retry_after_ms.div_ceil(1000);
    if let Ok(value) = HeaderValue::from_str(&retry_after_seconds.to_string()) {
        response.headers_mut().insert(RETRY_AFTER, value);
    }
}