feat(api): add MCP user-identity forwarding (#36839)

Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>

api/services/data_migration/import_service.py

@@ -18,7 +18,7 @@ import yaml
 from sqlalchemy import or_
 from sqlalchemy.orm import Session, sessionmaker
 
-from core.entities.mcp_provider import MCPAuthentication, MCPConfiguration
+from core.entities.mcp_provider import IdentityMode, MCPAuthentication, MCPConfiguration
 from core.tools.entities.tool_entities import ApiProviderSchemaType, WorkflowToolParameterConfiguration
 from extensions.ext_database import db
 from libs.datetime_utils import naive_utc_now
@@ -748,6 +748,9 @@ class MigrationImportService:
                     headers=mcp_data.get("headers") if isinstance(mcp_data.get("headers"), dict) else {},
                     configuration=configuration,
                     authentication=authentication,
+                    # Re-import must not silently reset forwarding: preserve the
+                    # stored mode (update_provider now defaults to OFF when omitted).
+                    identity_mode=IdentityMode(existing.identity_mode),
                 )
                 db.session.commit()
                 status = "updated"

api/services/enterprise/base.py

@@ -12,8 +12,28 @@ from services.errors.enterprise import (
     EnterpriseAPIForbiddenError,
     EnterpriseAPINotFoundError,
     EnterpriseAPIUnauthorizedError,
+    EnterpriseServiceError,
 )
 
+
+class MCPTokenError(EnterpriseServiceError):
+    """Generic failure of the IssueMCPToken RPC."""
+
+
+class MCPNoRefreshTokenError(MCPTokenError):
+    """User has no stored SSO refresh_token; ask them to re-authenticate."""
+
+    def __init__(self, description: str = ""):
+        super().__init__(description, status_code=428)
+
+
+class MCPIdentityRefreshError(MCPTokenError):
+    """IdP rejected the refresh attempt (revoked/expired session)."""
+
+    def __init__(self, description: str = ""):
+        super().__init__(description, status_code=401)
+
+
 logger = logging.getLogger(__name__)
 
 

api/services/enterprise/enterprise_service.py

@@ -11,7 +11,15 @@ from pydantic import BaseModel, ConfigDict, Field, model_validator
 
 from configs import dify_config
 from extensions.ext_redis import redis_client
-from services.enterprise.base import EnterpriseRequest
+from services.enterprise.base import (
+    EnterpriseRequest,
+    MCPIdentityRefreshError,
+    MCPNoRefreshTokenError,
+    MCPTokenError,
+)
+from services.errors.enterprise import (
+    EnterpriseServiceError,
+)
 
 if TYPE_CHECKING:
     from services.feature_service import LicenseStatus
@@ -121,6 +129,77 @@ class EnterpriseService:
     def get_workspace_info(cls, tenant_id: str):
         return EnterpriseRequest.send_request("GET", f"/workspace/{tenant_id}/info")
 
+    @classmethod
+    def issue_mcp_token(
+        cls,
+        user_id: str,
+        tenant_id: str,
+        app_id: str | None,
+        audience: str,
+    ) -> tuple[str, int]:
+        """Mint a short-lived SSO id_token (or OAuth2 access_token) representing
+        the calling Dify user, audience-scoped to the given MCP server identifier.
+
+        Used by MCPTool.invoke_remote_mcp_tool to stamp the
+        X-Dify-SSO-Access-Token header on outbound MCP requests when the
+        provider's identity_mode is set to "idp_token".
+
+        Returns:
+            (token, expires_at_unix_seconds)
+
+        Raises:
+            MCPNoRefreshTokenError: user has no stored SSO refresh_token on the
+                enterprise side; surface to the workflow as "please log in via SSO".
+            MCPIdentityRefreshError: enterprise tried to refresh against the IdP
+                and the IdP rejected (revoked/expired session).
+            MCPTokenError: any other failure of the enterprise endpoint.
+        """
+        try:
+            response = EnterpriseRequest.send_request(
+                "POST",
+                "/mcp/issue-token",
+                json={
+                    "user_id": user_id,
+                    "tenant_id": tenant_id,
+                    "app_id": app_id or "",
+                    "audience": audience,
+                },
+            )
+        except EnterpriseServiceError as e:
+            # The HTTP-status subclasses (400/401/403/404) inherit directly
+            # from EnterpriseServiceError, not EnterpriseAPIError, so we
+            # must catch the base class to route them all.
+            status = getattr(e, "status_code", None)
+            if status == 401:
+                # Enterprise side returns 401 when the IdP rejected the refresh.
+                raise MCPIdentityRefreshError(str(e) or "identity refresh failed; please re-authenticate") from e
+            if status == 428:
+                raise MCPNoRefreshTokenError(
+                    str(e) or "user has no stored SSO refresh token; please re-authenticate"
+                ) from e
+            if status == 403:
+                # 403 most often means the tenant isn't licensed for MCP
+                # identity-forwarding. Surface as identity-refresh-failure so
+                # the workflow halts loudly rather than retrying.
+                raise MCPIdentityRefreshError(
+                    str(e) or "enterprise refused to issue an MCP identity token (license or policy)"
+                ) from e
+            raise MCPTokenError(f"issue_mcp_token failed (status={status}): {e}") from e
+
+        if not isinstance(response, dict):
+            raise MCPTokenError("invalid response shape from enterprise /mcp/issue-token")
+
+        token = response.get("token")
+        expires_at = response.get("expires_at")
+        # Accept int or float for expires_at (some clocks emit float
+        # seconds-since-epoch). Reject bools explicitly because `bool` is
+        # an `int` subclass in Python and would pass isinstance(_, int).
+        if not isinstance(token, str) or not token:
+            raise MCPTokenError(f"missing or non-string token in enterprise response: {response!r}")
+        if isinstance(expires_at, bool) or not isinstance(expires_at, (int, float)):
+            raise MCPTokenError(f"missing or non-numeric expires_at in enterprise response: {response!r}")
+        return token, int(expires_at)
+
     @classmethod
     def initiate_device_flow_sso(cls, signed_state: str) -> dict:
         return EnterpriseRequest.send_request(

api/services/tools/mcp_tools_manage_service.py

@@ -12,7 +12,7 @@ from sqlalchemy import or_, select
 from sqlalchemy.exc import IntegrityError
 from sqlalchemy.orm import Session
 
-from core.entities.mcp_provider import MCPAuthentication, MCPConfiguration, MCPProviderEntity
+from core.entities.mcp_provider import IdentityMode, MCPAuthentication, MCPConfiguration, MCPProviderEntity
 from core.helper import encrypter
 from core.helper.provider_cache import NoOpProviderCredentialCache
 from core.mcp.auth.auth_flow import auth
@@ -136,6 +136,7 @@ class MCPToolManageService:
         configuration: MCPConfiguration,
         authentication: MCPAuthentication | None = None,
         headers: dict[str, str] | None = None,
+        identity_mode: IdentityMode = IdentityMode.OFF,
     ) -> ToolProviderApiEntity:
         """Create a new MCP provider."""
         # Validate URL format
@@ -171,6 +172,7 @@ class MCPToolManageService:
             sse_read_timeout=configuration.sse_read_timeout,
             encrypted_headers=encrypted_headers,
             encrypted_credentials=encrypted_credentials,
+            identity_mode=identity_mode,
         )
 
         self._session.add(mcp_tool)
@@ -194,6 +196,7 @@ class MCPToolManageService:
         configuration: MCPConfiguration,
         authentication: MCPAuthentication | None = None,
         validation_result: ServerUrlValidationResult | None = None,
+        identity_mode: IdentityMode = IdentityMode.OFF,
     ) -> None:
         """
         Update an MCP provider.
@@ -255,6 +258,11 @@ class MCPToolManageService:
             if authentication and authentication.client_id:
                 mcp_provider.encrypted_credentials = self._process_credentials(authentication, mcp_provider, tenant_id)
 
+            # Update user-identity forwarding mode. The controller has already
+            # resolved "leave unchanged" and applied the ENTERPRISE_ENABLED gate,
+            # so this is always a concrete, vetted value.
+            mcp_provider.identity_mode = identity_mode
+
             # Flush changes to database
             self._session.flush()