diff --git a/api/services/hit_testing_service.py b/api/services/hit_testing_service.py index 2f9d5c3677b..71bbf621cae 100644 --- a/api/services/hit_testing_service.py +++ b/api/services/hit_testing_service.py @@ -183,7 +183,10 @@ class HitTestingService: created_by=account.id, ) session.add(dataset_query) - session.commit() + # Stage the query log without committing: the caller owns the session + # (e.g. `with_session`), and committing here would close a caller-managed + # transaction before the document lookup below runs. See #38998. + session.flush() return cls.compact_retrieve_response(query, all_documents, session=session) diff --git a/api/tests/unit_tests/services/test_hit_testing_service.py b/api/tests/unit_tests/services/test_hit_testing_service.py index c251ac5147d..68edd7486a9 100644 --- a/api/tests/unit_tests/services/test_hit_testing_service.py +++ b/api/tests/unit_tests/services/test_hit_testing_service.py @@ -1,7 +1,14 @@ """Unit tests for the database-independent hit-testing helpers.""" -import pytest +from unittest.mock import patch +import pytest +from sqlalchemy import select +from sqlalchemy.orm import Session + +from core.rag.models.document import Document +from models import Account +from models.dataset import Dataset, DatasetQuery from services.hit_testing_service import HitTestingService @@ -30,3 +37,64 @@ class TestHitTestingServiceArguments: def test_hit_testing_args_check_should_raise_error_when_attachments_not_list(self) -> None: with pytest.raises(ValueError, match="Attachment_ids must be a list"): HitTestingService.hit_testing_args_check({"attachment_ids": "not a list"}) + + +class _FakeSegment: + """Minimal stand-in for a retrieval segment carrying a document id.""" + + def __init__(self, document_id: str) -> None: + self.document_id = document_id + + +class _FakeRetrievalRecord: + """Minimal stand-in for a RetrievalSegments record. + + Implements just the two members ``_dump_retrieval_records`` reads: + ``model_dump()`` and ``segment.document_id``. + """ + + def __init__(self, document_id: str, content: str, score: float) -> None: + self.segment = _FakeSegment(document_id) + self._content = content + self._score = score + + def model_dump(self) -> dict[str, str | float]: + return {"content": self._content, "score": self._score} + + +class TestHitTestingServiceRetrieveTransaction: + """Regression tests for https://github.com/langgenius/dify/issues/38998.""" + + def test_retrieve_should_not_close_caller_managed_transaction(self, sqlite_session: Session) -> None: + """retrieve() must not commit the caller-managed session mid-flow. + + When the caller manages the session inside a ``session.begin()`` block, an + explicit ``session.commit()`` inside retrieve() closes that transaction and + the follow-up document lookup then fails with "Can't operate on closed + transaction inside context manager". The query log must be staged with + ``flush()`` so the caller commits it. + """ + dataset = Dataset(id="dataset-123", tenant_id="tenant-123", provider="vendor", retrieval_model=None) + account = Account(name="Test User", email="test@example.com") + account.id = "user-789" + query = "test query" + + record = _FakeRetrievalRecord(document_id="doc-1", content="Doc 1", score=0.95) + + with ( + patch("services.hit_testing_service.RetrievalService.retrieve", autospec=True) as mock_retrieve, + patch( + "services.hit_testing_service.RetrievalService.format_retrieval_documents", autospec=True + ) as mock_format, + ): + mock_retrieve.return_value = [Document(page_content="Doc 1", metadata={})] + mock_format.return_value = [record] + + with sqlite_session.begin(): + result = HitTestingService.retrieve(dataset, query, account, None, {}, session=sqlite_session) + + assert result["query"]["content"] == query + query_log = sqlite_session.scalar(select(DatasetQuery)) + assert query_log is not None + assert query_log.dataset_id == dataset.id + assert query_log.created_by == account.id