наблюдаемость: завершить базовый контур Community
Добавить структурированные журналы, метрики, трассировку и безопасный канал критических ошибок. Усилить границы рантайма, тесты, проверку зависимостей и сценарии развёртывания.
This commit is contained in:
@@ -3,6 +3,7 @@ name = "crank-community-mcp"
|
||||
edition.workspace = true
|
||||
license.workspace = true
|
||||
rust-version.workspace = true
|
||||
publish.workspace = true
|
||||
version.workspace = true
|
||||
|
||||
[dependencies]
|
||||
@@ -11,10 +12,13 @@ axum.workspace = true
|
||||
base64.workspace = true
|
||||
crank-adapter-rest = { path = "../crank-adapter-rest" }
|
||||
crank-core = { path = "../crank-core" }
|
||||
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
|
||||
@@ -29,3 +33,7 @@ uuid.workspace = true
|
||||
[dev-dependencies]
|
||||
crank-mapping = { path = "../crank-mapping" }
|
||||
crank-test-support = { path = "../crank-test-support" }
|
||||
opentelemetry.workspace = true
|
||||
opentelemetry_sdk.workspace = true
|
||||
tracing-opentelemetry.workspace = true
|
||||
tracing-subscriber.workspace = true
|
||||
|
||||
@@ -1,27 +1,54 @@
|
||||
use std::sync::Arc;
|
||||
|
||||
use axum::http::{HeaderMap, StatusCode, header::AUTHORIZATION};
|
||||
use axum::{
|
||||
http::{HeaderMap, StatusCode, header::AUTHORIZATION},
|
||||
response::{IntoResponse, Response},
|
||||
};
|
||||
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 sha2::{Digest, Sha256};
|
||||
use time::OffsetDateTime;
|
||||
use tracing::Instrument;
|
||||
|
||||
use crate::{
|
||||
app::{AgentRoutePath, AppState},
|
||||
auth::VerifiedMachineCredential,
|
||||
};
|
||||
|
||||
#[derive(Clone, Copy, Debug)]
|
||||
pub(super) enum MachineAccessError {
|
||||
Denied(StatusCode),
|
||||
Unavailable,
|
||||
}
|
||||
|
||||
impl MachineAccessError {
|
||||
pub(super) fn is_denied(self) -> bool {
|
||||
matches!(self, Self::Denied(_))
|
||||
}
|
||||
}
|
||||
|
||||
impl IntoResponse for MachineAccessError {
|
||||
fn into_response(self) -> Response {
|
||||
match self {
|
||||
Self::Denied(status) => status.into_response(),
|
||||
Self::Unavailable => StatusCode::INTERNAL_SERVER_ERROR.into_response(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) async fn require_machine_access(
|
||||
state: &Arc<AppState>,
|
||||
path: &AgentRoutePath,
|
||||
headers: &HeaderMap,
|
||||
required_scope: PlatformApiKeyScope,
|
||||
) -> Result<VerifiedMachineCredential, StatusCode> {
|
||||
let secret = bearer_token(headers).ok_or(StatusCode::UNAUTHORIZED)?;
|
||||
) -> Result<VerifiedMachineCredential, MachineAccessError> {
|
||||
let secret =
|
||||
bearer_token(headers).ok_or(MachineAccessError::Denied(StatusCode::UNAUTHORIZED))?;
|
||||
let credential = resolve_machine_credential(state, path, secret).await?;
|
||||
|
||||
if !allows_scope(&credential.scopes, required_scope) {
|
||||
return Err(StatusCode::FORBIDDEN);
|
||||
return Err(MachineAccessError::Denied(StatusCode::FORBIDDEN));
|
||||
}
|
||||
|
||||
Ok(credential)
|
||||
@@ -35,15 +62,18 @@ pub(super) async fn require_approval_access(
|
||||
) -> Result<crank_registry::PlatformApiKeyRecord, StatusCode> {
|
||||
let secret = bearer_token(headers).ok_or(StatusCode::UNAUTHORIZED)?;
|
||||
let secret_hash = hash_access_secret(secret);
|
||||
let Some(api_key) = state
|
||||
.registry
|
||||
.get_approval_api_key_by_secret_for_agent_slug(
|
||||
&path.workspace_slug,
|
||||
&path.agent_slug,
|
||||
&secret_hash,
|
||||
)
|
||||
.await
|
||||
.map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?
|
||||
let Some(api_key) = observe_db_query(
|
||||
DbOperation::MachineAccessRead,
|
||||
state
|
||||
.registry
|
||||
.get_approval_api_key_by_secret_for_agent_slug(
|
||||
&path.workspace_slug,
|
||||
&path.agent_slug,
|
||||
&secret_hash,
|
||||
),
|
||||
)
|
||||
.await
|
||||
.map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?
|
||||
else {
|
||||
return Err(StatusCode::UNAUTHORIZED);
|
||||
};
|
||||
@@ -53,11 +83,16 @@ pub(super) async fn require_approval_access(
|
||||
}
|
||||
|
||||
let used_at = OffsetDateTime::now_utc();
|
||||
state
|
||||
.registry
|
||||
.touch_platform_api_key(&api_key.api_key.workspace_id, &api_key.api_key.id, &used_at)
|
||||
.await
|
||||
.map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
|
||||
observe_db_query(
|
||||
DbOperation::MachineAccessTouch,
|
||||
state.registry.touch_platform_api_key(
|
||||
&api_key.api_key.workspace_id,
|
||||
&api_key.api_key.id,
|
||||
&used_at,
|
||||
),
|
||||
)
|
||||
.await
|
||||
.map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
|
||||
|
||||
Ok(api_key)
|
||||
}
|
||||
@@ -100,7 +135,7 @@ async fn resolve_machine_credential(
|
||||
state: &Arc<AppState>,
|
||||
path: &AgentRoutePath,
|
||||
token: &str,
|
||||
) -> Result<VerifiedMachineCredential, StatusCode> {
|
||||
) -> Result<VerifiedMachineCredential, MachineAccessError> {
|
||||
if let Some(credential) = verify_static_agent_key(state, path, token).await? {
|
||||
return Ok(credential);
|
||||
}
|
||||
@@ -109,35 +144,68 @@ async fn resolve_machine_credential(
|
||||
.credential_verifier
|
||||
.verify_bearer_token(&path.workspace_slug, &path.agent_slug, token)
|
||||
.await
|
||||
.map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?
|
||||
.ok_or(StatusCode::UNAUTHORIZED)
|
||||
.map_err(|_| MachineAccessError::Unavailable)?
|
||||
.ok_or(MachineAccessError::Denied(StatusCode::UNAUTHORIZED))
|
||||
}
|
||||
|
||||
async fn verify_static_agent_key(
|
||||
state: &Arc<AppState>,
|
||||
path: &AgentRoutePath,
|
||||
secret: &str,
|
||||
) -> Result<Option<VerifiedMachineCredential>, StatusCode> {
|
||||
) -> Result<Option<VerifiedMachineCredential>, MachineAccessError> {
|
||||
let secret_hash = hash_access_secret(secret);
|
||||
let Some(api_key) = state
|
||||
let read_span = Stage::DbQuery
|
||||
.db_span(DbOperation::MachineAccessRead)
|
||||
.expect("database stage");
|
||||
let api_key_result = state
|
||||
.registry
|
||||
.get_platform_api_key_by_secret_for_agent_slug(
|
||||
&path.workspace_slug,
|
||||
&path.agent_slug,
|
||||
&secret_hash,
|
||||
)
|
||||
.instrument(read_span.clone())
|
||||
.await
|
||||
.map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?
|
||||
else {
|
||||
.map_err(|_| MachineAccessError::Unavailable);
|
||||
let api_key = match api_key_result {
|
||||
Ok(api_key) => {
|
||||
StageOutcome::Success.record(&read_span);
|
||||
drop(read_span);
|
||||
api_key
|
||||
}
|
||||
Err(status) => {
|
||||
StageOutcome::Error.record(&read_span);
|
||||
ErrorCategory::Database.record(&read_span);
|
||||
drop(read_span);
|
||||
return Err(status);
|
||||
}
|
||||
};
|
||||
let Some(api_key) = api_key else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
let used_at = OffsetDateTime::now_utc();
|
||||
state
|
||||
let touch_span = Stage::DbQuery
|
||||
.db_span(DbOperation::MachineAccessTouch)
|
||||
.expect("database stage");
|
||||
let touch_result = state
|
||||
.registry
|
||||
.touch_platform_api_key(&api_key.api_key.workspace_id, &api_key.api_key.id, &used_at)
|
||||
.instrument(touch_span.clone())
|
||||
.await
|
||||
.map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
|
||||
.map_err(|_| MachineAccessError::Unavailable);
|
||||
match touch_result {
|
||||
Ok(()) => {
|
||||
StageOutcome::Success.record(&touch_span);
|
||||
drop(touch_span);
|
||||
}
|
||||
Err(status) => {
|
||||
StageOutcome::Error.record(&touch_span);
|
||||
ErrorCategory::Database.record(&touch_span);
|
||||
drop(touch_span);
|
||||
return Err(status);
|
||||
}
|
||||
}
|
||||
|
||||
Ok(Some(VerifiedMachineCredential {
|
||||
machine_access_mode: crank_core::MachineAccessMode::StaticAgentKey,
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,117 @@
|
||||
use std::{sync::Arc, time::Duration};
|
||||
|
||||
use crank_core::{
|
||||
InvocationLevel, InvocationLog, InvocationLogId, InvocationSource, InvocationStatus,
|
||||
};
|
||||
use crank_registry::{
|
||||
CreateInvocationLogRequest, InvocationHistoryWriteOutcome, PublishedAgentTool,
|
||||
};
|
||||
use crank_trace::{DbOperation, ErrorCategory, Stage, StageOutcome};
|
||||
use serde_json::Value;
|
||||
use time::OffsetDateTime;
|
||||
use tracing::{Instrument, warn};
|
||||
|
||||
use super::AppState;
|
||||
|
||||
pub(crate) struct InvocationRecord<'a> {
|
||||
pub(crate) request_id: Option<&'a str>,
|
||||
pub(crate) tool_name: &'a str,
|
||||
pub(crate) status: InvocationStatus,
|
||||
pub(crate) level: InvocationLevel,
|
||||
pub(crate) message: &'a str,
|
||||
pub(crate) status_code: Option<u16>,
|
||||
pub(crate) error_kind: Option<&'a str>,
|
||||
pub(crate) duration: Duration,
|
||||
pub(crate) request_preview: Value,
|
||||
pub(crate) response_preview: Value,
|
||||
}
|
||||
|
||||
pub(crate) async fn persist_invocation(
|
||||
state: &Arc<AppState>,
|
||||
tool: &PublishedAgentTool,
|
||||
record: InvocationRecord<'_>,
|
||||
) -> InvocationHistoryWriteOutcome {
|
||||
let log = InvocationLog {
|
||||
id: InvocationLogId::new(format!("log_{}", uuid::Uuid::now_v7().simple())),
|
||||
workspace_id: tool.workspace_id.clone(),
|
||||
agent_id: Some(tool.agent_id.clone()),
|
||||
operation_id: tool.operation.id.clone(),
|
||||
source: InvocationSource::AgentToolCall,
|
||||
level: record.level,
|
||||
status: record.status,
|
||||
tool_name: record.tool_name.to_owned(),
|
||||
message: record.message.to_owned(),
|
||||
request_id: record.request_id.map(ToOwned::to_owned),
|
||||
status_code: record.status_code,
|
||||
duration_ms: u64::try_from(record.duration.as_millis()).unwrap_or(u64::MAX),
|
||||
error_kind: record.error_kind.map(ToOwned::to_owned),
|
||||
request_preview: record.request_preview,
|
||||
response_preview: record.response_preview,
|
||||
created_at: OffsetDateTime::now_utc(),
|
||||
};
|
||||
|
||||
let history_span = Stage::HistoryWrite.span();
|
||||
let (outcome, db_span) = async {
|
||||
let db_span = Stage::DbQuery
|
||||
.db_span(DbOperation::InvocationHistoryWrite)
|
||||
.expect("database stage");
|
||||
let outcome = state
|
||||
.registry
|
||||
.create_invocation_log(CreateInvocationLogRequest { log: &log })
|
||||
.instrument(db_span.clone())
|
||||
.await;
|
||||
(outcome, db_span)
|
||||
}
|
||||
.instrument(history_span.clone())
|
||||
.await;
|
||||
match outcome {
|
||||
InvocationHistoryWriteOutcome::Recorded => {
|
||||
StageOutcome::Success.record(&db_span);
|
||||
StageOutcome::Success.record(&history_span);
|
||||
}
|
||||
InvocationHistoryWriteOutcome::Lost(_) => {
|
||||
StageOutcome::Error.record(&db_span);
|
||||
ErrorCategory::Database.record(&db_span);
|
||||
StageOutcome::Error.record(&history_span);
|
||||
ErrorCategory::History.record(&history_span);
|
||||
}
|
||||
}
|
||||
drop(db_span);
|
||||
drop(history_span);
|
||||
observe_invocation_history_outcome(
|
||||
outcome,
|
||||
record.request_id,
|
||||
record.status,
|
||||
"agent_tool_call",
|
||||
);
|
||||
outcome
|
||||
}
|
||||
|
||||
pub(super) fn observe_invocation_history_outcome(
|
||||
outcome: InvocationHistoryWriteOutcome,
|
||||
request_id: Option<&str>,
|
||||
status: InvocationStatus,
|
||||
source: &'static str,
|
||||
) {
|
||||
let Some(loss) = outcome.loss() else {
|
||||
return;
|
||||
};
|
||||
crank_observability::record_operational_incident(
|
||||
crank_observability::OperationalIncident::InvocationHistoryLost,
|
||||
);
|
||||
warn!(
|
||||
name: "mcp.invocation_history.lost",
|
||||
request_id = request_id.unwrap_or_default(),
|
||||
source,
|
||||
invocation_status = invocation_status_label(status),
|
||||
error_category = loss.category.as_str(),
|
||||
"invocation history was not recorded"
|
||||
);
|
||||
}
|
||||
|
||||
fn invocation_status_label(status: InvocationStatus) -> &'static str {
|
||||
match status {
|
||||
InvocationStatus::Ok => "ok",
|
||||
InvocationStatus::Error => "error",
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,93 @@
|
||||
use std::sync::Arc;
|
||||
|
||||
use axum::http::StatusCode;
|
||||
use serde_json::Value;
|
||||
use tokio::sync::{OwnedSemaphorePermit, Semaphore};
|
||||
|
||||
use crate::{
|
||||
jsonrpc::{is_notification, is_response, method_name},
|
||||
transport::ResponseMode,
|
||||
};
|
||||
|
||||
pub(super) struct McpRequestMetrics {
|
||||
method: &'static str,
|
||||
response_mode: &'static str,
|
||||
outcome: &'static str,
|
||||
}
|
||||
|
||||
impl McpRequestMetrics {
|
||||
pub(super) fn new(message: &Value) -> Self {
|
||||
Self {
|
||||
method: normalized_mcp_method(message),
|
||||
response_mode: "unknown",
|
||||
outcome: "rejected",
|
||||
}
|
||||
}
|
||||
|
||||
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 complete(&mut self, status: StatusCode) {
|
||||
self.outcome = match status.as_u16() {
|
||||
200..=299 => "success",
|
||||
400..=499 => "client_error",
|
||||
500..=599 => "server_error",
|
||||
_ => "other",
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
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);
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn normalized_mcp_method(message: &Value) -> &'static str {
|
||||
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",
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) struct ActiveSessionGuard {
|
||||
_permit: OwnedSemaphorePermit,
|
||||
}
|
||||
|
||||
impl ActiveSessionGuard {
|
||||
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);
|
||||
})?;
|
||||
::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);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,86 @@
|
||||
use std::sync::Arc;
|
||||
|
||||
use axum::http::{HeaderMap, StatusCode};
|
||||
use crank_core::PlatformApiKeyScope;
|
||||
use crank_registry::PlatformApiKeyRecord;
|
||||
use crank_runtime::RateLimitCheckError;
|
||||
use crank_trace::{ErrorCategory, Stage, StageOutcome};
|
||||
use tracing::Instrument;
|
||||
|
||||
use crate::{
|
||||
access::{MachineAccessError, require_approval_access, require_machine_access},
|
||||
app::{AgentRoutePath, AppState},
|
||||
auth::VerifiedMachineCredential,
|
||||
rate_limit::enforce_transport_rate_limit,
|
||||
};
|
||||
|
||||
pub(super) async fn enforce_traced_rate_limit(
|
||||
state: &Arc<AppState>,
|
||||
path: &AgentRoutePath,
|
||||
headers: &HeaderMap,
|
||||
) -> Result<(), RateLimitCheckError> {
|
||||
let span = Stage::McpRateLimit.span();
|
||||
let result = enforce_transport_rate_limit(state, path, headers)
|
||||
.instrument(span.clone())
|
||||
.await;
|
||||
match &result {
|
||||
Ok(()) => StageOutcome::Allowed.record(&span),
|
||||
Err(RateLimitCheckError::Rejected(_)) => {
|
||||
StageOutcome::Denied.record(&span);
|
||||
ErrorCategory::RateLimit.record(&span);
|
||||
}
|
||||
Err(RateLimitCheckError::StoreUnavailable) => {
|
||||
StageOutcome::Error.record(&span);
|
||||
ErrorCategory::Internal.record(&span);
|
||||
}
|
||||
}
|
||||
result
|
||||
}
|
||||
|
||||
pub(super) async fn require_traced_machine_access(
|
||||
state: &Arc<AppState>,
|
||||
path: &AgentRoutePath,
|
||||
headers: &HeaderMap,
|
||||
required_scope: PlatformApiKeyScope,
|
||||
) -> Result<VerifiedMachineCredential, MachineAccessError> {
|
||||
let span = Stage::McpAccessCheck.span();
|
||||
let result = require_machine_access(state, path, headers, required_scope)
|
||||
.instrument(span.clone())
|
||||
.await;
|
||||
match &result {
|
||||
Ok(_) => StageOutcome::Allowed.record(&span),
|
||||
Err(error) if error.is_denied() => {
|
||||
StageOutcome::Denied.record(&span);
|
||||
ErrorCategory::Access.record(&span);
|
||||
}
|
||||
Err(_) => {
|
||||
StageOutcome::Error.record(&span);
|
||||
ErrorCategory::Internal.record(&span);
|
||||
}
|
||||
}
|
||||
result
|
||||
}
|
||||
|
||||
pub(super) async fn require_traced_approval_access(
|
||||
state: &Arc<AppState>,
|
||||
path: &AgentRoutePath,
|
||||
headers: &HeaderMap,
|
||||
required_scope: PlatformApiKeyScope,
|
||||
) -> Result<PlatformApiKeyRecord, StatusCode> {
|
||||
let span = Stage::McpAccessCheck.span();
|
||||
let result = require_approval_access(state, path, headers, required_scope)
|
||||
.instrument(span.clone())
|
||||
.await;
|
||||
match &result {
|
||||
Ok(_) => StageOutcome::Allowed.record(&span),
|
||||
Err(status) if *status == StatusCode::UNAUTHORIZED || *status == StatusCode::FORBIDDEN => {
|
||||
StageOutcome::Denied.record(&span);
|
||||
ErrorCategory::Access.record(&span);
|
||||
}
|
||||
Err(_) => {
|
||||
StageOutcome::Error.record(&span);
|
||||
ErrorCategory::Internal.record(&span);
|
||||
}
|
||||
}
|
||||
result
|
||||
}
|
||||
@@ -0,0 +1,149 @@
|
||||
use std::{
|
||||
io,
|
||||
sync::{Arc, Mutex},
|
||||
};
|
||||
|
||||
use axum::body::to_bytes;
|
||||
use crank_core::InvocationStatus;
|
||||
use crank_observability::{
|
||||
ObservabilityConfig, OperationalIncident, RedactionLimits, ServiceIdentity,
|
||||
operational_incident_total,
|
||||
};
|
||||
use crank_registry::{
|
||||
InvocationHistoryLoss, InvocationHistoryLossCategory, InvocationHistoryWriteOutcome,
|
||||
};
|
||||
use serde_json::{Value, json};
|
||||
use tracing_subscriber::fmt::MakeWriter;
|
||||
|
||||
use super::{
|
||||
ResponseMode, metrics::normalized_mcp_method, observe_invocation_history_outcome,
|
||||
tool_error_response,
|
||||
};
|
||||
use crate::jsonrpc::CURRENT_PROTOCOL_VERSION;
|
||||
use crate::tool_error::generic_tool_error_contract;
|
||||
|
||||
#[tokio::test]
|
||||
async fn tool_error_response_includes_structured_context() {
|
||||
let response = tool_error_response(
|
||||
&json!({"jsonrpc": "2.0", "id": "req-1"}),
|
||||
ResponseMode::Json,
|
||||
CURRENT_PROTOCOL_VERSION,
|
||||
generic_tool_error_contract(
|
||||
"streaming_payload_error",
|
||||
"request root must be an object",
|
||||
"req-1",
|
||||
false,
|
||||
Some("Проверьте параметры вызова инструмента."),
|
||||
),
|
||||
);
|
||||
|
||||
let body = to_bytes(response.into_body(), usize::MAX).await.unwrap();
|
||||
let payload: Value = serde_json::from_slice(&body).unwrap();
|
||||
|
||||
assert_eq!(
|
||||
payload["result"]["structuredContent"]["error"],
|
||||
json!({
|
||||
"code": "streaming_payload_error",
|
||||
"error_code": "streaming_payload_error",
|
||||
"message": "request root must be an object",
|
||||
"recoverable": false,
|
||||
"request_id": "req-1",
|
||||
"suggested_action": "Проверьте параметры вызова инструмента."
|
||||
})
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn emits_bounded_history_loss_incident() {
|
||||
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();
|
||||
let before = operational_incident_total(OperationalIncident::InvocationHistoryLost);
|
||||
let dispatch = tracing::Dispatch::new(subscriber);
|
||||
let _guard = tracing::dispatcher::set_default(&dispatch);
|
||||
|
||||
observe_invocation_history_outcome(
|
||||
InvocationHistoryWriteOutcome::Lost(InvocationHistoryLoss {
|
||||
category: InvocationHistoryLossCategory::Unavailable,
|
||||
}),
|
||||
Some("req_mcp_dc08"),
|
||||
InvocationStatus::Ok,
|
||||
"agent_tool_call",
|
||||
);
|
||||
|
||||
let output = writer.output();
|
||||
assert!(!output.contains("dc08-canary-secret"));
|
||||
let event: Value = output
|
||||
.lines()
|
||||
.map(|line| serde_json::from_str(line).unwrap())
|
||||
.find(|event: &Value| event["event"] == "mcp.invocation_history.lost")
|
||||
.unwrap();
|
||||
assert_eq!(event["request_id"], "req_mcp_dc08");
|
||||
assert_eq!(event["fields"]["source"], "agent_tool_call");
|
||||
assert_eq!(event["fields"]["invocation_status"], "ok");
|
||||
assert_eq!(event["fields"]["error_category"], "unavailable");
|
||||
assert!(operational_incident_total(OperationalIncident::InvocationHistoryLost) > before);
|
||||
}
|
||||
|
||||
#[test]
|
||||
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"
|
||||
);
|
||||
assert_eq!(
|
||||
normalized_mcp_method(
|
||||
&json!({"jsonrpc": "2.0", "id": 2, "method": "customer-controlled-method"})
|
||||
),
|
||||
"unsupported"
|
||||
);
|
||||
assert_eq!(
|
||||
normalized_mcp_method(
|
||||
&json!({"jsonrpc": "2.0", "method": "customer-controlled-notification"})
|
||||
),
|
||||
"notification"
|
||||
);
|
||||
}
|
||||
|
||||
#[derive(Clone, Default)]
|
||||
struct SharedLogWriter {
|
||||
buffer: Arc<Mutex<Vec<u8>>>,
|
||||
}
|
||||
|
||||
impl SharedLogWriter {
|
||||
fn output(&self) -> String {
|
||||
String::from_utf8(self.buffer.lock().unwrap().clone()).unwrap()
|
||||
}
|
||||
}
|
||||
|
||||
impl<'a> MakeWriter<'a> for SharedLogWriter {
|
||||
type Writer = SharedLogGuard;
|
||||
|
||||
fn make_writer(&'a self) -> Self::Writer {
|
||||
SharedLogGuard {
|
||||
buffer: Arc::clone(&self.buffer),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
struct SharedLogGuard {
|
||||
buffer: Arc<Mutex<Vec<u8>>>,
|
||||
}
|
||||
|
||||
impl io::Write for SharedLogGuard {
|
||||
fn write(&mut self, bytes: &[u8]) -> io::Result<usize> {
|
||||
self.buffer.lock().unwrap().extend_from_slice(bytes);
|
||||
Ok(bytes.len())
|
||||
}
|
||||
|
||||
fn flush(&mut self) -> io::Result<()> {
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
@@ -5,11 +5,13 @@ use axum::{
|
||||
response::{IntoResponse, Response},
|
||||
};
|
||||
use crank_core::{ApprovalRequestStatus, InvocationLevel, InvocationSource, InvocationStatus};
|
||||
use crank_observability::RequestId;
|
||||
use crank_registry::{ApprovalRequestRecord, FinishApprovalRequest};
|
||||
use crank_runtime::{RuntimeExecutionRequest, RuntimeRequestContext};
|
||||
use crank_trace::{DbOperation, ErrorCategory, Stage, StageOutcome, observe_db_query};
|
||||
use serde_json::json;
|
||||
use time::OffsetDateTime;
|
||||
use tracing::warn;
|
||||
use tracing::{Instrument, warn};
|
||||
|
||||
use crate::{
|
||||
app::{
|
||||
@@ -35,32 +37,81 @@ pub(super) fn spawn_approval_recovery(state: Arc<AppState>) {
|
||||
}
|
||||
|
||||
async fn recover_approved_requests(state: &Arc<AppState>) {
|
||||
fail_interrupted_requests(state).await;
|
||||
|
||||
for _ in 0..32 {
|
||||
let now = OffsetDateTime::now_utc();
|
||||
let approval = match state
|
||||
.registry
|
||||
.claim_next_recoverable_approval_request(
|
||||
now,
|
||||
now - RECOVERY_GRACE,
|
||||
now - EXECUTION_LEASE,
|
||||
)
|
||||
.await
|
||||
let approval = match observe_db_query(
|
||||
DbOperation::ApprovalWrite,
|
||||
state
|
||||
.registry
|
||||
.claim_next_recoverable_approval_request(now, now - RECOVERY_GRACE),
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(Some(approval)) => approval,
|
||||
Ok(None) => break,
|
||||
Err(error) => {
|
||||
warn!(error = %error, "approval recovery query failed");
|
||||
Err(_) => {
|
||||
warn!(
|
||||
name: "mcp.approval_recovery.query_failed",
|
||||
error_category = "registry",
|
||||
"approval recovery query failed"
|
||||
);
|
||||
break;
|
||||
}
|
||||
};
|
||||
let Some(path) = approval_agent_path(state, &approval).await else {
|
||||
continue;
|
||||
};
|
||||
if execute_approved_request(state, &path, approval)
|
||||
.await
|
||||
.is_err()
|
||||
let recovery_span = Stage::ApprovalRecovery.span();
|
||||
let result = execute_approved_request(state, &path, approval, None)
|
||||
.instrument(recovery_span.clone())
|
||||
.await;
|
||||
match &result {
|
||||
Ok(_) => StageOutcome::Success.record(&recovery_span),
|
||||
Err(_) => {
|
||||
StageOutcome::Error.record(&recovery_span);
|
||||
ErrorCategory::Approval.record(&recovery_span);
|
||||
}
|
||||
}
|
||||
drop(recovery_span);
|
||||
if result.is_err() {
|
||||
warn!(
|
||||
name: "mcp.approval_recovery.execution_failed",
|
||||
error_category = "runtime",
|
||||
"recovered approval execution did not finish"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn fail_interrupted_requests(state: &Arc<AppState>) {
|
||||
for _ in 0..32 {
|
||||
let stale_before = OffsetDateTime::now_utc() - EXECUTION_LEASE;
|
||||
match observe_db_query(
|
||||
DbOperation::ApprovalWrite,
|
||||
state
|
||||
.registry
|
||||
.fail_next_interrupted_approval_request(stale_before),
|
||||
)
|
||||
.await
|
||||
{
|
||||
warn!("recovered approval execution did not finish");
|
||||
Ok(Some(approval)) => {
|
||||
warn!(
|
||||
name: "mcp.approval_recovery.interrupted",
|
||||
approval_id = approval.approval.id.as_str(),
|
||||
"interrupted approval execution was not retried because its outcome is unknown"
|
||||
);
|
||||
}
|
||||
Ok(None) => break,
|
||||
Err(_) => {
|
||||
warn!(
|
||||
name: "mcp.approval_recovery.interrupted_query_failed",
|
||||
error_category = "registry",
|
||||
"interrupted approval recovery query failed"
|
||||
);
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -76,8 +127,12 @@ async fn approval_agent_path(
|
||||
{
|
||||
Ok(Some(workspace)) => workspace,
|
||||
Ok(None) => return None,
|
||||
Err(error) => {
|
||||
warn!(error = %error, "approval workspace lookup failed");
|
||||
Err(_) => {
|
||||
warn!(
|
||||
name: "mcp.approval_recovery.workspace_lookup_failed",
|
||||
error_category = "registry",
|
||||
"approval workspace lookup failed"
|
||||
);
|
||||
return None;
|
||||
}
|
||||
};
|
||||
@@ -88,8 +143,12 @@ async fn approval_agent_path(
|
||||
{
|
||||
Ok(Some(agent)) => agent,
|
||||
Ok(None) => return None,
|
||||
Err(error) => {
|
||||
warn!(error = %error, "approval agent lookup failed");
|
||||
Err(_) => {
|
||||
warn!(
|
||||
name: "mcp.approval_recovery.agent_lookup_failed",
|
||||
error_category = "registry",
|
||||
"approval agent lookup failed"
|
||||
);
|
||||
return None;
|
||||
}
|
||||
};
|
||||
@@ -103,7 +162,9 @@ pub(super) async fn execute_approved_request(
|
||||
state: &Arc<AppState>,
|
||||
path: &AgentRoutePath,
|
||||
approval: ApprovalRequestRecord,
|
||||
request_id: Option<&str>,
|
||||
) -> Result<ApprovalRequestRecord, Response> {
|
||||
let request_id = RequestId::resolve(request_id).into_string();
|
||||
let tools = state
|
||||
.catalog
|
||||
.list_tools(&path.workspace_slug, &path.agent_slug)
|
||||
@@ -123,18 +184,17 @@ pub(super) async fn execute_approved_request(
|
||||
&approval.approval.request_payload,
|
||||
);
|
||||
let started_at = Instant::now();
|
||||
let runtime_request_context =
|
||||
RuntimeRequestContext::from_request_id(approval.approval.id.as_str().to_owned())
|
||||
.with_response_cache_scope(
|
||||
tool.workspace_id.as_str().to_owned(),
|
||||
tool.agent_id.as_str().to_owned(),
|
||||
)
|
||||
.with_metering_context(
|
||||
tool.workspace_id.clone(),
|
||||
Some(tool.agent_id.clone()),
|
||||
InvocationSource::AgentToolCall,
|
||||
)
|
||||
.with_approval_granted();
|
||||
let runtime_request_context = RuntimeRequestContext::from_request_id(request_id.clone())
|
||||
.with_response_cache_scope(
|
||||
tool.workspace_id.as_str().to_owned(),
|
||||
tool.agent_id.as_str().to_owned(),
|
||||
)
|
||||
.with_metering_context(
|
||||
tool.workspace_id.clone(),
|
||||
Some(tool.agent_id.clone()),
|
||||
InvocationSource::AgentToolCall,
|
||||
)
|
||||
.with_approval_granted();
|
||||
let resolved_auth =
|
||||
resolve_operation_auth(state, &tool.workspace_id, &operation.execution_config).await;
|
||||
let result = match resolved_auth {
|
||||
@@ -176,11 +236,11 @@ pub(super) async fn execute_approved_request(
|
||||
),
|
||||
};
|
||||
|
||||
if let Err(error) = persist_invocation(
|
||||
persist_invocation(
|
||||
state,
|
||||
&tool,
|
||||
InvocationRecord {
|
||||
request_id: Some(approval.approval.id.as_str()),
|
||||
request_id: Some(&request_id),
|
||||
tool_name: &tool.tool_name,
|
||||
status: invocation_status,
|
||||
level: invocation_level,
|
||||
@@ -192,46 +252,49 @@ pub(super) async fn execute_approved_request(
|
||||
response_preview: response_payload.clone(),
|
||||
},
|
||||
)
|
||||
.await
|
||||
{
|
||||
warn!(error = %error, "approved invocation log write failed");
|
||||
}
|
||||
.await;
|
||||
|
||||
state
|
||||
.registry
|
||||
.finish_approval_request(FinishApprovalRequest {
|
||||
workspace_id: &approval.approval.workspace_id,
|
||||
agent_id: &approval.approval.agent_id,
|
||||
approval_id: &approval.approval.id,
|
||||
status,
|
||||
response_payload: Some(response_payload),
|
||||
decision_note: None,
|
||||
})
|
||||
.await
|
||||
.map_err(|_| StatusCode::INTERNAL_SERVER_ERROR.into_response())?
|
||||
.ok_or_else(|| StatusCode::CONFLICT.into_response())
|
||||
observe_db_query(
|
||||
DbOperation::ApprovalWrite,
|
||||
state
|
||||
.registry
|
||||
.finish_approval_request(FinishApprovalRequest {
|
||||
workspace_id: &approval.approval.workspace_id,
|
||||
agent_id: &approval.approval.agent_id,
|
||||
approval_id: &approval.approval.id,
|
||||
status,
|
||||
response_payload: Some(response_payload),
|
||||
decision_note: None,
|
||||
}),
|
||||
)
|
||||
.await
|
||||
.map_err(|_| StatusCode::INTERNAL_SERVER_ERROR.into_response())?
|
||||
.ok_or_else(|| StatusCode::CONFLICT.into_response())
|
||||
}
|
||||
|
||||
async fn finish_unavailable_approval(
|
||||
state: &Arc<AppState>,
|
||||
approval: &ApprovalRequestRecord,
|
||||
) -> Result<ApprovalRequestRecord, Response> {
|
||||
state
|
||||
.registry
|
||||
.finish_approval_request(FinishApprovalRequest {
|
||||
workspace_id: &approval.approval.workspace_id,
|
||||
agent_id: &approval.approval.agent_id,
|
||||
approval_id: &approval.approval.id,
|
||||
status: ApprovalRequestStatus::Failed,
|
||||
response_payload: Some(json!({
|
||||
"error": {
|
||||
"code": "approved_operation_unavailable",
|
||||
"message": "the approved operation version is no longer published"
|
||||
}
|
||||
})),
|
||||
decision_note: None,
|
||||
})
|
||||
.await
|
||||
.map_err(|_| StatusCode::INTERNAL_SERVER_ERROR.into_response())?
|
||||
.ok_or_else(|| StatusCode::CONFLICT.into_response())
|
||||
observe_db_query(
|
||||
DbOperation::ApprovalWrite,
|
||||
state
|
||||
.registry
|
||||
.finish_approval_request(FinishApprovalRequest {
|
||||
workspace_id: &approval.approval.workspace_id,
|
||||
agent_id: &approval.approval.agent_id,
|
||||
approval_id: &approval.approval.id,
|
||||
status: ApprovalRequestStatus::Failed,
|
||||
response_payload: Some(json!({
|
||||
"error": {
|
||||
"code": "approved_operation_unavailable",
|
||||
"message": "the approved operation version is no longer published"
|
||||
}
|
||||
})),
|
||||
decision_note: None,
|
||||
}),
|
||||
)
|
||||
.await
|
||||
.map_err(|_| StatusCode::INTERNAL_SERVER_ERROR.into_response())?
|
||||
.ok_or_else(|| StatusCode::CONFLICT.into_response())
|
||||
}
|
||||
|
||||
@@ -0,0 +1,38 @@
|
||||
use crank_core::{ApprovalRequest, OperationApprovalPolicy};
|
||||
use crank_registry::PublishedAgentTool;
|
||||
use serde_json::{Value, json};
|
||||
|
||||
pub(super) fn approval_required_response(
|
||||
tool: &PublishedAgentTool,
|
||||
approval: &ApprovalRequest,
|
||||
policy: &OperationApprovalPolicy,
|
||||
) -> Value {
|
||||
let approval_url = format!(
|
||||
"/v1/{}/{}/approvals/{}",
|
||||
tool.workspace_slug,
|
||||
tool.agent_slug,
|
||||
approval.id.as_str()
|
||||
);
|
||||
json!({
|
||||
"status": "approval_required",
|
||||
"approval_id": approval.id.as_str(),
|
||||
"approval_url": approval_url,
|
||||
"approve": {
|
||||
"method": "POST",
|
||||
"url": format!("{approval_url}/approve"),
|
||||
"body": { "approve": "yes" }
|
||||
},
|
||||
"deny": {
|
||||
"method": "POST",
|
||||
"url": format!("{approval_url}/deny"),
|
||||
"body": { "approve": "no" }
|
||||
},
|
||||
"expires_at": approval.expires_at,
|
||||
"risk_level": approval.risk_level,
|
||||
"payload_preview": if policy.show_payload_preview {
|
||||
approval.request_payload.clone()
|
||||
} else {
|
||||
Value::Null
|
||||
},
|
||||
})
|
||||
}
|
||||
@@ -1,24 +1,27 @@
|
||||
use std::{
|
||||
collections::HashMap,
|
||||
sync::Arc,
|
||||
sync::{Arc, Weak},
|
||||
time::{Duration, Instant, SystemTime, UNIX_EPOCH},
|
||||
};
|
||||
|
||||
use crank_core::{CacheScope, CoordinationStateStore, CoordinationStateValue};
|
||||
use crank_registry::{PostgresRegistry, PublishedAgentCatalog, PublishedAgentTool, RegistryError};
|
||||
use crank_trace::{DbOperation, ErrorCategory, Stage, StageOutcome};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use tokio::sync::{Mutex, RwLock};
|
||||
use tracing::{info, warn};
|
||||
use tracing::{Instrument, info, warn};
|
||||
|
||||
use crate::manifest::analyze_published_tool_catalog;
|
||||
|
||||
const MAX_LOCAL_CATALOGS: usize = 1_024;
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct PublishedToolCatalog {
|
||||
registry: PostgresRegistry,
|
||||
refresh_interval: Duration,
|
||||
coordination_store: Arc<dyn CoordinationStateStore>,
|
||||
cached: Arc<RwLock<HashMap<CatalogKey, CachedCatalog>>>,
|
||||
refresh_locks: Arc<Mutex<HashMap<CatalogKey, Arc<Mutex<()>>>>>,
|
||||
refresh_locks: Arc<Mutex<HashMap<CatalogKey, Weak<Mutex<()>>>>>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
|
||||
@@ -30,6 +33,14 @@ struct CatalogKey {
|
||||
struct CachedCatalog {
|
||||
loaded_at: Option<Instant>,
|
||||
catalog: PublishedAgentCatalog,
|
||||
metrics: CatalogMetrics,
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, Default)]
|
||||
struct CatalogMetrics {
|
||||
tool_count: usize,
|
||||
estimated_context_tokens: usize,
|
||||
warning_count: usize,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
|
||||
@@ -66,15 +77,29 @@ impl PublishedToolCatalog {
|
||||
workspace_slug: &str,
|
||||
agent_slug: &str,
|
||||
) -> Result<PublishedAgentCatalog, RegistryError> {
|
||||
self.refresh_if_stale(workspace_slug, agent_slug).await?;
|
||||
let guard = self.cached.read().await;
|
||||
guard
|
||||
.get(&CatalogKey::new(workspace_slug, agent_slug))
|
||||
.map(|entry| entry.catalog.clone())
|
||||
.ok_or_else(|| RegistryError::PublishedAgentNotFound {
|
||||
workspace_slug: workspace_slug.to_owned(),
|
||||
agent_slug: agent_slug.to_owned(),
|
||||
})
|
||||
let span = Stage::McpCatalogLoad.span();
|
||||
let result = async {
|
||||
self.refresh_if_stale(workspace_slug, agent_slug).await?;
|
||||
let guard = self.cached.read().await;
|
||||
guard
|
||||
.get(&CatalogKey::new(workspace_slug, agent_slug))
|
||||
.map(|entry| entry.catalog.clone())
|
||||
.ok_or_else(|| RegistryError::PublishedAgentNotFound {
|
||||
workspace_slug: workspace_slug.to_owned(),
|
||||
agent_slug: agent_slug.to_owned(),
|
||||
})
|
||||
}
|
||||
.instrument(span.clone())
|
||||
.await;
|
||||
match &result {
|
||||
Ok(_) => StageOutcome::Success.record(&span),
|
||||
Err(_) => {
|
||||
StageOutcome::Error.record(&span);
|
||||
ErrorCategory::Catalog.record(&span);
|
||||
}
|
||||
}
|
||||
drop(span);
|
||||
result
|
||||
}
|
||||
|
||||
async fn refresh_if_stale(
|
||||
@@ -98,11 +123,14 @@ impl PublishedToolCatalog {
|
||||
|
||||
let refresh_lock = {
|
||||
let mut locks = self.refresh_locks.lock().await;
|
||||
Arc::clone(
|
||||
locks
|
||||
.entry(key.clone())
|
||||
.or_insert_with(|| Arc::new(Mutex::new(()))),
|
||||
)
|
||||
locks.retain(|_, lock| lock.strong_count() > 0);
|
||||
if let Some(lock) = locks.get(&key).and_then(Weak::upgrade) {
|
||||
lock
|
||||
} else {
|
||||
let lock = Arc::new(Mutex::new(()));
|
||||
locks.insert(key.clone(), Arc::downgrade(&lock));
|
||||
lock
|
||||
}
|
||||
};
|
||||
let _refresh_guard = refresh_lock.lock().await;
|
||||
let still_stale = {
|
||||
@@ -117,50 +145,59 @@ impl PublishedToolCatalog {
|
||||
}
|
||||
|
||||
if let Some((catalog, age)) = self.load_shared_snapshot(workspace_slug, agent_slug).await {
|
||||
log_catalog_analysis(workspace_slug, agent_slug, "shared_cache", &catalog.tools);
|
||||
let mut guard = self.cached.write().await;
|
||||
guard.insert(
|
||||
let metrics =
|
||||
log_catalog_analysis(workspace_slug, agent_slug, "shared_cache", &catalog.tools);
|
||||
self.store_local_catalog(
|
||||
key,
|
||||
CachedCatalog {
|
||||
loaded_at: Instant::now().checked_sub(age),
|
||||
catalog,
|
||||
metrics,
|
||||
},
|
||||
);
|
||||
)
|
||||
.await;
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
let catalog = match self
|
||||
let db_span = Stage::DbQuery
|
||||
.db_span(DbOperation::CatalogLoad)
|
||||
.expect("database stage");
|
||||
let catalog_result = self
|
||||
.registry
|
||||
.get_published_agent_catalog_by_slug(workspace_slug, agent_slug)
|
||||
.await
|
||||
{
|
||||
.instrument(db_span.clone())
|
||||
.await;
|
||||
let catalog = match catalog_result {
|
||||
Ok(catalog) => catalog,
|
||||
Err(error) => return Err(error),
|
||||
Err(error) => {
|
||||
StageOutcome::Error.record(&db_span);
|
||||
ErrorCategory::Database.record(&db_span);
|
||||
drop(db_span);
|
||||
return Err(error);
|
||||
}
|
||||
};
|
||||
log_catalog_analysis(workspace_slug, agent_slug, "postgres", &catalog.tools);
|
||||
StageOutcome::Success.record(&db_span);
|
||||
drop(db_span);
|
||||
let metrics = log_catalog_analysis(workspace_slug, agent_slug, "postgres", &catalog.tools);
|
||||
self.store_shared_snapshot(workspace_slug, agent_slug, &catalog)
|
||||
.await;
|
||||
let mut guard = self.cached.write().await;
|
||||
let previous_count = guard
|
||||
.get(&key)
|
||||
.map(|entry| entry.catalog.tools.len())
|
||||
.unwrap_or_default();
|
||||
|
||||
guard.insert(
|
||||
key,
|
||||
CachedCatalog {
|
||||
loaded_at: Some(Instant::now()),
|
||||
catalog,
|
||||
},
|
||||
);
|
||||
let published_tool_count = catalog.tools.len();
|
||||
let previous_count = self
|
||||
.store_local_catalog(
|
||||
key,
|
||||
CachedCatalog {
|
||||
loaded_at: Some(Instant::now()),
|
||||
catalog,
|
||||
metrics,
|
||||
},
|
||||
)
|
||||
.await;
|
||||
|
||||
info!(
|
||||
name: "mcp.catalog.refreshed",
|
||||
workspace_slug,
|
||||
agent_slug,
|
||||
published_tool_count = guard
|
||||
.get(&CatalogKey::new(workspace_slug, agent_slug))
|
||||
.map(|entry| entry.catalog.tools.len())
|
||||
.unwrap_or_default(),
|
||||
published_tool_count,
|
||||
previous_published_tool_count = previous_count,
|
||||
"published agent catalog refreshed"
|
||||
);
|
||||
@@ -168,6 +205,27 @@ impl PublishedToolCatalog {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn store_local_catalog(&self, key: CatalogKey, entry: CachedCatalog) -> usize {
|
||||
let mut guard = self.cached.write().await;
|
||||
let previous_count = guard
|
||||
.get(&key)
|
||||
.map(|current| current.catalog.tools.len())
|
||||
.unwrap_or_default();
|
||||
|
||||
if guard.len() >= MAX_LOCAL_CATALOGS && !guard.contains_key(&key) {
|
||||
let oldest = guard
|
||||
.iter()
|
||||
.min_by_key(|(_, current)| current.loaded_at)
|
||||
.map(|(candidate, _)| candidate.clone());
|
||||
if let Some(oldest) = oldest {
|
||||
guard.remove(&oldest);
|
||||
}
|
||||
}
|
||||
guard.insert(key, entry);
|
||||
record_catalog_metrics(guard.values().map(|entry| entry.metrics));
|
||||
previous_count
|
||||
}
|
||||
|
||||
async fn load_shared_snapshot(
|
||||
&self,
|
||||
workspace_slug: &str,
|
||||
@@ -226,12 +284,19 @@ fn log_catalog_analysis(
|
||||
agent_slug: &str,
|
||||
source: &str,
|
||||
tools: &[PublishedAgentTool],
|
||||
) {
|
||||
) -> CatalogMetrics {
|
||||
let analysis = match analyze_published_tool_catalog(tools) {
|
||||
Ok(analysis) => analysis,
|
||||
Err(error) => {
|
||||
warn!(workspace_slug, agent_slug, source, %error, "published catalog analysis failed");
|
||||
return;
|
||||
Err(_) => {
|
||||
warn!(
|
||||
name: "mcp.catalog.analysis_failed",
|
||||
workspace_slug,
|
||||
agent_slug,
|
||||
source,
|
||||
error_category = "catalog_validation",
|
||||
"published catalog analysis failed"
|
||||
);
|
||||
return CatalogMetrics::default();
|
||||
}
|
||||
};
|
||||
let warning_count = analysis
|
||||
@@ -242,6 +307,7 @@ fn log_catalog_analysis(
|
||||
.count();
|
||||
|
||||
info!(
|
||||
name: "mcp.catalog.analyzed",
|
||||
workspace_slug,
|
||||
agent_slug,
|
||||
source,
|
||||
@@ -255,6 +321,30 @@ fn log_catalog_analysis(
|
||||
catalog_quality_warning_count = warning_count,
|
||||
"published agent catalog analyzed"
|
||||
);
|
||||
|
||||
CatalogMetrics {
|
||||
tool_count: analysis.budget.tool_count,
|
||||
estimated_context_tokens: analysis.budget.estimated_context_tokens,
|
||||
warning_count,
|
||||
}
|
||||
}
|
||||
|
||||
fn record_catalog_metrics(metrics: impl Iterator<Item = CatalogMetrics>) {
|
||||
let aggregate = metrics.fold(CatalogMetrics::default(), |mut aggregate, current| {
|
||||
aggregate.tool_count = aggregate.tool_count.saturating_add(current.tool_count);
|
||||
aggregate.estimated_context_tokens = aggregate
|
||||
.estimated_context_tokens
|
||||
.saturating_add(current.estimated_context_tokens);
|
||||
aggregate.warning_count = aggregate
|
||||
.warning_count
|
||||
.saturating_add(current.warning_count);
|
||||
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);
|
||||
}
|
||||
|
||||
fn now_unix_ms() -> u64 {
|
||||
|
||||
@@ -1,14 +1,18 @@
|
||||
mod access;
|
||||
mod app;
|
||||
mod approval_execution;
|
||||
mod approval_response;
|
||||
pub mod auth;
|
||||
pub mod catalog;
|
||||
pub mod jsonrpc;
|
||||
pub mod manifest;
|
||||
mod rate_limit;
|
||||
mod request_context;
|
||||
pub mod session;
|
||||
pub mod tool_error;
|
||||
mod tool_search;
|
||||
mod transport;
|
||||
|
||||
pub use app::{build_app, build_app_with_background_workers};
|
||||
pub use app::{
|
||||
build_app, build_app_with_background_workers, build_app_with_background_workers_and_limits,
|
||||
};
|
||||
|
||||
@@ -4,7 +4,7 @@ use axum::{
|
||||
http::{HeaderMap, HeaderValue, StatusCode, header::RETRY_AFTER},
|
||||
response::{IntoResponse, Response},
|
||||
};
|
||||
use crank_runtime::RateLimitRejection;
|
||||
use crank_runtime::RateLimitCheckError;
|
||||
use serde_json::{Value, json};
|
||||
|
||||
use crate::{
|
||||
@@ -14,19 +14,11 @@ use crate::{
|
||||
transport::{ResponseMode, session_id_from_headers, transport_response},
|
||||
};
|
||||
|
||||
pub(super) async fn enforce_post_rate_limit(
|
||||
state: &Arc<AppState>,
|
||||
path: &AgentRoutePath,
|
||||
headers: &HeaderMap,
|
||||
) -> Result<(), RateLimitRejection> {
|
||||
enforce_transport_rate_limit(state, path, headers).await
|
||||
}
|
||||
|
||||
pub(super) async fn enforce_transport_rate_limit(
|
||||
state: &Arc<AppState>,
|
||||
path: &AgentRoutePath,
|
||||
headers: &HeaderMap,
|
||||
) -> Result<(), RateLimitRejection> {
|
||||
) -> Result<(), RateLimitCheckError> {
|
||||
let key = rate_limit_key(path, headers);
|
||||
state.api_rate_limiter.check(&key).await
|
||||
}
|
||||
@@ -35,8 +27,25 @@ pub(super) fn rate_limited_jsonrpc_response(
|
||||
message: &Value,
|
||||
response_mode: ResponseMode,
|
||||
protocol_version: &str,
|
||||
rejection: RateLimitRejection,
|
||||
error: RateLimitCheckError,
|
||||
) -> Response {
|
||||
let RateLimitCheckError::Rejected(rejection) = error else {
|
||||
return transport_response(
|
||||
StatusCode::SERVICE_UNAVAILABLE,
|
||||
json!({
|
||||
"jsonrpc": "2.0",
|
||||
"id": request_id(message),
|
||||
"error": {
|
||||
"code": -32603,
|
||||
"message": "rate limit service unavailable",
|
||||
"data": { "code": "rate_limit_unavailable" }
|
||||
}
|
||||
}),
|
||||
response_mode,
|
||||
None,
|
||||
Some(protocol_version),
|
||||
);
|
||||
};
|
||||
let payload = json!({
|
||||
"jsonrpc": "2.0",
|
||||
"id": request_id(message),
|
||||
@@ -61,10 +70,15 @@ pub(super) fn rate_limited_jsonrpc_response(
|
||||
response
|
||||
}
|
||||
|
||||
pub(super) fn rate_limited_status_response(rejection: RateLimitRejection) -> Response {
|
||||
let mut response = StatusCode::TOO_MANY_REQUESTS.into_response();
|
||||
attach_retry_after_header(&mut response, rejection.retry_after_ms);
|
||||
response
|
||||
pub(super) fn rate_limited_status_response(error: RateLimitCheckError) -> Response {
|
||||
match error {
|
||||
RateLimitCheckError::Rejected(rejection) => {
|
||||
let mut response = StatusCode::TOO_MANY_REQUESTS.into_response();
|
||||
attach_retry_after_header(&mut response, rejection.retry_after_ms);
|
||||
response
|
||||
}
|
||||
RateLimitCheckError::StoreUnavailable => StatusCode::SERVICE_UNAVAILABLE.into_response(),
|
||||
}
|
||||
}
|
||||
|
||||
fn rate_limit_key(path: &AgentRoutePath, headers: &HeaderMap) -> String {
|
||||
|
||||
@@ -0,0 +1,39 @@
|
||||
use axum::{extract::Request, http::HeaderValue, middleware::Next, response::Response};
|
||||
use crank_observability::{RequestId, set_remote_trace_parent, with_request_correlation};
|
||||
use tracing::{Instrument, info_span};
|
||||
|
||||
use crate::transport::HEADER_X_REQUEST_ID;
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
pub(super) struct RequestContext {
|
||||
pub(super) request_id: String,
|
||||
}
|
||||
|
||||
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 context = RequestContext {
|
||||
request_id: request_id.clone(),
|
||||
};
|
||||
let span = info_span!(
|
||||
target: "crank::trace",
|
||||
"mcp.request",
|
||||
request_id = %request_id,
|
||||
);
|
||||
set_remote_trace_parent(&span, request.headers());
|
||||
request.extensions_mut().insert(context);
|
||||
|
||||
with_request_correlation(request_id.clone(), async move {
|
||||
let mut response = next.run(request).instrument(span).await;
|
||||
if let Ok(value) = HeaderValue::from_str(&request_id) {
|
||||
response.headers_mut().insert(HEADER_X_REQUEST_ID, value);
|
||||
}
|
||||
response
|
||||
})
|
||||
.await
|
||||
}
|
||||
@@ -52,6 +52,8 @@ pub trait TransportSessionStore: Send + Sync {
|
||||
) -> Result<bool, SessionStoreError>;
|
||||
|
||||
async fn delete(&self, session_id: &str) -> Result<bool, SessionStoreError>;
|
||||
|
||||
async fn cleanup_expired(&self, now: OffsetDateTime) -> Result<u64, SessionStoreError>;
|
||||
}
|
||||
|
||||
pub type SharedSessionStore = Arc<dyn TransportSessionStore>;
|
||||
@@ -62,6 +64,11 @@ pub struct PostgresTransportSessionStore {
|
||||
}
|
||||
|
||||
impl PostgresTransportSessionStore {
|
||||
pub async fn from_pool(pool: PgPool) -> Result<Self, SessionStoreError> {
|
||||
apply_postgres_migrations(&pool).await?;
|
||||
Ok(Self { pool })
|
||||
}
|
||||
|
||||
pub async fn connect_with_options_and_pool_config(
|
||||
connect_options: PgConnectOptions,
|
||||
pool_config: PostgresPoolConfig,
|
||||
@@ -84,9 +91,7 @@ impl PostgresTransportSessionStore {
|
||||
details: error.to_string(),
|
||||
})?;
|
||||
|
||||
apply_postgres_migrations(&pool).await?;
|
||||
|
||||
Ok(Self { pool })
|
||||
Self::from_pool(pool).await
|
||||
}
|
||||
}
|
||||
|
||||
@@ -164,6 +169,13 @@ impl TransportSessionStore for InMemorySessionStore {
|
||||
let mut guard = self.inner.write().await;
|
||||
Ok(guard.remove(session_id).is_some())
|
||||
}
|
||||
|
||||
async fn cleanup_expired(&self, now: OffsetDateTime) -> Result<u64, SessionStoreError> {
|
||||
let mut guard = self.inner.write().await;
|
||||
let before = guard.len();
|
||||
guard.retain(|_, session| !is_expired(session, now));
|
||||
Ok(u64::try_from(before.saturating_sub(guard.len())).unwrap_or(u64::MAX))
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
@@ -292,9 +304,68 @@ impl TransportSessionStore for PostgresTransportSessionStore {
|
||||
|
||||
Ok(result.rows_affected() > 0)
|
||||
}
|
||||
|
||||
async fn cleanup_expired(&self, now: OffsetDateTime) -> Result<u64, SessionStoreError> {
|
||||
let result = query(
|
||||
"delete from mcp_transport_sessions
|
||||
where expires_at is not null and expires_at <= $1::timestamptz",
|
||||
)
|
||||
.bind(now)
|
||||
.execute(&self.pool)
|
||||
.await
|
||||
.map_err(|error| SessionStoreError {
|
||||
details: error.to_string(),
|
||||
})?;
|
||||
|
||||
Ok(result.rows_affected())
|
||||
}
|
||||
}
|
||||
|
||||
async fn apply_postgres_migrations(pool: &PgPool) -> Result<(), SessionStoreError> {
|
||||
let mut transaction = pool.begin().await.map_err(|error| SessionStoreError {
|
||||
details: error.to_string(),
|
||||
})?;
|
||||
query("select pg_advisory_xact_lock($1)")
|
||||
.bind(0x4352_414E_4B4D_4350_i64)
|
||||
.execute(&mut *transaction)
|
||||
.await
|
||||
.map_err(|error| SessionStoreError {
|
||||
details: error.to_string(),
|
||||
})?;
|
||||
query(
|
||||
"create table if not exists __crank_mcp_migrations (
|
||||
version integer primary key,
|
||||
checksum text not null,
|
||||
applied_at timestamptz not null default now()
|
||||
)",
|
||||
)
|
||||
.execute(&mut *transaction)
|
||||
.await
|
||||
.map_err(|error| SessionStoreError {
|
||||
details: error.to_string(),
|
||||
})?;
|
||||
let applied = query("select checksum from __crank_mcp_migrations where version = 1")
|
||||
.fetch_optional(&mut *transaction)
|
||||
.await
|
||||
.map_err(|error| SessionStoreError {
|
||||
details: error.to_string(),
|
||||
})?;
|
||||
if let Some(row) = applied {
|
||||
let checksum = row.get::<String, _>("checksum");
|
||||
if checksum != "mcp-transport-sessions-v1" {
|
||||
return Err(SessionStoreError {
|
||||
details: format!("modified MCP migration version 1: {checksum}"),
|
||||
});
|
||||
}
|
||||
transaction
|
||||
.commit()
|
||||
.await
|
||||
.map_err(|error| SessionStoreError {
|
||||
details: error.to_string(),
|
||||
})?;
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
query(
|
||||
"create table if not exists mcp_transport_sessions (
|
||||
id text primary key,
|
||||
@@ -308,14 +379,14 @@ async fn apply_postgres_migrations(pool: &PgPool) -> Result<(), SessionStoreErro
|
||||
expires_at timestamptz null
|
||||
)",
|
||||
)
|
||||
.execute(pool)
|
||||
.execute(&mut *transaction)
|
||||
.await
|
||||
.map_err(|error| SessionStoreError {
|
||||
details: error.to_string(),
|
||||
})?;
|
||||
|
||||
query("alter table mcp_transport_sessions add column if not exists supports_elicitation boolean not null default false")
|
||||
.execute(pool)
|
||||
.execute(&mut *transaction)
|
||||
.await
|
||||
.map_err(|error| SessionStoreError {
|
||||
details: error.to_string(),
|
||||
@@ -324,7 +395,7 @@ async fn apply_postgres_migrations(pool: &PgPool) -> Result<(), SessionStoreErro
|
||||
query(
|
||||
"alter table mcp_transport_sessions add column if not exists expires_at timestamptz null",
|
||||
)
|
||||
.execute(pool)
|
||||
.execute(&mut *transaction)
|
||||
.await
|
||||
.map_err(|error| SessionStoreError {
|
||||
details: error.to_string(),
|
||||
@@ -334,12 +405,37 @@ async fn apply_postgres_migrations(pool: &PgPool) -> Result<(), SessionStoreErro
|
||||
"create index if not exists mcp_transport_sessions_workspace_agent_idx
|
||||
on mcp_transport_sessions(workspace_slug, agent_slug, updated_at desc)",
|
||||
)
|
||||
.execute(pool)
|
||||
.execute(&mut *transaction)
|
||||
.await
|
||||
.map_err(|error| SessionStoreError {
|
||||
details: error.to_string(),
|
||||
})?;
|
||||
|
||||
query(
|
||||
"create index if not exists mcp_transport_sessions_expires_at_idx
|
||||
on mcp_transport_sessions(expires_at)
|
||||
where expires_at is not null",
|
||||
)
|
||||
.execute(&mut *transaction)
|
||||
.await
|
||||
.map_err(|error| SessionStoreError {
|
||||
details: error.to_string(),
|
||||
})?;
|
||||
|
||||
query("insert into __crank_mcp_migrations (version, checksum) values (1, $1)")
|
||||
.bind("mcp-transport-sessions-v1")
|
||||
.execute(&mut *transaction)
|
||||
.await
|
||||
.map_err(|error| SessionStoreError {
|
||||
details: error.to_string(),
|
||||
})?;
|
||||
transaction
|
||||
.commit()
|
||||
.await
|
||||
.map_err(|error| SessionStoreError {
|
||||
details: error.to_string(),
|
||||
})?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
|
||||
@@ -86,6 +86,10 @@ pub fn runtime_error_code(error: &RuntimeError) -> &'static str {
|
||||
RuntimeError::ConfirmationRequired { .. } => "confirmation_required",
|
||||
RuntimeError::InvalidConfirmationToken { .. } => "invalid_confirmation_token",
|
||||
RuntimeError::ConfirmationStoreUnavailable { .. } => "confirmation_unavailable",
|
||||
RuntimeError::IdempotencyStoreUnavailable { .. } => "idempotency_unavailable",
|
||||
RuntimeError::IdempotencyInProgress { .. } => "idempotency_in_progress",
|
||||
RuntimeError::IdempotencyConflict { .. } => "idempotency_conflict",
|
||||
RuntimeError::IdempotencyOutcomeUnknown { .. } => "idempotency_outcome_unknown",
|
||||
RuntimeError::MissingAuthProfile { .. } => "auth_profile_not_found",
|
||||
RuntimeError::MissingSecret { .. } | RuntimeError::MissingSecretVersion { .. } => {
|
||||
"secret_not_found"
|
||||
@@ -146,6 +150,19 @@ fn safe_runtime_error_message(error: &RuntimeError) -> String {
|
||||
RuntimeError::ConfirmationStoreUnavailable { .. } => {
|
||||
"Хранилище подтверждений временно недоступно.".to_owned()
|
||||
}
|
||||
RuntimeError::IdempotencyStoreUnavailable { .. } => {
|
||||
"Хранилище идемпотентности временно недоступно.".to_owned()
|
||||
}
|
||||
RuntimeError::IdempotencyInProgress { .. } => {
|
||||
"Операция с этим ключом идемпотентности уже выполняется.".to_owned()
|
||||
}
|
||||
RuntimeError::IdempotencyConflict { .. } => {
|
||||
"Ключ идемпотентности уже использован с другими параметрами.".to_owned()
|
||||
}
|
||||
RuntimeError::IdempotencyOutcomeUnknown { .. } => {
|
||||
"Результат предыдущего выполнения неизвестен; автоматический повтор заблокирован."
|
||||
.to_owned()
|
||||
}
|
||||
RuntimeError::MissingAuthProfile { .. } => "Профиль авторизации не найден.".to_owned(),
|
||||
RuntimeError::MissingSecret { .. } | RuntimeError::MissingSecretVersion { .. } => {
|
||||
"Секрет авторизации не найден.".to_owned()
|
||||
@@ -169,6 +186,8 @@ fn is_recoverable_runtime_error(error: &RuntimeError) -> bool {
|
||||
| RuntimeError::ConcurrencyLimitExceeded { .. }
|
||||
| RuntimeError::SecretCrypto { .. }
|
||||
| RuntimeError::ConfirmationRequired { .. }
|
||||
| RuntimeError::IdempotencyStoreUnavailable { .. }
|
||||
| RuntimeError::IdempotencyInProgress { .. }
|
||||
)
|
||||
}
|
||||
|
||||
@@ -198,6 +217,14 @@ fn suggested_action(error: &RuntimeError) -> Option<&'static str> {
|
||||
Some("Запросите новый токен подтверждения.")
|
||||
}
|
||||
RuntimeError::ConfirmationStoreUnavailable { .. } => Some("Повторите запрос позже."),
|
||||
RuntimeError::IdempotencyStoreUnavailable { .. }
|
||||
| RuntimeError::IdempotencyInProgress { .. } => Some("Повторите запрос позже."),
|
||||
RuntimeError::IdempotencyConflict { .. } => {
|
||||
Some("Используйте новый ключ идемпотентности для изменённого запроса.")
|
||||
}
|
||||
RuntimeError::IdempotencyOutcomeUnknown { .. } => {
|
||||
Some("Проверьте результат во внешней системе перед ручным повтором.")
|
||||
}
|
||||
RuntimeError::MissingAuthProfile { .. }
|
||||
| RuntimeError::MissingSecret { .. }
|
||||
| RuntimeError::MissingSecretVersion { .. }
|
||||
|
||||
@@ -3,6 +3,7 @@ use std::{collections::BTreeSet, sync::Arc};
|
||||
use axum::{http::StatusCode, response::Response};
|
||||
use crank_core::{ToolAccessMode, search_tool_catalog};
|
||||
use crank_registry::PublishedAgentCatalog;
|
||||
use crank_trace::{ErrorCategory, Stage, StageOutcome};
|
||||
use serde::Deserialize;
|
||||
use serde_json::{Value, json};
|
||||
|
||||
@@ -127,13 +128,25 @@ async fn execute_catalog_tool(
|
||||
mut arguments: Value,
|
||||
transport_request_id: &str,
|
||||
) -> Response {
|
||||
let Some(resolved) = resolve_generated_tool(&catalog.tools, tool_name) else {
|
||||
return tool_not_found_response(
|
||||
message,
|
||||
response_mode,
|
||||
&session.protocol_version,
|
||||
tool_name,
|
||||
);
|
||||
let resolve_span = Stage::McpToolsResolve.span();
|
||||
let resolved = resolve_span.in_scope(|| resolve_generated_tool(&catalog.tools, tool_name));
|
||||
let resolved = match resolved {
|
||||
Some(resolved) => {
|
||||
StageOutcome::Success.record(&resolve_span);
|
||||
drop(resolve_span);
|
||||
resolved
|
||||
}
|
||||
None => {
|
||||
StageOutcome::Error.record(&resolve_span);
|
||||
ErrorCategory::Catalog.record(&resolve_span);
|
||||
drop(resolve_span);
|
||||
return tool_not_found_response(
|
||||
message,
|
||||
response_mode,
|
||||
&session.protocol_version,
|
||||
tool_name,
|
||||
);
|
||||
}
|
||||
};
|
||||
let confirmation_token = take_confirmation_token(&mut arguments);
|
||||
handle_tool_call(
|
||||
|
||||
@@ -22,7 +22,6 @@ use crate::jsonrpc::{
|
||||
pub(super) const HEADER_MCP_SESSION_ID: &str = "MCP-Session-Id";
|
||||
pub(super) const HEADER_MCP_PROTOCOL_VERSION: &str = "MCP-Protocol-Version";
|
||||
pub(super) const HEADER_X_REQUEST_ID: HeaderName = HeaderName::from_static("x-request-id");
|
||||
const MAX_REQUEST_ID_LEN: usize = 128;
|
||||
|
||||
#[derive(Clone, Copy)]
|
||||
pub(super) enum ResponseMode {
|
||||
@@ -303,24 +302,6 @@ where
|
||||
response
|
||||
}
|
||||
|
||||
pub(super) fn resolve_request_id(headers: &HeaderMap) -> String {
|
||||
headers
|
||||
.get(&HEADER_X_REQUEST_ID)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.map(str::trim)
|
||||
.filter(|value| is_valid_request_id(value))
|
||||
.map(ToOwned::to_owned)
|
||||
.unwrap_or_else(|| uuid::Uuid::now_v7().to_string())
|
||||
}
|
||||
|
||||
pub(super) fn is_valid_request_id(value: &str) -> bool {
|
||||
!value.is_empty()
|
||||
&& value.len() <= MAX_REQUEST_ID_LEN
|
||||
&& value
|
||||
.bytes()
|
||||
.all(|byte| matches!(byte, 0x21..=0x7e) && byte != b',' && byte != b';')
|
||||
}
|
||||
|
||||
pub(super) fn extract_origin(url: &str) -> Option<String> {
|
||||
Some(parse_origin(url, false)?.origin().ascii_serialization())
|
||||
}
|
||||
|
||||
@@ -97,3 +97,41 @@ async fn postgres_transport_sessions_evict_expired_rows_on_read() {
|
||||
|
||||
assert_eq!(remaining, 0);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn postgres_transport_session_cleanup_removes_abandoned_expired_rows() {
|
||||
let database_url = crank_test_support::postgres_schema_url("test_mcp_cleanup").await;
|
||||
let store = PostgresTransportSessionStore::connect_with_options_and_pool_config(
|
||||
database_url.parse::<PgConnectOptions>().unwrap(),
|
||||
PostgresPoolConfig::default(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
let now = OffsetDateTime::now_utc();
|
||||
|
||||
store
|
||||
.create(
|
||||
"2025-11-25",
|
||||
"default",
|
||||
"sales",
|
||||
false,
|
||||
now - time::Duration::hours(2),
|
||||
Some(now - time::Duration::hours(1)),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
let active = store
|
||||
.create(
|
||||
"2025-11-25",
|
||||
"default",
|
||||
"sales",
|
||||
false,
|
||||
now,
|
||||
Some(now + time::Duration::hours(1)),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(store.cleanup_expired(now).await.unwrap(), 1);
|
||||
assert!(store.get(&active).await.unwrap().is_some());
|
||||
}
|
||||
|
||||
@@ -72,6 +72,31 @@ async fn drops_expired_in_memory_transport_sessions_on_read() {
|
||||
assert!(store.get(&session_id).await.unwrap().is_none());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn cleanup_removes_only_expired_sessions() {
|
||||
let store = InMemorySessionStore::default();
|
||||
let now = time::OffsetDateTime::now_utc();
|
||||
let expired = store
|
||||
.create("2025-11-25", "default", "sales", false, now, Some(now))
|
||||
.await
|
||||
.unwrap();
|
||||
let active = store
|
||||
.create(
|
||||
"2025-11-25",
|
||||
"default",
|
||||
"sales",
|
||||
false,
|
||||
now,
|
||||
Some(now + time::Duration::hours(1)),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(store.cleanup_expired(now).await.unwrap(), 1);
|
||||
assert!(store.get(&expired).await.unwrap().is_none());
|
||||
assert!(store.get(&active).await.unwrap().is_some());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn formats_transport_session_store_error() {
|
||||
let error = SessionStoreError {
|
||||
|
||||
Reference in New Issue
Block a user