From 9a7d60593ae173f675c735f8486a31c1ab5f552b Mon Sep 17 00:00:00 2001 From: bsodfather Date: Fri, 31 Jul 2026 02:37:45 +0300 Subject: [PATCH] =?UTF-8?q?=D0=B8=D1=81=D0=BF=D1=80=D0=B0=D0=B2=D0=B8?= =?UTF-8?q?=D1=82=D1=8C:=20=D0=B7=D0=B0=D0=BA=D1=80=D1=8B=D1=82=D1=8C=20?= =?UTF-8?q?=D1=80=D0=B5=D0=B2=D1=8C=D1=8E=20=D1=81=D0=BA=D0=B2=D0=BE=D0=B7?= =?UTF-8?q?=D0=BD=D0=BE=D0=B9=20=D0=BA=D0=BE=D1=80=D1=80=D0=B5=D0=BB=D1=8F?= =?UTF-8?q?=D1=86=D0=B8=D0=B8?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- apps/admin-api/src/request_context.rs | 13 +- apps/admin-api/src/service.rs | 35 +-- apps/admin-api/src/service/api_keys.rs | 5 +- .../history_loss.rs => dc08.rs} | 69 +++++- .../integration/community_access_usage.rs | 2 - .../tests/integration/request_context.rs | 43 +++- apps/mcp-server/tests/dc08.rs | 227 ++++++++++++++++++ .../catalog_access/approval_access.rs | 49 ++++ .../tests/integration/request_context.rs | 35 ++- crates/crank-community-mcp/src/access.rs | 10 +- .../src/app/invocation_history.rs | 17 +- crates/crank-community-mcp/src/app/tests.rs | 2 +- crates/crank-community-mcp/src/catalog.rs | 4 +- crates/crank-community-mcp/src/rate_limit.rs | 10 +- .../src/request_context.rs | 8 +- crates/crank-core/src/access.rs | 15 ++ crates/crank-observability/src/correlation.rs | 12 + crates/crank-observability/src/otlp.rs | 10 +- .../crank-observability/tests/correlation.rs | 32 +++ crates/crank-observability/tests/otlp.rs | 2 +- crates/crank-runtime/src/executor.rs | 7 +- crates/crank-runtime/tests/integration.rs | 1 - .../crank-runtime/tests/integration/valkey.rs | 6 +- .../tests/{integration => }/stages.rs | 75 ++++++ crates/crank-trace/src/lib.rs | 19 +- crates/crank-trace/tests/contract.rs | 12 +- scripts/check-community-scope.py | 13 + tests/unit/test_check_community_scope.py | 32 +++ 28 files changed, 667 insertions(+), 98 deletions(-) rename apps/admin-api/tests/{integration/community_access_usage/history_loss.rs => dc08.rs} (68%) create mode 100644 apps/mcp-server/tests/dc08.rs rename crates/crank-runtime/tests/{integration => }/stages.rs (84%) diff --git a/apps/admin-api/src/request_context.rs b/apps/admin-api/src/request_context.rs index 8af3ab1..efa2818 100644 --- a/apps/admin-api/src/request_context.rs +++ b/apps/admin-api/src/request_context.rs @@ -1,6 +1,6 @@ use axum::{ extract::Request, - http::{HeaderMap, HeaderName, HeaderValue}, + http::{HeaderName, HeaderValue}, middleware::Next, response::Response, }; @@ -16,7 +16,7 @@ pub struct RequestContext { pub async fn apply_request_context(mut request: Request, next: Next) -> Response { let context = RequestContext { - request_id: resolve_request_id(request.headers()), + request_id: RequestId::resolve_from_headers(request.headers()).into_string(), }; let method = request.method().clone(); let path = request.uri().path().to_owned(); @@ -46,15 +46,6 @@ pub async fn apply_request_context(mut request: Request, next: Next) -> Response .await } -fn resolve_request_id(headers: &HeaderMap) -> String { - RequestId::resolve( - headers - .get(&REQUEST_ID_HEADER) - .and_then(|value| value.to_str().ok()), - ) - .into_string() -} - #[cfg(test)] mod tests { #[test] diff --git a/apps/admin-api/src/service.rs b/apps/admin-api/src/service.rs index aabdfac..d712690 100644 --- a/apps/admin-api/src/service.rs +++ b/apps/admin-api/src/service.rs @@ -5,10 +5,10 @@ use std::sync::Arc; use base64::{Engine as _, engine::general_purpose::URL_SAFE_NO_PAD}; use crank_core::{ AuditSink, AuthProfile, CapabilityProfile, CommunityCapabilityProfile, EditionCapabilities, - ExecutionMode, IdentityError, IdentityProvider, InvocationLog, InvocationLogId, NoopAuditSink, - OperationSecurityLevel, OwnerOnlyPolicyEngine, PolicyEngine, ProductEdition, Protocol, - ToolQualityMappingRule, ToolQualityMappingSet, ToolQualitySchemaKind, ToolQualitySchemaNode, - UsagePeriod, WorkspaceId, + ExecutionMode, IdentityError, IdentityProvider, InvocationLog, InvocationLogId, + InvocationSource, NoopAuditSink, OperationSecurityLevel, OwnerOnlyPolicyEngine, PolicyEngine, + ProductEdition, Protocol, ToolQualityMappingRule, ToolQualityMappingSet, ToolQualitySchemaKind, + ToolQualitySchemaNode, UsagePeriod, WorkspaceId, }; use crank_mapping::{MappingRule, MappingSet}; use crank_registry::{ @@ -479,9 +479,7 @@ impl AdminService { let history_span = crank_trace::Stage::HistoryWrite.span(); let (outcome, db_span) = async { - let db_span = crank_trace::Stage::DbQuery - .db_span(crank_trace::DbOperation::InvocationHistoryWrite) - .expect("database stage"); + let db_span = crank_trace::DbOperation::InvocationHistoryWrite.span(); let outcome = self .registry .create_invocation_log(CreateInvocationLogRequest { log: &log }) @@ -509,7 +507,7 @@ impl AdminService { outcome, request.request_id, request.status, - "admin_test_run", + request.source, ); outcome } @@ -519,7 +517,7 @@ fn observe_invocation_history_outcome( outcome: InvocationHistoryWriteOutcome, request_id: Option<&str>, status: crank_core::InvocationStatus, - source: &'static str, + source: InvocationSource, ) { let Some(loss) = outcome.loss() else { return; @@ -530,13 +528,20 @@ fn observe_invocation_history_outcome( tracing::warn!( name: "admin.invocation_history.lost", request_id = request_id.unwrap_or_default(), - source, + source = invocation_source_label(source), invocation_status = invocation_status_label(status), error_category = loss.category.as_str(), "invocation history was not recorded" ); } +fn invocation_source_label(source: InvocationSource) -> &'static str { + match source { + InvocationSource::AdminTestRun => "admin_test_run", + InvocationSource::AgentToolCall => "agent_tool_call", + } +} + fn invocation_status_label(status: crank_core::InvocationStatus) -> &'static str { match status { crank_core::InvocationStatus::Ok => "ok", @@ -606,9 +611,9 @@ fn new_prefixed_id(prefix: &str) -> String { format!("{prefix}_{}", Uuid::now_v7().simple()) } -fn generate_access_secret(prefix: &str) -> String { +fn generate_access_secret(marker: &str) -> String { let random = URL_SAFE_NO_PAD.encode(Uuid::now_v7().as_bytes()); - format!("{prefix}_{random}") + format!("{marker}{random}") } fn hash_access_secret(secret: &str) -> String { @@ -874,7 +879,7 @@ mod tests { sync::{Arc, Mutex}, }; - use crank_core::InvocationStatus; + use crank_core::{InvocationSource, InvocationStatus}; use crank_observability::{ ObservabilityConfig, OperationalIncident, RedactionLimits, ServiceIdentity, operational_incident_total, @@ -931,7 +936,7 @@ mod tests { }), Some("req_admin_dc08"), InvocationStatus::Error, - "admin_test_run", + InvocationSource::AgentToolCall, ); let output = writer.output(); @@ -942,7 +947,7 @@ mod tests { .find(|event: &Value| event["event"] == "admin.invocation_history.lost") .unwrap(); assert_eq!(event["request_id"], "req_admin_dc08"); - assert_eq!(event["fields"]["source"], "admin_test_run"); + assert_eq!(event["fields"]["source"], "agent_tool_call"); assert_eq!(event["fields"]["invocation_status"], "error"); assert_eq!(event["fields"]["error_category"], "invalid_record"); assert!(operational_incident_total(OperationalIncident::InvocationHistoryLost) > before); diff --git a/apps/admin-api/src/service/api_keys.rs b/apps/admin-api/src/service/api_keys.rs index d92c300..744e73a 100644 --- a/apps/admin-api/src/service/api_keys.rs +++ b/apps/admin-api/src/service/api_keys.rs @@ -65,10 +65,7 @@ impl AdminService { ), None => None, }; - let secret = generate_access_secret(match payload.key_kind { - PlatformApiKeyKind::McpClient => "crk", - PlatformApiKeyKind::Approval => "crk_appr", - }); + let secret = generate_access_secret(payload.key_kind.secret_marker()); let api_key = PlatformApiKeyRecord { api_key: PlatformApiKey { id: PlatformApiKeyId::new(new_prefixed_id("pk")), diff --git a/apps/admin-api/tests/integration/community_access_usage/history_loss.rs b/apps/admin-api/tests/dc08.rs similarity index 68% rename from apps/admin-api/tests/integration/community_access_usage/history_loss.rs rename to apps/admin-api/tests/dc08.rs index 93da145..ba027f3 100644 --- a/apps/admin-api/tests/integration/community_access_usage/history_loss.rs +++ b/apps/admin-api/tests/dc08.rs @@ -2,11 +2,21 @@ use std::sync::{ Arc, atomic::{AtomicUsize, Ordering}, }; +use std::{io, sync::Mutex}; use axum::{Json, Router, extract::State, routing::post}; +use crank_core::{OperationId, WorkspaceId}; +use serde_json::{Value, json}; +use serial_test::serial; use tokio::{net::TcpListener, sync::Notify}; +use tracing_subscriber::fmt::MakeWriter; -use super::*; +#[path = "integration/common.rs"] +mod common; + +use common::*; + +const DEFAULT_WORKSPACE_ID: &str = "ws_default"; #[tokio::test(flavor = "multi_thread")] #[serial] @@ -15,6 +25,17 @@ async fn preserves_external_success_when_invocation_history_is_lost() { let registry_for_failure = registry.clone(); let storage_root = test_storage_root("observability_history_loss"); let upstream = spawn_blocking_upstream_server().await; + let log_writer = SharedLogWriter::default(); + let subscriber = crank_observability::build_subscriber( + crank_observability::ObservabilityConfig::new( + crank_observability::ServiceIdentity::try_new("admin-api", "test", "test").unwrap(), + "info", + crank_observability::RedactionLimits::default(), + ), + log_writer.clone(), + ) + .unwrap(); + tracing::subscriber::set_global_default(subscriber).unwrap(); let base_url = spawn_admin_api(build_test_app(registry, storage_root)).await; let client = authorized_client(&base_url).await; @@ -85,6 +106,16 @@ async fn preserves_external_success_when_invocation_history_is_lost() { .await .unwrap(); assert!(logs["items"].as_array().unwrap().is_empty()); + + let output = log_writer.output(); + assert!(!output.contains("dc08-canary-secret")); + let incident = output + .lines() + .filter_map(|line| serde_json::from_str::(line).ok()) + .find(|event| event["event"] == "admin.invocation_history.lost") + .expect("DC-08 incident"); + assert_eq!(incident["request_id"], "req_dc08_admin"); + assert_eq!(incident["fields"]["source"], "admin_test_run"); } struct BlockingUpstream { @@ -139,3 +170,39 @@ async fn blocking_create_lead( "email": payload["email"] })) } + +#[derive(Clone, Default)] +struct SharedLogWriter { + buffer: Arc>>, +} + +impl SharedLogWriter { + fn output(&self) -> String { + String::from_utf8(self.buffer.lock().unwrap().clone()).unwrap() + } +} + +impl<'a> MakeWriter<'a> for SharedLogWriter { + type Writer = SharedLogGuard; + + fn make_writer(&'a self) -> Self::Writer { + SharedLogGuard { + buffer: Arc::clone(&self.buffer), + } + } +} + +struct SharedLogGuard { + buffer: Arc>>, +} + +impl io::Write for SharedLogGuard { + fn write(&mut self, bytes: &[u8]) -> io::Result { + self.buffer.lock().unwrap().extend_from_slice(bytes); + Ok(bytes.len()) + } + + fn flush(&mut self) -> io::Result<()> { + Ok(()) + } +} diff --git a/apps/admin-api/tests/integration/community_access_usage.rs b/apps/admin-api/tests/integration/community_access_usage.rs index 313423d..02665fb 100644 --- a/apps/admin-api/tests/integration/community_access_usage.rs +++ b/apps/admin-api/tests/integration/community_access_usage.rs @@ -41,8 +41,6 @@ const TEST_PASSWORD_PEPPER: &str = "test-password-pepper"; const TEST_SESSION_SECRET: &str = "test-session-secret"; const TEST_MASTER_KEY: &str = "test-master-key"; -mod history_loss; - struct TestServer { base_url: String, shutdown: Option>, diff --git a/apps/admin-api/tests/integration/request_context.rs b/apps/admin-api/tests/integration/request_context.rs index dd3d9e2..08864fc 100644 --- a/apps/admin-api/tests/integration/request_context.rs +++ b/apps/admin-api/tests/integration/request_context.rs @@ -7,7 +7,7 @@ use admin_api::request_context::{REQUEST_ID_HEADER, apply_request_context}; use axum::{ Router, body::Body, - http::{HeaderMap, Request, StatusCode}, + http::{HeaderMap, HeaderValue, Request, StatusCode}, routing::get, }; use crank_observability::{ @@ -16,12 +16,14 @@ use crank_observability::{ use opentelemetry::{global, trace::TracerProvider as _}; use opentelemetry_sdk::{propagation::TraceContextPropagator, trace::SdkTracerProvider}; use tower::ServiceExt; -use tracing::instrument::WithSubscriber; use tracing_subscriber::{fmt::MakeWriter, layer::SubscriberExt}; use uuid::Version; +static TRACING_TEST_LOCK: tokio::sync::Mutex<()> = tokio::sync::Mutex::const_new(()); + #[tokio::test(flavor = "current_thread")] async fn logs_request_completion_and_rejects_untrusted_values() { + let _tracing_test_guard = TRACING_TEST_LOCK.lock().await; let writer = SharedLogWriter::default(); let subscriber = crank_observability::build_subscriber( ObservabilityConfig::new( @@ -33,6 +35,7 @@ async fn logs_request_completion_and_rejects_untrusted_values() { ) .unwrap(); let dispatch = tracing::Dispatch::new(subscriber); + let _dispatch_guard = tracing::dispatcher::set_default(&dispatch); let app = probe_app(); let response = app @@ -44,7 +47,6 @@ async fn logs_request_completion_and_rejects_untrusted_values() { .body(Body::empty()) .unwrap(), ) - .with_subscriber(dispatch.clone()) .await .unwrap(); @@ -73,7 +75,6 @@ async fn logs_request_completion_and_rejects_untrusted_values() { .body(Body::empty()) .unwrap(), ) - .with_subscriber(dispatch) .await .unwrap(); let generated = invalid_response.headers()[REQUEST_ID_HEADER.as_str()] @@ -88,12 +89,14 @@ async fn logs_request_completion_and_rejects_untrusted_values() { #[tokio::test(flavor = "current_thread")] async fn covers_valid_invalid_and_absent_traceparent() { + let _tracing_test_guard = TRACING_TEST_LOCK.lock().await; global::set_text_map_propagator(TraceContextPropagator::new()); let provider = SdkTracerProvider::builder().build(); let tracer = provider.tracer("admin-request-context-test"); let subscriber = tracing_subscriber::registry().with(tracing_opentelemetry::layer().with_tracer(tracer)); let dispatch = tracing::Dispatch::new(subscriber); + let _dispatch_guard = tracing::dispatcher::set_default(&dispatch); let app = trace_probe_app(); let valid = observed_trace_id( @@ -109,7 +112,6 @@ async fn covers_valid_invalid_and_absent_traceparent() { .body(Body::empty()) .unwrap(), ) - .with_subscriber(dispatch.clone()) .await .unwrap(), ); @@ -122,7 +124,6 @@ async fn covers_valid_invalid_and_absent_traceparent() { .body(Body::empty()) .unwrap(), ) - .with_subscriber(dispatch.clone()) .await .unwrap(), ); @@ -133,7 +134,6 @@ async fn covers_valid_invalid_and_absent_traceparent() { .body(Body::empty()) .unwrap(), ) - .with_subscriber(dispatch) .await .unwrap(), ); @@ -145,6 +145,35 @@ async fn covers_valid_invalid_and_absent_traceparent() { provider.shutdown().unwrap(); } +#[tokio::test(flavor = "current_thread")] +async fn replaces_multiple_request_id_headers_with_one_uuid_v7() { + let _tracing_test_guard = TRACING_TEST_LOCK.lock().await; + let dispatch = tracing::Dispatch::new(tracing_subscriber::registry()); + let _dispatch_guard = tracing::dispatcher::set_default(&dispatch); + let mut request = Request::builder() + .uri("/probe") + .body(Body::empty()) + .unwrap(); + request.headers_mut().append( + REQUEST_ID_HEADER, + HeaderValue::from_static("first-request-id"), + ); + request.headers_mut().append( + REQUEST_ID_HEADER, + HeaderValue::from_static("second-request-id"), + ); + + let response = probe_app().oneshot(request).await.unwrap(); + let generated = response.headers()[REQUEST_ID_HEADER].to_str().unwrap(); + + assert_ne!(generated, "first-request-id"); + assert_ne!(generated, "second-request-id"); + assert_eq!( + uuid::Uuid::parse_str(generated).unwrap().get_version(), + Some(Version::SortRand) + ); +} + fn probe_app() -> Router { Router::new() .route("/probe", get(|| async { "ok" })) diff --git a/apps/mcp-server/tests/dc08.rs b/apps/mcp-server/tests/dc08.rs new file mode 100644 index 0000000..e3b3ed3 --- /dev/null +++ b/apps/mcp-server/tests/dc08.rs @@ -0,0 +1,227 @@ +use std::{ + io, + sync::{ + Arc, Mutex, + atomic::{AtomicUsize, Ordering}, + }, + time::Duration, +}; + +use axum::{Json, Router, routing::post}; +use crank_core::{PlatformApiKeyScope, WorkspaceId}; +use crank_observability::{ + ObservabilityConfig, OperationalIncident, RedactionLimits, ServiceIdentity, + operational_incident_total, +}; +use crank_registry::{ListInvocationLogsQuery, PublishRequest}; +use serde_json::{Value, json}; +use time::{OffsetDateTime, format_description::well_known::Rfc3339}; +use tokio::net::TcpListener; +use tracing_subscriber::fmt::MakeWriter; + +#[path = "integration/common.rs"] +mod common; + +use common::*; + +const CANARY_SECRET: &str = "dc08-canary-secret"; + +#[tokio::test] +async fn preserves_mcp_result_when_postgres_rejects_invocation_history() { + let registry = test_registry().await; + let upstream_calls = Arc::new(AtomicUsize::new(0)); + let upstream_base_url = spawn_counted_upstream(Arc::clone(&upstream_calls)).await; + let operation = test_operation(&upstream_base_url, "crm_dc08"); + + 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-dc08").await; + let api_key = create_platform_api_key( + ®istry, + "sales-dc08", + "mcp-dc08", + &[PlatformApiKeyScope::Read, PlatformApiKeyScope::Write], + ) + .await; + install_history_failure_trigger(®istry).await; + + let writer = SharedLogWriter::default(); + let subscriber = crank_observability::build_subscriber( + ObservabilityConfig::new( + ServiceIdentity::try_new("mcp-server", "test", "test").unwrap(), + "info", + RedactionLimits::default(), + ), + writer.clone(), + ) + .unwrap(); + tracing::subscriber::set_global_default(subscriber).unwrap(); + 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-dc08"); + let session_id = initialize_session(&client, &mcp_url, &api_key).await; + let before = operational_incident_total(OperationalIncident::InvocationHistoryLost); + + let result = post_jsonrpc_response( + &client, + &mcp_url, + &api_key, + Some(&session_id), + Some("req_mcp_dc08"), + json!({ + "jsonrpc": "2.0", + "id": 3, + "method": "tools/call", + "params": { + "name": "crm_dc08", + "arguments": { + "email": format!("{CANARY_SECRET}@example.com") + } + } + }), + ) + .await + .json::() + .await + .unwrap(); + + assert_eq!( + result["result"]["structuredContent"], + json!({"id": "lead_123"}) + ); + assert_eq!(result["result"]["isError"], false); + assert_eq!(upstream_calls.load(Ordering::SeqCst), 1); + assert!(operational_incident_total(OperationalIncident::InvocationHistoryLost) > before); + + let logs = registry + .list_invocation_logs(ListInvocationLogsQuery { + workspace_id: &WorkspaceId::new("ws_default"), + level: None, + search_text: None, + source: None, + operation_id: Some(&operation.id), + agent_id: None, + created_after: None, + limit: 10, + }) + .await + .unwrap(); + assert!(logs.is_empty()); + + let output = writer.output(); + assert!(!output.contains(CANARY_SECRET)); + let incidents = output + .lines() + .filter_map(|line| serde_json::from_str::(line).ok()) + .filter(|event| event["event"] == "mcp.invocation_history.lost") + .collect::>(); + assert_eq!(incidents.len(), 1); + assert_eq!(incidents[0]["request_id"], "req_mcp_dc08"); + assert_eq!(incidents[0]["fields"]["source"], "agent_tool_call"); + assert_eq!(incidents[0]["fields"]["invocation_status"], "ok"); +} + +async fn install_history_failure_trigger(registry: &crank_registry::PostgresRegistry) { + sqlx::query( + r#" + create function fail_dc08_invocation_history() returns trigger + language plpgsql + as $$ + begin + if new.request_id = 'req_mcp_dc08' then + raise exception 'forced DC-08 invocation history failure'; + end if; + return new; + end; + $$ + "#, + ) + .execute(registry.pool()) + .await + .unwrap(); + sqlx::query( + r#" + create trigger fail_dc08_invocation_history + before insert on invocation_logs + for each row execute function fail_dc08_invocation_history() + "#, + ) + .execute(registry.pool()) + .await + .unwrap(); +} + +async fn spawn_counted_upstream(calls: Arc) -> String { + let app = Router::new().route( + "/crm/leads", + post(move |Json(payload): Json| { + let calls = Arc::clone(&calls); + async move { + calls.fetch_add(1, 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}") +} + +#[derive(Clone, Default)] +struct SharedLogWriter { + buffer: Arc>>, +} + +impl SharedLogWriter { + fn output(&self) -> String { + String::from_utf8(self.buffer.lock().unwrap().clone()).unwrap() + } +} + +impl<'a> MakeWriter<'a> for SharedLogWriter { + type Writer = SharedLogGuard; + + fn make_writer(&'a self) -> Self::Writer { + SharedLogGuard { + buffer: Arc::clone(&self.buffer), + } + } +} + +struct SharedLogGuard { + buffer: Arc>>, +} + +impl io::Write for SharedLogGuard { + fn write(&mut self, bytes: &[u8]) -> io::Result { + self.buffer.lock().unwrap().extend_from_slice(bytes); + Ok(bytes.len()) + } + + fn flush(&mut self) -> io::Result<()> { + Ok(()) + } +} diff --git a/apps/mcp-server/tests/integration/catalog_access/approval_access.rs b/apps/mcp-server/tests/integration/catalog_access/approval_access.rs index 3e27055..a7affa3 100644 --- a/apps/mcp-server/tests/integration/catalog_access/approval_access.rs +++ b/apps/mcp-server/tests/integration/catalog_access/approval_access.rs @@ -518,6 +518,55 @@ async fn approval_http_endpoints_enforce_request_rate_limit() { assert!(limited.headers().contains_key(header::RETRY_AFTER)); } +#[tokio::test] +async fn unverified_session_ids_do_not_create_approval_rate_limit_buckets() { + let registry = test_registry().await; + let upstream_base_url = spawn_upstream_server().await; + let operation = test_operation(&upstream_base_url, "crm_approval_session_rate_limit"); + registry + .create_operation(&test_workspace_id(), &operation, Some("alice")) + .await + .unwrap(); + publish_agent_for_operation(®istry, &operation, "sales-approval-session-rate-limit").await; + let approval_key = create_approval_platform_api_key( + ®istry, + "sales-approval-session-rate-limit", + "approval-session-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-session-rate-limit") + ); + let client = reqwest::Client::new(); + + let allowed = client + .get(&approvals_url) + .header(header::AUTHORIZATION, format!("Bearer {approval_key}")) + .header("MCP-Session-Id", "unverified-session-a") + .send() + .await + .unwrap(); + assert_eq!(allowed.status(), reqwest::StatusCode::OK); + + let limited = client + .get(&approvals_url) + .header(header::AUTHORIZATION, format!("Bearer {approval_key}")) + .header("MCP-Session-Id", "unverified-session-b") + .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; diff --git a/apps/mcp-server/tests/integration/request_context.rs b/apps/mcp-server/tests/integration/request_context.rs index 7837b73..df44614 100644 --- a/apps/mcp-server/tests/integration/request_context.rs +++ b/apps/mcp-server/tests/integration/request_context.rs @@ -5,7 +5,7 @@ use std::{ use axum::{ body::Body, - http::{Request, StatusCode}, + http::{HeaderValue, Request, StatusCode}, }; use opentelemetry::{ global, @@ -23,9 +23,11 @@ use tracing_subscriber::layer::SubscriberExt; use super::common::{build_test_app, test_registry}; const REMOTE_TRACE_ID: &str = "0af7651916cd43dd8448eb211c80319c"; +static TRACING_TEST_LOCK: tokio::sync::Mutex<()> = tokio::sync::Mutex::const_new(()); #[tokio::test(flavor = "current_thread")] async fn covers_valid_invalid_and_absent_traceparent_on_mcp_boundary() { + let _tracing_test_guard = TRACING_TEST_LOCK.lock().await; global::set_text_map_propagator(TraceContextPropagator::new()); let exported = Arc::new(Mutex::new(Vec::new())); let provider = SdkTracerProvider::builder() @@ -81,6 +83,37 @@ async fn covers_valid_invalid_and_absent_traceparent_on_mcp_boundary() { provider.shutdown().unwrap(); } +#[tokio::test] +async fn replaces_multiple_request_id_headers_with_one_uuid_v7() { + let _tracing_test_guard = TRACING_TEST_LOCK.lock().await; + let app = build_test_app(test_registry().await, Duration::ZERO, None); + let mut request = Request::builder() + .uri("/health") + .body(Body::empty()) + .unwrap(); + request + .headers_mut() + .append("x-request-id", HeaderValue::from_static("first-request-id")); + request.headers_mut().append( + "x-request-id", + HeaderValue::from_static("second-request-id"), + ); + + let response = app + .oneshot(request) + .with_subscriber(tracing_subscriber::registry()) + .await + .unwrap(); + let generated = response.headers()["x-request-id"].to_str().unwrap(); + + assert_ne!(generated, "first-request-id"); + assert_ne!(generated, "second-request-id"); + assert_eq!( + uuid::Uuid::parse_str(generated).unwrap().get_version(), + Some(uuid::Version::SortRand) + ); +} + async fn send_health( app: axum::Router, traceparent: Option<&str>, diff --git a/crates/crank-community-mcp/src/access.rs b/crates/crank-community-mcp/src/access.rs index 90d0504..de8c624 100644 --- a/crates/crank-community-mcp/src/access.rs +++ b/crates/crank-community-mcp/src/access.rs @@ -6,7 +6,7 @@ use axum::{ }; use base64::{Engine as _, engine::general_purpose::URL_SAFE_NO_PAD}; use crank_core::{OperationSecurityLevel, PlatformApiKeyScope}; -use crank_trace::{DbOperation, ErrorCategory, Stage, StageOutcome, observe_db_query}; +use crank_trace::{DbOperation, ErrorCategory, StageOutcome, observe_db_query}; use sha2::{Digest, Sha256}; use time::OffsetDateTime; use tracing::Instrument; @@ -154,9 +154,7 @@ async fn verify_static_agent_key( secret: &str, ) -> Result, MachineAccessError> { let secret_hash = hash_access_secret(secret); - let read_span = Stage::DbQuery - .db_span(DbOperation::MachineAccessRead) - .expect("database stage"); + let read_span = DbOperation::MachineAccessRead.span(); let api_key_result = state .registry .get_platform_api_key_by_secret_for_agent_slug( @@ -185,9 +183,7 @@ async fn verify_static_agent_key( }; let used_at = OffsetDateTime::now_utc(); - let touch_span = Stage::DbQuery - .db_span(DbOperation::MachineAccessTouch) - .expect("database stage"); + let touch_span = DbOperation::MachineAccessTouch.span(); let touch_result = state .registry .touch_platform_api_key(&api_key.api_key.workspace_id, &api_key.api_key.id, &used_at) diff --git a/crates/crank-community-mcp/src/app/invocation_history.rs b/crates/crank-community-mcp/src/app/invocation_history.rs index ee20fbd..9161d86 100644 --- a/crates/crank-community-mcp/src/app/invocation_history.rs +++ b/crates/crank-community-mcp/src/app/invocation_history.rs @@ -52,9 +52,7 @@ pub(crate) async fn persist_invocation( let history_span = Stage::HistoryWrite.span(); let (outcome, db_span) = async { - let db_span = Stage::DbQuery - .db_span(DbOperation::InvocationHistoryWrite) - .expect("database stage"); + let db_span = DbOperation::InvocationHistoryWrite.span(); let outcome = state .registry .create_invocation_log(CreateInvocationLogRequest { log: &log }) @@ -82,7 +80,7 @@ pub(crate) async fn persist_invocation( outcome, record.request_id, record.status, - "agent_tool_call", + InvocationSource::AgentToolCall, ); outcome } @@ -91,7 +89,7 @@ pub(super) fn observe_invocation_history_outcome( outcome: InvocationHistoryWriteOutcome, request_id: Option<&str>, status: InvocationStatus, - source: &'static str, + source: InvocationSource, ) { let Some(loss) = outcome.loss() else { return; @@ -102,13 +100,20 @@ pub(super) fn observe_invocation_history_outcome( warn!( name: "mcp.invocation_history.lost", request_id = request_id.unwrap_or_default(), - source, + source = invocation_source_label(source), invocation_status = invocation_status_label(status), error_category = loss.category.as_str(), "invocation history was not recorded" ); } +fn invocation_source_label(source: InvocationSource) -> &'static str { + match source { + InvocationSource::AdminTestRun => "admin_test_run", + InvocationSource::AgentToolCall => "agent_tool_call", + } +} + fn invocation_status_label(status: InvocationStatus) -> &'static str { match status { InvocationStatus::Ok => "ok", diff --git a/crates/crank-community-mcp/src/app/tests.rs b/crates/crank-community-mcp/src/app/tests.rs index 90e4ed5..cb30c63 100644 --- a/crates/crank-community-mcp/src/app/tests.rs +++ b/crates/crank-community-mcp/src/app/tests.rs @@ -75,7 +75,7 @@ fn emits_bounded_history_loss_incident() { }), Some("req_mcp_dc08"), InvocationStatus::Ok, - "agent_tool_call", + crank_core::InvocationSource::AgentToolCall, ); let output = writer.output(); diff --git a/crates/crank-community-mcp/src/catalog.rs b/crates/crank-community-mcp/src/catalog.rs index 534bcb8..3346e39 100644 --- a/crates/crank-community-mcp/src/catalog.rs +++ b/crates/crank-community-mcp/src/catalog.rs @@ -159,9 +159,7 @@ impl PublishedToolCatalog { return Ok(()); } - let db_span = Stage::DbQuery - .db_span(DbOperation::CatalogLoad) - .expect("database stage"); + let db_span = DbOperation::CatalogLoad.span(); let catalog_result = self .registry .get_published_agent_catalog_by_slug(workspace_slug, agent_slug) diff --git a/crates/crank-community-mcp/src/rate_limit.rs b/crates/crank-community-mcp/src/rate_limit.rs index 446a77f..90e7e93 100644 --- a/crates/crank-community-mcp/src/rate_limit.rs +++ b/crates/crank-community-mcp/src/rate_limit.rs @@ -4,6 +4,7 @@ use axum::{ http::{HeaderMap, HeaderValue, StatusCode, header::RETRY_AFTER}, response::{IntoResponse, Response}, }; +use crank_core::PlatformApiKeyKind; use crank_runtime::RateLimitCheckError; use serde_json::{Value, json}; @@ -82,6 +83,13 @@ pub(super) fn rate_limited_status_response(error: RateLimitCheckError) -> Respon } fn rate_limit_key(path: &AgentRoutePath, headers: &HeaderMap) -> String { + let access_secret = bearer_token(headers); + if let Some(secret) = access_secret + && secret.starts_with(PlatformApiKeyKind::Approval.secret_marker()) + { + return format!("api_key:{}", hash_access_secret(secret)); + } + if let Ok(Some(session_id)) = session_id_from_headers(headers) { return format!( "session:{}:{}:{}", @@ -89,7 +97,7 @@ fn rate_limit_key(path: &AgentRoutePath, headers: &HeaderMap) -> String { ); } - if let Some(secret) = bearer_token(headers) { + if let Some(secret) = access_secret { return format!("api_key:{}", hash_access_secret(secret)); } diff --git a/crates/crank-community-mcp/src/request_context.rs b/crates/crank-community-mcp/src/request_context.rs index 5895dc0..820eb36 100644 --- a/crates/crank-community-mcp/src/request_context.rs +++ b/crates/crank-community-mcp/src/request_context.rs @@ -10,13 +10,7 @@ pub(super) struct RequestContext { } pub(super) async fn apply_request_context(mut request: Request, next: Next) -> Response { - let request_id = RequestId::resolve( - request - .headers() - .get(&HEADER_X_REQUEST_ID) - .and_then(|value| value.to_str().ok()), - ) - .into_string(); + let request_id = RequestId::resolve_from_headers(request.headers()).into_string(); let context = RequestContext { request_id: request_id.clone(), }; diff --git a/crates/crank-core/src/access.rs b/crates/crank-core/src/access.rs index f3fe32f..575d593 100644 --- a/crates/crank-core/src/access.rs +++ b/crates/crank-core/src/access.rs @@ -42,6 +42,15 @@ pub enum PlatformApiKeyKind { Approval, } +impl PlatformApiKeyKind { + pub const fn secret_marker(self) -> &'static str { + match self { + Self::McpClient => "crk_", + Self::Approval => "crk_appr_", + } + } +} + #[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)] #[serde(rename_all = "snake_case")] pub enum PlatformApiKeyScope { @@ -123,6 +132,12 @@ mod tests { }; use crate::ids::{AgentId, PlatformApiKeyId, UserId, WorkspaceId}; + #[test] + fn api_key_kinds_own_their_secret_markers() { + assert_eq!(PlatformApiKeyKind::McpClient.secret_marker(), "crk_"); + assert_eq!(PlatformApiKeyKind::Approval.secret_marker(), "crk_appr_"); + } + #[test] fn user_serializes_created_at_as_rfc3339() { let user = User { diff --git a/crates/crank-observability/src/correlation.rs b/crates/crank-observability/src/correlation.rs index 4d969fe..beaa51d 100644 --- a/crates/crank-observability/src/correlation.rs +++ b/crates/crank-observability/src/correlation.rs @@ -1,5 +1,6 @@ use std::fmt; +use axum::http::HeaderMap; use uuid::Uuid; #[derive(Clone, Debug, Eq, Hash, PartialEq)] @@ -7,6 +8,7 @@ pub struct RequestId(String); impl RequestId { pub const MAX_LEN: usize = 128; + const HEADER_NAME: &'static str = "x-request-id"; pub fn resolve(candidate: Option<&str>) -> Self { candidate @@ -15,6 +17,16 @@ impl RequestId { .unwrap_or_else(|| Self(Uuid::now_v7().to_string())) } + pub fn resolve_from_headers(headers: &HeaderMap) -> Self { + let mut values = headers.get_all(Self::HEADER_NAME).iter(); + let candidate = values.next(); + if values.next().is_some() { + return Self::resolve(None); + } + + Self::resolve(candidate.and_then(|value| value.to_str().ok())) + } + pub fn is_valid(value: &str) -> bool { !value.is_empty() && value.len() <= Self::MAX_LEN diff --git a/crates/crank-observability/src/otlp.rs b/crates/crank-observability/src/otlp.rs index 96b3b46..845dedf 100644 --- a/crates/crank-observability/src/otlp.rs +++ b/crates/crank-observability/src/otlp.rs @@ -597,7 +597,7 @@ mod tests { traces_endpoint: Some("https://traces.example.test/custom".to_owned()), generic_endpoint: Some("https://generic.example.test/otel".to_owned()), traces_protocol: Some("http/protobuf".to_owned()), - generic_protocol: Some("grpc".to_owned()), + generic_protocol: Some("grpc".to_owned()), // community-scope: allow=grpc traces_timeout: Some("2500".to_owned()), generic_timeout: Some("invalid-unused-fallback".to_owned()), ..OtlpEnvSettings::default() @@ -615,8 +615,8 @@ mod tests { #[test] fn disabled_export_ignores_inactive_settings() { let config = OtlpEnvSettings { - traces_protocol: Some("grpc".to_owned()), - generic_protocol: Some("grpc".to_owned()), + traces_protocol: Some("grpc".to_owned()), // community-scope: allow=grpc + generic_protocol: Some("grpc".to_owned()), // community-scope: allow=grpc traces_timeout: Some("invalid".to_owned()), generic_timeout: Some("invalid".to_owned()), max_queue_size: Some("invalid".to_owned()), @@ -637,7 +637,7 @@ mod tests { traces_endpoint: Some("https://traces.example.test/v1/traces".to_owned()), traces_headers: Some(String::new()), generic_headers: Some( - "authorization=Bearer%20canary-token,x-tenant=community".to_owned(), + "authorization=Bearer%20canary-token,x-scope=community".to_owned(), ), ..OtlpEnvSettings::default() } @@ -645,7 +645,7 @@ mod tests { .unwrap(); assert_eq!(config.header("authorization"), Some("Bearer canary-token")); - assert_eq!(config.header("x-tenant"), Some("community")); + assert_eq!(config.header("x-scope"), Some("community")); assert!(!format!("{config:?}").contains("canary-token")); } diff --git a/crates/crank-observability/tests/correlation.rs b/crates/crank-observability/tests/correlation.rs index c5c85c9..1404c56 100644 --- a/crates/crank-observability/tests/correlation.rs +++ b/crates/crank-observability/tests/correlation.rs @@ -1,3 +1,4 @@ +use axum::http::{HeaderMap, HeaderValue}; use crank_observability::RequestId; use uuid::Version; @@ -40,3 +41,34 @@ fn rejects_values_over_the_shared_limit() { Some(Version::SortRand) ); } + +#[test] +fn resolves_exactly_one_header_value_and_rejects_ambiguous_values() { + let mut single = HeaderMap::new(); + single.insert( + "x-request-id", + HeaderValue::from_static("opaque-request-id"), + ); + + assert_eq!( + RequestId::resolve_from_headers(&single).as_str(), + "opaque-request-id" + ); + + let mut ambiguous = HeaderMap::new(); + ambiguous.append("x-request-id", HeaderValue::from_static("first-request-id")); + ambiguous.append( + "x-request-id", + HeaderValue::from_static("second-request-id"), + ); + + let generated = RequestId::resolve_from_headers(&ambiguous); + assert_ne!(generated.as_str(), "first-request-id"); + assert_ne!(generated.as_str(), "second-request-id"); + assert_eq!( + uuid::Uuid::parse_str(generated.as_str()) + .expect("generated UUID") + .get_version(), + Some(Version::SortRand) + ); +} diff --git a/crates/crank-observability/tests/otlp.rs b/crates/crank-observability/tests/otlp.rs index 17cd498..adf3a8c 100644 --- a/crates/crank-observability/tests/otlp.rs +++ b/crates/crank-observability/tests/otlp.rs @@ -37,7 +37,7 @@ fn invalid_values_return_safe_typed_errors() { let secret_endpoint = "https://user:canary-secret@collector.example.test/v1/traces"; let error = OtlpTraceConfig::try_new( Some(secret_endpoint.to_owned()), - Some("grpc".to_owned()), + Some("grpc".to_owned()), // community-scope: allow=grpc Duration::ZERO, OtlpBatchConfig::default(), ) diff --git a/crates/crank-runtime/src/executor.rs b/crates/crank-runtime/src/executor.rs index 06294ca..bce29f2 100644 --- a/crates/crank-runtime/src/executor.rs +++ b/crates/crank-runtime/src/executor.rs @@ -11,6 +11,7 @@ use serde_json::{Map, Value, json}; use time::OffsetDateTime; use tokio::sync::{OwnedSemaphorePermit, Semaphore}; use tracing::{Instrument, Span, debug}; +use uuid::Uuid; use crate::{ AdapterResponse, PreparedRequest, ResolvedAuth, RuntimeError, RuntimeLimits, RuntimeOperation, @@ -517,9 +518,9 @@ impl RuntimeExecutor { fn adapter_request_context( request_context: Option<&RuntimeRequestContext>, ) -> crank_core::RuntimeRequestContext { - request_context - .map(Into::into) - .unwrap_or_else(|| crank_core::RuntimeRequestContext::new(String::new(), String::new())) + request_context.map(Into::into).unwrap_or_else(|| { + crank_core::RuntimeRequestContext::from_request_id(Uuid::now_v7().to_string()) + }) } fn map_protocol_adapter_error( diff --git a/crates/crank-runtime/tests/integration.rs b/crates/crank-runtime/tests/integration.rs index aa706b8..f3ffa21 100644 --- a/crates/crank-runtime/tests/integration.rs +++ b/crates/crank-runtime/tests/integration.rs @@ -2,6 +2,5 @@ mod integration { mod confirmation; mod idempotency; mod no_input_get; - mod stages; mod valkey; } diff --git a/crates/crank-runtime/tests/integration/valkey.rs b/crates/crank-runtime/tests/integration/valkey.rs index 3d952c5..ee780c1 100644 --- a/crates/crank-runtime/tests/integration/valkey.rs +++ b/crates/crank-runtime/tests/integration/valkey.rs @@ -25,8 +25,12 @@ async fn valkey_coordination_and_rate_limit_operations_are_atomic() { .get_host_port_ipv4(6379.tcp()) .await .expect("Valkey port must be mapped"); + let host = container + .get_host() + .await + .expect("Docker host must be resolved"); let store = Arc::new( - RedisCacheStore::connect(CacheBackend::Valkey, &format!("redis://127.0.0.1:{port}/0")) + RedisCacheStore::connect(CacheBackend::Valkey, &format!("redis://{host}:{port}/0")) .await .expect("runtime store must connect to Valkey"), ); diff --git a/crates/crank-runtime/tests/integration/stages.rs b/crates/crank-runtime/tests/stages.rs similarity index 84% rename from crates/crank-runtime/tests/integration/stages.rs rename to crates/crank-runtime/tests/stages.rs index 0bd0348..a52f859 100644 --- a/crates/crank-runtime/tests/integration/stages.rs +++ b/crates/crank-runtime/tests/stages.rs @@ -20,9 +20,13 @@ use serde_json::json; use time::OffsetDateTime; use tracing::{Id, Instrument, Subscriber, field::Visit, instrument::WithSubscriber}; use tracing_subscriber::{Layer, layer::SubscriberExt, registry::LookupSpan}; +use uuid::Version; + +static TRACING_TEST_LOCK: tokio::sync::Mutex<()> = tokio::sync::Mutex::const_new(()); #[tokio::test] async fn successful_execution_has_real_stages_and_omits_inapplicable_ones() { + let _tracing_test_guard = TRACING_TEST_LOCK.lock().await; let capture = TraceCapture::default(); let subscriber = tracing_subscriber::registry().with(capture.clone()); let executor = RuntimeExecutorBuilder::new() @@ -82,6 +86,7 @@ async fn successful_execution_has_real_stages_and_omits_inapplicable_ones() { #[tokio::test] async fn failed_mapping_records_closed_category_and_stops_later_stages() { + let _tracing_test_guard = TRACING_TEST_LOCK.lock().await; let capture = TraceCapture::default(); let subscriber = tracing_subscriber::registry().with(capture.clone()); let executor = RuntimeExecutorBuilder::new() @@ -110,8 +115,46 @@ async fn failed_mapping_records_closed_category_and_stops_later_stages() { ); } +#[tokio::test] +async fn execution_without_context_sends_generated_correlation_headers() { + let _tracing_test_guard = TRACING_TEST_LOCK.lock().await; + let captured_headers = Arc::new(Mutex::new(None)); + let executor = RuntimeExecutorBuilder::new() + .register_adapter(Arc::new(ContextCapturingAdapter { + captured_headers: Arc::clone(&captured_headers), + })) + .build(); + let operation = operation().into(); + + let result = async { + executor + .execute(&operation, &json!({"name": "without-context"})) + .await + } + .with_subscriber(tracing_subscriber::registry()) + .await; + + assert_eq!(result.unwrap(), json!({"accepted": true})); + let headers = captured_headers + .lock() + .expect("captured headers lock") + .clone() + .expect("adapter headers"); + let request_id = headers.get("x-request-id").expect("x-request-id"); + let correlation_id = headers.get("x-correlation-id").expect("x-correlation-id"); + assert!(!request_id.is_empty()); + assert_eq!(correlation_id, request_id); + assert_eq!( + uuid::Uuid::parse_str(request_id) + .expect("generated UUID") + .get_version(), + Some(Version::SortRand) + ); +} + #[tokio::test] async fn approval_stage_is_present_only_when_confirmation_is_required() { + let _tracing_test_guard = TRACING_TEST_LOCK.lock().await; let capture = TraceCapture::default(); let subscriber = tracing_subscriber::registry().with(capture.clone()); let executor = RuntimeExecutorBuilder::new() @@ -151,6 +194,7 @@ async fn approval_stage_is_present_only_when_confirmation_is_required() { #[tokio::test] async fn idempotency_stage_distinguishes_execution_from_replay() { + let _tracing_test_guard = TRACING_TEST_LOCK.lock().await; let capture = TraceCapture::default(); let subscriber = tracing_subscriber::registry().with(capture.clone()); let executor = RuntimeExecutorBuilder::new() @@ -238,6 +282,37 @@ impl ProtocolAdapter for SuccessAdapter { } } +struct ContextCapturingAdapter { + captured_headers: Arc>>>, +} + +#[async_trait] +impl ProtocolAdapter for ContextCapturingAdapter { + fn protocol(&self) -> Protocol { + Protocol::Rest + } + + fn supports_mode(&self, mode: ExecutionMode) -> bool { + mode == ExecutionMode::Unary + } + + async fn invoke_unary( + &self, + _target: &Target, + _prepared: &crank_core::PreparedRequest, + context: &crank_core::RuntimeRequestContext, + ) -> Result { + *self.captured_headers.lock().expect("captured headers lock") = + Some(context.outbound_headers()); + Ok(AdapterResponse { + status_code: 200, + headers: BTreeMap::new(), + body: json!({"accepted": true}), + data: json!({"accepted": true}), + }) + } +} + fn operation() -> Operation { Operation { id: OperationId::new("op_stage_test"), diff --git a/crates/crank-trace/src/lib.rs b/crates/crank-trace/src/lib.rs index e4c6d27..761e5c5 100644 --- a/crates/crank-trace/src/lib.rs +++ b/crates/crank-trace/src/lib.rs @@ -63,24 +63,13 @@ impl Stage { ), } } - - pub fn db_span(self, operation: DbOperation) -> Option { - if self != Self::DbQuery { - return None; - } - let span = self.span(); - span.record("db.operation", operation.as_str()); - Some(span) - } } pub async fn observe_db_query( operation: DbOperation, future: impl Future>, ) -> Result { - let span = Stage::DbQuery - .db_span(operation) - .expect("database operation requires db.query stage"); + let span = operation.span(); let result = future.instrument(span.clone()).await; match &result { Ok(_) => StageOutcome::Success.record(&span), @@ -182,6 +171,12 @@ pub enum DbOperation { } impl DbOperation { + pub fn span(self) -> Span { + let span = Stage::DbQuery.span(); + span.record("db.operation", self.as_str()); + span + } + pub const fn as_str(self) -> &'static str { match self { Self::MachineAccessRead => "machine_access.read", diff --git a/crates/crank-trace/tests/contract.rs b/crates/crank-trace/tests/contract.rs index 7a54b36..ae1edc2 100644 --- a/crates/crank-trace/tests/contract.rs +++ b/crates/crank-trace/tests/contract.rs @@ -15,9 +15,7 @@ fn stage_names_and_attributes_are_closed() { ErrorCategory::Mapping.record(&span); drop(span); - let db_span = Stage::DbQuery - .db_span(DbOperation::InvocationHistoryWrite) - .expect("db stage accepts a db operation"); + let db_span = DbOperation::InvocationHistoryWrite.span(); StageOutcome::Error.record(&db_span); drop(db_span); }); @@ -32,12 +30,8 @@ fn stage_names_and_attributes_are_closed() { } #[test] -fn non_database_stage_rejects_database_attributes() { - assert!( - Stage::RuntimeExecute - .db_span(DbOperation::CatalogLoad) - .is_none() - ); +fn database_operation_names_are_closed() { + assert_eq!(DbOperation::CatalogLoad.as_str(), "catalog.load"); } #[derive(Clone)] diff --git a/scripts/check-community-scope.py b/scripts/check-community-scope.py index 09bd9f1..7f79dac 100755 --- a/scripts/check-community-scope.py +++ b/scripts/check-community-scope.py @@ -20,6 +20,11 @@ DEFAULT_EXCLUDED_PREFIXES = { "apps/ui/dist/", } +ALLOW_DIRECTIVE = re.compile( + r"community-scope:\s*allow=([a-z0-9-]+(?:,[a-z0-9-]+)*)", + re.IGNORECASE, +) + FORBIDDEN_PATTERNS = [ ("enterprise", re.compile(r"\benterprise\b", re.IGNORECASE)), ("cloud", re.compile(r"\bcloud\b", re.IGNORECASE)), @@ -100,7 +105,15 @@ def scan_file(root: Path, relative_path: str) -> list[str]: text = data.decode("utf-8", errors="replace") findings: list[str] = [] for line_no, line in enumerate(text.splitlines(), start=1): + directive = ALLOW_DIRECTIVE.search(line) + allowed_markers = ( + {label.lower() for label in directive.group(1).split(",")} + if directive + else set() + ) for label, pattern in FORBIDDEN_PATTERNS: + if label in allowed_markers: + continue if pattern.search(line): findings.append(f"{relative_path}:{line_no}: forbidden marker `{label}`") return findings diff --git a/tests/unit/test_check_community_scope.py b/tests/unit/test_check_community_scope.py index f8195cc..73b03ee 100644 --- a/tests/unit/test_check_community_scope.py +++ b/tests/unit/test_check_community_scope.py @@ -47,6 +47,38 @@ class CommunityScopeCheckTests(unittest.TestCase): self.assertIn("apps/ui.md:1", result.stderr) self.assertIn("graphql", result.stderr.lower()) + def test_inline_directive_allows_only_declared_marker(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + root = Path(tmp) + (root / "protocol.rs").write_text( + 'assert_rejected_protocol("grpc"); ' + "// community-scope: allow=grpc\n" + 'assert_rejected_protocol("grpc");\n' + 'assert_rejected_protocol("graphql"); ' + "// community-scope: allow=grpc\n", + encoding="utf-8", + ) + + result = self.run_checker(root, ["protocol.rs"]) + + self.assertNotEqual(result.returncode, 0) + self.assertIn("protocol.rs:2: forbidden marker `grpc`", result.stderr) + self.assertIn("protocol.rs:3: forbidden marker `graphql`", result.stderr) + self.assertNotIn("protocol.rs:1:", result.stderr) + + def test_rejects_marker_without_inline_directive(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + root = Path(tmp) + (root / "README.md").write_text( + "Enable grpc transport.\n", + encoding="utf-8", + ) + + result = self.run_checker(root, ["README.md"]) + + self.assertNotEqual(result.returncode, 0) + self.assertIn("forbidden marker `grpc`", result.stderr) + def test_ignores_default_excluded_paths(self) -> None: with tempfile.TemporaryDirectory() as tmp: root = Path(tmp)