github-ops
1901 lines
65 KiB
Rust
1901 lines
65 KiB
Rust
use std::{env, net::SocketAddr, time::Duration};
|
|
|
|
use crank_community_mcp::{
|
|
auth::CommunityMachineCredentialVerifier, build_app, session::PostgresTransportSessionStore,
|
|
};
|
|
use crank_registry::{PostgresPoolConfig, PostgresRegistry};
|
|
use crank_runtime::{
|
|
RequestRateLimitConfig, RequestRateLimiter, RuntimeCacheConfig, RuntimeCacheStores,
|
|
RuntimeLimits, SecretCrypto, community_default,
|
|
};
|
|
use sqlx::postgres::PgConnectOptions;
|
|
use tokio::net::TcpListener;
|
|
use tracing::info;
|
|
|
|
#[tokio::main]
|
|
async fn main() -> Result<(), Box<dyn std::error::Error>> {
|
|
tracing_subscriber::fmt()
|
|
.with_env_filter(
|
|
env::var("CRANK_LOG_LEVEL")
|
|
.unwrap_or_else(|_| "mcp_server=info,tower_http=info".into()),
|
|
)
|
|
.init();
|
|
|
|
let bind_addr = env::var("CRANK_MCP_BIND").unwrap_or_else(|_| "0.0.0.0:3002".into());
|
|
let base_url = env::var("CRANK_BASE_URL").ok();
|
|
let refresh_interval = env::var("CRANK_MCP_REFRESH_MS")
|
|
.ok()
|
|
.and_then(|value| value.parse::<u64>().ok())
|
|
.map(Duration::from_millis)
|
|
.unwrap_or_else(|| Duration::from_secs(5));
|
|
let socket_addr: SocketAddr = bind_addr.parse()?;
|
|
let pool_config = PostgresPoolConfig::from_env()?;
|
|
let runtime_limits = RuntimeLimits::from_env()?;
|
|
let cache_config = RuntimeCacheConfig::from_env()?;
|
|
let cache_stores = RuntimeCacheStores::from_config(&cache_config).await?;
|
|
let api_rate_limit = mcp_api_rate_limit_config_from_env()?;
|
|
let database_options = database_options_from_env()?;
|
|
let registry = PostgresRegistry::connect_with_options_and_pool_config(
|
|
database_options.clone(),
|
|
pool_config,
|
|
)
|
|
.await?;
|
|
let session_store = PostgresTransportSessionStore::connect_with_options_and_pool_config(
|
|
database_options,
|
|
pool_config,
|
|
)
|
|
.await?;
|
|
let secret_crypto = SecretCrypto::new(&env::var("CRANK_MASTER_KEY")?)?;
|
|
let runtime = community_default()
|
|
.with_limits(runtime_limits)
|
|
.with_response_cache(cache_stores.response.clone())
|
|
.build();
|
|
let app = build_app(
|
|
registry,
|
|
refresh_interval,
|
|
base_url,
|
|
secret_crypto,
|
|
runtime,
|
|
if cache_config.backend.is_external() {
|
|
RequestRateLimiter::new_shared(api_rate_limit, cache_stores.rate_limit.clone())
|
|
} else {
|
|
RequestRateLimiter::new(api_rate_limit)
|
|
},
|
|
cache_stores.coordination.clone(),
|
|
std::sync::Arc::new(session_store),
|
|
std::sync::Arc::new(CommunityMachineCredentialVerifier),
|
|
);
|
|
let listener = TcpListener::bind(socket_addr).await?;
|
|
|
|
info!(
|
|
runtime_max_concurrent_unary = runtime_limits.max_concurrent_unary,
|
|
mcp_rate_limit_rps = api_rate_limit.requests_per_second,
|
|
mcp_rate_limit_burst = api_rate_limit.burst,
|
|
cache_backend = %cache_config.backend,
|
|
max_connections = pool_config.max_connections,
|
|
min_connections = pool_config.min_connections,
|
|
acquire_timeout_ms = pool_config.acquire_timeout_ms,
|
|
idle_timeout_ms = pool_config.idle_timeout_ms,
|
|
max_lifetime_ms = pool_config.max_lifetime_ms,
|
|
"postgres pool configured"
|
|
);
|
|
info!("mcp-server listening on {}", socket_addr);
|
|
|
|
axum::serve(listener, app).await?;
|
|
|
|
Ok(())
|
|
}
|
|
|
|
fn database_options_from_env() -> Result<PgConnectOptions, Box<dyn std::error::Error>> {
|
|
if let Ok(database_url) = env::var("CRANK_DATABASE_URL") {
|
|
return Ok(database_url.parse::<PgConnectOptions>()?);
|
|
}
|
|
|
|
let host = env::var("POSTGRES_HOST").unwrap_or_else(|_| "postgres".into());
|
|
let port = env::var("POSTGRES_PORT")
|
|
.ok()
|
|
.and_then(|value| value.parse::<u16>().ok())
|
|
.unwrap_or(5432);
|
|
let database = env::var("POSTGRES_DB").unwrap_or_else(|_| "crank".into());
|
|
let username = env::var("POSTGRES_USER").unwrap_or_else(|_| "crank".into());
|
|
let password = env::var("POSTGRES_PASSWORD").unwrap_or_else(|_| "crank".into());
|
|
|
|
Ok(PgConnectOptions::new()
|
|
.host(&host)
|
|
.port(port)
|
|
.database(&database)
|
|
.username(&username)
|
|
.password(&password))
|
|
}
|
|
|
|
fn mcp_api_rate_limit_config_from_env() -> Result<RequestRateLimitConfig, Box<dyn std::error::Error>>
|
|
{
|
|
let requests_per_second = env::var("CRANK_MCP_RATE_LIMIT_RPS")
|
|
.ok()
|
|
.and_then(|value| value.parse::<u32>().ok())
|
|
.unwrap_or(60);
|
|
let burst = env::var("CRANK_MCP_RATE_LIMIT_BURST")
|
|
.ok()
|
|
.and_then(|value| value.parse::<u32>().ok())
|
|
.unwrap_or(120);
|
|
|
|
Ok(RequestRateLimitConfig::new(requests_per_second, burst)?)
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use std::{
|
|
collections::BTreeMap,
|
|
env, 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 sqlx::{Executor, postgres::PgPoolOptions};
|
|
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,
|
|
)
|
|
}
|
|
|
|
#[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::<Value>().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::<Value>().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 dispatch = tracing::Dispatch::new(subscriber);
|
|
|
|
let _guard = tracing::dispatcher::set_default(&dispatch);
|
|
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"));
|
|
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::<Value>()
|
|
.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);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn refreshes_published_tools_without_restart() {
|
|
let registry = test_registry().await;
|
|
let upstream_base_url = spawn_upstream_server().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-refresh");
|
|
publish_agent_with_bindings(®istry, "sales-refresh", vec![]).await;
|
|
let api_key = create_platform_api_key(
|
|
®istry,
|
|
"sales-refresh",
|
|
"mcp-refresh",
|
|
&[PlatformApiKeyScope::Read, PlatformApiKeyScope::Write],
|
|
)
|
|
.await;
|
|
let initialized_session = initialize_session(&client, &mcp_url, &api_key).await;
|
|
|
|
let before_publish = post_jsonrpc(
|
|
&client,
|
|
&mcp_url,
|
|
&api_key,
|
|
Some(&initialized_session),
|
|
json!({
|
|
"jsonrpc": "2.0",
|
|
"id": 2,
|
|
"method": "tools/list",
|
|
"params": {}
|
|
}),
|
|
)
|
|
.await;
|
|
|
|
let operation = test_operation(&upstream_base_url, "crm_publish_later");
|
|
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();
|
|
registry
|
|
.save_agent_bindings(SaveAgentBindingsRequest {
|
|
workspace_id: &test_workspace_id(),
|
|
agent_id: &test_agent_id("sales-refresh"),
|
|
agent_version: 1,
|
|
bindings: &[binding_for_operation(&operation)],
|
|
})
|
|
.await
|
|
.unwrap();
|
|
|
|
let after_publish = post_jsonrpc(
|
|
&client,
|
|
&mcp_url,
|
|
&api_key,
|
|
Some(&initialized_session),
|
|
json!({
|
|
"jsonrpc": "2.0",
|
|
"id": 3,
|
|
"method": "tools/list",
|
|
"params": {}
|
|
}),
|
|
)
|
|
.await;
|
|
|
|
assert_eq!(before_publish["result"]["tools"], json!([]));
|
|
assert_eq!(
|
|
after_publish["result"]["tools"][0]["name"],
|
|
"crm_publish_later"
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn shares_published_catalog_snapshot_across_instances() {
|
|
let registry = test_registry().await;
|
|
let upstream_base_url = spawn_upstream_server().await;
|
|
let operation = test_operation(&upstream_base_url, "crm_catalog_shared");
|
|
let coordination_store = std::sync::Arc::new(InMemoryCoordinationStateStore::default());
|
|
|
|
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-shared-catalog").await;
|
|
|
|
let catalog_a = PublishedToolCatalog::new(
|
|
registry.clone(),
|
|
Duration::from_secs(60),
|
|
coordination_store.clone(),
|
|
);
|
|
let tools_a = catalog_a
|
|
.list_tools(test_workspace_slug(), "sales-shared-catalog")
|
|
.await
|
|
.unwrap();
|
|
assert_eq!(tools_a.len(), 1);
|
|
|
|
registry
|
|
.unpublish_agent(
|
|
&test_workspace_id(),
|
|
&test_agent_id("sales-shared-catalog"),
|
|
&OffsetDateTime::parse("2026-03-26T10:05:00Z", &Rfc3339).unwrap(),
|
|
)
|
|
.await
|
|
.unwrap();
|
|
|
|
let catalog_b = PublishedToolCatalog::new(
|
|
registry.clone(),
|
|
Duration::from_secs(60),
|
|
coordination_store,
|
|
);
|
|
let tools_b = catalog_b
|
|
.list_tools(test_workspace_slug(), "sales-shared-catalog")
|
|
.await
|
|
.unwrap();
|
|
|
|
assert_eq!(tools_b.len(), 1);
|
|
assert_eq!(tools_b[0].tool_name, operation.name);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn lists_only_bound_tools_for_agent_context() {
|
|
let registry = test_registry().await;
|
|
let upstream_base_url = spawn_upstream_server().await;
|
|
let operation_a = test_operation(&upstream_base_url, "crm_create_lead");
|
|
let operation_b = test_operation(&upstream_base_url, "crm_update_lead");
|
|
|
|
for operation in [&operation_a, &operation_b] {
|
|
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_with_bindings(
|
|
®istry,
|
|
"sales-a",
|
|
vec![binding_for_operation(&operation_a)],
|
|
)
|
|
.await;
|
|
publish_agent_with_bindings(
|
|
®istry,
|
|
"sales-b",
|
|
vec![binding_for_operation(&operation_b)],
|
|
)
|
|
.await;
|
|
let api_key_a = create_platform_api_key(
|
|
®istry,
|
|
"sales-a",
|
|
"mcp-agent-a",
|
|
&[PlatformApiKeyScope::Read, PlatformApiKeyScope::Write],
|
|
)
|
|
.await;
|
|
let api_key_b = create_platform_api_key(
|
|
®istry,
|
|
"sales-b",
|
|
"mcp-agent-b",
|
|
&[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 agent_a_url = agent_mcp_url(&base_url, "sales-a");
|
|
let agent_b_url = agent_mcp_url(&base_url, "sales-b");
|
|
let session_a = initialize_session(&client, &agent_a_url, &api_key_a).await;
|
|
let session_b = initialize_session(&client, &agent_b_url, &api_key_b).await;
|
|
|
|
let tools_a = post_jsonrpc(
|
|
&client,
|
|
&agent_a_url,
|
|
&api_key_a,
|
|
Some(&session_a),
|
|
json!({
|
|
"jsonrpc": "2.0",
|
|
"id": 2,
|
|
"method": "tools/list",
|
|
"params": {}
|
|
}),
|
|
)
|
|
.await;
|
|
let tools_b = post_jsonrpc(
|
|
&client,
|
|
&agent_b_url,
|
|
&api_key_b,
|
|
Some(&session_b),
|
|
json!({
|
|
"jsonrpc": "2.0",
|
|
"id": 2,
|
|
"method": "tools/list",
|
|
"params": {}
|
|
}),
|
|
)
|
|
.await;
|
|
|
|
assert_eq!(tools_a["result"]["tools"][0]["name"], "crm_create_lead");
|
|
assert_eq!(tools_b["result"]["tools"][0]["name"], "crm_update_lead");
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn rejects_initialize_with_key_from_different_agent() {
|
|
let registry = test_registry().await;
|
|
let upstream_base_url = spawn_upstream_server().await;
|
|
let operation_a = test_operation(&upstream_base_url, "crm_create_lead");
|
|
let operation_b = test_operation(&upstream_base_url, "crm_update_lead");
|
|
|
|
for operation in [&operation_a, &operation_b] {
|
|
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_with_bindings(
|
|
®istry,
|
|
"sales-a",
|
|
vec![binding_for_operation(&operation_a)],
|
|
)
|
|
.await;
|
|
publish_agent_with_bindings(
|
|
®istry,
|
|
"sales-b",
|
|
vec![binding_for_operation(&operation_b)],
|
|
)
|
|
.await;
|
|
let api_key_a = create_platform_api_key(
|
|
®istry,
|
|
"sales-a",
|
|
"mcp-agent-a-only",
|
|
&[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-b"))
|
|
.header(header::ACCEPT, "application/json, text/event-stream")
|
|
.header(header::AUTHORIZATION, format!("Bearer {api_key_a}"))
|
|
.json(&json!({
|
|
"jsonrpc": "2.0",
|
|
"id": 1,
|
|
"method": "initialize",
|
|
"params": {
|
|
"protocolVersion": "2025-11-25"
|
|
}
|
|
}))
|
|
.send()
|
|
.await
|
|
.unwrap();
|
|
|
|
assert_eq!(response.status(), reqwest::StatusCode::UNAUTHORIZED);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn rejects_initialize_without_platform_api_key() {
|
|
let registry = test_registry().await;
|
|
publish_agent_with_bindings(®istry, "sales-auth", vec![]).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-auth"))
|
|
.header(header::ACCEPT, "application/json, text/event-stream")
|
|
.json(&json!({
|
|
"jsonrpc": "2.0",
|
|
"id": 1,
|
|
"method": "initialize",
|
|
"params": {
|
|
"protocolVersion": "2025-11-25"
|
|
}
|
|
}))
|
|
.send()
|
|
.await
|
|
.unwrap();
|
|
|
|
assert_eq!(response.status(), reqwest::StatusCode::UNAUTHORIZED);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn rejects_tool_call_with_read_only_platform_api_key() {
|
|
let registry = test_registry().await;
|
|
let upstream_base_url = spawn_upstream_server().await;
|
|
let operation = test_operation(&upstream_base_url, "crm_read_only");
|
|
|
|
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-read-only").await;
|
|
let api_key = create_platform_api_key(
|
|
®istry,
|
|
"sales-read-only",
|
|
"mcp-read",
|
|
&[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-read-only");
|
|
let initialized_session = initialize_session(&client, &mcp_url, &api_key).await;
|
|
let response = client
|
|
.post(&mcp_url)
|
|
.header(header::ACCEPT, "application/json, text/event-stream")
|
|
.header(header::AUTHORIZATION, format!("Bearer {api_key}"))
|
|
.header("MCP-Session-Id", initialized_session)
|
|
.header("MCP-Protocol-Version", "2025-11-25")
|
|
.json(&json!({
|
|
"jsonrpc": "2.0",
|
|
"id": 3,
|
|
"method": "tools/call",
|
|
"params": {
|
|
"name": "crm_read_only",
|
|
"arguments": {
|
|
"email": "user@example.com"
|
|
}
|
|
}
|
|
}))
|
|
.send()
|
|
.await
|
|
.unwrap();
|
|
|
|
assert_eq!(response.status(), reqwest::StatusCode::FORBIDDEN);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn rejects_rapid_initialize_requests_with_429() {
|
|
let registry = test_registry().await;
|
|
let upstream_base_url = spawn_upstream_server().await;
|
|
let operation = test_operation(&upstream_base_url, "crm_rate_limited");
|
|
|
|
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-rate-limited").await;
|
|
let api_key = create_platform_api_key(
|
|
®istry,
|
|
"sales-rate-limited",
|
|
"mcp-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, 1).unwrap(),
|
|
))
|
|
.await;
|
|
let client = reqwest::Client::new();
|
|
let mcp_url = agent_mcp_url(&base_url, "sales-rate-limited");
|
|
|
|
let first_response = client
|
|
.post(&mcp_url)
|
|
.header(header::ACCEPT, "application/json, text/event-stream")
|
|
.header(header::AUTHORIZATION, format!("Bearer {api_key}"))
|
|
.header("x-request-id", "req_rate_limit_01")
|
|
.json(&json!({
|
|
"jsonrpc": "2.0",
|
|
"id": 1,
|
|
"method": "initialize",
|
|
"params": {
|
|
"protocolVersion": "2025-11-25"
|
|
}
|
|
}))
|
|
.send()
|
|
.await
|
|
.unwrap();
|
|
assert_eq!(first_response.status(), reqwest::StatusCode::OK);
|
|
|
|
let second_response = client
|
|
.post(&mcp_url)
|
|
.header(header::ACCEPT, "application/json, text/event-stream")
|
|
.header(header::AUTHORIZATION, format!("Bearer {api_key}"))
|
|
.header("x-request-id", "req_rate_limit_02")
|
|
.json(&json!({
|
|
"jsonrpc": "2.0",
|
|
"id": 2,
|
|
"method": "initialize",
|
|
"params": {
|
|
"protocolVersion": "2025-11-25"
|
|
}
|
|
}))
|
|
.send()
|
|
.await
|
|
.unwrap();
|
|
|
|
assert_eq!(
|
|
second_response.status(),
|
|
reqwest::StatusCode::TOO_MANY_REQUESTS
|
|
);
|
|
assert_eq!(
|
|
second_response.headers()["x-request-id"].to_str().unwrap(),
|
|
"req_rate_limit_02"
|
|
);
|
|
assert_eq!(
|
|
second_response.headers()["retry-after"].to_str().unwrap(),
|
|
"1"
|
|
);
|
|
|
|
let payload = second_response.json::<Value>().await.unwrap();
|
|
assert_eq!(payload["error"]["code"], json!(-32029));
|
|
assert_eq!(
|
|
payload["error"]["data"]["code"],
|
|
json!("request_rate_limited")
|
|
);
|
|
let retry_after_ms = payload["error"]["data"]["retry_after_ms"].as_u64().unwrap();
|
|
assert!((1..=1000).contains(&retry_after_ms));
|
|
}
|
|
|
|
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
|
|
}
|
|
|
|
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()
|
|
}
|
|
|
|
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()
|
|
}
|
|
|
|
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
|
|
}
|
|
|
|
fn hash_access_secret(secret: &str) -> String {
|
|
let digest = Sha256::digest(secret.as_bytes());
|
|
URL_SAFE_NO_PAD.encode(digest)
|
|
}
|
|
|
|
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)
|
|
}
|
|
|
|
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)
|
|
}
|
|
|
|
fn agent_mcp_url(base_url: &str, agent_slug: &str) -> String {
|
|
format!("{base_url}/v1/{}/{}", test_workspace_slug(), agent_slug)
|
|
}
|
|
|
|
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;
|
|
}
|
|
|
|
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();
|
|
}
|
|
|
|
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,
|
|
}
|
|
}
|
|
|
|
async fn create_lead(Json(payload): Json<Value>) -> Json<Value> {
|
|
Json(json!({
|
|
"id": "lead_123",
|
|
"email": payload["email"]
|
|
}))
|
|
}
|
|
|
|
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"]
|
|
}))
|
|
}
|
|
|
|
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)))
|
|
}
|
|
|
|
async fn test_registry() -> PostgresRegistry {
|
|
let database_url = env::var("TEST_DATABASE_URL")
|
|
.unwrap_or_else(|_| "postgres://crank:crank@127.0.0.1:15432/crank".to_owned());
|
|
let admin_pool = PgPoolOptions::new()
|
|
.max_connections(1)
|
|
.connect(&database_url)
|
|
.await
|
|
.unwrap();
|
|
let schema = format!("test_mcp_server_{}", uuid::Uuid::now_v7().simple());
|
|
|
|
admin_pool
|
|
.execute(sqlx::query(&format!("create schema {schema}")))
|
|
.await
|
|
.unwrap();
|
|
|
|
PostgresRegistry::connect(&format!("{database_url}?options=-csearch_path%3D{schema}"))
|
|
.await
|
|
.unwrap()
|
|
}
|
|
|
|
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,
|
|
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()),
|
|
}
|
|
}
|
|
|
|
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(),
|
|
}
|
|
}
|
|
}
|