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::{Value, json};
use time::OffsetDateTime;
use tracing::{Instrument, warn};

use crate::app::{
    AgentRoutePath, AppState, InvocationRecord, persist_invocation, resolve_operation_auth,
    runtime_operation,
};

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 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) => {
            match RuntimeExecutionRequest::try_new(
                &tool.workspace_id,
                crank_core::ExecutionOrigin::AgentSnapshot,
                Some(&tool.agent_id),
                &operation,
                &approval.approval.request_payload,
                crank_runtime::ExecutionAuthorization::Authorized,
                resolved_auth.as_ref(),
                &runtime_request_context,
                Instant::now()
                    + std::time::Duration::from_millis(
                        operation.execution_config.timeout_ms.max(1),
                    ),
            ) {
                Ok(request) => state.runtime.execute_outcome(request).await,
                Err(_) => Err(crank_core::ExecutionFailure::new(
                    crank_core::ExecutionErrorCode::RuntimeInternal,
                    correlation.clone(),
                )),
            }
        }
        Err(error) => Err(crank_runtime::normalize_runtime_error(&error, &correlation)),
    };

    let (
        status,
        response_payload,
        invocation_status,
        invocation_level,
        message,
        error_kind,
        execution_stage,
        execution_error_code,
        retryability,
        outcome_certainty,
        upstream_status,
        request_preview,
    ) = match result {
        Ok(success) => {
            let request_preview = success.request_preview;
            (
                ApprovalRequestStatus::Completed,
                success.output,
                InvocationStatus::Ok,
                InvocationLevel::Info,
                "approved tool call completed",
                None,
                Some(crank_core::ExecutionStage::Runtime),
                None,
                Some(crank_core::Retryability::Never),
                Some(crank_core::OutcomeCertainty::Certain),
                None,
                request_preview,
            )
        }
        Err(failure) => (
            ApprovalRequestStatus::Failed,
            json!({
                "error": {
                    "code": failure.error_code().as_str(),
                    "message": failure.error_code().message(crank_core::ExecutionLocale::Ru),
                    "stage": failure.stage().as_str(),
                    "retryability": failure.retryability().as_str(),
                    "outcome_certainty": failure.outcome_certainty().as_str(),
                    "request_id": request_id,
                    "trace_id": correlation.trace_id().as_str(),
                }
            }),
            InvocationStatus::Error,
            InvocationLevel::Error,
            "approved tool call failed",
            Some(failure.error_code().as_str()),
            Some(failure.stage()),
            Some(failure.error_code()),
            Some(failure.retryability()),
            Some(failure.outcome_certainty()),
            failure.upstream_status(),
            Value::Null,
        ),
    };

    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: upstream_status,
            error_kind,
            execution_stage,
            execution_error_code,
            retryability,
            outcome_certainty,
            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(|_| approved_completion_persistence_error(&correlation))?
    .ok_or_else(|| approved_completion_conflict_error(&correlation))
}

fn approved_completion_conflict_error(correlation: &CorrelationContext) -> Response {
    (
        StatusCode::CONFLICT,
        axum::Json(json!({
            "error": {
                "code": "approval_state_conflict",
                "stage": "mandatory_persistence",
                "retryability": "manual_reconcile",
                "outcome_certainty": "outcome_unknown",
                "request_id": correlation.request_id().as_str(),
                "trace_id": correlation.trace_id().as_str(),
            }
        })),
    )
        .into_response()
}

fn approved_completion_persistence_error(correlation: &CorrelationContext) -> Response {
    let failure = crank_core::ExecutionFailure::new(
        crank_core::ExecutionErrorCode::PersistenceUnavailable,
        correlation.clone(),
    )
    .with_dispatch_uncertainty();
    (
        StatusCode::SERVICE_UNAVAILABLE,
        axum::Json(json!({
            "error": {
                "code": failure.error_code().as_str(),
                "stage": failure.stage().as_str(),
                "retryability": failure.retryability().as_str(),
                "outcome_certainty": failure.outcome_certainty().as_str(),
                "request_id": failure.correlation().request_id().as_str(),
                "trace_id": failure.correlation().trace_id().as_str(),
            }
        })),
    )
        .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())
}