diff --git a/plugins/memory/hindsight/__init__.py b/plugins/memory/hindsight/__init__.py index 2f94c08da383..a1b5465daac9 100644 --- a/plugins/memory/hindsight/__init__.py +++ b/plugins/memory/hindsight/__init__.py @@ -1583,6 +1583,86 @@ def handle_tool_call(self, tool_name: str, args: dict, **kwargs) -> str: return tool_error(f"Unknown tool: {tool_name}") + # -- Session lifecycle hooks --------------------------------------------- + + def on_session_end(self, messages: List[Dict[str, Any]]) -> None: + """Flush any buffered turns before the session is torn down. + + Called at actual session boundaries — CLI exit, /reset, gateway + session expiry — via ``MemoryManager.on_session_end``, which runs + *before* ``shutdown_all()`` / ``shutdown()``. This is the right + place to flush because the writer thread is still alive and can + process the enqueued retain. + + Without this, users who set ``retain_every_n_turns > 1`` silently + lose whatever turns are buffered when a session ends without + hitting the modulo boundary. ``on_session_switch`` already + handles mid-process session rotations (/new, /resume, compression); + this method covers the final boundary where no new session follows. + """ + if not self._auto_retain: + return + if not self._session_turns: + return + if self._shutting_down.is_set(): + return + + old_turns = list(self._session_turns) + old_content = "[" + ",".join(old_turns) + "]" + old_metadata = self._build_metadata( + message_count=len(old_turns) * 2, + turn_index=self._turn_index, + ) + old_lineage_tags: list[str] = [] + if self._session_id: + old_lineage_tags.append(f"session:{self._session_id}") + if self._parent_session_id: + old_lineage_tags.append(f"parent:{self._parent_session_id}") + document_id, update_mode = self._resolve_retain_target(self._document_id) + bank_id = self._bank_id + retain_async_flag = self._retain_async + retain_context = self._retain_context + + def _flush_on_end() -> None: + try: + item = self._build_retain_kwargs( + old_content, + context=retain_context, + metadata=old_metadata, + tags=old_lineage_tags or None, + ) + item.pop("bank_id", None) + item.pop("retain_async", None) + if update_mode is not None: + item["update_mode"] = update_mode + logger.debug( + "Hindsight flush-on-session-end: bank=%s, doc=%s, mode=%s, num_turns=%d", + bank_id, document_id, update_mode, len(old_turns), + ) + self._run_hindsight_operation( + lambda client: client.aretain_batch( + bank_id=bank_id, + items=[item], + document_id=document_id, + retain_async=retain_async_flag, + ) + ) + except Exception as e: + logger.warning("Hindsight flush-on-session-end failed: %s", e, exc_info=True) + + self._ensure_writer() + self._register_atexit() + self._retain_queue.put(_flush_on_end) + + # Clear the buffer so a subsequent shutdown() doesn't re-flush. + self._session_turns = [] + self._turn_counter = 0 + self._turn_index = 0 + logger.debug( + "Hindsight on_session_end: flushed %d buffered turns for session %s", + len(old_turns), self._session_id, + ) + def on_session_switch( self, new_session_id: str, diff --git a/tests/agent/test_memory_session_switch.py b/tests/agent/test_memory_session_switch.py index a40654fa5793..707f890e64be 100644 --- a/tests/agent/test_memory_session_switch.py +++ b/tests/agent/test_memory_session_switch.py @@ -1,10 +1,14 @@ -"""Tests for the on_session_switch hook and session_id propagation. +"""Tests for the on_session_switch hook, on_session_end hook, and session_id propagation. Covers #6672: memory providers must be notified when AIAgent.session_id rotates mid-process (via /resume, /branch, /reset, /new, or context compression). Without the notification, providers that cache per-session state in initialize() (Hindsight, and any plugin that stores session_id for scoped writes) keep writing into the old session's record. + +Also covers: Hindsight on_session_end must flush any buffered turns when +a session is torn down (CLI exit, gateway idle expiry), preventing silent +data loss for users with retain_every_n_turns > 1. """ @@ -325,3 +329,82 @@ def test_hindsight_preserves_parent_across_empty_parent_arg(): provider._parent_session_id = "original-parent" provider.on_session_switch("new-sid") # no parent passed assert provider._parent_session_id == "original-parent" + + +# --------------------------------------------------------------------------- +# Hindsight on_session_end — flush buffered turns on session teardown +# --------------------------------------------------------------------------- + + +def test_hindsight_on_session_end_flushes_buffered_turns(): + """Buffered turns must be flushed via the writer queue on session end. + + Users with retain_every_n_turns > 1 accumulate turns in _session_turns + until the modulo boundary is hit. If the session ends before that + boundary (e.g. 5 turns with retain_every_n_turns=12), on_session_end + must enqueue a retain for whatever is buffered. + """ + provider = _make_hindsight_provider() + provider._auto_retain = True + assert len(provider._session_turns) == 2 # pre-seeded by helper + + provider.on_session_end([]) + + # Buffer must be cleared after flush. + assert provider._session_turns == [] + assert provider._turn_counter == 0 + assert provider._turn_index == 0 + # A flush closure must have been enqueued on the writer queue. + assert not provider._retain_queue.empty() + + +def test_hindsight_on_session_end_noop_when_buffer_empty(): + """No flush should be enqueued when there are no buffered turns.""" + provider = _make_hindsight_provider() + provider._auto_retain = True + provider._session_turns = [] + provider._turn_counter = 0 + + provider.on_session_end([]) + + assert provider._retain_queue.empty() + + +def test_hindsight_on_session_end_noop_when_auto_retain_disabled(): + """Providers with auto_retain=False must not flush on session end.""" + provider = _make_hindsight_provider() + provider._auto_retain = False + assert len(provider._session_turns) == 2 # pre-seeded + + provider.on_session_end([]) + + # Turns should NOT be cleared — auto_retain is off, nothing happened. + assert len(provider._session_turns) == 2 + assert provider._retain_queue.empty() + + +def test_hindsight_on_session_end_noop_when_shutting_down(): + """If shutdown has already fired, on_session_end must not enqueue.""" + provider = _make_hindsight_provider() + provider._auto_retain = True + provider._shutting_down.set() # simulate shutdown already in progress + + provider.on_session_end([]) + + # Buffer untouched — we can't enqueue after shutdown. + assert len(provider._session_turns) == 2 + assert provider._retain_queue.empty() + + +def test_hindsight_on_session_end_then_shutdown_no_double_flush(): + """on_session_end clears the buffer, so a subsequent shutdown() + draining the writer queue must not re-flush the same turns.""" + provider = _make_hindsight_provider() + provider._auto_retain = True + + provider.on_session_end([]) + queue_size_after_end = provider._retain_queue.qsize() + + # Simulate what shutdown() does — it should find an empty buffer. + assert provider._session_turns == [] + assert queue_size_after_end == 1 # exactly one flush from on_session_end