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/agent_framework/_workflows/_checkpoint_encoding.py b/python/packages/core/agent_framework/_workflows/_checkpoint_encoding.py index 209d1ffea9d..af916b14a12 100644 --- a/python/packages/core/agent_framework/_workflows/_checkpoint_encoding.py +++ b/python/packages/core/agent_framework/_workflows/_checkpoint_encoding.py @@ -54,6 +54,7 @@ from enum import EnumMeta from typing import Any, cast +from .._settings import SecretString from ..exceptions import WorkflowCheckpointException logger = logging.getLogger("agent_framework") @@ -370,6 +371,18 @@ def _decode(value: Any, *, allowed_types: frozenset[str] | None = None) -> Any: def _encode_pickle(value: Any) -> dict[str, str]: """Encode a value as a pickle envelope.""" + # SecretString is deliberately unpicklable: protocols >= 2 fail at restore + # on its immutability guard, and a working pickle would write the + # credential to durable storage in cleartext. Refuse at encode time with a + # precise error instead of the opaque restore failure (#9514); the session + # state path already rejects it the same way (#8874). + if isinstance(value, SecretString): + raise WorkflowCheckpointException( + "SecretString cannot be checkpointed: it is unpicklable by design, and " + "persisting it would write the credential to durable storage in cleartext. " + "Keep credentials out of checkpointed state and re-resolve them from " + "settings or configuration on restore." + ) return { _PICKLE_MARKER: _pickle_to_base64(value), _TYPE_MARKER: _value_type_to_key(value), 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/core/tests/workflow/test_checkpoint.py b/python/packages/core/tests/workflow/test_checkpoint.py index 6e13a414f8b..5ecb5024cc1 100644 --- a/python/packages/core/tests/workflow/test_checkpoint.py +++ b/python/packages/core/tests/workflow/test_checkpoint.py @@ -1170,6 +1170,20 @@ def _build_workflow() -> Any: # region FileCheckpointStorage +async def test_file_checkpoint_storage_save_refuses_secret_string(tmp_path): + """The precise encode error surfaces unwrapped at save time (#9514).""" + from agent_framework._settings import SecretString + + storage = FileCheckpointStorage(str(tmp_path)) + checkpoint = WorkflowCheckpoint( + workflow_name="w", + graph_signature_hash="h", + state={"credential": SecretString("sk-secret")}, + ) + with pytest.raises(WorkflowCheckpointException, match="SecretString cannot be checkpointed"): + await storage.save(checkpoint) + + async def test_file_checkpoint_storage_save_and_load(): with tempfile.TemporaryDirectory() as temp_dir: diff --git a/python/packages/core/tests/workflow/test_checkpoint_encode.py b/python/packages/core/tests/workflow/test_checkpoint_encode.py index 3b2a9149b9e..0a050e68cf1 100644 --- a/python/packages/core/tests/workflow/test_checkpoint_encode.py +++ b/python/packages/core/tests/workflow/test_checkpoint_encode.py @@ -6,6 +6,10 @@ from datetime import datetime, timezone from typing import Any, cast +import pytest + +from agent_framework._settings import SecretString +from agent_framework.exceptions import WorkflowCheckpointException from agent_framework._workflows._checkpoint_encoding import ( _PICKLE_MARKER, # pyright: ignore[reportPrivateUsage] _TYPE_MARKER, # pyright: ignore[reportPrivateUsage] @@ -437,3 +441,19 @@ def test_decode_preserves_user_dict_matching_old_escape_shape() -> None: """Test that user data shaped like an old escape envelope remains unchanged.""" data = {"__agent_framework_checkpoint_dict__": True, "value": {"safe": "data"}} assert decode_checkpoint_value(data) == data + + +def test_encode_secret_string_rejected_precisely() -> None: + """A credential in checkpointed state fails at encode, not at restore (#9514).""" + with pytest.raises(WorkflowCheckpointException, match="SecretString cannot be checkpointed"): + encode_checkpoint_value(SecretString("sk-secret")) + + +def test_encode_secret_string_nested_in_dict_rejected() -> None: + with pytest.raises(WorkflowCheckpointException, match="SecretString cannot be checkpointed"): + encode_checkpoint_value({"api_key": SecretString("sk-secret")}) + + +def test_encode_secret_string_nested_in_list_rejected() -> None: + with pytest.raises(WorkflowCheckpointException, match="SecretString cannot be checkpointed"): + encode_checkpoint_value({"keys": [SecretString("sk-secret")]}) diff --git a/python/packages/orchestrations/agent_framework_orchestrations/_magentic.py b/python/packages/orchestrations/agent_framework_orchestrations/_magentic.py index 52e74d9eed6..98d84017a67 100644 --- a/python/packages/orchestrations/agent_framework_orchestrations/_magentic.py +++ b/python/packages/orchestrations/agent_framework_orchestrations/_magentic.py @@ -942,7 +942,10 @@ def __init__( # Termination related state self._terminated: bool = False - self._max_rounds = manager.max_round_count + # Coerce like the base orchestrator contract ("smaller than 1 will be + # coerced to 1") instead of overriding the coerced value with the raw + # manager one, which would terminate the run before any round. + self._max_rounds = max(1, manager.max_round_count) if manager.max_round_count is not None else None @override async def _handle_messages( diff --git a/python/packages/orchestrations/tests/test_magentic.py b/python/packages/orchestrations/tests/test_magentic.py index a93e8b51328..0e24b4457c5 100644 --- a/python/packages/orchestrations/tests/test_magentic.py +++ b/python/packages/orchestrations/tests/test_magentic.py @@ -436,6 +436,23 @@ async def test_magentic_orchestrator_round_limit_produces_partial_result(): assert data.role == "assistant" +async def test_magentic_zero_max_round_count_is_coerced_to_one(): + """`max_round_count` below 1 is coerced to 1 (base orchestrator contract), so + the run gets its one coordination round and reaches the manager's final + answer instead of terminating before any participant is invoked.""" + manager = FakeManager(max_round_count=0) + wf = MagenticBuilder( + participants=[StubAgent(manager.next_speaker_name, "first round ran")], manager=manager + ).build() + + events = await wf.run("zero round limit") + outputs = [ev for ev in events if ev.type == "output"] + + assert len(outputs) == 1 + assert isinstance(outputs[0].data, AgentResponse) + assert outputs[0].data.messages[-1].text == manager.FINAL_ANSWER + + async def test_magentic_checkpoint_resume_round_trip(): storage = InMemoryCheckpointStorage() 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: