bsodfather/crank
1
0
Fork 0
Code Issues Pull Requests 3 Actions Packages Projects Releases Wiki Activity

feat: harden community production foundation through story 1.5

Browse Source
This commit is contained in:
bsodfather
2026-08-14 00:21:59 +03:00
parent c30461cc92
commit f6fc2e5c9b
161 changed files with 16758 additions and 2515 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/mcp-server/Cargo.toml
+1
View File
@@ -15,6 +15,7 @@ async-trait = "0.1"
axum.workspace = true
base64.workspace = true
crank-community-mcp = { path = "../../crates/crank-community-mcp" }
crank-config = { path = "../../crates/crank-config" }
crank-core = { path = "../../crates/crank-core" }
crank-observability = { path = "../../crates/crank-observability" }
crank-registry = { path = "../../crates/crank-registry" }
apps/mcp-server/Dockerfile
+6 -6
View File
@@ -31,9 +31,9 @@ RUN mkdir -p \
&& printf 'pub fn placeholder() {}\n' > crates/crank-runtime/src/lib.rs \
&& printf 'pub fn placeholder() {}\n' > crates/crank-adapter-rest/src/lib.rs
RUN --mount=type=cache,target=/usr/local/cargo/registry \
--mount=type=cache,target=/usr/local/cargo/git/db \
--mount=type=cache,target=/app/target \
RUN --mount=type=cache,id=crank-mcp-cargo-registry,target=/usr/local/cargo/registry \
--mount=type=cache,id=crank-mcp-cargo-git,target=/usr/local/cargo/git/db \
--mount=type=cache,id=crank-mcp-target,target=/app/target \
SQLX_OFFLINE=true cargo build --release -p mcp-server
FROM rust:1.96.1-bookworm AS builder
@@ -45,9 +45,9 @@ COPY .sqlx ./.sqlx
COPY apps ./apps
COPY crates ./crates
RUN --mount=type=cache,target=/usr/local/cargo/registry \
--mount=type=cache,target=/usr/local/cargo/git/db \
--mount=type=cache,target=/app/target \
RUN --mount=type=cache,id=crank-mcp-cargo-registry,target=/usr/local/cargo/registry \
--mount=type=cache,id=crank-mcp-cargo-git,target=/usr/local/cargo/git/db \
--mount=type=cache,id=crank-mcp-target,target=/app/target \
SQLX_OFFLINE=true cargo build --release -p mcp-server \
&& cp /app/target/release/mcp-server /tmp/mcp-server
apps/mcp-server/src/main.rs
+280 -71
View File
@@ -1,12 +1,17 @@
use std::{env, net::SocketAddr, time::Duration};
use std::{io, process::ExitCode, time::Duration};
use crank_community_mcp::{
auth::CommunityMachineCredentialVerifier, build_app_with_background_workers_and_limits,
session::PostgresTransportSessionStore,
};
use crank_config::{
CacheBackend as ConfigCacheBackend, ConfigSource, DatabaseSettings, DiagnosticCode,
McpProcessConfig, ObservabilitySettings, ProcessKind, parse_process,
};
use crank_core::CacheBackend;
use crank_observability::{
CriticalErrorCategory, MetricsConfig, ObservabilityConfig, ObservabilityLifecycle,
capture_critical_error,
OtlpTraceConfig, RedactionLimits, SentryConfig, ServiceIdentity, capture_critical_error,
};
use crank_registry::{PostgresPoolConfig, PostgresRegistry};
use crank_runtime::{
@@ -19,14 +24,76 @@ use tokio::net::TcpListener;
use tracing::info;
#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
let observability = crank_observability::init(ObservabilityConfig::from_env(
"mcp-server",
env!("CARGO_PKG_VERSION"),
"mcp_server=info,tower_http=info",
)?)?;
async fn main() -> ExitCode {
match main_result().await {
Ok(()) => ExitCode::SUCCESS,
Err(error) => {
eprintln!("{}", safe_startup_diagnostic(error.as_ref()));
ExitCode::FAILURE
}
}
}
fn safe_startup_diagnostic(error: &(dyn std::error::Error + 'static)) -> String {
let mut current = Some(error);
while let Some(cause) = current {
if let Some(config) = cause.downcast_ref::<crank_config::ConfigError>() {
return config.to_json();
}
if let Some(migration) = cause.downcast_ref::<crank_registry::MigrationError>() {
return serde_json::json!({
"status": "error",
"code": migration.code(),
"stage": migration.stage(),
"version": migration.version(),
"recovery": migration.recovery(),
})
.to_string();
}
if let Some(crank_registry::RegistryError::Migration(migration)) =
cause.downcast_ref::<crank_registry::RegistryError>()
{
return serde_json::json!({
"status": "error",
"code": migration.code(),
"stage": migration.stage(),
"version": migration.version(),
"recovery": migration.recovery(),
})
.to_string();
}
current = cause.source();
}
serde_json::json!({
"status": "error",
"code": "startup_failed",
"stage": "startup",
"version": null,
"recovery": "contact_operator",
})
.to_string()
}
async fn main_result() -> Result<(), Box<dyn std::error::Error>> {
let effective = parse_process(ProcessKind::McpServer, ConfigSource::from_os()?)?;
let config = effective
.mcp()
.cloned()
.ok_or_else(|| io::Error::other("MCP configuration projection is unavailable"))?;
preflight_config(&config)?;
let observability = init_observability(&config.observability)?;
for deprecation in effective.deprecations() {
tracing::warn!(
name: "config.deprecated",
field = deprecation.field,
source_class = deprecation.source_class,
replacement = deprecation.replacement,
removal_window = deprecation.removal_window,
"deprecated configuration accepted"
);
}
let mut startup_completed = false;
let result = run(&observability, &mut startup_completed).await;
let result = run(config, &observability, &mut startup_completed).await;
if result.is_err() {
capture_critical_error(if startup_completed {
CriticalErrorCategory::Internal
@@ -38,49 +105,61 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
}
async fn run(
config: McpProcessConfig,
observability: &ObservabilityLifecycle,
startup_completed: &mut bool,
) -> Result<(), Box<dyn std::error::Error>> {
let metrics_config =
MetricsConfig::from_env("CRANK_MCP_METRICS_BIND", "127.0.0.1:9465".parse()?)?;
let metrics_config = MetricsConfig::new(
config.observability.metrics.enabled,
config.observability.metrics.bind_addr,
config
.observability
.metrics
.bearer_token
.as_ref()
.map(|token| token.expose_secret().to_owned()),
)?;
let metrics_enabled = metrics_config.enabled();
let metrics_server = if metrics_config.enabled() {
let pool_config = postgres_pool_config(&config.database)?;
let runtime_limits = RuntimeLimits::try_new(
config.runtime.max_concurrent_unary,
config.runtime.max_concurrent_sessions,
)?;
let cache_config = runtime_cache_config(&config)?;
let cache_stores = RuntimeCacheStores::from_config(&cache_config).await?;
let api_rate_limit = RequestRateLimitConfig::new(
config.rate_limit.requests_per_second,
config.rate_limit.burst,
)?;
let registry = PostgresRegistry::connect_with_options_and_pool_config(
database_options(&config.database)?,
pool_config,
)
.await?;
let metrics_server = if metrics_enabled {
Some(observability.metrics_surface(metrics_config).bind().await?)
} else {
None
};
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, pool_config)
.await?;
if metrics_enabled {
spawn_postgres_pool_metrics(registry.pool().clone());
}
let session_store = PostgresTransportSessionStore::from_pool(registry.pool().clone()).await?;
let secret_crypto = SecretCrypto::new(&env::var("CRANK_MASTER_KEY")?)?;
let runtime = crank_runtime::community_from_env()?
let secret_crypto = SecretCrypto::new(config.runtime.master_key.expose_secret())?;
let outbound_http_policy = crank_runtime::OutboundHttpPolicy::try_new(
config.runtime.outbound.allowed_hosts.clone(),
config.runtime.outbound.denied_hosts.clone(),
config.runtime.outbound.max_response_bytes,
)?;
let runtime = crank_runtime::community_with_outbound_policy(outbound_http_policy)
.with_limits(runtime_limits)
.with_response_cache(cache_stores.response.clone())
.with_coordination_store(cache_stores.coordination.clone())
.build();
let app = build_app_with_background_workers_and_limits(
registry,
refresh_interval,
base_url,
Duration::from_millis(config.refresh_ms),
config.runtime.base_url.clone(),
secret_crypto,
runtime,
if cache_config.backend.is_external() {
@@ -93,7 +172,7 @@ async fn run(
std::sync::Arc::new(CommunityMachineCredentialVerifier),
runtime_limits.max_concurrent_sessions,
);
let listener = TcpListener::bind(socket_addr).await?;
let listener = TcpListener::bind(config.bind_addr).await?;
info!(
name: "mcp.postgres_pool.configured",
@@ -109,11 +188,7 @@ async fn run(
max_lifetime_ms = pool_config.max_lifetime_ms,
"postgres pool configured"
);
info!(
name: "mcp.server.listening",
bind_address = %socket_addr,
"mcp-server listening"
);
info!(name: "mcp.server.listening", bind_address = %config.bind_addr, "mcp-server listening");
*startup_completed = true;
if let Some(metrics_server) = metrics_server {
@@ -124,42 +199,176 @@ async fn run(
} else {
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>()?);
fn init_observability(
config: &ObservabilitySettings,
) -> Result<ObservabilityLifecycle, Box<dyn std::error::Error>> {
let identity = ServiceIdentity::try_new(
"mcp-server",
env!("CARGO_PKG_VERSION"),
config.environment.clone(),
)?;
let base = ObservabilityConfig::try_new(
identity,
config.log_filter.clone(),
RedactionLimits::default(),
)?;
let sentry = SentryConfig::parse(
config
.sentry_dsn
.as_ref()
.map(|value| value.expose_secret()),
)?;
let values = &config.otlp;
let otlp = OtlpTraceConfig::from_values(
values.endpoint.clone(),
values.traces_endpoint.clone(),
values.protocol.clone(),
values.traces_protocol.clone(),
values.timeout.clone(),
values.traces_timeout.clone(),
values
.headers
.as_ref()
.map(|value| value.expose_secret().to_owned()),
values
.traces_headers
.as_ref()
.map(|value| value.expose_secret().to_owned()),
values.max_queue_size,
values.max_export_batch_size,
values.schedule_delay.clone(),
values.export_timeout.clone(),
)?;
Ok(ObservabilityLifecycle::init_with_exporters(
base, sentry, otlp,
)?)
}
fn preflight_config(config: &McpProcessConfig) -> Result<(), crank_config::ConfigError> {
let invalid = |field| crank_config::ConfigError::single(DiagnosticCode::InvalidType, field);
database_options(&config.database).map_err(|_| invalid("database.source"))?;
postgres_pool_config(&config.database).map_err(|_| invalid("database.pool"))?;
MetricsConfig::new(
config.observability.metrics.enabled,
config.observability.metrics.bind_addr,
config
.observability
.metrics
.bearer_token
.as_ref()
.map(|v| v.expose_secret().to_owned()),
)
.map_err(|_| invalid("observability.metrics"))?;
RuntimeLimits::try_new(
config.runtime.max_concurrent_unary,
config.runtime.max_concurrent_sessions,
)
.map_err(|_| invalid("runtime.limits"))?;
runtime_cache_config(config).map_err(|_| invalid("cache"))?;
RequestRateLimitConfig::new(
config.rate_limit.requests_per_second,
config.rate_limit.burst,
)
.map_err(|_| invalid("mcp.rate_limit"))?;
SecretCrypto::new(config.runtime.master_key.expose_secret())
.map_err(|_| invalid("runtime.master_key"))?;
crank_runtime::OutboundHttpPolicy::try_new(
config.runtime.outbound.allowed_hosts.clone(),
config.runtime.outbound.denied_hosts.clone(),
config.runtime.outbound.max_response_bytes,
)
.map_err(|_| invalid("runtime.outbound"))?;
let identity = ServiceIdentity::try_new(
"mcp-server",
env!("CARGO_PKG_VERSION"),
config.observability.environment.clone(),
)
.map_err(|_| invalid("observability.environment"))?;
ObservabilityConfig::try_new(
identity,
config.observability.log_filter.clone(),
RedactionLimits::default(),
)
.map_err(|_| invalid("observability.log_filter"))?;
SentryConfig::parse(
config
.observability
.sentry_dsn
.as_ref()
.map(|v| v.expose_secret()),
)
.map_err(|_| invalid("observability.sentry_dsn"))?;
let values = &config.observability.otlp;
OtlpTraceConfig::from_values(
values.endpoint.clone(),
values.traces_endpoint.clone(),
values.protocol.clone(),
values.traces_protocol.clone(),
values.timeout.clone(),
values.traces_timeout.clone(),
values
.headers
.as_ref()
.map(|v| v.expose_secret().to_owned()),
values
.traces_headers
.as_ref()
.map(|v| v.expose_secret().to_owned()),
values.max_queue_size,
values.max_export_batch_size,
values.schedule_delay.clone(),
values.export_timeout.clone(),
)
.map_err(|_| invalid("observability.otlp"))?;
Ok(())
}
fn postgres_pool_config(
config: &DatabaseSettings,
) -> Result<PostgresPoolConfig, crank_registry::PostgresPoolConfigError> {
PostgresPoolConfig::try_new(
config.pool.max_connections,
config.pool.min_connections,
config.pool.acquire_timeout_ms,
config.pool.idle_timeout_ms,
config.pool.max_lifetime_ms,
)
}
fn database_options(
config: &DatabaseSettings,
) -> Result<PgConnectOptions, Box<dyn std::error::Error>> {
if let Some(url) = &config.url {
return url
.expose_secret()
.parse::<PgConnectOptions>()
.map_err(|_| io::Error::other("database URL is invalid").into());
}
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))
.host(&config.host)
.port(config.port)
.database(&config.database)
.username(&config.username)
.password(config.password.expose_secret()))
}
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)?)
fn runtime_cache_config(
config: &McpProcessConfig,
) -> Result<RuntimeCacheConfig, crank_runtime::RuntimeCacheConfigError> {
RuntimeCacheConfig::try_new(
match config.runtime.cache.backend {
ConfigCacheBackend::Memory => CacheBackend::Memory,
ConfigCacheBackend::Valkey => CacheBackend::Valkey,
ConfigCacheBackend::Redis => CacheBackend::Redis,
},
config
.runtime
.cache
.url
.as_ref()
.map(|value| value.expose_secret().to_owned()),
)
}
apps/mcp-server/tests/config_startup.rs
+65
View File
@@ -0,0 +1,65 @@
use std::process::Command;
fn run_with(entries: &[(&str, &str)]) -> String {
let mut command = Command::new(env!("CARGO_BIN_EXE_mcp-server"));
for field in crank_config::field_registry() {
command.env_remove(field.env_name);
}
command.env("CRANK_MASTER_KEY", "master");
for (name, value) in entries {
command.env(name, value);
}
let output = command.output().expect("MCP binary executes");
assert!(!output.status.success());
let stderr = String::from_utf8(output.stderr).expect("stderr is UTF-8");
assert!(!stderr.contains("CANARY_SECRET_VALUE"));
assert!(!stderr.contains("connection refused"));
assert!(stderr.len() <= 65_536);
stderr
}
#[test]
fn invalid_config_fails_before_database_or_listener_side_effects() {
for (entries, code) in [
(
vec![("CRANK_CONFIG_CANARY_UNKNOWN", "CANARY_SECRET_VALUE")],
"config.unknown_field",
),
(vec![("CRANK_MCP_REFRESH_MS", "bad")], "config.invalid_type"),
(
vec![
("CRANK_DATABASE_URL", "postgres://db/crank"),
("POSTGRES_HOST", "other"),
],
"config.conflict",
),
(vec![("CRANK_LOG_LEVEL", "[")], "config.invalid_type"),
] {
let stderr = run_with(&entries);
assert!(stderr.contains(code), "{stderr}");
}
}
#[test]
fn database_driver_failures_are_normalized_and_redacted() {
let stderr = run_with(&[(
"CRANK_DATABASE_URL",
"postgres://CANARY_SECRET_VALUE:CANARY_SECRET_VALUE@127.0.0.1:1/crank",
)]);
assert!(stderr.contains("startup_failed"), "{stderr}");
}
#[tokio::test]
async fn fresh_database_startup_is_read_only() {
let database_url = crank_test_support::postgres_schema_url("mcp_startup_read_only").await;
let stderr = run_with(&[("CRANK_DATABASE_URL", &database_url)]);
assert!(stderr.contains("schema_missing"), "{stderr}");
let pool = sqlx::PgPool::connect(&database_url).await.unwrap();
let present: bool = sqlx::query_scalar(
"select to_regclass(format('%I.%I', current_schema(), '__crank_core_migrations')) is not null",
)
.fetch_one(&pool)
.await
.unwrap();
assert!(!present);
}
apps/mcp-server/tests/integration.rs
+3
View File
@@ -1,6 +1,9 @@
mod integration {
mod catalog_access;
mod common;
mod execution_stages;
mod jsonrpc_correlation;
mod request_context;
mod tool_search;
mod transport_protocol;
}
apps/mcp-server/tests/integration/common.rs
+4
View File
@@ -466,6 +466,10 @@ pub(super) async fn stream_logs()
pub(super) async fn test_registry() -> PostgresRegistry {
let database_url = crank_test_support::postgres_schema_url("test_mcp_server").await;
let pool = sqlx::PgPool::connect(&database_url).await.unwrap();
crank_registry::MigrationAuthority::apply(&pool)
.await
.unwrap();
PostgresRegistry::connect(&database_url).await.unwrap()
}
apps/mcp-server/tests/integration/execution_stages.rs
+33
View File
@@ -154,6 +154,20 @@ async fn exports_real_tool_stages_without_sensitive_data() {
.await;
assert_eq!(call_result.status(), StatusCode::OK);
assert_eq!(
call_result
.headers()
.get("x-request-id")
.and_then(|value| value.to_str().ok()),
Some(REQUEST_ID),
);
assert_eq!(
call_result
.headers()
.get("x-trace-id")
.and_then(|value| value.to_str().ok()),
Some(REMOTE_TRACE_ID),
);
let body = to_bytes(call_result.into_body(), 1024 * 1024)
.await
.unwrap();
@@ -235,6 +249,16 @@ async fn exports_real_tool_stages_without_sensitive_data() {
assert!(!runtime.parent_span_id.is_empty());
assert_eq!(runtime.parent_span_id, root.span_id);
let upstream = trace_spans
.iter()
.find(|span| span.name == "upstream.http")
.expect("upstream attempt");
assert_eq!(
decode_span_id(&traceparent[36..52]).as_slice(),
upstream.span_id.as_slice(),
"outbound traceparent must identify the actual client attempt span",
);
let history = trace_spans
.iter()
.find(|span| span.name == "history.write")
@@ -263,6 +287,7 @@ async fn exports_real_tool_stages_without_sensitive_data() {
.unwrap();
assert_eq!(logs.len(), 1);
assert_eq!(logs[0].log.request_id.as_deref(), Some(REQUEST_ID));
assert_eq!(logs[0].log.trace_id.as_deref(), Some(REMOTE_TRACE_ID));
provider.shutdown().unwrap();
}
@@ -310,6 +335,14 @@ fn decode_trace_id(value: &str) -> [u8; 16] {
bytes
}
fn decode_span_id(value: &str) -> [u8; 8] {
let mut bytes = [0_u8; 8];
for (index, byte) in bytes.iter_mut().enumerate() {
*byte = u8::from_str_radix(&value[index * 2..index * 2 + 2], 16).unwrap();
}
bytes
}
fn string_attribute<'a>(span: &'a Span, key: &str) -> Option<&'a str> {
span.attributes.iter().find_map(|attribute| {
let value = attribute.value.as_ref()?.value.as_ref()?;
apps/mcp-server/tests/integration/jsonrpc_correlation.rs
+73
View File
@@ -0,0 +1,73 @@
use std::time::Duration;
use crank_core::PlatformApiKeyScope;
use crank_registry::PublishRequest;
use serde_json::{Value, json};
use time::{OffsetDateTime, format_description::well_known::Rfc3339};
use super::common::{
agent_mcp_url, build_test_app, create_platform_api_key, initialize_session,
post_jsonrpc_response, publish_agent_for_operation, spawn_mcp_server, spawn_upstream_server,
test_operation, test_registry, test_workspace_id,
};
#[tokio::test]
async fn generic_jsonrpc_errors_carry_the_same_safe_response_ids() {
let registry = test_registry().await;
let upstream_base_url = spawn_upstream_server().await;
let operation = test_operation(&upstream_base_url, "crm_error_identity");
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-error-identity").await;
let api_key = create_platform_api_key(
&registry,
"sales-error-identity",
"mcp-error-identity",
&[PlatformApiKeyScope::Read, PlatformApiKeyScope::Write],
)
.await;
let base_url = spawn_mcp_server(build_test_app(registry, Duration::ZERO, None)).await;
let client = reqwest::Client::new();
let mcp_url = agent_mcp_url(&base_url, "sales-error-identity");
let session = initialize_session(&client, &mcp_url, &api_key).await;
let response = post_jsonrpc_response(
&client,
&mcp_url,
&api_key,
Some(&session),
Some("jsonrpc-error-request"),
json!({
"jsonrpc": "2.0",
"id": 91,
"method": "unsupported/method",
"params": {}
}),
)
.await;
let request_id = response.headers()["x-request-id"]
.to_str()
.unwrap()
.to_owned();
let trace_id = response.headers()["x-trace-id"]
.to_str()
.unwrap()
.to_owned();
let payload = response.json::<Value>().await.unwrap();
assert_eq!(request_id, "jsonrpc-error-request");
assert_eq!(payload["error"]["data"]["request_id"], request_id);
assert_eq!(payload["error"]["data"]["trace_id"], trace_id);
}
apps/mcp-server/tests/integration/request_context.rs
+54 -1
View File
@@ -9,7 +9,7 @@ use axum::{
};
use opentelemetry::{
global,
trace::{TraceId, TracerProvider as _},
trace::{SpanId, TraceId, TracerProvider as _},
};
use opentelemetry_sdk::{
error::OTelSdkResult,
@@ -61,6 +61,7 @@ async fn covers_valid_invalid_and_absent_traceparent_on_mcp_boundary() {
assert_eq!(valid.status, StatusCode::OK);
assert_eq!(valid.request_id.as_deref(), Some("request-id-is-separate"));
assert_eq!(valid.trace_id.as_deref(), Some(REMOTE_TRACE_ID));
assert_eq!(invalid.status, StatusCode::OK);
assert_eq!(absent.status, StatusCode::OK);
assert!(valid.traceparent_response.is_none());
@@ -80,6 +81,41 @@ async fn covers_valid_invalid_and_absent_traceparent_on_mcp_boundary() {
assert_ne!(trace_ids[2], trace_ids[0]);
assert_ne!(trace_ids[1], trace_ids[2]);
assert!(!trace_ids.contains(&TraceId::INVALID));
let request_spans = exported
.lock()
.unwrap()
.iter()
.filter(|span| span.name.as_ref() == "mcp.request")
.cloned()
.collect::<Vec<_>>();
assert_eq!(
request_spans[0].parent_span_id.to_string(),
"b7ad6b7169203331"
);
assert_eq!(request_spans[1].parent_span_id, SpanId::INVALID);
assert_eq!(request_spans[2].parent_span_id, SpanId::INVALID);
provider.shutdown().unwrap();
}
#[tokio::test(flavor = "current_thread")]
async fn sampling_off_still_returns_a_local_trace_identity() {
let _tracing_test_guard = TRACING_TEST_LOCK.lock().await;
global::set_text_map_propagator(TraceContextPropagator::new());
let provider = SdkTracerProvider::builder()
.with_sampler(opentelemetry_sdk::trace::Sampler::AlwaysOff)
.build();
let tracer = provider.tracer("mcp-request-context-sampling-off-test");
let subscriber =
tracing_subscriber::registry().with(tracing_opentelemetry::layer().with_tracer(tracer));
let app = build_test_app(test_registry().await, Duration::ZERO, None);
let response = send_health(app, None, None)
.with_subscriber(subscriber)
.await;
let trace_id = response.trace_id.expect("local trace id");
assert_eq!(trace_id.len(), 32);
assert_ne!(trace_id, "00000000000000000000000000000000");
provider.shutdown().unwrap();
}
@@ -98,6 +134,14 @@ async fn replaces_multiple_request_id_headers_with_one_uuid_v7() {
"x-request-id",
HeaderValue::from_static("second-request-id"),
);
request.headers_mut().append(
"traceparent",
HeaderValue::from_static("00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01"),
);
request.headers_mut().append(
"traceparent",
HeaderValue::from_static("00-1af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01"),
);
let response = app
.oneshot(request)
@@ -112,6 +156,9 @@ async fn replaces_multiple_request_id_headers_with_one_uuid_v7() {
uuid::Uuid::parse_str(generated).unwrap().get_version(),
Some(uuid::Version::SortRand)
);
let trace_id = response.headers()["x-trace-id"].to_str().unwrap();
assert_ne!(trace_id, "0af7651916cd43dd8448eb211c80319c");
assert_ne!(trace_id, "1af7651916cd43dd8448eb211c80319c");
}
async fn send_health(
@@ -143,6 +190,11 @@ async fn send_health(
.get("traceparent")
.and_then(|value| value.to_str().ok())
.map(str::to_owned),
trace_id: response
.headers()
.get("x-trace-id")
.and_then(|value| value.to_str().ok())
.map(str::to_owned),
}
}
@@ -150,6 +202,7 @@ struct ProbeResponse {
status: StatusCode,
request_id: Option<String>,
traceparent_response: Option<String>,
trace_id: Option<String>,
}
#[derive(Clone, Debug)]
apps/mcp-server/tests/integration/transport_protocol.rs
+13
View File
@@ -322,6 +322,10 @@ async fn preserves_request_id_for_tool_call_invocations() {
response.headers()["x-request-id"].to_str().unwrap(),
"req_test_123"
);
let trace_id = response.headers()["x-trace-id"]
.to_str()
.unwrap()
.to_owned();
let call_result = response.json::<Value>().await.unwrap();
assert_eq!(call_result["result"]["isError"], false);
@@ -341,6 +345,7 @@ async fn preserves_request_id_for_tool_call_invocations() {
assert_eq!(logs.len(), 1);
assert_eq!(logs[0].log.request_id.as_deref(), Some("req_test_123"));
assert_eq!(logs[0].log.trace_id.as_deref(), Some(trace_id.as_str()));
}
#[tokio::test]
@@ -408,6 +413,13 @@ async fn generates_request_id_for_tool_call_responses_and_logs() {
.to_str()
.unwrap()
.to_owned();
let trace_id = response
.headers()
.get("x-trace-id")
.unwrap()
.to_str()
.unwrap()
.to_owned();
assert_eq!(
uuid::Uuid::parse_str(&request_id).unwrap().get_version(),
Some(Version::SortRand)
@@ -432,6 +444,7 @@ async fn generates_request_id_for_tool_call_responses_and_logs() {
assert_eq!(logs.len(), 1);
assert_eq!(logs[0].log.request_id.as_deref(), Some(request_id.as_str()));
assert_eq!(logs[0].log.trace_id.as_deref(), Some(trace_id.as_str()));
}
#[tokio::test(flavor = "current_thread")]