diff --git a/.github/workflows/test-migrations.yml b/.github/workflows/test-migrations.yml index 104429cc33..4cf968fee2 100644 --- a/.github/workflows/test-migrations.yml +++ b/.github/workflows/test-migrations.yml @@ -52,6 +52,11 @@ on: - 'tests/web/test_trusted_actor_oauth_postgresql.py' - 'tests/shared/postgres_disposable.py' - 'tests/web/services/checkpoint_anchor_shared.py' + - 'src/xagent/web/api/mcp.py' + - 'src/xagent/web/models/mcp.py' + - 'src/xagent/web/models/mcp_oauth.py' + - 'src/xagent/web/models/public_mcp.py' + - 'tests/web/api/test_mcp_app_teardown_postgresql.py' pull_request: branches: [main] # Required by the merge queue: without this the two required contexts below @@ -140,6 +145,11 @@ jobs: tests/web/test_trusted_actor_oauth_postgresql.py tests/shared/postgres_disposable.py tests/web/services/checkpoint_anchor_shared.py + src/xagent/web/api/mcp.py + src/xagent/web/models/mcp.py + src/xagent/web/models/mcp_oauth.py + src/xagent/web/models/public_mcp.py + tests/web/api/test_mcp_app_teardown_postgresql.py ) case "$EVENT_NAME" in @@ -346,6 +356,13 @@ jobs: env: XAGENT_TEST_POSTGRES_URL: postgresql://xagent:xagent@localhost:5432/xagent_test + - name: Test app-scoped MCP teardown concurrency (Postgres-only) + if: needs.detect-migration-changes.outputs.should-test == 'true' + run: | + pytest tests/web/api/test_mcp_app_teardown_postgresql.py -m postgresql -q + env: + XAGENT_TEST_POSTGRES_URL: postgresql://xagent:xagent@localhost:5432/xagent_test + - name: Test Gmail provisioning Postgres-only behavior if: needs.detect-migration-changes.outputs.should-test == 'true' run: | diff --git a/src/xagent/web/api/mcp.py b/src/xagent/web/api/mcp.py index 501293feb6..5a54ec6d9d 100644 --- a/src/xagent/web/api/mcp.py +++ b/src/xagent/web/api/mcp.py @@ -12,7 +12,7 @@ import logging import secrets import shlex -from collections.abc import Collection +from collections.abc import Collection, Mapping from dataclasses import dataclass from datetime import datetime, timedelta, timezone from typing import Annotated, Any, Callable, Dict, List, Literal, Optional, Union, cast @@ -22,6 +22,7 @@ from fastapi import APIRouter, Depends, Header, HTTPException, Request, status from fastapi.responses import JSONResponse, RedirectResponse, Response from pydantic import BaseModel, ConfigDict, Field +from sqlalchemy import text from sqlalchemy.exc import IntegrityError from sqlalchemy.orm import Session @@ -50,6 +51,7 @@ mcp_oauth_client_registration_lookup_hash, mcp_oauth_grant_lookup_hash, ) +from ..models.public_mcp import PublicMCPApp from ..models.user import User from ..services.mcp_oauth import ( MCP_OAUTH_HTTP_TIMEOUT_SECONDS, @@ -492,6 +494,19 @@ def _mcp_oauth_callback_error_redirect( raw_redirect_after = ( str(flow_state.redirect_after) if flow_state.redirect_after else None ) + return _mcp_oauth_callback_error_redirect_for_path( + raw_redirect_after, + error_code=error_code, + message=message, + ) + + +def _mcp_oauth_callback_error_redirect_for_path( + raw_redirect_after: str | None, + *, + error_code: str, + message: str, +) -> RedirectResponse: redirect_path = _redirect_after_with_params( raw_redirect_after, ( @@ -965,37 +980,80 @@ async def _exchange_mcp_oauth_code( return payload -async def _revoke_mcp_oauth_grant_externally( +@dataclass(frozen=True) +class _MCPOAuthRevocationSnapshot: + grant_id: int + metadata: dict[str, Any] + client_id: str + client_secret: str | None + token_endpoint_auth_method: str + access_token: str + refresh_token: str | None + + +def _mcp_oauth_revocation_snapshot( + *, client: MCPOAuthClient, grant: MCPOAuthGrant +) -> _MCPOAuthRevocationSnapshot: + """Capture the encrypted values needed after a teardown commits.""" + return _MCPOAuthRevocationSnapshot( + grant_id=int(grant.id), + metadata=( + dict(client.metadata_json) if isinstance(client.metadata_json, dict) else {} + ), + client_id=str(client.client_id), + client_secret=(str(client.client_secret) if client.client_secret else None), + token_endpoint_auth_method=str(client.token_endpoint_auth_method or "none"), + access_token=str(grant.access_token), + refresh_token=(str(grant.refresh_token) if grant.refresh_token else None), + ) + + +def _mcp_oauth_token_revocation_snapshot( *, client: MCPOAuthClient, - grant: MCPOAuthGrant, -) -> None: - metadata: dict[str, Any] = ( - client.metadata_json if isinstance(client.metadata_json, dict) else {} + token_data: dict[str, Any], + reference_id: int, +) -> _MCPOAuthRevocationSnapshot: + """Capture a just-issued token before lifecycle revalidation.""" + refresh_token = token_data.get("refresh_token") + return _MCPOAuthRevocationSnapshot( + grant_id=reference_id, + metadata=( + dict(client.metadata_json) if isinstance(client.metadata_json, dict) else {} + ), + client_id=str(client.client_id), + client_secret=(str(client.client_secret) if client.client_secret else None), + token_endpoint_auth_method=str(client.token_endpoint_auth_method or "none"), + access_token=encrypt_value(str(token_data["access_token"])), + refresh_token=(encrypt_value(str(refresh_token)) if refresh_token else None), ) - revocation_endpoint = metadata.get("revocation_endpoint") + + +async def _revoke_mcp_oauth_snapshot_externally( + snapshot: _MCPOAuthRevocationSnapshot, +) -> None: + revocation_endpoint = snapshot.metadata.get("revocation_endpoint") if not isinstance(revocation_endpoint, str) or not revocation_endpoint: return try: client_secret = ( - decrypt_value(str(client.client_secret)) if client.client_secret else "" + decrypt_value(snapshot.client_secret) if snapshot.client_secret else "" ) - except Exception as exc: + except Exception: logger.warning( "Skipping MCP OAuth token revocation for grant %s because client secret " - "could not be decrypted: %s", - grant.id, - exc, + "could not be decrypted", + snapshot.grant_id, ) return - auth_method = str(client.token_endpoint_auth_method or "none") + auth_method = snapshot.token_endpoint_auth_method auth: httpx.Auth | None = None - base_data: dict[str, str] = {"client_id": str(client.client_id)} + base_data: dict[str, str] = {"client_id": snapshot.client_id} if auth_method == "client_secret_post" and client_secret: base_data["client_secret"] = client_secret elif auth_method == "client_secret_basic" and client_secret: - auth = httpx.BasicAuth(str(client.client_id), client_secret) + auth = httpx.BasicAuth(snapshot.client_id, client_secret) elif auth_method not in {"none", "client_secret_post", "client_secret_basic"}: logger.warning( "Skipping MCP OAuth token revocation for unsupported auth method %s", @@ -1004,8 +1062,8 @@ async def _revoke_mcp_oauth_grant_externally( return encrypted_tokens = ( - (grant.access_token, "access_token"), - (grant.refresh_token, "refresh_token"), + (snapshot.access_token, "access_token"), + (snapshot.refresh_token, "refresh_token"), ) async with create_mcp_oauth_http_client( timeout=MCP_OAUTH_HTTP_TIMEOUT_SECONDS, @@ -1015,13 +1073,12 @@ async def _revoke_mcp_oauth_grant_externally( continue try: decrypted_token = decrypt_value(str(encrypted_token)) - except Exception as exc: + except Exception: logger.warning( "Skipping MCP OAuth %s revocation for grant %s because token " - "could not be decrypted: %s", + "could not be decrypted", token_type_hint, - grant.id, - exc, + snapshot.grant_id, ) continue data = { @@ -1045,16 +1102,40 @@ async def _revoke_mcp_oauth_grant_externally( logger.warning( "MCP OAuth token revocation returned HTTP %s for grant %s", response.status_code, - grant.id, + snapshot.grant_id, ) - except (MCPOAuthDiscoveryError, httpx.HTTPError) as exc: + except (MCPOAuthDiscoveryError, httpx.HTTPError): logger.warning( - "MCP OAuth token revocation failed for grant %s: %s", - grant.id, - exc, + "MCP OAuth token revocation failed for grant %s", + snapshot.grant_id, ) +async def _revoke_failed_mcp_oauth_callback_token( + snapshot: _MCPOAuthRevocationSnapshot, + *, + flow_id: int, +) -> None: + """Best-effort revoke a token that could not be persisted locally.""" + try: + await _revoke_mcp_oauth_snapshot_externally(snapshot) + except Exception: + logger.warning( + "MCP OAuth token revocation failed after callback persistence for flow %s", + flow_id, + ) + + +async def _revoke_mcp_oauth_grant_externally( + *, + client: MCPOAuthClient, + grant: MCPOAuthGrant, +) -> None: + await _revoke_mcp_oauth_snapshot_externally( + _mcp_oauth_revocation_snapshot(client=client, grant=grant) + ) + + def _upsert_mcp_oauth_grant( db: Session, *, @@ -3530,6 +3611,517 @@ def _catalog_server_has_platform_key(db: Session, server: MCPServer) -> bool: return False +def _lock_catalog_for_app_teardown(db: Session) -> None: + """Serialize catalog ownership changes for one destructive teardown. + + Locking only the expected catalog row is insufficient for legacy MCP + servers whose owner is encoded in ``MCPServer.name``: another catalog row + can be inserted or renamed onto that name without touching the expected + row. PostgreSQL's table-level SHARE lock excludes those catalog writes for + the rest of this transaction. SQLite needs a real write-intent transaction + before these reads because pysqlite's legacy transaction mode does not emit + BEGIN for SELECT and SQLite compiles FOR UPDATE away. + """ + dialect = db.get_bind().dialect.name + if dialect == "postgresql": + db.execute(text("LOCK TABLE public_mcp_apps IN SHARE MODE")) + elif dialect == "sqlite": + _begin_sqlite_write_intent(db) + + +class _SQLiteLifecycleTransactionError(RuntimeError): + """The caller already owns SQLite writes that this operation cannot reset.""" + + +def _begin_sqlite_write_intent(db: Session) -> None: + """Start an owned SQLite write transaction without discarding caller writes.""" + if db.new or db.dirty or db.deleted: + raise _SQLiteLifecycleTransactionError( + "SQLite lifecycle fencing requires a read-only preflight" + ) + + connection = db.connection() + driver_connection = connection.connection.driver_connection + if bool(getattr(driver_connection, "in_transaction", False)): + # A real DBAPI transaction may already contain flushed writes that the + # ORM collections above cannot see. Never roll it back or commit it on + # the caller's behalf. + raise _SQLiteLifecycleTransactionError( + "SQLite lifecycle fencing requires an owned transaction" + ) + + # SQLAlchemy has logically autobegun for the caller's read-only preflight, + # even though pysqlite has not sent BEGIN. End that logical transaction, + # then acquire SQLite's single-writer reservation before any identity read. + db.rollback() + connection = db.connection() + driver_connection = connection.connection.driver_connection + if bool(getattr(driver_connection, "in_transaction", False)): + raise RuntimeError("SQLite lifecycle fencing could not reset preflight") + connection.exec_driver_sql("BEGIN IMMEDIATE") + + +@dataclass(frozen=True) +class _MCPOAuthFlowIdentity: + id: int + state: str + server_id: int + user_id: int + client_id: int + + +def _lock_active_mcp_oauth_lifecycle( + db: Session, + *, + server_id: int, + user_id: int, + flow_identity: _MCPOAuthFlowIdentity | None = None, +) -> tuple[MCPServer, UserMCPServer, MCPOAuthFlowState | None] | None: + """Fence final OAuth persistence against disconnect. + + The shared lock order is server, current-user association, then flow. No + caller may invoke this helper until all provider network I/O has completed. + """ + if db.get_bind().dialect.name == "sqlite": + _begin_sqlite_write_intent(db) + db.expire_all() + + server = ( + db.query(MCPServer) + .filter(MCPServer.id == server_id) + .with_for_update() + .one_or_none() + ) + if server is None: + return None + association = ( + db.query(UserMCPServer) + .filter( + UserMCPServer.user_id == user_id, + UserMCPServer.mcpserver_id == server_id, + UserMCPServer.is_active.is_(True), + ) + .with_for_update() + .one_or_none() + ) + if association is None: + return None + + flow_state: MCPOAuthFlowState | None = None + if flow_identity is not None: + flow_state = ( + db.query(MCPOAuthFlowState) + .filter( + MCPOAuthFlowState.id == flow_identity.id, + MCPOAuthFlowState.state == flow_identity.state, + MCPOAuthFlowState.mcp_server_id == flow_identity.server_id, + MCPOAuthFlowState.user_id == flow_identity.user_id, + MCPOAuthFlowState.mcp_oauth_client_id == flow_identity.client_id, + ) + .with_for_update() + .one_or_none() + ) + if flow_state is None: + return None + return server, association, flow_state + + +def _persist_mcp_oauth_connect_flow( + db: Session, + *, + server_id: int, + user_id: int, + discovery: Any, + client_id: str, + client_secret: str | None, + token_endpoint_auth_method: str, + redirect_uri: str, + registration_lookup_hash: str | None, + resource_owner_key: str, + selected_issuer: str, + selected_resource: str, + selected_scope: str, + redirect_after: str | None, +) -> tuple[MCPOAuthClient, str, str] | None: + """Persist a connect flow only while its original lifecycle remains active.""" + lifecycle = _lock_active_mcp_oauth_lifecycle( + db, + server_id=server_id, + user_id=user_id, + ) + if lifecycle is None: + db.rollback() + return None + + try: + oauth_client = _upsert_mcp_oauth_client( + db, + server_id=server_id, + discovery=discovery, + client_id=client_id, + client_secret=client_secret, + token_endpoint_auth_method=token_endpoint_auth_method, + redirect_uri=redirect_uri, + registration_lookup_hash=registration_lookup_hash, + ) + + # Sweep this table's dead rows before adding another, mirroring the Slack + # OAuth flow-state ledger in channel.py. Every abandoned, denied or + # double-submitted authorization leaves a row that is permanently unusable + # once it expires (the claim query requires consumed_at IS NULL and + # expires_at > now), but nothing deleted it: the only other deletes are the + # per-user purge on disconnect and the FK cascade on server deletion, so + # the table grew without bound. Deliberately global rather than scoped to + # this user/server — a user who never reconnects would otherwise keep their + # rows forever, which is the leak this is meant to close. Filtering on + # expires_at alone (an indexed column) also covers consumed rows, since + # every row expires within MCP_OAUTH_STATE_TTL of being created. + # + # Bounded per request. The index narrows the scan but does not cap the work, + # and this table is precisely the one that has never been swept — the first + # connect after this ships meets whatever backlog has accumulated, inside a + # user-facing transaction. Draining a fixed batch per connect keeps that + # transaction short; the remainder is already dead, so later connects + # clearing it costs nothing. Deleting by a bounded id subquery rather than + # an IN list of fetched ids keeps it to one statement and clear of any + # bound-parameter limit. + stale_flow_state_ids = ( + db.query(MCPOAuthFlowState.id) + .filter( + MCPOAuthFlowState.expires_at + < _utc_now() - MCP_OAUTH_FLOW_STATE_RETENTION + ) + .limit(MCP_OAUTH_FLOW_STATE_SWEEP_BATCH) + .scalar_subquery() + ) + db.query(MCPOAuthFlowState).filter( + MCPOAuthFlowState.id.in_(stale_flow_state_ids) + ).delete(synchronize_session=False) + + state_value = secrets.token_urlsafe(32) + code_verifier = secrets.token_urlsafe(64) + flow_state = MCPOAuthFlowState( + state=state_value, + mcp_server_id=server_id, + user_id=user_id, + mcp_oauth_client_id=oauth_client.id, + resource_owner_key=resource_owner_key, + issuer=selected_issuer, + resource=selected_resource, + scope=selected_scope, + code_verifier=encrypt_value(code_verifier), + redirect_after=_safe_mcp_oauth_redirect_after(redirect_after), + expires_at=_utc_now() + MCP_OAUTH_STATE_TTL, + ) + db.add(flow_state) + db.commit() + return oauth_client, state_value, code_verifier + except Exception: + db.rollback() + raise + + +def _locked_catalog_app_for_server( + db: Session, + *, + server: MCPServer, + expected_app: PublicMCPApp, +) -> PublicMCPApp | None: + """Return the locked catalog row when the server has that exact owner.""" + app_id = str(expected_app.app_id) + if str(server.transport or "").lower() != "oauth": + # Catalog provisioners name non-builtin and remote-MCP-OAuth rows by + # exact app id. Their caller-authored auth blob is not an owner stamp. + return expected_app if str(server.name or "") == app_id else None + + auth = server.auth + if isinstance(auth, Mapping) and "app_id" in auth: + return expected_app if auth.get("app_id") == app_id else None + + # Legacy builtin OAuth rows can be named by exact app id or mutable display + # name. Both namespaces are legal, so every matching row must identify the + # same owner. The catalog table lock above makes this set stable until + # commit; without it a concurrent rename/reassignment is a phantom. + server_name = str(server.name or "") + owners = { + str(candidate.app_id) + for candidate in db.query(PublicMCPApp) + .filter( + (PublicMCPApp.app_id == server_name) | (PublicMCPApp.name == server_name) + ) + .all() + } + return expected_app if owners == {app_id} else None + + +async def teardown_mcp_app_server( + server_id: int, + *, + app_id: str, + expected_catalog_app_id: int, + expected_provider_name: str | None, + current_user: User, + db: Session, +) -> None: + """Atomically disconnect one exact catalog app for one user. + + ``expected_catalog_app_id`` is the immutable primary key observed by the + caller's catalog preflight. The exact string ``app_id`` alone cannot + distinguish that row from a delete-and-recreate using the same public id. + The caller must also pin the catalog row's exact ``provider_name`` because + that mutable value selects which user credential is deleted. All three + values are revalidated under catalog/server locks before any local + credential, grant, flow, association, or server row is changed. + + A missing association/server is a 404 idempotency race. A missing, + replaced, or ambiguous catalog owner is a 403 fail-closed refusal. Other + failures are sanitized as 500. All local deletion is committed once, so a + failed final server delete rolls every preceding cleanup back for retry. + """ + revocations: list[_MCPOAuthRevocationSnapshot] = [] + try: + with db.no_autoflush: + user_id = int(current_user.id) + current_user_is_admin = bool(getattr(current_user, "is_admin", False)) + if ( + not isinstance(app_id, str) + or not app_id + or app_id != app_id.strip() + or isinstance(expected_catalog_app_id, bool) + or not isinstance(expected_catalog_app_id, int) + or ( + expected_provider_name is not None + and not isinstance(expected_provider_name, str) + ) + ): + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail="MCP app teardown identity could not be verified", + ) + + _lock_catalog_for_app_teardown(db) + # The SaaS preflight deliberately uses this same Session and may have + # materialized catalog/server rows before a concurrent admin commit. + # Row/table locks serialize future writes but do not refresh SQLAlchemy's + # identity map, so expire it before the locked revalidation reads. + db.expire_all() + with db.no_autoflush: + expected_app = ( + db.query(PublicMCPApp) + .filter( + PublicMCPApp.id == expected_catalog_app_id, + PublicMCPApp.app_id == app_id, + PublicMCPApp.provider_name == expected_provider_name, + ) + .with_for_update() + .one_or_none() + ) + if expected_app is None: + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail="MCP app teardown owner changed", + ) + + server = ( + db.query(MCPServer) + .filter(MCPServer.id == server_id) + .with_for_update() + .one_or_none() + ) + if server is None: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail="MCP server not found", + ) + + user_mcp = ( + db.query(UserMCPServer) + .filter( + UserMCPServer.user_id == user_id, + UserMCPServer.mcpserver_id == server_id, + ) + .with_for_update() + .one_or_none() + ) + if user_mcp is None: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail="MCP server not found", + ) + + validated_app = _locked_catalog_app_for_server( + db, server=server, expected_app=expected_app + ) + if validated_app is None: + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail="MCP app teardown owner changed", + ) + + if not _check_mcp_permission(user_mcp, current_user_is_admin, require="delete"): + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail="You do not have permission to delete this MCP server", + ) + + from ..services.connector_team_scope import delete_team_connector + + team_delete = delete_team_connector(db, user_id, "mcp", server_id) + if team_delete.blocked_reason: + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail=team_delete.blocked_reason, + ) + if team_delete.team_owned and not team_delete.authorized: + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail="Only a team admin can delete a team MCP server", + ) + + if str(server.transport or "").lower() == "oauth": + # Use the row revalidated under the catalog/server locks above. + # Calling the generic resolver again would re-read mutable catalog + # state and could select a provider different from the preflight. + provider = validated_app.provider_name + providers_to_delete = restrict_to_app_scoped_oauth_grant( + app_id, [provider, app_id] + ) + if providers_to_delete: + delete_scoped_user_oauth_accounts( + db, + user_id=user_id, + resource_owner_key=None, + providers=providers_to_delete, + ) + + if provider and provider not in providers_to_delete: + other_servers = ( + db.query(MCPServer) + .join( + UserMCPServer, + UserMCPServer.mcpserver_id == MCPServer.id, + ) + .filter( + UserMCPServer.user_id == user_id, + MCPServer.id != server_id, + ) + .all() + ) + normalized_provider = _normalize_app_key(provider) + sibling_still_connected = any( + (sibling_app := get_app_for_mcp_server(db, other_server)) + and _normalize_app_key(sibling_app.get("provider")) + == normalized_provider + for other_server in other_servers + ) + if not sibling_still_connected: + delete_scoped_user_oauth_accounts( + db, + user_id=user_id, + resource_owner_key=None, + providers=[provider], + ) + + for grant in ( + db.query(MCPOAuthGrant) + .filter( + MCPOAuthGrant.mcp_server_id == server_id, + MCPOAuthGrant.user_id == user_id, + ) + .with_for_update() + .all() + ): + if str(grant.status) == "active" and isinstance( + grant.oauth_client, MCPOAuthClient + ): + revocations.append( + _mcp_oauth_revocation_snapshot( + client=grant.oauth_client, grant=grant + ) + ) + db.delete(grant) + + ( + db.query(MCPOAuthFlowState) + .filter( + MCPOAuthFlowState.mcp_server_id == server_id, + MCPOAuthFlowState.user_id == user_id, + ) + .delete(synchronize_session=False) + ) + db.delete(user_mcp) + db.flush() + + other_users = ( + db.query(UserMCPServer) + .filter(UserMCPServer.mcpserver_id == server_id) + .with_for_update() + .first() + ) + if other_users is None: + if team_delete.team_owned and not team_delete.delete_definition: + logger.info( + "Kept team-owned MCP server %s after app teardown", server_id + ) + elif _catalog_server_has_platform_key(db, server): + logger.info( + "Kept platform-key MCP server %s after app teardown", server_id + ) + else: + # Delete the locked identity directly. Database FK cascades + # remove MCP OAuth clients (including client_secret), grants, + # flows, and associations in this same transaction. + db.delete(server) + + db.commit() + # Network revocation is deliberately outside the locked transaction. + # It is best-effort, just like the generic teardown, while the local + # encrypted grant/client material is already durably gone. A provider + # timeout therefore cannot hold catalog/server locks or turn a complete + # local teardown back into a retryable partial commit. + for revocation in revocations: + try: + await _revoke_mcp_oauth_snapshot_externally(revocation) + except Exception: + logger.warning( + "MCP OAuth token revocation failed after teardown for grant %s", + revocation.grant_id, + ) + logger.info( + "Completed app-scoped MCP teardown for app %r, server %s, user %s", + app_id, + server_id, + user_id, + ) + except HTTPException: + db.rollback() + raise + except _SQLiteLifecycleTransactionError: + # The caller's pending/flushed writes are explicitly not ours to roll + # back. No teardown write or lock has happened when this error is raised. + logger.error( + "App-scoped MCP teardown requires an owned SQLite transaction for " + "app %r, server %s", + app_id, + server_id, + ) + raise HTTPException( + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, + detail="Failed to delete MCP server", + ) from None + except Exception: + db.rollback() + logger.error( + "App-scoped MCP teardown failed for app %r, server %s", + app_id, + server_id, + ) + raise HTTPException( + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, + detail="Failed to delete MCP server", + ) from None + + @mcp_router.delete("/servers/{server_id}", status_code=status.HTTP_204_NO_CONTENT) async def delete_mcp_server( server_id: int, @@ -4114,6 +4706,16 @@ async def mcp_oauth_callback( status_code=status.HTTP_400_BAD_REQUEST, detail={"code": "invalid_state", "message": "Invalid OAuth state"}, ) + flow_identity = _MCPOAuthFlowIdentity( + id=int(flow_state.id), + state=str(flow_state.state), + server_id=int(flow_state.mcp_server_id), + user_id=int(flow_state.user_id), + client_id=int(flow_state.mcp_oauth_client_id), + ) + callback_redirect_after = ( + str(flow_state.redirect_after) if flow_state.redirect_after else None + ) try: _validate_mcp_oauth_state_cookie(request, state_value) except HTTPException as exc: @@ -4131,8 +4733,8 @@ async def mcp_oauth_callback( state_error = _mcp_oauth_flow_state_error(db, flow_state) if state_error is not None: error_code, message = state_error - return _mcp_oauth_callback_error_redirect( - flow_state, + return _mcp_oauth_callback_error_redirect_for_path( + callback_redirect_after, error_code=error_code, message=message, ) @@ -4170,25 +4772,26 @@ async def mcp_oauth_callback( claim_error = _claim_mcp_oauth_flow_state(db, flow_state) if claim_error is not None: error_code, message = claim_error - return _mcp_oauth_callback_error_redirect( - flow_state, + return _mcp_oauth_callback_error_redirect_for_path( + callback_redirect_after, error_code=error_code, message=message, ) if error: - return _mcp_oauth_callback_error_redirect( - flow_state, + return _mcp_oauth_callback_error_redirect_for_path( + callback_redirect_after, error_code="token_exchange_failed", message=oauth_error_message( {"error": error}, "MCP OAuth authorization failed" ), ) if not code: - return _mcp_oauth_callback_error_redirect( - flow_state, + return _mcp_oauth_callback_error_redirect_for_path( + callback_redirect_after, error_code="invalid_state", message="Missing authorization code", ) + failed_callback_revocation: _MCPOAuthRevocationSnapshot | None = None try: token_data = await _exchange_mcp_oauth_code( client=client, @@ -4196,23 +4799,59 @@ async def mcp_oauth_callback( code_verifier=decrypt_value(str(flow_state.code_verifier)), resource=str(flow_state.resource), ) - _upsert_mcp_oauth_grant(db, flow_state=flow_state, token_data=token_data) + failed_callback_revocation = _mcp_oauth_token_revocation_snapshot( + client=client, + token_data=token_data, + reference_id=flow_identity.id, + ) + lifecycle = _lock_active_mcp_oauth_lifecycle( + db, + server_id=flow_identity.server_id, + user_id=flow_identity.user_id, + flow_identity=flow_identity, + ) + if lifecycle is None: + db.rollback() + await _revoke_failed_mcp_oauth_callback_token( + failed_callback_revocation, + flow_id=flow_identity.id, + ) + failed_callback_revocation = None + return _mcp_oauth_callback_error_redirect_for_path( + callback_redirect_after, + error_code="invalid_state", + message="OAuth connection changed while authorization completed", + ) + locked_flow_state = lifecycle[2] + if locked_flow_state is None: + raise RuntimeError("MCP OAuth callback lifecycle omitted its flow") + _upsert_mcp_oauth_grant(db, flow_state=locked_flow_state, token_data=token_data) db.commit() except HTTPException as exc: db.rollback() + if failed_callback_revocation is not None: + await _revoke_failed_mcp_oauth_callback_token( + failed_callback_revocation, + flow_id=flow_identity.id, + ) detail: dict[str, Any] = exc.detail if isinstance(exc.detail, dict) else {} error_code = str(detail.get("code") or "token_exchange_failed") message = str(detail.get("message") or "MCP OAuth authorization failed") - return _mcp_oauth_callback_error_redirect( - flow_state, + return _mcp_oauth_callback_error_redirect_for_path( + callback_redirect_after, error_code=error_code, message=message, ) except Exception: db.rollback() - logger.exception("MCP OAuth callback failed after state claim") - return _mcp_oauth_callback_error_redirect( - flow_state, + if failed_callback_revocation is not None: + await _revoke_failed_mcp_oauth_callback_token( + failed_callback_revocation, + flow_id=flow_identity.id, + ) + logger.error("MCP OAuth callback failed after state claim") + return _mcp_oauth_callback_error_redirect_for_path( + callback_redirect_after, error_code="token_exchange_failed", message="MCP OAuth authorization failed", ) @@ -4224,7 +4863,7 @@ async def mcp_oauth_callback( response = RedirectResponse( _mcp_oauth_redirect_after_url( _redirect_after_with_params( - str(flow_state.redirect_after) if flow_state.redirect_after else None, + callback_redirect_after, (("mcp_oauth_success", "1"),), ) ) @@ -4343,66 +4982,31 @@ async def connect_mcp_oauth( str(discovery.resource), field_name="resource" ) - oauth_client = _upsert_mcp_oauth_client( + persisted_flow = _persist_mcp_oauth_connect_flow( db, server_id=server_id, + user_id=user_id, discovery=discovery, client_id=client_id, client_secret=client_secret, token_endpoint_auth_method=token_endpoint_auth_method, redirect_uri=redirect_uri, registration_lookup_hash=registration_lookup_hash, - ) - - # Sweep this table's dead rows before adding another, mirroring the Slack - # OAuth flow-state ledger in channel.py. Every abandoned, denied or - # double-submitted authorization leaves a row that is permanently unusable - # once it expires (the claim query requires consumed_at IS NULL and - # expires_at > now), but nothing deleted it: the only other deletes are the - # per-user purge on disconnect and the FK cascade on server deletion, so - # the table grew without bound. Deliberately global rather than scoped to - # this user/server — a user who never reconnects would otherwise keep their - # rows forever, which is the leak this is meant to close. Filtering on - # expires_at alone (an indexed column) also covers consumed rows, since - # every row expires within MCP_OAUTH_STATE_TTL of being created. - # - # Bounded per request. The index narrows the scan but does not cap the work, - # and this table is precisely the one that has never been swept — the first - # connect after this ships meets whatever backlog has accumulated, inside a - # user-facing transaction. Draining a fixed batch per connect keeps that - # transaction short; the remainder is already dead, so later connects - # clearing it costs nothing. Deleting by a bounded id subquery rather than - # an IN list of fetched ids keeps it to one statement and clear of any - # bound-parameter limit. - stale_flow_state_ids = ( - db.query(MCPOAuthFlowState.id) - .filter( - MCPOAuthFlowState.expires_at < _utc_now() - MCP_OAUTH_FLOW_STATE_RETENTION - ) - .limit(MCP_OAUTH_FLOW_STATE_SWEEP_BATCH) - .scalar_subquery() - ) - db.query(MCPOAuthFlowState).filter( - MCPOAuthFlowState.id.in_(stale_flow_state_ids) - ).delete(synchronize_session=False) - - state_value = secrets.token_urlsafe(32) - code_verifier = secrets.token_urlsafe(64) - flow_state = MCPOAuthFlowState( - state=state_value, - mcp_server_id=server_id, - user_id=user_id, - mcp_oauth_client_id=oauth_client.id, resource_owner_key=resource_owner_key, - issuer=selected_issuer, - resource=selected_resource, - scope=selected_scope, - code_verifier=encrypt_value(code_verifier), - redirect_after=_safe_mcp_oauth_redirect_after(request_data.redirect_after), - expires_at=_utc_now() + MCP_OAUTH_STATE_TTL, + selected_issuer=selected_issuer, + selected_resource=selected_resource, + selected_scope=selected_scope, + redirect_after=request_data.redirect_after, ) - db.add(flow_state) - db.commit() + if persisted_flow is None: + raise HTTPException( + status_code=status.HTTP_409_CONFLICT, + detail={ + "code": "oauth_lifecycle_changed", + "message": "MCP server connection changed during OAuth setup", + }, + ) + oauth_client, state_value, code_verifier = persisted_flow params = { "response_type": "code", diff --git a/tests/web/api/test_mcp_app_teardown.py b/tests/web/api/test_mcp_app_teardown.py new file mode 100644 index 0000000000..0ec9d0ff71 --- /dev/null +++ b/tests/web/api/test_mcp_app_teardown.py @@ -0,0 +1,742 @@ +from __future__ import annotations + +import asyncio +import logging +import threading +from concurrent.futures import ThreadPoolExecutor +from datetime import datetime, timedelta, timezone + +import pytest +from fastapi import HTTPException +from sqlalchemy import create_engine, event +from sqlalchemy.orm import Session, sessionmaker +from sqlalchemy.pool import StaticPool + +from xagent.core.utils.encryption import encrypt_value +from xagent.db.sqlite import apply_sqlite_concurrency_pragmas +from xagent.web.api import mcp as mcp_api +from xagent.web.models.database import Base +from xagent.web.models.mcp import MCPServer, UserMCPServer +from xagent.web.models.mcp_oauth import ( + MCPOAuthClient, + MCPOAuthFlowState, + MCPOAuthGrant, +) +from xagent.web.models.public_mcp import PublicMCPApp +from xagent.web.models.user import User +from xagent.web.models.user_oauth import UserOAuth + + +@pytest.fixture +def db() -> Session: + engine = create_engine( + "sqlite://", + connect_args={"check_same_thread": False}, + poolclass=StaticPool, + ) + + @event.listens_for(engine, "connect") + def enable_foreign_keys(dbapi_connection, _connection_record) -> None: + cursor = dbapi_connection.cursor() + cursor.execute("PRAGMA foreign_keys=ON") + cursor.close() + + Base.metadata.create_all(engine) + with Session(engine) as session: + yield session + engine.dispose() + + +def _user(db: Session, name: str = "workspace-account") -> User: + user = User(username=name, password_hash="hash", is_admin=False) + db.add(user) + db.flush() + return user + + +def _app( + db: Session, + *, + app_id: str, + name: str, + transport: str, + provider: str | None = None, + launch_config: dict | None = None, +) -> PublicMCPApp: + app = PublicMCPApp( + app_id=app_id, + name=name, + transport=transport, + provider_name=provider, + oauth_scopes=[], + launch_config=launch_config or {}, + is_visible_in_connector=True, + ) + db.add(app) + db.flush() + return app + + +def _associate(db: Session, *, user: User, server: MCPServer) -> UserMCPServer: + association = UserMCPServer( + user_id=user.id, + mcpserver_id=server.id, + is_owner=True, + is_active=True, + ) + db.add(association) + db.flush() + return association + + +def _remote_oauth_state( + db: Session, *, user: User, app_id: str = "remote-notes" +) -> tuple[PublicMCPApp, MCPServer, MCPOAuthClient]: + app = _app( + db, + app_id=app_id, + name="Remote Notes", + transport="streamable_http", + launch_config={ + "url": "https://mcp.example/mcp", + "auth": {"type": "mcp_oauth"}, + }, + ) + server = MCPServer.from_config( + { + "name": app_id, + "managed": "external", + "transport": "streamable_http", + "url": "https://mcp.example/mcp", + "auth": { + "type": "mcp_oauth", + "resource": "https://mcp.example/mcp", + }, + } + ) + db.add(server) + db.flush() + _associate(db, user=user, server=server) + client = MCPOAuthClient( + mcp_server_id=server.id, + issuer="https://auth.example", + authorization_endpoint="https://auth.example/authorize", + token_endpoint="https://auth.example/token", + client_id="client-id", + client_secret=encrypt_value("client-secret"), + token_endpoint_auth_method="client_secret_post", + redirect_uri="https://xagent.example/callback", + metadata_json={"revocation_endpoint": "https://auth.example/revoke"}, + ) + db.add(client) + db.flush() + db.add_all( + [ + MCPOAuthGrant( + mcp_server_id=server.id, + user_id=user.id, + mcp_oauth_client_id=client.id, + resource_owner_key=f"xagent:user:{user.id}", + issuer="https://auth.example", + resource="https://mcp.example/mcp", + scope="notes.read", + access_token=encrypt_value("access-token"), + refresh_token=encrypt_value("refresh-token"), + status="active", + ), + MCPOAuthFlowState( + state=f"state-{app_id}", + mcp_server_id=server.id, + user_id=user.id, + mcp_oauth_client_id=client.id, + resource_owner_key=f"xagent:user:{user.id}", + issuer="https://auth.example", + resource="https://mcp.example/mcp", + scope="notes.read", + code_verifier=encrypt_value("verifier"), + expires_at=datetime.now(timezone.utc) + timedelta(minutes=10), + ), + ] + ) + db.commit() + return app, server, client + + +@pytest.mark.asyncio +async def test_builtin_oauth_teardown_deletes_exact_credential_and_last_server( + db: Session, +) -> None: + user = _user(db) + app = _app( + db, + app_id="calendar", + name="Calendar", + transport="oauth", + provider="custom-calendar", + ) + server = MCPServer( + name="Calendar", + managed="external", + transport="oauth", + auth={"app_id": "calendar", "provider": "custom-calendar"}, + ) + db.add(server) + db.flush() + _associate(db, user=user, server=server) + db.add_all( + [ + UserOAuth( + user_id=user.id, + provider="custom-calendar", + access_token=encrypt_value("provider-token"), + ), + UserOAuth( + user_id=user.id, + provider="unrelated", + access_token=encrypt_value("keep-token"), + ), + ] + ) + db.commit() + app_pk, server_id = int(app.id), int(server.id) + + await mcp_api.teardown_mcp_app_server( + server_id, + app_id="calendar", + expected_catalog_app_id=app_pk, + expected_provider_name="custom-calendar", + current_user=user, + db=db, + ) + + assert db.get(MCPServer, server_id) is None + assert db.query(UserMCPServer).count() == 0 + assert {row.provider for row in db.query(UserOAuth).all()} == {"unrelated"} + + +@pytest.mark.asyncio +async def test_remote_oauth_last_user_cascades_client_secret_grant_and_flow( + db: Session, monkeypatch: pytest.MonkeyPatch +) -> None: + user = _user(db) + app, server, client = _remote_oauth_state(db, user=user) + app_pk, server_id, client_id = int(app.id), int(server.id), int(client.id) + observed: list[tuple[int, int, int]] = [] + + async def observe_after_commit(snapshot) -> None: + observed.append( + ( + db.query(MCPServer).count(), + db.query(MCPOAuthClient).count(), + snapshot.grant_id, + ) + ) + + monkeypatch.setattr( + mcp_api, "_revoke_mcp_oauth_snapshot_externally", observe_after_commit + ) + + await mcp_api.teardown_mcp_app_server( + server_id, + app_id="remote-notes", + expected_catalog_app_id=app_pk, + expected_provider_name=None, + current_user=user, + db=db, + ) + + assert observed and observed[0][:2] == (0, 0) + assert db.get(MCPServer, server_id) is None + assert db.get(MCPOAuthClient, client_id) is None + assert db.query(MCPOAuthGrant).count() == 0 + assert db.query(MCPOAuthFlowState).count() == 0 + assert db.query(UserMCPServer).count() == 0 + + +@pytest.mark.asyncio +async def test_final_server_delete_failure_rolls_back_every_cleanup_for_retry( + db: Session, monkeypatch: pytest.MonkeyPatch +) -> None: + user = _user(db) + app, server, client = _remote_oauth_state(db, user=user) + app_pk, server_id, client_id = int(app.id), int(server.id), int(client.id) + failures = {"remaining": 1} + + def fail_once(_mapper, _connection, target) -> None: + if int(target.id) == server_id and failures["remaining"]: + failures["remaining"] -= 1 + raise RuntimeError("simulated final delete failure") + + event.listen(MCPServer, "before_delete", fail_once) + try: + with pytest.raises(HTTPException) as exc_info: + await mcp_api.teardown_mcp_app_server( + server_id, + app_id="remote-notes", + expected_catalog_app_id=app_pk, + expected_provider_name=None, + current_user=user, + db=db, + ) + assert exc_info.value.status_code == 500 + + assert db.get(MCPServer, server_id) is not None + assert db.get(MCPOAuthClient, client_id) is not None + assert db.query(MCPOAuthGrant).count() == 1 + assert db.query(MCPOAuthFlowState).count() == 1 + assert db.query(UserMCPServer).count() == 1 + + async def ignore_revoke(_snapshot) -> None: + return None + + monkeypatch.setattr( + mcp_api, "_revoke_mcp_oauth_snapshot_externally", ignore_revoke + ) + await mcp_api.teardown_mcp_app_server( + server_id, + app_id="remote-notes", + expected_catalog_app_id=app_pk, + expected_provider_name=None, + current_user=user, + db=db, + ) + finally: + event.remove(MCPServer, "before_delete", fail_once) + + assert db.get(MCPServer, server_id) is None + assert db.get(MCPOAuthClient, client_id) is None + assert db.query(MCPOAuthGrant).count() == 0 + assert db.query(MCPOAuthFlowState).count() == 0 + assert db.query(UserMCPServer).count() == 0 + + +@pytest.mark.asyncio +async def test_recreated_catalog_owner_fails_closed_before_cleanup(db: Session) -> None: + user = _user(db) + app = _app( + db, + app_id="legacy-mail", + name="Legacy Mail", + transport="oauth", + provider="legacy-provider", + ) + server = MCPServer( + name="Legacy Mail", managed="external", transport="oauth", auth=None + ) + db.add(server) + db.flush() + _associate(db, user=user, server=server) + credential = UserOAuth( + user_id=user.id, + provider="legacy-provider", + access_token=encrypt_value("must-survive"), + ) + db.add(credential) + _app( + db, + app_id="pk-keeper", + name="PK Keeper", + transport="stdio", + ) + db.commit() + expected_pk, server_id = int(app.id), int(server.id) + + db.delete(app) + db.commit() + replacement = _app( + db, + app_id="legacy-mail", + name="Legacy Mail", + transport="oauth", + provider="replacement-provider", + ) + db.commit() + assert int(replacement.id) != expected_pk + + with pytest.raises(HTTPException) as exc_info: + await mcp_api.teardown_mcp_app_server( + server_id, + app_id="legacy-mail", + expected_catalog_app_id=expected_pk, + expected_provider_name="legacy-provider", + current_user=user, + db=db, + ) + + assert exc_info.value.status_code == 403 + assert db.get(MCPServer, server_id) is not None + assert db.query(UserMCPServer).count() == 1 + assert db.query(UserOAuth).filter_by(provider="legacy-provider").count() == 1 + + +@pytest.mark.asyncio +async def test_provider_drift_fails_closed_before_cleanup(db: Session) -> None: + user = _user(db) + app = _app( + db, + app_id="calendar", + name="Calendar", + transport="oauth", + provider="original-provider", + ) + server = MCPServer( + name="Calendar", + managed="external", + transport="oauth", + auth={"app_id": "calendar", "provider": "original-provider"}, + ) + db.add(server) + db.flush() + _associate(db, user=user, server=server) + db.add_all( + [ + UserOAuth( + user_id=user.id, + provider="original-provider", + access_token=encrypt_value("original-secret"), + ), + UserOAuth( + user_id=user.id, + provider="replacement-provider", + access_token=encrypt_value("replacement-secret"), + ), + ] + ) + db.commit() + app_pk, server_id = int(app.id), int(server.id) + + app.provider_name = "replacement-provider" + db.commit() + + with pytest.raises(HTTPException) as exc_info: + await mcp_api.teardown_mcp_app_server( + server_id, + app_id="calendar", + expected_catalog_app_id=app_pk, + expected_provider_name="original-provider", + current_user=user, + db=db, + ) + + assert exc_info.value.status_code == 403 + assert db.get(MCPServer, server_id) is not None + assert db.query(UserMCPServer).count() == 1 + assert {row.provider for row in db.query(UserOAuth).all()} == { + "original-provider", + "replacement-provider", + } + + +@pytest.mark.asyncio +async def test_multi_user_teardown_deletes_all_current_user_grant_statuses( + db: Session, monkeypatch: pytest.MonkeyPatch +) -> None: + user = _user(db) + other_user = _user(db, "other-account") + app, server, client = _remote_oauth_state(db, user=user) + app_pk, server_id = int(app.id), int(server.id) + _associate(db, user=other_user, server=server) + db.add_all( + [ + MCPOAuthGrant( + mcp_server_id=server_id, + user_id=user.id, + mcp_oauth_client_id=client.id, + resource_owner_key=f"xagent:user:{user.id}", + issuer="https://auth.example", + resource="https://mcp.example/mcp", + scope="notes.revoked", + access_token=encrypt_value("revoked-access-token"), + refresh_token=encrypt_value("revoked-refresh-token"), + status="revoked", + ), + MCPOAuthGrant( + mcp_server_id=server_id, + user_id=user.id, + mcp_oauth_client_id=client.id, + resource_owner_key=f"xagent:user:{user.id}", + issuer="https://auth.example", + resource="https://mcp.example/mcp", + scope="notes.inactive", + access_token=encrypt_value("inactive-access-token"), + refresh_token=encrypt_value("inactive-refresh-token"), + status="inactive", + ), + MCPOAuthGrant( + mcp_server_id=server_id, + user_id=other_user.id, + mcp_oauth_client_id=client.id, + resource_owner_key=f"xagent:user:{other_user.id}", + issuer="https://auth.example", + resource="https://mcp.example/mcp", + scope="notes.read", + access_token=encrypt_value("other-access-token"), + refresh_token=encrypt_value("other-refresh-token"), + status="active", + ), + ] + ) + db.commit() + active_grant_id = int( + db.query(MCPOAuthGrant.id).filter_by(user_id=user.id, status="active").scalar() + ) + revoked_grant_ids: list[int] = [] + + async def observe_revoke(snapshot) -> None: + revoked_grant_ids.append(snapshot.grant_id) + + monkeypatch.setattr( + mcp_api, "_revoke_mcp_oauth_snapshot_externally", observe_revoke + ) + + await mcp_api.teardown_mcp_app_server( + server_id, + app_id="remote-notes", + expected_catalog_app_id=app_pk, + expected_provider_name=None, + current_user=user, + db=db, + ) + + assert revoked_grant_ids == [active_grant_id] + assert db.get(MCPServer, server_id) is not None + assert db.get(MCPOAuthClient, int(client.id)) is not None + assert db.query(MCPOAuthGrant).filter_by(user_id=user.id).count() == 0 + assert db.query(MCPOAuthGrant).filter_by(user_id=other_user.id).count() == 1 + assert db.query(MCPOAuthFlowState).filter_by(user_id=user.id).count() == 0 + assert db.query(UserMCPServer).filter_by(user_id=user.id).count() == 0 + assert db.query(UserMCPServer).filter_by(user_id=other_user.id).count() == 1 + + +@pytest.mark.asyncio +async def test_unexpected_teardown_failure_logs_no_exception_detail( + db: Session, + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, +) -> None: + user = _user(db) + db.commit() + secret_detail = "raw upstream detail access_token=audit-secret" + + def fail_lock(_db: Session) -> None: + raise RuntimeError(secret_detail) + + monkeypatch.setattr(mcp_api, "_lock_catalog_for_app_teardown", fail_lock) + with caplog.at_level(logging.ERROR, logger=mcp_api.logger.name): + with pytest.raises(HTTPException) as exc_info: + await mcp_api.teardown_mcp_app_server( + 42, + app_id="safe-app-id", + expected_catalog_app_id=7, + expected_provider_name=None, + current_user=user, + db=db, + ) + + assert exc_info.value.status_code == 500 + assert exc_info.value.detail == "Failed to delete MCP server" + assert "safe-app-id" in caplog.text + assert "server 42" in caplog.text + assert secret_detail not in caplog.text + assert "audit-secret" not in caplog.text + + +@pytest.mark.asyncio +async def test_sqlite_teardown_rejects_without_discarding_caller_writes( + db: Session, +) -> None: + user = _user(db) + app = _app( + db, + app_id="calendar", + name="Calendar", + transport="oauth", + provider="calendar-provider", + ) + server = MCPServer( + name="Calendar", + managed="external", + transport="oauth", + auth={"app_id": "calendar", "provider": "calendar-provider"}, + ) + db.add(server) + db.flush() + _associate(db, user=user, server=server) + db.commit() + db.refresh(user) + app_pk, server_id = int(app.id), int(server.id) + pending = UserOAuth( + user_id=user.id, + provider="unrelated-pending-write", + access_token=encrypt_value("must-remain-pending"), + ) + db.add(pending) + + with pytest.raises(HTTPException) as exc_info: + await mcp_api.teardown_mcp_app_server( + server_id, + app_id="calendar", + expected_catalog_app_id=app_pk, + expected_provider_name="calendar-provider", + current_user=user, + db=db, + ) + + assert exc_info.value.status_code == 500 + assert pending in db.new + with db.no_autoflush: + assert db.get(MCPServer, server_id) is not None + assert db.query(UserMCPServer).filter_by(mcpserver_id=server_id).count() == 1 + + +@pytest.mark.parametrize("mutation", ["provider-drift", "delete-recreate"]) +def test_sqlite_catalog_mutation_waits_for_teardown_write_fence( + tmp_path, + monkeypatch: pytest.MonkeyPatch, + mutation: str, +) -> None: + engine = create_engine( + f"sqlite:///{tmp_path / 'teardown-race.db'}", + connect_args={"check_same_thread": False}, + ) + apply_sqlite_concurrency_pragmas(engine, busy_timeout_ms=10_000) + Base.metadata.create_all(engine) + factory = sessionmaker(bind=engine, autoflush=False, autocommit=False) + with factory() as seed_db: + user = _user(seed_db) + app = _app( + seed_db, + app_id="calendar", + name="Calendar", + transport="oauth", + provider="original-provider", + ) + _app(seed_db, app_id="pk-keeper", name="PK Keeper", transport="stdio") + server = MCPServer( + name="Calendar", + managed="external", + transport="oauth", + auth={"app_id": "calendar", "provider": "original-provider"}, + ) + seed_db.add(server) + seed_db.flush() + _associate(seed_db, user=user, server=server) + seed_db.add_all( + [ + UserOAuth( + user_id=user.id, + provider="original-provider", + access_token=encrypt_value("delete-me"), + ), + UserOAuth( + user_id=user.id, + provider="replacement-provider", + access_token=encrypt_value("keep-me"), + ), + ] + ) + seed_db.commit() + user_id, app_pk, server_id = int(user.id), int(app.id), int(server.id) + + identity_checked = threading.Event() + release_teardown = threading.Event() + mutation_sent = threading.Event() + mutation_finished = threading.Event() + mutation_thread_id: list[int] = [] + real_gate = mcp_api._locked_catalog_app_for_server + + def gated_identity(*args, **kwargs): + result = real_gate(*args, **kwargs) + identity_checked.set() + assert release_teardown.wait(timeout=10) + return result + + monkeypatch.setattr(mcp_api, "_locked_catalog_app_for_server", gated_identity) + + expected_prefix = ( + "UPDATE public_mcp_apps" + if mutation == "provider-drift" + else "DELETE FROM public_mcp_apps" + ) + + def observe_mutation( + _connection, _cursor, statement, _parameters, _context, _executemany + ) -> None: + if ( + mutation_thread_id + and threading.get_ident() == mutation_thread_id[0] + and statement.lstrip().startswith(expected_prefix) + ): + mutation_sent.set() + + event.listen(engine, "before_cursor_execute", observe_mutation) + + def teardown() -> None: + with factory() as teardown_db: + current_user = teardown_db.get(User, user_id) + assert current_user is not None + asyncio.run( + mcp_api.teardown_mcp_app_server( + server_id, + app_id="calendar", + expected_catalog_app_id=app_pk, + expected_provider_name="original-provider", + current_user=current_user, + db=teardown_db, + ) + ) + + def mutate_catalog() -> None: + mutation_thread_id.append(threading.get_ident()) + try: + with factory() as mutation_db: + expected = mutation_db.get(PublicMCPApp, app_pk) + assert expected is not None + if mutation == "provider-drift": + expected.provider_name = "replacement-provider" + mutation_db.commit() + else: + mutation_db.delete(expected) + mutation_db.commit() + mutation_db.add( + PublicMCPApp( + app_id="calendar", + name="Calendar Replacement", + transport="oauth", + provider_name="replacement-provider", + oauth_scopes=[], + launch_config={}, + is_visible_in_connector=True, + ) + ) + mutation_db.commit() + finally: + mutation_finished.set() + + try: + with ThreadPoolExecutor(max_workers=2) as executor: + teardown_future = executor.submit(teardown) + assert identity_checked.wait(timeout=10) + mutation_future = executor.submit(mutate_catalog) + assert mutation_sent.wait(timeout=10) + assert not mutation_finished.wait(timeout=0.25) + release_teardown.set() + teardown_future.result(timeout=10) + mutation_future.result(timeout=10) + finally: + release_teardown.set() + event.remove(engine, "before_cursor_execute", observe_mutation) + + with factory() as verify_db: + assert verify_db.get(MCPServer, server_id) is None + assert ( + verify_db.query(UserOAuth).filter_by(provider="original-provider").count() + == 0 + ) + assert ( + verify_db.query(UserOAuth) + .filter_by(provider="replacement-provider") + .count() + == 1 + ) + engine.dispose() diff --git a/tests/web/api/test_mcp_app_teardown_postgresql.py b/tests/web/api/test_mcp_app_teardown_postgresql.py new file mode 100644 index 0000000000..3728387804 --- /dev/null +++ b/tests/web/api/test_mcp_app_teardown_postgresql.py @@ -0,0 +1,719 @@ +from __future__ import annotations + +import asyncio +import threading +from concurrent.futures import ThreadPoolExecutor +from datetime import datetime, timedelta, timezone +from types import SimpleNamespace +from urllib.parse import parse_qs, urlparse + +import pytest +from fastapi import HTTPException +from sqlalchemy import event +from sqlalchemy.exc import IntegrityError +from sqlalchemy.orm import sessionmaker +from starlette.requests import Request + +from tests.shared.postgres_disposable import disposable_database_factory +from xagent.core.utils.encryption import decrypt_value, encrypt_value +from xagent.web.api import mcp as mcp_api +from xagent.web.models.database import Base +from xagent.web.models.mcp import MCPServer, UserMCPServer +from xagent.web.models.mcp_oauth import MCPOAuthClient, MCPOAuthFlowState, MCPOAuthGrant +from xagent.web.models.public_mcp import PublicMCPApp +from xagent.web.models.user import User +from xagent.web.models.user_oauth import UserOAuth + +pytestmark = pytest.mark.postgresql + +POSTGRES_TABLES = [ + User.__table__, + PublicMCPApp.__table__, + MCPServer.__table__, + UserMCPServer.__table__, + UserOAuth.__table__, + MCPOAuthClient.__table__, + MCPOAuthGrant.__table__, + MCPOAuthFlowState.__table__, +] + + +@pytest.fixture +def postgresql_engine(): + with disposable_database_factory("xagent_mcp_teardown") as make: + yield make("owner_race") + + +@pytest.fixture +def postgresql_context(postgresql_engine): + Base.metadata.create_all(postgresql_engine, tables=POSTGRES_TABLES) + factory = sessionmaker( + bind=postgresql_engine, + autoflush=False, + autocommit=False, + ) + return postgresql_engine, factory + + +def _seed(factory) -> tuple[int, int, int]: + with factory() as db: + user = User(username="workspace-account", password_hash="hash") + expected = PublicMCPApp( + app_id="legacy-mail", + name="Legacy Mail", + transport="oauth", + provider_name="legacy-provider", + oauth_scopes=[], + launch_config={}, + is_visible_in_connector=True, + ) + other = PublicMCPApp( + app_id="other-mail", + name="Other Mail", + transport="oauth", + provider_name="other-provider", + oauth_scopes=[], + launch_config={}, + is_visible_in_connector=True, + ) + server = MCPServer( + name="Legacy Mail", + managed="external", + transport="oauth", + auth=None, + ) + db.add_all([user, expected, other, server]) + db.flush() + db.add_all( + [ + UserMCPServer( + user_id=user.id, + mcpserver_id=server.id, + is_owner=True, + is_active=True, + ), + UserOAuth( + user_id=user.id, + provider="legacy-provider", + access_token=encrypt_value("legacy-token"), + ), + UserOAuth( + user_id=user.id, + provider="other-provider", + access_token=encrypt_value("other-token"), + ), + ] + ) + db.commit() + return int(user.id), int(expected.id), int(server.id) + + +def _oauth_callback_request(state: str) -> Request: + cookie = ( + f"{mcp_api.MCP_OAUTH_STATE_COOKIE}=" + f"{mcp_api._mcp_oauth_state_cookie_value(state)}" + ) + return Request( + { + "type": "http", + "method": "GET", + "path": "/api/mcp/oauth/callback", + "query_string": f"code=code-123&state={state}".encode(), + "headers": [(b"cookie", cookie.encode())], + } + ) + + +def _seed_oauth_lifecycle(factory) -> tuple[int, int, int, int, int, int]: + with factory() as db: + user = User(username="oauth-owner", password_hash="hash") + other_user = User(username="oauth-other", password_hash="hash") + app = PublicMCPApp( + app_id="remote-notes", + name="Remote Notes", + transport="streamable_http", + provider_name=None, + oauth_scopes=[], + launch_config={ + "url": "https://mcp.example/mcp", + "auth": {"type": "mcp_oauth"}, + }, + is_visible_in_connector=True, + ) + server = MCPServer( + name="remote-notes", + managed="external", + transport="streamable_http", + url="https://mcp.example/mcp", + auth={ + "type": "mcp_oauth", + "client_id": "client-id", + "client_secret": "client-secret", + "resource": "https://mcp.example/mcp", + "issuer": "https://auth.example", + "scope": "notes.read", + "redirect_uri": "https://xagent.example/api/mcp/oauth/callback", + "token_endpoint_auth_method": "client_secret_post", + }, + ) + db.add_all([user, other_user, app, server]) + db.flush() + db.add_all( + [ + UserMCPServer( + user_id=user.id, + mcpserver_id=server.id, + is_owner=True, + is_active=True, + ), + UserMCPServer( + user_id=other_user.id, + mcpserver_id=server.id, + is_owner=False, + is_active=True, + ), + ] + ) + client = MCPOAuthClient( + mcp_server_id=server.id, + issuer="https://auth.example", + authorization_endpoint="https://auth.example/authorize", + token_endpoint="https://auth.example/token", + client_id="client-id", + client_secret=encrypt_value("client-secret"), + token_endpoint_auth_method="client_secret_post", + redirect_uri="https://xagent.example/api/mcp/oauth/callback", + metadata_json={"revocation_endpoint": "https://auth.example/revoke"}, + ) + db.add(client) + db.flush() + flow = MCPOAuthFlowState( + state="old-flow-state", + mcp_server_id=server.id, + user_id=user.id, + mcp_oauth_client_id=client.id, + resource_owner_key=f"xagent:user:{user.id}", + issuer="https://auth.example", + resource="https://mcp.example/mcp", + scope="notes.read", + code_verifier=encrypt_value("verifier"), + redirect_after="/mcp", + expires_at=datetime.now(timezone.utc) + timedelta(minutes=10), + ) + db.add(flow) + db.commit() + return ( + int(user.id), + int(other_user.id), + int(app.id), + int(server.id), + int(client.id), + int(flow.id), + ) + + +def _run_oauth_teardown(factory, *, user_id: int, app_pk: int, server_id: int) -> None: + with factory() as teardown_db: + current_user = teardown_db.get(User, user_id) + assert current_user is not None + asyncio.run( + mcp_api.teardown_mcp_app_server( + server_id, + app_id="remote-notes", + expected_catalog_app_id=app_pk, + expected_provider_name=None, + current_user=current_user, + db=teardown_db, + ) + ) + + +def _run_oauth_callback(factory, state: str = "old-flow-state"): + with factory() as callback_db: + return asyncio.run( + mcp_api.mcp_oauth_callback(_oauth_callback_request(state), callback_db) + ) + + +@pytest.mark.parametrize("mutation", ["delete", "rename-reassign", "provider-drift"]) +def test_owner_mutation_between_preflight_and_teardown_fails_closed( + postgresql_context, mutation: str +) -> None: + _, factory = postgresql_context + user_id, expected_pk, server_id = _seed(factory) + barrier = threading.Barrier(2) + + def preflight_then_teardown() -> int: + with factory() as teardown_db: + current_user = teardown_db.get(User, user_id) + assert current_user is not None + # This is the SaaS preflight boundary: exact app and immutable row + # identity are read before a second Session mutates the catalog. + expected = teardown_db.get(PublicMCPApp, expected_pk) + assert expected is not None and expected.app_id == "legacy-mail" + barrier.wait(timeout=10) + barrier.wait(timeout=10) + try: + asyncio.run( + mcp_api.teardown_mcp_app_server( + server_id, + app_id="legacy-mail", + expected_catalog_app_id=expected_pk, + expected_provider_name="legacy-provider", + current_user=current_user, + db=teardown_db, + ) + ) + except HTTPException as exc: + return exc.status_code + raise AssertionError("teardown unexpectedly accepted a changed owner") + + with ThreadPoolExecutor(max_workers=1) as executor: + future = executor.submit(preflight_then_teardown) + barrier.wait(timeout=10) + with factory() as mutation_db: + expected = mutation_db.get(PublicMCPApp, expected_pk) + assert expected is not None + if mutation == "delete": + mutation_db.delete(expected) + elif mutation == "provider-drift": + expected.provider_name = "replacement-provider" + else: + expected.name = "Moved Legacy Mail" + other = ( + mutation_db.query(PublicMCPApp) + .filter(PublicMCPApp.app_id == "other-mail") + .one() + ) + other.name = "Legacy Mail" + mutation_db.commit() + barrier.wait(timeout=10) + assert future.result(timeout=10) == 403 + + with factory() as verify_db: + assert verify_db.get(MCPServer, server_id) is not None + assert ( + verify_db.query(UserMCPServer) + .filter(UserMCPServer.mcpserver_id == server_id) + .count() + == 1 + ) + assert {row.provider for row in verify_db.query(UserOAuth).all()} == { + "legacy-provider", + "other-provider", + } + + +def test_catalog_mutation_waits_while_teardown_holds_identity_locks( + postgresql_context, monkeypatch: pytest.MonkeyPatch +) -> None: + postgresql_engine, factory = postgresql_context + user_id, expected_pk, server_id = _seed(factory) + identity_checked = threading.Event() + release_teardown = threading.Event() + mutation_sent = threading.Event() + mutation_committed = threading.Event() + mutation_thread_id: list[int] = [] + real_gate = mcp_api._locked_catalog_app_for_server + + def gated_identity(*args, **kwargs): + answer = real_gate(*args, **kwargs) + identity_checked.set() + assert release_teardown.wait(timeout=10) + return answer + + monkeypatch.setattr(mcp_api, "_locked_catalog_app_for_server", gated_identity) + + def observe_catalog_update( + _connection, _cursor, statement, _parameters, _context, _executemany + ) -> None: + if ( + mutation_thread_id + and threading.get_ident() == mutation_thread_id[0] + and statement.lstrip().startswith("UPDATE public_mcp_apps") + ): + mutation_sent.set() + + event.listen(postgresql_engine, "before_cursor_execute", observe_catalog_update) + + def teardown() -> None: + with factory() as teardown_db: + current_user = teardown_db.get(User, user_id) + assert current_user is not None + asyncio.run( + mcp_api.teardown_mcp_app_server( + server_id, + app_id="legacy-mail", + expected_catalog_app_id=expected_pk, + expected_provider_name="legacy-provider", + current_user=current_user, + db=teardown_db, + ) + ) + + def rename() -> None: + mutation_thread_id.append(threading.get_ident()) + with factory() as mutation_db: + app = mutation_db.get(PublicMCPApp, expected_pk) + assert app is not None + app.name = "Renamed After Teardown" + mutation_db.commit() + mutation_committed.set() + + try: + with ThreadPoolExecutor(max_workers=2) as executor: + teardown_future = executor.submit(teardown) + assert identity_checked.wait(timeout=10) + rename_future = executor.submit(rename) + assert mutation_sent.wait(timeout=10) + assert not mutation_committed.wait(timeout=0.25) + release_teardown.set() + teardown_future.result(timeout=10) + rename_future.result(timeout=10) + finally: + release_teardown.set() + event.remove( + postgresql_engine, + "before_cursor_execute", + observe_catalog_update, + ) + + assert mutation_committed.is_set() + with factory() as verify_db: + assert verify_db.get(MCPServer, server_id) is None + assert ( + verify_db.query(UserOAuth).filter_by(provider="legacy-provider").count() + == 0 + ) + + +def test_concurrent_association_insert_serializes_with_parent_teardown( + postgresql_context, monkeypatch: pytest.MonkeyPatch +) -> None: + postgresql_engine, factory = postgresql_context + user_id, expected_pk, server_id = _seed(factory) + with factory() as setup_db: + other_user = User(username="concurrent-account", password_hash="hash") + setup_db.add(other_user) + setup_db.commit() + other_user_id = int(other_user.id) + + identity_checked = threading.Event() + release_teardown = threading.Event() + association_insert_sent = threading.Event() + insert_finished = threading.Event() + real_gate = mcp_api._locked_catalog_app_for_server + + def gated_identity(*args, **kwargs): + answer = real_gate(*args, **kwargs) + identity_checked.set() + assert release_teardown.wait(timeout=10) + return answer + + monkeypatch.setattr(mcp_api, "_locked_catalog_app_for_server", gated_identity) + + def observe_association_insert( + _connection, _cursor, statement, _parameters, _context, _executemany + ) -> None: + if statement.lstrip().startswith("INSERT INTO user_mcpservers"): + association_insert_sent.set() + + event.listen(postgresql_engine, "before_cursor_execute", observe_association_insert) + + def teardown() -> None: + with factory() as teardown_db: + current_user = teardown_db.get(User, user_id) + assert current_user is not None + asyncio.run( + mcp_api.teardown_mcp_app_server( + server_id, + app_id="legacy-mail", + expected_catalog_app_id=expected_pk, + expected_provider_name="legacy-provider", + current_user=current_user, + db=teardown_db, + ) + ) + + def insert_association() -> str: + with factory() as insert_db: + insert_db.add( + UserMCPServer( + user_id=other_user_id, + mcpserver_id=server_id, + is_owner=False, + is_active=True, + ) + ) + try: + insert_db.commit() + except IntegrityError: + insert_db.rollback() + return "foreign-key-rejected" + finally: + insert_finished.set() + return "committed" + + with ThreadPoolExecutor(max_workers=2) as executor: + try: + teardown_future = executor.submit(teardown) + assert identity_checked.wait(timeout=10) + insert_future = executor.submit(insert_association) + assert association_insert_sent.wait(timeout=10) + assert not insert_finished.wait(timeout=0.25) + release_teardown.set() + teardown_future.result(timeout=10) + assert insert_future.result(timeout=10) == "foreign-key-rejected" + finally: + release_teardown.set() + event.remove( + postgresql_engine, + "before_cursor_execute", + observe_association_insert, + ) + + with factory() as verify_db: + assert verify_db.get(MCPServer, server_id) is None + assert ( + verify_db.query(UserMCPServer) + .filter(UserMCPServer.mcpserver_id == server_id) + .count() + == 0 + ) + + +def test_callback_cannot_resurrect_grant_after_teardown_and_fast_reconnect( + postgresql_context, monkeypatch: pytest.MonkeyPatch +) -> None: + _, factory = postgresql_context + user_id, _, app_pk, server_id, client_id, old_flow_id = _seed_oauth_lifecycle( + factory + ) + exchange_started = threading.Event() + release_exchange = threading.Event() + revocations: list[mcp_api._MCPOAuthRevocationSnapshot] = [] + + async def exchange_after_teardown(**_kwargs): + exchange_started.set() + assert release_exchange.wait(timeout=10) + return { + "access_token": "new-access-token", + "refresh_token": "new-refresh-token", + "token_type": "Bearer", + "scope": "notes.read", + } + + async def observe_revocation(snapshot) -> None: + revocations.append(snapshot) + + monkeypatch.setattr(mcp_api, "_exchange_mcp_oauth_code", exchange_after_teardown) + monkeypatch.setattr( + mcp_api, "_revoke_mcp_oauth_snapshot_externally", observe_revocation + ) + + with ThreadPoolExecutor(max_workers=2) as executor: + callback_future = executor.submit(_run_oauth_callback, factory) + assert exchange_started.wait(timeout=10) + teardown_future = executor.submit( + _run_oauth_teardown, + factory, + user_id=user_id, + app_pk=app_pk, + server_id=server_id, + ) + teardown_future.result(timeout=10) + + # A new lifecycle may reconnect while the old provider request is still + # returning. The old callback must require its exact deleted flow, not + # merely accept the new active association. + with factory() as reconnect_db: + reconnect_db.add( + UserMCPServer( + user_id=user_id, + mcpserver_id=server_id, + is_owner=False, + is_active=True, + ) + ) + reconnect_db.add( + MCPOAuthFlowState( + state="new-flow-state", + mcp_server_id=server_id, + user_id=user_id, + mcp_oauth_client_id=client_id, + resource_owner_key=f"xagent:user:{user_id}", + issuer="https://auth.example", + resource="https://mcp.example/mcp", + scope="notes.read", + code_verifier=encrypt_value("new-verifier"), + expires_at=datetime.now(timezone.utc) + timedelta(minutes=10), + ) + ) + reconnect_db.commit() + + release_exchange.set() + response = callback_future.result(timeout=10) + + query = parse_qs(urlparse(response.headers["location"]).query) + assert query["mcp_oauth_error"] == ["invalid_state"] + assert len(revocations) == 1 + assert decrypt_value(revocations[0].access_token) == "new-access-token" + assert decrypt_value(str(revocations[0].refresh_token)) == "new-refresh-token" + with factory() as verify_db: + assert verify_db.get(MCPServer, server_id) is not None + assert verify_db.get(MCPOAuthFlowState, old_flow_id) is None + assert verify_db.query(MCPOAuthGrant).filter_by(user_id=user_id).count() == 0 + assert ( + verify_db.query(MCPOAuthFlowState) + .filter_by(user_id=user_id, state="new-flow-state") + .count() + == 1 + ) + + +def test_connect_cannot_create_flow_after_teardown_wins_during_discovery( + postgresql_context, monkeypatch: pytest.MonkeyPatch +) -> None: + _, factory = postgresql_context + user_id, _, app_pk, server_id, _, _ = _seed_oauth_lifecycle(factory) + discovery_started = threading.Event() + release_discovery = threading.Event() + + async def discovery_after_teardown(_server, _auth_config): + discovery_started.set() + assert release_discovery.wait(timeout=10) + return SimpleNamespace( + resource="https://mcp.example/mcp", + scopes=("notes.read",), + protected_resource=SimpleNamespace( + authorization_servers=("https://auth.example",), + ), + authorization_server=SimpleNamespace( + issuer="https://auth.example", + authorization_endpoint="https://auth.example/authorize", + token_endpoint="https://auth.example/token", + registration_endpoint=None, + client_id_metadata_document_supported=False, + raw={"issuer": "https://auth.example"}, + ), + ) + + monkeypatch.setattr( + mcp_api, "_discover_mcp_oauth_for_server", discovery_after_teardown + ) + + def connect() -> int: + with factory() as connect_db: + current_user = connect_db.get(User, user_id) + assert current_user is not None + try: + asyncio.run( + mcp_api.connect_mcp_oauth( + server_id, + mcp_api.MCPOAuthConnectRequest(), + current_user, + connect_db, + ) + ) + except HTTPException as exc: + return exc.status_code + raise AssertionError("OAuth connect unexpectedly persisted a flow") + + with ThreadPoolExecutor(max_workers=2) as executor: + connect_future = executor.submit(connect) + assert discovery_started.wait(timeout=10) + teardown_future = executor.submit( + _run_oauth_teardown, + factory, + user_id=user_id, + app_pk=app_pk, + server_id=server_id, + ) + teardown_future.result(timeout=10) + release_discovery.set() + assert connect_future.result(timeout=10) == 409 + + with factory() as verify_db: + assert verify_db.get(MCPServer, server_id) is not None + assert verify_db.query(UserMCPServer).filter_by(user_id=user_id).count() == 0 + assert ( + verify_db.query(MCPOAuthFlowState).filter_by(user_id=user_id).count() == 0 + ) + + +def test_callback_persistence_serializes_before_teardown_cleanup( + postgresql_context, monkeypatch: pytest.MonkeyPatch +) -> None: + postgresql_engine, factory = postgresql_context + user_id, _, app_pk, server_id, _, _ = _seed_oauth_lifecycle(factory) + grant_staged = threading.Event() + release_callback = threading.Event() + teardown_lock_sent = threading.Event() + teardown_thread_id: list[int] = [] + real_upsert = mcp_api._upsert_mcp_oauth_grant + + async def exchange_immediately(**_kwargs): + return { + "access_token": "producer-first-access", + "refresh_token": "producer-first-refresh", + "token_type": "Bearer", + "scope": "notes.read", + } + + def gated_upsert(*args, **kwargs): + grant = real_upsert(*args, **kwargs) + grant_staged.set() + assert release_callback.wait(timeout=10) + return grant + + def observe_teardown_server_lock( + _connection, _cursor, statement, _parameters, _context, _executemany + ) -> None: + normalized = " ".join(statement.split()) + if ( + teardown_thread_id + and threading.get_ident() == teardown_thread_id[0] + and "FROM mcp_servers" in normalized + and "FOR UPDATE" in normalized + ): + teardown_lock_sent.set() + + monkeypatch.setattr(mcp_api, "_exchange_mcp_oauth_code", exchange_immediately) + monkeypatch.setattr(mcp_api, "_upsert_mcp_oauth_grant", gated_upsert) + event.listen( + postgresql_engine, "before_cursor_execute", observe_teardown_server_lock + ) + + def teardown() -> None: + teardown_thread_id.append(threading.get_ident()) + _run_oauth_teardown( + factory, user_id=user_id, app_pk=app_pk, server_id=server_id + ) + + try: + with ThreadPoolExecutor(max_workers=2) as executor: + callback_future = executor.submit(_run_oauth_callback, factory) + assert grant_staged.wait(timeout=10) + teardown_future = executor.submit(teardown) + assert teardown_lock_sent.wait(timeout=10) + assert not teardown_future.done() + release_callback.set() + callback_response = callback_future.result(timeout=10) + teardown_future.result(timeout=10) + finally: + release_callback.set() + event.remove( + postgresql_engine, + "before_cursor_execute", + observe_teardown_server_lock, + ) + + query = parse_qs(urlparse(callback_response.headers["location"]).query) + assert query["mcp_oauth_success"] == ["1"] + with factory() as verify_db: + assert verify_db.get(MCPServer, server_id) is not None + assert verify_db.query(MCPOAuthGrant).filter_by(user_id=user_id).count() == 0 + assert ( + verify_db.query(MCPOAuthFlowState).filter_by(user_id=user_id).count() == 0 + ) + assert verify_db.query(UserMCPServer).filter_by(user_id=user_id).count() == 0 diff --git a/tests/web/api/test_mcp_oauth_flow.py b/tests/web/api/test_mcp_oauth_flow.py index 303a75633c..a4e22ec58f 100644 --- a/tests/web/api/test_mcp_oauth_flow.py +++ b/tests/web/api/test_mcp_oauth_flow.py @@ -1,6 +1,7 @@ from __future__ import annotations import json +import logging from datetime import timedelta from types import SimpleNamespace from urllib.parse import parse_qs, urlparse @@ -656,6 +657,39 @@ async def fake_discover(*args, **kwargs): assert flow_state.mcp_oauth_client_id == client.id +@pytest.mark.asyncio +async def test_connect_revalidates_active_association_after_discovery( + db_session, monkeypatch +): + db, user, _ = db_session + server = _add_mcp_oauth_server(db, user) + factory = sessionmaker(bind=db.get_bind(), autoflush=False, autocommit=False) + + async def disconnect_during_discovery(_server, _auth_config): + with factory() as teardown_db: + teardown_db.query(UserMCPServer).filter( + UserMCPServer.user_id == user.id, + UserMCPServer.mcpserver_id == server.id, + ).delete(synchronize_session=False) + teardown_db.commit() + return _discovery() + + monkeypatch.setattr( + mcp_api, "_discover_mcp_oauth_for_server", disconnect_during_discovery + ) + + with pytest.raises(HTTPException) as exc_info: + await connect_mcp_oauth( + server.id, + MCPOAuthConnectRequest(redirect_after="/settings/mcp"), + user, + db, + ) + + assert exc_info.value.status_code == 409 + assert db.query(MCPOAuthFlowState).count() == 0 + + @pytest.mark.asyncio async def test_connect_sweeps_flow_states_expired_past_the_retention_window( db_session, monkeypatch @@ -2269,6 +2303,131 @@ def async_client_factory(*args, **kwargs): assert db.query(MCPOAuthFlowState).one().consumed_at is not None +@pytest.mark.asyncio +async def test_callback_rejects_old_flow_after_disconnect_and_fast_reconnect( + db_session, monkeypatch +): + db, user, _ = db_session + server, client, old_flow = _add_callback_client_and_state( + db, + user, + state="old-lifecycle-state", + metadata_json={"revocation_endpoint": "https://auth.example.com/revoke"}, + ) + factory = sessionmaker(bind=db.get_bind(), autoflush=False, autocommit=False) + revocations = [] + + async def exchange_after_reconnect(**_kwargs): + with factory() as lifecycle_db: + lifecycle_db.query(MCPOAuthFlowState).filter_by(id=old_flow.id).delete( + synchronize_session=False + ) + lifecycle_db.query(UserMCPServer).filter_by( + user_id=user.id, mcpserver_id=server.id + ).delete(synchronize_session=False) + lifecycle_db.add( + UserMCPServer( + user_id=user.id, + mcpserver_id=server.id, + is_owner=False, + is_active=True, + ) + ) + lifecycle_db.add( + MCPOAuthFlowState( + state="new-lifecycle-state", + mcp_server_id=server.id, + user_id=user.id, + mcp_oauth_client_id=client.id, + resource_owner_key="resource-owner-a", + issuer="https://auth.example.com", + resource="https://mcp.example.com/mcp", + scope="records.read", + code_verifier=mcp_api.encrypt_value("new-verifier"), + expires_at=mcp_api._utc_now() + timedelta(minutes=10), + ) + ) + lifecycle_db.commit() + return { + "access_token": "issued-access-token", + "refresh_token": "issued-refresh-token", + "token_type": "Bearer", + "scope": "records.read", + } + + async def observe_revocation(snapshot): + revocations.append(snapshot) + + monkeypatch.setattr(mcp_api, "_exchange_mcp_oauth_code", exchange_after_reconnect) + monkeypatch.setattr( + mcp_api, "_revoke_mcp_oauth_snapshot_externally", observe_revocation + ) + + response = await mcp_oauth_callback( + _request("/api/mcp/oauth/callback?code=auth-code&state=old-lifecycle-state"), + db, + ) + + assert _redirect_query(response)["mcp_oauth_error"] == ["invalid_state"] + assert db.query(MCPOAuthGrant).count() == 0 + assert db.query(MCPOAuthFlowState).one().state == "new-lifecycle-state" + assert len(revocations) == 1 + assert decrypt_value(revocations[0].access_token) == "issued-access-token" + assert decrypt_value(revocations[0].refresh_token) == "issued-refresh-token" + + +@pytest.mark.asyncio +async def test_callback_persistence_failure_revokes_token_without_logging_detail( + db_session, monkeypatch, caplog +): + db, user, _ = db_session + _add_callback_client_and_state( + db, + user, + state="persistence-failure-state", + metadata_json={"revocation_endpoint": "https://auth.example.com/revoke"}, + ) + secret_detail = "issued-access-token raw-upstream-detail" + revocations = [] + + async def exchange(**_kwargs): + return { + "access_token": "issued-access-token", + "refresh_token": "issued-refresh-token", + "token_type": "Bearer", + "scope": "records.read", + } + + def fail_persistence(*_args, **_kwargs): + raise RuntimeError(secret_detail) + + async def observe_revocation(snapshot): + revocations.append(snapshot) + + monkeypatch.setattr(mcp_api, "_exchange_mcp_oauth_code", exchange) + monkeypatch.setattr(mcp_api, "_upsert_mcp_oauth_grant", fail_persistence) + monkeypatch.setattr( + mcp_api, "_revoke_mcp_oauth_snapshot_externally", observe_revocation + ) + caplog.set_level(logging.ERROR, logger=mcp_api.__name__) + + response = await mcp_oauth_callback( + _request( + "/api/mcp/oauth/callback?code=auth-code&state=persistence-failure-state" + ), + db, + ) + + assert _redirect_query(response)["mcp_oauth_error"] == ["token_exchange_failed"] + assert db.query(MCPOAuthGrant).count() == 0 + assert len(revocations) == 1 + assert decrypt_value(revocations[0].access_token) == "issued-access-token" + assert decrypt_value(revocations[0].refresh_token) == "issued-refresh-token" + assert "MCP OAuth callback failed after state claim" in caplog.text + assert secret_detail not in caplog.text + assert "raw-upstream-detail" not in caplog.text + + @pytest.mark.asyncio async def test_exchange_code_sanitizes_transport_exception(db_session, monkeypatch): db, user, _ = db_session