From 4a53aa064987fefdb374fadd42dc1f5142621241 Mon Sep 17 00:00:00 2001 From: Shivam Sharma <192242476+ShivamSharma43@users.noreply.github.com> Date: Mon, 12 Oct 2026 16:05:56 +0530 Subject: [PATCH 1/3] Python: Keep sessions pending when a consolidation pass did not cover them Automatic consolidation cleared every pending session when a pass finished, including sessions whose memories were extracted after the pass took its topic snapshot. Those sessions stopped counting toward consolidation_min_sessions, so their memories waited for unrelated sessions before the next pass. A pass now retires only the sessions that were pending when it took its snapshot. after_run registers the session after extraction has written its topics, so a pass that sees a session as pending has also seen its memories. Co-Authored-By: Claude Opus 5.5 --- .../core/agent_framework/_harness/_memory.py | 26 ++- .../core/tests/core/test_harness_memory.py | 181 ++++++++++++++++++ 2 files changed, 198 insertions(+), 9 deletions(-) diff --git a/python/packages/core/agent_framework/_harness/_memory.py b/python/packages/core/agent_framework/_harness/_memory.py index 66163abed05..bca5bb7ba4e 100644 --- a/python/packages/core/agent_framework/_harness/_memory.py +++ b/python/packages/core/agent_framework/_harness/_memory.py @@ -1411,14 +1411,6 @@ 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, @@ -1433,6 +1425,16 @@ async def after_run( line_length=self.index_line_length, ) + # Register the session only after its extraction has written its topics, so a consolidation + # pass that sees this session as pending has also seen the memories it extracted. + 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) + await self._run_consolidation( client=self._chat_client(agent, override=self.consolidation_client), session=session, @@ -1568,6 +1570,8 @@ async def _run_consolidation( maintenance_state["sessions_since_consolidation"] = [] self.store.write_state(session, maintenance_state, source_id=self.source_id) return 0 + # Sessions registered after this snapshot are not covered by this pass and stay pending. + covered_session_ids = set(cast(list[str], maintenance_state["sessions_since_consolidation"])) success_count = 0 for record in topic_records: @@ -1590,7 +1594,11 @@ async def _run_consolidation( async with self._state_async_lock(session): maintenance_state = self.store.read_state(session, source_id=self.source_id) maintenance_state["last_consolidated_at"] = _timestamp(now) - maintenance_state["sessions_since_consolidation"] = [] + maintenance_state["sessions_since_consolidation"] = [ + session_id + for session_id in cast(list[str], maintenance_state["sessions_since_consolidation"]) + if session_id not in covered_session_ids + ] self.store.write_state(session, maintenance_state, source_id=self.source_id) self.store.rebuild_index( session, diff --git a/python/packages/core/tests/core/test_harness_memory.py b/python/packages/core/tests/core/test_harness_memory.py index d1c7b3dac51..d354f4fc388 100644 --- a/python/packages/core/tests/core/test_harness_memory.py +++ b/python/packages/core/tests/core/test_harness_memory.py @@ -734,6 +734,187 @@ async def test_memory_context_provider_preserves_concurrent_writes_to_same_topic assert sorted(topic.memories) == sorted(memories) +class _GatedMemoryClient: + """Chat client that can hold extraction for one marked turn, and consolidation, until released.""" + + additional_properties: dict[str, Any] + + def __init__(self, *, extraction_marker: str) -> None: + self.additional_properties = {} + self.extraction_marker = extraction_marker + self.extraction_started = asyncio.Event() + self.release_extraction = asyncio.Event() + self.consolidation_started = asyncio.Event() + self.release_consolidation = asyncio.Event() + self.release_consolidation.set() + self.consolidated_topics: list[str] = [] + + async def get_response(self, messages: Sequence[Message], **kwargs: object) -> ChatResponse[Any]: + del kwargs + system_text = messages[0].text.lower() if messages and messages[0].role == "system" else "" + if "extract durable memory candidates" in system_text: + memories: list[dict[str, str]] = [] + if self.extraction_marker in messages[1].text: + self.extraction_started.set() + await self.release_extraction.wait() + memories = [{"topic": "New", "memory": f"{self.extraction_marker} fact"}] + return ChatResponse(messages=[Message(role="assistant", contents=[json.dumps({"memories": memories})])]) + if "consolidate one topic memory file" in system_text: + self.consolidated_topics.append(messages[1].text) + self.consolidation_started.set() + await self.release_consolidation.wait() + payload = {"summary": "Old summary.", "memories": ["Old fact"]} + return ChatResponse(messages=[Message(role="assistant", contents=[json.dumps(payload)])]) + return ChatResponse(messages=[Message(role="assistant", contents=["Assistant reply."])]) + + +def _gated_consolidation_setup(tmp_path: Path, client: _GatedMemoryClient) -> tuple[MemoryFileStore, Agent]: + """Seed an Old topic and one pending session so the next finished session triggers consolidation.""" + store = MemoryFileStore(tmp_path, owner_state_key="owner_id") + provider = MemoryContextProvider(store=store, consolidation_min_sessions=2, consolidation_interval=timedelta(0)) + seed_session = _owned_session("session-a") + store.write_topic( + seed_session, + MemoryTopicRecord( + topic="Old", + summary="Old summary.", + memories=["Old fact"], + updated_at="2026-10-01T00:00:00+00:00", + session_ids=["session-a"], + ), + source_id=DEFAULT_MEMORY_SOURCE_ID, + ) + store.write_state( + seed_session, + {"last_consolidated_at": None, "sessions_since_consolidation": ["session-a"]}, + source_id=DEFAULT_MEMORY_SOURCE_ID, + ) + agent = Agent(client=client, context_providers=[provider], default_options=_no_store_options()) # type: ignore[arg-type] # pyrefly: ignore[bad-argument-type] # ty: ignore[invalid-argument-type] + return store, agent + + +def _owned_session(session_id: str) -> AgentSession: + session = AgentSession(session_id=session_id) + session.state["owner_id"] = "alice" + return session + + +def _pending_sessions(store: MemoryFileStore) -> list[str]: + pending: list[str] = store.read_state(_owned_session("reader"), source_id=DEFAULT_MEMORY_SOURCE_ID)[ + "sessions_since_consolidation" + ] + return pending + + +async def _run_turn(agent: Agent, text: str, session_id: str) -> None: + await agent.run(text, session=_owned_session(session_id)) + + +async def test_automatic_consolidation_keeps_session_extracting_during_pass(tmp_path: Path) -> None: + """A session still extracting when a consolidation pass finishes stays pending for the next pass.""" + client = _GatedMemoryClient(extraction_marker="B turn") + store, agent = _gated_consolidation_setup(tmp_path, client) + client.release_consolidation.clear() + + consolidating_run = asyncio.create_task(_run_turn(agent, "Second turn.", "session-2")) + await client.consolidation_started.wait() + extracting_run = asyncio.create_task(_run_turn(agent, "B turn.", "session-b")) + await client.extraction_started.wait() + + client.release_consolidation.set() + await consolidating_run + client.release_extraction.set() + await extracting_run + + assert _pending_sessions(store) == ["session-b"] + + await agent.run("Third turn.", session=_owned_session("session-c")) + assert _pending_sessions(store) == [] + assert len(client.consolidated_topics) == 3 + assert any("B turn fact" in topic for topic in client.consolidated_topics[1:]) + + +async def test_automatic_consolidation_keeps_session_extracting_before_snapshot(tmp_path: Path) -> None: + """A session that starts extracting before a pass snapshots the topics stays pending for the next pass.""" + client = _GatedMemoryClient(extraction_marker="B turn") + store, agent = _gated_consolidation_setup(tmp_path, client) + + extracting_run = asyncio.create_task(_run_turn(agent, "B turn.", "session-b")) + await client.extraction_started.wait() + + await agent.run("Second turn.", session=_owned_session("session-2")) + assert len(client.consolidated_topics) == 1 + assert _pending_sessions(store) == [] + + client.release_extraction.set() + await extracting_run + + assert _pending_sessions(store) == ["session-b"] + new_topic = store.get_topic(_owned_session("reader"), source_id=DEFAULT_MEMORY_SOURCE_ID, topic="New") + assert new_topic.memories == ["B turn fact"] + + await agent.run("Third turn.", session=_owned_session("session-c")) + assert _pending_sessions(store) == [] + assert any("B turn fact" in topic for topic in client.consolidated_topics[1:]) + + +class _RegisteringConsolidationClient: + """Consolidation client that registers another finished session while the pass is running.""" + + additional_properties: dict[str, Any] + + def __init__(self, store: MemoryFileStore, session: AgentSession, session_id: str) -> None: + self.additional_properties = {} + self.store = store + self.session = session + self.session_id = session_id + + async def get_response(self, *args: object, **kwargs: object) -> ChatResponse[Any]: + del args, kwargs + state = self.store.read_state(self.session, source_id=DEFAULT_MEMORY_SOURCE_ID) + state["sessions_since_consolidation"].append(self.session_id) + self.store.write_state(self.session, state, source_id=DEFAULT_MEMORY_SOURCE_ID) + payload = {"summary": "Old summary.", "memories": ["Old fact"]} + return ChatResponse(messages=[Message(role="assistant", contents=[json.dumps(payload)])]) + + +async def test_memory_consolidation_retires_only_sessions_in_its_snapshot(tmp_path: Path) -> None: + """Completing a pass clears the sessions it covered and keeps sessions registered after its snapshot.""" + session = _owned_session("session-2") + store = MemoryFileStore(tmp_path, owner_state_key="owner_id") + client = _RegisteringConsolidationClient(store, session, "session-b") + provider = MemoryContextProvider(store=store, consolidation_min_sessions=2, consolidation_interval=timedelta(0)) + store.write_state( + session, + {"last_consolidated_at": None, "sessions_since_consolidation": ["session-a", "session-2"]}, + source_id=DEFAULT_MEMORY_SOURCE_ID, + ) + store.write_topic( + session, + MemoryTopicRecord( + topic="Old", + summary="Old summary.", + memories=["Old fact"], + updated_at="2026-10-01T00:00:00+00:00", + session_ids=["session-a"], + ), + source_id=DEFAULT_MEMORY_SOURCE_ID, + ) + + consolidated_count = await provider._run_consolidation( # pyright: ignore[reportPrivateUsage] + client=client, # type: ignore[arg-type] # pyrefly: ignore[bad-argument-type] # ty: ignore[invalid-argument-type] + session=session, + force=False, + now=datetime(2026, 10, 8, tzinfo=timezone.utc), + ) + + assert consolidated_count == 1 + assert store.read_state(session, source_id=DEFAULT_MEMORY_SOURCE_ID) == { + "last_consolidated_at": "2026-10-08T00:00:00+00:00", + "sessions_since_consolidation": ["session-b"], + } + + def test_memory_harness_classes_are_marked_experimental() -> None: """Memory harness public classes should expose HARNESS experimental metadata.""" assert MemoryIndexEntry.__feature_id__ == ExperimentalFeature.HARNESS.value # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] From 480eea44c3d975bc1b0ef8c6ad83a8f9d2b457c5 Mon Sep 17 00:00:00 2001 From: Shivam Sharma <192242476+ShivamSharma43@users.noreply.github.com> Date: Tue, 13 Oct 2026 13:14:26 +0530 Subject: [PATCH 2/3] Python: Register a memory session even when its extraction fails or is cancelled Registering the session after extraction left it out of the pending list when a later topic write or the index rebuild failed, or the run was cancelled, after an earlier topic had already been written. Those topics then waited for unrelated sessions to trigger consolidation. Register the session in a `finally` once extraction is over, so it is still never visible as pending while its extraction is running. Co-Authored-By: Claude Opus 5.5 --- .../core/agent_framework/_harness/_memory.py | 46 +++++----- .../core/tests/core/test_harness_memory.py | 88 +++++++++++++++++++ 2 files changed, 112 insertions(+), 22 deletions(-) diff --git a/python/packages/core/agent_framework/_harness/_memory.py b/python/packages/core/agent_framework/_harness/_memory.py index bca5bb7ba4e..d3585d3c83b 100644 --- a/python/packages/core/agent_framework/_harness/_memory.py +++ b/python/packages/core/agent_framework/_harness/_memory.py @@ -1411,29 +1411,31 @@ 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) - 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, ) - - # Register the session only after its extraction has written its topics, so a consolidation - # pass that sees this session as pending has also seen the memories it extracted. - 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) + 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. + 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) await self._run_consolidation( client=self._chat_client(agent, override=self.consolidation_client), diff --git a/python/packages/core/tests/core/test_harness_memory.py b/python/packages/core/tests/core/test_harness_memory.py index d354f4fc388..2696997e3b5 100644 --- a/python/packages/core/tests/core/test_harness_memory.py +++ b/python/packages/core/tests/core/test_harness_memory.py @@ -915,6 +915,94 @@ async def test_memory_consolidation_retires_only_sessions_in_its_snapshot(tmp_pa } +class _TwoTopicExtractionClient: + """Chat client whose extraction returns two memories in different topics.""" + + additional_properties: dict[str, Any] + + def __init__(self) -> None: + self.additional_properties = {} + + async def get_response(self, messages: Sequence[Message], **kwargs: object) -> ChatResponse[Any]: + del kwargs + system_text = messages[0].text.lower() if messages and messages[0].role == "system" else "" + if "extract durable memory candidates" in system_text: + memories = [{"topic": "First", "memory": "First fact"}, {"topic": "Second", "memory": "Second fact"}] + return ChatResponse(messages=[Message(role="assistant", contents=[json.dumps({"memories": memories})])]) + return ChatResponse(messages=[Message(role="assistant", contents=["Assistant reply."])]) + + +def _two_topic_setup(tmp_path: Path) -> tuple[MemoryFileStore, MemoryContextProvider, Agent]: + store = MemoryFileStore(tmp_path, owner_state_key="owner_id") + provider = MemoryContextProvider(store=store, consolidation_min_sessions=10) + agent = Agent(client=_TwoTopicExtractionClient(), context_providers=[provider], default_options=_no_store_options()) # type: ignore[arg-type] # pyrefly: ignore[bad-argument-type] # ty: ignore[invalid-argument-type] + return store, provider, agent + + +@pytest.mark.parametrize("failing_step", ["second_topic_write", "rebuild_index"]) +async def test_memory_session_stays_pending_when_extraction_fails_after_a_topic_write( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, failing_step: str +) -> None: + """A run that persisted a topic and then failed still registers its session for consolidation.""" + store, _, agent = _two_topic_setup(tmp_path) + if failing_step == "second_topic_write": + original_write_topic = store.write_topic + written_topics: list[str] = [] + + def write_topic(session: AgentSession, record: MemoryTopicRecord, **kwargs: Any) -> None: + if written_topics: + raise OSError("disk full") + written_topics.append(record.topic) + original_write_topic(session, record, **kwargs) + + monkeypatch.setattr(store, "write_topic", write_topic) + else: + original_rebuild_index = store.rebuild_index + + def rebuild_index(session: AgentSession, **kwargs: Any) -> list[MemoryIndexEntry]: + # ``before_run`` also rebuilds the index; fail only the rebuild after extraction wrote topics. + if store.list_topics(session, source_id=DEFAULT_MEMORY_SOURCE_ID): + raise OSError("disk full") + return original_rebuild_index(session, **kwargs) + + monkeypatch.setattr(store, "rebuild_index", rebuild_index) + + with pytest.raises(OSError, match="disk full"): + await _run_turn(agent, "Remember this.", "session-1") + + first_topic = store.get_topic(_owned_session("reader"), source_id=DEFAULT_MEMORY_SOURCE_ID, topic="First") + assert first_topic.memories == ["First fact"] + assert _pending_sessions(store) == ["session-1"] + + +async def test_memory_session_stays_pending_when_extraction_is_cancelled_after_a_topic_write( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """A run cancelled after persisting a topic still registers its session for consolidation.""" + store, provider, agent = _two_topic_setup(tmp_path) + original_merge_memory = provider._merge_memory # pyright: ignore[reportPrivateUsage] + second_merge_started = asyncio.Event() + + async def merge_memory(**kwargs: Any) -> MemoryTopicRecord: + if kwargs["topic"] == "Second": + second_merge_started.set() + await asyncio.Event().wait() + return await original_merge_memory(**kwargs) + + monkeypatch.setattr(provider, "_merge_memory", merge_memory) + + run = asyncio.create_task(_run_turn(agent, "Remember this.", "session-1")) + await second_merge_started.wait() + assert _pending_sessions(store) == [] + run.cancel() + with pytest.raises(asyncio.CancelledError): + await run + + first_topic = store.get_topic(_owned_session("reader"), source_id=DEFAULT_MEMORY_SOURCE_ID, topic="First") + assert first_topic.memories == ["First fact"] + assert _pending_sessions(store) == ["session-1"] + + def test_memory_harness_classes_are_marked_experimental() -> None: """Memory harness public classes should expose HARNESS experimental metadata.""" assert MemoryIndexEntry.__feature_id__ == ExperimentalFeature.HARNESS.value # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] From 069de848892b1239ff2c7ab1b5ae0a9ba2c1b170 Mon Sep 17 00:00:00 2001 From: Shivam Sharma <192242476+ShivamSharma43@users.noreply.github.com> Date: Thu, 15 Oct 2026 13:44:10 +0530 Subject: [PATCH 3/3] Python: Record a distinct pending marker for each finished memory turn The pending-consolidation marker was the session id alone, de-duplicated. A second turn of a session that a running consolidation pass had already snapshotted as pending recorded nothing new, so when that pass succeeded it retired the session and dropped the later turn's topic out of the pending window. Agent.run does not serialize runs per AgentSession, so this can happen. Each finished turn now stores a fresh id for its session in `pending_turn_ids`. A pass snapshots those ids and retires a session only if its id is unchanged when the pass completes. `sessions_since_consolidation` keeps its format and still counts distinct sessions for `consolidation_min_sessions`; state without `pending_turn_ids` still loads, and a session pending without an id is retired by a pass that saw it. Co-Authored-By: Claude Opus 5.5 --- .../core/agent_framework/_harness/_memory.py | 33 ++++- .../core/tests/core/test_harness_memory.py | 124 +++++++++++++++++- 2 files changed, 145 insertions(+), 12 deletions(-) diff --git a/python/packages/core/agent_framework/_harness/_memory.py b/python/packages/core/agent_framework/_harness/_memory.py index d3585d3c83b..007ea99dfce 100644 --- a/python/packages/core/agent_framework/_harness/_memory.py +++ b/python/packages/core/agent_framework/_harness/_memory.py @@ -8,6 +8,7 @@ import os import re import threading +import uuid import weakref from abc import ABC, abstractmethod from base64 import b32decode @@ -236,7 +237,10 @@ 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 _format_search_results(results: Sequence[Mapping[str, Any]]) -> str: @@ -906,6 +910,8 @@ def read_state(self, session: AgentSession, *, source_id: str) -> dict[str, Any] 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("pending_turn_ids"), dict): + state["pending_turn_ids"] = {} if not isinstance(state.get("last_consolidated_at"), (str, type(None))): state["last_consolidated_at"] = None return state @@ -1429,12 +1435,15 @@ async def after_run( # 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.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 + 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( @@ -1570,10 +1579,16 @@ async def _run_consolidation( 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 - # Sessions registered after this snapshot are not covered by this pass and stay pending. - covered_session_ids = set(cast(list[str], maintenance_state["sessions_since_consolidation"])) + # 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: @@ -1596,11 +1611,19 @@ async def _run_consolidation( async with self._state_async_lock(session): maintenance_state = self.store.read_state(session, source_id=self.source_id) 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_session_ids + 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, diff --git a/python/packages/core/tests/core/test_harness_memory.py b/python/packages/core/tests/core/test_harness_memory.py index 2696997e3b5..a041e03a9d9 100644 --- a/python/packages/core/tests/core/test_harness_memory.py +++ b/python/packages/core/tests/core/test_harness_memory.py @@ -183,6 +183,7 @@ async def test_memory_file_store_writes_topics_index_state_and_transcripts(tmp_p assert store.read_state(session, source_id=DEFAULT_MEMORY_SOURCE_ID) == { "last_consolidated_at": None, "sessions_since_consolidation": [], + "pending_turn_ids": {}, } store.write_state( session, @@ -858,8 +859,97 @@ async def test_automatic_consolidation_keeps_session_extracting_before_snapshot( assert any("B turn fact" in topic for topic in client.consolidated_topics[1:]) +class _OrderedConsolidationClient: + """Chat client that holds each consolidation call until the test releases it.""" + + additional_properties: dict[str, Any] + + def __init__(self, *, extraction_marker: str) -> None: + self.additional_properties = {} + self.extraction_marker = extraction_marker + self.consolidated_topics: list[str] = [] + self.consolidation_gates: list[asyncio.Event] = [] + self.consolidation_call_started = asyncio.Event() + + async def wait_for_consolidation_calls(self, count: int) -> None: + while len(self.consolidation_gates) < count: + self.consolidation_call_started.clear() + await self.consolidation_call_started.wait() + + async def get_response(self, messages: Sequence[Message], **kwargs: object) -> ChatResponse[Any]: + del kwargs + system_text = messages[0].text.lower() if messages and messages[0].role == "system" else "" + if "extract durable memory candidates" in system_text: + memories: list[dict[str, str]] = [] + if self.extraction_marker in messages[1].text: + memories = [{"topic": "New", "memory": f"{self.extraction_marker} fact"}] + return ChatResponse(messages=[Message(role="assistant", contents=[json.dumps({"memories": memories})])]) + if "consolidate one topic memory file" in system_text: + self.consolidated_topics.append(messages[1].text) + gate = asyncio.Event() + self.consolidation_gates.append(gate) + self.consolidation_call_started.set() + await gate.wait() + topic = "New" if self.extraction_marker in messages[1].text else "Old" + payload = {"summary": f"{topic} summary.", "memories": [f"{topic} fact"]} + return ChatResponse(messages=[Message(role="assistant", contents=[json.dumps(payload)])]) + return ChatResponse(messages=[Message(role="assistant", contents=["Assistant reply."])]) + + +async def test_automatic_consolidation_keeps_overlapping_turn_of_covered_session_pending(tmp_path: Path) -> None: + """A later turn of a session that a running pass already covers stays pending after that pass. + + ``Agent.run`` does not serialize runs per session, so a second ``session-1`` turn can write a new + topic and register itself while a pass that snapshotted ``session-1`` as pending is still running. + """ + client = _OrderedConsolidationClient(extraction_marker="Later turn") + store = MemoryFileStore(tmp_path, owner_state_key="owner_id") + provider = MemoryContextProvider(store=store, consolidation_min_sessions=2, consolidation_interval=timedelta(0)) + seed_session = _owned_session("session-a") + store.write_topic( + seed_session, + MemoryTopicRecord( + topic="Old", + summary="Old summary.", + memories=["Old fact"], + updated_at="2026-10-01T00:00:00+00:00", + session_ids=["session-a"], + ), + source_id=DEFAULT_MEMORY_SOURCE_ID, + ) + store.write_state( + seed_session, + {"last_consolidated_at": None, "sessions_since_consolidation": ["session-a"]}, + source_id=DEFAULT_MEMORY_SOURCE_ID, + ) + agent = Agent(client=client, context_providers=[provider], default_options=_no_store_options()) # type: ignore[arg-type] # pyrefly: ignore[bad-argument-type] # ty: ignore[invalid-argument-type] + + # The first turn registers session-1 and starts a pass whose snapshot covers it and the Old topic. + first_run = asyncio.create_task(_run_turn(agent, "First turn.", "session-1")) + await client.wait_for_consolidation_calls(1) + + # The overlapping second turn writes the New topic, registers session-1 again and starts its own + # pass, which holds on the New topic. + second_run = asyncio.create_task(_run_turn(agent, "Later turn.", "session-1")) + await client.wait_for_consolidation_calls(2) + assert "Later turn fact" in client.consolidated_topics[1] + assert _pending_sessions(store) == ["session-a", "session-1"] + + # The first pass finishes: its snapshot did not include the second turn, so session-1 stays pending. + client.consolidation_gates[0].set() + await first_run + assert _pending_sessions(store) == ["session-1"] + + # The second pass covered the second turn and retires it. + client.consolidation_gates[1].set() + await client.wait_for_consolidation_calls(3) + client.consolidation_gates[2].set() + await second_run + assert _pending_sessions(store) == [] + + class _RegisteringConsolidationClient: - """Consolidation client that registers another finished session while the pass is running.""" + """Consolidation client that registers another finished turn while the pass is running.""" additional_properties: dict[str, Any] @@ -872,21 +962,38 @@ def __init__(self, store: MemoryFileStore, session: AgentSession, session_id: st async def get_response(self, *args: object, **kwargs: object) -> ChatResponse[Any]: del args, kwargs state = self.store.read_state(self.session, source_id=DEFAULT_MEMORY_SOURCE_ID) - state["sessions_since_consolidation"].append(self.session_id) + if self.session_id not in state["sessions_since_consolidation"]: + state["sessions_since_consolidation"].append(self.session_id) + state["pending_turn_ids"][self.session_id] = "later-turn" self.store.write_state(self.session, state, source_id=DEFAULT_MEMORY_SOURCE_ID) payload = {"summary": "Old summary.", "memories": ["Old fact"]} return ChatResponse(messages=[Message(role="assistant", contents=[json.dumps(payload)])]) -async def test_memory_consolidation_retires_only_sessions_in_its_snapshot(tmp_path: Path) -> None: - """Completing a pass clears the sessions it covered and keeps sessions registered after its snapshot.""" +@pytest.mark.parametrize( + "later_session_id", + [ + pytest.param("session-b", id="new_session"), + pytest.param("session-a", id="pending_session_without_turn_id"), + pytest.param("session-2", id="pending_session_with_turn_id"), + ], +) +async def test_memory_consolidation_retires_only_turns_in_its_snapshot(tmp_path: Path, later_session_id: str) -> None: + """Completing a pass clears the turns it covered and keeps turns registered after its snapshot. + + ``session-a`` is pending without a turn id, like state written before turn ids were recorded. + """ session = _owned_session("session-2") store = MemoryFileStore(tmp_path, owner_state_key="owner_id") - client = _RegisteringConsolidationClient(store, session, "session-b") + client = _RegisteringConsolidationClient(store, session, later_session_id) provider = MemoryContextProvider(store=store, consolidation_min_sessions=2, consolidation_interval=timedelta(0)) store.write_state( session, - {"last_consolidated_at": None, "sessions_since_consolidation": ["session-a", "session-2"]}, + { + "last_consolidated_at": None, + "sessions_since_consolidation": ["session-a", "session-2"], + "pending_turn_ids": {"session-2": "earlier-turn"}, + }, source_id=DEFAULT_MEMORY_SOURCE_ID, ) store.write_topic( @@ -911,7 +1018,8 @@ async def test_memory_consolidation_retires_only_sessions_in_its_snapshot(tmp_pa assert consolidated_count == 1 assert store.read_state(session, source_id=DEFAULT_MEMORY_SOURCE_ID) == { "last_consolidated_at": "2026-10-08T00:00:00+00:00", - "sessions_since_consolidation": ["session-b"], + "sessions_since_consolidation": [later_session_id], + "pending_turn_ids": {later_session_id: "later-turn"}, } @@ -1088,6 +1196,7 @@ async def test_memory_file_store_does_not_mkdir_on_pure_read_paths(tmp_path) -> assert store.read_state(session, source_id=DEFAULT_MEMORY_SOURCE_ID) == { "last_consolidated_at": None, "sessions_since_consolidation": [], + "pending_turn_ids": {}, } assert store.search_transcripts(session, source_id=DEFAULT_MEMORY_SOURCE_ID, query="anything") == [] @@ -1147,6 +1256,7 @@ async def test_memory_consolidation_transient_failure_preserves_state(tmp_path) pre_state = { "last_consolidated_at": "2026-04-20T09:00:00+00:00", "sessions_since_consolidation": ["queued-session"], + "pending_turn_ids": {"queued-session": "turn-1"}, } store.write_state(session, pre_state, source_id=DEFAULT_MEMORY_SOURCE_ID) store.write_topic(