Python: retry atomic session snapshot replacement on transient PermissionError

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)
 

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"))