наблюдаемость: завершить базовый контур Community
Добавить структурированные журналы, метрики, трассировку и безопасный канал критических ошибок. Усилить границы рантайма, тесты, проверку зависимостей и сценарии развёртывания.
This commit is contained in:
@@ -3,6 +3,7 @@ name = "mcp-server"
|
||||
edition.workspace = true
|
||||
license.workspace = true
|
||||
rust-version.workspace = true
|
||||
publish.workspace = true
|
||||
version.workspace = true
|
||||
|
||||
[[bin]]
|
||||
@@ -15,10 +16,12 @@ axum.workspace = true
|
||||
base64.workspace = true
|
||||
crank-community-mcp = { path = "../../crates/crank-community-mcp" }
|
||||
crank-core = { path = "../../crates/crank-core" }
|
||||
crank-observability = { path = "../../crates/crank-observability" }
|
||||
crank-registry = { path = "../../crates/crank-registry" }
|
||||
crank-runtime = { path = "../../crates/crank-runtime" }
|
||||
crank-schema = { path = "../../crates/crank-schema" }
|
||||
futures-util = "0.3"
|
||||
metrics.workspace = true
|
||||
serde.workspace = true
|
||||
serde_json.workspace = true
|
||||
sha2.workspace = true
|
||||
@@ -34,4 +37,10 @@ uuid.workspace = true
|
||||
crank-mapping = { path = "../../crates/crank-mapping" }
|
||||
crank-schema = { path = "../../crates/crank-schema" }
|
||||
crank-test-support = { path = "../../crates/crank-test-support" }
|
||||
opentelemetry.workspace = true
|
||||
opentelemetry-proto.workspace = true
|
||||
opentelemetry_sdk.workspace = true
|
||||
prost.workspace = true
|
||||
reqwest.workspace = true
|
||||
tower.workspace = true
|
||||
tracing-opentelemetry.workspace = true
|
||||
|
||||
+70
-21
@@ -1,26 +1,53 @@
|
||||
use std::{env, net::SocketAddr, time::Duration};
|
||||
|
||||
use crank_community_mcp::{
|
||||
auth::CommunityMachineCredentialVerifier, build_app_with_background_workers,
|
||||
auth::CommunityMachineCredentialVerifier, build_app_with_background_workers_and_limits,
|
||||
session::PostgresTransportSessionStore,
|
||||
};
|
||||
use crank_observability::{
|
||||
CriticalErrorCategory, MetricsConfig, ObservabilityConfig, ObservabilityLifecycle,
|
||||
capture_critical_error,
|
||||
};
|
||||
use crank_registry::{PostgresPoolConfig, PostgresRegistry};
|
||||
use crank_runtime::{
|
||||
RequestRateLimitConfig, RequestRateLimiter, RuntimeCacheConfig, RuntimeCacheStores,
|
||||
RuntimeLimits, SecretCrypto,
|
||||
};
|
||||
use sqlx::postgres::PgConnectOptions;
|
||||
use sqlx::{PgPool, 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 observability = crank_observability::init(ObservabilityConfig::from_env(
|
||||
"mcp-server",
|
||||
env!("CARGO_PKG_VERSION"),
|
||||
"mcp_server=info,tower_http=info",
|
||||
)?)?;
|
||||
let mut startup_completed = false;
|
||||
let result = run(&observability, &mut startup_completed).await;
|
||||
if result.is_err() {
|
||||
capture_critical_error(if startup_completed {
|
||||
CriticalErrorCategory::Internal
|
||||
} else {
|
||||
CriticalErrorCategory::Startup
|
||||
});
|
||||
}
|
||||
result
|
||||
}
|
||||
|
||||
async fn run(
|
||||
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_enabled = metrics_config.enabled();
|
||||
let metrics_server = if metrics_config.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();
|
||||
@@ -36,23 +63,20 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
|
||||
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 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()?
|
||||
.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(
|
||||
let app = build_app_with_background_workers_and_limits(
|
||||
registry,
|
||||
refresh_interval,
|
||||
base_url,
|
||||
@@ -66,11 +90,14 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
|
||||
cache_stores.coordination.clone(),
|
||||
std::sync::Arc::new(session_store),
|
||||
std::sync::Arc::new(CommunityMachineCredentialVerifier),
|
||||
runtime_limits.max_concurrent_sessions,
|
||||
);
|
||||
let listener = TcpListener::bind(socket_addr).await?;
|
||||
|
||||
info!(
|
||||
name: "mcp.postgres_pool.configured",
|
||||
runtime_max_concurrent_unary = runtime_limits.max_concurrent_unary,
|
||||
runtime_max_concurrent_sessions = runtime_limits.max_concurrent_sessions,
|
||||
mcp_rate_limit_rps = api_rate_limit.requests_per_second,
|
||||
mcp_rate_limit_burst = api_rate_limit.burst,
|
||||
cache_backend = %cache_config.backend,
|
||||
@@ -81,9 +108,21 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
|
||||
max_lifetime_ms = pool_config.max_lifetime_ms,
|
||||
"postgres pool configured"
|
||||
);
|
||||
info!("mcp-server listening on {}", socket_addr);
|
||||
info!(
|
||||
name: "mcp.server.listening",
|
||||
bind_address = %socket_addr,
|
||||
"mcp-server listening"
|
||||
);
|
||||
*startup_completed = true;
|
||||
|
||||
axum::serve(listener, app).await?;
|
||||
if let Some(metrics_server) = metrics_server {
|
||||
tokio::select! {
|
||||
result = axum::serve(listener, app) => result?,
|
||||
result = metrics_server.serve() => result?,
|
||||
}
|
||||
} else {
|
||||
axum::serve(listener, app).await?;
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
@@ -123,3 +162,13 @@ fn mcp_api_rate_limit_config_from_env() -> Result<RequestRateLimitConfig, Box<dy
|
||||
|
||||
Ok(RequestRateLimitConfig::new(requests_per_second, burst)?)
|
||||
}
|
||||
|
||||
fn spawn_postgres_pool_metrics(pool: PgPool) {
|
||||
tokio::spawn(async move {
|
||||
let mut interval = tokio::time::interval(Duration::from_secs(5));
|
||||
loop {
|
||||
interval.tick().await;
|
||||
crank_observability::record_db_pool_connections(pool.size(), pool.num_idle());
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
@@ -0,0 +1,4 @@
|
||||
#[path = "integration/common.rs"]
|
||||
mod common;
|
||||
#[path = "integration/execution_stages.rs"]
|
||||
mod execution_stages;
|
||||
@@ -89,11 +89,16 @@ async fn approval_key_lists_and_decides_pending_requests() {
|
||||
let approved = client
|
||||
.post(&approve_url)
|
||||
.header(header::AUTHORIZATION, format!("Bearer {approval_key}"))
|
||||
.header("x-request-id", "req_approval_execute_123")
|
||||
.json(&json!({ "approve": "yes", "note": "confirmed by test" }))
|
||||
.send()
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(approved.status(), reqwest::StatusCode::OK);
|
||||
assert_eq!(
|
||||
approved.headers()["x-request-id"].to_str().unwrap(),
|
||||
"req_approval_execute_123"
|
||||
);
|
||||
let approved_body = approved.json::<Value>().await.unwrap();
|
||||
assert_eq!(
|
||||
approved_body["approval"]["status"],
|
||||
@@ -169,6 +174,10 @@ async fn approval_key_lists_and_decides_pending_requests() {
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(logs.len(), 1);
|
||||
assert_eq!(
|
||||
logs[0].log.request_id.as_deref(),
|
||||
Some("req_approval_execute_123")
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
@@ -401,22 +410,23 @@ async fn tool_call_with_approval_policy_creates_pending_request() {
|
||||
let mcp_url = agent_mcp_url(&base_url, "sales-gated");
|
||||
let initialized_session = initialize_session(&client, &mcp_url, &api_key).await;
|
||||
|
||||
let tool_call = json!({
|
||||
"jsonrpc": "2.0",
|
||||
"id": 9,
|
||||
"method": "tools/call",
|
||||
"params": {
|
||||
"name": "crm_requires_human_approval",
|
||||
"arguments": {
|
||||
"email": "ada@example.com"
|
||||
}
|
||||
}
|
||||
});
|
||||
let tool_result = post_jsonrpc(
|
||||
&client,
|
||||
&mcp_url,
|
||||
&api_key,
|
||||
Some(&initialized_session),
|
||||
json!({
|
||||
"jsonrpc": "2.0",
|
||||
"id": 9,
|
||||
"method": "tools/call",
|
||||
"params": {
|
||||
"name": "crm_requires_human_approval",
|
||||
"arguments": {
|
||||
"email": "ada@example.com"
|
||||
}
|
||||
}
|
||||
}),
|
||||
tool_call.clone(),
|
||||
)
|
||||
.await;
|
||||
|
||||
@@ -430,6 +440,19 @@ async fn tool_call_with_approval_policy_creates_pending_request() {
|
||||
.unwrap();
|
||||
assert!(approval_id.starts_with("approval_"));
|
||||
|
||||
let repeated_tool_result = post_jsonrpc(
|
||||
&client,
|
||||
&mcp_url,
|
||||
&api_key,
|
||||
Some(&initialized_session),
|
||||
tool_call,
|
||||
)
|
||||
.await;
|
||||
assert_eq!(
|
||||
repeated_tool_result["result"]["structuredContent"]["approval_id"], approval_id,
|
||||
"deduplicated tools/call must return the persisted approval id",
|
||||
);
|
||||
|
||||
let approvals_url = format!("{}/approvals", agent_mcp_url(&base_url, "sales-gated"));
|
||||
let pending = client
|
||||
.get(&approvals_url)
|
||||
@@ -448,6 +471,189 @@ async fn tool_call_with_approval_policy_creates_pending_request() {
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn approval_http_endpoints_enforce_request_rate_limit() {
|
||||
let registry = test_registry().await;
|
||||
let upstream_base_url = spawn_upstream_server().await;
|
||||
let operation = test_operation(&upstream_base_url, "crm_approval_rate_limit");
|
||||
registry
|
||||
.create_operation(&test_workspace_id(), &operation, Some("alice"))
|
||||
.await
|
||||
.unwrap();
|
||||
publish_agent_for_operation(®istry, &operation, "sales-approval-rate-limit").await;
|
||||
let approval_key = create_approval_platform_api_key(
|
||||
®istry,
|
||||
"sales-approval-rate-limit",
|
||||
"approval-rate-limit",
|
||||
)
|
||||
.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 approvals_url = format!(
|
||||
"{}/approvals",
|
||||
agent_mcp_url(&base_url, "sales-approval-rate-limit")
|
||||
);
|
||||
let client = reqwest::Client::new();
|
||||
|
||||
let allowed = client
|
||||
.get(&approvals_url)
|
||||
.header(header::AUTHORIZATION, format!("Bearer {approval_key}"))
|
||||
.send()
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(allowed.status(), reqwest::StatusCode::OK);
|
||||
|
||||
let limited = client
|
||||
.get(&approvals_url)
|
||||
.header(header::AUTHORIZATION, format!("Bearer {approval_key}"))
|
||||
.send()
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(limited.status(), reqwest::StatusCode::TOO_MANY_REQUESTS);
|
||||
assert!(limited.headers().contains_key(header::RETRY_AFTER));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn recovery_does_not_repeat_interrupted_mutating_approval() {
|
||||
let registry = test_registry().await;
|
||||
let (upstream_base_url, upstream_calls) = spawn_counted_approval_upstream().await;
|
||||
let operation = test_operation(&upstream_base_url, "crm_interrupted_approval");
|
||||
registry
|
||||
.create_operation(&test_workspace_id(), &operation, Some("alice"))
|
||||
.await
|
||||
.unwrap();
|
||||
publish_agent_for_operation(®istry, &operation, "sales-interrupted-approval").await;
|
||||
let approval_key_name = "approval-interrupted";
|
||||
create_approval_platform_api_key(®istry, "sales-interrupted-approval", approval_key_name)
|
||||
.await;
|
||||
let now = OffsetDateTime::now_utc();
|
||||
let approval = ApprovalRequest {
|
||||
id: ApprovalRequestId::new("approval_interrupted_mutation"),
|
||||
workspace_id: test_workspace_id(),
|
||||
agent_id: test_agent_id("sales-interrupted-approval"),
|
||||
operation_id: operation.id.clone(),
|
||||
operation_version: operation.version,
|
||||
status: ApprovalRequestStatus::Pending,
|
||||
risk_level: OperationApprovalRiskLevel::Dangerous,
|
||||
request_payload: json!({"email": "interrupted@example.com"}),
|
||||
response_payload: None,
|
||||
created_at: now - time::Duration::minutes(10),
|
||||
expires_at: now + time::Duration::minutes(5),
|
||||
decided_at: None,
|
||||
decided_by_key_id: None,
|
||||
decision_note: None,
|
||||
};
|
||||
registry
|
||||
.create_approval_request(CreateApprovalRequest {
|
||||
approval: &approval,
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
let approval_key_id = PlatformApiKeyId::new(format!("pk_{approval_key_name}"));
|
||||
registry
|
||||
.decide_approval_request(crank_registry::DecideApprovalRequest {
|
||||
workspace_id: &approval.workspace_id,
|
||||
agent_id: &approval.agent_id,
|
||||
approval_id: &approval.id,
|
||||
status: ApprovalRequestStatus::Approved,
|
||||
decided_at: now - time::Duration::minutes(10),
|
||||
decided_by_key_id: &approval_key_id,
|
||||
response_payload: Some(json!({"approve": "yes"})),
|
||||
decision_note: None,
|
||||
})
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
registry
|
||||
.claim_approval_request(
|
||||
&approval.workspace_id,
|
||||
&approval.agent_id,
|
||||
&approval.id,
|
||||
now - time::Duration::minutes(7),
|
||||
)
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
|
||||
let _app = build_test_app_with_approval_recovery(registry.clone());
|
||||
let failed = tokio::time::timeout(Duration::from_secs(2), async {
|
||||
loop {
|
||||
let current = registry
|
||||
.get_approval_request_for_agent(
|
||||
&approval.workspace_id,
|
||||
&approval.agent_id,
|
||||
&approval.id,
|
||||
)
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
if current.approval.status == ApprovalRequestStatus::Failed {
|
||||
break current;
|
||||
}
|
||||
tokio::time::sleep(Duration::from_millis(20)).await;
|
||||
}
|
||||
})
|
||||
.await
|
||||
.expect("recovery must quarantine interrupted execution");
|
||||
|
||||
assert_eq!(
|
||||
failed.approval.response_payload.unwrap()["error"]["code"],
|
||||
"approval_execution_outcome_unknown"
|
||||
);
|
||||
assert_eq!(
|
||||
upstream_calls.load(std::sync::atomic::Ordering::SeqCst),
|
||||
0,
|
||||
"recovery must not repeat a mutating upstream request"
|
||||
);
|
||||
}
|
||||
|
||||
fn build_test_app_with_approval_recovery(registry: PostgresRegistry) -> Router {
|
||||
crank_community_mcp::build_app_with_background_workers(
|
||||
registry,
|
||||
Duration::from_millis(0),
|
||||
Some("https://crank.example.com".to_owned()),
|
||||
SecretCrypto::new("test-master-key").unwrap(),
|
||||
crank_runtime::community_with_outbound_policy(
|
||||
crank_runtime::OutboundHttpPolicy::allowing_hosts(["127.0.0.1"]),
|
||||
)
|
||||
.build(),
|
||||
RequestRateLimiter::new(RequestRateLimitConfig::new(10_000, 10_000).unwrap()),
|
||||
Arc::new(InMemoryCoordinationStateStore::default()),
|
||||
Arc::new(InMemorySessionStore::default()),
|
||||
Arc::new(CommunityMachineCredentialVerifier),
|
||||
)
|
||||
}
|
||||
|
||||
async fn spawn_counted_approval_upstream() -> (String, Arc<std::sync::atomic::AtomicUsize>) {
|
||||
let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
|
||||
let handler_calls = Arc::clone(&calls);
|
||||
let app = Router::new().route(
|
||||
"/crm/leads",
|
||||
post(move |Json(payload): Json<Value>| {
|
||||
let calls = Arc::clone(&handler_calls);
|
||||
async move {
|
||||
calls.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
|
||||
Json(json!({
|
||||
"id": "lead_123",
|
||||
"email": payload["email"]
|
||||
}))
|
||||
}
|
||||
}),
|
||||
);
|
||||
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}"), calls)
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn elicitation_approval_requires_client_capability() {
|
||||
let registry = test_registry().await;
|
||||
|
||||
@@ -0,0 +1,376 @@
|
||||
use std::{
|
||||
sync::{Arc, Mutex},
|
||||
time::Duration,
|
||||
};
|
||||
|
||||
use axum::{
|
||||
Json, Router,
|
||||
body::{Body, Bytes, to_bytes},
|
||||
extract::State,
|
||||
http::{HeaderMap, Request, StatusCode, header},
|
||||
routing::post,
|
||||
};
|
||||
use crank_core::{InvocationSource, PlatformApiKeyScope};
|
||||
use crank_observability::{
|
||||
OtlpBatchConfig, OtlpTraceConfig, ServiceIdentity, build_tracer_provider,
|
||||
};
|
||||
use crank_registry::{ListInvocationLogsQuery, PublishRequest};
|
||||
use opentelemetry::global;
|
||||
use opentelemetry_proto::tonic::{
|
||||
collector::trace::v1::ExportTraceServiceRequest, common::v1::any_value, trace::v1::Span,
|
||||
};
|
||||
use opentelemetry_sdk::propagation::TraceContextPropagator;
|
||||
use prost::Message;
|
||||
use serde_json::{Value, json};
|
||||
use time::{OffsetDateTime, format_description::well_known::Rfc3339};
|
||||
use tokio::net::TcpListener;
|
||||
use tower::ServiceExt;
|
||||
use tracing::instrument::WithSubscriber;
|
||||
use tracing_subscriber::{Layer, filter::filter_fn, layer::SubscriberExt};
|
||||
|
||||
use super::common::{
|
||||
build_test_app, create_platform_api_key, publish_agent_for_operation, test_operation,
|
||||
test_registry, test_workspace_id,
|
||||
};
|
||||
|
||||
const REMOTE_TRACE_ID: &str = "0af7651916cd43dd8448eb211c80319c";
|
||||
const REQUEST_ID: &str = "req_stage_end_to_end";
|
||||
const CANARY: &str = "dc-stage-canary-secret";
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn exports_real_tool_stages_without_sensitive_data() {
|
||||
global::set_text_map_propagator(TraceContextPropagator::new());
|
||||
let (otlp_endpoint, request_rx, collector) = spawn_otlp_collector().await;
|
||||
let trace_config = OtlpTraceConfig::try_new(
|
||||
Some(otlp_endpoint),
|
||||
Some("http/protobuf".to_owned()),
|
||||
Duration::from_secs(2),
|
||||
OtlpBatchConfig::try_new(128, 128, Duration::from_secs(300), Duration::from_secs(2))
|
||||
.unwrap(),
|
||||
)
|
||||
.unwrap();
|
||||
let identity = ServiceIdentity::try_new("mcp-server", "0.3.1", "integration-test").unwrap();
|
||||
let (provider, tracer) = build_tracer_provider(&identity, &trace_config)
|
||||
.unwrap()
|
||||
.expect("enabled OTLP provider");
|
||||
let subscriber = tracing_subscriber::registry().with(
|
||||
tracing_opentelemetry::layer()
|
||||
.with_tracer(tracer)
|
||||
.with_filter(filter_fn(|metadata| {
|
||||
metadata.is_span() && metadata.target() == "crank::trace"
|
||||
})),
|
||||
);
|
||||
let dispatch = tracing::Dispatch::new(subscriber);
|
||||
|
||||
let registry = test_registry().await;
|
||||
let observed_traceparent = Arc::new(Mutex::new(None));
|
||||
let upstream_base_url = spawn_upstream(Arc::clone(&observed_traceparent)).await;
|
||||
let operation = test_operation(&upstream_base_url, "stage_end_to_end");
|
||||
registry
|
||||
.create_operation(&test_workspace_id(), &operation, Some("test"))
|
||||
.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("test"),
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
publish_agent_for_operation(®istry, &operation, "stage-agent").await;
|
||||
let api_key = create_platform_api_key(
|
||||
®istry,
|
||||
"stage-agent",
|
||||
"stage-key",
|
||||
&[PlatformApiKeyScope::Read, PlatformApiKeyScope::Write],
|
||||
)
|
||||
.await;
|
||||
let app = build_test_app(registry.clone(), Duration::ZERO, None);
|
||||
|
||||
let call_result = async {
|
||||
let initialized = send_jsonrpc(
|
||||
app.clone(),
|
||||
&api_key,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
json!({
|
||||
"jsonrpc": "2.0",
|
||||
"id": 1,
|
||||
"method": "initialize",
|
||||
"params": {
|
||||
"protocolVersion": "2025-11-25",
|
||||
"capabilities": {}
|
||||
}
|
||||
}),
|
||||
)
|
||||
.await;
|
||||
assert_eq!(initialized.status(), StatusCode::OK);
|
||||
let session_id = initialized
|
||||
.headers()
|
||||
.get("MCP-Session-Id")
|
||||
.unwrap()
|
||||
.to_str()
|
||||
.unwrap()
|
||||
.to_owned();
|
||||
|
||||
let notification = send_jsonrpc(
|
||||
app.clone(),
|
||||
&api_key,
|
||||
Some(&session_id),
|
||||
None,
|
||||
None,
|
||||
json!({
|
||||
"jsonrpc": "2.0",
|
||||
"method": "notifications/initialized",
|
||||
"params": {}
|
||||
}),
|
||||
)
|
||||
.await;
|
||||
assert_eq!(notification.status(), StatusCode::ACCEPTED);
|
||||
|
||||
send_jsonrpc(
|
||||
app,
|
||||
&api_key,
|
||||
Some(&session_id),
|
||||
Some(REQUEST_ID),
|
||||
Some("00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01"),
|
||||
json!({
|
||||
"jsonrpc": "2.0",
|
||||
"id": 2,
|
||||
"method": "tools/call",
|
||||
"params": {
|
||||
"name": "stage_end_to_end",
|
||||
"arguments": { "email": CANARY }
|
||||
}
|
||||
}),
|
||||
)
|
||||
.await
|
||||
}
|
||||
.with_subscriber(dispatch)
|
||||
.await;
|
||||
|
||||
assert_eq!(call_result.status(), StatusCode::OK);
|
||||
let body = to_bytes(call_result.into_body(), 1024 * 1024)
|
||||
.await
|
||||
.unwrap();
|
||||
let body: Value = serde_json::from_slice(&body).unwrap();
|
||||
assert_eq!(body["result"]["isError"], false);
|
||||
provider.force_flush().unwrap();
|
||||
let request = tokio::time::timeout(Duration::from_secs(2), request_rx)
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
collector.abort();
|
||||
assert_eq!(
|
||||
request
|
||||
.headers
|
||||
.get(header::CONTENT_TYPE)
|
||||
.and_then(|value| value.to_str().ok()),
|
||||
Some("application/x-protobuf")
|
||||
);
|
||||
assert!(
|
||||
!request
|
||||
.body
|
||||
.windows(CANARY.len())
|
||||
.any(|window| window == CANARY.as_bytes())
|
||||
);
|
||||
let export = ExportTraceServiceRequest::decode(request.body).unwrap();
|
||||
let spans = export
|
||||
.resource_spans
|
||||
.iter()
|
||||
.flat_map(|resource| &resource.scope_spans)
|
||||
.flat_map(|scope| &scope.spans)
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
let traceparent = observed_traceparent
|
||||
.lock()
|
||||
.unwrap()
|
||||
.clone()
|
||||
.expect("upstream traceparent");
|
||||
assert_eq!(&traceparent[3..35], REMOTE_TRACE_ID);
|
||||
|
||||
let expected_trace_id = decode_trace_id(REMOTE_TRACE_ID);
|
||||
let trace_spans = spans
|
||||
.iter()
|
||||
.copied()
|
||||
.filter(|span| span.trace_id.as_slice() == expected_trace_id)
|
||||
.collect::<Vec<_>>();
|
||||
for expected in [
|
||||
"mcp.request",
|
||||
"mcp.rate_limit",
|
||||
"mcp.access.check",
|
||||
"mcp.catalog.load",
|
||||
"mcp.tools.resolve",
|
||||
"runtime.execute",
|
||||
"runtime.arguments.map",
|
||||
"upstream.http",
|
||||
"runtime.response.transform",
|
||||
"history.write",
|
||||
"db.query",
|
||||
] {
|
||||
assert!(
|
||||
trace_spans.iter().any(|span| span.name == expected),
|
||||
"missing span {expected}"
|
||||
);
|
||||
}
|
||||
assert!(!trace_spans.iter().any(|span| span.name == "approval.check"));
|
||||
assert!(
|
||||
!trace_spans
|
||||
.iter()
|
||||
.any(|span| span.name == "runtime.idempotency")
|
||||
);
|
||||
|
||||
let root = trace_spans
|
||||
.iter()
|
||||
.find(|span| span.name == "mcp.request")
|
||||
.expect("mcp root");
|
||||
let runtime = trace_spans
|
||||
.iter()
|
||||
.find(|span| span.name == "runtime.execute")
|
||||
.expect("runtime");
|
||||
assert!(!runtime.parent_span_id.is_empty());
|
||||
assert_eq!(runtime.parent_span_id, root.span_id);
|
||||
|
||||
let history = trace_spans
|
||||
.iter()
|
||||
.find(|span| span.name == "history.write")
|
||||
.expect("history write");
|
||||
let history_db = trace_spans
|
||||
.iter()
|
||||
.find(|span| {
|
||||
span.name == "db.query"
|
||||
&& string_attribute(span, "db.operation") == Some("invocation_history.write")
|
||||
})
|
||||
.expect("history PostgreSQL write");
|
||||
assert_eq!(history_db.parent_span_id, history.span_id);
|
||||
|
||||
let logs = registry
|
||||
.list_invocation_logs(ListInvocationLogsQuery {
|
||||
workspace_id: &test_workspace_id(),
|
||||
level: None,
|
||||
search_text: None,
|
||||
source: Some(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));
|
||||
|
||||
provider.shutdown().unwrap();
|
||||
}
|
||||
|
||||
struct OtlpRequest {
|
||||
headers: HeaderMap,
|
||||
body: Bytes,
|
||||
}
|
||||
|
||||
async fn spawn_otlp_collector() -> (
|
||||
String,
|
||||
tokio::sync::oneshot::Receiver<OtlpRequest>,
|
||||
tokio::task::JoinHandle<()>,
|
||||
) {
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let address = listener.local_addr().unwrap();
|
||||
let (request_tx, request_rx) = tokio::sync::oneshot::channel();
|
||||
let sender = Arc::new(Mutex::new(Some(request_tx)));
|
||||
let app = Router::new().route(
|
||||
"/v1/traces",
|
||||
post({
|
||||
let sender = Arc::clone(&sender);
|
||||
move |headers: HeaderMap, body: Bytes| {
|
||||
let sender = Arc::clone(&sender);
|
||||
async move {
|
||||
if let Some(sender) = sender.lock().unwrap().take() {
|
||||
let _ = sender.send(OtlpRequest { headers, body });
|
||||
}
|
||||
StatusCode::OK
|
||||
}
|
||||
}
|
||||
}),
|
||||
);
|
||||
let collector = tokio::spawn(async move {
|
||||
axum::serve(listener, app).await.unwrap();
|
||||
});
|
||||
(format!("http://{address}/v1/traces"), request_rx, collector)
|
||||
}
|
||||
|
||||
fn decode_trace_id(value: &str) -> [u8; 16] {
|
||||
let mut bytes = [0_u8; 16];
|
||||
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()?;
|
||||
(attribute.key == key)
|
||||
.then_some(value)
|
||||
.and_then(|value| match value {
|
||||
any_value::Value::StringValue(value) => Some(value.as_str()),
|
||||
_ => None,
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
async fn send_jsonrpc(
|
||||
app: Router,
|
||||
api_key: &str,
|
||||
session_id: Option<&str>,
|
||||
request_id: Option<&str>,
|
||||
traceparent: Option<&str>,
|
||||
payload: Value,
|
||||
) -> axum::response::Response {
|
||||
let mut request = Request::builder()
|
||||
.method("POST")
|
||||
.uri("/v1/default/stage-agent")
|
||||
.header(header::CONTENT_TYPE, "application/json")
|
||||
.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);
|
||||
}
|
||||
if let Some(traceparent) = traceparent {
|
||||
request = request.header("traceparent", traceparent);
|
||||
}
|
||||
app.oneshot(request.body(Body::from(payload.to_string())).unwrap())
|
||||
.await
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
async fn spawn_upstream(observed: Arc<Mutex<Option<String>>>) -> String {
|
||||
async fn create_lead(
|
||||
State(observed): State<Arc<Mutex<Option<String>>>>,
|
||||
headers: HeaderMap,
|
||||
Json(_payload): Json<Value>,
|
||||
) -> Json<Value> {
|
||||
*observed.lock().unwrap() = headers
|
||||
.get("traceparent")
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.map(str::to_owned);
|
||||
Json(json!({ "id": "lead_123" }))
|
||||
}
|
||||
|
||||
let app = Router::new()
|
||||
.route("/crm/leads", post(create_lead))
|
||||
.with_state(observed);
|
||||
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}")
|
||||
}
|
||||
@@ -0,0 +1,130 @@
|
||||
use std::{
|
||||
sync::{Arc, Mutex},
|
||||
time::Duration,
|
||||
};
|
||||
|
||||
use axum::{
|
||||
body::Body,
|
||||
http::{Request, StatusCode},
|
||||
};
|
||||
use opentelemetry::{
|
||||
global,
|
||||
trace::{TraceId, TracerProvider as _},
|
||||
};
|
||||
use opentelemetry_sdk::{
|
||||
error::OTelSdkResult,
|
||||
propagation::TraceContextPropagator,
|
||||
trace::{SdkTracerProvider, SpanData, SpanExporter},
|
||||
};
|
||||
use tower::ServiceExt;
|
||||
use tracing::instrument::WithSubscriber;
|
||||
use tracing_subscriber::layer::SubscriberExt;
|
||||
|
||||
use super::common::{build_test_app, test_registry};
|
||||
|
||||
const REMOTE_TRACE_ID: &str = "0af7651916cd43dd8448eb211c80319c";
|
||||
|
||||
#[tokio::test(flavor = "current_thread")]
|
||||
async fn covers_valid_invalid_and_absent_traceparent_on_mcp_boundary() {
|
||||
global::set_text_map_propagator(TraceContextPropagator::new());
|
||||
let exported = Arc::new(Mutex::new(Vec::new()));
|
||||
let provider = SdkTracerProvider::builder()
|
||||
.with_simple_exporter(CapturingExporter(Arc::clone(&exported)))
|
||||
.build();
|
||||
let tracer = provider.tracer("mcp-request-context-test");
|
||||
let subscriber =
|
||||
tracing_subscriber::registry().with(tracing_opentelemetry::layer().with_tracer(tracer));
|
||||
let dispatch = tracing::Dispatch::new(subscriber);
|
||||
let app = build_test_app(test_registry().await, Duration::ZERO, None);
|
||||
|
||||
let (valid, invalid, absent) = async {
|
||||
let valid = send_health(
|
||||
app.clone(),
|
||||
Some("00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01"),
|
||||
Some("request-id-is-separate"),
|
||||
)
|
||||
.await;
|
||||
let invalid = send_health(
|
||||
app.clone(),
|
||||
Some("canary-invalid-traceparent"),
|
||||
Some("bad,value"),
|
||||
)
|
||||
.await;
|
||||
let absent = send_health(app, None, None).await;
|
||||
(valid, invalid, absent)
|
||||
}
|
||||
.with_subscriber(dispatch)
|
||||
.await;
|
||||
provider.force_flush().unwrap();
|
||||
|
||||
assert_eq!(valid.status, StatusCode::OK);
|
||||
assert_eq!(valid.request_id.as_deref(), Some("request-id-is-separate"));
|
||||
assert_eq!(invalid.status, StatusCode::OK);
|
||||
assert_eq!(absent.status, StatusCode::OK);
|
||||
assert!(valid.traceparent_response.is_none());
|
||||
assert!(invalid.traceparent_response.is_none());
|
||||
assert!(absent.traceparent_response.is_none());
|
||||
|
||||
let trace_ids: Vec<_> = exported
|
||||
.lock()
|
||||
.unwrap()
|
||||
.iter()
|
||||
.filter(|span| span.name.as_ref() == "mcp.request")
|
||||
.map(|span| span.span_context.trace_id())
|
||||
.collect();
|
||||
assert_eq!(trace_ids.len(), 3);
|
||||
assert_eq!(trace_ids[0].to_string(), REMOTE_TRACE_ID);
|
||||
assert_ne!(trace_ids[1], trace_ids[0]);
|
||||
assert_ne!(trace_ids[2], trace_ids[0]);
|
||||
assert_ne!(trace_ids[1], trace_ids[2]);
|
||||
assert!(!trace_ids.contains(&TraceId::INVALID));
|
||||
provider.shutdown().unwrap();
|
||||
}
|
||||
|
||||
async fn send_health(
|
||||
app: axum::Router,
|
||||
traceparent: Option<&str>,
|
||||
request_id: Option<&str>,
|
||||
) -> ProbeResponse {
|
||||
let mut request = Request::builder().uri("/health");
|
||||
if let Some(traceparent) = traceparent {
|
||||
request = request.header("traceparent", traceparent);
|
||||
}
|
||||
if let Some(request_id) = request_id {
|
||||
request = request.header("x-request-id", request_id);
|
||||
}
|
||||
let response = app
|
||||
.oneshot(request.body(Body::empty()).unwrap())
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
ProbeResponse {
|
||||
status: response.status(),
|
||||
request_id: response
|
||||
.headers()
|
||||
.get("x-request-id")
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.map(str::to_owned),
|
||||
traceparent_response: response
|
||||
.headers()
|
||||
.get("traceparent")
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.map(str::to_owned),
|
||||
}
|
||||
}
|
||||
|
||||
struct ProbeResponse {
|
||||
status: StatusCode,
|
||||
request_id: Option<String>,
|
||||
traceparent_response: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
struct CapturingExporter(Arc<Mutex<Vec<SpanData>>>);
|
||||
|
||||
impl SpanExporter for CapturingExporter {
|
||||
async fn export(&self, batch: Vec<SpanData>) -> OTelSdkResult {
|
||||
self.0.lock().unwrap().extend(batch);
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
@@ -38,7 +38,8 @@ 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 tracing_subscriber::fmt::MakeWriter;
|
||||
use uuid::Version;
|
||||
|
||||
use crank_community_mcp::{
|
||||
auth::{CommunityMachineCredentialVerifier, SharedMachineCredentialVerifier},
|
||||
@@ -46,6 +47,9 @@ use crank_community_mcp::{
|
||||
catalog::PublishedToolCatalog,
|
||||
session::{InMemorySessionStore, SharedSessionStore, TransportSessionStore},
|
||||
};
|
||||
use crank_observability::{
|
||||
ObservabilityConfig, RedactionLimits, ServiceIdentity, build_subscriber,
|
||||
};
|
||||
|
||||
fn test_workspace_id() -> WorkspaceId {
|
||||
WorkspaceId::new("ws_default")
|
||||
@@ -404,7 +408,10 @@ async fn generates_request_id_for_tool_call_responses_and_logs() {
|
||||
.to_str()
|
||||
.unwrap()
|
||||
.to_owned();
|
||||
assert!(!request_id.is_empty());
|
||||
assert_eq!(
|
||||
uuid::Uuid::parse_str(&request_id).unwrap().get_version(),
|
||||
Some(Version::SortRand)
|
||||
);
|
||||
|
||||
let call_result = response.json::<Value>().await.unwrap();
|
||||
assert_eq!(call_result["result"]["isError"], false);
|
||||
@@ -427,7 +434,7 @@ async fn generates_request_id_for_tool_call_responses_and_logs() {
|
||||
assert_eq!(logs[0].log.request_id.as_deref(), Some(request_id.as_str()));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
#[tokio::test(flavor = "current_thread")]
|
||||
async fn emits_request_id_in_mcp_ingress_logs() {
|
||||
let registry = test_registry().await;
|
||||
let upstream_base_url = spawn_upstream_server().await;
|
||||
@@ -456,6 +463,18 @@ async fn emits_request_id_in_mcp_ingress_logs() {
|
||||
)
|
||||
.await;
|
||||
|
||||
let writer = SharedLogWriter::default();
|
||||
let subscriber = build_subscriber(
|
||||
ObservabilityConfig::new(
|
||||
ServiceIdentity::try_new("mcp-server", "test", "test").unwrap(),
|
||||
"info",
|
||||
RedactionLimits::default(),
|
||||
),
|
||||
writer.clone(),
|
||||
)
|
||||
.unwrap();
|
||||
let dispatch = tracing::Dispatch::new(subscriber);
|
||||
let _guard = tracing::dispatcher::set_default(&dispatch);
|
||||
let base_url = spawn_mcp_server(build_test_app(
|
||||
registry,
|
||||
Duration::from_millis(0),
|
||||
@@ -464,18 +483,7 @@ async fn emits_request_id_in_mcp_ingress_logs() {
|
||||
.await;
|
||||
let client = reqwest::Client::new();
|
||||
let mcp_url = agent_mcp_url(&base_url, "sales-request-trace");
|
||||
let writer = SharedLogWriter::default();
|
||||
let subscriber = tracing_subscriber::registry().with(
|
||||
tracing_subscriber::fmt::layer()
|
||||
.with_writer(writer.clone())
|
||||
.without_time()
|
||||
.with_ansi(false)
|
||||
.with_target(false)
|
||||
.compact()
|
||||
.with_filter(LevelFilter::INFO),
|
||||
);
|
||||
|
||||
let _ = tracing::subscriber::set_global_default(subscriber);
|
||||
let response = post_jsonrpc_response(
|
||||
&client,
|
||||
&mcp_url,
|
||||
@@ -500,14 +508,16 @@ async fn emits_request_id_in_mcp_ingress_logs() {
|
||||
assert_eq!(response.status(), reqwest::StatusCode::OK);
|
||||
|
||||
let logs = writer.output();
|
||||
assert!(
|
||||
logs.contains("mcp request received"),
|
||||
"captured logs did not include ingress marker: {logs}"
|
||||
);
|
||||
assert!(logs.contains("req_mcp_trace_123"));
|
||||
assert!(logs.contains("sales-request-trace"));
|
||||
assert!(logs.contains("default"));
|
||||
assert!(logs.contains("initialize"));
|
||||
let event = logs
|
||||
.lines()
|
||||
.filter_map(|line| serde_json::from_str::<Value>(line).ok())
|
||||
.find(|event| event["event"] == "mcp.request.received")
|
||||
.unwrap();
|
||||
assert_eq!(event["service"], "mcp-server");
|
||||
assert_eq!(event["request_id"], "req_mcp_trace_123");
|
||||
assert_eq!(event["fields"]["agent_slug"], "sales-request-trace");
|
||||
assert_eq!(event["fields"]["workspace_slug"], "default");
|
||||
assert_eq!(event["fields"]["jsonrpc_method"], "initialize");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
@@ -803,11 +813,30 @@ async fn get_requires_session_header() {
|
||||
.get(agent_mcp_url(&base_url, "sales-get-sse-missing"))
|
||||
.header(header::ACCEPT, "text/event-stream")
|
||||
.header(header::AUTHORIZATION, format!("Bearer {api_key}"))
|
||||
.header("x-request-id", "req_early_mcp_error")
|
||||
.send()
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(response.status(), reqwest::StatusCode::BAD_REQUEST);
|
||||
assert_eq!(
|
||||
response.headers()["x-request-id"].to_str().unwrap(),
|
||||
"req_early_mcp_error"
|
||||
);
|
||||
|
||||
let invalid_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}"))
|
||||
.header("x-request-id", "bad,value")
|
||||
.send()
|
||||
.await
|
||||
.unwrap();
|
||||
let generated = invalid_response.headers()["x-request-id"].to_str().unwrap();
|
||||
assert_eq!(
|
||||
uuid::Uuid::parse_str(generated).unwrap().get_version(),
|
||||
Some(Version::SortRand)
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
|
||||
@@ -0,0 +1,4 @@
|
||||
#[path = "integration/common.rs"]
|
||||
mod common;
|
||||
#[path = "integration/request_context.rs"]
|
||||
mod request_context;
|
||||
Reference in New Issue
Block a user