bsodfather/crank
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
Files
46f8e6cbebef20c26d9c094b89a80bb92c797a99
crank/apps/mcp-server/src/app.rs
T
a.tolmachev 46f8e6cbeb mcp: add bearer token verification seam
2026-05-03 18:43:48 +00:00

2825 lines
89 KiB
Rust
Raw Blame History

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<String>,
}
#[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<String>,
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<Value> {
Json(json!({
"service": "mcp-server",
"status": "ok"
}))
}
async fn mcp_get(
Path(path): Path<AgentRoutePath>,
State(state): State<Arc<AppState>>,
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::<Result<Event, Infallible>>(),
Some(&session_id),
Some(&session.protocol_version),
)
}
async fn mcp_delete(
Path(path): Path<AgentRoutePath>,
State(state): State<Arc<AppState>>,
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<AgentRoutePath>,
State(state): State<Arc<AppState>>,
headers: HeaderMap,
Json(message): Json<Value>,
) -> 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("<unknown>"),
"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::<Vec<_>>();
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<AppState>,
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<Value>,
next_index: usize,
batch_size: usize,
}
async fn resolve_operation_auth(
state: &Arc<AppState>,
workspace_id: &crank_core::WorkspaceId,
execution_config: &crank_core::ExecutionConfig,
) -> Result<Option<ResolvedAuth>, 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<Option<ResolvedAuth>, 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<ResolvedAuth, RuntimeError> {
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<AppState>,
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<AppState>,
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<AppState>,
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::<StoredSessionState>(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<AppState>,
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<AppState>,
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(
&registry,
&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<AppState>,
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<AppState>,
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<AppState>,
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<AppState>,
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<AppState>,
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<AppState>,
path: &AgentRoutePath,
headers: &HeaderMap,
message: &Value,
) -> Result<SessionState, Response> {
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<AppState>,
path: &AgentRoutePath,
headers: &HeaderMap,
required_scope: PlatformApiKeyScope,
) -> Result<VerifiedMachineCredential, StatusCode> {
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<AppState>,
path: &AgentRoutePath,
token: &str,
) -> Result<VerifiedMachineCredential, StatusCode> {
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<AppState>,
path: &AgentRoutePath,
secret: &str,
) -> Result<Option<VerifiedMachineCredential>, 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<AppState>,
path: &AgentRoutePath,
headers: &HeaderMap,
) -> Result<(), RateLimitRejection> {
enforce_transport_rate_limit(state, path, headers)
}
fn enforce_transport_rate_limit(
state: &Arc<AppState>,
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<ResponseMode, 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);
};
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<String, StatusCode> {
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<Option<String>, 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<u16>,
error_kind: Option<&'a str>,
duration: Duration,
request_preview: Value,
response_preview: Value,
}
async fn persist_invocation(
state: &Arc<AppState>,
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<Value>,
) -> 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<Value> {
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<S>(
status: StatusCode,
stream: S,
session_id: Option<&str>,
protocol_version: Option<&str>,
) -> Response
where
S: futures_util::stream::Stream<Item = Result<Event, Infallible>> + 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<Value> {
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<ResolvedToolCall> {
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::<Vec<_>>()
}),
}
}
impl AllowedOrigins {
fn new(public_base_url: Option<String>) -> 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<String> {
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"
}
})
);
}
}
Reference in New Issue View Git Blame Copy Permalink