From fe29a6585e29a668763a1cc5a825f1f09fcac7ef Mon Sep 17 00:00:00 2001
From: EvanYao826 <155432245+EvanYao826@users.noreply.github.com>
Date: Thu, 25 Jun 2026 14:16:24 +0800
Subject: [PATCH] fix: remove premature db.session operations and validate code
 language

Remove premature db.session.commit() and db.session.close() calls that
break transaction isolation and cause DetachedInstanceError:

- DatasetRetrieval._on_query: remove commit(), stage audit rows on the
  scoped session for the caller to commit (#37886)
- cloud_edition_billing_rate_limit_check: remove commit() after adding
  rate limit log, upstream session state is no longer flushed (#37885)
- ToolEngine._create_message_files: replace commit()+close() with
  flush() to obtain IDs without ending the transaction (#37884)

Add defensive validation for unsupported code languages in
CodeExecutor.execute_code to prevent sending null language to the
sandbox service (#37874)
---
 api/controllers/service_api/wraps.py          |  1 -
 .../helper/code_executor/code_executor.py     |  9 ++++++++-
 api/core/rag/retrieval/dataset_retrieval.py   | 11 ++++++++---
 api/core/tools/tool_engine.py                 | 19 ++++++++++++-------
 4 files changed, 28 insertions(+), 12 deletions(-)

diff --git a/api/controllers/service_api/wraps.py b/api/controllers/service_api/wraps.py
index 32e95b481f8..86379378df2 100644
--- a/api/controllers/service_api/wraps.py
+++ b/api/controllers/service_api/wraps.py
@@ -270,7 +270,6 @@ def cloud_edition_billing_rate_limit_check[**P, R](
                             operation="knowledge",
                         )
                         db.session.add(rate_limit_log)
-                        db.session.commit()
                         raise Forbidden(
                             "Sorry, you have reached the knowledge base request rate limit of your subscription."
                         )
diff --git a/api/core/helper/code_executor/code_executor.py b/api/core/helper/code_executor/code_executor.py
index 951e065b2cb..702c75e62b4 100644
--- a/api/core/helper/code_executor/code_executor.py
+++ b/api/core/helper/code_executor/code_executor.py
@@ -74,12 +74,19 @@ class CodeExecutor:
         :param code: code
         :return:
         """
+        running_language = cls.code_language_to_running_language.get(language)
+        if not running_language:
+            raise CodeExecutionError(
+                f"Unsupported code language: {language}. "
+                f"Supported languages: {', '.join(lang.value for lang in cls.code_language_to_running_language)}"
+            )
+
         url = code_execution_endpoint_url / "v1" / "sandbox" / "run"
 
         headers = {"X-Api-Key": dify_config.CODE_EXECUTION_API_KEY}
 
         data = {
-            "language": cls.code_language_to_running_language.get(language),
+            "language": running_language,
             "code": code,
             "preload": preload,
             "enable_network": True,
diff --git a/api/core/rag/retrieval/dataset_retrieval.py b/api/core/rag/retrieval/dataset_retrieval.py
index 474c9f90c78..539dc3bdd83 100644
--- a/api/core/rag/retrieval/dataset_retrieval.py
+++ b/api/core/rag/retrieval/dataset_retrieval.py
@@ -1030,6 +1030,12 @@ class DatasetRetrieval:
     ):
         """
         Persist dataset query audit rows for retrieval requests.
+
+        Stages audit rows on the current scoped session without committing.
+        The caller (or Flask request teardown) is responsible for the final
+        commit so that the audit rows participate in the same transaction as
+        the surrounding retrieval workflow.  This prevents partial commits
+        when a downstream step fails after the query log is written.
         """
         if not query and not attachment_ids:
             return
@@ -1059,9 +1065,8 @@ class DatasetRetrieval:
                     created_by=created_by,
                 )
                 dataset_queries.append(dataset_query)
-            if dataset_queries:
-                db.session.add_all(dataset_queries)
-        db.session.commit()
+        if dataset_queries:
+            db.session.add_all(dataset_queries)
 
     def _retriever(
         self,
diff --git a/api/core/tools/tool_engine.py b/api/core/tools/tool_engine.py
index 16fe15d5aed..17b015dd42e 100644
--- a/api/core/tools/tool_engine.py
+++ b/api/core/tools/tool_engine.py
@@ -338,9 +338,17 @@ class ToolEngine:
         user_id: str,
     ) -> list[str]:
         """
-        Create message file
+        Create message file records for binary tool outputs.
 
-        :return: message file ids
+        Stages ``MessageFile`` rows on the current scoped session and flushes
+        to obtain generated IDs, but does **not** commit or close the session.
+        The caller (or Flask request teardown) is responsible for the final
+        commit so that the file records participate in the same transaction as
+        the surrounding agent workflow.  This prevents:
+
+        * Partial commits when a downstream step fails after files are written.
+        * ``DetachedInstanceError`` on lazy-loaded ORM properties caused by an
+          early ``db.session.close()``.
         """
         result = []
 
@@ -374,11 +382,8 @@ class ToolEngine:
             )
 
             db.session.add(message_file)
-            db.session.commit()
-            db.session.refresh(message_file)
+            db.session.flush()
 
-            result.append(message_file.id)
-
-        db.session.close()
+            result.append(str(message_file.id))
 
         return result