From 0fe078d25e2ea64d6df473f338b5b93c63b28950 Mon Sep 17 00:00:00 2001
From: horochx <32632779+horochx@users.noreply.github.com>
Date: Sat, 30 Aug 2025 00:59:21 +0800
Subject: [PATCH] =?UTF-8?q?fix:=20workflow=5Ffinish=5Fto=5Fstream=5Frespon?=
 =?UTF-8?q?se=20assert=20exception=20with=20celery=20=E2=80=A6=20(#24674)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
---
 .../advanced_chat/generate_task_pipeline.py   |  1 +
 .../common/workflow_response_converter.py     | 39 ++++++++-----------
 .../apps/workflow/generate_task_pipeline.py   |  1 +
 3 files changed, 18 insertions(+), 23 deletions(-)

diff --git a/api/core/app/apps/advanced_chat/generate_task_pipeline.py b/api/core/app/apps/advanced_chat/generate_task_pipeline.py
index a61bba512..fb61b4c35 100644
--- a/api/core/app/apps/advanced_chat/generate_task_pipeline.py
+++ b/api/core/app/apps/advanced_chat/generate_task_pipeline.py
@@ -144,6 +144,7 @@ class AdvancedChatAppGenerateTaskPipeline:
 
         self._workflow_response_converter = WorkflowResponseConverter(
             application_generate_entity=application_generate_entity,
+            user=user,
         )
 
         self._task_state = WorkflowTaskState()
diff --git a/api/core/app/apps/common/workflow_response_converter.py b/api/core/app/apps/common/workflow_response_converter.py
index 1a8923733..c8760d3cf 100644
--- a/api/core/app/apps/common/workflow_response_converter.py
+++ b/api/core/app/apps/common/workflow_response_converter.py
@@ -3,7 +3,6 @@ from collections.abc import Mapping, Sequence
 from datetime import UTC, datetime
 from typing import Any, Optional, Union, cast
 
-from sqlalchemy import select
 from sqlalchemy.orm import Session
 
 from core.app.entities.app_invoke_entities import AdvancedChatAppGenerateEntity, WorkflowAppGenerateEntity
@@ -53,9 +52,7 @@ from core.workflow.workflow_type_encoder import WorkflowRuntimeTypeConverter
 from libs.datetime_utils import naive_utc_now
 from models import (
     Account,
-    CreatorUserRole,
     EndUser,
-    WorkflowRun,
 )
 
 
@@ -64,8 +61,10 @@ class WorkflowResponseConverter:
         self,
         *,
         application_generate_entity: Union[AdvancedChatAppGenerateEntity, WorkflowAppGenerateEntity],
+        user: Union[Account, EndUser],
     ) -> None:
         self._application_generate_entity = application_generate_entity
+        self._user = user
 
     def workflow_start_to_stream_response(
         self,
@@ -92,27 +91,21 @@ class WorkflowResponseConverter:
         workflow_execution: WorkflowExecution,
     ) -> WorkflowFinishStreamResponse:
         created_by = None
-        workflow_run = session.scalar(select(WorkflowRun).where(WorkflowRun.id == workflow_execution.id_))
-        assert workflow_run is not None
-        if workflow_run.created_by_role == CreatorUserRole.ACCOUNT:
-            stmt = select(Account).where(Account.id == workflow_run.created_by)
-            account = session.scalar(stmt)
-            if account:
-                created_by = {
-                    "id": account.id,
-                    "name": account.name,
-                    "email": account.email,
-                }
-        elif workflow_run.created_by_role == CreatorUserRole.END_USER:
-            stmt = select(EndUser).where(EndUser.id == workflow_run.created_by)
-            end_user = session.scalar(stmt)
-            if end_user:
-                created_by = {
-                    "id": end_user.id,
-                    "user": end_user.session_id,
-                }
+
+        user = self._user
+        if isinstance(user, Account):
+            created_by = {
+                "id": user.id,
+                "name": user.name,
+                "email": user.email,
+            }
+        elif isinstance(user, EndUser):
+            created_by = {
+                "id": user.id,
+                "user": user.session_id,
+            }
         else:
-            raise NotImplementedError(f"unknown created_by_role: {workflow_run.created_by_role}")
+            raise NotImplementedError(f"User type not supported: {type(user)}")
 
         # Handle the case where finished_at is None by using current time as default
         finished_at_timestamp = (
diff --git a/api/core/app/apps/workflow/generate_task_pipeline.py b/api/core/app/apps/workflow/generate_task_pipeline.py
index 537c070ad..58e51ccca 100644
--- a/api/core/app/apps/workflow/generate_task_pipeline.py
+++ b/api/core/app/apps/workflow/generate_task_pipeline.py
@@ -131,6 +131,7 @@ class WorkflowAppGenerateTaskPipeline:
 
         self._workflow_response_converter = WorkflowResponseConverter(
             application_generate_entity=application_generate_entity,
+            user=user,
         )
 
         self._application_generate_entity = application_generate_entity