fix: Preserve Session transformations during generation tracking

Author: seratch

src/agents/memory/openai_responses_compaction_session.py

@@ -303,11 +303,12 @@ async def get_items(self, limit: int | None = None) -> list[TResponseInputItem]:
         return await self.underlying_session.get_items(limit)
 
     async def _get_items_with_generation(
-        self, limit: int | None = None
+        self, read_items: Callable[[], Awaitable[list[TResponseInputItem]]]
     ) -> tuple[list[TResponseInputItem], int]:
         """Read one Runner snapshot with its exact wrapper generation."""
         async with self._mutation_lock:
-            items = await self.underlying_session.get_items(limit)
+            # Read through the outer Session so its decryption and filtering still apply.
+            items = await read_items()
             return items, self._mutation_generation
 
     async def _get_all_underlying_session_items(self) -> list[TResponseInputItem]:
@@ -460,15 +461,18 @@ async def add_items(self, items: list[TResponseInputItem]) -> None:
 
     async def _add_items_with_generation(
         self,
-        items: list[TResponseInputItem],
+        write_items: Callable[[], Awaitable[None]],
         *,
         expected_generation: int | None,
     ) -> int | None:
         """Append one Runner batch and retain ownership only when its read stayed current."""
-        async with self._mutation_lock:
-            owns_generation = expected_generation == self._mutation_generation
-            await self._add_items_locked(items)
-            return self._mutation_generation if owns_generation else None
+        # The outer Session applies its transformations before our public add_items acquires
+        # the mutation lock. Any intervening mutation revokes ownership, including a write
+        # that completes while the outer Session is awaiting acknowledgement after our append.
+        await write_items()
+        if expected_generation is not None and self._mutation_generation == expected_generation + 1:
+            return self._mutation_generation
+        return None
 
     async def _add_items_locked(self, items: list[TResponseInputItem]) -> None:
         try:

src/agents/run_internal/session_persistence.py

@@ -202,20 +202,23 @@ async def _session_get_items(
     capture_compaction_generation: bool = False,
 ) -> list[TResponseInputItem]:
     """Read session items while preserving the legacy method call shape."""
-    get_with_generation = getattr(session, "_get_items_with_generation", None)
-    if capture_compaction_generation and wrapper is not None and callable(get_with_generation):
+    session_wrapper = _get_session_wrapper(session, wrapper)
+
+    async def read_items() -> list[TResponseInputItem]:
         if limit is _SESSION_LIMIT_UNSET:
-            result, generation = await _call_session_method(get_with_generation)
+            result = await _call_session_method(session.get_items, wrapper=session_wrapper)
         else:
-            result, generation = await _call_session_method(get_with_generation, limit=limit)
+            result = await _call_session_method(
+                session.get_items, limit=limit, wrapper=session_wrapper
+            )
+        return cast(list[TResponseInputItem], result)
+
+    get_with_generation = getattr(session, "_get_items_with_generation", None)
+    if capture_compaction_generation and wrapper is not None and callable(get_with_generation):
+        result, generation = await _call_session_method(get_with_generation, read_items)
         wrapper._session_compaction_generation = generation  # type: ignore[attr-defined]
         return cast(list[TResponseInputItem], result)
-    wrapper = _get_session_wrapper(session, wrapper)
-    if limit is _SESSION_LIMIT_UNSET:
-        result = await _call_session_method(session.get_items, wrapper=wrapper)
-    else:
-        result = await _call_session_method(session.get_items, limit=limit, wrapper=wrapper)
-    return cast(list[TResponseInputItem], result)
+    return await read_items()
 
 
 async def _session_add_items(
@@ -225,18 +228,22 @@ async def _session_add_items(
     wrapper: RunContextWrapper[Any] | None = None,
 ) -> None:
     """Append session items while preserving the legacy method call shape."""
+    session_wrapper = _get_session_wrapper(session, wrapper)
+
+    async def write_items() -> None:
+        await _call_session_method(session.add_items, items, wrapper=session_wrapper)
+
     add_with_generation = getattr(session, "_add_items_with_generation", None)
     if wrapper is not None and callable(add_with_generation):
         expected_generation = getattr(wrapper, "_session_compaction_generation", None)
         generation = await _call_session_method(
             add_with_generation,
-            items,
+            write_items,
             expected_generation=expected_generation,
         )
         wrapper._session_compaction_generation = generation  # type: ignore[attr-defined]
         return
-    wrapper = _get_session_wrapper(session, wrapper)
-    await _call_session_method(session.add_items, items, wrapper=wrapper)
+    await write_items()
 
 
 async def _session_pop_item(
@@ -834,7 +841,7 @@ def digests(items: Sequence[TResponseInputItem]) -> list[str]:
             if wrapper is not None and callable(get_with_generation):
                 tail, committed_generation = await _call_session_method(
                     get_with_generation,
-                    limit=len(expected),
+                    lambda: _session_get_items(session, limit=len(expected), wrapper=wrapper),
                 )
             else:
                 tail = await _session_get_items(session, limit=len(expected), wrapper=wrapper)

tests/extensions/memory/test_encrypt_session.py

@@ -1,8 +1,12 @@
 from __future__ import annotations
 
+import asyncio
+import json
 import tempfile
 from pathlib import Path
+from types import SimpleNamespace
 from typing import Any, cast
+from unittest.mock import AsyncMock, MagicMock
 
 import pytest
 
@@ -12,15 +16,19 @@
 
 from agents import (
     Agent,
+    RunConfig,
     RunContextWrapper,
     Runner,
+    RunState,
     SessionSettings,
     SQLiteSession,
     TResponseInputItem,
 )
+from agents.decorators import tool
 from agents.extensions.memory.encrypt_session import EncryptedSession
-from agents.testing import ScriptedModel
-from tests.test_responses import get_text_message
+from agents.memory import OpenAIResponsesCompactionSession
+from agents.testing import ModelStep, ScriptedModel
+from tests.test_responses import get_function_tool_call, get_text_message
 
 # Mark all tests in this file as asyncio
 pytestmark = pytest.mark.asyncio
@@ -126,6 +134,228 @@ async def test_encrypted_session_with_runner(
     underlying_session.close()
 
 
+async def _run_encrypted_session(
+    agent: Agent[Any], value: str | RunState[Any], session: EncryptedSession, streamed: bool
+):
+    config = RunConfig(tracing_disabled=True)
+    if not streamed:
+        return await Runner.run(agent, value, session=session, run_config=config)
+    result = Runner.run_streamed(agent, value, session=session, run_config=config)
+    async for _ in result.stream_events():
+        pass
+    return result
+
+
+def _decrypt_stored_items(
+    session: EncryptedSession, stored: list[TResponseInputItem]
+) -> list[dict[str, Any]]:
+    # Inspect the storage envelope and decrypt directly, independently of Session.get_items.
+    envelopes = cast(list[dict[str, Any]], stored)
+    assert all(set(item) == {"__enc__", "v", "kid", "payload"} for item in envelopes)
+    assert all(item["__enc__"] == 1 for item in envelopes)
+    return [json.loads(session.cipher.decrypt(item["payload"].encode())) for item in envelopes]
+
+
+@pytest.mark.parametrize("streamed", [False, True])
+async def test_runner_encrypts_items_around_compaction(
+    streamed: bool, encryption_key: str, tmp_path: Path
+) -> None:
+    backend = SQLiteSession("encrypted-compaction", tmp_path / "history.db")
+    session = EncryptedSession(
+        backend.session_id,
+        OpenAIResponsesCompactionSession(
+            backend.session_id, backend, should_trigger_compaction=lambda _: False
+        ),
+        encryption_key,
+    )
+    output = get_text_message("private assistant answer")
+    model = ScriptedModel([[output], [get_text_message("followup answer")]])
+    agent = Agent(name="test", model=model)
+    expected = [
+        {"role": "user", "content": "private user input"},
+        output.model_dump(exclude_unset=True),
+    ]
+    try:
+        result = await _run_encrypted_session(agent, "private user input", session, streamed)
+        assert result.final_output == "private assistant answer"
+        assert _decrypt_stored_items(session, await backend.get_items()) == expected
+
+        await _run_encrypted_session(agent, "followup", session, streamed)
+        assert model.calls[-1].input == expected + [{"role": "user", "content": "followup"}]
+    finally:
+        backend.close()
+
+
+@pytest.mark.parametrize("streamed", [False, True])
+async def test_runner_decrypts_existing_compaction_history_with_ttl_and_limit(
+    streamed: bool, encryption_key: str, tmp_path: Path, set_fernet_time: Any
+) -> None:
+    backend = SQLiteSession("encrypted-history", tmp_path / "history.db")
+    session = EncryptedSession(
+        backend.session_id,
+        OpenAIResponsesCompactionSession(
+            backend.session_id, backend, should_trigger_compaction=lambda _: False
+        ),
+        encryption_key,
+        ttl=10,
+    )
+    expected_history = [{"role": "user", "content": "retained history"}]
+    try:
+        set_fernet_time(1_000)
+        await session.add_items([{"role": "assistant", "content": "expired history"}])
+        expired = (await backend.get_items())[0]
+        set_fernet_time(1_020)
+        await session.add_items(cast(list[TResponseInputItem], expected_history))
+        # An expired tail forces EncryptedSession to expand its one-item retrieval window.
+        await backend.add_items([expired])
+        session.session_settings = SessionSettings(limit=1)
+        model = ScriptedModel([[get_text_message("answer")]])
+
+        await _run_encrypted_session(Agent(name="test", model=model), "next", session, streamed)
+
+        assert model.calls[0].input == expected_history + [{"role": "user", "content": "next"}]
+    finally:
+        backend.close()
+
+
+@pytest.mark.parametrize("streamed", [False, True])
+async def test_runner_recovers_encrypted_resumed_append(
+    streamed: bool, encryption_key: str, tmp_path: Path
+) -> None:
+    class LostAckSQLiteSession(SQLiteSession):
+        """Fail after a real committed batch to exercise public Runner recovery."""
+
+        fail_after_commit = False
+
+        async def add_items(self, items: list[TResponseInputItem]) -> None:
+            await super().add_items(items)
+            if self.fail_after_commit:
+                self.fail_after_commit = False
+                raise RuntimeError("append acknowledgement lost")
+
+    backend = LostAckSQLiteSession("encrypted-resume", tmp_path / "history.db")
+    session = EncryptedSession(
+        backend.session_id,
+        OpenAIResponsesCompactionSession(
+            backend.session_id, backend, should_trigger_compaction=lambda _: False
+        ),
+        encryption_key,
+    )
+    effects: list[str] = []
+
+    @tool(needs_approval=True)
+    async def record() -> str:
+        effects.append("recorded")
+        return "private tool receipt"
+
+    call = get_function_tool_call("record", "{}", call_id="record-1")
+    output = get_text_message("private final answer")
+    model = ScriptedModel([[call], [output]])
+    agent = Agent(name="test", model=model, tools=[record])
+    try:
+        paused = await _run_encrypted_session(agent, "private request", session, streamed)
+        state = paused.to_state()
+        state.approve(state.get_interruptions()[0])
+        backend.fail_after_commit = True
+        with pytest.raises(RuntimeError, match="append acknowledgement lost"):
+            await _run_encrypted_session(agent, state, session, streamed)
+        assert effects == ["recorded"]
+        assert len(model.calls) == 1
+        state = await RunState.from_json(agent, state.to_json())
+
+        result = await _run_encrypted_session(agent, state, session, streamed)
+
+        assert result.final_output == "private final answer"
+        assert effects == ["recorded"]
+        assert "pending_session_write" not in result.to_state().to_json()
+        assert _decrypt_stored_items(session, await backend.get_items()) == [
+            {"role": "user", "content": "private request"},
+            call.model_dump(exclude_unset=True),
+            {
+                "type": "function_call_output",
+                "call_id": "record-1",
+                "output": "private tool receipt",
+            },
+            output.model_dump(exclude_unset=True),
+        ]
+    finally:
+        backend.close()
+
+
+@pytest.mark.parametrize(
+    "pause_before_append", [False, True], ids=["after-append", "before-append"]
+)
+async def test_encrypted_runner_skips_compaction_after_interleaved_outer_write(
+    pause_before_append: bool, encryption_key: str, tmp_path: Path
+) -> None:
+    paused = asyncio.Event()
+    release = asyncio.Event()
+
+    class PausingEncryptedSession(EncryptedSession):
+        """Suspend at the outer public append boundary while another Runner completes."""
+
+        async def add_items(
+            self,
+            items: list[TResponseInputItem],
+            *,
+            wrapper: RunContextWrapper[Any] | None = None,
+        ) -> None:
+            is_run_a_output = "answer-A" in json.dumps(items)
+            if is_run_a_output and pause_before_append:
+                paused.set()
+                await release.wait()
+            await super().add_items(items, wrapper=wrapper)
+            if is_run_a_output and not pause_before_append:
+                paused.set()
+                await release.wait()
+
+    backend = SQLiteSession("encrypted-concurrent", tmp_path / "history.db")
+    client = MagicMock()
+    client.responses.compact = AsyncMock(return_value=SimpleNamespace(output=[]))
+    session = PausingEncryptedSession(
+        backend.session_id,
+        OpenAIResponsesCompactionSession(
+            backend.session_id,
+            backend,
+            client=client,
+            should_trigger_compaction=lambda context: context["response_id"] == "resp-a",
+        ),
+        encryption_key,
+    )
+    output_a = get_text_message("answer-A")
+    output_b = get_text_message("answer-B")
+    agent_a = Agent(
+        name="A", model=ScriptedModel([ModelStep(output=[output_a], response_id="resp-a")])
+    )
+    agent_b = Agent(
+        name="B", model=ScriptedModel([ModelStep(output=[output_b], response_id="resp-b")])
+    )
+    run_a = asyncio.create_task(_run_encrypted_session(agent_a, "input-A", session, False))
+    try:
+        await asyncio.wait_for(paused.wait(), timeout=2)
+        result_b = await asyncio.wait_for(
+            _run_encrypted_session(agent_b, "input-B", session, False), timeout=2
+        )
+        assert result_b.final_output == "answer-B"
+        release.set()
+        result_a = await asyncio.wait_for(run_a, timeout=2)
+        assert result_a.final_output == "answer-A"
+        client.responses.compact.assert_not_awaited()
+        expected = [
+            {"role": "user", "content": "input-A"},
+            {"role": "user", "content": "input-B"},
+            output_b.model_dump(exclude_unset=True),
+        ]
+        expected.insert(3 if pause_before_append else 1, output_a.model_dump(exclude_unset=True))
+        assert _decrypt_stored_items(session, await backend.get_items()) == expected
+    finally:
+        release.set()
+        if not run_a.done():
+            run_a.cancel()
+        await asyncio.gather(run_a, return_exceptions=True)
+        backend.close()
+
+
 async def test_encrypted_session_pop_item(encryption_key: str, underlying_session: SQLiteSession):
     """Test pop_item functionality."""
     session = EncryptedSession(