Python: Keep sessions pending when a consolidation pass did not cover them

python/packages/core/agent_framework/_harness/_memory.py

@@ -7,6 +7,7 @@
 import logging
 import re
 import threading
+import uuid
 import weakref
 from abc import ABC, abstractmethod
 from base64 import b32decode
@@ -222,7 +223,25 @@ def _coerce_summary(summary: str, memories: Sequence[str]) -> str:
 
 
 def _default_state() -> dict[str, Any]:
-    return {"last_consolidated_at": None, "sessions_since_consolidation": []}
+    # ``sessions_since_consolidation`` lists the distinct sessions with turns not yet covered by a
+    # consolidation pass. ``pending_turn_ids`` maps each of them to a unique id of its latest finished
+    # turn, so a pass only retires a session whose latest turn finished before the pass took its snapshot.
+    return {"last_consolidated_at": None, "sessions_since_consolidation": [], "pending_turn_ids": {}}
+
+
+def _normalize_state(state: Mapping[str, Any]) -> dict[str, Any]:
+    """Return ``state`` with every maintenance field present and of the expected type.
+
+    Stores may return state saved before a field existed, such as state without ``pending_turn_ids``.
+    """
+    normalized = {**_default_state(), **state}
+    if not isinstance(normalized.get("sessions_since_consolidation"), list):
+        normalized["sessions_since_consolidation"] = []
+    if not isinstance(normalized.get("pending_turn_ids"), dict):
+        normalized["pending_turn_ids"] = {}
+    if not isinstance(normalized.get("last_consolidated_at"), (str, type(None))):
+        normalized["last_consolidated_at"] = None
+    return normalized
 
 
 def _format_search_results(results: Sequence[Mapping[str, Any]]) -> str:
@@ -893,12 +912,7 @@ def read_state(self, session: AgentSession, *, source_id: str) -> dict[str, Any]
         raw_state = self.loads(state_path.read_text(encoding="utf-8"))
         if not isinstance(raw_state, dict):
             raise ValueError("Memory state file must contain a JSON object.")
-        state = {**_default_state(), **cast(dict[str, Any], raw_state)}
-        if not isinstance(state.get("sessions_since_consolidation"), list):
-            state["sessions_since_consolidation"] = []
-        if not isinstance(state.get("last_consolidated_at"), (str, type(None))):
-            state["last_consolidated_at"] = None
-        return state
+        return _normalize_state(cast(dict[str, Any], raw_state))
 
     def write_state(self, session: AgentSession, state: Mapping[str, Any], *, source_id: str) -> None:
         """Persist the maintenance state for the current owner."""
@@ -1084,6 +1098,10 @@ def _topic_async_lock(self, session: AgentSession, topic: str) -> asyncio.Lock:
     def _state_async_lock(self, session: AgentSession) -> asyncio.Lock:
         return self._store_async_lock(session, kind="state", key="maintenance")
 
+    def _read_maintenance_state(self, session: AgentSession) -> dict[str, Any]:
+        # Custom stores may return state saved before newer fields existed, so fill them in here.
+        return _normalize_state(self.store.read_state(session, source_id=self.source_id))
+
     def _create_history_provider(self, session: AgentSession) -> FileHistoryProvider:
         return FileHistoryProvider(
             self.store.get_transcripts_directory(session, source_id=self.source_id),
@@ -1401,27 +1419,34 @@ async def after_run(
             await self.save_messages(session.session_id, messages_to_store, state=state)
 
         current_time = datetime.now(timezone.utc).replace(microsecond=0)
-        async with self._state_async_lock(session):
-            maintenance_state = self.store.read_state(session, source_id=self.source_id)
-            session_ids_since_consolidation = cast(list[str], maintenance_state["sessions_since_consolidation"])
-            if session.session_id not in session_ids_since_consolidation:
-                session_ids_since_consolidation.append(session.session_id)
-            maintenance_state["sessions_since_consolidation"] = session_ids_since_consolidation
-            self.store.write_state(session, maintenance_state, source_id=self.source_id)
-
-        extracted_topics = await self._extract_memories(
-            client=self._chat_client(agent),
-            session=session,
-            context=context,
-            now=current_time,
-        )
-        if extracted_topics:
-            self.store.rebuild_index(
-                session,
-                source_id=self.source_id,
-                line_limit=self.index_line_limit,
-                line_length=self.index_line_length,
+        try:
+            extracted_topics = await self._extract_memories(
+                client=self._chat_client(agent),
+                session=session,
+                context=context,
+                now=current_time,
             )
+            if extracted_topics:
+                self.store.rebuild_index(
+                    session,
+                    source_id=self.source_id,
+                    line_limit=self.index_line_limit,
+                    line_length=self.index_line_length,
+                )
+        finally:
+            # Register the session only once its extraction is over, so a consolidation pass that sees
+            # this session as pending has also seen the memories it extracted. Registering in `finally`
+            # keeps topics already written by a run that then failed or was cancelled pending too.
+            # A fresh turn id is recorded even when the session is already pending: runs on the same
+            # session can overlap, and a pass that snapshotted an earlier turn must not retire this one.
+            async with self._state_async_lock(session):
+                maintenance_state = self._read_maintenance_state(session)
+                session_ids_since_consolidation = cast(list[str], maintenance_state["sessions_since_consolidation"])
+                if session.session_id not in session_ids_since_consolidation:
+                    session_ids_since_consolidation.append(session.session_id)
+                maintenance_state["sessions_since_consolidation"] = session_ids_since_consolidation
+                cast(dict[str, str], maintenance_state["pending_turn_ids"])[session.session_id] = uuid.uuid4().hex
+                self.store.write_state(session, maintenance_state, source_id=self.source_id)
 
         await self._run_consolidation(
             client=self._chat_client(agent, override=self.consolidation_client),
@@ -1549,15 +1574,23 @@ async def _run_consolidation(
         # making LLM calls so concurrent before/after_run invocations don't have to wait through
         # multi-second consolidation calls.
         async with self._state_async_lock(session):
-            maintenance_state = self.store.read_state(session, source_id=self.source_id)
+            maintenance_state = self._read_maintenance_state(session)
             if not force and not self._should_consolidate(maintenance_state, now=now):
                 return 0
             topic_records = self.store.list_topics(session, source_id=self.source_id)
             if not topic_records:
                 maintenance_state["last_consolidated_at"] = _timestamp(now)
                 maintenance_state["sessions_since_consolidation"] = []
+                maintenance_state["pending_turn_ids"] = {}
                 self.store.write_state(session, maintenance_state, source_id=self.source_id)
                 return 0
+            # Turns registered after this snapshot are not covered by this pass and stay pending, even
+            # when their session was already pending at snapshot time.
+            pending_turn_ids = cast(dict[str, str], maintenance_state["pending_turn_ids"])
+            covered_turn_ids = {
+                session_id: pending_turn_ids.get(session_id)
+                for session_id in cast(list[str], maintenance_state["sessions_since_consolidation"])
+            }
 
         success_count = 0
         for record in topic_records:
@@ -1578,9 +1611,21 @@ async def _run_consolidation(
         # session IDs so the next after_run will retry instead of silently sliding the window.
         if success_count > 0:
             async with self._state_async_lock(session):
-                maintenance_state = self.store.read_state(session, source_id=self.source_id)
+                maintenance_state = self._read_maintenance_state(session)
                 maintenance_state["last_consolidated_at"] = _timestamp(now)
-                maintenance_state["sessions_since_consolidation"] = []
+                pending_turn_ids = cast(dict[str, str], maintenance_state["pending_turn_ids"])
+                remaining_session_ids = [
+                    session_id
+                    for session_id in cast(list[str], maintenance_state["sessions_since_consolidation"])
+                    if session_id not in covered_turn_ids
+                    or pending_turn_ids.get(session_id) != covered_turn_ids[session_id]
+                ]
+                maintenance_state["sessions_since_consolidation"] = remaining_session_ids
+                maintenance_state["pending_turn_ids"] = {
+                    session_id: pending_turn_ids[session_id]
+                    for session_id in remaining_session_ids
+                    if session_id in pending_turn_ids
+                }
                 self.store.write_state(session, maintenance_state, source_id=self.source_id)
                 self.store.rebuild_index(
                     session,