From 0e14d07adb187a1758b0ac3e562e911c6db7fc6a Mon Sep 17 00:00:00 2001 From: Charles Yao Date: Fri, 12 Jun 2026 09:30:15 +0200 Subject: [PATCH] feat(api): forward user_type for MCP identity forwarding (webapp end-users) (#37347) Co-authored-by: Claude Opus 4.8 Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com> --- api/core/app/apps/base_app_queue_manager.py | 2 +- api/core/app/apps/base_app_runner.py | 4 +-- api/core/app/apps/workflow_app_runner.py | 2 +- api/core/app/entities/app_invoke_entities.py | 8 +++++ api/core/tools/mcp_tool/tool.py | 10 ++++++ api/services/enterprise/enterprise_service.py | 2 ++ .../unit_tests/core/app/test_invoke_from.py | 16 ++++++++++ .../unit_tests/core/tools/test_mcp_tool.py | 32 +++++++++++++++++++ .../enterprise/test_enterprise_service.py | 15 +++++++++ 9 files changed, 86 insertions(+), 5 deletions(-) diff --git a/api/core/app/apps/base_app_queue_manager.py b/api/core/app/apps/base_app_queue_manager.py index b959987078e..9551a5e38c4 100644 --- a/api/core/app/apps/base_app_queue_manager.py +++ b/api/core/app/apps/base_app_queue_manager.py @@ -41,7 +41,7 @@ class AppQueueManager(ABC): self._invoke_from = invoke_from self.invoke_from = invoke_from # Public accessor for invoke_from - user_prefix = "account" if self._invoke_from in {InvokeFrom.EXPLORE, InvokeFrom.DEBUGGER} else "end-user" + user_prefix = "account" if self._invoke_from.runs_as_account() else "end-user" self._task_belong_cache_key = AppQueueManager._generate_task_belong_cache_key(self._task_id) redis_client.setex(self._task_belong_cache_key, 1800, f"{user_prefix}-{self._user_id}") diff --git a/api/core/app/apps/base_app_runner.py b/api/core/app/apps/base_app_runner.py index 15f6359929a..a89a0cf70db 100644 --- a/api/core/app/apps/base_app_runner.py +++ b/api/core/app/apps/base_app_runner.py @@ -431,9 +431,7 @@ class AppRunner: url=f"/files/tools/{tool_file.id}", upload_file_id=tool_file.id, created_by_role=( - CreatorUserRole.ACCOUNT - if queue_manager.invoke_from in {InvokeFrom.DEBUGGER, InvokeFrom.EXPLORE} - else CreatorUserRole.END_USER + CreatorUserRole.ACCOUNT if queue_manager.invoke_from.runs_as_account() else CreatorUserRole.END_USER ), created_by=user_id, ) diff --git a/api/core/app/apps/workflow_app_runner.py b/api/core/app/apps/workflow_app_runner.py index 944860ee39c..69f6c5b69b7 100644 --- a/api/core/app/apps/workflow_app_runner.py +++ b/api/core/app/apps/workflow_app_runner.py @@ -104,7 +104,7 @@ class WorkflowBasedAppRunner: @staticmethod def _resolve_user_from(invoke_from: InvokeFrom) -> UserFrom: - if invoke_from in {InvokeFrom.EXPLORE, InvokeFrom.DEBUGGER}: + if invoke_from.runs_as_account(): return UserFrom.ACCOUNT return UserFrom.END_USER diff --git a/api/core/app/entities/app_invoke_entities.py b/api/core/app/entities/app_invoke_entities.py index 08ecc2097b3..2153289e0e6 100644 --- a/api/core/app/entities/app_invoke_entities.py +++ b/api/core/app/entities/app_invoke_entities.py @@ -47,6 +47,14 @@ class InvokeFrom(StrEnum): } return source_mapping.get(self, "dev") + def runs_as_account(self) -> bool: + """Whether a run from this entry point is attributed to a workspace + Account rather than an end user. Console contexts (debugger/explore) + run as the signed-in Account; webapp/service-api/trigger run as an + EndUser. Single source of truth for the created-by-role / user-type + split shared by the app runners and MCP identity forwarding.""" + return self in (InvokeFrom.DEBUGGER, InvokeFrom.EXPLORE) + class DifyRunContext(BaseModel): tenant_id: str diff --git a/api/core/tools/mcp_tool/tool.py b/api/core/tools/mcp_tool/tool.py index dc34264fb41..7a1553a4b15 100644 --- a/api/core/tools/mcp_tool/tool.py +++ b/api/core/tools/mcp_tool/tool.py @@ -358,7 +358,17 @@ 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.""" + invoke_from = self.runtime.invoke_from + if invoke_from is not None and invoke_from.runs_as_account(): + return "account" + return "end_user" diff --git a/api/services/enterprise/enterprise_service.py b/api/services/enterprise/enterprise_service.py index ccfae848d1a..d7dbd12973d 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/app/test_invoke_from.py b/api/tests/unit_tests/core/app/test_invoke_from.py index e0a8344d2f6..33a4d2edf11 100644 --- a/api/tests/unit_tests/core/app/test_invoke_from.py +++ b/api/tests/unit_tests/core/app/test_invoke_from.py @@ -7,3 +7,19 @@ def test_openapi_variant_present(): def test_openapi_distinct_from_service_api(): assert InvokeFrom.OPENAPI != InvokeFrom.SERVICE_API + + +def test_runs_as_account_only_for_console_contexts(): + # Console contexts (studio debugger / explore) run as the signed-in Account. + assert InvokeFrom.DEBUGGER.runs_as_account() is True + assert InvokeFrom.EXPLORE.runs_as_account() is True + # Everything else is attributed to an end user. + for invoke_from in ( + InvokeFrom.WEB_APP, + InvokeFrom.SERVICE_API, + InvokeFrom.OPENAPI, + InvokeFrom.TRIGGER, + InvokeFrom.PUBLISHED_PIPELINE, + InvokeFrom.VALIDATION, + ): + assert invoke_from.runs_as_account() is False 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 cdcb628972c..2e2b961bf2a 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,38 @@ 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 e7efe79af00..f64c7233b9d 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