#![allow(dead_code, unused_imports)] use super::common::*; 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, ) } #[tokio::test] async fn initializes_lists_and_calls_published_tool_via_mcp() { let registry = test_registry().await; let upstream_base_url = spawn_upstream_server().await; let operation = test_operation(&upstream_base_url, "crm_create_lead"); registry .create_operation(&test_workspace_id(), &operation, Some("alice")) .await .unwrap(); registry .publish_operation(PublishRequest { workspace_id: &test_workspace_id(), operation_id: &operation.id, version: 1, published_at: &OffsetDateTime::parse("2026-03-26T10:00:00Z", &Rfc3339).unwrap(), published_by: Some("alice"), }) .await .unwrap(); publish_agent_for_operation(®istry, &operation, "sales-rest").await; let api_key = create_platform_api_key( ®istry, "sales-rest", "mcp-rest", &[PlatformApiKeyScope::Read, PlatformApiKeyScope::Write], ) .await; let base_url = spawn_mcp_server(build_test_app( registry.clone(), Duration::from_millis(0), Some("https://crank.example.com".to_owned()), )) .await; let client = reqwest::Client::new(); let mcp_url = agent_mcp_url(&base_url, "sales-rest"); let initialized_session = initialize_session(&client, &mcp_url, &api_key).await; let tools = post_jsonrpc( &client, &mcp_url, &api_key, Some(&initialized_session), json!({ "jsonrpc": "2.0", "id": 2, "method": "tools/list", "params": {} }), ) .await; let call_result = post_jsonrpc( &client, &mcp_url, &api_key, Some(&initialized_session), json!({ "jsonrpc": "2.0", "id": 3, "method": "tools/call", "params": { "name": "crm_create_lead", "arguments": { "email": "user@example.com" } } }), ) .await; assert_eq!(tools["result"]["tools"][0]["name"], "crm_create_lead"); assert_eq!( call_result["result"]["structuredContent"], json!({ "id": "lead_123" }) ); assert_eq!(call_result["result"]["isError"], false); let logs = registry .list_invocation_logs(ListInvocationLogsQuery { workspace_id: &test_workspace_id(), level: None, search_text: None, source: Some(crank_core::InvocationSource::AgentToolCall), operation_id: Some(&operation.id), agent_id: None, created_after: None, limit: 10, }) .await .unwrap(); assert_eq!(logs.len(), 1); assert_eq!( logs[0].log.source, crank_core::InvocationSource::AgentToolCall ); assert_eq!(logs[0].log.status, crank_core::InvocationStatus::Ok); assert_eq!(logs[0].log.tool_name, "crm_create_lead"); let keys = registry .list_platform_api_keys(&test_workspace_id()) .await .unwrap(); assert!(keys[0].api_key.last_used_at.is_some()); } #[tokio::test] async fn preserves_request_id_for_tool_call_invocations() { let registry = test_registry().await; let upstream_base_url = spawn_upstream_server().await; let operation = test_operation(&upstream_base_url, "crm_request_id"); registry .create_operation(&test_workspace_id(), &operation, Some("alice")) .await .unwrap(); registry .publish_operation(PublishRequest { workspace_id: &test_workspace_id(), operation_id: &operation.id, version: 1, published_at: &OffsetDateTime::parse("2026-03-26T10:00:00Z", &Rfc3339).unwrap(), published_by: Some("alice"), }) .await .unwrap(); publish_agent_for_operation(®istry, &operation, "sales-request-id").await; let api_key = create_platform_api_key( ®istry, "sales-request-id", "mcp-request-id", &[PlatformApiKeyScope::Read, PlatformApiKeyScope::Write], ) .await; let base_url = spawn_mcp_server(build_test_app( registry.clone(), Duration::from_millis(0), Some("https://crank.example.com".to_owned()), )) .await; let client = reqwest::Client::new(); let mcp_url = agent_mcp_url(&base_url, "sales-request-id"); let initialized_session = initialize_session(&client, &mcp_url, &api_key).await; let response = post_jsonrpc_response( &client, &mcp_url, &api_key, Some(&initialized_session), Some("req_test_123"), json!({ "jsonrpc": "2.0", "id": 3, "method": "tools/call", "params": { "name": "crm_request_id", "arguments": { "email": "user@example.com" } } }), ) .await; assert_eq!( response.headers()["x-request-id"].to_str().unwrap(), "req_test_123" ); let call_result = response.json::().await.unwrap(); assert_eq!(call_result["result"]["isError"], false); let logs = registry .list_invocation_logs(ListInvocationLogsQuery { workspace_id: &test_workspace_id(), level: None, search_text: None, source: Some(crank_core::InvocationSource::AgentToolCall), operation_id: Some(&operation.id), agent_id: None, created_after: None, limit: 10, }) .await .unwrap(); assert_eq!(logs.len(), 1); assert_eq!(logs[0].log.request_id.as_deref(), Some("req_test_123")); } #[tokio::test] async fn generates_request_id_for_tool_call_responses_and_logs() { let registry = test_registry().await; let upstream_base_url = spawn_upstream_server().await; let operation = test_operation(&upstream_base_url, "crm_generated_request_id"); registry .create_operation(&test_workspace_id(), &operation, Some("alice")) .await .unwrap(); registry .publish_operation(PublishRequest { workspace_id: &test_workspace_id(), operation_id: &operation.id, version: 1, published_at: &OffsetDateTime::parse("2026-03-26T10:00:00Z", &Rfc3339).unwrap(), published_by: Some("alice"), }) .await .unwrap(); publish_agent_for_operation(®istry, &operation, "sales-generated-request-id").await; let api_key = create_platform_api_key( ®istry, "sales-generated-request-id", "mcp-generated-request-id", &[PlatformApiKeyScope::Read, PlatformApiKeyScope::Write], ) .await; let base_url = spawn_mcp_server(build_test_app( registry.clone(), Duration::from_millis(0), Some("https://crank.example.com".to_owned()), )) .await; let client = reqwest::Client::new(); let mcp_url = agent_mcp_url(&base_url, "sales-generated-request-id"); let initialized_session = initialize_session(&client, &mcp_url, &api_key).await; let response = post_jsonrpc_response( &client, &mcp_url, &api_key, Some(&initialized_session), None, json!({ "jsonrpc": "2.0", "id": 3, "method": "tools/call", "params": { "name": "crm_generated_request_id", "arguments": { "email": "user@example.com" } } }), ) .await; let request_id = response .headers() .get("x-request-id") .unwrap() .to_str() .unwrap() .to_owned(); assert!(!request_id.is_empty()); let call_result = response.json::().await.unwrap(); assert_eq!(call_result["result"]["isError"], false); let logs = registry .list_invocation_logs(ListInvocationLogsQuery { workspace_id: &test_workspace_id(), level: None, search_text: None, source: Some(crank_core::InvocationSource::AgentToolCall), operation_id: Some(&operation.id), agent_id: None, created_after: None, limit: 10, }) .await .unwrap(); assert_eq!(logs.len(), 1); assert_eq!(logs[0].log.request_id.as_deref(), Some(request_id.as_str())); } #[tokio::test] async fn emits_request_id_in_mcp_ingress_logs() { let registry = test_registry().await; let upstream_base_url = spawn_upstream_server().await; let operation = test_operation(&upstream_base_url, "crm_request_trace"); registry .create_operation(&test_workspace_id(), &operation, Some("alice")) .await .unwrap(); registry .publish_operation(PublishRequest { workspace_id: &test_workspace_id(), operation_id: &operation.id, version: 1, published_at: &OffsetDateTime::parse("2026-03-26T10:00:00Z", &Rfc3339).unwrap(), published_by: Some("alice"), }) .await .unwrap(); publish_agent_for_operation(®istry, &operation, "sales-request-trace").await; let api_key = create_platform_api_key( ®istry, "sales-request-trace", "mcp-request-trace", &[PlatformApiKeyScope::Read, PlatformApiKeyScope::Write], ) .await; let base_url = spawn_mcp_server(build_test_app( registry, Duration::from_millis(0), Some("https://crank.example.com".to_owned()), )) .await; let client = reqwest::Client::new(); let mcp_url = agent_mcp_url(&base_url, "sales-request-trace"); let writer = SharedLogWriter::default(); let subscriber = tracing_subscriber::registry().with( tracing_subscriber::fmt::layer() .with_writer(writer.clone()) .without_time() .with_ansi(false) .with_target(false) .compact() .with_filter(LevelFilter::INFO), ); let _ = tracing::subscriber::set_global_default(subscriber); let response = post_jsonrpc_response( &client, &mcp_url, &api_key, None, Some("req_mcp_trace_123"), json!({ "jsonrpc": "2.0", "id": 1, "method": "initialize", "params": { "protocolVersion": "2025-03-26" } }), ) .await; assert_eq!( response.headers()["x-request-id"].to_str().unwrap(), "req_mcp_trace_123" ); assert_eq!(response.status(), reqwest::StatusCode::OK); let logs = writer.output(); assert!( logs.contains("mcp request received"), "captured logs did not include ingress marker: {logs}" ); assert!(logs.contains("req_mcp_trace_123")); assert!(logs.contains("sales-request-trace")); assert!(logs.contains("default")); assert!(logs.contains("initialize")); } #[tokio::test] async fn requires_initialized_notification_before_tool_methods() { let registry = test_registry().await; publish_agent_with_bindings(®istry, "sales-init", vec![]).await; let api_key = create_platform_api_key( ®istry, "sales-init", "mcp-init", &[PlatformApiKeyScope::Read], ) .await; let base_url = spawn_mcp_server(build_test_app( registry, Duration::from_millis(0), Some("https://crank.example.com".to_owned()), )) .await; let client = reqwest::Client::new(); let mcp_url = agent_mcp_url(&base_url, "sales-init"); 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(); let session_id = initialize_response .headers() .get("MCP-Session-Id") .unwrap() .to_str() .unwrap() .to_owned(); let tools_list = 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", "id": 2, "method": "tools/list", "params": {} })) .send() .await .unwrap() .json::() .await .unwrap(); assert_eq!(tools_list["error"]["code"], -32002); } #[tokio::test] async fn initialize_can_return_sse_response_when_client_prefers_event_stream() { let registry = test_registry().await; publish_agent_with_bindings(®istry, "sales-sse-init", vec![]).await; let api_key = create_platform_api_key( ®istry, "sales-sse-init", "mcp-sse-init", &[PlatformApiKeyScope::Read, PlatformApiKeyScope::Write], ) .await; let base_url = spawn_mcp_server(build_test_app( registry, Duration::from_millis(0), Some("https://crank.example.com".to_owned()), )) .await; let client = reqwest::Client::new(); let response = client .post(agent_mcp_url(&base_url, "sales-sse-init")) .header(header::ACCEPT, "text/event-stream, application/json") .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!(response.status(), reqwest::StatusCode::OK); assert_eq!( response .headers() .get(header::CONTENT_TYPE) .unwrap() .to_str() .unwrap(), "text/event-stream" ); assert!(response.headers().get("MCP-Session-Id").is_some()); let body = response.text().await.unwrap(); assert!(body.contains("\"jsonrpc\":\"2.0\"")); assert!(body.contains("\"protocolVersion\":\"2025-11-25\"")); } #[tokio::test] async fn get_opens_sse_stream_for_initialized_session() { let registry = test_registry().await; publish_agent_with_bindings(®istry, "sales-get-sse", vec![]).await; let api_key = create_platform_api_key( ®istry, "sales-get-sse", "mcp-get-sse", &[PlatformApiKeyScope::Read, PlatformApiKeyScope::Write], ) .await; let base_url = spawn_mcp_server(build_test_app( registry, Duration::from_millis(0), Some("https://crank.example.com".to_owned()), )) .await; let client = reqwest::Client::new(); let mcp_url = agent_mcp_url(&base_url, "sales-get-sse"); let initialized_session = initialize_session(&client, &mcp_url, &api_key).await; let response = client .get(&mcp_url) .header(header::ACCEPT, "text/event-stream") .header(header::AUTHORIZATION, format!("Bearer {api_key}")) .header("MCP-Session-Id", &initialized_session) .header("MCP-Protocol-Version", "2025-11-25") .send() .await .unwrap(); assert_eq!(response.status(), reqwest::StatusCode::OK); assert_eq!( response .headers() .get(header::CONTENT_TYPE) .unwrap() .to_str() .unwrap(), "text/event-stream" ); assert_eq!( response .headers() .get("MCP-Session-Id") .unwrap() .to_str() .unwrap(), initialized_session ); } #[tokio::test] async fn get_returns_not_found_for_expired_transport_session() { let registry = test_registry().await; publish_agent_with_bindings(®istry, "sales-expired-session", vec![]).await; let api_key = create_platform_api_key( ®istry, "sales-expired-session", "mcp-expired-session", &[PlatformApiKeyScope::Read], ) .await; let session_store = Arc::new(InMemorySessionStore::default()); let session_id = session_store .create( "2025-11-25", test_workspace_slug(), "sales-expired-session", OffsetDateTime::parse("2026-05-01T10:00:00Z", &Rfc3339).unwrap(), Some(OffsetDateTime::parse("2026-05-01T10:00:01Z", &Rfc3339).unwrap()), ) .await .unwrap(); session_store .mark_initialized( &session_id, OffsetDateTime::parse("2026-05-01T10:00:01Z", &Rfc3339).unwrap(), ) .await .unwrap(); let base_url = spawn_mcp_server(build_test_app_with_store( registry, Duration::from_millis(0), Some("https://crank.example.com".to_owned()), RequestRateLimitConfig::new(10_000, 10_000).unwrap(), session_store, std::sync::Arc::new(CommunityMachineCredentialVerifier), )) .await; let client = reqwest::Client::new(); let response = client .get(agent_mcp_url(&base_url, "sales-expired-session")) .header(header::ACCEPT, "text/event-stream") .header(header::AUTHORIZATION, format!("Bearer {api_key}")) .header("MCP-Session-Id", &session_id) .header("MCP-Protocol-Version", "2025-11-25") .send() .await .unwrap(); assert_eq!(response.status(), reqwest::StatusCode::NOT_FOUND); } #[tokio::test] async fn rejects_rapid_transport_get_requests_with_429() { let registry = test_registry().await; publish_agent_with_bindings(®istry, "sales-get-rate-limit", vec![]).await; let api_key = create_platform_api_key( ®istry, "sales-get-rate-limit", "mcp-get-rate-limit", &[PlatformApiKeyScope::Read, PlatformApiKeyScope::Write], ) .await; let base_url = spawn_mcp_server(build_test_app_with_rate_limit( registry, Duration::from_millis(0), Some("https://crank.example.com".to_owned()), RequestRateLimitConfig::new(1, 2).unwrap(), )) .await; let client = reqwest::Client::new(); let mcp_url = agent_mcp_url(&base_url, "sales-get-rate-limit"); let initialized_session = initialize_session(&client, &mcp_url, &api_key).await; let first_response = client .get(&mcp_url) .header(header::ACCEPT, "text/event-stream") .header(header::AUTHORIZATION, format!("Bearer {api_key}")) .header("MCP-Session-Id", &initialized_session) .header("MCP-Protocol-Version", "2025-11-25") .send() .await .unwrap(); assert_eq!(first_response.status(), reqwest::StatusCode::OK); let second_response = client .get(&mcp_url) .header(header::ACCEPT, "text/event-stream") .header(header::AUTHORIZATION, format!("Bearer {api_key}")) .header("MCP-Session-Id", &initialized_session) .header("MCP-Protocol-Version", "2025-11-25") .send() .await .unwrap(); assert_eq!( second_response.status(), reqwest::StatusCode::TOO_MANY_REQUESTS ); assert!(second_response.headers().get(header::RETRY_AFTER).is_some()); } #[tokio::test] async fn get_requires_session_header() { let registry = test_registry().await; publish_agent_with_bindings(®istry, "sales-get-sse-missing", vec![]).await; let api_key = create_platform_api_key( ®istry, "sales-get-sse-missing", "mcp-get-sse-missing", &[PlatformApiKeyScope::Read], ) .await; let base_url = spawn_mcp_server(build_test_app( registry, Duration::from_millis(0), Some("https://crank.example.com".to_owned()), )) .await; let client = reqwest::Client::new(); let response = client .get(agent_mcp_url(&base_url, "sales-get-sse-missing")) .header(header::ACCEPT, "text/event-stream") .header(header::AUTHORIZATION, format!("Bearer {api_key}")) .send() .await .unwrap(); assert_eq!(response.status(), reqwest::StatusCode::BAD_REQUEST); } #[tokio::test] async fn delete_terminates_transport_session() { let registry = test_registry().await; publish_agent_with_bindings(®istry, "sales-delete-session", vec![]).await; let api_key = create_platform_api_key( ®istry, "sales-delete-session", "mcp-delete-session", &[PlatformApiKeyScope::Read], ) .await; let base_url = spawn_mcp_server(build_test_app( registry, Duration::from_millis(0), Some("https://crank.example.com".to_owned()), )) .await; let client = reqwest::Client::new(); let mcp_url = agent_mcp_url(&base_url, "sales-delete-session"); let initialized_session = initialize_session(&client, &mcp_url, &api_key).await; let delete_response = client .delete(&mcp_url) .header(header::AUTHORIZATION, format!("Bearer {api_key}")) .header("MCP-Session-Id", &initialized_session) .header("MCP-Protocol-Version", "2025-11-25") .send() .await .unwrap(); assert_eq!(delete_response.status(), reqwest::StatusCode::NO_CONTENT); let after_delete = client .get(&mcp_url) .header(header::ACCEPT, "text/event-stream") .header(header::AUTHORIZATION, format!("Bearer {api_key}")) .header("MCP-Session-Id", &initialized_session) .header("MCP-Protocol-Version", "2025-11-25") .send() .await .unwrap(); assert_eq!(after_delete.status(), reqwest::StatusCode::NOT_FOUND); } #[tokio::test] async fn initialize_accepts_json_only_response_negotiation() { let registry = test_registry().await; publish_agent_with_bindings(®istry, "sales-json-accept", vec![]).await; let api_key = create_platform_api_key( ®istry, "sales-json-accept", "mcp-json-accept", &[PlatformApiKeyScope::Read], ) .await; let base_url = spawn_mcp_server(build_test_app( registry, Duration::from_millis(0), Some("https://crank.example.com".to_owned()), )) .await; let client = reqwest::Client::new(); let response = client .post(agent_mcp_url(&base_url, "sales-json-accept")) .header(header::ACCEPT, "application/json") .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!(response.status(), reqwest::StatusCode::OK); assert_eq!( response .headers() .get(header::CONTENT_TYPE) .and_then(|value| value.to_str().ok()) .unwrap(), "application/json" ); } #[tokio::test] async fn rejects_get_with_protocol_version_mismatch() { let registry = test_registry().await; publish_agent_with_bindings(®istry, "sales-get-bad-version", vec![]).await; let api_key = create_platform_api_key( ®istry, "sales-get-bad-version", "mcp-get-bad-version", &[PlatformApiKeyScope::Read], ) .await; let base_url = spawn_mcp_server(build_test_app( registry, Duration::from_millis(0), Some("https://crank.example.com".to_owned()), )) .await; let client = reqwest::Client::new(); let mcp_url = agent_mcp_url(&base_url, "sales-get-bad-version"); let initialized_session = initialize_session(&client, &mcp_url, &api_key).await; let response = client .get(&mcp_url) .header(header::ACCEPT, "text/event-stream") .header(header::AUTHORIZATION, format!("Bearer {api_key}")) .header("MCP-Session-Id", &initialized_session) .header("MCP-Protocol-Version", "2025-06-18") .send() .await .unwrap(); assert_eq!(response.status(), reqwest::StatusCode::BAD_REQUEST); }