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

use std::{sync::Arc, time::Instant};

use axum::{
    http::StatusCode,
    response::{IntoResponse, Response},
};
use crank_core::{ApprovalRequestStatus, InvocationLevel, InvocationSource, InvocationStatus};
use crank_core::{CorrelationContext, RequestId, TraceContext};
use crank_registry::{ApprovalRequestRecord, FinishApprovalRequest};
use crank_runtime::{RuntimeExecutionRequest, RuntimeRequestContext};
use crank_trace::{DbOperation, ErrorCategory, Stage, StageOutcome, observe_db_query};
use serde_json::json;
use time::OffsetDateTime;
use tracing::{Instrument, warn};

use crate::{
    app::{
        AgentRoutePath, AppState, InvocationRecord, build_request_preview, persist_invocation,
        resolve_operation_auth, runtime_operation,
    },
    tool_error::{runtime_error_code, safe_runtime_error_message},
};

const RECOVERY_INTERVAL: std::time::Duration = std::time::Duration::from_secs(5);
const RECOVERY_GRACE: time::Duration = time::Duration::seconds(5);
const EXECUTION_LEASE: time::Duration = time::Duration::minutes(6);

pub(super) fn spawn_approval_recovery(state: Arc<AppState>) {
    tokio::spawn(async move {
        let mut interval = tokio::time::interval(RECOVERY_INTERVAL);
        interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
        loop {
            interval.tick().await;
            recover_approved_requests(&state).await;
        }
    });
}

async fn recover_approved_requests(state: &Arc<AppState>) {
    fail_interrupted_requests(state).await;

    for _ in 0..32 {
        let now = OffsetDateTime::now_utc();
        let approval = match observe_db_query(
            DbOperation::ApprovalWrite,
            state
                .registry
                .claim_next_recoverable_approval_request(now, now - RECOVERY_GRACE),
        )
        .await
        {
            Ok(Some(approval)) => approval,
            Ok(None) => break,
            Err(_) => {
                warn!(
                    name: "mcp.approval_recovery.query_failed",
                    error_category = "registry",
                    "approval recovery query failed"
                );
                break;
            }
        };
        let Some(path) = approval_agent_path(state, &approval).await else {
            continue;
        };
        let recovery_span = Stage::ApprovalRecovery.span();
        let trace_context = crank_trace::trace_context_for_span(&recovery_span)
            .unwrap_or_else(TraceContext::generate);
        let correlation = CorrelationContext::new(RequestId::generate(), trace_context);
        let result = execute_approved_request(state, &path, approval, Some(&correlation))
            .instrument(recovery_span.clone())
            .await;
        match &result {
            Ok(_) => StageOutcome::Success.record(&recovery_span),
            Err(_) => {
                StageOutcome::Error.record(&recovery_span);
                ErrorCategory::Approval.record(&recovery_span);
            }
        }
        drop(recovery_span);
        if result.is_err() {
            warn!(
                name: "mcp.approval_recovery.execution_failed",
                error_category = "runtime",
                "recovered approval execution did not finish"
            );
        }
    }
}

async fn fail_interrupted_requests(state: &Arc<AppState>) {
    for _ in 0..32 {
        let stale_before = OffsetDateTime::now_utc() - EXECUTION_LEASE;
        match observe_db_query(
            DbOperation::ApprovalWrite,
            state
                .registry
                .fail_next_interrupted_approval_request(stale_before),
        )
        .await
        {
            Ok(Some(approval)) => {
                warn!(
                    name: "mcp.approval_recovery.interrupted",
                    approval_id = approval.approval.id.as_str(),
                    "interrupted approval execution was not retried because its outcome is unknown"
                );
            }
            Ok(None) => break,
            Err(_) => {
                warn!(
                    name: "mcp.approval_recovery.interrupted_query_failed",
                    error_category = "registry",
                    "interrupted approval recovery query failed"
                );
                break;
            }
        }
    }
}

async fn approval_agent_path(
    state: &Arc<AppState>,
    approval: &ApprovalRequestRecord,
) -> Option<AgentRoutePath> {
    let workspace = match state
        .registry
        .get_workspace(&approval.approval.workspace_id)
        .await
    {
        Ok(Some(workspace)) => workspace,
        Ok(None) => return None,
        Err(_) => {
            warn!(
                name: "mcp.approval_recovery.workspace_lookup_failed",
                error_category = "registry",
                "approval workspace lookup failed"
            );
            return None;
        }
    };
    let agent = match state
        .registry
        .get_agent_summary(&approval.approval.workspace_id, &approval.approval.agent_id)
        .await
    {
        Ok(Some(agent)) => agent,
        Ok(None) => return None,
        Err(_) => {
            warn!(
                name: "mcp.approval_recovery.agent_lookup_failed",
                error_category = "registry",
                "approval agent lookup failed"
            );
            return None;
        }
    };
    Some(AgentRoutePath {
        workspace_slug: workspace.workspace.slug,
        agent_slug: agent.slug,
    })
}

pub(super) async fn execute_approved_request(
    state: &Arc<AppState>,
    path: &AgentRoutePath,
    approval: ApprovalRequestRecord,
    correlation: Option<&CorrelationContext>,
) -> Result<ApprovalRequestRecord, Response> {
    let correlation = correlation
        .cloned()
        .unwrap_or_else(CorrelationContext::generate);
    let request_id = correlation.request_id().as_str();
    let tools = state
        .catalog
        .list_tools(&path.workspace_slug, &path.agent_slug)
        .await
        .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR.into_response())?;
    let Some(tool) = tools.into_iter().find(|tool| {
        tool.operation.id == approval.approval.operation_id
            && tool.operation.version == approval.approval.operation_version
    }) else {
        return finish_unavailable_approval(state, &approval).await;
    };

    let operation = runtime_operation(&tool);
    let request_preview = build_request_preview(
        &state.runtime,
        &operation,
        &approval.approval.request_payload,
    );
    let started_at = Instant::now();
    let runtime_request_context = RuntimeRequestContext::from_correlation(&correlation)
        .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,
        )
        .with_approval_granted();
    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, &approval.approval.request_payload)
                        .with_optional_auth(resolved_auth.as_ref())
                        .with_context(&runtime_request_context),
                )
                .await
        }
        Err(error) => Err(error),
    };

    let (status, response_payload, invocation_status, invocation_level, message, error_kind) =
        match result {
            Ok(output) => (
                ApprovalRequestStatus::Completed,
                output,
                InvocationStatus::Ok,
                InvocationLevel::Info,
                "approved tool call completed",
                None,
            ),
            Err(error) => (
                ApprovalRequestStatus::Failed,
                json!({
                    "error": {
                        "code": runtime_error_code(&error),
                        "message": safe_runtime_error_message(&error),
                    }
                }),
                InvocationStatus::Error,
                InvocationLevel::Error,
                "approved tool call failed",
                Some(runtime_error_code(&error)),
            ),
        };

    persist_invocation(
        state,
        &tool,
        InvocationRecord {
            request_id: Some(request_id),
            trace_id: Some(correlation.trace_id().as_str()),
            tool_name: &tool.tool_name,
            status: invocation_status,
            level: invocation_level,
            message,
            status_code: None,
            error_kind,
            duration: started_at.elapsed(),
            request_preview,
            response_preview: response_payload.clone(),
        },
    )
    .await;

    observe_db_query(
        DbOperation::ApprovalWrite,
        state
            .registry
            .finish_approval_request(FinishApprovalRequest {
                workspace_id: &approval.approval.workspace_id,
                agent_id: &approval.approval.agent_id,
                approval_id: &approval.approval.id,
                status,
                response_payload: Some(response_payload),
                decision_note: None,
            }),
    )
    .await
    .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR.into_response())?
    .ok_or_else(|| StatusCode::CONFLICT.into_response())
}

async fn finish_unavailable_approval(
    state: &Arc<AppState>,
    approval: &ApprovalRequestRecord,
) -> Result<ApprovalRequestRecord, Response> {
    observe_db_query(
        DbOperation::ApprovalWrite,
        state
            .registry
            .finish_approval_request(FinishApprovalRequest {
                workspace_id: &approval.approval.workspace_id,
                agent_id: &approval.approval.agent_id,
                approval_id: &approval.approval.id,
                status: ApprovalRequestStatus::Failed,
                response_payload: Some(json!({
                    "error": {
                        "code": "approved_operation_unavailable",
                        "message": "the approved operation version is no longer published"
                    }
                })),
                decision_note: None,
            }),
    )
    .await
    .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR.into_response())?
    .ok_or_else(|| StatusCode::CONFLICT.into_response())
}