From d7f0a31e2464881d49da0c776ba27107f14c6f39 Mon Sep 17 00:00:00 2001
From: quicksand <quicksandzn@gmail.com>
Date: Fri, 17 Oct 2025 09:09:45 +0800
Subject: [PATCH] =?UTF-8?q?Fix:=20User=20Context=20Loss=20When=20Invoking?=
 =?UTF-8?q?=20Workflow=20Tool=20Node=20in=20Knowledge=20=E2=80=A6=20(#2649?=
 =?UTF-8?q?5)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 api/core/tools/workflow_as_tool/tool.py       | 58 ++++++++++++++++++-
 .../core/tools/workflow_as_tool/test_tool.py  |  7 ++-
 2 files changed, 61 insertions(+), 4 deletions(-)

diff --git a/api/core/tools/workflow_as_tool/tool.py b/api/core/tools/workflow_as_tool/tool.py
index 5adf04611d..50c2327004 100644
--- a/api/core/tools/workflow_as_tool/tool.py
+++ b/api/core/tools/workflow_as_tool/tool.py
@@ -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
diff --git a/api/tests/unit_tests/core/tools/workflow_as_tool/test_tool.py b/api/tests/unit_tests/core/tools/workflow_as_tool/test_tool.py
index 17e3ebeea0..c68aad0b22 100644
--- a/api/tests/unit_tests/core/tools/workflow_as_tool/test_tool.py
+++ b/api/tests/unit_tests/core/tools/workflow_as_tool/test_tool.py
@@ -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