наблюдаемость: ввести безопасный контракт метрик
CI / Rust Checks (pull_request) Successful in 6m15s
CI / UI Checks (pull_request) Successful in 5s
CI / Community Image Smoke (pull_request) Successful in 4m25s
CI / Frontend E2E (pull_request) Successful in 5m17s
CI / Deploy (pull_request) Has been skipped
CI / Rust Checks (push) Successful in 6m9s
CI / UI Checks (push) Successful in 5s
CI / Community Image Smoke (push) Successful in 1m3s
CI / Frontend E2E (push) Successful in 3m47s
CI / Deploy (push) Failing after 3s

This commit is contained in:
2026-07-31 05:04:01 +03:00
parent ec2453c00f
commit 9b1a739e39
50 changed files with 3066 additions and 433 deletions
+2 -1
View File
@@ -12,13 +12,13 @@ axum.workspace = true
base64.workspace = true
crank-adapter-rest = { path = "../crank-adapter-rest" }
crank-core = { path = "../crank-core" }
crank-metrics = { path = "../crank-metrics" }
crank-observability = { path = "../crank-observability" }
crank-registry = { path = "../crank-registry" }
crank-runtime = { path = "../crank-runtime" }
crank-schema = { path = "../crank-schema" }
crank-trace = { path = "../crank-trace" }
futures-util = "0.3"
metrics.workspace = true
reqwest.workspace = true
serde.workspace = true
serde_json.workspace = true
@@ -33,6 +33,7 @@ uuid.workspace = true
[dev-dependencies]
crank-mapping = { path = "../crank-mapping" }
crank-test-support = { path = "../crank-test-support" }
metrics-util = "0.20.4"
opentelemetry.workspace = true
opentelemetry_sdk.workspace = true
tracing-opentelemetry.workspace = true
+89 -97
View File
@@ -7,7 +7,7 @@ use std::{
use axum::{
Json, Router,
extract::{Extension, Path, State},
extract::{Extension, Path, State, rejection::JsonRejection},
http::{HeaderMap, StatusCode},
response::{IntoResponse, Response, sse::Event},
routing::{get, post},
@@ -31,7 +31,7 @@ use serde::{Deserialize, Serialize};
use serde_json::{Value, json};
use time::OffsetDateTime;
use tokio::sync::Semaphore;
use tracing::{Instrument, info, warn};
use tracing::{Instrument, info};
use crate::{
access::{
@@ -48,23 +48,23 @@ use crate::{
manifest::catalog_tool_definitions,
rate_limit::{rate_limited_jsonrpc_response, rate_limited_status_response},
request_context::{RequestContext, apply_request_context},
session::{SessionState, SharedSessionStore},
session::{ActiveSessionMetrics, SessionState, SharedSessionStore, spawn_session_cleanup},
tool_error::{
ToolErrorContract, generic_tool_error_contract, runtime_error_code,
tool_error_contract_from_runtime, tool_error_text, tool_error_value,
},
tool_search::handle_catalog_tool_call,
transport::{
AllowedOrigins, HEADER_MCP_SESSION_ID, ResponseMode, json_response,
negotiate_post_response_mode, protocol_version_from_headers, session_id_from_headers,
sse_response, transport_response, validate_get_accept_header, validate_origin,
validate_session_protocol_version, with_request_id_header,
AllowedOrigins, ResponseMode, json_response, negotiate_post_response_mode,
protocol_version_from_headers, session_id_from_headers, sse_response, transport_response,
validate_get_accept_header, validate_origin, validate_session_protocol_version,
with_request_id_header,
},
};
mod invocation_history;
mod metrics;
mod stages;
use self::metrics::{ActiveSessionGuard, McpRequestMetrics};
use self::metrics::{ActiveStreamGuard, McpRequestMetrics};
use self::stages::{
enforce_traced_rate_limit, require_traced_approval_access, require_traced_machine_access,
};
@@ -84,6 +84,7 @@ pub(super) struct AppState {
pub(super) api_rate_limiter: RequestRateLimiter,
secret_crypto: SecretCrypto,
sessions: SharedSessionStore,
session_metrics: ActiveSessionMetrics,
session_slots: Arc<Semaphore>,
pub(super) credential_verifier: SharedMachineCredentialVerifier,
allowed_origins: AllowedOrigins,
@@ -230,6 +231,7 @@ fn build_app_inner(
max_concurrent_sessions: usize,
start_background_workers: bool,
) -> Router {
let session_metrics = ActiveSessionMetrics::start(Arc::clone(&sessions));
let state = Arc::new(AppState {
registry: registry.clone(),
catalog: PublishedToolCatalog::new(registry, refresh_interval, coordination_store),
@@ -237,13 +239,18 @@ fn build_app_inner(
api_rate_limiter,
secret_crypto,
sessions,
session_metrics,
session_slots: Arc::new(Semaphore::new(max_concurrent_sessions)),
credential_verifier,
allowed_origins: AllowedOrigins::new(public_base_url),
});
if start_background_workers {
spawn_approval_recovery(Arc::clone(&state));
spawn_session_cleanup(Arc::clone(&state));
spawn_session_cleanup(
Arc::clone(&state.sessions),
state.session_metrics.clone(),
SESSION_CLEANUP_INTERVAL,
);
}
Router::new()
@@ -283,28 +290,6 @@ async fn health() -> Json<Value> {
}))
}
fn spawn_session_cleanup(state: Arc<AppState>) {
tokio::spawn(async move {
let mut interval = tokio::time::interval(SESSION_CLEANUP_INTERVAL);
loop {
interval.tick().await;
match state
.sessions
.cleanup_expired(OffsetDateTime::now_utc())
.await
{
Ok(removed) if removed > 0 => {
info!(name: "mcp.session_cleanup.completed", removed);
}
Ok(_) => {}
Err(_) => {
warn!(name: "mcp.session_cleanup.failed", error_category = "session_store");
}
}
}
});
}
async fn readiness(State(state): State<Arc<AppState>>) -> Response {
match state.registry.ping().await {
Ok(()) => Json(json!({
@@ -664,7 +649,7 @@ async fn mcp_get(
return status.into_response();
}
let Ok(permit) = ActiveSessionGuard::try_acquire(&state.session_slots) else {
let Ok(permit) = ActiveStreamGuard::try_acquire(&state.session_slots) else {
return StatusCode::TOO_MANY_REQUESTS.into_response();
};
@@ -704,7 +689,10 @@ async fn mcp_delete(
&& session.agent_slug == path.agent_slug =>
{
match state.sessions.delete(&session_id).await {
Ok(true) => StatusCode::NO_CONTENT.into_response(),
Ok(true) => {
state.session_metrics.refresh();
StatusCode::NO_CONTENT.into_response()
}
Ok(false) => StatusCode::NOT_FOUND.into_response(),
Err(_) => StatusCode::INTERNAL_SERVER_ERROR.into_response(),
}
@@ -723,8 +711,17 @@ async fn mcp_post(
State(state): State<Arc<AppState>>,
Extension(request_context): Extension<RequestContext>,
headers: HeaderMap,
Json(message): Json<Value>,
payload: Result<Json<Value>, JsonRejection>,
) -> Response {
let message = match payload {
Ok(Json(message)) => message,
Err(rejection) => {
let mut request_metrics = McpRequestMetrics::invalid();
let response = rejection.into_response();
request_metrics.complete(&response);
return with_request_id_header(response, &request_context.request_id);
}
};
let mut request_metrics = McpRequestMetrics::new(&message);
let transport_request_id = request_context.request_id;
info!(
@@ -736,78 +733,77 @@ async fn mcp_post(
"mcp request received"
);
if let Err(status) = validate_origin(&state.allowed_origins, &headers) {
return with_request_id_header(status.into_response(), &transport_request_id);
let response = mcp_post_response(&path, state, &headers, &message, &transport_request_id).await;
request_metrics.complete(&response);
with_request_id_header(response, &transport_request_id)
}
async fn mcp_post_response(
path: &AgentRoutePath,
state: Arc<AppState>,
headers: &HeaderMap,
message: &Value,
transport_request_id: &str,
) -> Response {
if let Err(status) = validate_origin(&state.allowed_origins, headers) {
return status.into_response();
}
let response_mode = match negotiate_post_response_mode(&headers) {
Ok(mode) => request_metrics.set_response_mode(mode),
Err(status) => {
return with_request_id_header(status.into_response(), &transport_request_id);
}
let response_mode = match negotiate_post_response_mode(headers) {
Ok(mode) => mode,
Err(status) => return status.into_response(),
};
if is_response(&message) || is_notification(&message) && method_name(&message).is_none() {
return with_request_id_header(StatusCode::ACCEPTED.into_response(), &transport_request_id);
let request_session_id = match session_id_from_headers(headers) {
Ok(session_id) => session_id,
Err(status) => return status.into_response(),
};
if is_response(message) || is_notification(message) && method_name(message).is_none() {
return StatusCode::ACCEPTED.into_response();
}
let protocol_version = match protocol_version_from_headers(&headers) {
let protocol_version = match protocol_version_from_headers(headers) {
Ok(value) => value,
Err(status) => {
return with_request_id_header(status.into_response(), &transport_request_id);
}
Err(status) => return status.into_response(),
};
let rate_limit_result = enforce_traced_rate_limit(&state, &path, &headers).await;
let rate_limit_result = enforce_traced_rate_limit(&state, path, headers).await;
if let Err(error) = rate_limit_result {
return with_request_id_header(
rate_limited_jsonrpc_response(&message, response_mode, &protocol_version, error),
&transport_request_id,
);
return rate_limited_jsonrpc_response(message, response_mode, &protocol_version, error);
}
if let Some(session_id) = headers.get(HEADER_MCP_SESSION_ID)
&& let Ok(session_id) = session_id.to_str()
{
if let Some(session_id) = request_session_id.as_deref() {
let session = match state.sessions.get(session_id).await {
Ok(session) => session,
Err(_) => {
return with_request_id_header(
StatusCode::INTERNAL_SERVER_ERROR.into_response(),
&transport_request_id,
);
}
Err(_) => return StatusCode::INTERNAL_SERVER_ERROR.into_response(),
};
if let Some(session) = session
&& let Err(status) =
validate_session_protocol_version(&headers, &session.protocol_version)
validate_session_protocol_version(headers, &session.protocol_version)
{
return with_request_id_header(status.into_response(), &transport_request_id);
return status.into_response();
}
}
let required_scope = match method_name(&message) {
let required_scope = match method_name(message) {
Some("tools/call") => PlatformApiKeyScope::Write,
_ => PlatformApiKeyScope::Read,
};
let access_result =
require_traced_machine_access(&state, &path, &headers, required_scope).await;
let access_result = require_traced_machine_access(&state, path, headers, required_scope).await;
let credential = match access_result {
Ok(credential) => credential,
Err(error) => {
return with_request_id_header(error.into_response(), &transport_request_id);
}
Err(error) => return error.into_response(),
};
let response = match method_name(&message) {
Some("initialize") if is_request(&message) => {
handle_initialize(state, &path, &message, response_mode).await
match method_name(message) {
Some("initialize") if is_request(message) => {
handle_initialize(state, path, message, response_mode).await
}
Some("notifications/initialized") if is_notification(&message) => {
handle_initialized_notification(state, &path, &headers).await
Some("notifications/initialized") if is_notification(message) => {
handle_initialized_notification(state, path, headers).await
}
Some("ping") if is_request(&message) => {
let session = match require_initialized_session(&state, &path, &headers, &message).await
{
Some("ping") if is_request(message) => {
let session = match require_initialized_session(&state, path, headers, message).await {
Ok(session) => session,
Err(response) => return response,
};
@@ -815,7 +811,7 @@ async fn mcp_post(
transport_response(
StatusCode::OK,
jsonrpc_result(
request_id(&message),
request_id(message),
json!({ "protocolVersion": session.protocol_version }),
),
response_mode,
@@ -823,9 +819,8 @@ async fn mcp_post(
Some(&session.protocol_version),
)
}
Some("tools/list") if is_request(&message) => {
let session = match require_initialized_session(&state, &path, &headers, &message).await
{
Some("tools/list") if is_request(message) => {
let session = match require_initialized_session(&state, path, headers, message).await {
Ok(session) => session,
Err(response) => return response,
};
@@ -840,27 +835,26 @@ async fn mcp_post(
transport_response(
StatusCode::OK,
jsonrpc_result(request_id(&message), json!({ "tools": definitions })),
jsonrpc_result(request_id(message), json!({ "tools": definitions })),
response_mode,
None,
Some(&session.protocol_version),
)
}
Err(error) => internal_jsonrpc_error(&message, error),
Err(error) => internal_jsonrpc_error(message, error),
}
}
Some("tools/call") if is_request(&message) => {
let session = match require_initialized_session(&state, &path, &headers, &message).await
{
Some("tools/call") if is_request(message) => {
let session = match require_initialized_session(&state, path, headers, message).await {
Ok(session) => session,
Err(response) => return response,
};
let tool_call_params: ToolCallParams = match serde_json::from_value(params(&message)) {
let tool_call_params: ToolCallParams = match serde_json::from_value(params(message)) {
Ok(value) => value,
Err(error) => {
return transport_response(
StatusCode::OK,
jsonrpc_error(request_id(&message), -32602, error.to_string()),
jsonrpc_error(request_id(message), -32602, error.to_string()),
response_mode,
None,
Some(&session.protocol_version),
@@ -883,27 +877,27 @@ async fn mcp_post(
handle_catalog_tool_call(
state.clone(),
&session,
&message,
message,
response_mode,
&credential,
&catalog,
&tool_call_params.name,
arguments,
&transport_request_id,
transport_request_id,
)
.await
}
Err(error) => internal_jsonrpc_error(&message, error),
Err(error) => internal_jsonrpc_error(message, error),
}
}
Some(method) if is_notification(&message) => {
Some(method) if is_notification(message) => {
let _ = method;
StatusCode::ACCEPTED.into_response()
}
Some(method) => transport_response(
StatusCode::OK,
jsonrpc_error(
request_id(&message),
request_id(message),
-32601,
format!("method {method} is not supported"),
),
@@ -918,10 +912,7 @@ async fn mcp_post(
None,
Some(&protocol_version),
),
};
request_metrics.complete(response.status());
with_request_id_header(response, &transport_request_id)
}
}
#[allow(clippy::too_many_arguments)]
@@ -1432,6 +1423,7 @@ async fn handle_initialize(
Ok(session_id) => session_id,
Err(error) => return internal_jsonrpc_error(message, error),
};
state.session_metrics.refresh();
transport_response(
StatusCode::OK,
+70 -56
View File
@@ -1,93 +1,107 @@
use std::sync::Arc;
use axum::http::StatusCode;
use axum::{http::header::CONTENT_TYPE, response::Response};
use crank_metrics::{
InFlightGuard, LimitStage, McpMethod, McpOutcome, McpResponseMode, record_limit_rejection,
record_mcp_request,
};
use serde_json::Value;
use tokio::sync::{OwnedSemaphorePermit, Semaphore};
use crate::{
jsonrpc::{is_notification, is_response, method_name},
transport::ResponseMode,
};
use crate::jsonrpc::{is_notification, is_response, method_name};
pub(super) struct McpRequestMetrics {
method: &'static str,
response_mode: &'static str,
outcome: &'static str,
method: McpMethod,
response_mode: McpResponseMode,
outcome: McpOutcome,
}
impl McpRequestMetrics {
pub(super) fn new(message: &Value) -> Self {
pub(super) const fn invalid() -> Self {
Self {
method: normalized_mcp_method(message),
response_mode: "unknown",
outcome: "rejected",
method: McpMethod::Invalid,
response_mode: McpResponseMode::Unknown,
outcome: McpOutcome::Aborted,
}
}
pub(super) fn set_response_mode(&mut self, mode: ResponseMode) -> ResponseMode {
self.response_mode = match mode {
ResponseMode::Json => "json",
ResponseMode::Sse => "sse",
};
mode
pub(super) fn new(message: &Value) -> Self {
Self {
method: normalized_mcp_method(message),
response_mode: McpResponseMode::Unknown,
outcome: McpOutcome::Aborted,
}
}
pub(super) fn complete(&mut self, status: StatusCode) {
self.outcome = match status.as_u16() {
200..=299 => "success",
400..=499 => "client_error",
500..=599 => "server_error",
_ => "other",
pub(super) fn complete(&mut self, response: &Response) {
self.response_mode = response
.headers()
.get(CONTENT_TYPE)
.and_then(|value| value.to_str().ok())
.map_or(McpResponseMode::Unknown, |content_type| {
if content_type.starts_with("application/json") {
McpResponseMode::Json
} else if content_type.starts_with("text/event-stream") {
McpResponseMode::Sse
} else {
McpResponseMode::Unknown
}
});
self.outcome = if response.status().is_success() {
response
.extensions()
.get::<McpOutcome>()
.copied()
.unwrap_or_else(|| McpOutcome::from_http_status(response.status().as_u16()))
} else {
McpOutcome::from_http_status(response.status().as_u16())
};
}
#[cfg(test)]
pub(super) const fn outcome(&self) -> McpOutcome {
self.outcome
}
#[cfg(test)]
pub(super) const fn response_mode(&self) -> McpResponseMode {
self.response_mode
}
}
impl Drop for McpRequestMetrics {
fn drop(&mut self) {
::metrics::counter!(
"crank_mcp_requests_total",
"method" => self.method,
"response_mode" => self.response_mode,
"outcome" => self.outcome
)
.increment(1);
record_mcp_request(self.method, self.response_mode, self.outcome);
}
}
pub(super) fn normalized_mcp_method(message: &Value) -> &'static str {
pub(super) fn normalized_mcp_method(message: &Value) -> McpMethod {
match method_name(message) {
Some("initialize") => "initialize",
Some("notifications/initialized") => "initialized",
Some("ping") => "ping",
Some("tools/list") => "tools_list",
Some("tools/call") => "tools_call",
Some(_) if is_notification(message) => "notification",
Some(_) => "unsupported",
None if is_response(message) => "response",
None => "invalid",
Some("initialize") => McpMethod::Initialize,
Some("notifications/initialized") => McpMethod::Initialized,
Some("ping") => McpMethod::Ping,
Some("tools/list") => McpMethod::ToolsList,
Some("tools/call") => McpMethod::ToolsCall,
Some(_) if is_notification(message) => McpMethod::Notification,
Some(_) => McpMethod::Unsupported,
None if is_response(message) => McpMethod::Response,
None => McpMethod::Invalid,
}
}
pub(super) struct ActiveSessionGuard {
pub(super) struct ActiveStreamGuard {
_permit: OwnedSemaphorePermit,
_inflight: InFlightGuard,
}
impl ActiveSessionGuard {
impl ActiveStreamGuard {
pub(super) fn try_acquire(slots: &Arc<Semaphore>) -> Result<Self, ()> {
let permit = Arc::clone(slots).try_acquire_owned().map_err(|_| {
::metrics::counter!(
"crank_runtime_limit_rejections_total",
"stage" => "mcp_session"
)
.increment(1);
record_limit_rejection(LimitStage::McpStream);
})?;
::metrics::gauge!("crank_mcp_active_sessions").increment(1.0);
Ok(Self { _permit: permit })
}
}
impl Drop for ActiveSessionGuard {
fn drop(&mut self) {
::metrics::gauge!("crank_mcp_active_sessions").decrement(1.0);
Ok(Self {
_permit: permit,
_inflight: InFlightGuard::mcp_stream(),
})
}
}
+79 -7
View File
@@ -5,6 +5,7 @@ use std::{
use axum::body::to_bytes;
use crank_core::InvocationStatus;
use crank_metrics::{McpMethod, McpOutcome, McpResponseMode};
use crank_observability::{
ObservabilityConfig, OperationalIncident, RedactionLimits, ServiceIdentity,
operational_incident_total,
@@ -16,16 +17,20 @@ use serde_json::{Value, json};
use tracing_subscriber::fmt::MakeWriter;
use super::{
ResponseMode, metrics::normalized_mcp_method, observe_invocation_history_outcome,
tool_error_response,
ResponseMode,
metrics::{McpRequestMetrics, normalized_mcp_method},
observe_invocation_history_outcome, tool_error_response,
};
use crate::jsonrpc::CURRENT_PROTOCOL_VERSION;
use crate::jsonrpc::{CURRENT_PROTOCOL_VERSION, jsonrpc_error};
use crate::tool_error::generic_tool_error_contract;
use crate::transport::transport_response;
#[tokio::test]
async fn tool_error_response_includes_structured_context() {
let message = json!({"jsonrpc": "2.0", "id": "req-1", "method": "tools/call"});
let mut request_metrics = McpRequestMetrics::new(&message);
let response = tool_error_response(
&json!({"jsonrpc": "2.0", "id": "req-1"}),
&message,
ResponseMode::Json,
CURRENT_PROTOCOL_VERSION,
generic_tool_error_contract(
@@ -36,6 +41,12 @@ async fn tool_error_response_includes_structured_context() {
Some("Проверьте параметры вызова инструмента."),
),
);
assert_eq!(
response.extensions().get::<McpOutcome>(),
Some(&McpOutcome::ToolError)
);
request_metrics.complete(&response);
assert_eq!(request_metrics.outcome(), McpOutcome::ToolError);
let body = to_bytes(response.into_body(), usize::MAX).await.unwrap();
let payload: Value = serde_json::from_slice(&body).unwrap();
@@ -53,6 +64,67 @@ async fn tool_error_response_includes_structured_context() {
);
}
#[test]
fn jsonrpc_error_over_http_200_is_not_counted_as_success() {
let message = json!({"jsonrpc": "2.0", "id": 1, "method": "unsupported"});
let mut request_metrics = McpRequestMetrics::new(&message);
let response = transport_response(
axum::http::StatusCode::OK,
jsonrpc_error(json!(1), -32601, "unsupported"),
ResponseMode::Json,
None,
Some(CURRENT_PROTOCOL_VERSION),
);
request_metrics.complete(&response);
assert_eq!(request_metrics.outcome(), McpOutcome::JsonRpcError);
}
#[test]
fn metric_uses_actual_json_response_when_sse_request_falls_back() {
let message = json!({"jsonrpc": "2.0", "id": 1, "method": "unsupported"});
let mut request_metrics = McpRequestMetrics::new(&message);
let response = transport_response(
axum::http::StatusCode::INTERNAL_SERVER_ERROR,
jsonrpc_error(json!(1), -32603, "internal error"),
ResponseMode::Json,
None,
Some(CURRENT_PROTOCOL_VERSION),
);
request_metrics.complete(&response);
assert_eq!(request_metrics.response_mode(), McpResponseMode::Json);
assert_eq!(request_metrics.outcome(), McpOutcome::ServerError);
}
#[test]
fn transport_failure_takes_priority_over_jsonrpc_payload() {
let message = json!({"jsonrpc": "2.0", "id": 1, "method": "tools/call"});
let mut request_metrics = McpRequestMetrics::new(&message);
let response = transport_response(
axum::http::StatusCode::TOO_MANY_REQUESTS,
jsonrpc_error(json!(1), -32000, "rate limited"),
ResponseMode::Json,
None,
Some(CURRENT_PROTOCOL_VERSION),
);
request_metrics.complete(&response);
assert_eq!(request_metrics.outcome(), McpOutcome::ClientError);
}
#[test]
fn unfinished_request_is_classified_as_aborted() {
let message = json!({"jsonrpc": "2.0", "id": 1, "method": "tools/call"});
let request_metrics = McpRequestMetrics::new(&message);
assert_eq!(request_metrics.response_mode(), McpResponseMode::Unknown);
assert_eq!(request_metrics.outcome(), McpOutcome::Aborted);
}
#[test]
fn emits_bounded_history_loss_incident() {
let writer = SharedLogWriter::default();
@@ -96,19 +168,19 @@ fn emits_bounded_history_loss_incident() {
fn mcp_metric_method_is_always_from_a_closed_set() {
assert_eq!(
normalized_mcp_method(&json!({"jsonrpc": "2.0", "id": 1, "method": "tools/call"})),
"tools_call"
McpMethod::ToolsCall
);
assert_eq!(
normalized_mcp_method(
&json!({"jsonrpc": "2.0", "id": 2, "method": "customer-controlled-method"})
),
"unsupported"
McpMethod::Unsupported
);
assert_eq!(
normalized_mcp_method(
&json!({"jsonrpc": "2.0", "method": "customer-controlled-notification"})
),
"notification"
McpMethod::Notification
);
}
+5 -4
View File
@@ -339,10 +339,11 @@ fn record_catalog_metrics(metrics: impl Iterator<Item = CatalogMetrics>) {
aggregate
});
metrics::gauge!("crank_catalog_tools").set(aggregate.tool_count as f64);
metrics::gauge!("crank_catalog_estimated_context_tokens")
.set(aggregate.estimated_context_tokens as f64);
metrics::gauge!("crank_catalog_warnings").set(aggregate.warning_count as f64);
crank_metrics::set_catalog(
aggregate.tool_count,
aggregate.estimated_context_tokens,
aggregate.warning_count,
);
}
fn now_unix_ms() -> u64 {
+93 -1
View File
@@ -9,9 +9,12 @@ use sqlx::{
};
use thiserror::Error;
use time::OffsetDateTime;
use tokio::sync::RwLock;
use tokio::sync::{RwLock, mpsc};
use tracing::{info, warn};
use uuid::Uuid;
const ACTIVE_SESSION_COUNT_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(2);
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct SessionState {
pub id: String,
@@ -54,10 +57,72 @@ pub trait TransportSessionStore: Send + Sync {
async fn delete(&self, session_id: &str) -> Result<bool, SessionStoreError>;
async fn cleanup_expired(&self, now: OffsetDateTime) -> Result<u64, SessionStoreError>;
async fn active_count(&self, now: OffsetDateTime) -> Result<u64, SessionStoreError>;
}
pub type SharedSessionStore = Arc<dyn TransportSessionStore>;
#[derive(Clone)]
pub struct ActiveSessionMetrics {
refresh_tx: mpsc::Sender<()>,
}
impl ActiveSessionMetrics {
pub fn start(sessions: SharedSessionStore) -> Self {
let (refresh_tx, mut refresh_rx) = mpsc::channel(1);
tokio::spawn(async move {
while refresh_rx.recv().await.is_some() {
match tokio::time::timeout(
ACTIVE_SESSION_COUNT_TIMEOUT,
sessions.active_count(OffsetDateTime::now_utc()),
)
.await
{
Ok(Ok(count)) => crank_metrics::set_mcp_active_sessions(count),
Ok(Err(_)) | Err(_) => {
warn!(
name: "mcp.active_session_metrics.refresh_failed",
error_category = "session_store",
"active session metrics refresh failed"
);
}
}
}
});
let metrics = Self { refresh_tx };
metrics.refresh();
metrics
}
pub fn refresh(&self) {
let _ = self.refresh_tx.try_send(());
}
}
pub fn spawn_session_cleanup(
sessions: SharedSessionStore,
metrics: ActiveSessionMetrics,
cleanup_interval: std::time::Duration,
) {
tokio::spawn(async move {
let mut interval = tokio::time::interval(cleanup_interval);
loop {
interval.tick().await;
match sessions.cleanup_expired(OffsetDateTime::now_utc()).await {
Ok(removed) if removed > 0 => {
info!(name: "mcp.session_cleanup.completed", removed);
}
Ok(_) => {}
Err(_) => {
warn!(name: "mcp.session_cleanup.failed", error_category = "session_store");
}
}
metrics.refresh();
}
});
}
#[derive(Clone, Debug)]
pub struct PostgresTransportSessionStore {
pool: PgPool,
@@ -176,6 +241,17 @@ impl TransportSessionStore for InMemorySessionStore {
guard.retain(|_, session| !is_expired(session, now));
Ok(u64::try_from(before.saturating_sub(guard.len())).unwrap_or(u64::MAX))
}
async fn active_count(&self, now: OffsetDateTime) -> Result<u64, SessionStoreError> {
let guard = self.inner.read().await;
Ok(u64::try_from(
guard
.values()
.filter(|session| !is_expired(session, now))
.count(),
)
.unwrap_or(u64::MAX))
}
}
#[async_trait]
@@ -319,6 +395,22 @@ impl TransportSessionStore for PostgresTransportSessionStore {
Ok(result.rows_affected())
}
async fn active_count(&self, now: OffsetDateTime) -> Result<u64, SessionStoreError> {
let row = query(
"select count(*)::bigint as active_count
from mcp_transport_sessions
where expires_at is null or expires_at > $1::timestamptz",
)
.bind(now)
.fetch_one(&self.pool)
.await
.map_err(|error| SessionStoreError {
details: error.to_string(),
})?;
let count = row.get::<i64, _>("active_count");
Ok(u64::try_from(count).unwrap_or_default())
}
}
async fn apply_postgres_migrations(pool: &PgPool) -> Result<(), SessionStoreError> {
+25 -1
View File
@@ -11,6 +11,7 @@ use axum::{
sse::{Event, KeepAlive, Sse},
},
};
use crank_metrics::McpOutcome;
use futures_util::stream;
use reqwest::Url;
use serde_json::Value;
@@ -223,7 +224,11 @@ pub(super) fn json_response(
session_id: Option<&str>,
protocol_version: Option<&str>,
) -> Response {
let outcome = payload_outcome(&payload);
let mut response = (status, Json(payload)).into_response();
if let Some(outcome) = outcome {
response.extensions_mut().insert(outcome);
}
if let Some(session_id) = session_id {
response.headers_mut().insert(
@@ -251,16 +256,35 @@ pub(super) fn transport_response(
session_id: Option<&str>,
protocol_version: Option<&str>,
) -> Response {
let outcome = payload_outcome(&payload);
if status == StatusCode::OK && matches!(response_mode, ResponseMode::Sse) {
let payload = payload.to_string();
let stream = stream::once(async move { Ok(Event::default().data(payload)) });
return sse_response(status, stream, session_id, protocol_version);
let mut response = sse_response(status, stream, session_id, protocol_version);
if let Some(outcome) = outcome {
response.extensions_mut().insert(outcome);
}
return response;
}
json_response(status, payload, session_id, protocol_version)
}
fn payload_outcome(payload: &Value) -> Option<McpOutcome> {
if payload.get("error").is_some() {
return Some(McpOutcome::JsonRpcError);
}
if payload
.pointer("/result/isError")
.and_then(Value::as_bool)
.unwrap_or(false)
{
return Some(McpOutcome::ToolError);
}
None
}
pub(super) fn with_request_id_header(mut response: Response, request_id: &str) -> Response {
if let Ok(value) = HeaderValue::from_str(request_id) {
response.headers_mut().insert(HEADER_X_REQUEST_ID, value);
@@ -133,5 +133,9 @@ async fn postgres_transport_session_cleanup_removes_abandoned_expired_rows() {
.unwrap();
assert_eq!(store.cleanup_expired(now).await.unwrap(), 1);
assert_eq!(store.active_count(now).await.unwrap(), 1);
assert!(store.get(&active).await.unwrap().is_some());
assert!(store.delete(&active).await.unwrap());
assert_eq!(store.active_count(now).await.unwrap(), 0);
}
@@ -0,0 +1,76 @@
use std::{sync::Arc, time::Duration};
use crank_community_mcp::session::{
ActiveSessionMetrics, InMemorySessionStore, SharedSessionStore, TransportSessionStore,
spawn_session_cleanup,
};
use metrics_util::debugging::{DebugValue, DebuggingRecorder, Snapshotter};
use time::OffsetDateTime;
#[tokio::test]
async fn active_session_sampler_tracks_create_delete_and_expiry() {
let recorder = DebuggingRecorder::new();
let snapshotter = recorder.snapshotter();
recorder
.install()
.expect("isolated integration test recorder");
let store = Arc::new(InMemorySessionStore::default());
let sessions: SharedSessionStore = store.clone();
let sampler = ActiveSessionMetrics::start(Arc::clone(&sessions));
assert_gauge(&sampler, &snapshotter, 0.0).await;
let now = OffsetDateTime::now_utc();
store
.create(
"2025-11-25",
"default",
"agent",
false,
now,
Some(now + time::Duration::milliseconds(100)),
)
.await
.unwrap();
assert_gauge(&sampler, &snapshotter, 1.0).await;
spawn_session_cleanup(sessions, sampler, Duration::from_millis(10));
assert_gauge_without_refresh(&snapshotter, 0.0).await;
}
async fn assert_gauge_without_refresh(snapshotter: &Snapshotter, expected: f64) {
for _ in 0..100 {
tokio::time::sleep(Duration::from_millis(10)).await;
if gauge_value(snapshotter) == Some(expected) {
return;
}
}
panic!("active session gauge did not become {expected} through cleanup");
}
async fn assert_gauge(sampler: &ActiveSessionMetrics, snapshotter: &Snapshotter, expected: f64) {
for _ in 0..50 {
sampler.refresh();
tokio::time::sleep(Duration::from_millis(10)).await;
if gauge_value(snapshotter) == Some(expected) {
return;
}
}
panic!("active session gauge did not become {expected}");
}
fn gauge_value(snapshotter: &Snapshotter) -> Option<f64> {
snapshotter
.snapshot()
.into_vec()
.into_iter()
.find_map(|(key, _, _, value)| {
if key.key().name() != "crank_mcp_active_sessions" {
return None;
}
match value {
DebugValue::Gauge(value) => Some(value.into_inner()),
_ => None,
}
})
}
@@ -93,8 +93,12 @@ async fn cleanup_removes_only_expired_sessions() {
.unwrap();
assert_eq!(store.cleanup_expired(now).await.unwrap(), 1);
assert_eq!(store.active_count(now).await.unwrap(), 1);
assert!(store.get(&expired).await.unwrap().is_none());
assert!(store.get(&active).await.unwrap().is_some());
assert!(store.delete(&active).await.unwrap());
assert_eq!(store.active_count(now).await.unwrap(), 0);
}
#[test]