diff --git a/python/packages/core/agent_framework/_harness/_memory.py b/python/packages/core/agent_framework/_harness/_memory.py index c3f02f883b5..50731e82821 100644 --- a/python/packages/core/agent_framework/_harness/_memory.py +++ b/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 @@ -886,11 +887,38 @@ def get_index_text( return self._get_index_path(session, source_id=source_id).read_text(encoding="utf-8").strip() def read_state(self, session: AgentSession, *, source_id: str) -> dict[str, Any]: - """Return the maintenance state for the current owner.""" + """Return the maintenance state for the current owner. + + An unreadable state file (interrupted editor, partial external write) + must not fail every later turn's memory bookkeeping: quarantine it + next to the original, mirroring the session store's corrupt-snapshot + convention, and restart from a clean state. + """ state_path = self._get_state_path(session, source_id=source_id) if not state_path.exists(): return _default_state() - raw_state = self.loads(state_path.read_text(encoding="utf-8")) + try: + raw_state = self.loads(state_path.read_text(encoding="utf-8")) + except ValueError as exc: + quarantine_path = state_path.with_name(f".{state_path.name}.{uuid.uuid4().hex}.corrupt") + try: + state_path.rename(quarantine_path) + except OSError as rename_error: + LOGGER.warning( + "Memory state file %s is unreadable (%s) and could not be moved aside (%s); " + "using a clean state.", + state_path, + exc, + rename_error, + ) + else: + LOGGER.warning( + "Memory state file %s is unreadable (%s); moved aside to %s, using a clean state.", + state_path, + exc, + quarantine_path, + ) + return _default_state() 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)} diff --git a/python/packages/core/tests/core/test_harness_memory.py b/python/packages/core/tests/core/test_harness_memory.py index 464255d0b43..2f3ee6af96a 100644 --- a/python/packages/core/tests/core/test_harness_memory.py +++ b/python/packages/core/tests/core/test_harness_memory.py @@ -271,6 +271,44 @@ async def test_memory_file_store_writes_topics_index_state_and_transcripts(tmp_p ] +def test_memory_file_store_quarantines_corrupt_state_file(tmp_path: Path) -> None: + """An unreadable state.json is moved aside and the store restarts clean. + + Memory bookkeeping runs on every stored turn, so a corrupt state file + would otherwise raise out of every later turn. The corrupt file follows + the session store's quarantine convention and the next write recreates a + clean one. + """ + store = MemoryFileStore(tmp_path, owner_state_key="owner_id") + session = AgentSession(session_id="session-1") + session.state["owner_id"] = "alice" + store.write_state( + session, + {"sessions_since_consolidation": ["session-1"]}, + source_id=DEFAULT_MEMORY_SOURCE_ID, + ) + state_path = store._get_state_path(session, source_id=DEFAULT_MEMORY_SOURCE_ID) + state_path.write_text("{not json", encoding="utf-8") + + assert store.read_state(session, source_id=DEFAULT_MEMORY_SOURCE_ID) == { + "last_consolidated_at": None, + "sessions_since_consolidation": [], + } + assert not state_path.exists() + quarantined = list(state_path.parent.glob(".state.json.*.corrupt")) + assert len(quarantined) == 1 + assert quarantined[0].read_text(encoding="utf-8") == "{not json" + + store.write_state( + session, + {"sessions_since_consolidation": ["session-2"]}, + source_id=DEFAULT_MEMORY_SOURCE_ID, + ) + assert store.read_state(session, source_id=DEFAULT_MEMORY_SOURCE_ID)["sessions_since_consolidation"] == [ + "session-2" + ] + + async def _write_transcript(store: MemoryFileStore, session: AgentSession, session_id: str, text: str) -> None: """Append a one-message transcript for ``session_id`` through the real write path.""" provider = FileHistoryProvider(store.get_transcripts_directory(session, source_id=DEFAULT_MEMORY_SOURCE_ID)) diff --git a/python/packages/redis/agent_framework_redis/_context_provider.py b/python/packages/redis/agent_framework_redis/_context_provider.py index 4751bc16961..b02a057ab9a 100644 --- a/python/packages/redis/agent_framework_redis/_context_provider.py +++ b/python/packages/redis/agent_framework_redis/_context_provider.py @@ -115,6 +115,9 @@ def __init__( self._token_escaper: TokenEscaper = TokenEscaper() self._index_initialized: bool = False self._schema_dict: dict[str, Any] | None = None + # Ownership rule matches RedisHistoryProvider: only the index we built + # ourselves is ours to close; an injected index stays the caller's. + self._owns_index: bool = redis_index is None index = redis_index or AsyncSearchIndex.from_dict( # pyright: ignore[reportUnknownMemberType] self.schema_dict, redis_url=self.redis_url, validate_on_load=True ) @@ -440,8 +443,18 @@ async def __aenter__(self) -> Self: """Async context manager entry.""" return self + async def aclose(self) -> None: + """Close the underlying Redis connection when this provider built the index. + + An index injected via ``redis_index`` stays under the caller's control, + matching the ownership rule of ``RedisHistoryProvider``. + """ + if self._owns_index and self.redis_index is not None: + await self.redis_index.disconnect() + async def __aexit__(self, exc_type: type[BaseException] | None, exc_val: BaseException | None, exc_tb: Any) -> None: - """Async context manager exit.""" + """Async context manager exit; closes an owned index's connections.""" + await self.aclose() __all__ = ["RedisContextProvider"] diff --git a/python/packages/redis/tests/test_providers.py b/python/packages/redis/tests/test_providers.py index 711c63e7686..8885d7e21a8 100644 --- a/python/packages/redis/tests/test_providers.py +++ b/python/packages/redis/tests/test_providers.py @@ -490,6 +490,32 @@ async def test_borrowed_client_is_not_closed(self): await borrowed.aclose() +class TestRedisContextProviderClose: + """The context provider closes only the index it built itself, matching + the history provider's ownership rule.""" + + async def test_owned_index_is_disconnected(self, patch_index_from_dict: MagicMock, mock_index: AsyncMock): + provider = RedisContextProvider(source_id="ctx", user_id="u1") + + await provider.aclose() + + mock_index.disconnect.assert_awaited_once() + + async def test_borrowed_index_is_not_disconnected(self): + borrowed = AsyncMock() + provider = RedisContextProvider(source_id="ctx", redis_index=borrowed) + + await provider.aclose() + + borrowed.disconnect.assert_not_awaited() + + async def test_aexit_closes_owned_index(self, patch_index_from_dict: MagicMock, mock_index: AsyncMock): + async with RedisContextProvider(source_id="ctx", user_id="u1"): + pass + + mock_index.disconnect.assert_awaited_once() + + class TestRedisHistoryProviderRedisKey: def test_scoped_key_is_deterministic_and_encodes_opaque_components(self, mock_redis_client: MagicMock): with patch("agent_framework_redis._history_provider.redis.from_url") as mock_from_url: