mirror of
https://github.com/langgenius/dify.git
synced 2026-06-16 22:11:09 +08:00
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 <noreply@anthropic.com>
This commit is contained in:
parent
7e4b1cadd8
commit
90567c537f
@ -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"
|
||||
|
||||
@ -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:
|
||||
|
||||
@ -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."""
|
||||
|
||||
@ -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
|
||||
|
||||
Loading…
Reference in New Issue
Block a user