diff --git a/python/packages/core/agent_framework/_sessions.py b/python/packages/core/agent_framework/_sessions.py index 7ea242078a..d16f6da17a 100644 --- a/python/packages/core/agent_framework/_sessions.py +++ b/python/packages/core/agent_framework/_sessions.py @@ -23,6 +23,7 @@ import math import os import threading +import time import uuid import warnings import weakref @@ -2000,6 +2001,37 @@ async def delete(self, session_id: str) -> None: self._sessions.pop(session_id, None) +_SESSION_FILE_REPLACE_ATTEMPTS = 3 +_SESSION_FILE_REPLACE_RETRY_DELAY_SECONDS = 0.01 +# Only Windows reports a briefly-held destination as ``PermissionError``; on POSIX the +# same code means the caller genuinely may not replace the file. +_RETRY_REPLACE_ON_THIS_PLATFORM: bool = os.name == "nt" + + +def _replace_session_file(source: Path, target: Path) -> None: + """Replace ``target`` with ``source``, retrying briefly on transient lock contention. + + Windows raises :class:`PermissionError` when the destination file is briefly held + open (for example right after a previous replacement), even though callers already + serialize access per snapshot. A bounded retry preserves the documented + last-writer-wins behaviour instead of surfacing the transient failure. + + The retry is Windows-only: on POSIX a :class:`PermissionError` means the caller + really is not allowed to replace the file, so retrying would only delay the same + failure. The whole sequence is bounded at ~30 ms and re-raises the original error + once the attempts are exhausted. + """ + attempts = _SESSION_FILE_REPLACE_ATTEMPTS if _RETRY_REPLACE_ON_THIS_PLATFORM else 1 + for attempt in range(attempts): + try: + os.replace(source, target) + return + except PermissionError: + if attempt == attempts - 1: + raise + time.sleep(_SESSION_FILE_REPLACE_RETRY_DELAY_SECONDS * (attempt + 1)) + + @experimental(feature_id=ExperimentalFeature.SESSION_STORE) class FileSessionStore(SessionStore): """File-backed storage for Agent Framework session snapshots. @@ -2151,7 +2183,7 @@ def _write() -> None: temp_path = file_path.with_name(f".{file_path.name}.{uuid.uuid4().hex}.tmp") try: temp_path.write_bytes(serialized) - os.replace(temp_path, file_path) + _replace_session_file(temp_path, file_path) finally: temp_path.unlink(missing_ok=True) diff --git a/python/packages/core/tests/core/test_sessions.py b/python/packages/core/tests/core/test_sessions.py index 23eee40810..6ffdbb3651 100644 --- a/python/packages/core/tests/core/test_sessions.py +++ b/python/packages/core/tests/core/test_sessions.py @@ -3,6 +3,7 @@ import asyncio import json import math +import os import threading import time import tracemalloc @@ -1257,6 +1258,95 @@ async def test_set_replaces_atomically_without_leaving_temp_files(self, tmp_path temp_files = await asyncio.to_thread(lambda: [*tmp_path.glob("*.tmp"), *tmp_path.glob(".*.tmp")]) assert not temp_files + async def test_set_retries_transient_permission_error_during_atomic_replace( + self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + # The retry is Windows-only, so pin the platform flag to exercise it on every CI OS. + monkeypatch.setattr("agent_framework._sessions._RETRY_REPLACE_ON_THIS_PLATFORM", True) + store = FileSessionStore(tmp_path) + await store.set("session-1", AgentSession(session_id="first")) + + real_replace = os.replace + attempts = 0 + + def flaky_replace(source: Any, target: Any) -> None: + nonlocal attempts + attempts += 1 + if attempts == 1: + raise PermissionError(13, "Access is denied") + real_replace(source, target) + + monkeypatch.setattr(os, "replace", flaky_replace) + + await store.set("session-1", AgentSession(session_id="second")) + + assert attempts >= 2 + stored = await store.get("session-1") + assert stored is not None + assert stored.session_id == "second" + temp_files = await asyncio.to_thread(lambda: [*tmp_path.glob("*.tmp"), *tmp_path.glob(".*.tmp")]) + assert not temp_files + + async def test_set_retries_three_times_then_preserves_snapshot_on_persistent_permission_error( + self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + # Pin the Windows-only retry path so every CI OS exercises the same code. + monkeypatch.setattr("agent_framework._sessions._RETRY_REPLACE_ON_THIS_PLATFORM", True) + store = FileSessionStore(tmp_path) + await store.set("session-1", AgentSession(session_id="first")) + + attempts = 0 + delays: list[float] = [] + + def denied_replace(source: Any, target: Any) -> None: + nonlocal attempts + attempts += 1 + raise PermissionError(13, "Access is denied") + + def recording_sleep(seconds: float) -> None: + delays.append(seconds) + + monkeypatch.setattr(os, "replace", denied_replace) + monkeypatch.setattr("agent_framework._sessions.time.sleep", recording_sleep) + + with pytest.raises(PermissionError): + await store.set("session-1", AgentSession(session_id="second")) + + assert attempts == 3 + assert delays == [pytest.approx(0.01), pytest.approx(0.02)] + temp_files = await asyncio.to_thread(lambda: [*tmp_path.glob("*.tmp"), *tmp_path.glob(".*.tmp")]) + assert not temp_files + stored = await store.get("session-1") + assert stored is not None + assert stored.session_id == "first" + + async def test_set_reports_permission_error_without_retrying_off_windows( + self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + # On POSIX a PermissionError is a real one: retrying only delays the same failure. + monkeypatch.setattr("agent_framework._sessions._RETRY_REPLACE_ON_THIS_PLATFORM", False) + store = FileSessionStore(tmp_path) + await store.set("session-1", AgentSession(session_id="first")) + + real_replace = os.replace + attempts = 0 + + def denied_replace(source: Any, target: Any) -> None: + nonlocal attempts + attempts += 1 + raise PermissionError(13, "Access is denied") + + monkeypatch.setattr(os, "replace", denied_replace) + + with pytest.raises(PermissionError): + await store.set("session-1", AgentSession(session_id="second")) + + assert attempts == 1 + monkeypatch.setattr(os, "replace", real_replace) + stored = await store.get("session-1") + assert stored is not None + assert stored.session_id == "first" + async def test_corrupt_session_file_is_quarantined_and_retry_returns_none(self, tmp_path: Path) -> None: store = FileSessionStore(tmp_path) await store.set("session-1", AgentSession(session_id="session-1"))