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_registry::{ApprovalRequestRecord, FinishApprovalRequest};
use crank_runtime::{RuntimeExecutionRequest, RuntimeRequestContext};
use serde_json::json;
use time::OffsetDateTime;
use tracing::warn;

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

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>) {
    for _ in 0..32 {
        let now = OffsetDateTime::now_utc();
        let approval = match state
            .registry
            .claim_next_recoverable_approval_request(
                now,
                now - RECOVERY_GRACE,
                now - EXECUTION_LEASE,
            )
            .await
        {
            Ok(Some(approval)) => approval,
            Ok(None) => break,
            Err(error) => {
                warn!(error = %error, "approval recovery query failed");
                break;
            }
        };
        let Some(path) = approval_agent_path(state, &approval).await else {
            continue;
        };
        if execute_approved_request(state, &path, approval)
            .await
            .is_err()
        {
            warn!("recovered approval execution did not finish");
        }
    }
}

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(error) => {
            warn!(error = %error, "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(error) => {
            warn!(error = %error, "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,
) -> Result<ApprovalRequestRecord, Response> {
    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_request_id(approval.approval.id.as_str().to_owned())
            .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": error.to_string(),
                    }
                }),
                InvocationStatus::Error,
                InvocationLevel::Error,
                "approved tool call failed",
                Some(runtime_error_code(&error)),
            ),
        };

    if let Err(error) = persist_invocation(
        state,
        &tool,
        InvocationRecord {
            request_id: Some(approval.approval.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
    {
        warn!(error = %error, "approved invocation log write failed");
    }

    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> {
    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())
}