From ddbbc01fea3c7daa3db93bd83c7d94b8beb14e81 Mon Sep 17 00:00:00 2001 From: Wenxi Date: Thu, 13 Aug 2026 20:58:48 +0000 Subject: [PATCH] fix(mcp): request offline access from Google OAuth and surface dead grants (#13924) --- backend/onyx/auth/oauth_token_manager.py | 32 +++- backend/onyx/db/mcp.py | 24 ++- .../sandbox_proxy/resolvers/mcp_server.py | 6 +- backend/onyx/server/features/mcp/api.py | 2 +- backend/onyx/server/features/mcp/models.py | 17 ++ backend/onyx/server/features/mcp/oauth.py | 69 ++++++- .../tool_implementations/mcp/mcp_tool.py | 7 +- .../mcp/test_mcp_oauth_dead_grant.py | 170 ++++++++++++++++++ .../auth/test_offline_access_auth_params.py | 104 +++++++++++ .../features/mcp/test_mcp_oauth_refresh.py | 2 +- .../test_mcp_known_provider_oauth_helpers.py | 2 +- 11 files changed, 411 insertions(+), 24 deletions(-) create mode 100644 backend/tests/external_dependency_unit/mcp/test_mcp_oauth_dead_grant.py create mode 100644 backend/tests/unit/onyx/auth/test_offline_access_auth_params.py diff --git a/backend/onyx/auth/oauth_token_manager.py b/backend/onyx/auth/oauth_token_manager.py index 9604e79f66..be3a2f1278 100644 --- a/backend/onyx/auth/oauth_token_manager.py +++ b/backend/onyx/auth/oauth_token_manager.py @@ -1,6 +1,6 @@ import time from typing import Any -from urllib.parse import urlencode +from urllib.parse import parse_qsl, urlencode, urlparse, urlunparse from uuid import UUID import requests @@ -41,6 +41,32 @@ OAUTH_RESPONSE_TYPE_CODE = "code" OAUTH_GRANT_TYPE_AUTHORIZATION_CODE = "authorization_code" OAUTH_PKCE_CHALLENGE_METHOD_S256 = "S256" +# Providers that issue a refresh token only when the authorization request +# explicitly asks for offline access, keyed by authorization-endpoint host. +# Without these params the grant dies at access-token expiry with no way to +# refresh. Explicitly configured params always win over these defaults. +_OFFLINE_ACCESS_AUTH_PARAMS_BY_HOST: dict[str, dict[str, str]] = { + # `prompt=consent` because Google issues a refresh token only on flows + # that show the consent screen; a silent re-auth would leave a reconnected + # config without one. Matches Onyx's own Google login flow. + "accounts.google.com": {"access_type": "offline", "prompt": "consent"}, +} + + +def ensure_offline_access_auth_params(authorization_url: str) -> str: + """Merge missing offline-access params into a built authorize URL. + Params already in the URL win.""" + parsed = urlparse(authorization_url) + defaults = _OFFLINE_ACCESS_AUTH_PARAMS_BY_HOST.get((parsed.hostname or "").lower()) + if not defaults: + return authorization_url + query = dict(parse_qsl(parsed.query, keep_blank_values=True)) + missing = {key: value for key, value in defaults.items() if key not in query} + if not missing: + return authorization_url + query.update(missing) + return urlunparse(parsed._replace(query=urlencode(query))) + class OAuthFlowParams(BaseModel): """Stateless inputs for the OAuth 2.0 authorization-code grant, decoupled @@ -83,7 +109,9 @@ def build_oauth_authorization_url( query.update(params.additional_params) separator = "&" if "?" in params.authorization_url else "?" - return f"{params.authorization_url}{separator}{urlencode(query)}" + return ensure_offline_access_auth_params( + f"{params.authorization_url}{separator}{urlencode(query)}" + ) def exchange_oauth_code_for_token( diff --git a/backend/onyx/db/mcp.py b/backend/onyx/db/mcp.py index ae367ac457..923c5c94b1 100644 --- a/backend/onyx/db/mcp.py +++ b/backend/onyx/db/mcp.py @@ -34,6 +34,7 @@ from onyx.server.features.mcp.models import ( MCPAuthTemplate, MCPConnectionData, MCPOAuthKeys, + mcp_oauth_reauth_required, merge_mcp_headers, ) from onyx.utils.logger import setup_logger @@ -396,12 +397,14 @@ def extract_connection_data( def get_connection_config_by_id( - config_id: int, db_session: Session + config_id: int, db_session: Session, *, for_update: bool = False ) -> MCPConnectionConfig: - """Get connection config by ID""" - config = db_session.scalar( - select(MCPConnectionConfig).where(MCPConnectionConfig.id == config_id) - ) + """Get connection config by ID. ``for_update`` row-locks the config so a + read-compare-write over its JSON blob is atomic against concurrent writers.""" + stmt = select(MCPConnectionConfig).where(MCPConnectionConfig.id == config_id) + if for_update: + stmt = stmt.with_for_update() + config = db_session.scalar(stmt) if not config: raise ValueError("Connection config by specified id does not exist") return config @@ -511,14 +514,25 @@ class ResolvedMCPCredentials(BaseModel): ) return headers + def needs_reauth(self) -> bool: + """The stored OAuth grant is dead (expired, no refresh token); only a + user reconnect can restore it. Its bearer header still takes precedence + in ``build_headers``, so no caller-supplied header can work around it.""" + return self.auth_type == MCPAuthenticationType.OAUTH and ( + mcp_oauth_reauth_required(self._config_data()) + ) + def is_authenticated(self) -> bool: if self.auth_type in (None, MCPAuthenticationType.NONE): return self._has_required_substitutions() if self.auth_type == MCPAuthenticationType.PT_OAUTH: return bool(self.user_oauth_token) and self._has_required_substitutions() if self.auth_type == MCPAuthenticationType.OAUTH: + # A dead grant must read as unauthenticated so the UI prompts a + # reconnect and Craft configs exclude the server. return ( bool(self._generated_auth_headers()) + and not self.needs_reauth() and self._has_required_substitutions() ) return bool(self._configured_headers()) and self._has_required_substitutions() diff --git a/backend/onyx/sandbox_proxy/resolvers/mcp_server.py b/backend/onyx/sandbox_proxy/resolvers/mcp_server.py index d10819905c..80a3f720cc 100644 --- a/backend/onyx/sandbox_proxy/resolvers/mcp_server.py +++ b/backend/onyx/sandbox_proxy/resolvers/mcp_server.py @@ -47,10 +47,8 @@ from onyx.sandbox_proxy.resolvers.mcp_matching import ( normalized_request_path, parse_target, ) -from onyx.server.features.mcp.oauth import ( - mcp_token_expired, - refresh_mcp_oauth_token_if_expired, -) +from onyx.server.features.mcp.models import mcp_token_expired +from onyx.server.features.mcp.oauth import refresh_mcp_oauth_token_if_expired from onyx.utils.credential_audit import emit_credential_access from onyx.utils.logger import setup_logger from shared_configs.contextvars import CURRENT_TENANT_ID_CONTEXTVAR diff --git a/backend/onyx/server/features/mcp/api.py b/backend/onyx/server/features/mcp/api.py index 264fdaddd1..abf6c1ecb9 100644 --- a/backend/onyx/server/features/mcp/api.py +++ b/backend/onyx/server/features/mcp/api.py @@ -111,6 +111,7 @@ from onyx.server.features.mcp.models import ( MCPUserOAuthConnectRequest, MCPUserOAuthConnectResponse, contains_mcp_placeholder, + mcp_token_expired, merge_mcp_headers, ) from onyx.server.features.mcp.oauth import ( @@ -125,7 +126,6 @@ from onyx.server.features.mcp.oauth import ( key_state, key_tokens, make_oauth_provider, - mcp_token_expired, ) from onyx.server.features.mcp.ssrf import validate_mcp_outbound_url from onyx.server.features.tool.models import ToolSnapshot diff --git a/backend/onyx/server/features/mcp/models.py b/backend/onyx/server/features/mcp/models.py index ce6c77f92e..56d7ef9792 100644 --- a/backend/onyx/server/features/mcp/models.py +++ b/backend/onyx/server/features/mcp/models.py @@ -1,5 +1,6 @@ import datetime import re +import time from enum import Enum from typing import Any, List, Literal, NotRequired, Optional, TypedDict from uuid import UUID @@ -121,6 +122,22 @@ class MCPConnectionData(TypedDict): # from mcp.shared.auth import OAuthClientInformationFull, OAuthClientMetadata, OAuthToken +def mcp_token_expired(config_data: MCPConnectionData) -> bool: + """True iff the stored access token is past its persisted expiry.""" + expires_at = config_data.get(MCPOAuthKeys.TOKEN_EXPIRES_AT.value) + return expires_at is not None and float(expires_at) <= time.time() + + +def mcp_oauth_reauth_required(config_data: MCPConnectionData) -> bool: + """True when the stored OAuth grant can never authenticate again on its + own: the access token is expired and there is no refresh token to redeem. + Such a config must not count as authenticated — the user has to reconnect.""" + if not mcp_token_expired(config_data): + return False + tokens = config_data.get(MCPOAuthKeys.TOKENS.value) or {} + return not tokens.get("refresh_token") + + class MCPAuthTemplate(BaseModel): """Header template shared by every MCP authentication type.""" diff --git a/backend/onyx/server/features/mcp/oauth.py b/backend/onyx/server/features/mcp/oauth.py index 0b4abb0e12..22ad662bee 100644 --- a/backend/onyx/server/features/mcp/oauth.py +++ b/backend/onyx/server/features/mcp/oauth.py @@ -31,6 +31,7 @@ from mcp.shared.auth import ( from pydantic import AnyUrl, BaseModel, ValidationError from sqlalchemy.orm import Session +from onyx.auth.oauth_token_manager import ensure_offline_access_auth_params from onyx.cache.interface import CacheLockAcquisitionError from onyx.cache.locks import cache_shared_lock from onyx.db.engine.sql_engine import get_session_with_current_tenant @@ -51,8 +52,8 @@ from onyx.server.features.mcp.client_metadata import ( ) from onyx.server.features.mcp.models import ( DENYLISTED_MCP_HEADERS, - MCPConnectionData, MCPOAuthKeys, + mcp_token_expired, merge_mcp_headers, ) from onyx.server.features.mcp.ssrf import ( @@ -519,12 +520,6 @@ def refresh_mcp_oauth_token_if_expired( return _persisted_auth_header(connection_config_id) -def mcp_token_expired(config_data: MCPConnectionData) -> bool: - """True iff the stored access token is past its persisted expiry.""" - expires_at = config_data.get(MCPOAuthKeys.TOKEN_EXPIRES_AT.value) - return expires_at is not None and float(expires_at) <= time.time() - - def _persisted_auth_header(connection_config_id: int) -> str | None: """The stored ``Authorization`` header when the persisted token is still fresh, else None — used as the fallback when a concurrent refresh wins.""" @@ -580,8 +575,12 @@ class OnyxTokenStorage(TokenStorage): def bind_oauth_context(self, context: OAuthContext) -> None: self._oauth_context = context - def _ensure_connection_config(self, db_session: Session) -> MCPConnectionConfig: - config = get_connection_config_by_id(self.connection_config_id, db_session) + def _ensure_connection_config( + self, db_session: Session, *, for_update: bool = False + ) -> MCPConnectionConfig: + config = get_connection_config_by_id( + self.connection_config_id, db_session, for_update=for_update + ) if config is None: raise OnyxError(OnyxErrorCode.NOT_FOUND, "Connection config not found") return config @@ -682,6 +681,37 @@ class OnyxTokenStorage(TokenStorage): r.rpush(key_tokens(str(self.alt_config_id)), tokens.model_dump_json()) r.expire(key_tokens(str(self.alt_config_id)), OAUTH_WAIT_SECONDS) + async def discard_persisted_tokens( + self, redeemed_refresh_token: str | None + ) -> bool: + """Drop the persisted grant (tokens, expiry, Authorization header) so + the connection reads as unauthenticated until the user reconnects. + + No-ops when the stored refresh token no longer matches + ``redeemed_refresh_token``: a concurrent refresh or reconnect already + replaced the grant, and the replacement must survive. Returns whether + the grant was discarded.""" + with get_session_with_current_tenant() as db_session: + # Row lock: a concurrent set_tokens commit between an unlocked read + # and the update below would be clobbered by this stale read. + config = self._ensure_connection_config(db_session, for_update=True) + config_data = extract_connection_data(config) + stored_tokens = config_data.get(MCPOAuthKeys.TOKENS.value) or {} + if ( + redeemed_refresh_token is None + or stored_tokens.get("refresh_token") != redeemed_refresh_token + ): + return False + config_data.pop(MCPOAuthKeys.TOKENS.value, None) + config_data.pop(MCPOAuthKeys.TOKEN_EXPIRES_AT.value, None) + config_data["headers"] = { + key: value + for key, value in config_data.get("headers", {}).items() + if key.lower() != "authorization" + } + update_connection_config(config.id, db_session, config_data) + return True + async def get_client_info(self) -> OAuthClientInformationFull | None: with get_session_with_current_tenant() as db_session: config = self._ensure_connection_config(db_session) @@ -758,6 +788,9 @@ class OnyxOAuthClientProvider(OAuthClientProvider): self.refresh_log_context = refresh_log_context self.refresh_attempt_id: str | None = None self.token_expiry_before_refresh: float | None = None + # The refresh token the in-flight refresh attempt redeemed; guards the + # invalid_grant discard against wiping a concurrently stored grant. + self.redeemed_refresh_token: str | None = None def _refresh_log_fields(self, **fields: Any) -> dict[str, Any]: log_fields: dict[str, Any] = { @@ -771,6 +804,10 @@ class OnyxOAuthClientProvider(OAuthClientProvider): async def _refresh_token(self) -> httpx.Request: self.refresh_attempt_id = uuid4().hex self.token_expiry_before_refresh = self.context.token_expiry_time + current_tokens = self.context.current_tokens + self.redeemed_refresh_token = ( + current_tokens.refresh_token if current_tokens else None + ) request = await super()._refresh_token() logger.info( @@ -820,6 +857,17 @@ class OnyxOAuthClientProvider(OAuthClientProvider): if response.status_code != 200: logger.warning("mcp_oauth.refresh.failed", extra=response_fields) self.context.clear_tokens() + # invalid_grant means the refresh token itself is expired or + # revoked (RFC 6749) — discard the dead grant so the connection + # reads as unauthenticated instead of retrying forever. + storage = self.context.storage + if oauth_error == "invalid_grant" and isinstance(storage, OnyxTokenStorage): + if await storage.discard_persisted_tokens(self.redeemed_refresh_token): + logger.warning("mcp_oauth.tokens_discarded", extra=response_fields) + else: + logger.info( + "mcp_oauth.tokens_discard_skipped_stale", extra=response_fields + ) return False try: @@ -857,6 +905,9 @@ def make_oauth_provider( async def redirect_handler(auth_url: str) -> None: if return_path == UNUSED_RETURN_PATH: raise ValueError("Please Reconnect to the server") + # The SDK exposes no hook for provider-specific authorize params; add + # them here, where every auto-discovery consent URL passes through. + auth_url = ensure_offline_access_auth_params(auth_url) r = get_redis_client() # The SDK generated & embedded 'state' in the auth_url; extract & store it. parsed = urlparse(auth_url) diff --git a/backend/onyx/tools/tool_implementations/mcp/mcp_tool.py b/backend/onyx/tools/tool_implementations/mcp/mcp_tool.py index 2e870cfd44..9c3a999b82 100644 --- a/backend/onyx/tools/tool_implementations/mcp/mcp_tool.py +++ b/backend/onyx/tools/tool_implementations/mcp/mcp_tool.py @@ -176,7 +176,12 @@ class MCPTool(Tool[None]): credentials.build_headers(), ) - if not credentials.is_authenticated() and not self._additional_headers: + # Extra request headers can stand in for missing credentials, but + # not for a dead OAuth grant — its stale bearer wins the header + # merge, so the call can only fail upstream. + if not credentials.is_authenticated() and ( + credentials.needs_reauth() or not self._additional_headers + ): auth_error_msg = ( f"The {self._name} tool from {self.mcp_server.name} requires " "connection values. Tell the user to connect to the server " diff --git a/backend/tests/external_dependency_unit/mcp/test_mcp_oauth_dead_grant.py b/backend/tests/external_dependency_unit/mcp/test_mcp_oauth_dead_grant.py new file mode 100644 index 0000000000..44e53246c3 --- /dev/null +++ b/backend/tests/external_dependency_unit/mcp/test_mcp_oauth_dead_grant.py @@ -0,0 +1,170 @@ +"""OAuth MCP grants that can never refresh must read as unauthenticated (ext-dep): +`is_authenticated` semantics, Craft config exclusion, and the persisted token +discard on an `invalid_grant` refresh response, against a real DB.""" + +from __future__ import annotations + +import asyncio +import time +from typing import Any +from uuid import uuid4 + +import httpx +from sqlalchemy.orm import Session + +from onyx.db.enums import ( + MCPAuthenticationPerformer, + MCPAuthenticationType, + MCPTransport, +) +from onyx.db.mcp import ( + can_resolve_mcp_credentials, + create_mcp_server__no_commit, + extract_connection_data, + get_user_connection_config, + update_mcp_server__no_commit, + upsert_user_connection_config, +) +from onyx.db.models import MCPServer, User +from onyx.server.features.build.sandbox.util.mcp_config import ( + resolve_craft_mcp_servers, +) +from onyx.server.features.mcp.models import MCPConnectionData +from onyx.server.features.mcp.oauth import ( + UNUSED_RETURN_PATH, + make_oauth_provider, +) +from tests.external_dependency_unit.conftest import create_test_user + + +def _make_oauth_server(db_session: Session, *, owner_email: str) -> MCPServer: + server = create_mcp_server__no_commit( + owner_email=owner_email, + name=f"oauth-{uuid4().hex[:8]}", + description=None, + server_url=f"https://mcp-{uuid4().hex[:8]}.example.com/mcp", + auth_type=MCPAuthenticationType.OAUTH, + transport=MCPTransport.STREAMABLE_HTTP, + auth_performer=MCPAuthenticationPerformer.PER_USER, + db_session=db_session, + is_public=True, + ) + update_mcp_server__no_commit( + server_id=server.id, db_session=db_session, available_in_craft=True + ) + db_session.commit() + return server + + +def _store_grant( + db_session: Session, + server: MCPServer, + user: User, + *, + expires_at: float, + refresh_token: str | None, +) -> int: + tokens: dict[str, Any] = { + "access_token": "access-token", + "token_type": "Bearer", + "refresh_token": refresh_token, + } + config_data: MCPConnectionData = { + "headers": {"Authorization": "Bearer access-token"}, + "tokens": tokens, + "token_expires_at": expires_at, + } + config = upsert_user_connection_config( + server.id, user.email, config_data, db_session + ) + db_session.commit() + return config.id + + +def test_expired_token_with_refresh_token_is_authenticated( + db_session: Session, + tenant_context: None, # noqa: ARG001 +) -> None: + user = create_test_user(db_session, "oauth_refreshable") + server = _make_oauth_server(db_session, owner_email=user.email) + _store_grant( + db_session, + server, + user, + expires_at=time.time() - 60, + refresh_token="refresh-token", + ) + assert can_resolve_mcp_credentials(server, user, db_session) + + +def test_dead_grant_is_not_authenticated_and_excluded_from_craft( + db_session: Session, + tenant_context: None, # noqa: ARG001 +) -> None: + user = create_test_user(db_session, "oauth_dead") + server = _make_oauth_server(db_session, owner_email=user.email) + + def craft_server_ids() -> list[int]: + return [r.server_id for r in resolve_craft_mcp_servers(db_session, user)] + + # Positive control: a live grant authenticates and reaches Craft configs, + # so the exclusion below can't pass vacuously. + _store_grant( + db_session, server, user, expires_at=time.time() + 3600, refresh_token=None + ) + assert can_resolve_mcp_credentials(server, user, db_session) + assert server.id in craft_server_ids() + + _store_grant( + db_session, server, user, expires_at=time.time() - 60, refresh_token=None + ) + assert not can_resolve_mcp_credentials(server, user, db_session) + assert server.id not in craft_server_ids() + + +def _invalid_grant_response() -> httpx.Response: + return httpx.Response( + status_code=400, + json={"error": "invalid_grant"}, + request=httpx.Request("POST", "https://oauth2.example.com/token"), + ) + + +def test_invalid_grant_refresh_discards_persisted_tokens( + db_session: Session, + tenant_context: None, # noqa: ARG001 +) -> None: + user = create_test_user(db_session, "oauth_revoked") + server = _make_oauth_server(db_session, owner_email=user.email) + config_id = _store_grant( + db_session, + server, + user, + expires_at=time.time() - 60, + refresh_token="revoked-refresh-token", + ) + + provider = make_oauth_provider( + server, str(user.id), UNUSED_RETURN_PATH, config_id, None + ) + + # A stale attempt (its refresh token was already rotated away by a + # concurrent refresh or reconnect) must not wipe the stored grant. + provider.redeemed_refresh_token = "rotated-away-refresh-token" + assert not asyncio.run(provider._handle_refresh_response(_invalid_grant_response())) + db_session.expire_all() + config = get_user_connection_config(server.id, user.email, db_session) + assert config is not None + assert "tokens" in extract_connection_data(config, apply_mask=False) + + # The attempt that redeemed the stored refresh token discards the grant. + provider.redeemed_refresh_token = "revoked-refresh-token" + assert not asyncio.run(provider._handle_refresh_response(_invalid_grant_response())) + db_session.expire_all() + config = get_user_connection_config(server.id, user.email, db_session) + assert config is not None + config_value = extract_connection_data(config, apply_mask=False) + assert "tokens" not in config_value + assert "token_expires_at" not in config_value + assert "Authorization" not in config_value.get("headers", {}) + assert not can_resolve_mcp_credentials(server, user, db_session) diff --git a/backend/tests/unit/onyx/auth/test_offline_access_auth_params.py b/backend/tests/unit/onyx/auth/test_offline_access_auth_params.py new file mode 100644 index 0000000000..a63c3c8c18 --- /dev/null +++ b/backend/tests/unit/onyx/auth/test_offline_access_auth_params.py @@ -0,0 +1,104 @@ +"""Providers like Google issue refresh tokens only when the authorize request +asks for offline access; these params must be injected into every flow.""" + +from urllib.parse import parse_qs, urlparse + +from onyx.auth.oauth_token_manager import ( + OAuthFlowParams, + build_oauth_authorization_url, + ensure_offline_access_auth_params, +) + +_GOOGLE_AUTH_ENDPOINT = "https://accounts.google.com/o/oauth2/v2/auth" + + +def _query(url: str) -> dict[str, list[str]]: + return parse_qs(urlparse(url).query) + + +def test_google_prebuilt_url_gains_offline_access_params() -> None: + url = ensure_offline_access_auth_params( + f"{_GOOGLE_AUTH_ENDPOINT}?client_id=x&state=s&code_challenge=c" + ) + query = _query(url) + assert query["access_type"] == ["offline"] + assert query["prompt"] == ["consent"] + # Existing params survive the rewrite. + assert query["client_id"] == ["x"] + assert query["state"] == ["s"] + assert query["code_challenge"] == ["c"] + + +def test_prebuilt_url_existing_params_win() -> None: + url = ensure_offline_access_auth_params( + f"{_GOOGLE_AUTH_ENDPOINT}?client_id=x&prompt=none" + ) + query = _query(url) + assert query["prompt"] == ["none"] + assert query["access_type"] == ["offline"] + + +def test_non_google_prebuilt_url_unchanged() -> None: + url = "https://example.com/authorize?client_id=x" + assert ensure_offline_access_auth_params(url) == url + + +def test_build_authorization_url_adds_google_offline_access() -> None: + url = build_oauth_authorization_url( + OAuthFlowParams( + authorization_url=_GOOGLE_AUTH_ENDPOINT, + token_url="https://oauth2.googleapis.com/token", + client_id="x", + ), + redirect_uri="https://app.example.com/callback", + state="s", + ) + query = _query(url) + assert query["access_type"] == ["offline"] + assert query["prompt"] == ["consent"] + + +def test_build_authorization_url_additional_params_override_defaults() -> None: + url = build_oauth_authorization_url( + OAuthFlowParams( + authorization_url=_GOOGLE_AUTH_ENDPOINT, + token_url="https://oauth2.googleapis.com/token", + client_id="x", + additional_params={"prompt": "select_account"}, + ), + redirect_uri="https://app.example.com/callback", + state="s", + ) + query = _query(url) + assert query["prompt"] == ["select_account"] + assert query["access_type"] == ["offline"] + + +def test_build_authorization_url_endpoint_embedded_params_win() -> None: + url = build_oauth_authorization_url( + OAuthFlowParams( + authorization_url=f"{_GOOGLE_AUTH_ENDPOINT}?prompt=select_account", + token_url="https://oauth2.googleapis.com/token", + client_id="x", + ), + redirect_uri="https://app.example.com/callback", + state="s", + ) + query = _query(url) + assert query["prompt"] == ["select_account"] + assert query["access_type"] == ["offline"] + + +def test_build_authorization_url_non_google_untouched() -> None: + url = build_oauth_authorization_url( + OAuthFlowParams( + authorization_url="https://example.com/authorize", + token_url="https://example.com/token", + client_id="x", + ), + redirect_uri="https://app.example.com/callback", + state="s", + ) + query = _query(url) + assert "access_type" not in query + assert "prompt" not in query diff --git a/backend/tests/unit/onyx/server/features/mcp/test_mcp_oauth_refresh.py b/backend/tests/unit/onyx/server/features/mcp/test_mcp_oauth_refresh.py index 0dabf33368..8c4c87d30c 100644 --- a/backend/tests/unit/onyx/server/features/mcp/test_mcp_oauth_refresh.py +++ b/backend/tests/unit/onyx/server/features/mcp/test_mcp_oauth_refresh.py @@ -99,7 +99,7 @@ def _install_mocks( monkeypatch.setattr( mcp_oauth, "get_connection_config_by_id", - lambda config_id, _db_session: SimpleNamespace(id=config_id), + lambda config_id, _db_session, **_kwargs: SimpleNamespace(id=config_id), ) # extract_connection_data returns the same dict the SDK storage mutates. monkeypatch.setattr( diff --git a/backend/tests/unit/onyx/server/test_mcp_known_provider_oauth_helpers.py b/backend/tests/unit/onyx/server/test_mcp_known_provider_oauth_helpers.py index 50d43dd7f6..4bdfb58f5d 100644 --- a/backend/tests/unit/onyx/server/test_mcp_known_provider_oauth_helpers.py +++ b/backend/tests/unit/onyx/server/test_mcp_known_provider_oauth_helpers.py @@ -275,7 +275,7 @@ def _patch_config_read( monkeypatch.setattr( mcp_oauth.OnyxTokenStorage, "_ensure_connection_config", - lambda _self, _db: SimpleNamespace(id=1), + lambda _self, _db, **_kwargs: SimpleNamespace(id=1), ) monkeypatch.setattr( mcp_oauth, "extract_connection_data", lambda _config: config_data