fix: preserve legacy wrapper compatibility in EncryptedSession

Author: HuxleyHu98

src/agents/extensions/memory/encrypt_session.py

@@ -28,6 +28,7 @@
 from __future__ import annotations
 
 import base64
+import inspect
 import json
 from typing import Any, cast
 
@@ -97,6 +98,35 @@ def _is_encrypted_envelope(item: object) -> TypeGuard[EncryptedEnvelope]:
     )
 
 
+def _method_accepts_wrapper(method: Any) -> bool:
+    try:
+        parameters = tuple(inspect.signature(method).parameters.values())
+    except (TypeError, ValueError):
+        return False
+
+    return any(
+        parameter.kind is inspect.Parameter.VAR_KEYWORD or parameter.name == "wrapper"
+        for parameter in parameters
+    )
+
+
+def _method_accepts_limit(method: Any) -> bool:
+    try:
+        parameters = tuple(inspect.signature(method).parameters.values())
+    except (TypeError, ValueError):
+        return False
+
+    return any(
+        (
+            parameter.kind
+            in (inspect.Parameter.POSITIONAL_OR_KEYWORD, inspect.Parameter.KEYWORD_ONLY)
+            and parameter.name == "limit"
+        )
+        or parameter.kind is inspect.Parameter.VAR_KEYWORD
+        for parameter in parameters
+    )
+
+
 class EncryptedSession(SessionABC):
     """Encrypted wrapper for Session implementations with TTL-based expiration.
 
@@ -176,8 +206,10 @@ async def get_items(
         limit: int | None = None,
         wrapper: RunContextWrapper[Any] | None = None,
     ) -> list[TResponseInputItem]:
-        if wrapper is not None:
+        if wrapper is not None and _method_accepts_wrapper(self.underlying_session.get_items):
             encrypted_items = await self.underlying_session.get_items(limit, wrapper=wrapper)
+        elif limit is None and not _method_accepts_limit(self.underlying_session.get_items):
+            encrypted_items = await self.underlying_session.get_items()
         else:
             encrypted_items = await self.underlying_session.get_items(limit)
         valid_items: list[TResponseInputItem] = []
@@ -194,7 +226,7 @@ async def add_items(
     ) -> None:
         wrapped: list[EncryptedEnvelope] = [self._wrap(it) for it in items]
         wrapped_items = cast(list[TResponseInputItem], wrapped)
-        if wrapper is not None:
+        if wrapper is not None and _method_accepts_wrapper(self.underlying_session.add_items):
             await self.underlying_session.add_items(
                 wrapped_items,
                 wrapper=wrapper,

tests/extensions/memory/test_encrypt_session.py

@@ -11,6 +11,7 @@
 
 from agents import Agent, Runner, SQLiteSession, TResponseInputItem
 from agents.extensions.memory.encrypt_session import EncryptedSession
+from agents.memory.session import SessionABC
 from tests.fake_model import FakeModel
 from tests.test_responses import get_text_message
 
@@ -111,6 +112,46 @@ async def test_encrypted_session_with_runner(
     underlying_session.close()
 
 
+async def test_encrypted_session_preserves_legacy_underlying_signatures(
+    agent: Agent,
+    encryption_key: str,
+):
+    class LegacyUnderlyingSession(SessionABC):
+        def __init__(self) -> None:
+            self.session_id = "test_session"
+            self.items: list[TResponseInputItem] = []
+
+        async def get_items(self, limit: int | None = None) -> list[TResponseInputItem]:
+            if limit is None:
+                return list(self.items)
+            return list(self.items[-limit:]) if limit > 0 else []
+
+        async def add_items(self, items: list[TResponseInputItem]) -> None:
+            self.items.extend(items)
+
+        async def pop_item(self) -> TResponseInputItem | None:
+            if not self.items:
+                return None
+            return self.items.pop()
+
+        async def clear_session(self) -> None:
+            self.items.clear()
+
+    legacy_underlying = LegacyUnderlyingSession()
+    session = EncryptedSession(
+        session_id="test_session",
+        underlying_session=legacy_underlying,
+        encryption_key=encryption_key,
+    )
+
+    assert isinstance(agent.model, FakeModel)
+    agent.model.set_next_output([get_text_message("Hello")])
+    result = await Runner.run(agent, "Hi there", session=session)
+
+    assert result.final_output == "Hello"
+    assert legacy_underlying.items
+
+
 async def test_encrypted_session_pop_item(encryption_key: str, underlying_session: SQLiteSession):
     """Test pop_item functionality."""
     session = EncryptedSession(