From 0ef793f935cb8bc2c00fe3d96d9b307a286dbd27 Mon Sep 17 00:00:00 2001
From: Yunlu Wen <yunlu.wen@dify.ai>
Date: Fri, 15 May 2026 16:30:31 +0800
Subject: [PATCH] fix: performance optimization on TTFE critical path (#36185)

---
 .../app/apps/advanced_chat/generate_task_pipeline.py   | 10 ++++------
 api/core/app/llm/model_access.py                       |  8 ++++++--
 api/core/model_manager.py                              |  4 ++--
 3 files changed, 12 insertions(+), 10 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 f7b5030d33..cd775f54aa 100644
--- a/api/core/app/apps/advanced_chat/generate_task_pipeline.py
+++ b/api/core/app/apps/advanced_chat/generate_task_pipeline.py
@@ -7,7 +7,7 @@ from contextlib import contextmanager
 from threading import Thread
 from typing import Any, Union
 
-from sqlalchemy import select
+from sqlalchemy import select, update
 from sqlalchemy.orm import Session
 
 from constants.tts_auto_play_timeout import TTS_AUTO_PLAY_TIMEOUT, TTS_AUTO_PLAY_YIELD_CPU_TIME
@@ -311,11 +311,9 @@ class AdvancedChatAppGenerateTaskPipeline(GraphRuntimeStateSupport):
         self._workflow_run_id = run_id
 
         with self._database_session() as session:
-            message = self._get_message(session=session)
-            if not message:
-                raise ValueError(f"Message not found: {self._message_id}")
-
-            message.workflow_run_id = run_id
+            session.execute(
+                update(Message).where(Message.id == self._message_id).values(workflow_run_id=run_id)
+            )
 
         workflow_start_resp = self._workflow_response_converter.workflow_start_to_stream_response(
             task_id=self._application_generate_entity.task_id,
diff --git a/api/core/app/llm/model_access.py b/api/core/app/llm/model_access.py
index 1d6c377d8f..a81ca0ec71 100644
--- a/api/core/app/llm/model_access.py
+++ b/api/core/app/llm/model_access.py
@@ -63,9 +63,13 @@ class DifyModelFactory:
 
 
 def build_dify_model_access(tenant_id: str) -> tuple[CredentialsProvider, ModelFactory]:
+    provider_manager = ProviderManager()
     return (
-        DifyCredentialsProvider(tenant_id=tenant_id),
-        DifyModelFactory(tenant_id=tenant_id),
+        DifyCredentialsProvider(tenant_id=tenant_id, provider_manager=provider_manager),
+        DifyModelFactory(
+            tenant_id=tenant_id,
+            model_manager=ModelManager(enable_credentials_cache=True, provider_manager=provider_manager),
+        ),
     )
 
 
diff --git a/api/core/model_manager.py b/api/core/model_manager.py
index 2eaf24aede..abcc3dea41 100644
--- a/api/core/model_manager.py
+++ b/api/core/model_manager.py
@@ -480,8 +480,8 @@ class ModelInstance:
 
 
 class ModelManager:
-    def __init__(self, enable_credentials_cache: bool = False):
-        self._provider_manager = ProviderManager()
+    def __init__(self, enable_credentials_cache: bool = False, provider_manager: ProviderManager | None = None):
+        self._provider_manager = provider_manager or ProviderManager()
         self._credentials_cache: dict[tuple[str, str, str, str], Any] = {}
         self._enable_credentials_cache = enable_credentials_cache