mcp: extract community mcp crate

This commit is contained in:
github-ops
2026-05-14 20:10:44 +00:00
parent 7c2a8f4fac
commit c080692df5
12 changed files with 64 additions and 13 deletions
File diff suppressed because it is too large Load Diff
-22
View File
@@ -1,22 +0,0 @@
use async_trait::async_trait;
pub use crank_core::{
MachineCredentialVerifier, MachineCredentialVerifierError, SharedMachineCredentialVerifier,
VerifiedMachineCredential,
};
#[derive(Clone, Default)]
pub struct StaticAgentKeyVerifier;
#[async_trait]
impl MachineCredentialVerifier for StaticAgentKeyVerifier {
async fn verify_bearer_token(
&self,
_workspace_slug: &str,
_agent_slug: &str,
_token: &str,
) -> Result<Option<VerifiedMachineCredential>, MachineCredentialVerifierError> {
Ok(None)
}
}
pub use StaticAgentKeyVerifier as CommunityMachineCredentialVerifier;
-206
View File
@@ -1,206 +0,0 @@
use std::{
collections::HashMap,
sync::Arc,
time::{Duration, Instant},
};
use crank_core::{CacheScope, CoordinationStateStore, CoordinationStateValue};
use crank_registry::{PostgresRegistry, PublishedAgentTool, RegistryError};
use serde::{Deserialize, Serialize};
use tokio::sync::RwLock;
use tracing::info;
#[derive(Clone)]
pub struct PublishedToolCatalog {
registry: PostgresRegistry,
refresh_interval: Duration,
coordination_store: Arc<dyn CoordinationStateStore>,
cached: Arc<RwLock<HashMap<CatalogKey, CachedCatalog>>>,
}
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
struct CatalogKey {
workspace_slug: String,
agent_slug: String,
}
#[derive(Default)]
struct CachedCatalog {
loaded_at: Option<Instant>,
tools: Vec<PublishedAgentTool>,
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
struct CatalogSnapshot {
tools: Vec<PublishedAgentTool>,
}
impl PublishedToolCatalog {
pub fn new(
registry: PostgresRegistry,
refresh_interval: Duration,
coordination_store: Arc<dyn CoordinationStateStore>,
) -> Self {
Self {
registry,
refresh_interval,
coordination_store,
cached: Arc::new(RwLock::new(HashMap::new())),
}
}
pub async fn list_tools(
&self,
workspace_slug: &str,
agent_slug: &str,
) -> Result<Vec<PublishedAgentTool>, RegistryError> {
self.refresh_if_stale(workspace_slug, agent_slug).await?;
let guard = self.cached.read().await;
Ok(guard
.get(&CatalogKey::new(workspace_slug, agent_slug))
.map(|entry| entry.tools.clone())
.unwrap_or_default())
}
async fn refresh_if_stale(
&self,
workspace_slug: &str,
agent_slug: &str,
) -> Result<(), RegistryError> {
let key = CatalogKey::new(workspace_slug, agent_slug);
let should_refresh = {
let guard = self.cached.read().await;
match guard.get(&key).and_then(|entry| entry.loaded_at) {
Some(loaded_at) => loaded_at.elapsed() >= self.refresh_interval,
None => true,
}
};
if !should_refresh {
return Ok(());
}
if let Some(tools) = self.load_shared_snapshot(workspace_slug, agent_slug).await {
let mut guard = self.cached.write().await;
guard.insert(
key,
CachedCatalog {
loaded_at: Some(Instant::now()),
tools,
},
);
return Ok(());
}
let tools = match self
.registry
.get_published_agent_tools_by_slug(workspace_slug, agent_slug)
.await
{
Ok(tools) => tools,
Err(RegistryError::PublishedAgentNotFound { .. }) => Vec::new(),
Err(error) => return Err(error),
};
self.store_shared_snapshot(workspace_slug, agent_slug, &tools)
.await;
let mut guard = self.cached.write().await;
let previous_count = guard
.get(&key)
.map(|entry| entry.tools.len())
.unwrap_or_default();
guard.insert(
key,
CachedCatalog {
loaded_at: Some(Instant::now()),
tools,
},
);
info!(
workspace_slug,
agent_slug,
published_tool_count = guard
.get(&CatalogKey::new(workspace_slug, agent_slug))
.map(|entry| entry.tools.len())
.unwrap_or_default(),
previous_published_tool_count = previous_count,
"published agent catalog refreshed"
);
Ok(())
}
async fn load_shared_snapshot(
&self,
workspace_slug: &str,
agent_slug: &str,
) -> Option<Vec<PublishedAgentTool>> {
if self.refresh_interval.is_zero() {
return None;
}
let key = catalog_snapshot_key(workspace_slug, agent_slug);
let value = match self
.coordination_store
.get_value(CacheScope::Coordination, &key)
.await
{
Ok(value) => value?,
Err(_) => return None,
};
serde_json::from_value::<CatalogSnapshot>(value.payload)
.ok()
.map(|snapshot| snapshot.tools)
}
async fn store_shared_snapshot(
&self,
workspace_slug: &str,
agent_slug: &str,
tools: &[PublishedAgentTool],
) {
let Some(ttl) = catalog_snapshot_ttl(self.refresh_interval) else {
return;
};
let payload = match serde_json::to_value(CatalogSnapshot {
tools: tools.to_vec(),
}) {
Ok(payload) => payload,
Err(_) => return,
};
let key = catalog_snapshot_key(workspace_slug, agent_slug);
let _ = self
.coordination_store
.put_value(
CacheScope::Coordination,
&key,
CoordinationStateValue { payload },
ttl,
)
.await;
}
}
impl CatalogKey {
fn new(workspace_slug: &str, agent_slug: &str) -> Self {
Self {
workspace_slug: workspace_slug.to_owned(),
agent_slug: agent_slug.to_owned(),
}
}
}
fn catalog_snapshot_key(workspace_slug: &str, agent_slug: &str) -> String {
format!("published_catalog:{workspace_slug}:{agent_slug}")
}
fn catalog_snapshot_ttl(refresh_interval: Duration) -> Option<Duration> {
if refresh_interval.is_zero() {
return None;
}
refresh_interval.checked_mul(2).or(Some(refresh_interval))
}
-63
View File
@@ -1,63 +0,0 @@
use serde_json::{Value, json};
pub const CURRENT_PROTOCOL_VERSION: &str = "2025-11-25";
pub const DEFAULT_PROTOCOL_VERSION: &str = "2025-03-26";
pub const SUPPORTED_PROTOCOL_VERSIONS: &[&str] = &[
DEFAULT_PROTOCOL_VERSION,
"2025-06-18",
CURRENT_PROTOCOL_VERSION,
];
pub fn request_id(message: &Value) -> Value {
message.get("id").cloned().unwrap_or(Value::Null)
}
pub fn method_name(message: &Value) -> Option<&str> {
message.get("method").and_then(Value::as_str)
}
pub fn params(message: &Value) -> Value {
message
.get("params")
.cloned()
.unwrap_or_else(|| Value::Object(Default::default()))
}
pub fn is_notification(message: &Value) -> bool {
method_name(message).is_some() && message.get("id").is_none()
}
pub fn is_request(message: &Value) -> bool {
method_name(message).is_some() && message.get("id").is_some()
}
pub fn is_response(message: &Value) -> bool {
method_name(message).is_none()
&& (message.get("result").is_some() || message.get("error").is_some())
}
pub fn jsonrpc_result(id: Value, result: Value) -> Value {
json!({
"jsonrpc": "2.0",
"id": id,
"result": result
})
}
pub fn jsonrpc_error(id: Value, code: i64, message: impl Into<String>) -> Value {
json!({
"jsonrpc": "2.0",
"id": id,
"error": {
"code": code,
"message": message.into()
}
})
}
pub fn negotiated_protocol_version(requested: &str) -> Option<&'static str> {
SUPPORTED_PROTOCOL_VERSIONS
.iter()
.copied()
.find(|version| *version == requested)
}
+5 -13
View File
@@ -1,11 +1,8 @@
mod app;
mod auth;
mod catalog;
mod jsonrpc;
mod session;
use std::{env, net::SocketAddr, time::Duration};
use crank_community_mcp::{
auth::CommunityMachineCredentialVerifier, build_app, session::PostgresTransportSessionStore,
};
use crank_registry::{PostgresPoolConfig, PostgresRegistry};
use crank_runtime::{
RequestRateLimitConfig, RequestRateLimiter, RuntimeCacheConfig, RuntimeCacheStores,
@@ -15,11 +12,6 @@ use sqlx::postgres::PgConnectOptions;
use tokio::net::TcpListener;
use tracing::info;
use crate::{
app::build_app, auth::CommunityMachineCredentialVerifier,
session::PostgresTransportSessionStore,
};
#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
tracing_subscriber::fmt()
@@ -184,8 +176,8 @@ mod tests {
use tokio::time::sleep;
use tracing_subscriber::{filter::LevelFilter, fmt::MakeWriter, prelude::*};
use crate::{
app::build_app,
use crank_community_mcp::{
build_app,
auth::{
CommunityMachineCredentialVerifier, MachineCredentialVerifier,
MachineCredentialVerifierError, SharedMachineCredentialVerifier,
-550
View File
@@ -1,550 +0,0 @@
use std::{collections::HashMap, sync::Arc};
use async_trait::async_trait;
use crank_registry::PostgresPoolConfig;
use sqlx::{
PgPool,
postgres::{PgConnectOptions, PgPoolOptions},
query,
};
use thiserror::Error;
use time::OffsetDateTime;
use tokio::sync::RwLock;
use uuid::Uuid;
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct SessionState {
pub id: String,
pub protocol_version: String,
pub initialized: bool,
pub workspace_slug: String,
pub agent_slug: String,
pub created_at: OffsetDateTime,
pub updated_at: OffsetDateTime,
pub expires_at: Option<OffsetDateTime>,
}
#[derive(Debug, Error)]
#[error("transport session store is unavailable: {details}")]
pub struct SessionStoreError {
pub details: String,
}
#[async_trait]
pub trait TransportSessionStore: Send + Sync {
async fn create(
&self,
protocol_version: &str,
workspace_slug: &str,
agent_slug: &str,
now: OffsetDateTime,
expires_at: Option<OffsetDateTime>,
) -> Result<String, SessionStoreError>;
async fn get(&self, session_id: &str) -> Result<Option<SessionState>, SessionStoreError>;
async fn mark_initialized(
&self,
session_id: &str,
now: OffsetDateTime,
) -> Result<bool, SessionStoreError>;
async fn delete(&self, session_id: &str) -> Result<bool, SessionStoreError>;
}
pub type SharedSessionStore = Arc<dyn TransportSessionStore>;
#[derive(Clone, Debug)]
pub struct PostgresTransportSessionStore {
pool: PgPool,
}
impl PostgresTransportSessionStore {
pub async fn connect_with_options_and_pool_config(
connect_options: PgConnectOptions,
pool_config: PostgresPoolConfig,
) -> Result<Self, SessionStoreError> {
let pool = PgPoolOptions::new()
.max_connections(pool_config.max_connections)
.min_connections(pool_config.min_connections)
.acquire_timeout(std::time::Duration::from_millis(
pool_config.acquire_timeout_ms,
))
.idle_timeout(Some(std::time::Duration::from_millis(
pool_config.idle_timeout_ms,
)))
.max_lifetime(Some(std::time::Duration::from_millis(
pool_config.max_lifetime_ms,
)))
.connect_with(connect_options)
.await
.map_err(|error| SessionStoreError {
details: error.to_string(),
})?;
apply_postgres_migrations(&pool).await?;
Ok(Self { pool })
}
}
#[derive(Clone, Default)]
pub struct InMemorySessionStore {
inner: Arc<RwLock<HashMap<String, SessionState>>>,
}
#[async_trait]
impl TransportSessionStore for InMemorySessionStore {
async fn create(
&self,
protocol_version: &str,
workspace_slug: &str,
agent_slug: &str,
now: OffsetDateTime,
expires_at: Option<OffsetDateTime>,
) -> Result<String, SessionStoreError> {
let session_id = Uuid::now_v7().to_string();
let mut guard = self.inner.write().await;
guard.insert(
session_id.clone(),
SessionState {
id: session_id.clone(),
protocol_version: protocol_version.to_owned(),
initialized: false,
workspace_slug: workspace_slug.to_owned(),
agent_slug: agent_slug.to_owned(),
created_at: now,
updated_at: now,
expires_at,
},
);
Ok(session_id)
}
async fn get(&self, session_id: &str) -> Result<Option<SessionState>, SessionStoreError> {
{
let guard = self.inner.read().await;
if let Some(session) = guard.get(session_id) {
if !is_expired(session, OffsetDateTime::now_utc()) {
return Ok(Some(session.clone()));
}
} else {
return Ok(None);
}
}
let mut guard = self.inner.write().await;
guard.remove(session_id);
Ok(None)
}
async fn mark_initialized(
&self,
session_id: &str,
now: OffsetDateTime,
) -> Result<bool, SessionStoreError> {
let mut guard = self.inner.write().await;
if let Some(session) = guard.get_mut(session_id) {
session.initialized = true;
session.updated_at = now;
return Ok(true);
}
Ok(false)
}
async fn delete(&self, session_id: &str) -> Result<bool, SessionStoreError> {
let mut guard = self.inner.write().await;
Ok(guard.remove(session_id).is_some())
}
}
#[async_trait]
impl TransportSessionStore for PostgresTransportSessionStore {
async fn create(
&self,
protocol_version: &str,
workspace_slug: &str,
agent_slug: &str,
now: OffsetDateTime,
expires_at: Option<OffsetDateTime>,
) -> Result<String, SessionStoreError> {
let session_id = Uuid::now_v7().to_string();
query(
"insert into mcp_transport_sessions (
id,
protocol_version,
initialized,
workspace_slug,
agent_slug,
created_at,
updated_at,
expires_at
) values (
$1, $2, false, $3, $4, $5::timestamptz, $5::timestamptz, $6::timestamptz
)",
)
.bind(&session_id)
.bind(protocol_version)
.bind(workspace_slug)
.bind(agent_slug)
.bind(now)
.bind(expires_at)
.execute(&self.pool)
.await
.map_err(|error| SessionStoreError {
details: error.to_string(),
})?;
Ok(session_id)
}
async fn get(&self, session_id: &str) -> Result<Option<SessionState>, SessionStoreError> {
let row = sqlx::query!(
"select
id,
protocol_version,
initialized,
workspace_slug,
agent_slug,
created_at as \"created_at!: OffsetDateTime\",
updated_at as \"updated_at!: OffsetDateTime\",
expires_at as \"expires_at: OffsetDateTime\"
from mcp_transport_sessions
where id = $1",
session_id,
)
.fetch_optional(&self.pool)
.await
.map_err(|error| SessionStoreError {
details: error.to_string(),
})?;
let Some(session) = row.map(|row| SessionState {
id: row.id,
protocol_version: row.protocol_version,
initialized: row.initialized,
workspace_slug: row.workspace_slug,
agent_slug: row.agent_slug,
created_at: row.created_at,
updated_at: row.updated_at,
expires_at: row.expires_at,
}) else {
return Ok(None);
};
if is_expired(&session, OffsetDateTime::now_utc()) {
query("delete from mcp_transport_sessions where id = $1")
.bind(session_id)
.execute(&self.pool)
.await
.map_err(|error| SessionStoreError {
details: error.to_string(),
})?;
return Ok(None);
}
Ok(Some(session))
}
async fn mark_initialized(
&self,
session_id: &str,
now: OffsetDateTime,
) -> Result<bool, SessionStoreError> {
let result = query(
"update mcp_transport_sessions
set initialized = true,
updated_at = $2::timestamptz
where id = $1",
)
.bind(session_id)
.bind(now)
.execute(&self.pool)
.await
.map_err(|error| SessionStoreError {
details: error.to_string(),
})?;
Ok(result.rows_affected() > 0)
}
async fn delete(&self, session_id: &str) -> Result<bool, SessionStoreError> {
let result = query("delete from mcp_transport_sessions where id = $1")
.bind(session_id)
.execute(&self.pool)
.await
.map_err(|error| SessionStoreError {
details: error.to_string(),
})?;
Ok(result.rows_affected() > 0)
}
}
async fn apply_postgres_migrations(pool: &PgPool) -> Result<(), SessionStoreError> {
query(
"create table if not exists mcp_transport_sessions (
id text primary key,
protocol_version text not null,
initialized boolean not null default false,
workspace_slug text not null,
agent_slug text not null,
created_at timestamptz not null,
updated_at timestamptz not null,
expires_at timestamptz null
)",
)
.execute(pool)
.await
.map_err(|error| SessionStoreError {
details: error.to_string(),
})?;
query(
"alter table mcp_transport_sessions add column if not exists expires_at timestamptz null",
)
.execute(pool)
.await
.map_err(|error| SessionStoreError {
details: error.to_string(),
})?;
query(
"create index if not exists mcp_transport_sessions_workspace_agent_idx
on mcp_transport_sessions(workspace_slug, agent_slug, updated_at desc)",
)
.execute(pool)
.await
.map_err(|error| SessionStoreError {
details: error.to_string(),
})?;
Ok(())
}
fn is_expired(session: &SessionState, now: OffsetDateTime) -> bool {
session
.expires_at
.is_some_and(|expires_at| expires_at <= now)
}
#[cfg(test)]
mod tests {
use std::env;
use sqlx::{
Executor,
postgres::{PgConnectOptions, PgPoolOptions},
};
use time::{OffsetDateTime, format_description::well_known::Rfc3339};
use uuid::Uuid;
use super::{
InMemorySessionStore, PostgresTransportSessionStore, SessionStoreError,
TransportSessionStore,
};
use crank_registry::PostgresPoolConfig;
fn timestamp(value: &str) -> time::OffsetDateTime {
time::OffsetDateTime::parse(value, &Rfc3339).unwrap()
}
fn truncate_to_micros(value: OffsetDateTime) -> OffsetDateTime {
value
.replace_nanosecond((value.nanosecond() / 1_000) * 1_000)
.unwrap()
}
#[tokio::test]
async fn creates_and_reads_transport_sessions() {
let store = InMemorySessionStore::default();
let created_at = timestamp("2026-05-01T10:00:00Z");
let session_id = store
.create("2025-11-25", "default", "sales", created_at, None)
.await
.unwrap();
let session = store.get(&session_id).await.unwrap().unwrap();
assert_eq!(session.id, session_id);
assert_eq!(session.protocol_version, "2025-11-25");
assert_eq!(session.workspace_slug, "default");
assert_eq!(session.agent_slug, "sales");
assert!(!session.initialized);
assert_eq!(session.created_at, created_at);
assert_eq!(session.updated_at, created_at);
}
#[tokio::test]
async fn marks_transport_sessions_initialized() {
let store = InMemorySessionStore::default();
let created_at = timestamp("2026-05-01T10:00:00Z");
let initialized_at = timestamp("2026-05-01T10:00:05Z");
let session_id = store
.create("2025-11-25", "default", "sales", created_at, None)
.await
.unwrap();
assert!(
store
.mark_initialized(&session_id, initialized_at)
.await
.unwrap()
);
let session = store.get(&session_id).await.unwrap().unwrap();
assert!(session.initialized);
assert_eq!(session.updated_at, initialized_at);
}
#[tokio::test]
async fn drops_expired_in_memory_transport_sessions_on_read() {
let store = InMemorySessionStore::default();
let created_at = timestamp("2026-05-01T10:00:00Z");
let expires_at = timestamp("2026-05-01T10:00:01Z");
let session_id = store
.create(
"2025-11-25",
"default",
"sales",
created_at,
Some(expires_at),
)
.await
.unwrap();
assert!(store.get(&session_id).await.unwrap().is_none());
}
#[test]
fn formats_transport_session_store_error() {
let error = SessionStoreError {
details: "postgres is unavailable".to_owned(),
};
assert_eq!(
error.to_string(),
"transport session store is unavailable: postgres is unavailable"
);
}
#[tokio::test]
async fn postgres_transport_sessions_survive_store_reconnect() {
let database_url = env::var("TEST_DATABASE_URL")
.unwrap_or_else(|_| "postgres://crank:crank@127.0.0.1:15432/crank".to_owned());
let admin_pool = PgPoolOptions::new()
.max_connections(1)
.connect(&database_url)
.await
.unwrap();
let schema = format!("test_mcp_transport_{}", Uuid::now_v7().simple());
admin_pool
.execute(sqlx::query(&format!("create schema {schema}")))
.await
.unwrap();
let connect_options = format!("{database_url}?options=-csearch_path%3D{schema}")
.parse::<PgConnectOptions>()
.unwrap();
let pool_config = PostgresPoolConfig::default();
let store_a = PostgresTransportSessionStore::connect_with_options_and_pool_config(
connect_options.clone(),
pool_config,
)
.await
.unwrap();
let created_at = OffsetDateTime::now_utc() - time::Duration::hours(1);
let initialized_at = created_at + time::Duration::seconds(5);
let session_id = store_a
.create(
"2025-11-25",
"default",
"sales",
created_at,
Some(created_at + time::Duration::days(30)),
)
.await
.unwrap();
store_a
.mark_initialized(&session_id, initialized_at)
.await
.unwrap();
let store_b = PostgresTransportSessionStore::connect_with_options_and_pool_config(
connect_options,
pool_config,
)
.await
.unwrap();
let session = store_b.get(&session_id).await.unwrap().unwrap();
assert_eq!(session.id, session_id);
assert!(session.initialized);
assert_eq!(session.updated_at, truncate_to_micros(initialized_at));
assert_eq!(session.workspace_slug, "default");
assert_eq!(session.agent_slug, "sales");
}
#[tokio::test]
async fn postgres_transport_sessions_evict_expired_rows_on_read() {
let database_url = env::var("TEST_DATABASE_URL")
.unwrap_or_else(|_| "postgres://crank:crank@127.0.0.1:15432/crank".to_owned());
let admin_pool = PgPoolOptions::new()
.max_connections(1)
.connect(&database_url)
.await
.unwrap();
let schema = format!("test_mcp_transport_{}", Uuid::now_v7().simple());
admin_pool
.execute(sqlx::query(&format!("create schema {schema}")))
.await
.unwrap();
let connect_options = format!("{database_url}?options=-csearch_path%3D{schema}")
.parse::<PgConnectOptions>()
.unwrap();
let store = PostgresTransportSessionStore::connect_with_options_and_pool_config(
connect_options.clone(),
PostgresPoolConfig::default(),
)
.await
.unwrap();
let session_id = store
.create(
"2025-11-25",
"default",
"sales",
timestamp("2026-05-01T10:00:00Z"),
Some(timestamp("2026-05-01T10:00:01Z")),
)
.await
.unwrap();
assert!(store.get(&session_id).await.unwrap().is_none());
let pool = PgPoolOptions::new()
.max_connections(1)
.connect_with(connect_options)
.await
.unwrap();
let remaining = sqlx::query_scalar::<_, i64>(
"select count(*) from mcp_transport_sessions where id = $1",
)
.bind(&session_id)
.fetch_one(&pool)
.await
.unwrap();
assert_eq!(remaining, 0);
}
}