bsodfather/crank
1
0
Fork 0
Code Issues Pull Requests 3 Actions Packages Projects Releases Wiki Activity
Files
dc8171d052d01db84eb99856bc30268ee1c78ed0
crank/apps/mcp-server/tests/integration/catalog_access.rs
T

842 lines
25 KiB
Rust
Raw Blame History

#![allow(dead_code, unused_imports)]
mod approval_access;
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, ApprovalRequest,
ApprovalRequestId, ApprovalRequestStatus, ExecutionConfig, HttpMethod, InvocationSource,
Operation, OperationApprovalMode, OperationApprovalPayloadPreviewMode, OperationApprovalPolicy,
OperationApprovalRiskLevel, OperationId, OperationStatus, PlatformApiKey, PlatformApiKeyId,
PlatformApiKeyKind, PlatformApiKeyScope, PlatformApiKeyStatus, Protocol, RestTarget, Target,
ToolDescription, WorkspaceId,
};
use crank_mapping::{MappingRule, MappingSet};
use crank_registry::{
CreateAgentDraftVersionRequest, CreateAgentRequest, CreateApprovalRequest,
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-00000000000000000000000000000000").unwrap(),
crank_runtime::community_with_outbound_policy(
crank_runtime::OutboundHttpPolicy::allowing_hosts(["127.0.0.1"]),
)
.build(),
RequestRateLimiter::new(rate_limit_config),
std::sync::Arc::new(InMemoryCoordinationStateStore::default()),
sessions,
credential_verifier,
)
}
async fn assert_revoked_mcp_request_is_denied(
response: reqwest::Response,
api_key: &str,
key_id: &PlatformApiKeyId,
) {
assert_eq!(response.status(), reqwest::StatusCode::UNAUTHORIZED);
let denied_body = response.text().await.unwrap();
assert!(!denied_body.contains(api_key));
assert!(!denied_body.contains(key_id.as_str()));
}
#[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(&registry, "sales-refresh", vec![]).await;
let api_key = create_platform_api_key(
&registry,
"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();
let agent_id = test_agent_id("sales-refresh");
let draft_v2 = AgentVersion {
agent_id: agent_id.clone(),
version: 2,
status: AgentStatus::Draft,
instructions: json!({}),
tool_selection_policy: Default::default(),
created_at: OffsetDateTime::parse("2026-03-26T10:01:00Z", &Rfc3339).unwrap(),
};
registry
.create_agent_draft_version(CreateAgentDraftVersionRequest {
workspace_id: &test_workspace_id(),
agent_id: &agent_id,
version: &draft_v2,
bindings: &[binding_for_operation(&operation)],
updated_at: &OffsetDateTime::parse("2026-03-26T10:01:00Z", &Rfc3339).unwrap(),
})
.await
.unwrap();
registry
.publish_agent(PublishAgentRequest {
workspace_id: &test_workspace_id(),
agent_id: &agent_id,
version: 2,
published_at: &OffsetDateTime::parse("2026-03-26T10:02:00Z", &Rfc3339).unwrap(),
published_by: Some("alice"),
expected_state: None,
})
.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 shared_catalog_cache_does_not_serve_stale_snapshot_after_unpublish() {
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(&registry, &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(),
None,
)
.await
.unwrap();
let catalog_b = PublishedToolCatalog::new(
registry.clone(),
Duration::from_secs(60),
coordination_store,
);
let error = catalog_b
.list_tools(test_workspace_slug(), "sales-shared-catalog")
.await
.expect_err("unpublished Agent must not be served from shared cache");
assert!(matches!(
error,
crank_registry::RegistryError::PublishedAgentNotFound { .. }
));
}
#[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(
&registry,
"sales-a",
vec![binding_for_operation(&operation_a)],
)
.await;
publish_agent_with_bindings(
&registry,
"sales-b",
vec![binding_for_operation(&operation_b)],
)
.await;
let api_key_a = create_platform_api_key(
&registry,
"sales-a",
"mcp-agent-a",
&[PlatformApiKeyScope::Read, PlatformApiKeyScope::Write],
)
.await;
let api_key_b = create_platform_api_key(
&registry,
"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(
&registry,
"sales-a",
vec![binding_for_operation(&operation_a)],
)
.await;
publish_agent_with_bindings(
&registry,
"sales-b",
vec![binding_for_operation(&operation_b)],
)
.await;
let api_key_a = create_platform_api_key(
&registry,
"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 revoked_mcp_key_stops_existing_session_without_restart() {
let registry = test_registry().await;
let upstream_base_url = spawn_upstream_server().await;
let operation = test_operation(&upstream_base_url, "crm_revoked_key");
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(&registry, &operation, "sales-revoked-key").await;
let api_key = create_platform_api_key(
&registry,
"sales-revoked-key",
"mcp-revoked-boundary",
&[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-revoked-key");
let initialized_session = initialize_session(&client, &mcp_url, &api_key).await;
let before_revoke = post_jsonrpc(
&client,
&mcp_url,
&api_key,
Some(&initialized_session),
json!({
"jsonrpc": "2.0",
"id": 2,
"method": "tools/list",
"params": {}
}),
)
.await;
assert_eq!(
before_revoke["result"]["tools"][0]["name"],
"crm_revoked_key"
);
let key_id = registry
.list_platform_api_keys_for_agent(&test_workspace_id(), &test_agent_id("sales-revoked-key"))
.await
.unwrap()
.into_iter()
.find(|record| record.api_key.name == "mcp-revoked-boundary")
.unwrap()
.api_key
.id;
registry
.revoke_platform_api_key_for_agent(
&test_workspace_id(),
&test_agent_id("sales-revoked-key"),
&key_id,
&OffsetDateTime::now_utc(),
)
.await
.unwrap();
let rejected_initialize = post_jsonrpc_response(
&client,
&mcp_url,
&api_key,
None,
None,
json!({
"jsonrpc": "2.0",
"id": 3,
"method": "initialize",
"params": {
"protocolVersion": "2025-11-25"
}
}),
)
.await;
assert_revoked_mcp_request_is_denied(rejected_initialize, &api_key, &key_id).await;
let rejected_tools_list = post_jsonrpc_response(
&client,
&mcp_url,
&api_key,
Some(&initialized_session),
None,
json!({
"jsonrpc": "2.0",
"id": 4,
"method": "tools/list",
"params": {}
}),
)
.await;
assert_revoked_mcp_request_is_denied(rejected_tools_list, &api_key, &key_id).await;
let rejected_tools_call = post_jsonrpc_response(
&client,
&mcp_url,
&api_key,
Some(&initialized_session),
None,
json!({
"jsonrpc": "2.0",
"id": 5,
"method": "tools/call",
"params": {
"name": "crm_revoked_key",
"arguments": {}
}
}),
)
.await;
assert_revoked_mcp_request_is_denied(rejected_tools_call, &api_key, &key_id).await;
}
#[tokio::test]
async fn rejects_initialize_without_platform_api_key() {
let registry = test_registry().await;
publish_agent_with_bindings(&registry, "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_initialize_with_approval_platform_api_key() {
let registry = test_registry().await;
publish_agent_with_bindings(&registry, "sales-approval-key", vec![]).await;
let api_key =
create_approval_platform_api_key(&registry, "sales-approval-key", "approval-only").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-approval-key"))
.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!(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(&registry, &operation, "sales-read-only").await;
let api_key = create_platform_api_key(
&registry,
"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(&registry, &operation, "sales-rate-limited").await;
let api_key = create_platform_api_key(
&registry,
"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));
}
Reference in New Issue View Git Blame Copy Permalink