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

use std::{
    collections::BTreeMap,
    convert::Infallible,
    sync::Arc,
    time::{Duration, Instant},
};

use axum::{
    Json, Router,
    extract::{Extension, Path, State},
    http::{HeaderMap, StatusCode},
    response::{IntoResponse, Response, sse::Event},
    routing::{get, post},
};
use crank_core::{
    ApprovalRequest, ApprovalRequestId, ApprovalRequestStatus, AuthProfile, CoordinationStateStore,
    InvocationLevel, InvocationSource, InvocationStatus, OperationApprovalMode,
    PlatformApiKeyScope, SecretId,
};
use crank_registry::{
    CreateApprovalRequest, DecideApprovalRequest, ExpireApprovalRequest, PostgresRegistry,
    PublishedAgentTool,
};
use crank_runtime::{
    RequestRateLimiter, ResolvedAuth, RuntimeError, RuntimeExecutionRequest, RuntimeExecutor,
    RuntimeOperation, RuntimeRequestContext, SecretCrypto,
};
use crank_trace::{DbOperation, ErrorCategory, Stage, StageOutcome, observe_db_query};
use futures_util::stream;
use serde::{Deserialize, Serialize};
use serde_json::{Value, json};
use time::OffsetDateTime;
use tokio::sync::Semaphore;
use tracing::{Instrument, info, warn};

use crate::{
    access::{
        credential_allows_security_level, serialize_machine_access_mode, serialize_security_level,
    },
    approval_execution::{execute_approved_request, spawn_approval_recovery},
    approval_response::approval_required_response,
    auth::{SharedMachineCredentialVerifier, VerifiedMachineCredential},
    catalog::PublishedToolCatalog,
    jsonrpc::{
        DEFAULT_PROTOCOL_VERSION, is_notification, is_request, is_response, jsonrpc_error,
        jsonrpc_result, method_name, negotiated_protocol_version, params, request_id,
    },
    manifest::catalog_tool_definitions,
    rate_limit::{rate_limited_jsonrpc_response, rate_limited_status_response},
    request_context::{RequestContext, apply_request_context},
    session::{SessionState, SharedSessionStore},
    tool_error::{
        ToolErrorContract, generic_tool_error_contract, runtime_error_code,
        tool_error_contract_from_runtime, tool_error_text, tool_error_value,
    },
    tool_search::handle_catalog_tool_call,
    transport::{
        AllowedOrigins, HEADER_MCP_SESSION_ID, ResponseMode, json_response,
        negotiate_post_response_mode, protocol_version_from_headers, session_id_from_headers,
        sse_response, transport_response, validate_get_accept_header, validate_origin,
        validate_session_protocol_version, with_request_id_header,
    },
};
mod invocation_history;
mod metrics;
mod stages;
use self::metrics::{ActiveSessionGuard, McpRequestMetrics};
use self::stages::{
    enforce_traced_rate_limit, require_traced_approval_access, require_traced_machine_access,
};
#[cfg(test)]
use invocation_history::observe_invocation_history_outcome;
pub(super) use invocation_history::{InvocationRecord, persist_invocation};

const TRANSPORT_SESSION_TTL_MS: u64 = 86_400_000;
const DEFAULT_MAX_CONCURRENT_SESSIONS: usize = 16;
const SESSION_CLEANUP_INTERVAL: Duration = Duration::from_secs(60);

#[derive(Clone)]
pub(super) struct AppState {
    pub(super) registry: PostgresRegistry,
    pub(super) catalog: PublishedToolCatalog,
    pub(super) runtime: RuntimeExecutor,
    pub(super) api_rate_limiter: RequestRateLimiter,
    secret_crypto: SecretCrypto,
    sessions: SharedSessionStore,
    session_slots: Arc<Semaphore>,
    pub(super) credential_verifier: SharedMachineCredentialVerifier,
    allowed_origins: AllowedOrigins,
}

#[derive(Debug, Serialize, Deserialize)]
struct InitializeParams {
    #[serde(rename = "protocolVersion")]
    protocol_version: String,
    #[serde(default)]
    capabilities: Value,
}

#[derive(Debug, Serialize, Deserialize)]
struct ToolCallParams {
    name: String,
    #[serde(default)]
    arguments: Value,
}

#[derive(Debug, Deserialize)]
struct ApprovalDecisionPayload {
    approve: String,
    #[serde(default)]
    note: Option<String>,
}

#[derive(Clone)]
pub(super) struct ResolvedToolCall {
    tool: PublishedAgentTool,
}

struct ToolCallExecution {
    tool: PublishedAgentTool,
    arguments: Value,
    confirmation_token: Option<String>,
}

#[derive(Clone, Debug, Deserialize)]
pub(super) struct AgentRoutePath {
    pub(super) workspace_slug: String,
    pub(super) agent_slug: String,
}

#[derive(Clone, Debug, Deserialize)]
struct ApprovalRoutePath {
    workspace_slug: String,
    agent_slug: String,
    approval_id: String,
}

#[allow(clippy::too_many_arguments)]
pub fn build_app(
    registry: PostgresRegistry,
    refresh_interval: Duration,
    public_base_url: Option<String>,
    secret_crypto: SecretCrypto,
    runtime: RuntimeExecutor,
    api_rate_limiter: RequestRateLimiter,
    coordination_store: Arc<dyn CoordinationStateStore>,
    sessions: SharedSessionStore,
    credential_verifier: SharedMachineCredentialVerifier,
) -> Router {
    build_app_inner(
        registry,
        refresh_interval,
        public_base_url,
        secret_crypto,
        runtime,
        api_rate_limiter,
        coordination_store,
        sessions,
        credential_verifier,
        DEFAULT_MAX_CONCURRENT_SESSIONS,
        false,
    )
}

#[allow(clippy::too_many_arguments)]
pub fn build_app_with_background_workers(
    registry: PostgresRegistry,
    refresh_interval: Duration,
    public_base_url: Option<String>,
    secret_crypto: SecretCrypto,
    runtime: RuntimeExecutor,
    api_rate_limiter: RequestRateLimiter,
    coordination_store: Arc<dyn CoordinationStateStore>,
    sessions: SharedSessionStore,
    credential_verifier: SharedMachineCredentialVerifier,
) -> Router {
    build_app_with_background_workers_and_limits(
        registry,
        refresh_interval,
        public_base_url,
        secret_crypto,
        runtime,
        api_rate_limiter,
        coordination_store,
        sessions,
        credential_verifier,
        DEFAULT_MAX_CONCURRENT_SESSIONS,
    )
}

#[allow(clippy::too_many_arguments)]
pub fn build_app_with_background_workers_and_limits(
    registry: PostgresRegistry,
    refresh_interval: Duration,
    public_base_url: Option<String>,
    secret_crypto: SecretCrypto,
    runtime: RuntimeExecutor,
    api_rate_limiter: RequestRateLimiter,
    coordination_store: Arc<dyn CoordinationStateStore>,
    sessions: SharedSessionStore,
    credential_verifier: SharedMachineCredentialVerifier,
    max_concurrent_sessions: usize,
) -> Router {
    build_app_inner(
        registry,
        refresh_interval,
        public_base_url,
        secret_crypto,
        runtime,
        api_rate_limiter,
        coordination_store,
        sessions,
        credential_verifier,
        max_concurrent_sessions,
        true,
    )
}

#[allow(clippy::too_many_arguments)]
fn build_app_inner(
    registry: PostgresRegistry,
    refresh_interval: Duration,
    public_base_url: Option<String>,
    secret_crypto: SecretCrypto,
    runtime: RuntimeExecutor,
    api_rate_limiter: RequestRateLimiter,
    coordination_store: Arc<dyn CoordinationStateStore>,
    sessions: SharedSessionStore,
    credential_verifier: SharedMachineCredentialVerifier,
    max_concurrent_sessions: usize,
    start_background_workers: bool,
) -> Router {
    let state = Arc::new(AppState {
        registry: registry.clone(),
        catalog: PublishedToolCatalog::new(registry, refresh_interval, coordination_store),
        runtime,
        api_rate_limiter,
        secret_crypto,
        sessions,
        session_slots: Arc::new(Semaphore::new(max_concurrent_sessions)),
        credential_verifier,
        allowed_origins: AllowedOrigins::new(public_base_url),
    });
    if start_background_workers {
        spawn_approval_recovery(Arc::clone(&state));
        spawn_session_cleanup(Arc::clone(&state));
    }

    Router::new()
        .route("/health", get(health))
        .route("/ready", get(readiness))
        .route(
            "/v1/{workspace_slug}/{agent_slug}",
            get(mcp_get).post(mcp_post).delete(mcp_delete),
        )
        .route(
            "/v1/{workspace_slug}/{agent_slug}/approvals",
            get(list_pending_approvals),
        )
        .route(
            "/v1/{workspace_slug}/{agent_slug}/approvals/{approval_id}/approve",
            post(approve_request),
        )
        .route(
            "/v1/{workspace_slug}/{agent_slug}/approvals/{approval_id}",
            get(get_approval_request),
        )
        .route(
            "/v1/{workspace_slug}/{agent_slug}/approvals/{approval_id}/deny",
            post(deny_request),
        )
        .with_state(state)
        .layer(axum::middleware::from_fn(apply_request_context))
        .layer(axum::middleware::from_fn(
            crank_observability::record_http_request,
        ))
}

async fn health() -> Json<Value> {
    Json(json!({
        "service": "mcp-server",
        "status": "ok"
    }))
}

fn spawn_session_cleanup(state: Arc<AppState>) {
    tokio::spawn(async move {
        let mut interval = tokio::time::interval(SESSION_CLEANUP_INTERVAL);
        loop {
            interval.tick().await;
            match state
                .sessions
                .cleanup_expired(OffsetDateTime::now_utc())
                .await
            {
                Ok(removed) if removed > 0 => {
                    info!(name: "mcp.session_cleanup.completed", removed);
                }
                Ok(_) => {}
                Err(_) => {
                    warn!(name: "mcp.session_cleanup.failed", error_category = "session_store");
                }
            }
        }
    });
}

async fn readiness(State(state): State<Arc<AppState>>) -> Response {
    match state.registry.ping().await {
        Ok(()) => Json(json!({
            "service": "mcp-server",
            "status": "ready",
            "checks": { "postgres": "ready" }
        }))
        .into_response(),
        Err(error) => (
            StatusCode::SERVICE_UNAVAILABLE,
            Json(json!({
                "service": "mcp-server",
                "status": "not_ready",
                "checks": { "postgres": "not_ready" },
                "error": error.to_string()
            })),
        )
            .into_response(),
    }
}

async fn list_pending_approvals(
    Path(path): Path<AgentRoutePath>,
    State(state): State<Arc<AppState>>,
    headers: HeaderMap,
) -> Response {
    if let Err(rejection) = enforce_traced_rate_limit(&state, &path, &headers).await {
        return rate_limited_status_response(rejection);
    }

    let key = match require_traced_approval_access(
        &state,
        &path,
        &headers,
        PlatformApiKeyScope::ReadPending,
    )
    .await
    {
        Ok(key) => key,
        Err(status) => return status.into_response(),
    };

    let Some(agent_id) = key.api_key.agent_id.as_ref() else {
        return StatusCode::FORBIDDEN.into_response();
    };

    match observe_db_query(
        DbOperation::ApprovalRead,
        state
            .registry
            .list_pending_approval_requests_for_agent(&key.api_key.workspace_id, agent_id),
    )
    .await
    {
        Ok(items) => Json(json!({ "items": items })).into_response(),
        Err(_) => StatusCode::INTERNAL_SERVER_ERROR.into_response(),
    }
}

async fn approve_request(
    Path(path): Path<ApprovalRoutePath>,
    State(state): State<Arc<AppState>>,
    Extension(request_context): Extension<RequestContext>,
    headers: HeaderMap,
    Json(payload): Json<ApprovalDecisionPayload>,
) -> Response {
    decide_approval_request(
        path,
        state,
        headers,
        payload,
        PlatformApiKeyScope::Approve,
        ApprovalRequestStatus::Approved,
        Some(request_context.request_id),
    )
    .await
}

async fn get_approval_request(
    Path(path): Path<ApprovalRoutePath>,
    State(state): State<Arc<AppState>>,
    headers: HeaderMap,
) -> Response {
    let agent_path = AgentRoutePath {
        workspace_slug: path.workspace_slug,
        agent_slug: path.agent_slug,
    };
    if let Err(rejection) = enforce_traced_rate_limit(&state, &agent_path, &headers).await {
        return rate_limited_status_response(rejection);
    }

    let key = match require_traced_approval_access(
        &state,
        &agent_path,
        &headers,
        PlatformApiKeyScope::ReadPending,
    )
    .await
    {
        Ok(key) => key,
        Err(status) => return status.into_response(),
    };
    let Some(agent_id) = key.api_key.agent_id.as_ref() else {
        return StatusCode::FORBIDDEN.into_response();
    };
    let approval_id = ApprovalRequestId::new(path.approval_id);

    approval_record_response(&state, &key.api_key.workspace_id, agent_id, &approval_id).await
}

async fn deny_request(
    Path(path): Path<ApprovalRoutePath>,
    State(state): State<Arc<AppState>>,
    headers: HeaderMap,
    Json(payload): Json<ApprovalDecisionPayload>,
) -> Response {
    decide_approval_request(
        path,
        state,
        headers,
        payload,
        PlatformApiKeyScope::Deny,
        ApprovalRequestStatus::Denied,
        None,
    )
    .await
}

async fn decide_approval_request(
    path: ApprovalRoutePath,
    state: Arc<AppState>,
    headers: HeaderMap,
    payload: ApprovalDecisionPayload,
    required_scope: PlatformApiKeyScope,
    status: ApprovalRequestStatus,
    execution_request_id: Option<String>,
) -> Response {
    let agent_path = AgentRoutePath {
        workspace_slug: path.workspace_slug,
        agent_slug: path.agent_slug,
    };
    if let Err(rejection) = enforce_traced_rate_limit(&state, &agent_path, &headers).await {
        return rate_limited_status_response(rejection);
    }

    let key =
        match require_traced_approval_access(&state, &agent_path, &headers, required_scope).await {
            Ok(key) => key,
            Err(status) => return status.into_response(),
        };

    if (status == ApprovalRequestStatus::Approved && !payload.approve.eq_ignore_ascii_case("yes"))
        || (status == ApprovalRequestStatus::Denied && !payload.approve.eq_ignore_ascii_case("no"))
    {
        return (
            StatusCode::BAD_REQUEST,
            Json(json!({
                "error": "invalid_decision_payload",
                "message": "approve must be yes for approve endpoint and no for deny endpoint"
            })),
        )
            .into_response();
    }

    let Some(agent_id) = key.api_key.agent_id.as_ref() else {
        return StatusCode::FORBIDDEN.into_response();
    };
    let approval_id = ApprovalRequestId::new(path.approval_id);

    match observe_db_query(
        DbOperation::ApprovalWrite,
        state
            .registry
            .decide_approval_request(DecideApprovalRequest {
                workspace_id: &key.api_key.workspace_id,
                agent_id,
                approval_id: &approval_id,
                status,
                decided_at: OffsetDateTime::now_utc(),
                decided_by_key_id: &key.api_key.id,
                response_payload: Some(json!({ "approve": payload.approve })),
                decision_note: payload.note.as_deref(),
            }),
    )
    .await
    {
        Ok(Some(record)) if status == ApprovalRequestStatus::Approved => {
            let claimed = match observe_db_query(
                DbOperation::ApprovalWrite,
                state.registry.claim_approval_request(
                    &record.approval.workspace_id,
                    &record.approval.agent_id,
                    &record.approval.id,
                    OffsetDateTime::now_utc(),
                ),
            )
            .await
            {
                Ok(Some(claimed)) => claimed,
                Ok(None) => return StatusCode::CONFLICT.into_response(),
                Err(_) => return StatusCode::INTERNAL_SERVER_ERROR.into_response(),
            };
            match execute_approved_request(
                &state,
                &agent_path,
                claimed,
                execution_request_id.as_deref(),
            )
            .await
            {
                Ok(record) => Json(json!(record)).into_response(),
                Err(response) => response,
            }
        }
        Ok(Some(record)) => Json(json!(record)).into_response(),
        Ok(None) => {
            terminal_decision_response(&state, &key.api_key.workspace_id, agent_id, &approval_id)
                .await
        }
        Err(_) => StatusCode::INTERNAL_SERVER_ERROR.into_response(),
    }
}

async fn approval_record_response(
    state: &Arc<AppState>,
    workspace_id: &crank_core::WorkspaceId,
    agent_id: &crank_core::AgentId,
    approval_id: &ApprovalRequestId,
) -> Response {
    match observe_db_query(
        DbOperation::ApprovalRead,
        state
            .registry
            .get_approval_request_for_agent(workspace_id, agent_id, approval_id),
    )
    .await
    {
        Ok(Some(record))
            if record.approval.status == ApprovalRequestStatus::Pending
                && record.approval.expires_at <= OffsetDateTime::now_utc() =>
        {
            expire_approval_response(state, workspace_id, agent_id, approval_id).await
        }
        Ok(Some(record)) => Json(json!(record)).into_response(),
        Ok(None) => StatusCode::NOT_FOUND.into_response(),
        Err(_) => StatusCode::INTERNAL_SERVER_ERROR.into_response(),
    }
}

async fn terminal_decision_response(
    state: &Arc<AppState>,
    workspace_id: &crank_core::WorkspaceId,
    agent_id: &crank_core::AgentId,
    approval_id: &ApprovalRequestId,
) -> Response {
    match observe_db_query(
        DbOperation::ApprovalRead,
        state
            .registry
            .get_approval_request_for_agent(workspace_id, agent_id, approval_id),
    )
    .await
    {
        Ok(Some(record))
            if record.approval.status == ApprovalRequestStatus::Pending
                && record.approval.expires_at <= OffsetDateTime::now_utc() =>
        {
            expire_approval_response(state, workspace_id, agent_id, approval_id).await
        }
        Ok(Some(record))
            if matches!(
                record.approval.status,
                ApprovalRequestStatus::Completed
                    | ApprovalRequestStatus::Failed
                    | ApprovalRequestStatus::Denied
                    | ApprovalRequestStatus::Expired
            ) =>
        {
            Json(json!(record)).into_response()
        }
        Ok(Some(_)) => StatusCode::CONFLICT.into_response(),
        Ok(None) => StatusCode::NOT_FOUND.into_response(),
        Err(_) => StatusCode::INTERNAL_SERVER_ERROR.into_response(),
    }
}

async fn expire_approval_response(
    state: &Arc<AppState>,
    workspace_id: &crank_core::WorkspaceId,
    agent_id: &crank_core::AgentId,
    approval_id: &ApprovalRequestId,
) -> Response {
    match observe_db_query(
        DbOperation::ApprovalWrite,
        state
            .registry
            .expire_approval_request(ExpireApprovalRequest {
                workspace_id,
                agent_id,
                approval_id,
                expired_at: OffsetDateTime::now_utc(),
            }),
    )
    .await
    {
        Ok(Some(record)) => Json(json!(record)).into_response(),
        Ok(None) => StatusCode::CONFLICT.into_response(),
        Err(_) => StatusCode::INTERNAL_SERVER_ERROR.into_response(),
    }
}

async fn mcp_get(
    Path(path): Path<AgentRoutePath>,
    State(state): State<Arc<AppState>>,
    headers: HeaderMap,
) -> Response {
    if let Err(status) = validate_origin(&state.allowed_origins, &headers) {
        return status.into_response();
    }

    if let Err(status) = validate_get_accept_header(&headers) {
        return status.into_response();
    }

    if let Err(status) =
        require_traced_machine_access(&state, &path, &headers, PlatformApiKeyScope::Read).await
    {
        return status.into_response();
    }

    if let Err(rejection) = enforce_traced_rate_limit(&state, &path, &headers).await {
        return rate_limited_status_response(rejection);
    }

    let session_id = match session_id_from_headers(&headers) {
        Ok(Some(session_id)) => session_id,
        Ok(None) => return StatusCode::BAD_REQUEST.into_response(),
        Err(status) => return status.into_response(),
    };

    let Some(session) = (match state.sessions.get(&session_id).await {
        Ok(session) => session,
        Err(_) => return StatusCode::INTERNAL_SERVER_ERROR.into_response(),
    }) else {
        return StatusCode::NOT_FOUND.into_response();
    };

    if session.workspace_slug != path.workspace_slug || session.agent_slug != path.agent_slug {
        return StatusCode::NOT_FOUND.into_response();
    }

    if !session.initialized {
        return StatusCode::BAD_REQUEST.into_response();
    }

    if let Err(status) = validate_session_protocol_version(&headers, &session.protocol_version) {
        return status.into_response();
    }

    let Ok(permit) = ActiveSessionGuard::try_acquire(&state.session_slots) else {
        return StatusCode::TOO_MANY_REQUESTS.into_response();
    };

    let stream = stream::unfold(permit, |permit| async move {
        tokio::time::sleep(Duration::from_millis(TRANSPORT_SESSION_TTL_MS)).await;
        drop(permit);
        None::<(Result<Event, Infallible>, _)>
    });

    sse_response(
        StatusCode::OK,
        stream,
        Some(&session_id),
        Some(&session.protocol_version),
    )
}

async fn mcp_delete(
    Path(path): Path<AgentRoutePath>,
    State(state): State<Arc<AppState>>,
    headers: HeaderMap,
) -> Response {
    if let Err(status) =
        require_traced_machine_access(&state, &path, &headers, PlatformApiKeyScope::Read).await
    {
        return status.into_response();
    }

    if let Err(rejection) = enforce_traced_rate_limit(&state, &path, &headers).await {
        return rate_limited_status_response(rejection);
    }

    match session_id_from_headers(&headers) {
        Ok(Some(session_id)) => match state.sessions.get(&session_id).await {
            Ok(Some(session))
                if session.workspace_slug == path.workspace_slug
                    && session.agent_slug == path.agent_slug =>
            {
                match state.sessions.delete(&session_id).await {
                    Ok(true) => StatusCode::NO_CONTENT.into_response(),
                    Ok(false) => StatusCode::NOT_FOUND.into_response(),
                    Err(_) => StatusCode::INTERNAL_SERVER_ERROR.into_response(),
                }
            }
            Ok(Some(_)) => StatusCode::NOT_FOUND.into_response(),
            Ok(None) => StatusCode::NOT_FOUND.into_response(),
            Err(_) => StatusCode::INTERNAL_SERVER_ERROR.into_response(),
        },
        Ok(None) => StatusCode::BAD_REQUEST.into_response(),
        Err(status) => status.into_response(),
    }
}

async fn mcp_post(
    Path(path): Path<AgentRoutePath>,
    State(state): State<Arc<AppState>>,
    Extension(request_context): Extension<RequestContext>,
    headers: HeaderMap,
    Json(message): Json<Value>,
) -> Response {
    let mut request_metrics = McpRequestMetrics::new(&message);
    let transport_request_id = request_context.request_id;
    info!(
        name: "mcp.request.received",
        request_id = %transport_request_id,
        workspace_slug = %path.workspace_slug,
        agent_slug = %path.agent_slug,
        jsonrpc_method = method_name(&message).unwrap_or("<unknown>"),
        "mcp request received"
    );

    if let Err(status) = validate_origin(&state.allowed_origins, &headers) {
        return with_request_id_header(status.into_response(), &transport_request_id);
    }

    let response_mode = match negotiate_post_response_mode(&headers) {
        Ok(mode) => request_metrics.set_response_mode(mode),
        Err(status) => {
            return with_request_id_header(status.into_response(), &transport_request_id);
        }
    };

    if is_response(&message) || is_notification(&message) && method_name(&message).is_none() {
        return with_request_id_header(StatusCode::ACCEPTED.into_response(), &transport_request_id);
    }

    let protocol_version = match protocol_version_from_headers(&headers) {
        Ok(value) => value,
        Err(status) => {
            return with_request_id_header(status.into_response(), &transport_request_id);
        }
    };

    let rate_limit_result = enforce_traced_rate_limit(&state, &path, &headers).await;
    if let Err(error) = rate_limit_result {
        return with_request_id_header(
            rate_limited_jsonrpc_response(&message, response_mode, &protocol_version, error),
            &transport_request_id,
        );
    }
    if let Some(session_id) = headers.get(HEADER_MCP_SESSION_ID)
        && let Ok(session_id) = session_id.to_str()
    {
        let session = match state.sessions.get(session_id).await {
            Ok(session) => session,
            Err(_) => {
                return with_request_id_header(
                    StatusCode::INTERNAL_SERVER_ERROR.into_response(),
                    &transport_request_id,
                );
            }
        };
        if let Some(session) = session
            && let Err(status) =
                validate_session_protocol_version(&headers, &session.protocol_version)
        {
            return with_request_id_header(status.into_response(), &transport_request_id);
        }
    }

    let required_scope = match method_name(&message) {
        Some("tools/call") => PlatformApiKeyScope::Write,
        _ => PlatformApiKeyScope::Read,
    };
    let access_result =
        require_traced_machine_access(&state, &path, &headers, required_scope).await;
    let credential = match access_result {
        Ok(credential) => credential,
        Err(error) => {
            return with_request_id_header(error.into_response(), &transport_request_id);
        }
    };

    let response = match method_name(&message) {
        Some("initialize") if is_request(&message) => {
            handle_initialize(state, &path, &message, response_mode).await
        }
        Some("notifications/initialized") if is_notification(&message) => {
            handle_initialized_notification(state, &path, &headers).await
        }
        Some("ping") if is_request(&message) => {
            let session = match require_initialized_session(&state, &path, &headers, &message).await
            {
                Ok(session) => session,
                Err(response) => return response,
            };

            transport_response(
                StatusCode::OK,
                jsonrpc_result(
                    request_id(&message),
                    json!({ "protocolVersion": session.protocol_version }),
                ),
                response_mode,
                None,
                Some(&session.protocol_version),
            )
        }
        Some("tools/list") if is_request(&message) => {
            let session = match require_initialized_session(&state, &path, &headers, &message).await
            {
                Ok(session) => session,
                Err(response) => return response,
            };

            match state
                .catalog
                .get_catalog(&session.workspace_slug, &session.agent_slug)
                .await
            {
                Ok(catalog) => {
                    let definitions = catalog_tool_definitions(&catalog);

                    transport_response(
                        StatusCode::OK,
                        jsonrpc_result(request_id(&message), json!({ "tools": definitions })),
                        response_mode,
                        None,
                        Some(&session.protocol_version),
                    )
                }
                Err(error) => internal_jsonrpc_error(&message, error),
            }
        }
        Some("tools/call") if is_request(&message) => {
            let session = match require_initialized_session(&state, &path, &headers, &message).await
            {
                Ok(session) => session,
                Err(response) => return response,
            };
            let tool_call_params: ToolCallParams = match serde_json::from_value(params(&message)) {
                Ok(value) => value,
                Err(error) => {
                    return transport_response(
                        StatusCode::OK,
                        jsonrpc_error(request_id(&message), -32602, error.to_string()),
                        response_mode,
                        None,
                        Some(&session.protocol_version),
                    );
                }
            };

            let arguments = if tool_call_params.arguments.is_null() {
                json!({})
            } else {
                tool_call_params.arguments
            };

            match state
                .catalog
                .get_catalog(&session.workspace_slug, &session.agent_slug)
                .await
            {
                Ok(catalog) => {
                    handle_catalog_tool_call(
                        state.clone(),
                        &session,
                        &message,
                        response_mode,
                        &credential,
                        &catalog,
                        &tool_call_params.name,
                        arguments,
                        &transport_request_id,
                    )
                    .await
                }
                Err(error) => internal_jsonrpc_error(&message, error),
            }
        }
        Some(method) if is_notification(&message) => {
            let _ = method;
            StatusCode::ACCEPTED.into_response()
        }
        Some(method) => transport_response(
            StatusCode::OK,
            jsonrpc_error(
                request_id(&message),
                -32601,
                format!("method {method} is not supported"),
            ),
            response_mode,
            None,
            Some(&protocol_version),
        ),
        None => transport_response(
            StatusCode::BAD_REQUEST,
            jsonrpc_error(Value::Null, -32600, "invalid JSON-RPC message"),
            response_mode,
            None,
            Some(&protocol_version),
        ),
    };

    request_metrics.complete(response.status());
    with_request_id_header(response, &transport_request_id)
}

#[allow(clippy::too_many_arguments)]
pub(super) async fn handle_tool_call(
    state: Arc<AppState>,
    session: &SessionState,
    message: &Value,
    response_mode: ResponseMode,
    credential: &VerifiedMachineCredential,
    resolved: ResolvedToolCall,
    arguments: Value,
    confirmation_token: Option<String>,
    transport_request_id: &str,
) -> Response {
    if !credential_allows_security_level(credential, resolved.tool.operation.security_level) {
        return tool_error_response(
            message,
            response_mode,
            &session.protocol_version,
            generic_tool_error_contract(
                "machine_access_insufficient",
                format!(
                    "machine access mode {} does not satisfy {} operation security",
                    serialize_machine_access_mode(credential.machine_access_mode),
                    serialize_security_level(resolved.tool.operation.security_level),
                ),
                transport_request_id,
                false,
                Some("Используйте ключ агента с достаточным уровнем доступа."),
            ),
        );
    }

    handle_base_tool_call(
        state,
        session,
        message,
        response_mode,
        ToolCallExecution {
            tool: resolved.tool,
            arguments,
            confirmation_token,
        },
        transport_request_id,
    )
    .await
}

pub(super) async fn resolve_operation_auth(
    state: &Arc<AppState>,
    workspace_id: &crank_core::WorkspaceId,
    execution_config: &crank_core::ExecutionConfig,
) -> Result<Option<ResolvedAuth>, RuntimeError> {
    if execution_config.auth_profile_ref.is_none() {
        return Ok(None);
    }
    let span = Stage::AuthResolve.span();
    let result = resolve_runtime_auth_for_task(
        &state.registry,
        &state.secret_crypto,
        workspace_id,
        execution_config,
    )
    .instrument(span.clone())
    .await;
    match &result {
        Ok(_) => StageOutcome::Success.record(&span),
        Err(_) => {
            StageOutcome::Error.record(&span);
            ErrorCategory::Configuration.record(&span);
        }
    }
    result
}

async fn resolve_runtime_auth_for_task(
    registry: &PostgresRegistry,
    secret_crypto: &SecretCrypto,
    workspace_id: &crank_core::WorkspaceId,
    execution_config: &crank_core::ExecutionConfig,
) -> Result<Option<ResolvedAuth>, RuntimeError> {
    let Some(auth_profile_id) = execution_config.auth_profile_ref.as_ref() else {
        return Ok(None);
    };

    let auth_profile = observe_db_query(
        DbOperation::AuthProfileRead,
        registry.get_auth_profile(workspace_id, auth_profile_id),
    )
    .await
    .map_err(|error| RuntimeError::SecretCrypto {
        operation: "load auth profile",
        details: error.to_string(),
    })?
    .ok_or_else(|| RuntimeError::MissingAuthProfile {
        auth_profile_id: auth_profile_id.as_str().to_owned(),
    })?;

    resolve_auth_profile(registry, secret_crypto, workspace_id, &auth_profile)
        .await
        .map(Some)
}

async fn resolve_auth_profile(
    registry: &PostgresRegistry,
    secret_crypto: &SecretCrypto,
    workspace_id: &crank_core::WorkspaceId,
    auth_profile: &AuthProfile,
) -> Result<ResolvedAuth, RuntimeError> {
    let mut secrets = BTreeMap::new();
    let used_at = OffsetDateTime::now_utc();

    for secret_id in auth_profile.config.secret_ids() {
        let secret = observe_db_query(
            DbOperation::SecretRead,
            registry.get_secret(workspace_id, secret_id),
        )
        .await
        .map_err(|error| RuntimeError::SecretCrypto {
            operation: "load secret",
            details: error.to_string(),
        })?
        .ok_or_else(|| RuntimeError::MissingSecret {
            secret_id: secret_id.as_str().to_owned(),
        })?;
        let version = observe_db_query(
            DbOperation::SecretRead,
            registry.get_current_secret_version(workspace_id, secret_id),
        )
        .await
        .map_err(|error| RuntimeError::SecretCrypto {
            operation: "load current secret version",
            details: error.to_string(),
        })?
        .ok_or_else(|| RuntimeError::MissingSecretVersion {
            secret_id: secret_id.as_str().to_owned(),
            version: secret.secret.current_version,
        })?;
        let plaintext = secret_crypto.decrypt(
            &version.secret_version.key_version,
            &version.secret_version.ciphertext,
        )?;
        observe_db_query(
            DbOperation::SecretTouch,
            registry.touch_secret(workspace_id, secret_id, &used_at),
        )
        .await
        .map_err(|error| RuntimeError::SecretCrypto {
            operation: "touch secret",
            details: error.to_string(),
        })?;
        secrets.insert(SecretId::new(secret_id.as_str()), plaintext);
    }

    ResolvedAuth::from_profile(auth_profile, &secrets)
}

async fn handle_base_tool_call(
    state: Arc<AppState>,
    session: &SessionState,
    message: &Value,
    response_mode: ResponseMode,
    execution: ToolCallExecution,
    transport_request_id: &str,
) -> Response {
    let tool = execution.tool;
    let arguments = execution.arguments;
    let operation = runtime_operation(&tool);
    if tool
        .operation
        .execution_config
        .approval_policy
        .as_ref()
        .is_some_and(|policy| policy.required)
    {
        let approval_span = Stage::ApprovalCheck.span();
        let response = maybe_handle_approval_policy(
            &state,
            session,
            message,
            response_mode,
            &tool,
            &arguments,
            transport_request_id,
        )
        .instrument(approval_span.clone())
        .await;
        if let Some(result) = response {
            return match result {
                ApprovalPolicyResult::Required(response) => {
                    StageOutcome::Required.record(&approval_span);
                    response
                }
                ApprovalPolicyResult::Error(response) => {
                    StageOutcome::Error.record(&approval_span);
                    ErrorCategory::Approval.record(&approval_span);
                    response
                }
            };
        }
        StageOutcome::Allowed.record(&approval_span);
    }

    let mut runtime_request_context = RuntimeRequestContext::from_request_id(transport_request_id)
        .with_response_cache_scope(
            tool.workspace_id.as_str().to_owned(),
            tool.agent_id.as_str().to_owned(),
        )
        .with_metering_context(
            tool.workspace_id.clone(),
            Some(tool.agent_id.clone()),
            InvocationSource::AgentToolCall,
        );
    if let Some(token) = execution.confirmation_token {
        runtime_request_context = runtime_request_context.with_confirmation_token(token);
    }
    let request_preview = build_request_preview(&state.runtime, &operation, &arguments);
    let started_at = Instant::now();
    let resolved_auth =
        resolve_operation_auth(&state, &tool.workspace_id, &operation.execution_config).await;

    let result = match resolved_auth {
        Ok(resolved_auth) => {
            state
                .runtime
                .execute_request(
                    RuntimeExecutionRequest::new(&operation, &arguments)
                        .with_optional_auth(resolved_auth.as_ref())
                        .with_context(&runtime_request_context),
                )
                .await
        }
        Err(error) => Err(error),
    };

    match result {
        Ok(output) => {
            persist_invocation(
                &state,
                &tool,
                InvocationRecord {
                    request_id: Some(transport_request_id),
                    tool_name: &tool.tool_name,
                    status: InvocationStatus::Ok,
                    level: InvocationLevel::Info,
                    message: "agent tool call completed",
                    status_code: None,
                    error_kind: None,
                    duration: started_at.elapsed(),
                    request_preview,
                    response_preview: output.clone(),
                },
            )
            .await;

            success_tool_response(message, response_mode, &session.protocol_version, output)
        }
        Err(error) => {
            persist_invocation(
                &state,
                &tool,
                InvocationRecord {
                    request_id: Some(transport_request_id),
                    tool_name: &tool.tool_name,
                    status: InvocationStatus::Error,
                    level: InvocationLevel::Error,
                    message: &error.to_string(),
                    status_code: None,
                    error_kind: Some(runtime_error_code(&error)),
                    duration: started_at.elapsed(),
                    request_preview,
                    response_preview: Value::Null,
                },
            )
            .await;

            tool_error_response(
                message,
                response_mode,
                &session.protocol_version,
                tool_error_contract_from_runtime(&error, transport_request_id),
            )
        }
    }
}

enum ApprovalPolicyResult {
    Required(Response),
    Error(Response),
}

async fn maybe_handle_approval_policy(
    state: &Arc<AppState>,
    session: &SessionState,
    message: &Value,
    response_mode: ResponseMode,
    tool: &PublishedAgentTool,
    arguments: &Value,
    transport_request_id: &str,
) -> Option<ApprovalPolicyResult> {
    let policy = tool.operation.execution_config.approval_policy.as_ref()?;
    if !policy.required {
        return None;
    }

    match policy.mode {
        OperationApprovalMode::Custom => {
            maybe_create_custom_pending_approval(
                state,
                session,
                message,
                response_mode,
                tool,
                arguments,
                transport_request_id,
            )
            .await
        }
        OperationApprovalMode::Elicitation => Some(handle_elicitation_approval(
            session,
            message,
            response_mode,
            tool,
            arguments,
            policy.elicitation_message.as_deref(),
            transport_request_id,
        )),
    }
}

async fn maybe_create_custom_pending_approval(
    state: &Arc<AppState>,
    session: &SessionState,
    message: &Value,
    response_mode: ResponseMode,
    tool: &PublishedAgentTool,
    arguments: &Value,
    transport_request_id: &str,
) -> Option<ApprovalPolicyResult> {
    let policy = tool.operation.execution_config.approval_policy.as_ref()?;

    let approval_id = ApprovalRequestId::new(format!("approval_{}", uuid::Uuid::now_v7().simple()));
    let now = OffsetDateTime::now_utc();
    let expires_at = now + time::Duration::seconds(i64::from(policy.ttl_seconds));
    let approval = ApprovalRequest {
        id: approval_id,
        workspace_id: tool.workspace_id.clone(),
        agent_id: tool.agent_id.clone(),
        operation_id: tool.operation.id.clone(),
        operation_version: tool.operation.version,
        status: ApprovalRequestStatus::Pending,
        risk_level: policy.risk_level,
        request_payload: arguments.clone(),
        response_payload: None,
        created_at: now,
        expires_at,
        decided_at: None,
        decided_by_key_id: None,
        decision_note: None,
    };

    let persisted_approval = match observe_db_query(
        DbOperation::ApprovalWrite,
        state
            .registry
            .create_approval_request(CreateApprovalRequest {
                approval: &approval,
            }),
    )
    .await
    {
        Ok(approval) => approval,
        Err(error) => {
            return Some(ApprovalPolicyResult::Error(internal_jsonrpc_error(
                message, error,
            )));
        }
    };
    let response_payload = approval_required_response(tool, &persisted_approval.approval, policy);

    persist_invocation(
        state,
        tool,
        InvocationRecord {
            request_id: Some(transport_request_id),
            tool_name: &tool.tool_name,
            status: InvocationStatus::Ok,
            level: InvocationLevel::Info,
            message: "agent tool call is waiting for human approval",
            status_code: None,
            error_kind: None,
            duration: Duration::from_millis(0),
            request_preview: arguments.clone(),
            response_preview: response_payload.clone(),
        },
    )
    .await;

    Some(ApprovalPolicyResult::Required(success_tool_response(
        message,
        response_mode,
        &session.protocol_version,
        response_payload,
    )))
}

fn handle_elicitation_approval(
    session: &SessionState,
    message: &Value,
    response_mode: ResponseMode,
    tool: &PublishedAgentTool,
    arguments: &Value,
    elicitation_message: Option<&str>,
    transport_request_id: &str,
) -> ApprovalPolicyResult {
    if !session.supports_elicitation {
        return ApprovalPolicyResult::Error(tool_error_response(
            message,
            response_mode,
            &session.protocol_version,
            generic_tool_error_contract(
                "approval_elicitation_not_supported",
                "operation requires MCP Elicitation, but the MCP client did not advertise elicitation capability",
                transport_request_id,
                false,
                Some(
                    "Выберите Custom MCP Approval или подключите MCP-клиент с поддержкой elicitation.",
                ),
            ),
        ));
    }

    let payload_preview = tool
        .operation
        .execution_config
        .approval_policy
        .as_ref()
        .and_then(|policy| policy.show_payload_preview.then(|| arguments.clone()))
        .unwrap_or(Value::Null);

    ApprovalPolicyResult::Required(success_tool_response(
        message,
        response_mode,
        &session.protocol_version,
        json!({
            "status": "elicitation_required",
            "message": elicitation_message.unwrap_or("Confirm operation execution."),
            "tool": tool.tool_name,
            "payload_preview": payload_preview,
            "note": "This MCP client advertised elicitation support. Full elicitation/create continuation is handled by compatible client integrations.",
        }),
    ))
}

async fn handle_initialize(
    state: Arc<AppState>,
    path: &AgentRoutePath,
    message: &Value,
    response_mode: ResponseMode,
) -> Response {
    let initialize_params: InitializeParams = match serde_json::from_value(params(message)) {
        Ok(value) => value,
        Err(error) => {
            return transport_response(
                StatusCode::OK,
                jsonrpc_error(request_id(message), -32602, error.to_string()),
                response_mode,
                None,
                Some(DEFAULT_PROTOCOL_VERSION),
            );
        }
    };
    let Some(protocol_version) = negotiated_protocol_version(&initialize_params.protocol_version)
    else {
        return transport_response(
            StatusCode::OK,
            jsonrpc_error(
                request_id(message),
                -32602,
                format!(
                    "unsupported protocol version {}",
                    initialize_params.protocol_version
                ),
            ),
            response_mode,
            None,
            Some(DEFAULT_PROTOCOL_VERSION),
        );
    };
    let now = OffsetDateTime::now_utc();
    let expires_at = add_millis(now, TRANSPORT_SESSION_TTL_MS);
    let supports_elicitation = initialize_params
        .capabilities
        .get("elicitation")
        .is_some_and(Value::is_object);
    let session_id = match state
        .sessions
        .create(
            protocol_version,
            &path.workspace_slug,
            &path.agent_slug,
            supports_elicitation,
            now,
            Some(expires_at),
        )
        .await
    {
        Ok(session_id) => session_id,
        Err(error) => return internal_jsonrpc_error(message, error),
    };

    transport_response(
        StatusCode::OK,
        jsonrpc_result(
            request_id(message),
            json!({
                "protocolVersion": protocol_version,
                "capabilities": {
                    "tools": {
                        "listChanged": false
                    }
                },
                "serverInfo": {
                    "name": "crank-mcp-server",
                    "version": env!("CARGO_PKG_VERSION")
                }
            }),
        ),
        response_mode,
        Some(&session_id),
        Some(protocol_version),
    )
}

async fn handle_initialized_notification(
    state: Arc<AppState>,
    path: &AgentRoutePath,
    headers: &HeaderMap,
) -> Response {
    match session_id_from_headers(headers) {
        Ok(Some(session_id)) => match state.sessions.get(&session_id).await {
            Ok(Some(session))
                if session.workspace_slug == path.workspace_slug
                    && session.agent_slug == path.agent_slug =>
            {
                match state
                    .sessions
                    .mark_initialized(&session_id, OffsetDateTime::now_utc())
                    .await
                {
                    Ok(true) => StatusCode::ACCEPTED.into_response(),
                    Ok(false) => StatusCode::NOT_FOUND.into_response(),
                    Err(_) => StatusCode::INTERNAL_SERVER_ERROR.into_response(),
                }
            }
            Ok(Some(_)) => StatusCode::NOT_FOUND.into_response(),
            Ok(None) => StatusCode::NOT_FOUND.into_response(),
            Err(_) => StatusCode::INTERNAL_SERVER_ERROR.into_response(),
        },
        Ok(None) => StatusCode::BAD_REQUEST.into_response(),
        Err(status) => status.into_response(),
    }
}

async fn require_initialized_session(
    state: &Arc<AppState>,
    path: &AgentRoutePath,
    headers: &HeaderMap,
    message: &Value,
) -> Result<SessionState, Response> {
    let session_id = match session_id_from_headers(headers) {
        Ok(Some(session_id)) => session_id,
        Ok(None) => return Err(StatusCode::BAD_REQUEST.into_response()),
        Err(status) => return Err(status.into_response()),
    };

    let Some(session) = (match state.sessions.get(&session_id).await {
        Ok(session) => session,
        Err(_) => return Err(StatusCode::INTERNAL_SERVER_ERROR.into_response()),
    }) else {
        return Err(StatusCode::NOT_FOUND.into_response());
    };

    if session.workspace_slug != path.workspace_slug || session.agent_slug != path.agent_slug {
        return Err(StatusCode::NOT_FOUND.into_response());
    }

    if !session.initialized {
        return Err(json_response(
            StatusCode::OK,
            jsonrpc_error(request_id(message), -32002, "session is not initialized"),
            None,
            Some(&session.protocol_version),
        ));
    }

    Ok(session)
}

fn internal_jsonrpc_error(message: &Value, error: impl std::fmt::Display) -> Response {
    transport_response(
        StatusCode::INTERNAL_SERVER_ERROR,
        jsonrpc_error(request_id(message), -32603, error.to_string()),
        ResponseMode::Json,
        None,
        Some(DEFAULT_PROTOCOL_VERSION),
    )
}

pub(super) fn take_confirmation_token(arguments: &mut Value) -> Option<String> {
    let Value::Object(object) = arguments else {
        return None;
    };
    object
        .remove("_crank_confirmation_token")
        .and_then(|value| value.as_str().map(str::to_owned))
        .filter(|value| !value.trim().is_empty())
}

pub(super) fn build_request_preview(
    runtime: &RuntimeExecutor,
    operation: &RuntimeOperation,
    arguments: &Value,
) -> Value {
    match runtime.prepare_request(operation, arguments) {
        Ok(prepared) => json!({
            "path": prepared.path_params,
            "query": prepared.query_params,
            "headers": prepared.headers,
            "body": prepared.body.unwrap_or(Value::Null)
        }),
        Err(_) => Value::Null,
    }
}

fn success_tool_response(
    message: &Value,
    response_mode: ResponseMode,
    protocol_version: &str,
    output: Value,
) -> Response {
    transport_response(
        StatusCode::OK,
        jsonrpc_result(
            request_id(message),
            json!({
                "content": [
                    {
                        "type": "text",
                        "text": serde_json::to_string_pretty(&output).unwrap_or_else(|_| "{}".to_owned())
                    }
                ],
                "structuredContent": output,
                "isError": false
            }),
        ),
        response_mode,
        None,
        Some(protocol_version),
    )
}

pub(super) fn tool_error_response(
    message: &Value,
    response_mode: ResponseMode,
    protocol_version: &str,
    error: ToolErrorContract,
) -> Response {
    let error_message = tool_error_text(&error);
    let error_value = tool_error_value(&error);

    transport_response(
        StatusCode::OK,
        jsonrpc_result(
            request_id(message),
            json!({
                "content": [
                    {
                        "type": "text",
                        "text": error_message
                    }
                ],
                "structuredContent": {
                    "error": error_value
                },
                "isError": true
            }),
        ),
        response_mode,
        None,
        Some(protocol_version),
    )
}

fn add_millis(timestamp: OffsetDateTime, millis: u64) -> OffsetDateTime {
    let delta = time::Duration::milliseconds(i64::try_from(millis).unwrap_or(i64::MAX));

    timestamp + delta
}

pub(super) fn resolve_generated_tool(
    tools: &[PublishedAgentTool],
    tool_name: &str,
) -> Option<ResolvedToolCall> {
    for tool in tools {
        if tool.tool_name == tool_name {
            return Some(ResolvedToolCall { tool: tool.clone() });
        }
    }

    None
}

pub(super) fn runtime_operation(tool: &PublishedAgentTool) -> RuntimeOperation {
    let mut operation = RuntimeOperation::from(tool.operation.clone());
    operation.tool_name = tool.tool_name.clone();
    operation.tool_description.title = tool.tool_title.clone();
    operation.tool_description.description = tool.tool_description.clone();
    operation
}

#[cfg(test)]
mod tests;