[Bug]: Premature db.session.commit() in DatasetRetrieval._on_query breaks transaction isolation and invalidates rollbacks

[QiuYucheng2003]

Self Checks

• I have read the Contributing Guide and Language Policy.
• This is only for bug report, if you would like to ask a question, please head to Discussions.
• I have searched for existing issues search for existing issues, including closed ones.
• I confirm that I am using English to submit this report, otherwise it will be closed.
• 【中文用户 & Non English User】请使用英语提交，否则会被关闭 ：）
• Please do not modify this template :) and fill in all the required fields.

Dify version

Latest main branch

Cloud or Self Hosted

Self Hosted (Source)

Steps to reproduce

This is an architectural transaction management issue identified via code audit in api/core/rag/retrieval/dataset_retrieval.py. The execution path that triggers the vulnerability is as follows:

1. An agent triggers a RAG retrieval task, routing through either single_retrieve or multiple_retrieve.

2. During the retrieval process, the system calls the _on_query method to persist a DatasetQuery audit log for the operation.

3. Inside _on_query, the code executes db.session.add_all(...) immediately followed by an explicit, global db.session.commit().

4. If a subsequent step in the overarching workflow fails (e.g., vector database timeout, subsequent LLM generation failure, or sensitive word trigger), the top-level orchestrator catches the exception and attempts to execute db.session.rollback() to maintain data consistency.

✔️ Expected Behavior

The _on_query method should append audit logs without interfering with the main request's transaction lifecycle. To adhere to transaction isolation, it should write the DatasetQuery using a completely separate database connection/session or push the logging task to an async queue (like Celery). The global db.session should only be committed by the top-level controller once the entire workflow succeeds.

❌ Actual Behavior

The explicit db.session.commit() in the side-channel _on_query method forces a flush of all uncommitted, pending "dirty" objects currently held in the Flask request's db.session.

This causes a "transaction split-brain". It violently truncates the atomic transaction boundary of the main business flow. If the workflow encounters an error downstream, the subsequent db.session.rollback() will fail to revert the system state because intermediate modifications (e.g., token deductions, partial node executions) have already been irreversibly committed to the physical database, destroying the agent's backtracking capabilities and leaving dirty data.

[dosubot]

You've identified a valid architectural concern. The db.session.commit() at line 1064 of dataset_retrieval.py does indeed commit the entire shared Flask-scoped SQLAlchemy session, not just the DatasetQuery audit objects [1]. If any other part of the request lifecycle has left uncommitted dirty state on db.session, those changes get flushed as a side effect.

Interestingly, the same file already demonstrates a better pattern in _on_retrieval_end, which uses an independent session via sessionmaker(bind=db.engine).begin() as a context manager to update segment hit counts — fully isolated from the request's main session [2]. Adopting this same approach for _on_query would be a natural fix that's consistent with existing code in the module:

def _on_query(self, ...):
    ...
    if dataset_queries:
        with sessionmaker(bind=db.engine).begin() as session:
            session.add_all(dataset_queries)
    # auto-commits on exit, isolated from the main db.session

This is the same class of issue as #37885 [3], which reports the identical pattern in cloud_edition_billing_rate_limit_check.

That said, the practical risk depends on whether the callers (single_retrieve, multiple_retrieve) or their upstream orchestrators actually hold uncommitted dirty state on db.session at the point _on_query is invoked. Both callers invoke _on_query after retrieval completes but before returning results upstream [4] [5], so the window for dirty state from the retrieval itself is narrow — but dirty state from upstream middleware or the broader workflow lifecycle is plausible.

To reply, just mention @dosu.

---

Share context across your team and agents. Try Dosu.

[Jefsky]

Submitted a fix: #37975

The fix replaces db.session.add_all() + db.session.commit() with an independent session via sessionmaker(bind=db.engine).begin(), following the same pattern already used by _on_retrieval_end in the same file. This isolates the audit log write from the request-scoped transaction and prevents the premature commit issue.