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) { 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) { 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) { 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, approval: &ApprovalRequestRecord, ) -> Option { 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, path: &AgentRoutePath, approval: ApprovalRequestRecord, correlation: Option<&CorrelationContext>, ) -> Result { 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, approval: &ApprovalRequestRecord, ) -> Result { 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()) }