mcp: extract community mcp crate
This commit is contained in:
@@ -0,0 +1,25 @@
|
||||
[package]
|
||||
name = "crank-community-mcp"
|
||||
edition.workspace = true
|
||||
license.workspace = true
|
||||
rust-version.workspace = true
|
||||
version.workspace = true
|
||||
|
||||
[dependencies]
|
||||
async-trait = "0.1"
|
||||
axum.workspace = true
|
||||
base64.workspace = true
|
||||
crank-core = { path = "../crank-core" }
|
||||
crank-registry = { path = "../crank-registry" }
|
||||
crank-runtime = { path = "../crank-runtime" }
|
||||
crank-schema = { path = "../crank-schema" }
|
||||
futures-util = "0.3"
|
||||
serde.workspace = true
|
||||
serde_json.workspace = true
|
||||
sha2.workspace = true
|
||||
sqlx.workspace = true
|
||||
thiserror.workspace = true
|
||||
time.workspace = true
|
||||
tokio = { workspace = true, features = ["sync"] }
|
||||
tracing.workspace = true
|
||||
uuid.workspace = true
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,22 @@
|
||||
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;
|
||||
@@ -0,0 +1,206 @@
|
||||
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))
|
||||
}
|
||||
@@ -0,0 +1,63 @@
|
||||
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)
|
||||
}
|
||||
@@ -0,0 +1,7 @@
|
||||
mod app;
|
||||
pub mod auth;
|
||||
pub mod catalog;
|
||||
pub mod jsonrpc;
|
||||
pub mod session;
|
||||
|
||||
pub use app::build_app;
|
||||
@@ -0,0 +1,550 @@
|
||||
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);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user