From 6cf7adfa21b52efb05cef568cc25fef0c118f8c3 Mon Sep 17 00:00:00 2001
From: ViditOstwal <viditostwal@gmail.com>
Date: Fri, 18 Sep 2026 13:51:47 +0530
Subject: [PATCH] fix(flows): persist MongoDB feedback atomically

---
 .../src/crewai/flow/persistence/mongodb.py    | 72 ++++++++++++-------
 .../tests/test_flow_persistence_mongodb.py    |  9 ++-
 2 files changed, 52 insertions(+), 29 deletions(-)

diff --git a/lib/crewai/src/crewai/flow/persistence/mongodb.py b/lib/crewai/src/crewai/flow/persistence/mongodb.py
index 6bba66974..e1a5e03bd 100644
--- a/lib/crewai/src/crewai/flow/persistence/mongodb.py
+++ b/lib/crewai/src/crewai/flow/persistence/mongodb.py
@@ -24,6 +24,7 @@ cases.
 
 from __future__ import annotations
 
+from collections.abc import Callable
 from datetime import datetime, timezone
 import json
 import os
@@ -173,6 +174,33 @@ class MongoDbFlowPersistence(FlowPersistence):
         )
         return int(doc["seq"])
 
+    def _insert_state(
+        self,
+        flow_uuid: str,
+        method_name: str,
+        state_dict: dict[str, Any],
+        session: Any,
+    ) -> None:
+        """Insert a state snapshot using the transaction's sequence."""
+        self._db_ready()[self.states_collection].insert_one(
+            {
+                "flow_uuid": flow_uuid,
+                "method_name": method_name,
+                "timestamp": datetime.now(timezone.utc).isoformat(),
+                "state_json": json.dumps(state_dict, default=_json_default),
+                "seq": self._next_sequence(self.states_collection, session),
+            },
+            session=session,
+        )
+
+    def _run_transaction(self, operation: Callable[[Any], None]) -> None:
+        """Run a persistence write operation in a retrying Mongo transaction."""
+        self._db_ready()
+        if self._client is None:
+            raise RuntimeError("MongoDB client was not initialized.")
+        with self._client.start_session() as session:
+            session.with_transaction(operation)
+
     def save_state(
         self,
         flow_uuid: str,
@@ -186,24 +214,11 @@ class MongoDbFlowPersistence(FlowPersistence):
         ordering on ``seq`` rather than the client-generated ObjectId ``_id``.
         """
         state_dict = self._to_state_dict(state_data)
-        db = self._db_ready()
 
         def write_state(session: Any) -> None:
-            db[self.states_collection].insert_one(
-                {
-                    "flow_uuid": flow_uuid,
-                    "method_name": method_name,
-                    "timestamp": datetime.now(timezone.utc).isoformat(),
-                    "state_json": json.dumps(state_dict, default=_json_default),
-                    "seq": self._next_sequence(self.states_collection, session),
-                },
-                session=session,
-            )
+            self._insert_state(flow_uuid, method_name, state_dict, session)
 
-        if self._client is None:
-            raise RuntimeError("MongoDB client was not initialized.")
-        with self._client.start_session() as session:
-            session.with_transaction(write_state)
+        self._run_transaction(write_state)
 
     def load_state(self, flow_uuid: str) -> dict[str, Any] | None:
         """Load the most recent state for a given flow UUID."""
@@ -223,20 +238,23 @@ class MongoDbFlowPersistence(FlowPersistence):
     ) -> None:
         """Save state with a pending feedback marker (upsert per flow)."""
         state_dict = self._to_state_dict(state_data)
+        context_json = json.dumps(context.to_dict(), default=_json_default)
 
-        # Mirror SQLite: record the state snapshot, then upsert the pending row.
-        self.save_state(flow_uuid, context.method_name, state_data)
+        def write_pending_feedback(session: Any) -> None:
+            self._insert_state(flow_uuid, context.method_name, state_dict, session)
+            self._db_ready()[self.pending_collection].replace_one(
+                {"flow_uuid": flow_uuid},
+                {
+                    "flow_uuid": flow_uuid,
+                    "context_json": context_json,
+                    "state_json": json.dumps(state_dict, default=_json_default),
+                    "created_at": datetime.now(timezone.utc).isoformat(),
+                },
+                upsert=True,
+                session=session,
+            )
 
-        self._db_ready()[self.pending_collection].replace_one(
-            {"flow_uuid": flow_uuid},
-            {
-                "flow_uuid": flow_uuid,
-                "context_json": json.dumps(context.to_dict(), default=_json_default),
-                "state_json": json.dumps(state_dict, default=_json_default),
-                "created_at": datetime.now(timezone.utc).isoformat(),
-            },
-            upsert=True,
-        )
+        self._run_transaction(write_pending_feedback)
 
     def load_pending_feedback(
         self,
diff --git a/lib/crewai/tests/test_flow_persistence_mongodb.py b/lib/crewai/tests/test_flow_persistence_mongodb.py
index 3aeaacfb7..b01c88806 100644
--- a/lib/crewai/tests/test_flow_persistence_mongodb.py
+++ b/lib/crewai/tests/test_flow_persistence_mongodb.py
@@ -68,7 +68,11 @@ class _FakeCollection:
         return dict(row)
 
     def replace_one(
-        self, flt: dict[str, Any], doc: dict[str, Any], upsert: bool = False
+        self,
+        flt: dict[str, Any],
+        doc: dict[str, Any],
+        upsert: bool = False,
+        session: Any = None,
     ) -> None:
         for i, existing in enumerate(self.docs):
             if self._match(existing, flt):
@@ -267,7 +271,7 @@ def test_env_resolved_at_call_time(monkeypatch: pytest.MonkeyPatch) -> None:
 
 
 def test_pending_feedback_roundtrip(monkeypatch: pytest.MonkeyPatch) -> None:
-    _patch_client(monkeypatch)
+    created = _patch_client(monkeypatch)
     persistence = MongoDbFlowPersistence(CONN)
     context = PendingFeedbackContext(
         flow_id="flow-1",
@@ -288,6 +292,7 @@ def test_pending_feedback_roundtrip(monkeypatch: pytest.MonkeyPatch) -> None:
 
     persistence.clear_pending_feedback("flow-1")
     assert persistence.load_pending_feedback("flow-1") is None
+    assert created["client"].transactions_started == 1
 
 
 def test_missing_pymongo_raises_helpful_error(