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, 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, } #[derive(Clone)] pub(super) struct ResolvedToolCall { tool: PublishedAgentTool, } struct ToolCallExecution { tool: PublishedAgentTool, arguments: Value, confirmation_token: Option, } #[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, secret_crypto: SecretCrypto, runtime: RuntimeExecutor, api_rate_limiter: RequestRateLimiter, coordination_store: Arc, 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, secret_crypto: SecretCrypto, runtime: RuntimeExecutor, api_rate_limiter: RequestRateLimiter, coordination_store: Arc, 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, secret_crypto: SecretCrypto, runtime: RuntimeExecutor, api_rate_limiter: RequestRateLimiter, coordination_store: Arc, 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, secret_crypto: SecretCrypto, runtime: RuntimeExecutor, api_rate_limiter: RequestRateLimiter, coordination_store: Arc, 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 { Json(json!({ "service": "mcp-server", "status": "ok" })) } fn spawn_session_cleanup(state: Arc) { 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>) -> 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, State(state): State>, 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, State(state): State>, Extension(request_context): Extension, headers: HeaderMap, Json(payload): Json, ) -> 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, State(state): State>, 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, State(state): State>, headers: HeaderMap, Json(payload): Json, ) -> Response { decide_approval_request( path, state, headers, payload, PlatformApiKeyScope::Deny, ApprovalRequestStatus::Denied, None, ) .await } async fn decide_approval_request( path: ApprovalRoutePath, state: Arc, headers: HeaderMap, payload: ApprovalDecisionPayload, required_scope: PlatformApiKeyScope, status: ApprovalRequestStatus, execution_request_id: Option, ) -> 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, 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, 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, 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, State(state): State>, 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, _)> }); sse_response( StatusCode::OK, stream, Some(&session_id), Some(&session.protocol_version), ) } async fn mcp_delete( Path(path): Path, State(state): State>, 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, State(state): State>, Extension(request_context): Extension, headers: HeaderMap, Json(message): Json, ) -> 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(""), "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, session: &SessionState, message: &Value, response_mode: ResponseMode, credential: &VerifiedMachineCredential, resolved: ResolvedToolCall, arguments: Value, confirmation_token: Option, 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, workspace_id: &crank_core::WorkspaceId, execution_config: &crank_core::ExecutionConfig, ) -> Result, 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, 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 { 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, 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, session: &SessionState, message: &Value, response_mode: ResponseMode, tool: &PublishedAgentTool, arguments: &Value, transport_request_id: &str, ) -> Option { 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, session: &SessionState, message: &Value, response_mode: ResponseMode, tool: &PublishedAgentTool, arguments: &Value, transport_request_id: &str, ) -> Option { 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, 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, 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, path: &AgentRoutePath, headers: &HeaderMap, message: &Value, ) -> Result { 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 { 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 { 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;