fix: address typecheck issues in wrapper compatibility changes

Author: HuxleyHu98

src/agents/extensions/memory/advanced_sqlite_session.py

@@ -140,8 +140,9 @@ async def add_items(
     async def get_items(
         self,
         limit: int | None = None,
+        wrapper: RunContextWrapper[Any] | str | None = None,
+        *,
         branch_id: str | None = None,
-        wrapper: RunContextWrapper[Any] | None = None,
     ) -> list[TResponseInputItem]:
         """Get items from current or specified branch.
 
@@ -152,6 +153,10 @@ async def get_items(
         Returns:
             List of conversation items from the specified branch.
         """
+        if isinstance(wrapper, str) and branch_id is None:
+            branch_id = wrapper
+            wrapper = None
+
         session_limit = resolve_session_limit(limit, self.session_settings)
 
         if branch_id is None:

tests/extensions/memory/test_encrypt_session.py

@@ -7,6 +7,8 @@
 
 pytest.importorskip("cryptography")  # Skip tests if cryptography is not installed
 
+from typing import cast
+
 from cryptography.fernet import Fernet
 
 from agents import Agent, Runner, SQLiteSession, TResponseInputItem
@@ -116,7 +118,7 @@ async def test_encrypted_session_preserves_legacy_underlying_signatures(
     agent: Agent,
     encryption_key: str,
 ):
-    class LegacyUnderlyingSession(SessionABC):
+    class LegacyUnderlyingSession:
         def __init__(self) -> None:
             self.session_id = "test_session"
             self.items: list[TResponseInputItem] = []
@@ -140,7 +142,7 @@ async def clear_session(self) -> None:
     legacy_underlying = LegacyUnderlyingSession()
     session = EncryptedSession(
         session_id="test_session",
-        underlying_session=legacy_underlying,
+        underlying_session=cast(SessionABC, legacy_underlying),
         encryption_key=encryption_key,
     )
 

tests/test_session.py

@@ -9,6 +9,7 @@
 
 from agents import Agent, RunConfig, Runner, SQLiteSession, TResponseInputItem
 from agents.memory.session import SessionABC
+from agents.memory.session_settings import SessionSettings
 from agents.run_context import RunContextWrapper
 from agents.run_internal.session_persistence import prepare_input_with_session
 
@@ -136,16 +137,26 @@ async def clear_session(self) -> None:
 
 class DefaultLimitedSession:
     session_id = "default-limited"
+    session_settings: SessionSettings | None = None
 
     def __init__(self) -> None:
         self.items: list[TResponseInputItem] = []
         self.get_call_count = 0
 
-    async def get_items(self, limit: int = 1) -> list[TResponseInputItem]:
+    async def get_items(
+        self,
+        limit: int | None = None,
+        wrapper: RunContextWrapper[Any] | None = None,
+    ) -> list[TResponseInputItem]:
         self.get_call_count += 1
-        return list(self.items[-limit:]) if limit > 0 else []
+        effective_limit = 1 if limit is None else limit
+        return list(self.items[-effective_limit:]) if effective_limit > 0 else []
 
-    async def add_items(self, items: list[TResponseInputItem]) -> None:
+    async def add_items(
+        self,
+        items: list[TResponseInputItem],
+        wrapper: RunContextWrapper[Any] | None = None,
+    ) -> None:
         self.items.extend(items)
 
     async def pop_item(self) -> TResponseInputItem | None:
@@ -333,8 +344,10 @@ async def test_get_items_preserves_default_limit_when_none_is_unset() -> None:
     )
 
     assert isinstance(prepared_input, list)
-    assert prepared_input[0]["content"] == "two"
-    assert prepared_input[1]["content"] == "new"
+    first = prepared_input[0]
+    second = prepared_input[1]
+    assert isinstance(first, dict) and first.get("content") == "two"
+    assert isinstance(second, dict) and second.get("content") == "new"
     assert session.get_call_count == 1