use std::{ collections::BTreeMap, convert::Infallible, sync::Arc, time::{Duration, Instant}, }; use axum::{ Json, Router, extract::{Path, State}, http::{ HeaderMap, HeaderValue, StatusCode, header::{self, ACCEPT, AUTHORIZATION, HeaderName, RETRY_AFTER}, }, response::{ IntoResponse, Response, sse::{Event, KeepAlive, Sse}, }, routing::get, }; use base64::{Engine as _, engine::general_purpose::URL_SAFE_NO_PAD}; use crank_core::{ AsyncJobHandle, AsyncJobId, AuthProfile, InvocationLevel, InvocationLog, InvocationLogId, InvocationSource, InvocationStatus, JobStatus, PlatformApiKeyScope, SecretId, StreamSession, StreamSessionId, StreamStatus, }; use crank_registry::{ CreateAsyncJobRequest, CreateInvocationLogRequest, CreateStreamSessionRequest, PostgresRegistry, PublishedAgentTool, UpdateAsyncJobStatusRequest, UpdateStreamSessionStateRequest, }; use crank_runtime::{ RateLimitRejection, RequestRateLimiter, ResolvedAuth, RuntimeError, RuntimeExecutor, RuntimeOperation, RuntimeRequestContext, SecretCrypto, }; use futures_util::stream; use serde::{Deserialize, Serialize}; use serde_json::{Value, json}; use sha2::{Digest, Sha256}; use time::{OffsetDateTime, format_description::well_known::Rfc3339}; use tracing::info; use crate::{ auth::{SharedMachineCredentialVerifier, VerifiedMachineCredential}, catalog::PublishedToolCatalog, jsonrpc::{ CURRENT_PROTOCOL_VERSION, DEFAULT_PROTOCOL_VERSION, is_notification, is_request, is_response, jsonrpc_error, jsonrpc_result, method_name, negotiated_protocol_version, params, request_id, }, session::{SessionState, SharedSessionStore}, }; const HEADER_MCP_SESSION_ID: &str = "MCP-Session-Id"; const HEADER_MCP_PROTOCOL_VERSION: &str = "MCP-Protocol-Version"; const HEADER_X_REQUEST_ID: HeaderName = HeaderName::from_static("x-request-id"); const MAX_REQUEST_ID_LEN: usize = 128; const TRANSPORT_SESSION_TTL_MS: u64 = 86_400_000; #[derive(Clone, Copy)] enum ResponseMode { Json, Sse, } #[derive(Clone)] pub struct AppState { registry: PostgresRegistry, catalog: PublishedToolCatalog, runtime: RuntimeExecutor, api_rate_limiter: RequestRateLimiter, secret_crypto: SecretCrypto, sessions: SharedSessionStore, credential_verifier: SharedMachineCredentialVerifier, allowed_origins: AllowedOrigins, } #[derive(Clone)] pub struct AllowedOrigins { public_origin: Option, } #[derive(Debug, Serialize, Deserialize)] struct InitializeParams { #[serde(rename = "protocolVersion")] protocol_version: String, } #[derive(Debug, Serialize, Deserialize)] struct ToolCallParams { name: String, #[serde(default)] arguments: Value, } #[derive(Clone, Copy)] enum GeneratedToolKind { Base, SessionStart, SessionPoll, SessionStop, AsyncJobStart, AsyncJobStatus, AsyncJobResult, AsyncJobCancel, } #[derive(Clone)] struct ResolvedToolCall { tool: PublishedAgentTool, kind: GeneratedToolKind, } #[derive(Debug, Deserialize)] struct SessionControlArgs { session_id: String, } #[derive(Debug, Deserialize)] struct AsyncJobControlArgs { job_id: String, } #[derive(Clone, Debug, Deserialize)] struct AgentRoutePath { workspace_slug: String, agent_slug: String, } pub fn build_app( registry: PostgresRegistry, refresh_interval: Duration, public_base_url: Option, secret_crypto: SecretCrypto, runtime: RuntimeExecutor, api_rate_limiter: RequestRateLimiter, sessions: SharedSessionStore, credential_verifier: SharedMachineCredentialVerifier, ) -> Router { let state = Arc::new(AppState { registry: registry.clone(), catalog: PublishedToolCatalog::new(registry, refresh_interval), runtime, api_rate_limiter, secret_crypto, sessions, credential_verifier, allowed_origins: AllowedOrigins::new(public_base_url), }); Router::new() .route("/health", get(health)) .route( "/v1/{workspace_slug}/{agent_slug}", get(mcp_get).post(mcp_post).delete(mcp_delete), ) .with_state(state) } async fn health() -> Json { Json(json!({ "service": "mcp-server", "status": "ok" })) } 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_machine_access(&state, &path, &headers, PlatformApiKeyScope::Read).await { return status.into_response(); } if let Err(rejection) = enforce_transport_rate_limit(&state, &path, &headers) { 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(); } sse_response( StatusCode::OK, stream::pending::>(), 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_machine_access(&state, &path, &headers, PlatformApiKeyScope::Read).await { return status.into_response(); } if let Err(rejection) = enforce_transport_rate_limit(&state, &path, &headers) { 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>, headers: HeaderMap, Json(message): Json, ) -> Response { let transport_request_id = resolve_request_id(&headers); info!( 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) => 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); } }; if let Err(rejection) = enforce_post_rate_limit(&state, &path, &headers) { return with_request_id_header( rate_limited_jsonrpc_response(&message, response_mode, &protocol_version, rejection), &transport_request_id, ); } if let Some(session_id) = headers.get(HEADER_MCP_SESSION_ID) { if 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 { if 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, }; if let Err(status) = require_machine_access(&state, &path, &headers, required_scope).await { return with_request_id_header(status.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 .list_tools(&session.workspace_slug, &session.agent_slug) .await { Ok(tools) => { let definitions = tools.iter().flat_map(tool_definitions).collect::>(); 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 .list_tools(&session.workspace_slug, &session.agent_slug) .await { Ok(tools) => match resolve_generated_tool(&tools, &tool_call_params.name) { Some(resolved) => { handle_tool_call( state.clone(), &session, &message, response_mode, resolved, arguments, &transport_request_id, ) .await } None => transport_response( StatusCode::OK, jsonrpc_error( request_id(&message), -32602, format!("tool {} was not found", tool_call_params.name), ), response_mode, None, Some(&session.protocol_version), ), }, 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), ), }; with_request_id_header(response, &transport_request_id) } async fn handle_tool_call( state: Arc, session: &SessionState, message: &Value, response_mode: ResponseMode, resolved: ResolvedToolCall, arguments: Value, transport_request_id: &str, ) -> Response { match resolved.kind { GeneratedToolKind::Base => { handle_base_tool_call( state, session, message, response_mode, resolved.tool, arguments, transport_request_id, ) .await } GeneratedToolKind::SessionStart => { handle_session_start_call( state, session, message, response_mode, resolved.tool, arguments, transport_request_id, ) .await } GeneratedToolKind::SessionPoll => { handle_session_poll_call( state, session, message, response_mode, resolved.tool, arguments, transport_request_id, ) .await } GeneratedToolKind::SessionStop => { handle_session_stop_call( state, session, message, response_mode, resolved.tool, arguments, ) .await } GeneratedToolKind::AsyncJobStart => { handle_async_job_start_call( state, session, message, response_mode, resolved.tool, arguments, transport_request_id, ) .await } GeneratedToolKind::AsyncJobStatus => { handle_async_job_status_call( state, session, message, response_mode, resolved.tool, arguments, ) .await } GeneratedToolKind::AsyncJobResult => { handle_async_job_result_call( state, session, message, response_mode, resolved.tool, arguments, ) .await } GeneratedToolKind::AsyncJobCancel => { handle_async_job_cancel_call( state, session, message, response_mode, resolved.tool, arguments, ) .await } } } #[derive(Debug, Serialize, Deserialize)] struct StoredSessionState { input: Value, summary: Value, items: Vec, next_index: usize, batch_size: usize, } async fn resolve_operation_auth( state: &Arc, workspace_id: &crank_core::WorkspaceId, execution_config: &crank_core::ExecutionConfig, ) -> Result, RuntimeError> { resolve_runtime_auth_for_task( &state.registry, &state.secret_crypto, workspace_id, execution_config, ) .await } 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 = 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 = 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 = 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, )?; 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, tool: PublishedAgentTool, arguments: Value, transport_request_id: &str, ) -> Response { let operation = runtime_operation(&tool); let runtime_request_context = RuntimeRequestContext::from_request_id(transport_request_id); let request_preview = build_request_preview(&state.runtime, &operation, &arguments); let started_at = Instant::now(); let is_window_mode = matches!( operation .execution_config .streaming .as_ref() .map(|streaming| streaming.mode), Some(crank_core::ExecutionMode::Window) ); let resolved_auth = resolve_operation_auth(&state, &tool.workspace_id, &operation.execution_config).await; let result = match resolved_auth { Ok(resolved_auth) if is_window_mode => state .runtime .execute_window_with_auth_and_context( &operation, &arguments, resolved_auth.as_ref(), Some(&runtime_request_context), ) .await .map(|output| { json!({ "summary": output.summary, "items": output.items, "cursor": output.cursor, "window_complete": output.window_complete, "truncated": output.truncated, "has_more": output.has_more, }) }), Ok(resolved_auth) => { state .runtime .execute_with_auth_and_context( &operation, &arguments, resolved_auth.as_ref(), Some(&runtime_request_context), ) .await } Err(error) => Err(error), }; match result { Ok(output) => { let _ = 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) => { let _ = 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, runtime_error_code(&error), error.to_string(), runtime_error_context(&error), ) } } } async fn handle_session_start_call( state: Arc, session: &SessionState, message: &Value, response_mode: ResponseMode, tool: PublishedAgentTool, arguments: Value, transport_request_id: &str, ) -> Response { let runtime_operation = runtime_operation(&tool); let runtime_request_context = RuntimeRequestContext::from_request_id(transport_request_id); let request_preview = build_request_preview(&state.runtime, &runtime_operation, &arguments); let started_at = Instant::now(); let Some(streaming) = runtime_operation.execution_config.streaming.as_ref() else { return tool_error_response( message, response_mode, &session.protocol_version, "streaming_config_error", "streaming config is required for session tools".to_owned(), None, ); }; let resolved_auth = resolve_operation_auth( &state, &tool.workspace_id, &runtime_operation.execution_config, ) .await; match match resolved_auth { Ok(resolved_auth) => { state .runtime .execute_session_seed_with_auth_and_context( &runtime_operation, &arguments, resolved_auth.as_ref(), Some(&runtime_request_context), ) .await } Err(error) => Err(error), } { Ok(seed) => { let batch_size = streaming.max_items.unwrap_or(10).max(1) as usize; let preview_count = seed.items.len().min(batch_size); let preview_items = seed.items[..preview_count].to_vec(); let next_index = preview_count; let session_id = StreamSessionId::new(format!("sess_{}", uuid::Uuid::now_v7().simple())); let now = OffsetDateTime::now_utc(); let expires_at = add_millis(now, streaming.max_session_lifetime_ms.unwrap_or(60_000)); let session_record = StreamSession { id: session_id.clone(), workspace_id: tool.workspace_id.clone(), agent_id: Some(tool.agent_id.clone()), operation_id: tool.operation.id.clone(), protocol: tool.operation.protocol, mode: crank_core::ExecutionMode::Session, status: StreamStatus::Running, cursor: (next_index < seed.items.len()).then(|| json!(next_index)), state: json!(StoredSessionState { input: arguments.clone(), summary: seed.summary.clone(), items: seed.items.clone(), next_index, batch_size, }), expires_at, last_poll_at: None, created_at: now, closed_at: None, }; if let Err(error) = state .registry .create_stream_session(CreateStreamSessionRequest { session: &session_record, }) .await { return internal_jsonrpc_error(message, error); } let output = json!({ "session_id": session_id.as_str(), "status": "running", "expires_at": format_rfc3339(expires_at), "poll_after_ms": streaming.poll_interval_ms.unwrap_or(1000), "preview": { "summary": seed.summary, "items": preview_items, } }); let _ = persist_invocation( &state, &tool, InvocationRecord { request_id: Some(transport_request_id), tool_name: streaming .tool_family .start_tool_name .as_deref() .unwrap_or(&tool.tool_name), status: InvocationStatus::Ok, level: InvocationLevel::Info, message: "stream session started", 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) => { let _ = persist_invocation( &state, &tool, InvocationRecord { request_id: Some(transport_request_id), tool_name: streaming .tool_family .start_tool_name .as_deref() .unwrap_or(&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, runtime_error_code(&error), error.to_string(), runtime_error_context(&error), ) } } } async fn handle_session_poll_call( state: Arc, session: &SessionState, message: &Value, response_mode: ResponseMode, tool: PublishedAgentTool, arguments: Value, transport_request_id: &str, ) -> Response { let control: SessionControlArgs = match serde_json::from_value(arguments.clone()) { 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 Some(streaming) = tool.operation.execution_config.streaming.as_ref() else { return tool_error_response( message, response_mode, &session.protocol_version, "streaming_config_error", "streaming config is required for session tools".to_owned(), None, ); }; let now = OffsetDateTime::now_utc(); let loaded = match state .registry .get_stream_session(&StreamSessionId::new(control.session_id.clone())) .await { Ok(Some(session_record)) => session_record, Ok(None) => { return tool_error_response( message, response_mode, &session.protocol_version, "stream_session_not_found", format!("stream session {} was not found", control.session_id), None, ); } Err(error) => return internal_jsonrpc_error(message, error), }; if !stream_session_belongs_to_tool(&loaded, &tool) { return tool_error_response( message, response_mode, &session.protocol_version, "stream_session_not_found", format!("stream session {} was not found", control.session_id), None, ); } if loaded.is_expired(now) { let _ = state .registry .update_stream_session_state(UpdateStreamSessionStateRequest { session_id: &loaded.id, current_status: loaded.status, next_status: StreamStatus::Expired, cursor: loaded.cursor.as_ref(), state: &loaded.state, expires_at: None, last_poll_at: Some(&now), closed_at: Some(&now), }) .await; return tool_error_response( message, response_mode, &session.protocol_version, "stream_session_expired", format!("stream session {} has expired", control.session_id), None, ); } if !loaded.can_poll(now) { return tool_error_response( message, response_mode, &session.protocol_version, "stream_session_not_found", format!("stream session {} is not running", control.session_id), None, ); } let poll_after_ms = streaming.poll_interval_ms.unwrap_or(1_000); let remaining_delay_ms = loaded.remaining_poll_delay_ms(now, poll_after_ms); if remaining_delay_ms > 0 { return tool_error_response( message, response_mode, &session.protocol_version, "stream_session_poll_rate_limited", format!( "stream session {} must wait before the next poll", control.session_id ), Some(json!({ "poll_after_ms": remaining_delay_ms })), ); } let mut state_payload: StoredSessionState = match serde_json::from_value(loaded.state.clone()) { Ok(value) => value, Err(error) => { return tool_error_response( message, response_mode, &session.protocol_version, "streaming_payload_error", error.to_string(), None, ); } }; let next_end = (state_payload.next_index + state_payload.batch_size).min(state_payload.items.len()); let items = state_payload.items[state_payload.next_index..next_end].to_vec(); state_payload.next_index = next_end; let has_more = state_payload.next_index < state_payload.items.len(); let next_status = if has_more { StreamStatus::Running } else { StreamStatus::Stopped }; let cursor = has_more.then(|| json!(state_payload.next_index)); let updated = match state .registry .update_stream_session_state(UpdateStreamSessionStateRequest { session_id: &loaded.id, current_status: loaded.status, next_status, cursor: cursor.as_ref(), state: &json!(state_payload), expires_at: Some(&add_millis( now, streaming.max_session_lifetime_ms.unwrap_or(60_000), )), last_poll_at: Some(&now), closed_at: matches!(next_status, StreamStatus::Stopped).then_some(&now), }) .await { Ok(session_record) => session_record, Err(error) => return internal_jsonrpc_error(message, error), }; let output = json!({ "session_id": updated.id.as_str(), "status": serialize_stream_status(updated.status), "expires_at": format_rfc3339(updated.expires_at), "summary": serde_json::from_value::(updated.state.clone()).map(|value| value.summary).unwrap_or(Value::Null), "items": items, "cursor": updated.cursor, "has_more": updated.status == StreamStatus::Running, }); let _ = persist_invocation( &state, &tool, InvocationRecord { request_id: Some(transport_request_id), tool_name: streaming .tool_family .poll_tool_name .as_deref() .unwrap_or(&tool.tool_name), status: InvocationStatus::Ok, level: InvocationLevel::Info, message: "stream session polled", status_code: None, error_kind: None, duration: Duration::ZERO, request_preview: json!({ "session_id": control.session_id }), response_preview: output.clone(), }, ) .await; success_tool_response(message, response_mode, &session.protocol_version, output) } async fn handle_session_stop_call( state: Arc, session: &SessionState, message: &Value, response_mode: ResponseMode, tool: PublishedAgentTool, arguments: Value, ) -> Response { let control: SessionControlArgs = match serde_json::from_value(arguments.clone()) { 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 loaded = match state .registry .get_stream_session(&StreamSessionId::new(control.session_id.clone())) .await { Ok(Some(session_record)) => session_record, Ok(None) => { return tool_error_response( message, response_mode, &session.protocol_version, "stream_session_not_found", format!("stream session {} was not found", control.session_id), None, ); } Err(error) => return internal_jsonrpc_error(message, error), }; if !stream_session_belongs_to_tool(&loaded, &tool) { return tool_error_response( message, response_mode, &session.protocol_version, "stream_session_not_found", format!("stream session {} was not found", control.session_id), None, ); } match state .registry .close_stream_session( &StreamSessionId::new(control.session_id.clone()), &OffsetDateTime::now_utc(), ) .await { Ok(()) => success_tool_response( message, response_mode, &session.protocol_version, json!({ "session_id": control.session_id, "status": "stopped" }), ), Err(error) => internal_jsonrpc_error(message, error), } } async fn handle_async_job_start_call( state: Arc, session: &SessionState, message: &Value, response_mode: ResponseMode, tool: PublishedAgentTool, arguments: Value, transport_request_id: &str, ) -> Response { let operation = runtime_operation(&tool); let runtime_request_context = RuntimeRequestContext::from_request_id(transport_request_id); let request_preview = build_request_preview(&state.runtime, &operation, &arguments); let Some(streaming) = tool.operation.execution_config.streaming.as_ref() else { return tool_error_response( message, response_mode, &session.protocol_version, "streaming_config_error", "streaming config is required for async job tools".to_owned(), None, ); }; let now = OffsetDateTime::now_utc(); let job = AsyncJobHandle { id: AsyncJobId::new(format!("job_{}", uuid::Uuid::now_v7().simple())), workspace_id: tool.workspace_id.clone(), agent_id: Some(tool.agent_id.clone()), operation_id: tool.operation.id.clone(), status: JobStatus::Running, progress: json!({ "pct": 0 }), result: None, error: None, expires_at: Some(add_millis( now, streaming.max_session_lifetime_ms.unwrap_or(300_000), )), last_poll_at: None, created_at: now, updated_at: now, finished_at: None, }; if let Err(error) = state .registry .create_async_job(CreateAsyncJobRequest { job: &job }) .await { return internal_jsonrpc_error(message, error); } let registry = state.registry.clone(); let secret_crypto = state.secret_crypto.clone(); let tool_for_task = tool.clone(); let arguments_for_task = arguments.clone(); let request_context_for_task = runtime_request_context.clone(); let job_id = job.id.clone(); tokio::spawn(async move { let runtime = RuntimeExecutor::new(); let task_operation = runtime_operation(&tool_for_task); let resolved_auth = resolve_runtime_auth_for_task( ®istry, &secret_crypto, &tool_for_task.workspace_id, &task_operation.execution_config, ) .await; let result = match resolved_auth { Ok(resolved_auth) => { runtime .execute_with_auth_and_context( &task_operation, &arguments_for_task, resolved_auth.as_ref(), Some(&request_context_for_task), ) .await } Err(error) => Err(error), }; let finished_at = OffsetDateTime::now_utc(); let update_result = match result { Ok(output) => { registry .update_async_job_status(UpdateAsyncJobStatusRequest { job_id: &job_id, current_status: JobStatus::Running, next_status: JobStatus::Completed, progress: &json!({ "pct": 100 }), result: Some(&output), error: None, expires_at: None, updated_at: &finished_at, finished_at: Some(&finished_at), }) .await } Err(error) => { registry .update_async_job_status(UpdateAsyncJobStatusRequest { job_id: &job_id, current_status: JobStatus::Running, next_status: JobStatus::Failed, progress: &json!({ "pct": 100 }), result: None, error: Some(&json!({ "code": runtime_error_code(&error), "message": error.to_string(), "context": runtime_error_context(&error) })), expires_at: None, updated_at: &finished_at, finished_at: Some(&finished_at), }) .await } }; let _ = update_result; }); let output = json!({ "job_id": job.id.as_str(), "status": "running", "progress": job.progress, }); let _ = persist_invocation( &state, &tool, InvocationRecord { request_id: Some(transport_request_id), tool_name: streaming .tool_family .start_tool_name .as_deref() .unwrap_or(&tool.tool_name), status: InvocationStatus::Ok, level: InvocationLevel::Info, message: "async job started", status_code: None, error_kind: None, duration: Duration::ZERO, request_preview, response_preview: output.clone(), }, ) .await; success_tool_response(message, response_mode, &session.protocol_version, output) } async fn handle_async_job_status_call( state: Arc, session: &SessionState, message: &Value, response_mode: ResponseMode, tool: PublishedAgentTool, arguments: Value, ) -> Response { let control: AsyncJobControlArgs = match serde_json::from_value(arguments.clone()) { 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 job = match state .registry .get_async_job(&AsyncJobId::new(control.job_id.clone())) .await { Ok(Some(job)) => job, Ok(None) => { return tool_error_response( message, response_mode, &session.protocol_version, "async_job_not_found", format!("async job {} was not found", control.job_id), None, ); } Err(error) => return internal_jsonrpc_error(message, error), }; if !async_job_belongs_to_tool(&job, &tool) { return tool_error_response( message, response_mode, &session.protocol_version, "async_job_not_found", format!("async job {} was not found", control.job_id), None, ); } let Some(streaming) = tool.operation.execution_config.streaming.as_ref() else { return tool_error_response( message, response_mode, &session.protocol_version, "streaming_config_error", "streaming config is required for async job tools".to_owned(), None, ); }; let job = if matches!(job.status, JobStatus::Completed) { job } else { let now = OffsetDateTime::now_utc(); let poll_after_ms = streaming.poll_interval_ms.unwrap_or(1_000); let remaining_delay_ms = job.remaining_poll_delay_ms(now, poll_after_ms); if remaining_delay_ms > 0 { return tool_error_response( message, response_mode, &session.protocol_version, "async_job_poll_rate_limited", format!( "async job {} must wait before the next poll", control.job_id ), Some(json!({ "poll_after_ms": remaining_delay_ms, })), ); } match state.registry.touch_async_job_poll(&job.id, &now).await { Ok(job) => job, Err(error) => return internal_jsonrpc_error(message, error), } }; success_tool_response( message, response_mode, &session.protocol_version, json!({ "job_id": job.id.as_str(), "status": serialize_job_status(job.status), "progress": job.progress, "updated_at": job.updated_at, "finished_at": job.finished_at, }), ) } async fn handle_async_job_result_call( state: Arc, session: &SessionState, message: &Value, response_mode: ResponseMode, tool: PublishedAgentTool, arguments: Value, ) -> Response { let control: AsyncJobControlArgs = match serde_json::from_value(arguments.clone()) { 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 job = match state .registry .get_async_job(&AsyncJobId::new(control.job_id.clone())) .await { Ok(Some(job)) => job, Ok(None) => { return tool_error_response( message, response_mode, &session.protocol_version, "async_job_not_found", format!("async job {} was not found", control.job_id), None, ); } Err(error) => return internal_jsonrpc_error(message, error), }; if !async_job_belongs_to_tool(&job, &tool) { return tool_error_response( message, response_mode, &session.protocol_version, "async_job_not_found", format!("async job {} was not found", control.job_id), None, ); } let Some(streaming) = tool.operation.execution_config.streaming.as_ref() else { return tool_error_response( message, response_mode, &session.protocol_version, "streaming_config_error", "streaming config is required for async job tools".to_owned(), None, ); }; let job = if matches!(job.status, JobStatus::Completed) { job } else { let now = OffsetDateTime::now_utc(); let poll_after_ms = streaming.poll_interval_ms.unwrap_or(1_000); let remaining_delay_ms = job.remaining_poll_delay_ms(now, poll_after_ms); if remaining_delay_ms > 0 { return tool_error_response( message, response_mode, &session.protocol_version, "async_job_poll_rate_limited", format!( "async job {} must wait before the next poll", control.job_id ), Some(json!({ "poll_after_ms": remaining_delay_ms, })), ); } match state.registry.touch_async_job_poll(&job.id, &now).await { Ok(job) => job, Err(error) => return internal_jsonrpc_error(message, error), } }; match job.status { JobStatus::Completed => success_tool_response( message, response_mode, &session.protocol_version, job.result.unwrap_or(Value::Null), ), JobStatus::Failed => tool_error_response( message, response_mode, &session.protocol_version, "async_job_failed", job.error .as_ref() .and_then(|value| value.get("message")) .and_then(Value::as_str) .map(ToOwned::to_owned) .unwrap_or_else(|| "async job failed".to_owned()), job.error .as_ref() .and_then(|value| value.get("context")) .filter(|value| !value.is_null()) .cloned(), ), JobStatus::Cancelled => tool_error_response( message, response_mode, &session.protocol_version, "async_job_cancelled", "async job was cancelled".to_owned(), None, ), _ => tool_error_response( message, response_mode, &session.protocol_version, "async_job_not_ready", "async job result is not ready".to_owned(), None, ), } } async fn handle_async_job_cancel_call( state: Arc, session: &SessionState, message: &Value, response_mode: ResponseMode, _tool: PublishedAgentTool, arguments: Value, ) -> Response { let control: AsyncJobControlArgs = match serde_json::from_value(arguments.clone()) { 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), ); } }; match state .registry .get_async_job(&AsyncJobId::new(control.job_id.clone())) .await { Ok(Some(job)) => { if !async_job_belongs_to_tool(&job, &_tool) { return tool_error_response( message, response_mode, &session.protocol_version, "async_job_not_found", format!("async job {} was not found", control.job_id), None, ); } } Ok(None) => { return tool_error_response( message, response_mode, &session.protocol_version, "async_job_not_found", format!("async job {} was not found", control.job_id), None, ); } Err(error) => return internal_jsonrpc_error(message, error), }; match state .registry .cancel_async_job( &AsyncJobId::new(control.job_id.clone()), &OffsetDateTime::now_utc(), ) .await { Ok(()) => success_tool_response( message, response_mode, &session.protocol_version, json!({ "job_id": control.job_id, "status": "cancelled" }), ), Err(error) => internal_jsonrpc_error(message, error), } } 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 session_id = match state .sessions .create( protocol_version, &path.workspace_slug, &path.agent_slug, 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) } async fn require_machine_access( state: &Arc, path: &AgentRoutePath, headers: &HeaderMap, required_scope: PlatformApiKeyScope, ) -> Result { let secret = bearer_token(headers).ok_or(StatusCode::UNAUTHORIZED)?; let credential = resolve_machine_credential(state, path, secret).await?; if !allows_scope(&credential.scopes, required_scope) { return Err(StatusCode::FORBIDDEN); } Ok(credential) } async fn resolve_machine_credential( state: &Arc, path: &AgentRoutePath, token: &str, ) -> Result { if let Some(credential) = verify_static_agent_key(state, path, token).await? { return Ok(credential); } state .credential_verifier .verify_bearer_token(&path.workspace_slug, &path.agent_slug, token) .await .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)? .ok_or(StatusCode::UNAUTHORIZED) } async fn verify_static_agent_key( state: &Arc, path: &AgentRoutePath, secret: &str, ) -> Result, StatusCode> { let secret_hash = hash_access_secret(secret); let Some(api_key) = state .registry .get_platform_api_key_by_secret_for_agent_slug( &path.workspace_slug, &path.agent_slug, &secret_hash, ) .await .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)? else { return Ok(None); }; let used_at = OffsetDateTime::now_utc(); state .registry .touch_platform_api_key(&api_key.api_key.workspace_id, &api_key.api_key.id, &used_at) .await .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?; Ok(Some(VerifiedMachineCredential { machine_access_mode: crank_core::MachineAccessMode::StaticAgentKey, max_security_level: crank_core::OperationSecurityLevel::Standard, scopes: api_key.api_key.scopes, })) } fn bearer_token(headers: &HeaderMap) -> Option<&str> { let value = headers.get(AUTHORIZATION)?.to_str().ok()?; let (scheme, token) = value.split_once(' ')?; if !scheme.eq_ignore_ascii_case("Bearer") || token.is_empty() { return None; } Some(token) } fn enforce_post_rate_limit( state: &Arc, path: &AgentRoutePath, headers: &HeaderMap, ) -> Result<(), RateLimitRejection> { enforce_transport_rate_limit(state, path, headers) } fn enforce_transport_rate_limit( state: &Arc, path: &AgentRoutePath, headers: &HeaderMap, ) -> Result<(), RateLimitRejection> { let key = rate_limit_key(path, headers); state.api_rate_limiter.check(&key) } fn rate_limit_key(path: &AgentRoutePath, headers: &HeaderMap) -> String { if let Ok(Some(session_id)) = session_id_from_headers(headers) { return format!( "session:{}:{}:{}", path.workspace_slug, path.agent_slug, session_id ); } if let Some(secret) = bearer_token(headers) { return format!("api_key:{}", hash_access_secret(secret)); } format!("workspace:{}:anonymous", path.workspace_slug) } fn rate_limited_jsonrpc_response( message: &Value, response_mode: ResponseMode, protocol_version: &str, rejection: RateLimitRejection, ) -> Response { let payload = json!({ "jsonrpc": "2.0", "id": request_id(message), "error": { "code": -32029, "message": "request rate limit exceeded", "data": { "code": "request_rate_limited", "retry_after_ms": rejection.retry_after_ms, } } }); let mut response = transport_response( StatusCode::TOO_MANY_REQUESTS, payload, response_mode, None, Some(protocol_version), ); let retry_after_seconds = rejection.retry_after_ms.div_ceil(1000); if let Ok(value) = HeaderValue::from_str(&retry_after_seconds.to_string()) { response.headers_mut().insert(RETRY_AFTER, value); } response } fn rate_limited_status_response(rejection: RateLimitRejection) -> Response { let mut response = StatusCode::TOO_MANY_REQUESTS.into_response(); let retry_after_seconds = rejection.retry_after_ms.div_ceil(1000); if let Ok(value) = HeaderValue::from_str(&retry_after_seconds.to_string()) { response.headers_mut().insert(RETRY_AFTER, value); } response } fn allows_scope(scopes: &[PlatformApiKeyScope], required_scope: PlatformApiKeyScope) -> bool { match required_scope { PlatformApiKeyScope::Read => scopes.iter().any(|scope| { matches!( scope, PlatformApiKeyScope::Read | PlatformApiKeyScope::Write | PlatformApiKeyScope::Deploy ) }), PlatformApiKeyScope::Write => scopes.iter().any(|scope| { matches!( scope, PlatformApiKeyScope::Write | PlatformApiKeyScope::Deploy ) }), PlatformApiKeyScope::Deploy => scopes .iter() .any(|scope| matches!(scope, PlatformApiKeyScope::Deploy)), } } fn validate_origin( allowed_origins: &AllowedOrigins, headers: &HeaderMap, ) -> Result<(), StatusCode> { let Some(origin) = headers.get(header::ORIGIN) else { return Ok(()); }; let Ok(origin) = origin.to_str() else { return Err(StatusCode::FORBIDDEN); }; if allowed_origins.is_allowed(origin) { return Ok(()); } Err(StatusCode::FORBIDDEN) } fn negotiate_post_response_mode(headers: &HeaderMap) -> Result { let Some(accept) = headers.get(ACCEPT) else { return Err(StatusCode::BAD_REQUEST); }; let Ok(accept) = accept.to_str() else { return Err(StatusCode::BAD_REQUEST); }; let mut saw_json = false; let mut saw_sse = false; let mut preferred = None; for part in accept.split(',') { let media_type = part .split(';') .next() .unwrap_or_default() .trim() .to_ascii_lowercase(); match media_type.as_str() { "application/json" => { saw_json = true; if preferred.is_none() { preferred = Some(ResponseMode::Json); } } "text/event-stream" => { saw_sse = true; if preferred.is_none() { preferred = Some(ResponseMode::Sse); } } _ => {} } } if saw_json || saw_sse { return preferred.ok_or(StatusCode::NOT_ACCEPTABLE); } Err(StatusCode::NOT_ACCEPTABLE) } fn validate_get_accept_header(headers: &HeaderMap) -> Result<(), StatusCode> { let Some(accept) = headers.get(ACCEPT) else { return Err(StatusCode::BAD_REQUEST); }; let Ok(accept) = accept.to_str() else { return Err(StatusCode::BAD_REQUEST); }; if accept .split(',') .map(|part| { part.split(';') .next() .unwrap_or_default() .trim() .to_ascii_lowercase() }) .any(|media_type| media_type == "text/event-stream") { return Ok(()); } Err(StatusCode::NOT_ACCEPTABLE) } fn stream_session_belongs_to_tool( session_record: &StreamSession, tool: &PublishedAgentTool, ) -> bool { session_record.workspace_id == tool.workspace_id && session_record.agent_id.as_ref() == Some(&tool.agent_id) && session_record.operation_id == tool.operation.id } fn async_job_belongs_to_tool(job: &AsyncJobHandle, tool: &PublishedAgentTool) -> bool { job.workspace_id == tool.workspace_id && job.agent_id.as_ref() == Some(&tool.agent_id) && job.operation_id == tool.operation.id } fn protocol_version_from_headers(headers: &HeaderMap) -> Result { let Some(version) = headers.get(HEADER_MCP_PROTOCOL_VERSION) else { return Ok(DEFAULT_PROTOCOL_VERSION.to_owned()); }; let Ok(version) = version.to_str() else { return Err(StatusCode::BAD_REQUEST); }; if negotiated_protocol_version(version).is_some() { return Ok(version.to_owned()); } Err(StatusCode::BAD_REQUEST) } fn validate_session_protocol_version( headers: &HeaderMap, negotiated_session_version: &str, ) -> Result<(), StatusCode> { let Some(version) = headers.get(HEADER_MCP_PROTOCOL_VERSION) else { return Ok(()); }; let Ok(version) = version.to_str() else { return Err(StatusCode::BAD_REQUEST); }; if negotiated_protocol_version(version).is_none() { return Err(StatusCode::BAD_REQUEST); } if version != negotiated_session_version { return Err(StatusCode::BAD_REQUEST); } Ok(()) } fn session_id_from_headers(headers: &HeaderMap) -> Result, StatusCode> { let Some(session_id) = headers.get(HEADER_MCP_SESSION_ID) else { return Ok(None); }; let Ok(session_id) = session_id.to_str() else { return Err(StatusCode::BAD_REQUEST); }; Ok(Some(session_id.to_owned())) } 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), ) } 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, "grpc": prepared.grpc.unwrap_or(Value::Null), "variables": prepared.variables.unwrap_or(Value::Null), "body": prepared.body.unwrap_or(Value::Null) }), Err(_) => Value::Null, } } struct InvocationRecord<'a> { request_id: Option<&'a str>, tool_name: &'a str, status: InvocationStatus, level: InvocationLevel, message: &'a str, status_code: Option, error_kind: Option<&'a str>, duration: Duration, request_preview: Value, response_preview: Value, } async fn persist_invocation( state: &Arc, tool: &PublishedAgentTool, record: InvocationRecord<'_>, ) -> Result<(), crank_registry::RegistryError> { let created_at = OffsetDateTime::now_utc(); let duration_ms = u64::try_from(record.duration.as_millis()).unwrap_or(u64::MAX); let log = InvocationLog { id: InvocationLogId::new(format!("log_{}", uuid::Uuid::now_v7().simple())), workspace_id: tool.workspace_id.clone(), agent_id: Some(tool.agent_id.clone()), operation_id: tool.operation.id.clone(), source: InvocationSource::AgentToolCall, level: record.level, status: record.status, tool_name: record.tool_name.to_owned(), message: record.message.to_owned(), request_id: record.request_id.map(ToOwned::to_owned), status_code: record.status_code, duration_ms, error_kind: record.error_kind.map(ToOwned::to_owned), request_preview: record.request_preview, response_preview: record.response_preview, created_at, }; state .registry .create_invocation_log(CreateInvocationLogRequest { log: &log }) .await } 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), ) } fn tool_error_response( message: &Value, response_mode: ResponseMode, protocol_version: &str, code: &str, error_message: String, error_context: Option, ) -> Response { let mut error = json!({ "code": code, "message": error_message }); if let Some(context) = error_context { error["context"] = context; } transport_response( StatusCode::OK, jsonrpc_result( request_id(message), json!({ "content": [ { "type": "text", "text": error_message } ], "structuredContent": { "error": error }, "isError": true }), ), response_mode, None, Some(protocol_version), ) } fn format_rfc3339(timestamp: OffsetDateTime) -> String { timestamp .format(&Rfc3339) .unwrap_or_else(|_| "1970-01-01T00:00:00Z".to_owned()) } fn add_millis(timestamp: OffsetDateTime, millis: u64) -> OffsetDateTime { let delta = time::Duration::milliseconds(i64::try_from(millis).unwrap_or(i64::MAX)); timestamp + delta } fn serialize_stream_status(status: StreamStatus) -> &'static str { match status { StreamStatus::Created => "created", StreamStatus::Running => "running", StreamStatus::Stopped => "stopped", StreamStatus::Failed => "failed", StreamStatus::Expired => "expired", } } fn serialize_job_status(status: JobStatus) -> &'static str { match status { JobStatus::Created => "created", JobStatus::Running => "running", JobStatus::Completed => "completed", JobStatus::Failed => "failed", JobStatus::Cancelled => "cancelled", JobStatus::Expired => "expired", } } fn runtime_error_code(error: &RuntimeError) -> &'static str { match error { RuntimeError::Schema(_) => "schema_validation_error", RuntimeError::Mapping(_) => "mapping_error", RuntimeError::GraphqlAdapter(_) => "adapter_execution_error", RuntimeError::GrpcAdapter(_) => "adapter_execution_error", RuntimeError::RestAdapter(_) => "adapter_execution_error", RuntimeError::SoapAdapter(_) => "adapter_execution_error", RuntimeError::WebsocketAdapter(_) => "adapter_execution_error", RuntimeError::UnsupportedProtocol { .. } => "unsupported_protocol", RuntimeError::ConcurrencyLimitExceeded { .. } => "runtime_overloaded", RuntimeError::MissingStreamingConfig { .. } => "streaming_config_error", RuntimeError::UnsupportedExecutionMode { .. } => "streaming_mode_error", RuntimeError::InvalidPreparedRequest { .. } => "runtime_error", RuntimeError::InvalidStreamingPayload { .. } => "streaming_payload_error", RuntimeError::MissingAuthProfile { .. } => "auth_profile_not_found", RuntimeError::MissingSecret { .. } | RuntimeError::MissingSecretVersion { .. } => { "secret_not_found" } RuntimeError::InvalidAuthSecretValue { .. } => "secret_value_error", RuntimeError::SecretCrypto { .. } => "secret_crypto_error", } } fn runtime_error_context(error: &RuntimeError) -> Option { match error { RuntimeError::InvalidPreparedRequest { field, reason } => Some(json!({ "field": field, "reason": reason, })), RuntimeError::InvalidStreamingPayload { field, reason } => Some(json!({ "field": field, "reason": reason, })), RuntimeError::InvalidAuthSecretValue { secret_id, reason } => Some(json!({ "secret_id": secret_id, "reason": reason, })), RuntimeError::SecretCrypto { operation, details } => Some(json!({ "operation": operation, "details": details, })), RuntimeError::MissingAuthProfile { auth_profile_id } => Some(json!({ "auth_profile_id": auth_profile_id, })), RuntimeError::MissingSecret { secret_id } => Some(json!({ "secret_id": secret_id, })), RuntimeError::MissingSecretVersion { secret_id, version } => Some(json!({ "secret_id": secret_id, "version": version, })), RuntimeError::MissingStreamingConfig { operation_id } => Some(json!({ "operation_id": operation_id, })), RuntimeError::UnsupportedExecutionMode { operation_id, mode } => Some(json!({ "operation_id": operation_id, "mode": mode, })), RuntimeError::UnsupportedProtocol { protocol } => Some(json!({ "protocol": protocol, })), RuntimeError::ConcurrencyLimitExceeded { kind, limit } => Some(json!({ "kind": kind, "limit": limit, })), _ => None, } } fn hash_access_secret(secret: &str) -> String { let digest = Sha256::digest(secret.as_bytes()); URL_SAFE_NO_PAD.encode(digest) } fn json_response( status: StatusCode, payload: Value, session_id: Option<&str>, protocol_version: Option<&str>, ) -> Response { let mut response = (status, Json(payload)).into_response(); if let Some(session_id) = session_id { response.headers_mut().insert( HEADER_MCP_SESSION_ID, HeaderValue::from_str(session_id) .unwrap_or_else(|_| HeaderValue::from_static("invalid")), ); } if let Some(protocol_version) = protocol_version { response.headers_mut().insert( HEADER_MCP_PROTOCOL_VERSION, HeaderValue::from_str(protocol_version) .unwrap_or_else(|_| HeaderValue::from_static(CURRENT_PROTOCOL_VERSION)), ); } response } fn transport_response( status: StatusCode, payload: Value, response_mode: ResponseMode, session_id: Option<&str>, protocol_version: Option<&str>, ) -> Response { if status == StatusCode::OK && matches!(response_mode, ResponseMode::Sse) { let payload = payload.to_string(); let stream = stream::once(async move { Ok(Event::default().data(payload)) }); return sse_response(status, stream, session_id, protocol_version); } json_response(status, payload, session_id, protocol_version) } fn with_request_id_header(mut response: Response, request_id: &str) -> Response { if let Ok(value) = HeaderValue::from_str(request_id) { response.headers_mut().insert(HEADER_X_REQUEST_ID, value); } response } fn sse_response( status: StatusCode, stream: S, session_id: Option<&str>, protocol_version: Option<&str>, ) -> Response where S: futures_util::stream::Stream> + Send + 'static, { let mut response = ( status, Sse::new(stream).keep_alive(KeepAlive::new().interval(Duration::from_secs(15))), ) .into_response(); if let Some(session_id) = session_id { response.headers_mut().insert( HEADER_MCP_SESSION_ID, HeaderValue::from_str(session_id) .unwrap_or_else(|_| HeaderValue::from_static("invalid")), ); } if let Some(protocol_version) = protocol_version { response.headers_mut().insert( HEADER_MCP_PROTOCOL_VERSION, HeaderValue::from_str(protocol_version) .unwrap_or_else(|_| HeaderValue::from_static(CURRENT_PROTOCOL_VERSION)), ); } response } fn tool_definitions(tool: &PublishedAgentTool) -> Vec { let mut definitions = Vec::new(); let operation = &tool.operation; match operation .execution_config .streaming .as_ref() .map(|value| value.mode) { Some(crank_core::ExecutionMode::Session) => { let streaming = operation .execution_config .streaming .as_ref() .expect("streaming"); if let (Some(start), Some(poll), Some(stop)) = ( streaming.tool_family.start_tool_name.as_ref(), streaming.tool_family.poll_tool_name.as_ref(), streaming.tool_family.stop_tool_name.as_ref(), ) { definitions.push(tool_definition( start, &format!("{} Start", tool.tool_title), &format!("Start session for {}", tool.tool_description), schema_to_json_schema(&operation.input_schema), )); definitions.push(tool_definition( poll, &format!("{} Poll", tool.tool_title), &format!("Poll session for {}", tool.tool_description), id_input_schema("session_id"), )); definitions.push(tool_definition( stop, &format!("{} Stop", tool.tool_title), &format!("Stop session for {}", tool.tool_description), id_input_schema("session_id"), )); } } Some(crank_core::ExecutionMode::AsyncJob) => { let streaming = operation .execution_config .streaming .as_ref() .expect("streaming"); if let (Some(start), Some(status), Some(result), Some(cancel)) = ( streaming.tool_family.start_tool_name.as_ref(), streaming.tool_family.status_tool_name.as_ref(), streaming.tool_family.result_tool_name.as_ref(), streaming.tool_family.cancel_tool_name.as_ref(), ) { definitions.push(tool_definition( start, &format!("{} Start", tool.tool_title), &format!("Start async job for {}", tool.tool_description), schema_to_json_schema(&operation.input_schema), )); definitions.push(tool_definition( status, &format!("{} Status", tool.tool_title), &format!("Get job status for {}", tool.tool_description), id_input_schema("job_id"), )); definitions.push(tool_definition( result, &format!("{} Result", tool.tool_title), &format!("Get async job result for {}", tool.tool_description), id_input_schema("job_id"), )); definitions.push(tool_definition( cancel, &format!("{} Cancel", tool.tool_title), &format!("Cancel async job for {}", tool.tool_description), id_input_schema("job_id"), )); } } _ => { definitions.push(tool_definition( &tool.tool_name, &tool.tool_title, &tool.tool_description, schema_to_json_schema(&operation.input_schema), )); } } definitions } fn resolve_request_id(headers: &HeaderMap) -> String { headers .get(&HEADER_X_REQUEST_ID) .and_then(|value| value.to_str().ok()) .map(str::trim) .filter(|value| is_valid_request_id(value)) .map(ToOwned::to_owned) .unwrap_or_else(|| uuid::Uuid::now_v7().to_string()) } fn is_valid_request_id(value: &str) -> bool { !value.is_empty() && value.len() <= MAX_REQUEST_ID_LEN && value .bytes() .all(|byte| matches!(byte, 0x21..=0x7e) && byte != b',' && byte != b';') } fn tool_definition(name: &str, title: &str, description: &str, input_schema: Value) -> Value { json!({ "name": name, "title": title, "description": description, "inputSchema": input_schema }) } fn id_input_schema(field_name: &str) -> Value { json!({ "type": "object", "properties": { field_name: { "type": "string" } }, "required": [field_name] }) } 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(), kind: GeneratedToolKind::Base, }); } let Some(streaming) = tool.operation.execution_config.streaming.as_ref() else { continue; }; match streaming.mode { crank_core::ExecutionMode::Session => { if streaming.tool_family.start_tool_name.as_deref() == Some(tool_name) { return Some(ResolvedToolCall { tool: tool.clone(), kind: GeneratedToolKind::SessionStart, }); } if streaming.tool_family.poll_tool_name.as_deref() == Some(tool_name) { return Some(ResolvedToolCall { tool: tool.clone(), kind: GeneratedToolKind::SessionPoll, }); } if streaming.tool_family.stop_tool_name.as_deref() == Some(tool_name) { return Some(ResolvedToolCall { tool: tool.clone(), kind: GeneratedToolKind::SessionStop, }); } } crank_core::ExecutionMode::AsyncJob => { if streaming.tool_family.start_tool_name.as_deref() == Some(tool_name) { return Some(ResolvedToolCall { tool: tool.clone(), kind: GeneratedToolKind::AsyncJobStart, }); } if streaming.tool_family.status_tool_name.as_deref() == Some(tool_name) { return Some(ResolvedToolCall { tool: tool.clone(), kind: GeneratedToolKind::AsyncJobStatus, }); } if streaming.tool_family.result_tool_name.as_deref() == Some(tool_name) { return Some(ResolvedToolCall { tool: tool.clone(), kind: GeneratedToolKind::AsyncJobResult, }); } if streaming.tool_family.cancel_tool_name.as_deref() == Some(tool_name) { return Some(ResolvedToolCall { tool: tool.clone(), kind: GeneratedToolKind::AsyncJobCancel, }); } } _ => {} } } None } 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 } fn schema_to_json_schema(schema: &crank_schema::Schema) -> Value { match schema.kind { crank_schema::SchemaKind::Object => { let mut properties = serde_json::Map::new(); let mut required = Vec::new(); for (field_name, field_schema) in &schema.fields { properties.insert(field_name.clone(), schema_to_json_schema(field_schema)); if field_schema.required { required.push(Value::String(field_name.clone())); } } json!({ "type": "object", "properties": properties, "required": required }) } crank_schema::SchemaKind::Array => json!({ "type": "array", "items": schema .items .as_deref() .map(schema_to_json_schema) .unwrap_or_else(|| json!({})) }), crank_schema::SchemaKind::String => json!({ "type": "string" }), crank_schema::SchemaKind::Integer => json!({ "type": "integer" }), crank_schema::SchemaKind::Number => json!({ "type": "number" }), crank_schema::SchemaKind::Boolean => json!({ "type": "boolean" }), crank_schema::SchemaKind::Enum => json!({ "type": "string", "enum": schema.enum_values }), crank_schema::SchemaKind::Null => json!({ "type": "null" }), crank_schema::SchemaKind::Oneof => json!({ "anyOf": schema.variants.iter().map(schema_to_json_schema).collect::>() }), } } impl AllowedOrigins { fn new(public_base_url: Option) -> Self { Self { public_origin: public_base_url.and_then(|value| extract_origin(&value)), } } fn is_allowed(&self, origin: &str) -> bool { if origin.starts_with("http://localhost") || origin.starts_with("http://127.0.0.1") || origin.starts_with("https://localhost") || origin.starts_with("https://127.0.0.1") { return true; } match &self.public_origin { Some(public_origin) => public_origin == origin, None => false, } } } fn extract_origin(url: &str) -> Option { let mut parts = url.split('/'); let scheme = parts.next()?; let empty = parts.next()?; let authority = parts.next()?; if empty.is_empty() { return Some(format!("{scheme}//{authority}")); } None } #[cfg(test)] mod tests { use axum::body::to_bytes; use serde_json::{Value, json}; use super::{ResponseMode, runtime_error_context, tool_error_response}; use crate::jsonrpc::CURRENT_PROTOCOL_VERSION; use crank_runtime::RuntimeError; #[test] fn runtime_error_context_includes_secret_crypto_operation() { let context = runtime_error_context(&RuntimeError::SecretCrypto { operation: "decode secret envelope", details: "bad base64".to_owned(), }) .unwrap(); assert_eq!( context, json!({ "operation": "decode secret envelope", "details": "bad base64" }) ); } #[test] fn runtime_error_context_includes_runtime_overload_details() { let context = runtime_error_context(&RuntimeError::ConcurrencyLimitExceeded { kind: "async_job", limit: 8, }) .unwrap(); assert_eq!( context, json!({ "kind": "async_job", "limit": 8 }) ); } #[tokio::test] async fn tool_error_response_includes_structured_context() { let response = tool_error_response( &json!({"jsonrpc": "2.0", "id": "req-1"}), ResponseMode::Json, CURRENT_PROTOCOL_VERSION, "streaming_payload_error", "request root must be an object".to_owned(), Some(json!({ "field": "request", "reason": "must be an object" })), ); let body = to_bytes(response.into_body(), usize::MAX).await.unwrap(); let payload: Value = serde_json::from_slice(&body).unwrap(); assert_eq!( payload["result"]["structuredContent"]["error"], json!({ "code": "streaming_payload_error", "message": "request root must be an object", "context": { "field": "request", "reason": "must be an object" } }) ); } }