crank/apps/mcp-server/src/app.rs

use std::{
    convert::Infallible,
    sync::Arc,
    time::{Duration, Instant},
};

use axum::{
    Json, Router,
    extract::{Path, State},
    http::{
        HeaderMap, HeaderValue, StatusCode,
        header::{self, ACCEPT, AUTHORIZATION},
    },
    response::{
        IntoResponse, Response,
        sse::{Event, KeepAlive, Sse},
    },
    routing::get,
};
use base64::{Engine as _, engine::general_purpose::URL_SAFE_NO_PAD};
use crank_core::{
    InvocationLevel, InvocationLog, InvocationLogId, InvocationSource, InvocationStatus,
    PlatformApiKeyScope,
};
use crank_registry::{CreateInvocationLogRequest, PostgresRegistry, PublishedAgentTool};
use crank_runtime::{RuntimeError, RuntimeExecutor, RuntimeOperation};
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 crate::{
    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, SessionStore},
};

const HEADER_MCP_SESSION_ID: &str = "MCP-Session-Id";
const HEADER_MCP_PROTOCOL_VERSION: &str = "MCP-Protocol-Version";

#[derive(Clone, Copy)]
enum ResponseMode {
    Json,
    Sse,
}

#[derive(Clone)]
pub struct AppState {
    registry: PostgresRegistry,
    catalog: PublishedToolCatalog,
    runtime: RuntimeExecutor,
    sessions: SessionStore,
    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, Debug, Deserialize)]
struct AgentRoutePath {
    workspace_slug: String,
    agent_slug: String,
}

pub fn build_app(
    registry: PostgresRegistry,
    refresh_interval: Duration,
    public_base_url: Option<String>,
) -> Router {
    let state = Arc::new(AppState {
        registry: registry.clone(),
        catalog: PublishedToolCatalog::new(registry, refresh_interval),
        runtime: RuntimeExecutor::new(),
        sessions: SessionStore::new(),
        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_platform_api_key(&state, &path, &headers, PlatformApiKeyScope::Read).await
    {
        return status.into_response();
    }

    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) = state.sessions.get(&session_id).await 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_platform_api_key(&state, &path, &headers, PlatformApiKeyScope::Read).await
    {
        return status.into_response();
    }

    match session_id_from_headers(&headers) {
        Ok(Some(session_id)) => match state.sessions.get(&session_id).await {
            Some(session)
                if session.workspace_slug == path.workspace_slug
                    && session.agent_slug == path.agent_slug =>
            {
                if state.sessions.delete(&session_id).await {
                    StatusCode::NO_CONTENT.into_response()
                } else {
                    StatusCode::NOT_FOUND.into_response()
                }
            }
            Some(_) => StatusCode::NOT_FOUND.into_response(),
            None => StatusCode::NOT_FOUND.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 {
    if let Err(status) = validate_origin(&state.allowed_origins, &headers) {
        return status.into_response();
    }

    let response_mode = match negotiate_post_response_mode(&headers) {
        Ok(mode) => mode,
        Err(status) => return status.into_response(),
    };

    if is_response(&message) || is_notification(&message) && method_name(&message).is_none() {
        return StatusCode::ACCEPTED.into_response();
    }

    let protocol_version = match protocol_version_from_headers(&headers) {
        Ok(value) => value,
        Err(status) => return status.into_response(),
    };

    if let Some(session_id) = headers.get(HEADER_MCP_SESSION_ID) {
        if let Ok(session_id) = session_id.to_str() {
            if let Some(session) = state.sessions.get(session_id).await {
                if let Err(status) =
                    validate_session_protocol_version(&headers, &session.protocol_version)
                {
                    return status.into_response();
                }
            }
        }
    }

    let required_scope = match method_name(&message) {
        Some("tools/call") => PlatformApiKeyScope::Write,
        _ => PlatformApiKeyScope::Read,
    };
    if let Err(status) = require_platform_api_key(&state, &path, &headers, required_scope).await {
        return status.into_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().map(tool_definition).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
                .get_tool(
                    &session.workspace_slug,
                    &session.agent_slug,
                    &tool_call_params.name,
                )
                .await
            {
                Ok(Some(tool)) => {
                    let runtime_operation = runtime_operation(&tool);
                    let request_preview =
                        build_request_preview(&state.runtime, &runtime_operation, &arguments);
                    let started_at = Instant::now();

                    match state.runtime.execute(&runtime_operation, &arguments).await {
                        Ok(output) => transport_response(
                            StatusCode::OK,
                            {
                                let _ = persist_invocation(
                                    &state,
                                    &tool,
                                    InvocationRecord {
                                        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;
                                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(&session.protocol_version),
                        ),
                        Err(error) => {
                            let _ = persist_invocation(
                                &state,
                                &tool,
                                InvocationRecord {
                                    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;
                            transport_response(
                                StatusCode::OK,
                                jsonrpc_result(
                                    request_id(&message),
                                    json!({
                                        "content": [
                                            {
                                                "type": "text",
                                                "text": error.to_string()
                                            }
                                        ],
                                        "structuredContent": {
                                            "error": {
                                                "code": runtime_error_code(&error),
                                                "message": error.to_string()
                                            }
                                        },
                                        "isError": true
                                    }),
                                ),
                                response_mode,
                                None,
                                Some(&session.protocol_version),
                            )
                        }
                    }
                }
                Ok(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),
        ),
    }
}

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 session_id = state
        .sessions
        .create(protocol_version, &path.workspace_slug, &path.agent_slug)
        .await;

    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 {
            Some(session)
                if session.workspace_slug == path.workspace_slug
                    && session.agent_slug == path.agent_slug =>
            {
                if state.sessions.mark_initialized(&session_id).await {
                    StatusCode::ACCEPTED.into_response()
                } else {
                    StatusCode::NOT_FOUND.into_response()
                }
            }
            Some(_) => StatusCode::NOT_FOUND.into_response(),
            None => StatusCode::NOT_FOUND.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) = state.sessions.get(&session_id).await 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_platform_api_key(
    state: &Arc<AppState>,
    path: &AgentRoutePath,
    headers: &HeaderMap,
    required_scope: PlatformApiKeyScope,
) -> Result<(), StatusCode> {
    let secret = bearer_token(headers).ok_or(StatusCode::UNAUTHORIZED)?;
    let secret_hash = hash_access_secret(secret);
    let Some(api_key) = state
        .registry
        .get_platform_api_key_by_secret_for_workspace_slug(&path.workspace_slug, &secret_hash)
        .await
        .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?
    else {
        return Err(StatusCode::UNAUTHORIZED);
    };

    if !allows_scope(&api_key.api_key.scopes, required_scope) {
        return Err(StatusCode::FORBIDDEN);
    }

    let used_at = OffsetDateTime::now_utc()
        .format(&Rfc3339)
        .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
    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(())
}

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 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 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> {
    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()
        .format(&Rfc3339)
        .unwrap_or_else(|_| "1970-01-01T00:00:00Z".to_owned());
    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: tool.tool_name.clone(),
        message: record.message.to_owned(),
        request_id: None,
        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 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::UnsupportedProtocol { .. } => "unsupported_protocol",
        RuntimeError::InvalidPreparedRequest { .. } => "runtime_error",
    }
}

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 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_definition(tool: &PublishedAgentTool) -> Value {
    json!({
        "name": tool.tool_name,
        "title": tool.tool_title,
        "description": tool.tool_description,
        "inputSchema": schema_to_json_schema(&tool.operation.input_schema)
    })
}

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
}