diff --git a/api/services/tools/mcp_tools_manage_service.py b/api/services/tools/mcp_tools_manage_service.py index 2b52615338e307..763c7b44158b0e 100644 --- a/api/services/tools/mcp_tools_manage_service.py +++ b/api/services/tools/mcp_tools_manage_service.py @@ -800,12 +800,20 @@ def _process_headers(self, headers: dict[str, str], mcp_provider: MCPToolProvide def _process_credentials( self, authentication: MCPAuthentication, mcp_provider: MCPToolProvider, tenant_id: str ) -> str: - """Process credentials update, handling masked values.""" + """Process credentials update, handling masked values and keeping stored OAuth state when unchanged.""" # Merge with existing credentials final_client_id, final_client_secret = self._merge_credentials_with_masked( authentication.client_id, authentication.client_secret, mcp_provider ) + existing_client_info = mcp_provider.to_entity().retrieve_client_information() + if ( + existing_client_info + and final_client_id == existing_client_info.client_id + and final_client_secret == existing_client_info.client_secret + ): + return mcp_provider.encrypted_credentials or EMPTY_CREDENTIALS_JSON + # Build and encrypt return self._build_and_encrypt_credentials(final_client_id, final_client_secret, tenant_id) @@ -851,20 +859,24 @@ def _merge_credentials_with_masked( Tuple of (final_client_id, final_client_secret) """ mcp_provider_entity = mcp_provider.to_entity() - existing_decrypted = mcp_provider_entity.decrypt_credentials() + existing_client_info = mcp_provider_entity.retrieve_client_information() existing_masked = mcp_provider_entity.masked_credentials() # Check if client_id is masked and unchanged final_client_id = client_id - if existing_masked.get("client_id") and client_id == existing_masked["client_id"]: + if existing_client_info and existing_masked.get("client_id") and client_id == existing_masked["client_id"]: # Use existing decrypted value - final_client_id = existing_decrypted.get("client_id", client_id) + final_client_id = existing_client_info.client_id # Check if client_secret is masked and unchanged final_client_secret = client_secret - if existing_masked.get("client_secret") and client_secret == existing_masked["client_secret"]: + if ( + existing_client_info + and existing_masked.get("client_secret") + and client_secret == existing_masked["client_secret"] + ): # Use existing decrypted value - final_client_secret = existing_decrypted.get("client_secret", client_secret) + final_client_secret = existing_client_info.client_secret return final_client_id, final_client_secret diff --git a/api/tests/unit_tests/services/tools/test_mcp_tools_manage_service.py b/api/tests/unit_tests/services/tools/test_mcp_tools_manage_service.py index 95ef5320ed023f..cec0a005717197 100644 --- a/api/tests/unit_tests/services/tools/test_mcp_tools_manage_service.py +++ b/api/tests/unit_tests/services/tools/test_mcp_tools_manage_service.py @@ -1,14 +1,17 @@ """Exercise tenant-scoped MCP provider lookups against persisted rows. Primary keys and server identifiers are distinct, even when an identifier looks -like a UUID or collides with another provider's primary key. +like a UUID or collides with another provider's primary key. Unchanged masked +OAuth client credentials resolve to the stored client information on update. """ +import json from typing import Literal import pytest from sqlalchemy.orm import Session +from core.entities.mcp_provider import MCPAuthentication, MCPConfiguration from models.tools import MCPToolProvider from services.tools.mcp_tools_manage_service import MCPToolManageService @@ -158,3 +161,61 @@ def test_server_identifier_can_be_reused_across_tenants(service: MCPToolManageSe service.get_provider_by_server_identifier(server_identifier=SERVER_IDENTIFIER, tenant_id=TENANT_ID).id == provider.id ) + + +def test_merge_credentials_with_masked_restores_nested_client_information( + service: MCPToolManageService, sqlite_session: Session +) -> None: + provider = _provider() + provider.encrypted_credentials = json.dumps( + { + "client_information": { + "client_id": "ba0f3e1c-registered-client-c2", + "client_secret": "dynamic-client-secret", + "redirect_uris": ["https://example.com/callback"], + } + } + ) + _persist(sqlite_session, provider) + masked = provider.to_entity().masked_credentials() + + assert service._merge_credentials_with_masked(masked["client_id"], masked["client_secret"], provider) == ( + "ba0f3e1c-registered-client-c2", + "dynamic-client-secret", + ) + + +def test_update_provider_keeps_oauth_state_when_masked_credentials_are_unchanged( + service: MCPToolManageService, sqlite_session: Session +) -> None: + provider = _provider() + credentials = json.dumps( + { + "client_information": { + "client_id": "ba0f3e1c-registered-client-c2", + "client_secret": "dynamic-client-secret", + "redirect_uris": ["https://example.com/callback"], + "is_dynamic_registration": True, + }, + "access_token": "access-token", + "refresh_token": "refresh-token", + } + ) + provider.encrypted_credentials = credentials + _persist(sqlite_session, provider) + masked = provider.to_entity().masked_credentials() + + service.update_provider( + tenant_id=TENANT_ID, + provider_id=PROVIDER_UUID, + name="Renamed MCP", + server_url=provider.server_url, + icon="icon", + icon_type="link", + icon_background="", + server_identifier=SERVER_IDENTIFIER, + configuration=MCPConfiguration(), + authentication=MCPAuthentication(client_id=masked["client_id"], client_secret=masked["client_secret"]), + ) + + assert provider.encrypted_credentials == credentials