From 90567c537fb008403ab6ab64305b833a8e874df3 Mon Sep 17 00:00:00 2001 From: Charles Yao Date: Thu, 11 Jun 2026 22:44:43 +0200 Subject: [PATCH] feat(api): forward user_type for MCP identity forwarding issue_mcp_token gains a user_type field (account|end_user), derived from invoke_from ({DEBUGGER,EXPLORE}->account else end_user) so the enterprise side routes webapp end-users to the published-webapp token store. Co-Authored-By: Claude Opus 4.8 --- api/core/tools/mcp_tool/tool.py | 15 ++++++++++ api/services/enterprise/enterprise_service.py | 2 ++ .../unit_tests/core/tools/test_mcp_tool.py | 30 +++++++++++++++++++ .../enterprise/test_enterprise_service.py | 15 ++++++++++ 4 files changed, 62 insertions(+) diff --git a/api/core/tools/mcp_tool/tool.py b/api/core/tools/mcp_tool/tool.py index b0f1f7a5f2..6779e535fa 100644 --- a/api/core/tools/mcp_tool/tool.py +++ b/api/core/tools/mcp_tool/tool.py @@ -7,6 +7,7 @@ from collections.abc import Generator, Mapping from typing import Any, cast, override from configs import dify_config +from core.app.entities.app_invoke_entities import InvokeFrom from core.entities.mcp_provider import IdentityMode from core.mcp.auth_client import MCPClientWithAuthRetry from core.mcp.error import MCPConnectionError @@ -32,6 +33,11 @@ logger = logging.getLogger(__name__) # user-supplied custom credentials), which would silently break those flows. FORWARDED_IDENTITY_HEADER = "X-Dify-SSO-Access-Token" +# invoke_from values where the caller is a console Account rather than a webapp/ +# API end-user. Mirrors core/app/apps/base_app_runner.py and +# workflow_app_runner._resolve_user_from. +_ACCOUNT_INVOKE_FROMS = frozenset({InvokeFrom.DEBUGGER, InvokeFrom.EXPLORE}) + class MCPTool(Tool): def __init__( @@ -358,7 +364,16 @@ class MCPTool(Tool): tenant_id=self.tenant_id, app_id=app_id, audience=audience, + user_type=self._resolve_user_type(), ) except MCPTokenError as e: raise ToolInvokeError(f"Failed to obtain forwarded identity token: {e}") from e headers[FORWARDED_IDENTITY_HEADER] = token + + def _resolve_user_type(self) -> str: + """Return "account" for console-authenticated callers (debugger/explore), + "end_user" for webapp / service-api / trigger callers — so the enterprise + side routes to the console store vs the published-webapp store.""" + if self.runtime.invoke_from in _ACCOUNT_INVOKE_FROMS: + return "account" + return "end_user" diff --git a/api/services/enterprise/enterprise_service.py b/api/services/enterprise/enterprise_service.py index 50a39ef927..a0549f0dde 100644 --- a/api/services/enterprise/enterprise_service.py +++ b/api/services/enterprise/enterprise_service.py @@ -136,6 +136,7 @@ class EnterpriseService: tenant_id: str, app_id: str | None, audience: str, + user_type: str = "account", ) -> 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. @@ -163,6 +164,7 @@ class EnterpriseService: "tenant_id": tenant_id, "app_id": app_id or "", "audience": audience, + "user_type": user_type, }, ) except EnterpriseServiceError as e: diff --git a/api/tests/unit_tests/core/tools/test_mcp_tool.py b/api/tests/unit_tests/core/tools/test_mcp_tool.py index 1504889f01..3ac421dce5 100644 --- a/api/tests/unit_tests/core/tools/test_mcp_tool.py +++ b/api/tests/unit_tests/core/tools/test_mcp_tool.py @@ -219,6 +219,36 @@ def test_inject_forwarded_identity_translates_token_error_to_invoke_error(): assert "Authorization" not in headers +def test_inject_forwarded_identity_sends_end_user_type_for_webapp(): + """A WEB_APP run forwards user_type=end_user so enterprise routes to the + published-webapp token store.""" + tool = _build_forwarding_tool() + tool.runtime = ToolRuntime(tenant_id="tenant-1", invoke_from=InvokeFrom.WEB_APP) + headers: dict[str, str] = {} + + with patch( + "services.enterprise.enterprise_service.EnterpriseService.issue_mcp_token", + return_value=("forwarded.jwt", 1900000000), + ) as issue: + tool._inject_forwarded_identity(headers, user_id="eu-1", app_id="app-1", audience="https://mcp.example.com/mcp/") + + assert issue.call_args.kwargs["user_type"] == "end_user" + + +def test_inject_forwarded_identity_sends_account_type_for_debugger(): + """A DEBUGGER/console run forwards user_type=account (the existing behaviour).""" + tool = _build_forwarding_tool() # built with InvokeFrom.DEBUGGER + headers: dict[str, str] = {} + + with patch( + "services.enterprise.enterprise_service.EnterpriseService.issue_mcp_token", + return_value=("forwarded.jwt", 1900000000), + ) as issue: + tool._inject_forwarded_identity(headers, user_id="acc-1", app_id=None, audience="https://mcp.example.com/mcp/") + + assert issue.call_args.kwargs["user_type"] == "account" + + def test_invoke_remote_mcp_tool_fails_closed_when_user_id_missing(): """When forwarding is enabled AND the deployment is enterprise, missing user_id must raise — never silently invoke as the static identity.""" diff --git a/api/tests/unit_tests/services/enterprise/test_enterprise_service.py b/api/tests/unit_tests/services/enterprise/test_enterprise_service.py index e7efe79af0..f64c7233b9 100644 --- a/api/tests/unit_tests/services/enterprise/test_enterprise_service.py +++ b/api/tests/unit_tests/services/enterprise/test_enterprise_service.py @@ -500,9 +500,24 @@ class TestIssueMCPToken: "tenant_id": "tenant-uuid", "app_id": "app-uuid", "audience": "https://mcp.example.com/mcp/", + "user_type": "account", }, ) + def test_end_user_type_is_forwarded(self): + with patch(f"{MODULE}.EnterpriseRequest") as req: + req.send_request.return_value = {"token": "t", "expires_at": 1900000000} + EnterpriseService.issue_mcp_token( + user_id="end-user-uuid", + tenant_id="tenant-uuid", + app_id="app-uuid", + audience="https://mcp.example.com/mcp/", + user_type="end_user", + ) + body = req.send_request.call_args.kwargs["json"] + assert body["user_type"] == "end_user" + assert body["app_id"] == "app-uuid" + def test_401_maps_to_identity_refresh_error(self): from services.enterprise.base import MCPIdentityRefreshError from services.errors.enterprise import EnterpriseAPIUnauthorizedError