crank/apps/mcp-server/tests/integration/common.rs

#![allow(dead_code, unused_imports)]

use std::{
    collections::BTreeMap,
    io,
    sync::{Arc, Mutex},
    time::Duration,
};

use axum::{
    Json, Router,
    http::header,
    response::sse::{Event, KeepAlive, Sse},
    routing::{get, post},
};
use base64::{Engine as _, engine::general_purpose::URL_SAFE_NO_PAD};
use crank_core::{
    Agent, AgentId, AgentOperationBinding, AgentStatus, AgentVersion, ExecutionConfig, HttpMethod,
    Operation, OperationId, OperationStatus, PlatformApiKey, PlatformApiKeyId, PlatformApiKeyKind,
    PlatformApiKeyScope, PlatformApiKeyStatus, Protocol, RestTarget, Target, ToolDescription,
    WorkspaceId,
};
use crank_mapping::{MappingRule, MappingSet};
use crank_registry::{
    CreateAgentRequest, CreatePlatformApiKeyRequest, ListInvocationLogsQuery, PostgresRegistry,
    PublishAgentRequest, PublishRequest, SaveAgentBindingsRequest,
};
use crank_runtime::{
    InMemoryCoordinationStateStore, RequestRateLimitConfig, RequestRateLimiter, RuntimeExecutor,
    SecretCrypto,
};
use crank_schema::{Schema, SchemaKind};
use futures_util::stream;
use serde_json::{Value, json};
use sha2::{Digest, Sha256};
use time::{OffsetDateTime, format_description::well_known::Rfc3339};
use tokio::net::TcpListener;
use tokio::time::sleep;
use tracing_subscriber::{filter::LevelFilter, fmt::MakeWriter, prelude::*};

use crank_community_mcp::{
    auth::{CommunityMachineCredentialVerifier, SharedMachineCredentialVerifier},
    build_app,
    catalog::PublishedToolCatalog,
    session::{InMemorySessionStore, SharedSessionStore, TransportSessionStore},
};

fn test_workspace_id() -> WorkspaceId {
    WorkspaceId::new("ws_default")
}

#[derive(Clone, Default)]
struct SharedLogWriter {
    buffer: Arc<Mutex<Vec<u8>>>,
}

impl SharedLogWriter {
    fn output(&self) -> String {
        String::from_utf8(self.buffer.lock().unwrap().clone()).unwrap()
    }
}

impl<'a> MakeWriter<'a> for SharedLogWriter {
    type Writer = SharedLogGuard;

    fn make_writer(&'a self) -> Self::Writer {
        SharedLogGuard {
            buffer: Arc::clone(&self.buffer),
        }
    }
}

struct SharedLogGuard {
    buffer: Arc<Mutex<Vec<u8>>>,
}

impl io::Write for SharedLogGuard {
    fn write(&mut self, bytes: &[u8]) -> io::Result<usize> {
        self.buffer.lock().unwrap().extend_from_slice(bytes);
        Ok(bytes.len())
    }

    fn flush(&mut self) -> io::Result<()> {
        Ok(())
    }
}

fn test_workspace_slug() -> &'static str {
    "default"
}

fn test_agent_id(agent_slug: &str) -> AgentId {
    AgentId::new(format!("agent_{agent_slug}"))
}

fn build_test_app(
    registry: PostgresRegistry,
    refresh_interval: Duration,
    public_base_url: Option<String>,
) -> axum::Router {
    build_test_app_with_rate_limit(
        registry,
        refresh_interval,
        public_base_url,
        RequestRateLimitConfig::new(10_000, 10_000).unwrap(),
    )
}

fn build_test_app_with_rate_limit(
    registry: PostgresRegistry,
    refresh_interval: Duration,
    public_base_url: Option<String>,
    rate_limit_config: RequestRateLimitConfig,
) -> axum::Router {
    build_test_app_with_store(
        registry,
        refresh_interval,
        public_base_url,
        rate_limit_config,
        std::sync::Arc::new(InMemorySessionStore::default()),
        std::sync::Arc::new(CommunityMachineCredentialVerifier),
    )
}

fn build_test_app_with_store(
    registry: PostgresRegistry,
    refresh_interval: Duration,
    public_base_url: Option<String>,
    rate_limit_config: RequestRateLimitConfig,
    sessions: SharedSessionStore,
    credential_verifier: SharedMachineCredentialVerifier,
) -> axum::Router {
    build_app(
        registry,
        refresh_interval,
        public_base_url,
        SecretCrypto::new("test-master-key").unwrap(),
        RuntimeExecutor::new(),
        RequestRateLimiter::new(rate_limit_config),
        std::sync::Arc::new(InMemoryCoordinationStateStore::default()),
        sessions,
        credential_verifier,
    )
}

pub(super) async fn initialize_session(
    client: &reqwest::Client,
    mcp_url: &str,
    api_key: &str,
) -> String {
    let initialize_response = client
        .post(mcp_url)
        .header(header::ACCEPT, "application/json, text/event-stream")
        .header(header::AUTHORIZATION, format!("Bearer {api_key}"))
        .json(&json!({
            "jsonrpc": "2.0",
            "id": 1,
            "method": "initialize",
            "params": {
                "protocolVersion": "2025-11-25"
            }
        }))
        .send()
        .await
        .unwrap();
    assert_eq!(initialize_response.status(), reqwest::StatusCode::OK);
    let session_id = initialize_response
        .headers()
        .get("MCP-Session-Id")
        .unwrap()
        .to_str()
        .unwrap()
        .to_owned();

    let initialized_response = client
        .post(mcp_url)
        .header(header::ACCEPT, "application/json, text/event-stream")
        .header(header::AUTHORIZATION, format!("Bearer {api_key}"))
        .header("MCP-Session-Id", &session_id)
        .header("MCP-Protocol-Version", "2025-11-25")
        .json(&json!({
            "jsonrpc": "2.0",
            "method": "notifications/initialized",
            "params": {}
        }))
        .send()
        .await
        .unwrap();

    assert_eq!(initialized_response.status(), reqwest::StatusCode::ACCEPTED);

    session_id
}

pub(super) async fn post_jsonrpc(
    client: &reqwest::Client,
    mcp_url: &str,
    api_key: &str,
    session_id: Option<&str>,
    payload: Value,
) -> Value {
    post_jsonrpc_response(client, mcp_url, api_key, session_id, None, payload)
        .await
        .json::<Value>()
        .await
        .unwrap()
}

pub(super) async fn post_jsonrpc_response(
    client: &reqwest::Client,
    mcp_url: &str,
    api_key: &str,
    session_id: Option<&str>,
    request_id: Option<&str>,
    payload: Value,
) -> reqwest::Response {
    let mut request = client
        .post(mcp_url)
        .header(header::ACCEPT, "application/json, text/event-stream")
        .header(header::AUTHORIZATION, format!("Bearer {api_key}"))
        .header("MCP-Protocol-Version", "2025-11-25");

    if let Some(session_id) = session_id {
        request = request.header("MCP-Session-Id", session_id);
    }

    if let Some(request_id) = request_id {
        request = request.header("x-request-id", request_id);
    }

    request.json(&payload).send().await.unwrap()
}

pub(super) async fn create_platform_api_key(
    registry: &PostgresRegistry,
    agent_slug: &str,
    name: &str,
    scopes: &[PlatformApiKeyScope],
) -> String {
    create_platform_api_key_with_kind(
        registry,
        agent_slug,
        name,
        PlatformApiKeyKind::McpClient,
        scopes,
        "crk",
    )
    .await
}

pub(super) async fn create_approval_platform_api_key(
    registry: &PostgresRegistry,
    agent_slug: &str,
    name: &str,
) -> String {
    create_platform_api_key_with_kind(
        registry,
        agent_slug,
        name,
        PlatformApiKeyKind::Approval,
        &[PlatformApiKeyScope::Approve, PlatformApiKeyScope::Deny],
        "crk_appr",
    )
    .await
}

async fn create_platform_api_key_with_kind(
    registry: &PostgresRegistry,
    agent_slug: &str,
    name: &str,
    key_kind: PlatformApiKeyKind,
    scopes: &[PlatformApiKeyScope],
    prefix: &str,
) -> String {
    let secret = format!("{prefix}_{}_{}", name, uuid::Uuid::now_v7().simple());
    let api_key = PlatformApiKey {
        id: PlatformApiKeyId::new(format!("pk_{name}")),
        workspace_id: test_workspace_id(),
        agent_id: Some(test_agent_id(agent_slug)),
        key_kind,
        name: name.to_owned(),
        prefix: secret.chars().take(16).collect(),
        scopes: scopes.to_vec(),
        status: PlatformApiKeyStatus::Active,
        created_at: OffsetDateTime::parse("2026-03-26T10:00:00Z", &Rfc3339).unwrap(),
        last_used_at: None,
        expires_at: None,
        allowed_origins: Vec::new(),
    };

    registry
        .create_platform_api_key(CreatePlatformApiKeyRequest {
            api_key: &api_key,
            secret_hash: &hash_access_secret(&secret),
        })
        .await
        .unwrap();

    secret
}

pub(super) fn hash_access_secret(secret: &str) -> String {
    let digest = Sha256::digest(secret.as_bytes());
    URL_SAFE_NO_PAD.encode(digest)
}

pub(super) async fn spawn_upstream_server() -> String {
    let app = Router::new()
        .route("/sse/logs", get(stream_logs))
        .route("/crm/leads", post(create_lead))
        .route("/crm/slow-leads", post(create_slow_lead));
    let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
    let address = listener.local_addr().unwrap();

    tokio::spawn(async move {
        axum::serve(listener, app).await.unwrap();
    });

    format!("http://{}", address)
}

pub(super) async fn spawn_mcp_server(app: Router) -> String {
    let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
    let address = listener.local_addr().unwrap();

    tokio::spawn(async move {
        axum::serve(listener, app).await.unwrap();
    });

    format!("http://{}", address)
}

pub(super) fn agent_mcp_url(base_url: &str, agent_slug: &str) -> String {
    format!("{base_url}/v1/{}/{}", test_workspace_slug(), agent_slug)
}

pub(super) async fn publish_agent_for_operation(
    registry: &PostgresRegistry,
    operation: &Operation<Schema, MappingSet>,
    agent_slug: &str,
) {
    publish_agent_with_bindings(registry, agent_slug, vec![binding_for_operation(operation)]).await;
}

pub(super) async fn publish_agent_with_bindings(
    registry: &PostgresRegistry,
    agent_slug: &str,
    bindings: Vec<AgentOperationBinding>,
) {
    let agent_id = AgentId::new(format!("agent_{agent_slug}"));
    let agent = Agent {
        id: agent_id.clone(),
        workspace_id: test_workspace_id(),
        slug: agent_slug.to_owned(),
        display_name: format!("Agent {agent_slug}"),
        description: "Curated MCP toolset".to_owned(),
        status: AgentStatus::Draft,
        current_draft_version: 1,
        latest_published_version: None,
        created_at: OffsetDateTime::parse("2026-03-26T10:00:00Z", &Rfc3339).unwrap(),
        updated_at: OffsetDateTime::parse("2026-03-26T10:00:00Z", &Rfc3339).unwrap(),
        published_at: None,
    };
    let version = AgentVersion {
        agent_id: agent_id.clone(),
        version: 1,
        status: AgentStatus::Draft,
        instructions: json!({}),
        tool_selection_policy: json!({}),
        created_at: OffsetDateTime::parse("2026-03-26T10:00:00Z", &Rfc3339).unwrap(),
    };

    registry
        .create_agent(CreateAgentRequest {
            agent: &agent,
            version: &version,
            bindings: &bindings,
        })
        .await
        .unwrap();
    registry
        .publish_agent(PublishAgentRequest {
            workspace_id: &test_workspace_id(),
            agent_id: &agent_id,
            version: 1,
            published_at: &OffsetDateTime::parse("2026-03-26T10:00:00Z", &Rfc3339).unwrap(),
            published_by: Some("alice"),
        })
        .await
        .unwrap();
}

pub(super) fn binding_for_operation(
    operation: &Operation<Schema, MappingSet>,
) -> AgentOperationBinding {
    AgentOperationBinding {
        operation_id: operation.id.clone(),
        operation_version: 1,
        tool_name: operation.name.clone(),
        tool_title: operation.tool_description.title.clone(),
        tool_description_override: None,
        enabled: true,
    }
}

pub(super) async fn create_lead(Json(payload): Json<Value>) -> Json<Value> {
    Json(json!({
        "id": "lead_123",
        "email": payload["email"]
    }))
}

pub(super) async fn create_slow_lead(Json(payload): Json<Value>) -> Json<Value> {
    sleep(Duration::from_millis(250)).await;

    Json(json!({
        "id": "lead_123",
        "email": payload["email"]
    }))
}

pub(super) async fn stream_logs()
-> Sse<impl futures_util::Stream<Item = Result<Event, std::convert::Infallible>>> {
    let events = vec![
        json!({ "level": "info", "message": "sync started" }),
        json!({ "level": "warn", "message": "cache warmup slow" }),
        json!({ "level": "error", "message": "upstream timeout" }),
    ];
    let stream = stream::iter(events.into_iter().map(|payload| {
        Ok::<_, std::convert::Infallible>(Event::default().data(payload.to_string()))
    }));
    Sse::new(stream).keep_alive(KeepAlive::new().interval(Duration::from_secs(15)))
}

pub(super) async fn test_registry() -> PostgresRegistry {
    let database_url = crank_test_support::postgres_schema_url("test_mcp_server").await;
    PostgresRegistry::connect(&database_url).await.unwrap()
}

pub(super) fn test_operation(base_url: &str, name: &str) -> Operation<Schema, MappingSet> {
    Operation {
        id: OperationId::new(format!("op_{name}")),
        name: name.to_owned(),
        display_name: "Create Lead".to_owned(),
        category: "sales".to_owned(),
        protocol: Protocol::Rest,
        security_level: crank_core::OperationSecurityLevel::Standard,
        status: OperationStatus::Published,
        version: 1,
        target: Target::Rest(RestTarget {
            base_url: base_url.to_owned(),
            method: HttpMethod::Post,
            path_template: "/crm/leads".to_owned(),
            static_headers: BTreeMap::new(),
        }),
        input_schema: object_schema("email"),
        output_schema: object_schema("id"),
        input_mapping: MappingSet {
            rules: vec![MappingRule {
                source: "$.mcp.email".to_owned(),
                target: "$.request.body.email".to_owned(),
                required: true,
                default_value: None,
                transform: None,
                condition: None,
                notes: None,
            }],
        },
        output_mapping: MappingSet {
            rules: vec![MappingRule {
                source: "$.response.body.id".to_owned(),
                target: "$.output.id".to_owned(),
                required: true,
                default_value: None,
                transform: None,
                condition: None,
                notes: None,
            }],
        },
        execution_config: ExecutionConfig {
            timeout_ms: 1_000,
            retry_policy: None,
            response_cache: None,
            idempotency: None,
            safety: None,
            auth_profile_ref: None,
            headers: BTreeMap::new(),
        },
        tool_description: ToolDescription {
            title: "Create Lead".to_owned(),
            description: "Creates a CRM lead".to_owned(),
            tags: Vec::new(),
            examples: Vec::new(),
        },
        samples: None,
        generated_draft: None,
        config_export: None,
        wizard_state: None,
        created_at: OffsetDateTime::parse("2026-03-26T10:00:00Z", &Rfc3339).unwrap(),
        updated_at: OffsetDateTime::parse("2026-03-26T10:00:00Z", &Rfc3339).unwrap(),
        published_at: Some(OffsetDateTime::parse("2026-03-26T10:00:00Z", &Rfc3339).unwrap()),
    }
}

pub(super) fn object_schema(field_name: &str) -> Schema {
    Schema {
        kind: SchemaKind::Object,
        description: None,
        required: true,
        nullable: false,
        default_value: None,
        fields: BTreeMap::from([(
            field_name.to_owned(),
            Schema {
                kind: SchemaKind::String,
                description: None,
                required: true,
                nullable: false,
                default_value: None,
                fields: BTreeMap::new(),
                items: None,
                enum_values: Vec::new(),
                variants: Vec::new(),
            },
        )]),
        items: None,
        enum_values: Vec::new(),
        variants: Vec::new(),
    }
}