Fix: User Context Loss When Invoking Workflow Tool Node in Knowledge … (#26495)

This commit is contained in:
quicksand 2025-10-17 09:09:45 +08:00 committed by GitHub
parent 312974aa20
commit d7f0a31e24
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 61 additions and 4 deletions

View File

@ -3,6 +3,7 @@ import logging
from collections.abc import Generator
from typing import Any
from flask import has_request_context
from sqlalchemy import select
from core.file import FILE_MODEL_IDENTITY, File, FileTransferMethod
@ -18,7 +19,8 @@ from core.tools.errors import ToolInvokeError
from extensions.ext_database import db
from factories.file_factory import build_from_mapping
from libs.login import current_user
from models.model import App
from models import Account, Tenant
from models.model import App, EndUser
from models.workflow import Workflow
logger = logging.getLogger(__name__)
@ -79,11 +81,16 @@ class WorkflowTool(Tool):
generator = WorkflowAppGenerator()
assert self.runtime is not None
assert self.runtime.invoke_from is not None
assert current_user is not None
user = self._resolve_user(user_id=user_id)
if user is None:
raise ToolInvokeError("User not found")
result = generator.generate(
app_model=app,
workflow=workflow,
user=current_user,
user=user,
args={"inputs": tool_parameters, "files": files},
invoke_from=self.runtime.invoke_from,
streaming=False,
@ -123,6 +130,51 @@ class WorkflowTool(Tool):
label=self.label,
)
def _resolve_user(self, user_id: str) -> Account | EndUser | None:
"""
Resolve user object in both HTTP and worker contexts.
In HTTP context: dereference the current_user LocalProxy (can return Account or EndUser).
In worker context: load Account from database by user_id (only returns Account, never EndUser).
Returns:
Account | EndUser | None: The resolved user object, or None if resolution fails.
"""
if has_request_context():
return self._resolve_user_from_request()
else:
return self._resolve_user_from_database(user_id=user_id)
def _resolve_user_from_request(self) -> Account | EndUser | None:
"""
Resolve user from Flask request context.
"""
try:
# Note: `current_user` is a LocalProxy. Never compare it with None directly.
return getattr(current_user, "_get_current_object", lambda: current_user)()
except Exception as e:
logger.warning("Failed to resolve user from request context: %s", e)
return None
def _resolve_user_from_database(self, user_id: str) -> Account | None:
"""
Resolve user from database (worker/Celery context).
"""
user_stmt = select(Account).where(Account.id == user_id)
user = db.session.scalar(user_stmt)
if not user:
return None
tenant_stmt = select(Tenant).where(Tenant.id == self.runtime.tenant_id)
tenant = db.session.scalar(tenant_stmt)
if not tenant:
return None
user.current_tenant = tenant
return user
def _get_workflow(self, app_id: str, version: str) -> Workflow:
"""
get the workflow by app id and version

View File

@ -34,12 +34,17 @@ def test_workflow_tool_should_raise_tool_invoke_error_when_result_has_error_fiel
monkeypatch.setattr(tool, "_get_app", lambda *args, **kwargs: None)
monkeypatch.setattr(tool, "_get_workflow", lambda *args, **kwargs: None)
# Mock user resolution to avoid database access
from unittest.mock import Mock
mock_user = Mock()
monkeypatch.setattr(tool, "_resolve_user", lambda *args, **kwargs: mock_user)
# replace `WorkflowAppGenerator.generate` 's return value.
monkeypatch.setattr(
"core.app.apps.workflow.app_generator.WorkflowAppGenerator.generate",
lambda *args, **kwargs: {"data": {"error": "oops"}},
)
monkeypatch.setattr("libs.login.current_user", lambda *args, **kwargs: None)
with pytest.raises(ToolInvokeError) as exc_info:
# WorkflowTool always returns a generator, so we need to iterate to