From 50a30d4504c7f7bff747c08f97c92d6726e7b433 Mon Sep 17 00:00:00 2001 From: fylorn <249551762+fylorn@users.noreply.github.com> Date: Thu, 24 Sep 2026 16:48:51 +0800 Subject: [PATCH] refactor(server): move the MCP handlers' SQL into repositories The MCP server, connection, shared-credential, store and tool-catalog handlers no longer carry SQL. Their 50 statements move, text and binds unchanged, into four repositories under `services/`: - mcp_server_repository: mcp_servers reads and writes, the update and delete transactions (with their credential purge and install-count decrement), last_error updates, the unique-violation to 409 mapping - mcp_credential_repository: per-user and shared credentials, including the upsert, revoke and set-default transactions - mcp_store_repository: templates, installs, the registry-sync upsert and prune, the install advisory lock - mcp_tool_repository: the tool catalog count and page (McpToolRow moves here) Server create and registry sync keep their transaction in the handler and pass it to the repository, since they interleave non-SQL work. admin_mcp_catalog.rs covers the endpoints the suite didn't reach; it passes on the code before and after this change. Co-Authored-By: Claude Opus 5.5 --- crates/server/src/handlers/mcp_oauth.rs | 234 +--- .../server/src/handlers/mcp_oauth/shared.rs | 105 +- .../server/src/handlers/mcp_oauth/wizard.rs | 31 +- crates/server/src/handlers/mcp_servers.rs | 341 ++---- crates/server/src/handlers/mcp_store.rs | 139 +-- crates/server/src/handlers/mcp_tools.rs | 127 +-- crates/server/src/openapi.rs | 3 +- .../src/services/mcp_credential_repository.rs | 403 +++++++ .../src/services/mcp_server_repository.rs | 314 ++++++ .../src/services/mcp_store_repository.rs | 235 +++++ .../src/services/mcp_tool_repository.rs | 136 +++ crates/server/src/services/mod.rs | 4 + .../test-support/tests/admin_mcp_catalog.rs | 995 ++++++++++++++++++ 13 files changed, 2286 insertions(+), 781 deletions(-) create mode 100644 crates/server/src/services/mcp_credential_repository.rs create mode 100644 crates/server/src/services/mcp_server_repository.rs create mode 100644 crates/server/src/services/mcp_store_repository.rs create mode 100644 crates/server/src/services/mcp_tool_repository.rs create mode 100644 crates/test-support/tests/admin_mcp_catalog.rs diff --git a/crates/server/src/handlers/mcp_oauth.rs b/crates/server/src/handlers/mcp_oauth.rs index 0800f32a..88311499 100644 --- a/crates/server/src/handlers/mcp_oauth.rs +++ b/crates/server/src/handlers/mcp_oauth.rs @@ -32,8 +32,8 @@ pub use shared::{ }; pub use wizard::{ PoppedWizardCredential, WizardAuthorizeRequest, WizardCredentialStatus, - claim_wizard_credential, discard_wizard_credential, insert_shared_credential_from_wizard, - start_wizard_authorize, wizard_credential_status, + claim_wizard_credential, discard_wizard_credential, start_wizard_authorize, + wizard_credential_status, }; use axum::Json; @@ -55,6 +55,8 @@ use think_watch_common::models::McpServer; use crate::app::AppState; use crate::middleware::auth_guard::AuthUser; +use crate::services::mcp_credential_repository as credential_repo; +use crate::services::mcp_server_repository as server_repo; pub(super) const OAUTH_STATE_PREFIX: &str = "mcp_oauth:state:"; pub(super) const OAUTH_STATE_TTL_SECS: i64 = 600; @@ -193,36 +195,8 @@ pub async fn list_connections( ) -> Result>, AppError> { auth_user.require_permission("mcp:connect")?; - let servers = sqlx::query_as::<_, McpServer>( - r#"SELECT s.*, 0::bigint AS tools_count, 0::bigint AS call_count - FROM mcp_servers s - ORDER BY s.name"#, - ) - .fetch_all(&state.db) - .await?; - - #[derive(sqlx::FromRow)] - struct AccountRow { - mcp_server_id: Uuid, - account_label: String, - credential_type: String, - is_default: bool, - scopes: Vec, - expires_at: Option>, - upstream_subject: Option, - created_at: DateTime, - updated_at: DateTime, - } - let rows = sqlx::query_as::<_, AccountRow>( - r#"SELECT mcp_server_id, account_label, credential_type, is_default, - scopes, expires_at, upstream_subject, created_at, updated_at - FROM mcp_user_credentials - WHERE user_id = $1 - ORDER BY mcp_server_id, is_default DESC, account_label"#, - ) - .bind(auth_user.claims.sub) - .fetch_all(&state.db) - .await?; + let servers = server_repo::list_by_name(&state.db).await?; + let rows = credential_repo::list_user_accounts(&state.db, auth_user.claims.sub).await?; let mut out = Vec::with_capacity(servers.len()); for s in servers { @@ -577,8 +551,8 @@ pub async fn oauth_callback( } => { let server = load_server(&state, *server_id).await?; retry_pg_storage("per_user_credential", || { - upsert_credential( - &state, + credential_repo::upsert_user_credential( + &state.db, *server_id, *user_id, account_label, @@ -641,8 +615,8 @@ pub async fn oauth_callback( } => { let server = load_server(&state, *server_id).await?; retry_pg_storage("admin_shared_credential", || { - shared::upsert_shared_credential( - &state, + credential_repo::upsert_shared_credential( + &state.db, *server_id, "oauth_authcode", &access_encrypted, @@ -911,15 +885,12 @@ pub async fn revoke_connection( // Best-effort revoke at the upstream — only when we actually have // an access_token AND the server advertises a revocation endpoint. - let row: Option<(String, Vec)> = sqlx::query_as( - r#"SELECT credential_type, access_token_encrypted - FROM mcp_user_credentials - WHERE mcp_server_id = $1 AND user_id = $2 AND account_label = $3"#, + let row = credential_repo::find_user_token( + &state.db, + server_id, + auth_user.claims.sub, + &account_label, ) - .bind(server_id) - .bind(auth_user.claims.sub) - .bind(&account_label) - .fetch_optional(&state.db) .await?; let Some((credential_type, access_encrypted)) = row else { return Err(AppError::NotFound("Connection not found".into())); @@ -955,43 +926,14 @@ pub async fn revoke_connection( // call, even though the user clearly still has a usable connection. // Promote the most recently created remaining row as a graceful // fallback so the user keeps working without manually re-marking. - let mut tx = state.db.begin().await?; - let was_default: Option = sqlx::query_scalar( - r#"DELETE FROM mcp_user_credentials - WHERE mcp_server_id = $1 AND user_id = $2 AND account_label = $3 - RETURNING is_default"#, + credential_repo::delete_user_credential( + &state.db, + server_id, + auth_user.claims.sub, + &account_label, ) - .bind(server_id) - .bind(auth_user.claims.sub) - .bind(&account_label) - .fetch_optional(&mut *tx) .await?; - if matches!(was_default, Some(true)) { - // Promote the newest remaining credential for the same - // (server, user). Newest wins because a user juggling - // multiple credentials usually treats the latest one as - // "current" — same heuristic the connect-then-overwrite UX - // already nudges them toward. NULL `created_at` shouldn't - // exist (column is NOT NULL DEFAULT now()) but the ORDER BY - // is still safe under NULLS LAST. - sqlx::query( - r#"UPDATE mcp_user_credentials - SET is_default = true - WHERE id = ( - SELECT id FROM mcp_user_credentials - WHERE mcp_server_id = $1 AND user_id = $2 - ORDER BY created_at DESC NULLS LAST - LIMIT 1 - )"#, - ) - .bind(server_id) - .bind(auth_user.claims.sub) - .execute(&mut *tx) - .await?; - } - tx.commit().await?; - // Cached responses pinned to this credential are now serving an // identity that no longer has access. Wipe the user's lane for // this server so post-revoke reads can't tunnel back to the @@ -1034,42 +976,17 @@ pub async fn set_default_connection( ) .await?; - let mut tx = state.db.begin().await?; - let exists: Option = sqlx::query_scalar( - r#"SELECT 1 FROM mcp_user_credentials - WHERE mcp_server_id = $1 AND user_id = $2 AND account_label = $3"#, + let found = credential_repo::set_default_user_credential( + &state.db, + server_id, + auth_user.claims.sub, + &account_label, ) - .bind(server_id) - .bind(auth_user.claims.sub) - .bind(&account_label) - .fetch_optional(&mut *tx) .await?; - if exists.is_none() { + if !found { return Err(AppError::NotFound("Connection not found".into())); } - // Two-step toggle so the partial unique index never sees two - // is_default rows at once: clear the old default first, then mark - // the new one inside the same transaction. - sqlx::query( - r#"UPDATE mcp_user_credentials SET is_default = false, updated_at = now() - WHERE mcp_server_id = $1 AND user_id = $2 AND is_default"#, - ) - .bind(server_id) - .bind(auth_user.claims.sub) - .execute(&mut *tx) - .await?; - sqlx::query( - r#"UPDATE mcp_user_credentials SET is_default = true, updated_at = now() - WHERE mcp_server_id = $1 AND user_id = $2 AND account_label = $3"#, - ) - .bind(server_id) - .bind(auth_user.claims.sub) - .bind(&account_label) - .execute(&mut *tx) - .await?; - tx.commit().await?; - // Switching default flips which credential the resolver picks // when no API-key override is set. The no-override lane (`_`) // is now serving responses pinned to the *old* default's @@ -1146,8 +1063,8 @@ pub async fn paste_static_token( let access_encrypted = crypto::encrypt(req.token.as_bytes(), &enc_key) .map_err(|e| AppError::Internal(anyhow::anyhow!("encrypt token: {e}")))?; - upsert_credential( - &state, + credential_repo::upsert_user_credential( + &state.db, server_id, auth_user.claims.sub, account_label.trim(), @@ -1282,16 +1199,14 @@ pub async fn test_connection( // Confirm the credential exists before probing — saves a misleading // `NeedsUserCredentials` result for an account_label the user // never created (typo in the URL, stale UI cache, etc.). - let exists: Option = sqlx::query_scalar( - r#"SELECT 1 FROM mcp_user_credentials - WHERE mcp_server_id = $1 AND user_id = $2 AND account_label = $3"#, + let exists = credential_repo::user_account_exists( + &state.db, + server_id, + auth_user.claims.sub, + &account_label, ) - .bind(server_id) - .bind(auth_user.claims.sub) - .bind(&account_label) - .fetch_optional(&state.db) .await?; - if exists.is_none() { + if !exists { return Err(AppError::NotFound("Connection not found".into())); } @@ -1458,84 +1373,9 @@ pub(crate) async fn resolve_upstream_subject( // --------------------------------------------------------------------------- pub(super) async fn load_server(state: &AppState, server_id: Uuid) -> Result { - sqlx::query_as::<_, McpServer>( - r#"SELECT s.*, 0::bigint AS tools_count, 0::bigint AS call_count - FROM mcp_servers s WHERE s.id = $1"#, - ) - .bind(server_id) - .fetch_optional(&state.db) - .await? - .ok_or_else(|| AppError::NotFound("MCP server not found".into())) -} - -#[allow(clippy::too_many_arguments)] -pub(super) async fn upsert_credential( - state: &AppState, - server_id: Uuid, - user_id: Uuid, - account_label: &str, - credential_type: &str, - access_encrypted: &[u8], - refresh_encrypted: Option<&[u8]>, - expires_at: Option>, - scopes: &[String], - upstream_subject: Option<&str>, -) -> Result<(), AppError> { - // First credential for (server, user) becomes the default. - // SELECT-then-INSERT inside one tx is NOT enough on its own — - // two concurrent first-time inserts (admin opens authorize in two - // tabs, two account labels) would each read empty + each try - // is_default=true and the partial unique index - // `uq_mcp_user_credentials_default` would 23505 the loser into a - // user-facing 500. Take a per-(server, user) advisory lock so the - // decision is serialized. - let mut tx = state.db.begin().await?; - let lock_key = format!("mcp_user_default:{server_id}:{user_id}"); - sqlx::query("SELECT pg_advisory_xact_lock(hashtextextended($1, 0))") - .bind(&lock_key) - .execute(&mut *tx) - .await?; - let any_existing: Option = sqlx::query_scalar( - r#"SELECT 1 FROM mcp_user_credentials - WHERE mcp_server_id = $1 AND user_id = $2 LIMIT 1"#, - ) - .bind(server_id) - .bind(user_id) - .fetch_optional(&mut *tx) - .await?; - let new_default = any_existing.is_none(); - - sqlx::query( - r#"INSERT INTO mcp_user_credentials ( - mcp_server_id, user_id, account_label, credential_type, is_default, - access_token_encrypted, refresh_token_encrypted, - expires_at, scopes, upstream_subject - ) - VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10) - ON CONFLICT (mcp_server_id, user_id, account_label) DO UPDATE SET - credential_type = EXCLUDED.credential_type, - access_token_encrypted = EXCLUDED.access_token_encrypted, - refresh_token_encrypted = EXCLUDED.refresh_token_encrypted, - expires_at = EXCLUDED.expires_at, - scopes = EXCLUDED.scopes, - upstream_subject = EXCLUDED.upstream_subject, - updated_at = now()"#, - ) - .bind(server_id) - .bind(user_id) - .bind(account_label) - .bind(credential_type) - .bind(new_default) - .bind(access_encrypted) - .bind(refresh_encrypted) - .bind(expires_at) - .bind(scopes) - .bind(upstream_subject) - .execute(&mut *tx) - .await?; - - tx.commit().await?; - Ok(()) + server_repo::find_without_counts(&state.db, server_id) + .await? + .ok_or_else(|| AppError::NotFound("MCP server not found".into())) } // --------------------------------------------------------------------------- diff --git a/crates/server/src/handlers/mcp_oauth/shared.rs b/crates/server/src/handlers/mcp_oauth/shared.rs index 69e682ac..a82e9032 100644 --- a/crates/server/src/handlers/mcp_oauth/shared.rs +++ b/crates/server/src/handlers/mcp_oauth/shared.rs @@ -29,6 +29,8 @@ use think_watch_common::models::McpServer; use crate::app::AppState; use crate::middleware::auth_guard::AuthUser; +use crate::services::mcp_credential_repository as credential_repo; +use crate::services::mcp_server_repository as server_repo; use super::{ AuthorizeResponse, McpOauthState, OAUTH_STATE_PREFIX, OAUTH_STATE_TTL_SECS, OauthStateTarget, @@ -39,54 +41,6 @@ use super::{ // Admin: shared-credential storage // --------------------------------------------------------------------------- -/// UPSERT into `mcp_server_shared_credentials`. Single row per server -/// — when the admin rotates the credential the new row replaces the -/// previous one. Uses `INSERT … ON CONFLICT` keyed on the server_id -/// PK so the lifecycle code in [`UserTokenResolver`] sees a fresh -/// `(access_token_encrypted, expires_at)` after a rotation without -/// any extra coordination. -#[allow(clippy::too_many_arguments)] -pub(super) async fn upsert_shared_credential( - state: &AppState, - server_id: Uuid, - credential_type: &str, - access_encrypted: &[u8], - refresh_encrypted: Option<&[u8]>, - expires_at: Option>, - scopes: &[String], - upstream_subject: Option<&str>, - configured_by: Uuid, -) -> Result<(), AppError> { - sqlx::query( - r#"INSERT INTO mcp_server_shared_credentials ( - mcp_server_id, credential_type, - access_token_encrypted, refresh_token_encrypted, - expires_at, scopes, upstream_subject, configured_by - ) - VALUES ($1, $2, $3, $4, $5, $6, $7, $8) - ON CONFLICT (mcp_server_id) DO UPDATE SET - credential_type = EXCLUDED.credential_type, - access_token_encrypted = EXCLUDED.access_token_encrypted, - refresh_token_encrypted = EXCLUDED.refresh_token_encrypted, - expires_at = EXCLUDED.expires_at, - scopes = EXCLUDED.scopes, - upstream_subject = EXCLUDED.upstream_subject, - configured_by = EXCLUDED.configured_by, - updated_at = now()"#, - ) - .bind(server_id) - .bind(credential_type) - .bind(access_encrypted) - .bind(refresh_encrypted) - .bind(expires_at) - .bind(scopes) - .bind(upstream_subject) - .bind(configured_by) - .execute(&state.db) - .await?; - Ok(()) -} - /// Background tool-catalog refresh after a shared-credential write. /// Builds the auth header from the server's `auth_header_name` / /// `auth_value_template` so X-API-Key and other non-Bearer shapes @@ -117,21 +71,19 @@ pub(super) fn spawn_shared_tool_discovery( tools = n, "Shared-credential MCP tool discovery succeeded" ); - let _ = sqlx::query("UPDATE mcp_servers SET last_error = NULL WHERE id = $1") - .bind(server.id) - .execute(&db) - .await; + let _ = server_repo::clear_last_error(&db, server.id).await; } crate::mcp_runtime::SystemDiscoveryOutcome::AuthRequired => { tracing::warn!( mcp_server = %server.name, "Shared credential rejected by upstream tools/list (401/403)" ); - let _ = sqlx::query("UPDATE mcp_servers SET last_error = $1 WHERE id = $2") - .bind("Shared credential rejected by upstream — verify token / scopes") - .bind(server.id) - .execute(&db) - .await; + let _ = server_repo::set_last_error( + &db, + server.id, + "Shared credential rejected by upstream — verify token / scopes", + ) + .await; } crate::mcp_runtime::SystemDiscoveryOutcome::Failed(e) => { tracing::warn!( @@ -139,11 +91,7 @@ pub(super) fn spawn_shared_tool_discovery( error = %e, "Shared-credential MCP tool discovery failed" ); - let _ = sqlx::query("UPDATE mcp_servers SET last_error = $1 WHERE id = $2") - .bind(format!("{e}")) - .bind(server.id) - .execute(&db) - .await; + let _ = server_repo::set_last_error(&db, server.id, &format!("{e}")).await; } } }); @@ -187,8 +135,8 @@ pub async fn paste_shared_static_token( let access_encrypted = crypto::encrypt(req.token.as_bytes(), &enc_key) .map_err(|e| AppError::Internal(anyhow::anyhow!("encrypt token: {e}")))?; - upsert_shared_credential( - &state, + credential_repo::upsert_shared_credential( + &state.db, server_id, "static_token", &access_encrypted, @@ -345,21 +293,7 @@ pub async fn shared_credential_status( .require_global_permission(&state.db, "mcp_servers:read") .await?; - #[derive(sqlx::FromRow)] - struct Row { - credential_type: String, - expires_at: Option>, - upstream_subject: Option, - configured_by: Option, - updated_at: DateTime, - } - let row = sqlx::query_as::<_, Row>( - r#"SELECT credential_type, expires_at, upstream_subject, configured_by, updated_at - FROM mcp_server_shared_credentials WHERE mcp_server_id = $1"#, - ) - .bind(server_id) - .fetch_optional(&state.db) - .await?; + let row = credential_repo::find_shared_status(&state.db, server_id).await?; Ok(Json(match row { Some(r) => SharedCredentialStatus { @@ -401,10 +335,7 @@ pub async fn revoke_shared_credential( )); } - sqlx::query("DELETE FROM mcp_server_shared_credentials WHERE mcp_server_id = $1") - .bind(server_id) - .execute(&state.db) - .await?; + credential_repo::delete_shared_credential(&state.db, server_id).await?; // The shared bearer is gone — every cached response was minted // under it and is now serving against an identity that no longer @@ -438,13 +369,7 @@ pub async fn best_effort_revoke_shared_upstream( state: &AppState, server_id: Uuid, ) -> Result { - let row: Option<(String, Vec)> = sqlx::query_as( - r#"SELECT credential_type, access_token_encrypted - FROM mcp_server_shared_credentials WHERE mcp_server_id = $1"#, - ) - .bind(server_id) - .fetch_optional(&state.db) - .await?; + let row = credential_repo::find_shared_token(&state.db, server_id).await?; let Some((credential_type, access_encrypted)) = row else { return Ok(false); }; diff --git a/crates/server/src/handlers/mcp_oauth/wizard.rs b/crates/server/src/handlers/mcp_oauth/wizard.rs index 07145b8e..201f9bea 100644 --- a/crates/server/src/handlers/mcp_oauth/wizard.rs +++ b/crates/server/src/handlers/mcp_oauth/wizard.rs @@ -9,7 +9,7 @@ //! state blob and the resulting tokens land in Redis under //! `mcp_wizard:cred:{wizard_session_id}` instead of going straight //! to `mcp_server_shared_credentials`. The wizard's `Save` step -//! calls `claim_wizard_credential` + `insert_shared_credential_from_wizard` +//! calls `claim_wizard_credential` + `mcp_credential_repository::insert_shared_credential` //! to atomically promote the pending blob into the real //! shared-credential table at server-create time. @@ -311,32 +311,3 @@ pub struct PoppedWizardCredential { pub upstream_subject: Option, pub configured_by: Uuid, } - -/// Insert a popped wizard credential into `mcp_server_shared_credentials`. -/// Called by [`mcp_servers::create_server`] inside the same TX as the -/// server-row insert so the credential and the row land atomically. -pub async fn insert_shared_credential_from_wizard( - tx: &mut sqlx::Transaction<'_, sqlx::Postgres>, - server_id: Uuid, - cred: &PoppedWizardCredential, -) -> Result<(), AppError> { - sqlx::query( - r#"INSERT INTO mcp_server_shared_credentials ( - mcp_server_id, credential_type, - access_token_encrypted, refresh_token_encrypted, - expires_at, scopes, upstream_subject, configured_by - ) - VALUES ($1, $2, $3, $4, $5, $6, $7, $8)"#, - ) - .bind(server_id) - .bind(&cred.credential_type) - .bind(&cred.access_token_encrypted) - .bind(cred.refresh_token_encrypted.as_deref()) - .bind(cred.expires_at) - .bind(&cred.scopes) - .bind(cred.upstream_subject.as_deref()) - .bind(cred.configured_by) - .execute(&mut **tx) - .await?; - Ok(()) -} diff --git a/crates/server/src/handlers/mcp_servers.rs b/crates/server/src/handlers/mcp_servers.rs index bd283435..5bf34021 100644 --- a/crates/server/src/handlers/mcp_servers.rs +++ b/crates/server/src/handlers/mcp_servers.rs @@ -10,6 +10,9 @@ use think_watch_common::models::McpServer; use super::serde_util::deserialize_some; use crate::app::AppState; use crate::middleware::auth_guard::AuthUser; +use crate::services::mcp_credential_repository as credential_repo; +use crate::services::mcp_server_repository::{self as repo, McpServerFields}; +use crate::services::mcp_store_repository as store_repo; // `probe_mcp_endpoint`, `McpProbeOutcome`, `McpToolSummary`, and // `normalize_namespace_prefix` live in `super::mcp_shared` so @@ -17,18 +20,6 @@ use crate::middleware::auth_guard::AuthUser; // reaching across handlers. pub use super::mcp_shared::{McpToolSummary, normalize_namespace_prefix, probe_mcp_endpoint}; -/// Process-wide advisory-lock key for serializing template installs. -/// The literal spells "mcpStore" in ASCII so a DBA glancing at -/// `pg_locks` can tell what's holding it. Any new advisory lock -/// added elsewhere in the codebase MUST use a distinct constant — -/// collisions silently serialize unrelated work and can deadlock -/// under concurrent load. -/// -/// Reserved advisory lock keys (keep this list current): -/// * `MCP_STORE_INSTALL_LOCK_KEY` (here): template-install -/// serialization in `create_server` when `template_slug` is set. -const MCP_STORE_INSTALL_LOCK_KEY: i64 = 0x6D637053746F7265; - /// Find an available `(name, namespace_prefix)` pair by appending /// `_2`, `_3`, … when the base values are already taken. Runs inside /// the caller's tx so two concurrent installs of the same template @@ -36,7 +27,7 @@ const MCP_STORE_INSTALL_LOCK_KEY: i64 = 0x6D637053746F7265; /// path; non-template `create_server` calls just rely on UNIQUE to /// reject collisions and surface a 409 to the admin. async fn resolve_server_collisions( - tx: &mut sqlx::Transaction<'_, sqlx::Postgres>, + conn: &mut sqlx::PgConnection, base_name: &str, base_prefix: &str, ) -> Result<(String, String), AppError> { @@ -46,18 +37,7 @@ async fn resolve_server_collisions( } else { (format!("{base_name} #{i}"), format!("{base_prefix}_{i}")) }; - // `SELECT 1` is INT4 on the wire; binding into `Option` - // panics with a column-decode mismatch the moment a row - // comes back. We don't actually care about the value — only - // whether the row exists — so use Option. - let conflict: Option = sqlx::query_scalar( - "SELECT 1 FROM mcp_servers WHERE name = $1 OR namespace_prefix = $2 LIMIT 1", - ) - .bind(&n) - .bind(&p) - .fetch_optional(&mut **tx) - .await?; - if conflict.is_none() { + if !repo::name_or_prefix_taken(conn, &n, &p).await? { return Ok((n, p)); } } @@ -177,15 +157,7 @@ pub async fn list_servers( auth_user .require_global_permission(&state.db, "mcp_servers:read") .await?; - let mut servers = sqlx::query_as::<_, McpServer>( - r#"SELECT s.*, COALESCE(t.cnt, 0) AS tools_count - FROM mcp_servers s - LEFT JOIN (SELECT server_id, COUNT(*) AS cnt FROM mcp_tools WHERE is_active = true GROUP BY server_id) t - ON t.server_id = s.id - ORDER BY s.created_at DESC"#, - ) - .fetch_all(&state.db) - .await?; + let mut servers = repo::list_with_tool_counts(&state.db).await?; // Attach lifetime call counts from ClickHouse (mcp_logs) — best-effort: // if CH is unavailable we simply leave the counter at 0. @@ -501,16 +473,10 @@ pub async fn create_server( // window between snapshot fetch and INSERT. let (final_name, final_prefix, template_id) = match req.template_slug.as_deref() { Some(slug) if !slug.is_empty() => { - sqlx::query("SELECT pg_advisory_xact_lock($1)") - .bind(MCP_STORE_INSTALL_LOCK_KEY) - .execute(&mut *tx) - .await?; - let template_id: Uuid = - sqlx::query_scalar("SELECT id FROM mcp_store_templates WHERE slug = $1 FOR UPDATE") - .bind(slug) - .fetch_optional(&mut *tx) - .await? - .ok_or_else(|| AppError::NotFound(format!("Template '{slug}' not found")))?; + store_repo::lock_installs(&mut tx).await?; + let template_id: Uuid = store_repo::lock_template_by_slug(&mut tx, slug) + .await? + .ok_or_else(|| AppError::NotFound(format!("Template '{slug}' not found")))?; let (resolved_name, resolved_prefix) = resolve_server_collisions(&mut tx, &req.name, &namespace_prefix).await?; (resolved_name, resolved_prefix, Some(template_id)) @@ -527,78 +493,61 @@ pub async fn create_server( .filter(|s| !s.is_empty()) .map(String::from); - let server = sqlx::query_as::<_, McpServer>( - r#"INSERT INTO mcp_servers ( - name, namespace_prefix, display_label, description, endpoint_url, transport_type, - oauth_issuer, oauth_authorization_endpoint, oauth_token_endpoint, - oauth_revocation_endpoint, oauth_userinfo_endpoint, - oauth_client_id, oauth_client_secret_encrypted, - oauth_scopes, auth_shape, static_token_help_url, - auth_header_name, auth_value_template, credential_owner, - config_json - ) - VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, - $16, $17, $18, $19, $20) - RETURNING *"#, + let server = repo::insert( + &mut tx, + &McpServerFields { + name: &final_name, + namespace_prefix: &final_prefix, + display_label: display_label.as_deref(), + description: req.description.as_deref(), + endpoint_url: &req.endpoint_url, + transport_type: &transport_type, + oauth_issuer: req.oauth_issuer.as_deref(), + oauth_authorization_endpoint: req.oauth_authorization_endpoint.as_deref(), + oauth_token_endpoint: req.oauth_token_endpoint.as_deref(), + oauth_revocation_endpoint: req.oauth_revocation_endpoint.as_deref(), + oauth_userinfo_endpoint: req.oauth_userinfo_endpoint.as_deref(), + oauth_client_id: req.oauth_client_id.as_deref(), + oauth_client_secret_encrypted: oauth_client_secret_encrypted.as_deref(), + oauth_scopes: &oauth_scopes, + auth_shape: &auth_shape, + static_token_help_url: req.static_token_help_url.as_deref(), + auth_header_name: &auth_header_name, + auth_value_template: &auth_value_template, + credential_owner: &credential_owner, + config_json: &config_json, + }, ) - .bind(&final_name) - .bind(&final_prefix) - .bind(&display_label) - .bind(&req.description) - .bind(&req.endpoint_url) - .bind(&transport_type) - .bind(&req.oauth_issuer) - .bind(&req.oauth_authorization_endpoint) - .bind(&req.oauth_token_endpoint) - .bind(&req.oauth_revocation_endpoint) - .bind(&req.oauth_userinfo_endpoint) - .bind(&req.oauth_client_id) - .bind(&oauth_client_secret_encrypted) - .bind(&oauth_scopes) - .bind(&auth_shape) - .bind(&req.static_token_help_url) - .bind(&auth_header_name) - .bind(&auth_value_template) - .bind(&credential_owner) - .bind(&config_json) - .fetch_one(&mut *tx) - .await - .map_err(map_mcp_server_unique_violation)?; + .await?; // Template install audit row + install_count bump. Same TX as // the server INSERT so the count never drifts even if // mcp_store_installs FK violations rollback the whole thing. if let Some(tid) = template_id { - sqlx::query( - "INSERT INTO mcp_store_installs (template_id, server_id, installed_by) VALUES ($1, $2, $3)", - ) - .bind(tid) - .bind(server.id) - .bind(auth_user.claims.sub) - .execute(&mut *tx) - .await?; - sqlx::query( - "UPDATE mcp_store_templates SET install_count = install_count + 1 WHERE id = $1", - ) - .bind(tid) - .execute(&mut *tx) - .await?; + store_repo::record_install(&mut tx, tid, server.id, auth_user.claims.sub).await?; } // Atomic credential install for admin_shared mode. if let Some(cred) = &wizard_cred { - super::mcp_oauth::insert_shared_credential_from_wizard(&mut tx, server.id, cred).await?; + credential_repo::insert_shared_credential( + &mut tx, + server.id, + &cred.credential_type, + &cred.access_token_encrypted, + cred.refresh_token_encrypted.as_deref(), + cred.expires_at, + &cred.scopes, + cred.upstream_subject.as_deref(), + cred.configured_by, + ) + .await?; } else if let Some(encrypted) = &shared_static_token_encrypted { - sqlx::query( - r#"INSERT INTO mcp_server_shared_credentials ( - mcp_server_id, credential_type, access_token_encrypted, configured_by - ) - VALUES ($1, 'static_token', $2, $3)"#, + credential_repo::insert_shared_static_token( + &mut tx, + server.id, + encrypted, + auth_user.claims.sub, ) - .bind(server.id) - .bind(encrypted) - .bind(auth_user.claims.sub) - .execute(&mut *tx) .await?; } @@ -684,10 +633,7 @@ pub async fn create_server( tools = n, "MCP tool discovery completed for new server" ); - let _ = sqlx::query("UPDATE mcp_servers SET last_error = NULL WHERE id = $1") - .bind(server_id) - .execute(&db_for_err) - .await; + let _ = repo::clear_last_error(&db_for_err, server_id).await; if let Ok(updated) = crate::mcp_runtime::build_registered_server(&db, &server, &key).await { @@ -698,10 +644,7 @@ pub async fn create_server( // Server requires per-user auth — `mcp_tools` stays // empty by design. Clear last_error so the admin UI // doesn't show stale failure text. - let _ = sqlx::query("UPDATE mcp_servers SET last_error = NULL WHERE id = $1") - .bind(server_id) - .execute(&db_for_err) - .await; + let _ = repo::clear_last_error(&db_for_err, server_id).await; } SystemDiscoveryOutcome::Failed(e) => { tracing::warn!( @@ -709,11 +652,7 @@ pub async fn create_server( error = %e, "Initial MCP tool discovery failed" ); - let _ = sqlx::query("UPDATE mcp_servers SET last_error = $1 WHERE id = $2") - .bind(format!("{e}")) - .bind(server_id) - .execute(&db_for_err) - .await; + let _ = repo::set_last_error(&db_for_err, server_id, &format!("{e}")).await; } } }); @@ -823,9 +762,7 @@ pub async fn update_server( auth_user .require_global_permission(&state.db, "mcp_servers:update") .await?; - let existing = sqlx::query_as::<_, McpServer>("SELECT * FROM mcp_servers WHERE id = $1") - .bind(id) - .fetch_optional(&state.db) + let existing = repo::find(&state.db, id) .await? .ok_or(AppError::NotFound("MCP Server not found".into()))?; @@ -1015,74 +952,48 @@ pub async fn update_server( // credential_owner with old-shape credentials still attached — // the resolver would then mismatch. Wrapping both in a TX makes // the transition atomic. - let mut tx = state.db.begin().await?; - let updated = sqlx::query_as::<_, McpServer>( - r#"UPDATE mcp_servers SET - name = $2, namespace_prefix = $3, display_label = $4, - description = $5, endpoint_url = $6, - transport_type = $7, - oauth_issuer = $8, oauth_authorization_endpoint = $9, - oauth_token_endpoint = $10, oauth_revocation_endpoint = $11, - oauth_userinfo_endpoint = $12, - oauth_client_id = $13, oauth_client_secret_encrypted = $14, - oauth_scopes = $15, auth_shape = $16, static_token_help_url = $17, - auth_header_name = $18, auth_value_template = $19, credential_owner = $20, - config_json = $21 - WHERE id = $1 RETURNING *"#, - ) - .bind(id) - .bind(name) - .bind(&namespace_prefix) - .bind(display_label) - .bind(description) - .bind(endpoint_url) - .bind(transport_type) - .bind(oauth_issuer) - .bind(oauth_authorization_endpoint) - .bind(oauth_token_endpoint) - .bind(oauth_revocation_endpoint) - .bind(oauth_userinfo_endpoint) - .bind(oauth_client_id) - .bind(&oauth_client_secret_encrypted) - .bind(&oauth_scopes) - .bind(&auth_shape) - .bind(static_token_help_url) - .bind(&auth_header_name) - .bind(&auth_value_template) - .bind(&credential_owner) - .bind(&config_json) - .fetch_one(&mut *tx) - .await - .map_err(map_mcp_server_unique_violation)?; - + // // Credential cleanup on relevant transitions. Switching to // admin_shared makes per-user creds dead weight; flipping the // auth_shape (oauth ↔ static, or either ↔ anonymous) makes the // *previous shape's* tokens incompatible with the new resolver // path. Both cases purge per-user + shared rows for the server // so callers don't end up holding mismatched credentials. - if switching_to_admin_shared || auth_shape_changed { - sqlx::query("DELETE FROM mcp_user_credentials WHERE mcp_server_id = $1") - .bind(id) - .execute(&mut *tx) - .await?; - sqlx::query("DELETE FROM mcp_user_tools WHERE mcp_server_id = $1") - .bind(id) - .execute(&mut *tx) - .await?; - } - // Drop the shared-credential row when *either* the auth_shape + // + // The shared-credential row goes when *either* the auth_shape // changed (old token is wrong shape) OR we left admin_shared // entirely. Same DELETE either way; collapsing the two // conditions avoids running it twice on a combined transition // (e.g. admin_shared/oauth → per_user/static). - if auth_shape_changed || switching_off_admin_shared { - sqlx::query("DELETE FROM mcp_server_shared_credentials WHERE mcp_server_id = $1") - .bind(id) - .execute(&mut *tx) - .await?; - } - tx.commit().await?; + let updated = repo::update( + &state.db, + id, + &McpServerFields { + name, + namespace_prefix: &namespace_prefix, + display_label, + description, + endpoint_url, + transport_type: &transport_type, + oauth_issuer, + oauth_authorization_endpoint, + oauth_token_endpoint, + oauth_revocation_endpoint, + oauth_userinfo_endpoint, + oauth_client_id, + oauth_client_secret_encrypted: oauth_client_secret_encrypted.as_deref(), + oauth_scopes: &oauth_scopes, + auth_shape: &auth_shape, + static_token_help_url, + auth_header_name: &auth_header_name, + auth_value_template: &auth_value_template, + credential_owner: &credential_owner, + config_json: &config_json, + }, + switching_to_admin_shared || auth_shape_changed, + auth_shape_changed || switching_off_admin_shared, + ) + .await?; // Evict any cached connection first — the pool keys by id, so a // changed endpoint URL needs a fresh connection. @@ -1190,9 +1101,7 @@ pub async fn get_server( auth_user .require_global_permission(&state.db, "mcp_servers:read") .await?; - let server = sqlx::query_as::<_, McpServer>("SELECT * FROM mcp_servers WHERE id = $1") - .bind(id) - .fetch_optional(&state.db) + let server = repo::find(&state.db, id) .await? .ok_or(AppError::NotFound("MCP Server not found".into()))?; @@ -1223,11 +1132,9 @@ pub async fn delete_server( .require_global_permission(&state.db, "mcp_servers:delete") .await?; - let mut tx = state.db.begin().await?; - let name = delete_server_inner(&mut tx, id) + let name = repo::delete(&state.db, id) .await? .ok_or_else(|| AppError::NotFound("MCP Server not found".into()))?; - tx.commit().await?; // Drop from the in-memory registry and connection pool — otherwise the // gateway would keep a stale entry for a server that no longer exists @@ -1246,47 +1153,6 @@ pub async fn delete_server( Ok(Json(serde_json::json!({"status": "deleted"}))) } -/// Tear down a single MCP server inside the caller's transaction. -/// Performs the same DB-side work as [`delete_server`]: -/// * SELECT the server name (returned to the caller for audit detail) -/// * decrement the originating store template's `install_count` -/// * DELETE the server row (children CASCADE: `mcp_tools`, -/// `mcp_user_credentials`, `mcp_server_shared_credentials`, -/// `mcp_user_tools`, `mcp_store_installs`) -/// -/// Returns `Ok(Some(name))` on success, `Ok(None)` if the row doesn't -/// exist (caller maps that to a "not_found" skip). In-memory registry -/// / connection-pool eviction happens at the call site, *after* the -/// TX commits, so a rolled-back batch never desyncs the registry. -pub(super) async fn delete_server_inner( - tx: &mut sqlx::Transaction<'_, sqlx::Postgres>, - id: Uuid, -) -> Result, AppError> { - let name: Option = sqlx::query_scalar("SELECT name FROM mcp_servers WHERE id = $1") - .bind(id) - .fetch_optional(&mut **tx) - .await?; - if name.is_none() { - return Ok(None); - } - - // Decrement install_count if this server was installed from the store. - sqlx::query( - r#"UPDATE mcp_store_templates SET install_count = GREATEST(install_count - 1, 0) - WHERE id = (SELECT template_id FROM mcp_store_installs WHERE server_id = $1)"#, - ) - .bind(id) - .execute(&mut **tx) - .await?; - - sqlx::query("DELETE FROM mcp_servers WHERE id = $1") - .bind(id) - .execute(&mut **tx) - .await?; - - Ok(name) -} - /// Hard cap on `POST /api/mcp/servers/bulk-delete` batch size. Picked /// to keep the worst-case transaction short — every id triggers a /// SELECT + UPDATE + DELETE plus CASCADE work on @@ -1368,11 +1234,10 @@ pub async fn bulk_delete_servers( // separate TXs would let a mid-batch failure leave the DB in a // half-deleted state, which is exactly the footgun bulk-delete // is meant to avoid. - let mut tx = state.db.begin().await?; let mut deleted_pairs: Vec<(Uuid, String)> = Vec::new(); let mut skipped: Vec = Vec::new(); - for id in unique_ids { - match delete_server_inner(&mut tx, id).await? { + for (id, name) in repo::delete_many(&state.db, &unique_ids).await? { + match name { Some(name) => deleted_pairs.push((id, name)), None => skipped.push(BulkDeleteSkip { id, @@ -1380,7 +1245,6 @@ pub async fn bulk_delete_servers( }), } } - tx.commit().await?; // Post-commit cleanup + audit. Done outside the TX so an audit // emit that briefly blocks on the forwarder pool can't roll back @@ -1404,25 +1268,6 @@ pub async fn bulk_delete_servers( })) } -/// Translate PostgreSQL unique-constraint violations on `mcp_servers` into -/// user-facing conflict errors, so the UI shows "already in use" instead of -/// a generic 500. Other sqlx errors fall through unchanged. -fn map_mcp_server_unique_violation(e: sqlx::Error) -> AppError { - if let sqlx::Error::Database(db_err) = &e - && db_err.code().as_deref() == Some("23505") - { - let constraint = db_err.constraint().unwrap_or(""); - if constraint.contains("namespace_prefix") { - return AppError::Conflict("namespace_prefix already in use".into()); - } - if constraint.contains("name") { - return AppError::Conflict("server name already in use".into()); - } - return AppError::Conflict("duplicate server".into()); - } - AppError::from(e) -} - #[cfg(test)] mod tests { use super::*; diff --git a/crates/server/src/handlers/mcp_store.rs b/crates/server/src/handlers/mcp_store.rs index 639b9354..b57efbeb 100644 --- a/crates/server/src/handlers/mcp_store.rs +++ b/crates/server/src/handlers/mcp_store.rs @@ -8,6 +8,7 @@ use think_watch_common::models::McpStoreTemplate; use crate::app::AppState; use crate::middleware::auth_guard::AuthUser; +use crate::services::mcp_store_repository::{self as repo, TemplateUpsert}; // --------------------------------------------------------------------------- // DTOs @@ -43,19 +44,13 @@ pub async fn list_templates( Query(q): Query, ) -> Result>, AppError> { // Fetch all installed template IDs for this instance - let installed_ids: Vec = sqlx::query_scalar("SELECT template_id FROM mcp_store_installs") - .fetch_all(&state.db) - .await?; + let installed_ids = repo::installed_template_ids(&state.db).await?; let installed_set: std::collections::HashSet = installed_ids.into_iter().collect(); // Fetch all templates and filter in Rust — the store catalog is small // enough that dynamic SQL bind complexity isn't worth it. - let templates = sqlx::query_as::<_, McpStoreTemplate>( - "SELECT * FROM mcp_store_templates ORDER BY featured DESC, install_count DESC, name ASC", - ) - .fetch_all(&state.db) - .await?; + let templates = repo::list_templates(&state.db).await?; let results: Vec = templates .into_iter() @@ -110,12 +105,9 @@ pub async fn get_template( State(state): State, Path(slug): Path, ) -> Result, AppError> { - let template = - sqlx::query_as::<_, McpStoreTemplate>("SELECT * FROM mcp_store_templates WHERE slug = $1") - .bind(&slug) - .fetch_optional(&state.db) - .await? - .ok_or_else(|| AppError::NotFound(format!("Template '{slug}' not found")))?; + let template = repo::find_template_by_slug(&state.db, &slug) + .await? + .ok_or_else(|| AppError::NotFound(format!("Template '{slug}' not found")))?; Ok(Json(template)) } @@ -127,17 +119,7 @@ pub async fn list_categories( _auth_user: AuthUser, State(state): State, ) -> Result>, AppError> { - #[derive(sqlx::FromRow)] - struct Row { - category: Option, - count: Option, - } - - let rows = sqlx::query_as::<_, Row>( - "SELECT category, COUNT(*) as count FROM mcp_store_templates GROUP BY category ORDER BY count DESC", - ) - .fetch_all(&state.db) - .await?; + let rows = repo::category_counts(&state.db).await?; let categories = rows .into_iter() @@ -339,90 +321,43 @@ pub async fn sync_registry( _ => "anonymous".to_string(), }; - sqlx::query( - r#"INSERT INTO mcp_store_templates - (slug, name, description, category, tags, endpoint_template, - oauth_issuer, oauth_authorization_endpoint, oauth_token_endpoint, - oauth_revocation_endpoint, oauth_userinfo_endpoint, - oauth_default_scopes, - auth_shape, static_token_help_url, - auth_header_name, auth_value_template, - auth_instructions, deploy_type, - deploy_command, deploy_docs_url, homepage_url, repo_url, featured, updated_at) - VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, - $16, $17, $18, $19, $20, $21, $22, $23, now()) - ON CONFLICT (slug) DO UPDATE SET - name = EXCLUDED.name, - description = EXCLUDED.description, - category = EXCLUDED.category, - tags = EXCLUDED.tags, - endpoint_template = EXCLUDED.endpoint_template, - oauth_issuer = EXCLUDED.oauth_issuer, - oauth_authorization_endpoint = EXCLUDED.oauth_authorization_endpoint, - oauth_token_endpoint = EXCLUDED.oauth_token_endpoint, - oauth_revocation_endpoint = EXCLUDED.oauth_revocation_endpoint, - oauth_userinfo_endpoint = EXCLUDED.oauth_userinfo_endpoint, - oauth_default_scopes = EXCLUDED.oauth_default_scopes, - auth_shape = EXCLUDED.auth_shape, - static_token_help_url = EXCLUDED.static_token_help_url, - auth_header_name = EXCLUDED.auth_header_name, - auth_value_template = EXCLUDED.auth_value_template, - auth_instructions = EXCLUDED.auth_instructions, - deploy_type = EXCLUDED.deploy_type, - deploy_command = EXCLUDED.deploy_command, - deploy_docs_url = EXCLUDED.deploy_docs_url, - homepage_url = EXCLUDED.homepage_url, - repo_url = EXCLUDED.repo_url, - featured = EXCLUDED.featured, - updated_at = now()"#, - ) - .bind(&t.slug) - .bind(&t.name) - .bind(t.description.as_ref().and_then(flatten_i18n).as_deref()) - .bind(&t.category) - .bind(t.tags.as_deref().unwrap_or(&[])) - .bind(&t.endpoint_template) - .bind(&t.oauth_issuer) - .bind(&t.oauth_authorization_endpoint) - .bind(&t.oauth_token_endpoint) - .bind(&t.oauth_revocation_endpoint) - .bind(&t.oauth_userinfo_endpoint) - .bind(t.oauth_default_scopes.as_deref().unwrap_or(&[])) - .bind(&auth_shape) - .bind(&t.static_token_help_url) - .bind(&auth_header_name) - .bind(&auth_value_template) - .bind( - t.auth_instructions - .as_ref() - .and_then(flatten_i18n) - .as_deref(), + let description = t.description.as_ref().and_then(flatten_i18n); + let auth_instructions = t.auth_instructions.as_ref().and_then(flatten_i18n); + repo::upsert_template( + &mut tx, + &TemplateUpsert { + slug: &t.slug, + name: &t.name, + description: description.as_deref(), + category: t.category.as_deref(), + tags: t.tags.as_deref().unwrap_or(&[]), + endpoint_template: t.endpoint_template.as_deref(), + oauth_issuer: t.oauth_issuer.as_deref(), + oauth_authorization_endpoint: t.oauth_authorization_endpoint.as_deref(), + oauth_token_endpoint: t.oauth_token_endpoint.as_deref(), + oauth_revocation_endpoint: t.oauth_revocation_endpoint.as_deref(), + oauth_userinfo_endpoint: t.oauth_userinfo_endpoint.as_deref(), + oauth_default_scopes: t.oauth_default_scopes.as_deref().unwrap_or(&[]), + auth_shape: &auth_shape, + static_token_help_url: t.static_token_help_url.as_deref(), + auth_header_name: &auth_header_name, + auth_value_template: &auth_value_template, + auth_instructions: auth_instructions.as_deref(), + deploy_type: t.deploy_type.as_deref().unwrap_or("hosted"), + deploy_command: t.deploy_command.as_deref(), + deploy_docs_url: t.deploy_docs_url.as_deref(), + homepage_url: t.homepage_url.as_deref(), + repo_url: t.repo_url.as_deref(), + featured: t.featured.unwrap_or(false), + }, ) - .bind(t.deploy_type.as_deref().unwrap_or("hosted")) - .bind(&t.deploy_command) - .bind(&t.deploy_docs_url) - .bind(&t.homepage_url) - .bind(&t.repo_url) - .bind(t.featured.unwrap_or(false)) - .execute(&mut *tx) .await?; synced += 1; } // Remove templates that are no longer in the registry (but keep those with active installs) let registry_slugs: Vec<&str> = registry.templates.iter().map(|t| t.slug.as_str()).collect(); - let removed = sqlx::query_scalar::<_, i64>( - r#"WITH deleted AS ( - DELETE FROM mcp_store_templates - WHERE slug != ALL($1) - AND id NOT IN (SELECT template_id FROM mcp_store_installs) - RETURNING 1 - ) - SELECT COUNT(*) FROM deleted"#, - ) - .bind(®istry_slugs) - .fetch_one(&mut *tx) - .await?; + let removed = repo::delete_templates_not_in(&mut tx, ®istry_slugs).await?; tx.commit().await?; state.audit.log( diff --git a/crates/server/src/handlers/mcp_tools.rs b/crates/server/src/handlers/mcp_tools.rs index 92ff372e..dac1e176 100644 --- a/crates/server/src/handlers/mcp_tools.rs +++ b/crates/server/src/handlers/mcp_tools.rs @@ -1,26 +1,12 @@ use axum::Json; use axum::extract::{Query, State}; use serde::{Deserialize, Serialize}; -use sqlx::FromRow; use think_watch_common::errors::AppError; use crate::app::AppState; use crate::middleware::auth_guard::AuthUser; - -#[derive(Debug, Serialize, FromRow, utoipa::ToSchema)] -pub struct McpToolRow { - #[schema(value_type = String, format = Uuid)] - pub id: uuid::Uuid, - #[schema(value_type = String, format = Uuid)] - pub server_id: uuid::Uuid, - pub server_name: String, - pub name: String, - pub namespaced_name: String, - pub description: Option, - #[schema(value_type = Object)] - pub input_schema: Option, -} +use crate::services::mcp_tool_repository::{self as repo, CatalogQuery, McpToolRow}; #[derive(Debug, Deserialize)] pub struct McpToolListQuery { @@ -83,97 +69,16 @@ pub async fn list_tools( None }; - // Pre-namespace the per-user catalog the same way mcp_tools does - // (`__`) and union the two sources. `mcp_user_tools` - // doesn't carry an `id` column — synthesize a stable v5-style UUID - // from `(server_id, user_id, tool_name)` so the frontend's keying - // (`tool.id`) keeps working without a schema change. - let total: i64 = sqlx::query_scalar( - r#"WITH catalog AS ( - SELECT t.id, - t.server_id, - s.name AS server_name, - s.namespace_prefix, - t.tool_name, - t.description - FROM mcp_tools t - JOIN mcp_servers s ON s.id = t.server_id - WHERE t.is_active = true - UNION ALL - SELECT gen_random_uuid() AS id, - u.mcp_server_id AS server_id, - s.name AS server_name, - s.namespace_prefix, - u.tool_name, - u.description - FROM mcp_user_tools u - JOIN mcp_servers s ON s.id = u.mcp_server_id - WHERE $6::uuid IS NOT NULL AND u.user_id = $6::uuid - ) - SELECT COUNT(*) FROM catalog - WHERE ($3::uuid IS NULL OR server_id = $3) - AND ($1 = '' - OR tool_name ILIKE $2 - OR (namespace_prefix || '__' || tool_name) ILIKE $2 - OR COALESCE(description, '') ILIKE $2)"#, - ) - .bind(search) - .bind(&search_pattern) - .bind(query.server_id) - .bind(page_size) - .bind(offset) - .bind(user_filter) - .fetch_one(&state.db) - .await?; - - let items = sqlx::query_as::<_, McpToolRow>( - r#"WITH catalog AS ( - SELECT t.id, - t.server_id, - s.name AS server_name, - s.namespace_prefix, - t.tool_name, - t.description, - t.input_schema - FROM mcp_tools t - JOIN mcp_servers s ON s.id = t.server_id - WHERE t.is_active = true - UNION ALL - SELECT gen_random_uuid() AS id, - u.mcp_server_id AS server_id, - s.name AS server_name, - s.namespace_prefix, - u.tool_name, - u.description, - u.input_schema - FROM mcp_user_tools u - JOIN mcp_servers s ON s.id = u.mcp_server_id - WHERE $6::uuid IS NOT NULL AND u.user_id = $6::uuid - ) - SELECT id, - server_id, - server_name, - tool_name AS name, - namespace_prefix || '__' || tool_name AS namespaced_name, - description, - input_schema - FROM catalog - WHERE ($3::uuid IS NULL OR server_id = $3) - AND ($1 = '' - OR tool_name ILIKE $2 - OR (namespace_prefix || '__' || tool_name) ILIKE $2 - OR COALESCE(description, '') ILIKE $2) - ORDER BY server_name, tool_name - LIMIT $4 OFFSET $5"#, - ) - .bind(search) - .bind(&search_pattern) - .bind(query.server_id) - .bind(page_size) - .bind(offset) - .bind(user_filter) - .fetch_all(&state.db) - .await?; + let catalog = CatalogQuery { + search, + search_pattern: &search_pattern, + server_id: query.server_id, + page_size, + offset, + user_id: user_filter, + }; + let total = repo::count_catalog(&state.db, &catalog).await?; + let items = repo::list_catalog(&state.db, &catalog).await?; Ok(Json(McpToolListResponse { items, total })) } @@ -203,13 +108,9 @@ pub async fn discover_tools( axum::extract::Path(server_id): axum::extract::Path, ) -> Result, AppError> { auth_user.require_permission("mcp_servers:update")?; - let server = sqlx::query_as::<_, think_watch_common::models::McpServer>( - "SELECT * FROM mcp_servers WHERE id = $1", - ) - .bind(server_id) - .fetch_optional(&state.db) - .await? - .ok_or(AppError::NotFound("MCP Server not found".into()))?; + let server = crate::services::mcp_server_repository::find(&state.db, server_id) + .await? + .ok_or(AppError::NotFound("MCP Server not found".into()))?; use crate::mcp_runtime::SystemDiscoveryOutcome; let http = state.http_client.load(); diff --git a/crates/server/src/openapi.rs b/crates/server/src/openapi.rs index 7768cb6b..fcacd92b 100644 --- a/crates/server/src/openapi.rs +++ b/crates/server/src/openapi.rs @@ -27,7 +27,7 @@ use crate::handlers::{ BulkDeleteMcpServersRequest, BulkDeleteMcpServersResponse, BulkDeleteSkip, UpdateMcpServerRequest, }, - mcp_tools::{McpToolListResponse, McpToolRow}, + mcp_tools::McpToolListResponse, models::{ BatchWeightUpdate, BatchWeightsRequest, CreateModelRequest, RouteHistoryBucket, RouteHistoryResponse, UpdateModelRequest, @@ -46,6 +46,7 @@ use crate::handlers::{ ResetCounterResponse, UsageDay, }, }; +use crate::services::mcp_tool_repository::McpToolRow; use crate::services::model_repository::ModelRow; /// OpenAPI document covering the ThinkWatch console API (port 3001). diff --git a/crates/server/src/services/mcp_credential_repository.rs b/crates/server/src/services/mcp_credential_repository.rs new file mode 100644 index 00000000..f4217e87 --- /dev/null +++ b/crates/server/src/services/mcp_credential_repository.rs @@ -0,0 +1,403 @@ +//! MCP credential repository — `mcp_user_credentials` (one row per user +//! account on a per-user server) and `mcp_server_shared_credentials` +//! (the single admin-supplied credential of an admin-shared server). +//! +//! Tokens arrive and leave encrypted; encryption, upstream revocation, +//! cache invalidation and audit stay in `handlers::mcp_oauth`. + +use chrono::{DateTime, Utc}; +use sqlx::{PgConnection, PgPool}; +use think_watch_common::errors::AppError; +use uuid::Uuid; + +/// One of the caller's accounts, as listed on `/connections`. +#[derive(sqlx::FromRow)] +pub struct UserCredentialAccountRow { + pub mcp_server_id: Uuid, + pub account_label: String, + pub credential_type: String, + pub is_default: bool, + pub scopes: Vec, + pub expires_at: Option>, + pub upstream_subject: Option, + pub created_at: DateTime, + pub updated_at: DateTime, +} + +/// What the admin UI shows about a server's shared credential. +#[derive(sqlx::FromRow)] +pub struct SharedCredentialStatusRow { + pub credential_type: String, + pub expires_at: Option>, + pub upstream_subject: Option, + pub configured_by: Option, + pub updated_at: DateTime, +} + +// --------------------------------------------------------------------------- +// mcp_user_credentials +// --------------------------------------------------------------------------- + +/// Every account a user holds, grouped by server, default first. +pub async fn list_user_accounts( + pool: &PgPool, + user_id: Uuid, +) -> Result, AppError> { + Ok(sqlx::query_as::<_, UserCredentialAccountRow>( + r#"SELECT mcp_server_id, account_label, credential_type, is_default, + scopes, expires_at, upstream_subject, created_at, updated_at + FROM mcp_user_credentials + WHERE user_id = $1 + ORDER BY mcp_server_id, is_default DESC, account_label"#, + ) + .bind(user_id) + .fetch_all(pool) + .await?) +} + +/// An account's credential type and encrypted access token. +pub async fn find_user_token( + pool: &PgPool, + server_id: Uuid, + user_id: Uuid, + account_label: &str, +) -> Result)>, AppError> { + Ok(sqlx::query_as( + r#"SELECT credential_type, access_token_encrypted + FROM mcp_user_credentials + WHERE mcp_server_id = $1 AND user_id = $2 AND account_label = $3"#, + ) + .bind(server_id) + .bind(user_id) + .bind(account_label) + .fetch_optional(pool) + .await?) +} + +pub async fn user_account_exists( + pool: &PgPool, + server_id: Uuid, + user_id: Uuid, + account_label: &str, +) -> Result { + let exists: Option = sqlx::query_scalar( + r#"SELECT 1 FROM mcp_user_credentials + WHERE mcp_server_id = $1 AND user_id = $2 AND account_label = $3"#, + ) + .bind(server_id) + .bind(user_id) + .bind(account_label) + .fetch_optional(pool) + .await?; + Ok(exists.is_some()) +} + +/// Store an account's credential, replacing the one under the same +/// label. The first account a user holds on a server becomes the +/// default. +#[allow(clippy::too_many_arguments)] +pub async fn upsert_user_credential( + pool: &PgPool, + server_id: Uuid, + user_id: Uuid, + account_label: &str, + credential_type: &str, + access_encrypted: &[u8], + refresh_encrypted: Option<&[u8]>, + expires_at: Option>, + scopes: &[String], + upstream_subject: Option<&str>, +) -> Result<(), AppError> { + // First credential for (server, user) becomes the default. + // SELECT-then-INSERT inside one tx is NOT enough on its own — + // two concurrent first-time inserts (admin opens authorize in two + // tabs, two account labels) would each read empty + each try + // is_default=true and the partial unique index + // `uq_mcp_user_credentials_default` would 23505 the loser into a + // user-facing 500. Take a per-(server, user) advisory lock so the + // decision is serialized. + let mut tx = pool.begin().await?; + let lock_key = format!("mcp_user_default:{server_id}:{user_id}"); + sqlx::query("SELECT pg_advisory_xact_lock(hashtextextended($1, 0))") + .bind(&lock_key) + .execute(&mut *tx) + .await?; + let any_existing: Option = sqlx::query_scalar( + r#"SELECT 1 FROM mcp_user_credentials + WHERE mcp_server_id = $1 AND user_id = $2 LIMIT 1"#, + ) + .bind(server_id) + .bind(user_id) + .fetch_optional(&mut *tx) + .await?; + let new_default = any_existing.is_none(); + + sqlx::query( + r#"INSERT INTO mcp_user_credentials ( + mcp_server_id, user_id, account_label, credential_type, is_default, + access_token_encrypted, refresh_token_encrypted, + expires_at, scopes, upstream_subject + ) + VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10) + ON CONFLICT (mcp_server_id, user_id, account_label) DO UPDATE SET + credential_type = EXCLUDED.credential_type, + access_token_encrypted = EXCLUDED.access_token_encrypted, + refresh_token_encrypted = EXCLUDED.refresh_token_encrypted, + expires_at = EXCLUDED.expires_at, + scopes = EXCLUDED.scopes, + upstream_subject = EXCLUDED.upstream_subject, + updated_at = now()"#, + ) + .bind(server_id) + .bind(user_id) + .bind(account_label) + .bind(credential_type) + .bind(new_default) + .bind(access_encrypted) + .bind(refresh_encrypted) + .bind(expires_at) + .bind(scopes) + .bind(upstream_subject) + .execute(&mut *tx) + .await?; + + tx.commit().await?; + Ok(()) +} + +/// Delete an account and, if it was the default, promote the user's +/// newest remaining account on that server, in one transaction. +pub async fn delete_user_credential( + pool: &PgPool, + server_id: Uuid, + user_id: Uuid, + account_label: &str, +) -> Result<(), AppError> { + let mut tx = pool.begin().await?; + let was_default: Option = sqlx::query_scalar( + r#"DELETE FROM mcp_user_credentials + WHERE mcp_server_id = $1 AND user_id = $2 AND account_label = $3 + RETURNING is_default"#, + ) + .bind(server_id) + .bind(user_id) + .bind(account_label) + .fetch_optional(&mut *tx) + .await?; + + if matches!(was_default, Some(true)) { + // Promote the newest remaining credential for the same + // (server, user). Newest wins because a user juggling + // multiple credentials usually treats the latest one as + // "current" — same heuristic the connect-then-overwrite UX + // already nudges them toward. NULL `created_at` shouldn't + // exist (column is NOT NULL DEFAULT now()) but the ORDER BY + // is still safe under NULLS LAST. + sqlx::query( + r#"UPDATE mcp_user_credentials + SET is_default = true + WHERE id = ( + SELECT id FROM mcp_user_credentials + WHERE mcp_server_id = $1 AND user_id = $2 + ORDER BY created_at DESC NULLS LAST + LIMIT 1 + )"#, + ) + .bind(server_id) + .bind(user_id) + .execute(&mut *tx) + .await?; + } + tx.commit().await?; + Ok(()) +} + +/// Make an account the user's default on its server. Returns `false` +/// (and changes nothing) when the account doesn't exist. +pub async fn set_default_user_credential( + pool: &PgPool, + server_id: Uuid, + user_id: Uuid, + account_label: &str, +) -> Result { + let mut tx = pool.begin().await?; + let exists: Option = sqlx::query_scalar( + r#"SELECT 1 FROM mcp_user_credentials + WHERE mcp_server_id = $1 AND user_id = $2 AND account_label = $3"#, + ) + .bind(server_id) + .bind(user_id) + .bind(account_label) + .fetch_optional(&mut *tx) + .await?; + if exists.is_none() { + return Ok(false); + } + + // Two-step toggle so the partial unique index never sees two + // is_default rows at once: clear the old default first, then mark + // the new one inside the same transaction. + sqlx::query( + r#"UPDATE mcp_user_credentials SET is_default = false, updated_at = now() + WHERE mcp_server_id = $1 AND user_id = $2 AND is_default"#, + ) + .bind(server_id) + .bind(user_id) + .execute(&mut *tx) + .await?; + sqlx::query( + r#"UPDATE mcp_user_credentials SET is_default = true, updated_at = now() + WHERE mcp_server_id = $1 AND user_id = $2 AND account_label = $3"#, + ) + .bind(server_id) + .bind(user_id) + .bind(account_label) + .execute(&mut *tx) + .await?; + tx.commit().await?; + Ok(true) +} + +// --------------------------------------------------------------------------- +// mcp_server_shared_credentials +// --------------------------------------------------------------------------- + +/// UPSERT into `mcp_server_shared_credentials`. Single row per server +/// — when the admin rotates the credential the new row replaces the +/// previous one. Uses `INSERT … ON CONFLICT` keyed on the server_id +/// PK so the lifecycle code in `UserTokenResolver` sees a fresh +/// `(access_token_encrypted, expires_at)` after a rotation without +/// any extra coordination. +#[allow(clippy::too_many_arguments)] +pub async fn upsert_shared_credential( + pool: &PgPool, + server_id: Uuid, + credential_type: &str, + access_encrypted: &[u8], + refresh_encrypted: Option<&[u8]>, + expires_at: Option>, + scopes: &[String], + upstream_subject: Option<&str>, + configured_by: Uuid, +) -> Result<(), AppError> { + sqlx::query( + r#"INSERT INTO mcp_server_shared_credentials ( + mcp_server_id, credential_type, + access_token_encrypted, refresh_token_encrypted, + expires_at, scopes, upstream_subject, configured_by + ) + VALUES ($1, $2, $3, $4, $5, $6, $7, $8) + ON CONFLICT (mcp_server_id) DO UPDATE SET + credential_type = EXCLUDED.credential_type, + access_token_encrypted = EXCLUDED.access_token_encrypted, + refresh_token_encrypted = EXCLUDED.refresh_token_encrypted, + expires_at = EXCLUDED.expires_at, + scopes = EXCLUDED.scopes, + upstream_subject = EXCLUDED.upstream_subject, + configured_by = EXCLUDED.configured_by, + updated_at = now()"#, + ) + .bind(server_id) + .bind(credential_type) + .bind(access_encrypted) + .bind(refresh_encrypted) + .bind(expires_at) + .bind(scopes) + .bind(upstream_subject) + .bind(configured_by) + .execute(pool) + .await?; + Ok(()) +} + +/// Insert a server's shared credential inside the caller's transaction +/// (the one that inserts the server row). +#[allow(clippy::too_many_arguments)] +pub async fn insert_shared_credential( + conn: &mut PgConnection, + server_id: Uuid, + credential_type: &str, + access_encrypted: &[u8], + refresh_encrypted: Option<&[u8]>, + expires_at: Option>, + scopes: &[String], + upstream_subject: Option<&str>, + configured_by: Uuid, +) -> Result<(), AppError> { + sqlx::query( + r#"INSERT INTO mcp_server_shared_credentials ( + mcp_server_id, credential_type, + access_token_encrypted, refresh_token_encrypted, + expires_at, scopes, upstream_subject, configured_by + ) + VALUES ($1, $2, $3, $4, $5, $6, $7, $8)"#, + ) + .bind(server_id) + .bind(credential_type) + .bind(access_encrypted) + .bind(refresh_encrypted) + .bind(expires_at) + .bind(scopes) + .bind(upstream_subject) + .bind(configured_by) + .execute(conn) + .await?; + Ok(()) +} + +/// Insert a pasted static token as a server's shared credential, inside +/// the caller's transaction (the one that inserts the server row). +pub async fn insert_shared_static_token( + conn: &mut PgConnection, + server_id: Uuid, + access_encrypted: &[u8], + configured_by: Uuid, +) -> Result<(), AppError> { + sqlx::query( + r#"INSERT INTO mcp_server_shared_credentials ( + mcp_server_id, credential_type, access_token_encrypted, configured_by + ) + VALUES ($1, 'static_token', $2, $3)"#, + ) + .bind(server_id) + .bind(access_encrypted) + .bind(configured_by) + .execute(conn) + .await?; + Ok(()) +} + +pub async fn find_shared_status( + pool: &PgPool, + server_id: Uuid, +) -> Result, AppError> { + Ok(sqlx::query_as::<_, SharedCredentialStatusRow>( + r#"SELECT credential_type, expires_at, upstream_subject, configured_by, updated_at + FROM mcp_server_shared_credentials WHERE mcp_server_id = $1"#, + ) + .bind(server_id) + .fetch_optional(pool) + .await?) +} + +/// The shared credential's type and encrypted access token. +pub async fn find_shared_token( + pool: &PgPool, + server_id: Uuid, +) -> Result)>, AppError> { + Ok(sqlx::query_as( + r#"SELECT credential_type, access_token_encrypted + FROM mcp_server_shared_credentials WHERE mcp_server_id = $1"#, + ) + .bind(server_id) + .fetch_optional(pool) + .await?) +} + +pub async fn delete_shared_credential(pool: &PgPool, server_id: Uuid) -> Result<(), AppError> { + sqlx::query("DELETE FROM mcp_server_shared_credentials WHERE mcp_server_id = $1") + .bind(server_id) + .execute(pool) + .await?; + Ok(()) +} diff --git a/crates/server/src/services/mcp_server_repository.rs b/crates/server/src/services/mcp_server_repository.rs new file mode 100644 index 00000000..21a06543 --- /dev/null +++ b/crates/server/src/services/mcp_server_repository.rs @@ -0,0 +1,314 @@ +//! MCP server repository — the `mcp_servers` table, plus the credential +//! and store-count rows that change in the same transaction as a server +//! update or delete. +//! +//! Thin wrappers over sqlx, one statement (or one transaction) per +//! function; validation, secret encryption, registry sync and audit stay +//! in `handlers::mcp_servers`. + +use sqlx::{PgConnection, PgPool}; +use think_watch_common::errors::AppError; +use think_watch_common::models::McpServer; +use uuid::Uuid; + +/// The columns an admin writes when creating or updating a server. +pub struct McpServerFields<'a> { + pub name: &'a str, + pub namespace_prefix: &'a str, + pub display_label: Option<&'a str>, + pub description: Option<&'a str>, + pub endpoint_url: &'a str, + pub transport_type: &'a str, + pub oauth_issuer: Option<&'a str>, + pub oauth_authorization_endpoint: Option<&'a str>, + pub oauth_token_endpoint: Option<&'a str>, + pub oauth_revocation_endpoint: Option<&'a str>, + pub oauth_userinfo_endpoint: Option<&'a str>, + pub oauth_client_id: Option<&'a str>, + pub oauth_client_secret_encrypted: Option<&'a [u8]>, + pub oauth_scopes: &'a [String], + pub auth_shape: &'a str, + pub static_token_help_url: Option<&'a str>, + pub auth_header_name: &'a str, + pub auth_value_template: &'a str, + pub credential_owner: &'a str, + pub config_json: &'a serde_json::Value, +} + +/// Every server with its active-tool count, newest first. +pub async fn list_with_tool_counts(pool: &PgPool) -> Result, AppError> { + Ok(sqlx::query_as::<_, McpServer>( + r#"SELECT s.*, COALESCE(t.cnt, 0) AS tools_count + FROM mcp_servers s + LEFT JOIN (SELECT server_id, COUNT(*) AS cnt FROM mcp_tools WHERE is_active = true GROUP BY server_id) t + ON t.server_id = s.id + ORDER BY s.created_at DESC"#, + ) + .fetch_all(pool) + .await?) +} + +/// Every server by name, with zeroed tool and call counts. +pub async fn list_by_name(pool: &PgPool) -> Result, AppError> { + Ok(sqlx::query_as::<_, McpServer>( + r#"SELECT s.*, 0::bigint AS tools_count, 0::bigint AS call_count + FROM mcp_servers s + ORDER BY s.name"#, + ) + .fetch_all(pool) + .await?) +} + +pub async fn find(pool: &PgPool, id: Uuid) -> Result, AppError> { + Ok( + sqlx::query_as::<_, McpServer>("SELECT * FROM mcp_servers WHERE id = $1") + .bind(id) + .fetch_optional(pool) + .await?, + ) +} + +/// One server with zeroed tool and call counts. +pub async fn find_without_counts(pool: &PgPool, id: Uuid) -> Result, AppError> { + Ok(sqlx::query_as::<_, McpServer>( + r#"SELECT s.*, 0::bigint AS tools_count, 0::bigint AS call_count + FROM mcp_servers s WHERE s.id = $1"#, + ) + .bind(id) + .fetch_optional(pool) + .await?) +} + +/// Whether another server already uses this name or namespace prefix. +pub async fn name_or_prefix_taken( + conn: &mut PgConnection, + name: &str, + namespace_prefix: &str, +) -> Result { + // `SELECT 1` is INT4 on the wire; binding into `Option` + // panics with a column-decode mismatch the moment a row + // comes back. We don't actually care about the value — only + // whether the row exists — so use Option. + let conflict: Option = sqlx::query_scalar( + "SELECT 1 FROM mcp_servers WHERE name = $1 OR namespace_prefix = $2 LIMIT 1", + ) + .bind(name) + .bind(namespace_prefix) + .fetch_optional(conn) + .await?; + Ok(conflict.is_some()) +} + +/// Insert a server inside the caller's transaction. A taken name or +/// prefix comes back as a 409. +pub async fn insert( + conn: &mut PgConnection, + f: &McpServerFields<'_>, +) -> Result { + sqlx::query_as::<_, McpServer>( + r#"INSERT INTO mcp_servers ( + name, namespace_prefix, display_label, description, endpoint_url, transport_type, + oauth_issuer, oauth_authorization_endpoint, oauth_token_endpoint, + oauth_revocation_endpoint, oauth_userinfo_endpoint, + oauth_client_id, oauth_client_secret_encrypted, + oauth_scopes, auth_shape, static_token_help_url, + auth_header_name, auth_value_template, credential_owner, + config_json + ) + VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, + $16, $17, $18, $19, $20) + RETURNING *"#, + ) + .bind(f.name) + .bind(f.namespace_prefix) + .bind(f.display_label) + .bind(f.description) + .bind(f.endpoint_url) + .bind(f.transport_type) + .bind(f.oauth_issuer) + .bind(f.oauth_authorization_endpoint) + .bind(f.oauth_token_endpoint) + .bind(f.oauth_revocation_endpoint) + .bind(f.oauth_userinfo_endpoint) + .bind(f.oauth_client_id) + .bind(f.oauth_client_secret_encrypted) + .bind(f.oauth_scopes) + .bind(f.auth_shape) + .bind(f.static_token_help_url) + .bind(f.auth_header_name) + .bind(f.auth_value_template) + .bind(f.credential_owner) + .bind(f.config_json) + .fetch_one(conn) + .await + .map_err(map_unique_violation) +} + +/// Update a server and, in the same transaction, drop the credentials +/// the change made stale: per-user credentials and tool caches when +/// `purge_user_credentials`, the shared credential when +/// `purge_shared_credential`. A taken name or prefix comes back as a 409. +pub async fn update( + pool: &PgPool, + id: Uuid, + f: &McpServerFields<'_>, + purge_user_credentials: bool, + purge_shared_credential: bool, +) -> Result { + let mut tx = pool.begin().await?; + let updated = sqlx::query_as::<_, McpServer>( + r#"UPDATE mcp_servers SET + name = $2, namespace_prefix = $3, display_label = $4, + description = $5, endpoint_url = $6, + transport_type = $7, + oauth_issuer = $8, oauth_authorization_endpoint = $9, + oauth_token_endpoint = $10, oauth_revocation_endpoint = $11, + oauth_userinfo_endpoint = $12, + oauth_client_id = $13, oauth_client_secret_encrypted = $14, + oauth_scopes = $15, auth_shape = $16, static_token_help_url = $17, + auth_header_name = $18, auth_value_template = $19, credential_owner = $20, + config_json = $21 + WHERE id = $1 RETURNING *"#, + ) + .bind(id) + .bind(f.name) + .bind(f.namespace_prefix) + .bind(f.display_label) + .bind(f.description) + .bind(f.endpoint_url) + .bind(f.transport_type) + .bind(f.oauth_issuer) + .bind(f.oauth_authorization_endpoint) + .bind(f.oauth_token_endpoint) + .bind(f.oauth_revocation_endpoint) + .bind(f.oauth_userinfo_endpoint) + .bind(f.oauth_client_id) + .bind(f.oauth_client_secret_encrypted) + .bind(f.oauth_scopes) + .bind(f.auth_shape) + .bind(f.static_token_help_url) + .bind(f.auth_header_name) + .bind(f.auth_value_template) + .bind(f.credential_owner) + .bind(f.config_json) + .fetch_one(&mut *tx) + .await + .map_err(map_unique_violation)?; + + if purge_user_credentials { + sqlx::query("DELETE FROM mcp_user_credentials WHERE mcp_server_id = $1") + .bind(id) + .execute(&mut *tx) + .await?; + sqlx::query("DELETE FROM mcp_user_tools WHERE mcp_server_id = $1") + .bind(id) + .execute(&mut *tx) + .await?; + } + if purge_shared_credential { + sqlx::query("DELETE FROM mcp_server_shared_credentials WHERE mcp_server_id = $1") + .bind(id) + .execute(&mut *tx) + .await?; + } + tx.commit().await?; + Ok(updated) +} + +/// Delete one server. Returns its name, or `None` (and changes nothing) +/// when there is no such server. +pub async fn delete(pool: &PgPool, id: Uuid) -> Result, AppError> { + let mut tx = pool.begin().await?; + let Some(name) = delete_in_tx(&mut tx, id).await? else { + return Ok(None); + }; + tx.commit().await?; + Ok(Some(name)) +} + +/// Delete several servers in one transaction: all of them or none. +/// Returns each id with its name, or `None` where there was no such +/// server. +pub async fn delete_many( + pool: &PgPool, + ids: &[Uuid], +) -> Result)>, AppError> { + let mut tx = pool.begin().await?; + let mut out = Vec::with_capacity(ids.len()); + for &id in ids { + out.push((id, delete_in_tx(&mut tx, id).await?)); + } + tx.commit().await?; + Ok(out) +} + +/// Tear down a single MCP server inside the caller's transaction: +/// * SELECT the server name (returned to the caller for audit detail) +/// * decrement the originating store template's `install_count` +/// * DELETE the server row (children CASCADE: `mcp_tools`, +/// `mcp_user_credentials`, `mcp_server_shared_credentials`, +/// `mcp_user_tools`, `mcp_store_installs`) +/// +/// Returns `Ok(Some(name))` on success, `Ok(None)` if the row doesn't +/// exist. +async fn delete_in_tx(conn: &mut PgConnection, id: Uuid) -> Result, AppError> { + let name: Option = sqlx::query_scalar("SELECT name FROM mcp_servers WHERE id = $1") + .bind(id) + .fetch_optional(&mut *conn) + .await?; + if name.is_none() { + return Ok(None); + } + + // Decrement install_count if this server was installed from the store. + sqlx::query( + r#"UPDATE mcp_store_templates SET install_count = GREATEST(install_count - 1, 0) + WHERE id = (SELECT template_id FROM mcp_store_installs WHERE server_id = $1)"#, + ) + .bind(id) + .execute(&mut *conn) + .await?; + + sqlx::query("DELETE FROM mcp_servers WHERE id = $1") + .bind(id) + .execute(&mut *conn) + .await?; + + Ok(name) +} + +pub async fn clear_last_error(pool: &PgPool, id: Uuid) -> Result<(), AppError> { + sqlx::query("UPDATE mcp_servers SET last_error = NULL WHERE id = $1") + .bind(id) + .execute(pool) + .await?; + Ok(()) +} + +pub async fn set_last_error(pool: &PgPool, id: Uuid, error: &str) -> Result<(), AppError> { + sqlx::query("UPDATE mcp_servers SET last_error = $1 WHERE id = $2") + .bind(error) + .bind(id) + .execute(pool) + .await?; + Ok(()) +} + +/// Translate PostgreSQL unique-constraint violations on `mcp_servers` into +/// user-facing conflict errors, so the UI shows "already in use" instead of +/// a generic 500. Other sqlx errors fall through unchanged. +fn map_unique_violation(e: sqlx::Error) -> AppError { + if let sqlx::Error::Database(db_err) = &e + && db_err.code().as_deref() == Some("23505") + { + let constraint = db_err.constraint().unwrap_or(""); + if constraint.contains("namespace_prefix") { + return AppError::Conflict("namespace_prefix already in use".into()); + } + if constraint.contains("name") { + return AppError::Conflict("server name already in use".into()); + } + return AppError::Conflict("duplicate server".into()); + } + AppError::from(e) +} diff --git a/crates/server/src/services/mcp_store_repository.rs b/crates/server/src/services/mcp_store_repository.rs new file mode 100644 index 00000000..f445d384 --- /dev/null +++ b/crates/server/src/services/mcp_store_repository.rs @@ -0,0 +1,235 @@ +//! MCP store repository — `mcp_store_templates` (the catalog synced from +//! a remote registry) and `mcp_store_installs` (which server came from +//! which template). +//! +//! Registry fetching, template validation and audit stay in +//! `handlers::mcp_store`; installing runs inside the server-create +//! transaction in `handlers::mcp_servers`. + +use sqlx::{PgConnection, PgPool}; +use think_watch_common::errors::AppError; +use think_watch_common::models::McpStoreTemplate; +use uuid::Uuid; + +/// Process-wide advisory-lock key for serializing template installs. +/// The literal spells "mcpStore" in ASCII so a DBA glancing at +/// `pg_locks` can tell what's holding it. Any new advisory lock +/// added elsewhere in the codebase MUST use a distinct constant — +/// collisions silently serialize unrelated work and can deadlock +/// under concurrent load. +/// +/// Reserved advisory lock keys (keep this list current): +/// * `MCP_STORE_INSTALL_LOCK_KEY` (here): template-install +/// serialization in `create_server` when `template_slug` is set. +const MCP_STORE_INSTALL_LOCK_KEY: i64 = 0x6D637053746F7265; + +/// A category and how many templates are in it. +#[derive(sqlx::FromRow)] +pub struct CategoryCountRow { + pub category: Option, + pub count: Option, +} + +/// One registry template as written to `mcp_store_templates`. +pub struct TemplateUpsert<'a> { + pub slug: &'a str, + pub name: &'a str, + pub description: Option<&'a str>, + pub category: Option<&'a str>, + pub tags: &'a [String], + pub endpoint_template: Option<&'a str>, + pub oauth_issuer: Option<&'a str>, + pub oauth_authorization_endpoint: Option<&'a str>, + pub oauth_token_endpoint: Option<&'a str>, + pub oauth_revocation_endpoint: Option<&'a str>, + pub oauth_userinfo_endpoint: Option<&'a str>, + pub oauth_default_scopes: &'a [String], + pub auth_shape: &'a str, + pub static_token_help_url: Option<&'a str>, + pub auth_header_name: &'a str, + pub auth_value_template: &'a str, + pub auth_instructions: Option<&'a str>, + pub deploy_type: &'a str, + pub deploy_command: Option<&'a str>, + pub deploy_docs_url: Option<&'a str>, + pub homepage_url: Option<&'a str>, + pub repo_url: Option<&'a str>, + pub featured: bool, +} + +/// Every template that has been installed at least once. +pub async fn installed_template_ids(pool: &PgPool) -> Result, AppError> { + Ok( + sqlx::query_scalar("SELECT template_id FROM mcp_store_installs") + .fetch_all(pool) + .await?, + ) +} + +/// Every template, featured and most-installed first. +pub async fn list_templates(pool: &PgPool) -> Result, AppError> { + Ok(sqlx::query_as::<_, McpStoreTemplate>( + "SELECT * FROM mcp_store_templates ORDER BY featured DESC, install_count DESC, name ASC", + ) + .fetch_all(pool) + .await?) +} + +pub async fn find_template_by_slug( + pool: &PgPool, + slug: &str, +) -> Result, AppError> { + Ok( + sqlx::query_as::<_, McpStoreTemplate>("SELECT * FROM mcp_store_templates WHERE slug = $1") + .bind(slug) + .fetch_optional(pool) + .await?, + ) +} + +/// Template counts per category, largest first. +pub async fn category_counts(pool: &PgPool) -> Result, AppError> { + Ok(sqlx::query_as::<_, CategoryCountRow>( + "SELECT category, COUNT(*) as count FROM mcp_store_templates GROUP BY category ORDER BY count DESC", + ) + .fetch_all(pool) + .await?) +} + +/// Insert or refresh one template by slug, inside the caller's sync +/// transaction. +pub async fn upsert_template( + conn: &mut PgConnection, + t: &TemplateUpsert<'_>, +) -> Result<(), AppError> { + sqlx::query( + r#"INSERT INTO mcp_store_templates + (slug, name, description, category, tags, endpoint_template, + oauth_issuer, oauth_authorization_endpoint, oauth_token_endpoint, + oauth_revocation_endpoint, oauth_userinfo_endpoint, + oauth_default_scopes, + auth_shape, static_token_help_url, + auth_header_name, auth_value_template, + auth_instructions, deploy_type, + deploy_command, deploy_docs_url, homepage_url, repo_url, featured, updated_at) + VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, + $16, $17, $18, $19, $20, $21, $22, $23, now()) + ON CONFLICT (slug) DO UPDATE SET + name = EXCLUDED.name, + description = EXCLUDED.description, + category = EXCLUDED.category, + tags = EXCLUDED.tags, + endpoint_template = EXCLUDED.endpoint_template, + oauth_issuer = EXCLUDED.oauth_issuer, + oauth_authorization_endpoint = EXCLUDED.oauth_authorization_endpoint, + oauth_token_endpoint = EXCLUDED.oauth_token_endpoint, + oauth_revocation_endpoint = EXCLUDED.oauth_revocation_endpoint, + oauth_userinfo_endpoint = EXCLUDED.oauth_userinfo_endpoint, + oauth_default_scopes = EXCLUDED.oauth_default_scopes, + auth_shape = EXCLUDED.auth_shape, + static_token_help_url = EXCLUDED.static_token_help_url, + auth_header_name = EXCLUDED.auth_header_name, + auth_value_template = EXCLUDED.auth_value_template, + auth_instructions = EXCLUDED.auth_instructions, + deploy_type = EXCLUDED.deploy_type, + deploy_command = EXCLUDED.deploy_command, + deploy_docs_url = EXCLUDED.deploy_docs_url, + homepage_url = EXCLUDED.homepage_url, + repo_url = EXCLUDED.repo_url, + featured = EXCLUDED.featured, + updated_at = now()"#, + ) + .bind(t.slug) + .bind(t.name) + .bind(t.description) + .bind(t.category) + .bind(t.tags) + .bind(t.endpoint_template) + .bind(t.oauth_issuer) + .bind(t.oauth_authorization_endpoint) + .bind(t.oauth_token_endpoint) + .bind(t.oauth_revocation_endpoint) + .bind(t.oauth_userinfo_endpoint) + .bind(t.oauth_default_scopes) + .bind(t.auth_shape) + .bind(t.static_token_help_url) + .bind(t.auth_header_name) + .bind(t.auth_value_template) + .bind(t.auth_instructions) + .bind(t.deploy_type) + .bind(t.deploy_command) + .bind(t.deploy_docs_url) + .bind(t.homepage_url) + .bind(t.repo_url) + .bind(t.featured) + .execute(conn) + .await?; + Ok(()) +} + +/// Delete the templates whose slug isn't in `keep_slugs`, except those +/// with installs, inside the caller's sync transaction. Returns how many +/// went. +pub async fn delete_templates_not_in( + conn: &mut PgConnection, + keep_slugs: &[&str], +) -> Result { + Ok(sqlx::query_scalar::<_, i64>( + r#"WITH deleted AS ( + DELETE FROM mcp_store_templates + WHERE slug != ALL($1) + AND id NOT IN (SELECT template_id FROM mcp_store_installs) + RETURNING 1 + ) + SELECT COUNT(*) FROM deleted"#, + ) + .bind(keep_slugs) + .fetch_one(conn) + .await?) +} + +/// Serialize template installs for the rest of the caller's transaction, +/// so two concurrent installs can't resolve to the same server name. +pub async fn lock_installs(conn: &mut PgConnection) -> Result<(), AppError> { + sqlx::query("SELECT pg_advisory_xact_lock($1)") + .bind(MCP_STORE_INSTALL_LOCK_KEY) + .execute(conn) + .await?; + Ok(()) +} + +/// A template's id by slug, row-locked for the caller's transaction. +pub async fn lock_template_by_slug( + conn: &mut PgConnection, + slug: &str, +) -> Result, AppError> { + Ok( + sqlx::query_scalar("SELECT id FROM mcp_store_templates WHERE slug = $1 FOR UPDATE") + .bind(slug) + .fetch_optional(conn) + .await?, + ) +} + +/// Record that a server was installed from a template and bump the +/// template's `install_count`, inside the caller's transaction. +pub async fn record_install( + conn: &mut PgConnection, + template_id: Uuid, + server_id: Uuid, + installed_by: Uuid, +) -> Result<(), AppError> { + sqlx::query( + "INSERT INTO mcp_store_installs (template_id, server_id, installed_by) VALUES ($1, $2, $3)", + ) + .bind(template_id) + .bind(server_id) + .bind(installed_by) + .execute(&mut *conn) + .await?; + sqlx::query("UPDATE mcp_store_templates SET install_count = install_count + 1 WHERE id = $1") + .bind(template_id) + .execute(&mut *conn) + .await?; + Ok(()) +} diff --git a/crates/server/src/services/mcp_tool_repository.rs b/crates/server/src/services/mcp_tool_repository.rs new file mode 100644 index 00000000..e70765ea --- /dev/null +++ b/crates/server/src/services/mcp_tool_repository.rs @@ -0,0 +1,136 @@ +//! MCP tool catalog repository — `mcp_tools` (discovered per server) +//! unioned with a user's own `mcp_user_tools`, as listed on +//! `/api/mcp/tools`. +//! +//! The per-user catalog is pre-namespaced the same way mcp_tools is +//! (`__`) and unioned with it. `mcp_user_tools` doesn't +//! carry an `id` column — synthesize a stable v5-style UUID from +//! `(server_id, user_id, tool_name)` so the frontend's keying +//! (`tool.id`) keeps working without a schema change. + +use serde::Serialize; +use sqlx::{FromRow, PgPool}; +use think_watch_common::errors::AppError; +use uuid::Uuid; + +#[derive(Debug, Serialize, FromRow, utoipa::ToSchema)] +pub struct McpToolRow { + #[schema(value_type = String, format = Uuid)] + pub id: uuid::Uuid, + #[schema(value_type = String, format = Uuid)] + pub server_id: uuid::Uuid, + pub server_name: String, + pub name: String, + pub namespaced_name: String, + pub description: Option, + #[schema(value_type = Object)] + pub input_schema: Option, +} + +/// Filter and page for [`count_catalog`] / [`list_catalog`]. +pub struct CatalogQuery<'a> { + /// Trimmed search text; empty matches everything. + pub search: &'a str, + /// `%search%`, matched with ILIKE. + pub search_pattern: &'a str, + pub server_id: Option, + pub page_size: i64, + pub offset: i64, + /// Whose `mcp_user_tools` to include; `None` includes none. + pub user_id: Option, +} + +pub async fn count_catalog(pool: &PgPool, q: &CatalogQuery<'_>) -> Result { + Ok(sqlx::query_scalar( + r#"WITH catalog AS ( + SELECT t.id, + t.server_id, + s.name AS server_name, + s.namespace_prefix, + t.tool_name, + t.description + FROM mcp_tools t + JOIN mcp_servers s ON s.id = t.server_id + WHERE t.is_active = true + UNION ALL + SELECT gen_random_uuid() AS id, + u.mcp_server_id AS server_id, + s.name AS server_name, + s.namespace_prefix, + u.tool_name, + u.description + FROM mcp_user_tools u + JOIN mcp_servers s ON s.id = u.mcp_server_id + WHERE $6::uuid IS NOT NULL AND u.user_id = $6::uuid + ) + SELECT COUNT(*) FROM catalog + WHERE ($3::uuid IS NULL OR server_id = $3) + AND ($1 = '' + OR tool_name ILIKE $2 + OR (namespace_prefix || '__' || tool_name) ILIKE $2 + OR COALESCE(description, '') ILIKE $2)"#, + ) + .bind(q.search) + .bind(q.search_pattern) + .bind(q.server_id) + .bind(q.page_size) + .bind(q.offset) + .bind(q.user_id) + .fetch_one(pool) + .await?) +} + +/// One page of the catalog, by server name then tool name. +pub async fn list_catalog( + pool: &PgPool, + q: &CatalogQuery<'_>, +) -> Result, AppError> { + Ok(sqlx::query_as::<_, McpToolRow>( + r#"WITH catalog AS ( + SELECT t.id, + t.server_id, + s.name AS server_name, + s.namespace_prefix, + t.tool_name, + t.description, + t.input_schema + FROM mcp_tools t + JOIN mcp_servers s ON s.id = t.server_id + WHERE t.is_active = true + UNION ALL + SELECT gen_random_uuid() AS id, + u.mcp_server_id AS server_id, + s.name AS server_name, + s.namespace_prefix, + u.tool_name, + u.description, + u.input_schema + FROM mcp_user_tools u + JOIN mcp_servers s ON s.id = u.mcp_server_id + WHERE $6::uuid IS NOT NULL AND u.user_id = $6::uuid + ) + SELECT id, + server_id, + server_name, + tool_name AS name, + namespace_prefix || '__' || tool_name AS namespaced_name, + description, + input_schema + FROM catalog + WHERE ($3::uuid IS NULL OR server_id = $3) + AND ($1 = '' + OR tool_name ILIKE $2 + OR (namespace_prefix || '__' || tool_name) ILIKE $2 + OR COALESCE(description, '') ILIKE $2) + ORDER BY server_name, tool_name + LIMIT $4 OFFSET $5"#, + ) + .bind(q.search) + .bind(q.search_pattern) + .bind(q.server_id) + .bind(q.page_size) + .bind(q.offset) + .bind(q.user_id) + .fetch_all(pool) + .await?) +} diff --git a/crates/server/src/services/mod.rs b/crates/server/src/services/mod.rs index cd49e0db..daa58ac5 100644 --- a/crates/server/src/services/mod.rs +++ b/crates/server/src/services/mod.rs @@ -29,6 +29,10 @@ pub mod analytics_repository; pub mod auth_lockout; pub mod limits_repository; pub mod log_forwarder_repository; +pub mod mcp_credential_repository; +pub mod mcp_server_repository; +pub mod mcp_store_repository; +pub mod mcp_tool_repository; pub mod model_repository; pub mod observability_repository; pub mod pricing_repository; diff --git a/crates/test-support/tests/admin_mcp_catalog.rs b/crates/test-support/tests/admin_mcp_catalog.rs new file mode 100644 index 00000000..82ebd56a --- /dev/null +++ b/crates/test-support/tests/admin_mcp_catalog.rs @@ -0,0 +1,995 @@ +//! The MCP admin and connection endpoints end to end: servers, the +//! store, the tool catalog, shared and per-user credentials. +//! +//! The rest of the MCP suite covers the proxy and the credential +//! transitions; this file pins what the remaining endpoints read and +//! write — lookups, 404s, 409s, background error reporting, the registry +//! sync — so moving their SQL around (into `services::mcp_*_repository`) +//! is checked rather than assumed. + +use std::time::Duration; + +use serde_json::Value; +use think_watch_test_support::prelude::*; +use wiremock::matchers::{method, path}; +use wiremock::{Mock, MockServer, ResponseTemplate}; + +/// An MCP upstream whose `tools/list` returns one tool (`echo`). +async fn mcp_ok() -> MockServer { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/mcp")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "jsonrpc": "2.0", + "id": 1, + "result": { + "tools": [{ + "name": "echo", + "description": "Echo back the input", + "inputSchema": {"type": "object"} + }] + } + }))) + .mount(&server) + .await; + server +} + +/// An MCP upstream that answers every request with `status`. +async fn mcp_status(status: u16) -> MockServer { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/mcp")) + .respond_with(ResponseTemplate::new(status)) + .mount(&server) + .await; + server +} + +async fn get(con: &TestClient, path: &str) -> Value { + let resp = con.get(path).await.unwrap(); + resp.assert_ok(); + resp.json().unwrap() +} + +async fn create_server(con: &TestClient, body: Value) -> Value { + let resp = con.post("/api/mcp/servers", body).await.unwrap(); + resp.assert_ok(); + resp.json().unwrap() +} + +fn prefix() -> String { + format!("p_{}", &Uuid::new_v4().simple().to_string()[..12]) +} + +/// Poll `GET /api/mcp/servers` until the server's row satisfies `done` +/// (background discovery writes it after the request returns). +async fn wait_for_server(con: &TestClient, id: &str, done: impl Fn(&Value) -> bool) -> Value { + let mut last = Value::Null; + for _ in 0..100 { + let list = get(con, "/api/mcp/servers").await; + if let Some(row) = list.as_array().unwrap().iter().find(|s| s["id"] == id) { + if done(row) { + return row.clone(); + } + last = row.clone(); + } + tokio::time::sleep(Duration::from_millis(100)).await; + } + panic!("server {id} never reached the expected state: {last}"); +} + +async fn insert_tool(app: &TestApp, server_id: Uuid, name: &str, desc: &str, active: bool) { + sqlx::query( + "INSERT INTO mcp_tools (server_id, tool_name, description, input_schema, is_active) + VALUES ($1, $2, $3, '{}'::jsonb, $4)", + ) + .bind(server_id) + .bind(name) + .bind(desc) + .bind(active) + .execute(&app.db) + .await + .unwrap(); +} + +// --------------------------------------------------------------------------- +// Servers +// --------------------------------------------------------------------------- + +#[ignore = "integration test — run via `make test-it`"] +#[tokio::test] +async fn a_server_is_created_read_edited_and_deleted() { + let app = TestApp::spawn_reaching_loopback().await; + let con = admin_session(&app).await; + let upstream = mcp_status(500).await; + + let name = unique_name("srv"); + let pfx = prefix(); + let created = create_server( + &con, + json!({ + "name": name, + "namespace_prefix": pfx, + "display_label": " Shown ", + "description": "first", + "endpoint_url": format!("{}/mcp", upstream.uri()), + "transport_type": "streamable_http", + "oauth_scopes": ["a", "b"], + "custom_headers": {"X-Team": "{{user_id}}"}, + "cache_ttl_secs": 30, + }), + ) + .await; + let id = created["id"].as_str().unwrap().to_string(); + assert_eq!(created["name"], name.as_str()); + assert_eq!(created["namespace_prefix"], pfx.as_str()); + assert_eq!(created["display_label"], "Shown", "{created}"); + assert_eq!(created["auth_shape"], "anonymous"); + assert_eq!(created["credential_owner"], "per_user"); + assert_eq!(created["auth_header_name"], "Authorization"); + assert_eq!(created["auth_value_template"], "Bearer {{token}}"); + assert_eq!(created["oauth_scopes"], json!(["a", "b"])); + assert_eq!( + created["config_json"], + json!({"custom_headers": {"X-Team": "{{user_id}}"}, "cache_ttl_secs": 30}) + ); + + // The failed first discovery lands on the row. + let row = wait_for_server(&con, &id, |s| s["last_error"].is_string()).await; + assert!( + row["last_error"].as_str().unwrap().contains("HTTP 500"), + "{row}" + ); + + let got = get(&con, &format!("/api/mcp/servers/{id}")).await; + assert_eq!(got["description"], "first"); + assert_eq!(got["display_label"], "Shown"); + con.get(&format!("/api/mcp/servers/{}", Uuid::new_v4())) + .await + .unwrap() + .assert_status(404); + + // PATCH: JSON null clears, absent keeps, a value replaces. + let new_pfx = prefix(); + let resp = con + .patch( + &format!("/api/mcp/servers/{id}"), + json!({ + "display_label": null, + "description": "second", + "namespace_prefix": new_pfx, + "custom_headers": {"X-Other": "1"}, + }), + ) + .await + .unwrap(); + resp.assert_ok(); + let patched: Value = resp.json().unwrap(); + assert!(patched["display_label"].is_null(), "{patched}"); + assert_eq!(patched["description"], "second"); + assert_eq!(patched["namespace_prefix"], new_pfx.as_str()); + assert_eq!(patched["name"], name.as_str()); + assert_eq!(patched["oauth_scopes"], json!(["a", "b"])); + assert_eq!( + patched["config_json"], + json!({"custom_headers": {"X-Other": "1"}, "cache_ttl_secs": 30}) + ); + let got = get(&con, &format!("/api/mcp/servers/{id}")).await; + assert_eq!(got["description"], "second"); + assert!(got["display_label"].is_null()); + + con.patch( + &format!("/api/mcp/servers/{}", Uuid::new_v4()), + json!({"description": "x"}), + ) + .await + .unwrap() + .assert_status(404); + + con.delete(&format!("/api/mcp/servers/{id}")) + .await + .unwrap() + .assert_ok(); + con.get(&format!("/api/mcp/servers/{id}")) + .await + .unwrap() + .assert_status(404); + con.delete(&format!("/api/mcp/servers/{id}")) + .await + .unwrap() + .assert_status(404); +} + +#[ignore = "integration test — run via `make test-it`"] +#[tokio::test] +async fn a_taken_name_or_prefix_is_a_409() { + let app = TestApp::spawn_reaching_loopback().await; + let con = admin_session(&app).await; + let upstream = mcp_status(500).await; + let endpoint = format!("{}/mcp", upstream.uri()); + + let (name_a, pfx_a) = (unique_name("a"), prefix()); + create_server( + &con, + json!({"name": name_a, "namespace_prefix": pfx_a, "endpoint_url": endpoint, + "transport_type": "streamable_http"}), + ) + .await; + let b = create_server( + &con, + json!({"name": unique_name("b"), "namespace_prefix": prefix(), "endpoint_url": endpoint, + "transport_type": "streamable_http"}), + ) + .await; + let b_id = b["id"].as_str().unwrap(); + + let resp = con + .post( + "/api/mcp/servers", + json!({"name": name_a, "namespace_prefix": prefix(), "endpoint_url": endpoint, + "transport_type": "streamable_http"}), + ) + .await + .unwrap(); + resp.assert_status(409); + assert!( + resp.text().contains("server name already in use"), + "{}", + resp.text() + ); + + let resp = con + .post( + "/api/mcp/servers", + json!({"name": unique_name("c"), "namespace_prefix": pfx_a, "endpoint_url": endpoint, + "transport_type": "streamable_http"}), + ) + .await + .unwrap(); + resp.assert_status(409); + assert!( + resp.text().contains("namespace_prefix already in use"), + "{}", + resp.text() + ); + + let resp = con + .patch(&format!("/api/mcp/servers/{b_id}"), json!({"name": name_a})) + .await + .unwrap(); + resp.assert_status(409); + assert!( + resp.text().contains("server name already in use"), + "{}", + resp.text() + ); + + let resp = con + .patch( + &format!("/api/mcp/servers/{b_id}"), + json!({"namespace_prefix": pfx_a}), + ) + .await + .unwrap(); + resp.assert_status(409); + assert!( + resp.text().contains("namespace_prefix already in use"), + "{}", + resp.text() + ); + + // Unknown template on the install path. + con.post( + "/api/mcp/servers", + json!({"name": unique_name("t"), "namespace_prefix": prefix(), "endpoint_url": endpoint, + "transport_type": "streamable_http", "template_slug": unique_name("nope")}), + ) + .await + .unwrap() + .assert_status(404); +} + +#[ignore = "integration test — run via `make test-it`"] +#[tokio::test] +async fn the_server_list_counts_active_tools_only() { + let app = TestApp::spawn().await; + let con = admin_session(&app).await; + let id = fixtures::create_mcp_server( + &app.db, + &unique_name("count"), + &prefix(), + "https://example.com/mcp", + ) + .await + .unwrap(); + insert_tool(&app, id, "one", "", true).await; + insert_tool(&app, id, "two", "", true).await; + insert_tool(&app, id, "gone", "", false).await; + let bare = fixtures::create_mcp_server( + &app.db, + &unique_name("bare"), + &prefix(), + "https://example.com/mcp", + ) + .await + .unwrap(); + + let list = get(&con, "/api/mcp/servers").await; + let rows = list.as_array().unwrap(); + let row = |id: Uuid| { + rows.iter() + .find(|s| s["id"] == id.to_string()) + .unwrap_or_else(|| panic!("{id} missing: {list}")) + }; + assert_eq!(row(id)["tools_count"], 2); + assert_eq!(row(bare)["tools_count"], 0); + // Newest first. + let pos = |id: Uuid| rows.iter().position(|s| s["id"] == id.to_string()).unwrap(); + assert!(pos(bare) < pos(id), "{list}"); +} + +#[ignore = "integration test — run via `make test-it`"] +#[tokio::test] +async fn a_shared_static_token_at_create_lands_with_the_server() { + let app = TestApp::spawn_reaching_loopback().await; + let (con, admin) = admin_session_with_user(&app).await; + let upstream = mcp_ok().await; + + let created = create_server( + &con, + json!({ + "name": unique_name("shared"), + "namespace_prefix": prefix(), + "endpoint_url": format!("{}/mcp", upstream.uri()), + "transport_type": "streamable_http", + "auth_shape": "static", + "credential_owner": "admin_shared", + "shared_static_token": "tok-at-create", + }), + ) + .await; + let id = created["id"].as_str().unwrap(); + + let status = get( + &con, + &format!("/api/admin/mcp/servers/{id}/shared-credential"), + ) + .await; + assert_eq!(status["configured"], true, "{status}"); + assert_eq!(status["credential_type"], "static_token"); + assert_eq!(status["configured_by"], admin.user.id.to_string()); + + // Discovery ran with the shared bearer and found the tool. + wait_for_server(&con, id, |s| s["tools_count"] == 1).await; +} + +#[ignore = "integration test — run via `make test-it`"] +#[tokio::test] +async fn discover_on_an_unknown_server_is_a_404() { + let app = TestApp::spawn().await; + let con = admin_session(&app).await; + con.post( + &format!("/api/mcp/servers/{}/discover", Uuid::new_v4()), + json!({}), + ) + .await + .unwrap() + .assert_status(404); +} + +// --------------------------------------------------------------------------- +// Shared credentials +// --------------------------------------------------------------------------- + +#[ignore = "integration test — run via `make test-it`"] +#[tokio::test] +async fn a_shared_credential_is_pasted_reported_and_revoked() { + let app = TestApp::spawn().await; + let (con, admin) = admin_session_with_user(&app).await; + let upstream = mcp_ok().await; + let id = fixtures::create_mcp_server_with( + &app.db, + &unique_name("sc"), + &prefix(), + &format!("{}/mcp", upstream.uri()), + fixtures::McpServerOpts { + auth_shape: "static".into(), + credential_owner: "admin_shared".into(), + ..Default::default() + }, + ) + .await + .unwrap(); + let base = format!("/api/admin/mcp/servers/{id}/shared-credential"); + + let status = get(&con, &base).await; + assert_eq!( + status, + json!({"configured": false, "credential_type": null, "expires_at": null, + "upstream_subject": null, "configured_by": null, "updated_at": null}) + ); + con.delete(&base).await.unwrap().assert_status(404); + + // A stale error is cleared once discovery with the new token works. + sqlx::query("UPDATE mcp_servers SET last_error = 'stale' WHERE id = $1") + .bind(id) + .execute(&app.db) + .await + .unwrap(); + con.put(&format!("{base}/static-token"), json!({"token": "one"})) + .await + .unwrap() + .assert_ok(); + let row = wait_for_server(&con, &id.to_string(), |s| s["last_error"].is_null()).await; + assert_eq!(row["tools_count"], 1, "{row}"); + + let status = get(&con, &base).await; + assert_eq!(status["configured"], true); + assert_eq!(status["credential_type"], "static_token"); + assert_eq!(status["configured_by"], admin.user.id.to_string()); + assert!(status["updated_at"].is_string()); + assert!(status["expires_at"].is_null()); + + // Pasting again replaces the one row. + con.put(&format!("{base}/static-token"), json!({"token": "two"})) + .await + .unwrap() + .assert_ok(); + let rows: i64 = sqlx::query_scalar( + "SELECT COUNT(*) FROM mcp_server_shared_credentials WHERE mcp_server_id = $1", + ) + .bind(id) + .fetch_one(&app.db) + .await + .unwrap(); + assert_eq!(rows, 1); + + let resp = con.delete(&base).await.unwrap(); + resp.assert_ok(); + assert_eq!(resp.json::().unwrap(), json!({"status": "revoked"})); + assert_eq!(get(&con, &base).await["configured"], false); + con.delete(&base).await.unwrap().assert_status(404); +} + +#[ignore = "integration test — run via `make test-it`"] +#[tokio::test] +async fn a_rejected_shared_token_is_reported_on_the_server() { + let app = TestApp::spawn().await; + let con = admin_session(&app).await; + let upstream = mcp_status(401).await; + let id = fixtures::create_mcp_server_with( + &app.db, + &unique_name("rej"), + &prefix(), + &format!("{}/mcp", upstream.uri()), + fixtures::McpServerOpts { + auth_shape: "static".into(), + credential_owner: "admin_shared".into(), + ..Default::default() + }, + ) + .await + .unwrap(); + + con.put( + &format!("/api/admin/mcp/servers/{id}/shared-credential/static-token"), + json!({"token": "bad"}), + ) + .await + .unwrap() + .assert_ok(); + let row = wait_for_server(&con, &id.to_string(), |s| s["last_error"].is_string()).await; + assert_eq!( + row["last_error"], + "Shared credential rejected by upstream — verify token / scopes" + ); +} + +#[ignore = "integration test — run via `make test-it`"] +#[tokio::test] +async fn shared_authorize_needs_an_admin_shared_oauth_server() { + let app = TestApp::spawn().await; + let con = admin_session(&app).await; + let per_user = fixtures::create_mcp_server( + &app.db, + &unique_name("pu"), + &prefix(), + "https://example.com/mcp", + ) + .await + .unwrap(); + let static_shared = fixtures::create_mcp_server_with( + &app.db, + &unique_name("ss"), + &prefix(), + "https://example.com/mcp", + fixtures::McpServerOpts { + auth_shape: "static".into(), + credential_owner: "admin_shared".into(), + ..Default::default() + }, + ) + .await + .unwrap(); + let oauth_shared = fixtures::create_mcp_server_with( + &app.db, + &unique_name("os"), + &prefix(), + "https://example.com/mcp", + fixtures::McpServerOpts { + auth_shape: "oauth".into(), + credential_owner: "admin_shared".into(), + oauth_authorization_endpoint: Some("https://auth.example.com/authorize".into()), + oauth_token_endpoint: Some("https://auth.example.com/token".into()), + oauth_client_id: Some("cid".into()), + oauth_scopes: vec!["read".into()], + ..Default::default() + }, + ) + .await + .unwrap(); + let authorize = |id: Uuid| format!("/api/admin/mcp/servers/{id}/shared-credential/authorize"); + + con.post(&authorize(Uuid::new_v4()), json!({})) + .await + .unwrap() + .assert_status(404); + con.post(&authorize(per_user), json!({})) + .await + .unwrap() + .assert_status(400); + con.post(&authorize(static_shared), json!({})) + .await + .unwrap() + .assert_status(400); + con.put( + &format!("/api/admin/mcp/servers/{per_user}/shared-credential/static-token"), + json!({"token": "x"}), + ) + .await + .unwrap() + .assert_status(400); + + let resp = con.post(&authorize(oauth_shared), json!({})).await.unwrap(); + resp.assert_ok(); + let url = resp.json::().unwrap()["authorize_url"] + .as_str() + .unwrap() + .to_string(); + assert!( + url.starts_with("https://auth.example.com/authorize?"), + "{url}" + ); + assert!( + url.contains("client_id=cid") && url.contains("scope=read"), + "{url}" + ); +} + +// --------------------------------------------------------------------------- +// Per-user connections +// --------------------------------------------------------------------------- + +#[ignore = "integration test — run via `make test-it`"] +#[tokio::test] +async fn accounts_are_listed_switched_and_revoked() { + let app = TestApp::spawn().await; + let con = admin_session(&app).await; + let upstream = mcp_ok().await; + let id = fixtures::create_mcp_server_with( + &app.db, + &unique_name("conn"), + &prefix(), + &format!("{}/mcp", upstream.uri()), + fixtures::McpServerOpts { + auth_shape: "static".into(), + ..Default::default() + }, + ) + .await + .unwrap(); + for label in ["first", "second", "third"] { + con.put( + &format!("/api/mcp/connections/{id}/{label}/static-token"), + json!({"token": format!("tok-{label}")}), + ) + .await + .unwrap() + .assert_ok(); + } + + let accounts = |list: &Value| -> Vec<(String, bool)> { + let entry = list + .as_array() + .unwrap() + .iter() + .find(|s| s["server_id"] == id.to_string()) + .unwrap_or_else(|| panic!("server missing: {list}")); + entry["accounts"] + .as_array() + .unwrap() + .iter() + .map(|a| { + ( + a["account_label"].as_str().unwrap().to_string(), + a["is_default"].as_bool().unwrap(), + ) + }) + .collect() + }; + // Default first, then by label. + assert_eq!( + accounts(&get(&con, "/api/mcp/connections").await), + vec![ + ("first".to_string(), true), + ("second".to_string(), false), + ("third".to_string(), false) + ] + ); + + con.put( + &format!("/api/mcp/connections/{id}/second/default"), + json!({}), + ) + .await + .unwrap() + .assert_ok(); + assert_eq!( + accounts(&get(&con, "/api/mcp/connections").await), + vec![ + ("second".to_string(), true), + ("first".to_string(), false), + ("third".to_string(), false) + ] + ); + + // Revoking a non-default account leaves the default alone. + con.delete(&format!("/api/mcp/connections/{id}/first")) + .await + .unwrap() + .assert_ok(); + assert_eq!( + accounts(&get(&con, "/api/mcp/connections").await), + vec![("second".to_string(), true), ("third".to_string(), false)] + ); + + con.delete(&format!("/api/mcp/connections/{id}/first")) + .await + .unwrap() + .assert_status(404); + con.put( + &format!("/api/mcp/connections/{id}/missing/default"), + json!({}), + ) + .await + .unwrap() + .assert_status(404); +} + +#[ignore = "integration test — run via `make test-it`"] +#[tokio::test] +async fn connections_list_only_per_user_servers_that_need_a_credential() { + let app = TestApp::spawn().await; + let con = admin_session(&app).await; + let listed = fixtures::create_mcp_server_with( + &app.db, + &unique_name("listed"), + &prefix(), + "https://example.com/mcp", + fixtures::McpServerOpts { + auth_shape: "static".into(), + auth_header_name: "X-API-Key".into(), + auth_value_template: "{{token}}".into(), + ..Default::default() + }, + ) + .await + .unwrap(); + sqlx::query("UPDATE mcp_servers SET display_label = 'Label' WHERE id = $1") + .bind(listed) + .execute(&app.db) + .await + .unwrap(); + let anonymous = fixtures::create_mcp_server( + &app.db, + &unique_name("anon"), + &prefix(), + "https://example.com/mcp", + ) + .await + .unwrap(); + let shared = fixtures::create_mcp_server_with( + &app.db, + &unique_name("shared"), + &prefix(), + "https://example.com/mcp", + fixtures::McpServerOpts { + auth_shape: "static".into(), + credential_owner: "admin_shared".into(), + ..Default::default() + }, + ) + .await + .unwrap(); + + let list = get(&con, "/api/mcp/connections").await; + let rows = list.as_array().unwrap(); + let ids: Vec<&str> = rows + .iter() + .map(|r| r["server_id"].as_str().unwrap()) + .collect(); + assert!(ids.contains(&listed.to_string().as_str()), "{list}"); + assert!(!ids.contains(&anonymous.to_string().as_str()), "{list}"); + assert!(!ids.contains(&shared.to_string().as_str()), "{list}"); + let row = rows + .iter() + .find(|r| r["server_id"] == listed.to_string()) + .unwrap(); + assert_eq!(row["display_label"], "Label"); + assert_eq!(row["auth_shape"], "static"); + assert_eq!(row["auth_header_name"], "X-API-Key"); + assert_eq!(row["auth_value_template"], "{{token}}"); + assert_eq!(row["accounts"], json!([])); + + // Per-user endpoints on a server that doesn't exist. + let ghost = Uuid::new_v4(); + con.put( + &format!("/api/mcp/connections/{ghost}/x/static-token"), + json!({"token": "t"}), + ) + .await + .unwrap() + .assert_status(404); + con.post( + &format!("/api/mcp/connections/{ghost}/authorize"), + json!({"account_label": "x"}), + ) + .await + .unwrap() + .assert_status(404); +} + +// --------------------------------------------------------------------------- +// Store +// --------------------------------------------------------------------------- + +#[ignore = "integration test — run via `make test-it`"] +#[tokio::test] +async fn a_template_is_read_by_slug() { + let app = TestApp::spawn().await; + let con = admin_session(&app).await; + let slug = unique_name("tpl"); + sqlx::query( + "INSERT INTO mcp_store_templates (slug, name, category, endpoint_template, deploy_type) + VALUES ($1, 'Tpl', 'dev', 'https://example.com/mcp', 'hosted')", + ) + .bind(&slug) + .execute(&app.db) + .await + .unwrap(); + + let got = get(&con, &format!("/api/mcp/store/{slug}")).await; + assert_eq!(got["slug"], slug.as_str()); + assert_eq!(got["name"], "Tpl"); + assert_eq!(got["endpoint_template"], "https://example.com/mcp"); + let resp = con + .get(&format!("/api/mcp/store/{}", unique_name("nope"))) + .await + .unwrap(); + resp.assert_status(404); +} + +#[ignore = "integration test — run via `make test-it`"] +#[tokio::test] +async fn the_registry_sync_upserts_prunes_and_keeps_installed_templates() { + let app = TestApp::spawn_reaching_loopback().await; + let (con, admin) = admin_session_with_user(&app).await; + + // Something installed survives a registry that no longer lists it. + let installed_slug = unique_name("kept"); + let installed_id: Uuid = sqlx::query_scalar( + "INSERT INTO mcp_store_templates (slug, name, deploy_type) VALUES ($1, 'Kept', 'hosted') + RETURNING id", + ) + .bind(&installed_slug) + .fetch_one(&app.db) + .await + .unwrap(); + let server = fixtures::create_mcp_server( + &app.db, + &unique_name("inst"), + &prefix(), + "https://example.com/mcp", + ) + .await + .unwrap(); + sqlx::query( + "INSERT INTO mcp_store_installs (template_id, server_id, installed_by) VALUES ($1, $2, $3)", + ) + .bind(installed_id) + .bind(server) + .bind(admin.user.id) + .execute(&app.db) + .await + .unwrap(); + let before: i64 = sqlx::query_scalar("SELECT COUNT(*) FROM mcp_store_templates") + .fetch_one(&app.db) + .await + .unwrap(); + + let registry = MockServer::start().await; + let registry_body = |description: &str| { + json!({ + "version": 1, + "templates": [ + { + "slug": "sync-oauth", + "name": "Sync OAuth", + "description": {"en": description, "zh": "说明"}, + "category": "dev", + "tags": ["x", "y"], + "endpoint_template": "https://oauth.example.com/mcp", + "oauth_issuer": "https://oauth.example.com", + "oauth_default_scopes": ["repo"], + "featured": true + }, + { + "slug": "sync-static", + "name": "Sync Static", + "static_token_help_url": "https://example.com/token", + "auth_header_name": "X-API-Key", + "auth_value_template": "{{token}}", + "auth_instructions": "paste it", + "deploy_type": "manual" + }, + { + "slug": "sync-bad", + "name": "Bad", + "auth_value_template": "Bearer {{nope}}" + } + ] + }) + }; + Mock::given(method("GET")) + .and(path("/registry.json")) + .respond_with(ResponseTemplate::new(200).set_body_json(registry_body("first"))) + .up_to_n_times(1) + .mount(®istry) + .await; + Mock::given(method("GET")) + .and(path("/registry.json")) + .respond_with(ResponseTemplate::new(200).set_body_json(registry_body("second"))) + .mount(®istry) + .await; + let url = format!("{}/registry.json", registry.uri()); + + let resp = con + .post("/api/admin/mcp-store/sync", json!({"registry_url": url})) + .await + .unwrap(); + resp.assert_ok(); + let body: Value = resp.json().unwrap(); + assert_eq!(body["status"], "synced"); + assert_eq!(body["count"], 2, "the invalid template is skipped: {body}"); + // Every other seeded template went; the installed one stayed. + assert_eq!(body["removed"], before - 1, "{body}"); + + let slugs: Vec = + sqlx::query_scalar("SELECT slug FROM mcp_store_templates ORDER BY slug") + .fetch_all(&app.db) + .await + .unwrap(); + let mut want = vec![ + installed_slug.clone(), + "sync-oauth".to_string(), + "sync-static".to_string(), + ]; + want.sort(); + assert_eq!(slugs, want); + + let oauth = get(&con, "/api/mcp/store/sync-oauth").await; + assert_eq!(oauth["description"], "first\n---\n说明"); + assert_eq!(oauth["auth_shape"], "oauth"); + assert_eq!(oauth["tags"], json!(["x", "y"])); + assert_eq!(oauth["oauth_default_scopes"], json!(["repo"])); + assert_eq!(oauth["auth_header_name"], "Authorization"); + assert_eq!(oauth["auth_value_template"], "Bearer {{token}}"); + assert_eq!(oauth["deploy_type"], "hosted"); + assert_eq!(oauth["featured"], true); + let stat = get(&con, "/api/mcp/store/sync-static").await; + assert_eq!(stat["auth_shape"], "static"); + assert_eq!(stat["auth_header_name"], "X-API-Key"); + assert_eq!(stat["auth_value_template"], "{{token}}"); + assert_eq!(stat["auth_instructions"], "paste it"); + assert_eq!(stat["deploy_type"], "manual"); + assert_eq!(stat["tags"], json!([])); + assert_eq!(stat["featured"], false); + + // A second sync updates in place and removes nothing. + let body: Value = con + .post("/api/admin/mcp-store/sync", json!({"registry_url": url})) + .await + .unwrap() + .json() + .unwrap(); + assert_eq!(body["count"], 2); + assert_eq!(body["removed"], 0); + let oauth = get(&con, "/api/mcp/store/sync-oauth").await; + assert_eq!(oauth["description"], "second\n---\n说明"); + + let cats = get(&con, "/api/mcp/store/categories").await; + assert_eq!(cats, json!([{"category": "dev", "count": 1}])); +} + +// --------------------------------------------------------------------------- +// Tool catalog +// --------------------------------------------------------------------------- + +#[ignore = "integration test — run via `make test-it`"] +#[tokio::test] +async fn the_tool_catalog_filters_pages_and_includes_the_callers_own_tools() { + let app = TestApp::spawn().await; + let (con, admin) = admin_session_with_user(&app).await; + let other = fixtures::create_admin_user(&app.db).await.unwrap(); + + let pfx_a = prefix(); + let a = fixtures::create_mcp_server(&app.db, "cat-a", &pfx_a, "https://example.com/mcp") + .await + .unwrap(); + let b = fixtures::create_mcp_server(&app.db, "cat-b", &prefix(), "https://example.com/mcp") + .await + .unwrap(); + insert_tool(&app, a, "alpha", "first tool", true).await; + insert_tool(&app, a, "beta", "finds needles", true).await; + insert_tool(&app, a, "hidden", "", false).await; + insert_tool(&app, b, "gamma", "", true).await; + for (user, tool) in [(admin.user.id, "mine"), (other.user.id, "theirs")] { + sqlx::query( + "INSERT INTO mcp_user_tools (mcp_server_id, user_id, tool_name, description, input_schema) + VALUES ($1, $2, $3, 'personal', '{}'::jsonb)", + ) + .bind(b) + .bind(user) + .bind(tool) + .execute(&app.db) + .await + .unwrap(); + } + + let names = |v: &Value| -> Vec { + v["items"] + .as_array() + .unwrap() + .iter() + .map(|t| t["name"].as_str().unwrap().to_string()) + .collect() + }; + + let all = get(&con, "/api/mcp/tools").await; + assert_eq!(all["total"], 3, "{all}"); + assert_eq!(names(&all), ["alpha", "beta", "gamma"]); + let alpha = &all["items"][0]; + assert_eq!(alpha["server_name"], "cat-a"); + assert_eq!(alpha["server_id"], a.to_string()); + assert_eq!(alpha["namespaced_name"], format!("{pfx_a}__alpha")); + assert_eq!(alpha["description"], "first tool"); + assert_eq!(alpha["input_schema"], json!({})); + + let mine = get(&con, "/api/mcp/tools?include_user_tools=true").await; + assert_eq!(mine["total"], 4); + assert_eq!(names(&mine), ["alpha", "beta", "gamma", "mine"]); + + let by_server = get(&con, &format!("/api/mcp/tools?server_id={b}")).await; + assert_eq!(names(&by_server), ["gamma"]); + + // Search matches the name, the namespaced name and the description. + assert_eq!(names(&get(&con, "/api/mcp/tools?q=ALP").await), ["alpha"]); + assert_eq!(names(&get(&con, "/api/mcp/tools?q=needle").await), ["beta"]); + let by_prefix = get(&con, &format!("/api/mcp/tools?q={pfx_a}__b")).await; + assert_eq!(names(&by_prefix), ["beta"]); + + let page = get(&con, "/api/mcp/tools?page=2&page_size=2").await; + assert_eq!(page["total"], 3); + assert_eq!(names(&page), ["gamma"]); +}