fix(persistence): address review feedback on PR #1851

- Fix naive datetime.now() → datetime.now(UTC) in all ORM models
- Fix seq race condition in DbRunEventStore.put() with FOR UPDATE
  and UNIQUE(thread_id, seq) constraint
- Encapsulate _store access in RunManager.update_run_completion()
- Deduplicate _store.put() logic in RunManager via _persist_to_store()
- Add update_run_completion to RunStore ABC + MemoryRunStore
- Wire follow_up_to_run_id through the full create path
- Add error recovery to RunJournal._flush_sync() lost-event scenario
- Add migration note for search_threads breaking change
- Fix test_checkpointer_none_fix mock to set database=None

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

backend/packages/harness/deerflow/runtime/runs/manager.py

@@ -54,6 +54,33 @@ class RunManager:
         self._lock = asyncio.Lock()
         self._store = store
 
+    async def _persist_to_store(self, record: RunRecord, *, follow_up_to_run_id: str | None = None) -> None:
+        """Best-effort persist run record to backing store."""
+        if self._store is None:
+            return
+        try:
+            await self._store.put(
+                record.run_id,
+                thread_id=record.thread_id,
+                assistant_id=record.assistant_id,
+                status=record.status.value,
+                multitask_strategy=record.multitask_strategy,
+                metadata=record.metadata or {},
+                kwargs=record.kwargs or {},
+                created_at=record.created_at,
+                follow_up_to_run_id=follow_up_to_run_id,
+            )
+        except Exception:
+            logger.warning("Failed to persist run %s to store", record.run_id, exc_info=True)
+
+    async def update_run_completion(self, run_id: str, **kwargs) -> None:
+        """Persist token usage and completion data to the backing store."""
+        if self._store is not None:
+            try:
+                await self._store.update_run_completion(run_id, **kwargs)
+            except Exception:
+                logger.warning("Failed to persist run completion for %s", run_id, exc_info=True)
+
     async def create(
         self,
         thread_id: str,
@@ -63,6 +90,7 @@ class RunManager:
         metadata: dict | None = None,
         kwargs: dict | None = None,
         multitask_strategy: str = "reject",
+        follow_up_to_run_id: str | None = None,
     ) -> RunRecord:
         """Create a new pending run and register it."""
         run_id = str(uuid.uuid4())
@@ -81,20 +109,7 @@ class RunManager:
         )
         async with self._lock:
             self._runs[run_id] = record
-        if self._store is not None:
-            try:
-                await self._store.put(
-                    run_id,
-                    thread_id=thread_id,
-                    assistant_id=assistant_id,
-                    status=RunStatus.pending.value,
-                    multitask_strategy=multitask_strategy,
-                    metadata=metadata or {},
-                    kwargs=kwargs or {},
-                    created_at=now,
-                )
-            except Exception:
-                logger.warning("Failed to persist run %s to store", run_id, exc_info=True)
+        await self._persist_to_store(record, follow_up_to_run_id=follow_up_to_run_id)
         logger.info("Run created: run_id=%s thread_id=%s", run_id, thread_id)
         return record
 
@@ -161,6 +176,7 @@ class RunManager:
         metadata: dict | None = None,
         kwargs: dict | None = None,
         multitask_strategy: str = "reject",
+        follow_up_to_run_id: str | None = None,
     ) -> RunRecord:
         """Atomically check for inflight runs and create a new one.
 
@@ -214,21 +230,7 @@ class RunManager:
             )
             self._runs[run_id] = record
 
-        if self._store is not None:
-            try:
-                await self._store.put(
-                    run_id,
-                    thread_id=thread_id,
-                    assistant_id=assistant_id,
-                    status=RunStatus.pending.value,
-                    multitask_strategy=multitask_strategy,
-                    metadata=metadata or {},
-                    kwargs=kwargs or {},
-                    created_at=now,
-                )
-            except Exception:
-                logger.warning("Failed to persist run %s to store", run_id, exc_info=True)
-
+        await self._persist_to_store(record, follow_up_to_run_id=follow_up_to_run_id)
         logger.info("Run created: run_id=%s thread_id=%s", run_id, thread_id)
         return record
 

backend/packages/harness/deerflow/runtime/runs/store/base.py

@@ -29,6 +29,7 @@ class RunStore(abc.ABC):
         kwargs: dict[str, Any] | None = None,
         error: str | None = None,
         created_at: str | None = None,
+        follow_up_to_run_id: str | None = None,
     ) -> None: ...
 
     @abc.abstractmethod
@@ -55,5 +56,24 @@ class RunStore(abc.ABC):
     @abc.abstractmethod
     async def delete(self, run_id: str) -> None: ...
 
+    @abc.abstractmethod
+    async def update_run_completion(
+        self,
+        run_id: str,
+        *,
+        status: str,
+        total_input_tokens: int = 0,
+        total_output_tokens: int = 0,
+        total_tokens: int = 0,
+        llm_call_count: int = 0,
+        lead_agent_tokens: int = 0,
+        subagent_tokens: int = 0,
+        middleware_tokens: int = 0,
+        message_count: int = 0,
+        last_ai_message: str | None = None,
+        first_human_message: str | None = None,
+        error: str | None = None,
+    ) -> None: ...
+
     @abc.abstractmethod
     async def list_pending(self, *, before: str | None = None) -> list[dict[str, Any]]: ...

backend/packages/harness/deerflow/runtime/runs/store/memory.py

@@ -28,6 +28,7 @@ class MemoryRunStore(RunStore):
         kwargs=None,
         error=None,
         created_at=None,
+        follow_up_to_run_id=None,
     ):
         now = datetime.now(UTC).isoformat()
         self._runs[run_id] = {
@@ -40,6 +41,7 @@ class MemoryRunStore(RunStore):
             "metadata": metadata or {},
             "kwargs": kwargs or {},
             "error": error,
+            "follow_up_to_run_id": follow_up_to_run_id,
             "created_at": created_at or now,
             "updated_at": now,
         }
@@ -62,6 +64,14 @@ class MemoryRunStore(RunStore):
     async def delete(self, run_id):
         self._runs.pop(run_id, None)
 
+    async def update_run_completion(self, run_id, *, status, **kwargs):
+        if run_id in self._runs:
+            self._runs[run_id]["status"] = status
+            for key, value in kwargs.items():
+                if value is not None:
+                    self._runs[run_id][key] = value
+            self._runs[run_id]["updated_at"] = datetime.now(UTC).isoformat()
+
     async def list_pending(self, *, before=None):
         now = before or datetime.now(UTC).isoformat()
         results = [r for r in self._runs.values() if r["status"] == "pending" and r["created_at"] <= now]

backend/packages/harness/deerflow/runtime/runs/worker.py

@@ -257,16 +257,8 @@ async def run_agent(
                 logger.warning("Failed to flush journal for run %s", run_id, exc_info=True)
 
             # Persist token usage + convenience fields to RunStore
-            if run_manager._store is not None:
-                try:
-                    completion = journal.get_completion_data()
-                    await run_manager._store.update_run_completion(
-                        run_id,
-                        status=record.status.value,
-                        **completion,
-                    )
-                except Exception:
-                    logger.warning("Failed to persist run completion for %s", run_id, exc_info=True)
+            completion = journal.get_completion_data()
+            await run_manager.update_run_completion(run_id, status=record.status.value, **completion)
 
         # Sync title from checkpoint to threads_meta.display_name
         if thread_meta_repo is not None and checkpointer is not None: