#![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, 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>>, } 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>>, } impl io::Write for SharedLogGuard { fn write(&mut self, bytes: &[u8]) -> io::Result { 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, ) -> 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, 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, 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::() .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 { let secret = format!("crk_{}_{}", 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)), 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, }; 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, 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, ) { 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, ) -> 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) -> Json { Json(json!({ "id": "lead_123", "email": payload["email"] })) } pub(super) async fn create_slow_lead(Json(payload): Json) -> Json { sleep(Duration::from_millis(250)).await; Json(json!({ "id": "lead_123", "email": payload["email"] })) } pub(super) async fn stream_logs() -> Sse>> { 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 { 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(), } }