From f6cf6a5160fbdaf4c7ce657ccfc26392bd34c39c Mon Sep 17 00:00:00 2001 From: Limitless2023 Date: Sat, 14 Feb 2026 22:23:51 +0800 Subject: [PATCH 1/2] feat: add optional wrapper parameter to Session methods Add optional RunContextWrapper parameter to Session protocol methods (get_items, add_items, pop_item, clear_session) to allow custom sessions to access context data during history operations. This enables custom sessions to make decisions based on agent context, such as storing agent-specific information during handoff. Fixes #2072 --- .../memory/openai_conversations_session.py | 47 ++++++++++++++-- .../openai_responses_compaction_session.py | 53 ++++++++++++++++--- src/agents/memory/session.py | 48 +++++++++++++---- src/agents/memory/sqlite_session.py | 29 ++++++++-- 4 files changed, 151 insertions(+), 26 deletions(-) diff --git a/src/agents/memory/openai_conversations_session.py b/src/agents/memory/openai_conversations_session.py index f2345f22f2..6e572bafae 100644 --- a/src/agents/memory/openai_conversations_session.py +++ b/src/agents/memory/openai_conversations_session.py @@ -1,10 +1,13 @@ from __future__ import annotations +from typing import Any + from openai import AsyncOpenAI from agents.models._openai_shared import get_default_openai_client from ..items import TResponseInputItem +from ..run_context import RunContextWrapper from .session import SessionABC from .session_settings import SessionSettings, resolve_session_limit @@ -68,7 +71,18 @@ async def _get_session_id(self) -> str: async def _clear_session_id(self) -> None: self._session_id = None - async def get_items(self, limit: int | None = None) -> list[TResponseInputItem]: + async def get_items( + self, limit: int | None = None, wrapper: RunContextWrapper[Any] | None = None + ) -> list[TResponseInputItem]: + """Retrieve the conversation history for this session. + + Args: + limit: Maximum number of items to retrieve. If None, retrieves all items. + wrapper: Optional RunContextWrapper for accessing context data during retrieval. + + Returns: + List of input items representing the conversation history + """ session_id = await self._get_session_id() session_limit = resolve_session_limit(limit, self.session_settings) @@ -95,7 +109,15 @@ async def get_items(self, limit: int | None = None) -> list[TResponseInputItem]: return all_items # type: ignore - async def add_items(self, items: list[TResponseInputItem]) -> None: + async def add_items( + self, items: list[TResponseInputItem], wrapper: RunContextWrapper[Any] | None = None + ) -> None: + """Add new items to the conversation history. + + Args: + items: List of input items to add to the history + wrapper: Optional RunContextWrapper for accessing context data during addition + """ session_id = await self._get_session_id() if not items: return @@ -105,7 +127,17 @@ async def add_items(self, items: list[TResponseInputItem]) -> None: items=items, ) - async def pop_item(self) -> TResponseInputItem | None: + async def pop_item( + self, wrapper: RunContextWrapper[Any] | None = None + ) -> TResponseInputItem | None: + """Remove and return the most recent item from the session. + + Args: + wrapper: Optional RunContextWrapper for accessing context data during pop + + Returns: + The most recent item if it exists, None if the session is empty + """ session_id = await self._get_session_id() items = await self.get_items(limit=1) if not items: @@ -116,7 +148,14 @@ async def pop_item(self) -> TResponseInputItem | None: ) return items[0] - async def clear_session(self) -> None: + async def clear_session( + self, wrapper: RunContextWrapper[Any] | None = None + ) -> None: + """Clear all items for this session. + + Args: + wrapper: Optional RunContextWrapper for accessing context data during clear + """ session_id = await self._get_session_id() await self._openai_client.conversations.delete( conversation_id=session_id, diff --git a/src/agents/memory/openai_responses_compaction_session.py b/src/agents/memory/openai_responses_compaction_session.py index e2148f4868..7e89fb6dd5 100644 --- a/src/agents/memory/openai_responses_compaction_session.py +++ b/src/agents/memory/openai_responses_compaction_session.py @@ -6,6 +6,7 @@ from openai import AsyncOpenAI from ..models._openai_shared import get_default_openai_client +from ..run_context import RunContextWrapper from .openai_conversations_session import OpenAIConversationsSession from .session import ( OpenAIResponsesCompactionArgs, @@ -236,8 +237,19 @@ async def run_compaction(self, args: OpenAIResponsesCompactionArgs | None = None f"candidates={len(self._compaction_candidate_items)})" ) - async def get_items(self, limit: int | None = None) -> list[TResponseInputItem]: - return await self.underlying_session.get_items(limit) + async def get_items( + self, limit: int | None = None, wrapper: RunContextWrapper[Any] | None = None + ) -> list[TResponseInputItem]: + """Retrieve the conversation history for this session. + + Args: + limit: Maximum number of items to retrieve. + wrapper: Optional RunContextWrapper for accessing context data during retrieval. + + Returns: + List of input items representing the conversation history + """ + return await self.underlying_session.get_items(limit, wrapper) async def _defer_compaction(self, response_id: str, store: bool | None = None) -> None: if self._deferred_response_id is not None: @@ -265,8 +277,16 @@ def _get_deferred_compaction_response_id(self) -> str | None: def _clear_deferred_compaction(self) -> None: self._deferred_response_id = None - async def add_items(self, items: list[TResponseInputItem]) -> None: - await self.underlying_session.add_items(items) + async def add_items( + self, items: list[TResponseInputItem], wrapper: RunContextWrapper[Any] | None = None + ) -> None: + """Add new items to the conversation history. + + Args: + items: List of input items to add to the history + wrapper: Optional RunContextWrapper for accessing context data during addition + """ + await self.underlying_session.add_items(items, wrapper) if self._compaction_candidate_items is not None: new_candidates = select_compaction_candidate_items(items) if new_candidates: @@ -274,15 +294,32 @@ async def add_items(self, items: list[TResponseInputItem]) -> None: if self._session_items is not None: self._session_items.extend(items) - async def pop_item(self) -> TResponseInputItem | None: - popped = await self.underlying_session.pop_item() + async def pop_item( + self, wrapper: RunContextWrapper[Any] | None = None + ) -> TResponseInputItem | None: + """Remove and return the most recent item from the session. + + Args: + wrapper: Optional RunContextWrapper for accessing context data during pop + + Returns: + The most recent item if it exists, None if the session is empty + """ + popped = await self.underlying_session.pop_item(wrapper) if popped: self._compaction_candidate_items = None self._session_items = None return popped - async def clear_session(self) -> None: - await self.underlying_session.clear_session() + async def clear_session( + self, wrapper: RunContextWrapper[Any] | None = None + ) -> None: + """Clear all items for this session. + + Args: + wrapper: Optional RunContextWrapper for accessing context data during clear + """ + await self.underlying_session.clear_session(wrapper) self._compaction_candidate_items = [] self._session_items = [] self._deferred_response_id = None diff --git a/src/agents/memory/session.py b/src/agents/memory/session.py index 85a65a1690..ef5803f7a2 100644 --- a/src/agents/memory/session.py +++ b/src/agents/memory/session.py @@ -1,12 +1,14 @@ from __future__ import annotations from abc import ABC, abstractmethod -from typing import TYPE_CHECKING, Literal, Protocol, runtime_checkable +from typing import TYPE_CHECKING, Any, Literal, Protocol, runtime_checkable from typing_extensions import TypedDict, TypeGuard if TYPE_CHECKING: from ..items import TResponseInputItem + from ..run_context import RunContextWrapper + from .session_settings import SessionSettings from .session_settings import SessionSettings @@ -21,36 +23,53 @@ class Session(Protocol): session_id: str session_settings: SessionSettings | None = None - async def get_items(self, limit: int | None = None) -> list[TResponseInputItem]: + async def get_items( + self, limit: int | None = None, wrapper: RunContextWrapper[Any] | None = None + ) -> list[TResponseInputItem]: """Retrieve the conversation history for this session. Args: limit: Maximum number of items to retrieve. If None, retrieves all items. When specified, returns the latest N items in chronological order. + wrapper: Optional RunContextWrapper for accessing context data during retrieval. Returns: List of input items representing the conversation history """ ... - async def add_items(self, items: list[TResponseInputItem]) -> None: + async def add_items( + self, items: list[TResponseInputItem], wrapper: RunContextWrapper[Any] | None = None + ) -> None: """Add new items to the conversation history. Args: items: List of input items to add to the history + wrapper: Optional RunContextWrapper for accessing context data during addition """ ... - async def pop_item(self) -> TResponseInputItem | None: + async def pop_item( + self, wrapper: RunContextWrapper[Any] | None = None + ) -> TResponseInputItem | None: """Remove and return the most recent item from the session. + Args: + wrapper: Optional RunContextWrapper for accessing context data during pop + Returns: The most recent item if it exists, None if the session is empty """ ... - async def clear_session(self) -> None: - """Clear all items for this session.""" + async def clear_session( + self, wrapper: RunContextWrapper[Any] | None = None + ) -> None: + """Clear all items for this session. + + Args: + wrapper: Optional RunContextWrapper for accessing context data during clear + """ ... @@ -68,12 +87,15 @@ class SessionABC(ABC): session_settings: SessionSettings | None = None @abstractmethod - async def get_items(self, limit: int | None = None) -> list[TResponseInputItem]: + async def get_items( + self, limit: int | None = None, wrapper: RunContextWrapper[Any] | None = None + ) -> list[TResponseInputItem]: """Retrieve the conversation history for this session. Args: limit: Maximum number of items to retrieve. If None, retrieves all items. When specified, returns the latest N items in chronological order. + wrapper: Optional RunContextWrapper for accessing context data during retrieval. Returns: List of input items representing the conversation history @@ -81,18 +103,26 @@ async def get_items(self, limit: int | None = None) -> list[TResponseInputItem]: ... @abstractmethod - async def add_items(self, items: list[TResponseInputItem]) -> None: + async def add_items( + self, items: list[TResponseInputItem], wrapper: RunContextWrapper[Any] | None = None + ) -> None: """Add new items to the conversation history. Args: items: List of input items to add to the history + wrapper: Optional RunContextWrapper for accessing context data during addition """ ... @abstractmethod - async def pop_item(self) -> TResponseInputItem | None: + async def pop_item( + self, wrapper: RunContextWrapper[Any] | None = None + ) -> TResponseInputItem | None: """Remove and return the most recent item from the session. + Args: + wrapper: Optional RunContextWrapper for accessing context data during pop + Returns: The most recent item if it exists, None if the session is empty """ diff --git a/src/agents/memory/sqlite_session.py b/src/agents/memory/sqlite_session.py index 710993b549..0eced3b732 100644 --- a/src/agents/memory/sqlite_session.py +++ b/src/agents/memory/sqlite_session.py @@ -5,8 +5,10 @@ import sqlite3 import threading from pathlib import Path +from typing import Any from ..items import TResponseInputItem +from ..run_context import RunContextWrapper from .session import SessionABC from .session_settings import SessionSettings, resolve_session_limit @@ -112,12 +114,15 @@ def _init_db_for_connection(self, conn: sqlite3.Connection) -> None: conn.commit() - async def get_items(self, limit: int | None = None) -> list[TResponseInputItem]: + async def get_items( + self, limit: int | None = None, wrapper: RunContextWrapper[Any] | None = None + ) -> list[TResponseInputItem]: """Retrieve the conversation history for this session. Args: limit: Maximum number of items to retrieve. If None, uses session_settings.limit. When specified, returns the latest N items in chronological order. + wrapper: Optional RunContextWrapper for accessing context data during retrieval. Returns: List of input items representing the conversation history @@ -168,11 +173,14 @@ def _get_items_sync(): return await asyncio.to_thread(_get_items_sync) - async def add_items(self, items: list[TResponseInputItem]) -> None: + async def add_items( + self, items: list[TResponseInputItem], wrapper: RunContextWrapper[Any] | None = None + ) -> None: """Add new items to the conversation history. Args: items: List of input items to add to the history + wrapper: Optional RunContextWrapper for accessing context data during addition """ if not items: return @@ -212,9 +220,14 @@ def _add_items_sync(): await asyncio.to_thread(_add_items_sync) - async def pop_item(self) -> TResponseInputItem | None: + async def pop_item( + self, wrapper: RunContextWrapper[Any] | None = None + ) -> TResponseInputItem | None: """Remove and return the most recent item from the session. + Args: + wrapper: Optional RunContextWrapper for accessing context data during pop + Returns: The most recent item if it exists, None if the session is empty """ @@ -253,8 +266,14 @@ def _pop_item_sync(): return await asyncio.to_thread(_pop_item_sync) - async def clear_session(self) -> None: - """Clear all items for this session.""" + async def clear_session( + self, wrapper: RunContextWrapper[Any] | None = None + ) -> None: + """Clear all items for this session. + + Args: + wrapper: Optional RunContextWrapper for accessing context data during clear + """ def _clear_session_sync(): conn = self._get_connection() From 344cf8f7c8a20536b067c9914c0c889b2f09b940 Mon Sep 17 00:00:00 2001 From: Limitless2023 Date: Sat, 14 Feb 2026 23:05:10 +0800 Subject: [PATCH 2/2] fix: only pass wrapper when not None for backward compatibility --- .../openai_responses_compaction_session.py | 24 +++++++++++++++---- 1 file changed, 20 insertions(+), 4 deletions(-) diff --git a/src/agents/memory/openai_responses_compaction_session.py b/src/agents/memory/openai_responses_compaction_session.py index 7e89fb6dd5..b56405f6df 100644 --- a/src/agents/memory/openai_responses_compaction_session.py +++ b/src/agents/memory/openai_responses_compaction_session.py @@ -249,7 +249,11 @@ async def get_items( Returns: List of input items representing the conversation history """ - return await self.underlying_session.get_items(limit, wrapper) + # Only pass wrapper if provided, for backward compatibility with sessions + # that haven't updated their signatures yet + if wrapper is not None: + return await self.underlying_session.get_items(limit, wrapper) + return await self.underlying_session.get_items(limit) async def _defer_compaction(self, response_id: str, store: bool | None = None) -> None: if self._deferred_response_id is not None: @@ -286,7 +290,11 @@ async def add_items( items: List of input items to add to the history wrapper: Optional RunContextWrapper for accessing context data during addition """ - await self.underlying_session.add_items(items, wrapper) + # Only pass wrapper if provided, for backward compatibility + if wrapper is not None: + await self.underlying_session.add_items(items, wrapper) + else: + await self.underlying_session.add_items(items) if self._compaction_candidate_items is not None: new_candidates = select_compaction_candidate_items(items) if new_candidates: @@ -305,7 +313,11 @@ async def pop_item( Returns: The most recent item if it exists, None if the session is empty """ - popped = await self.underlying_session.pop_item(wrapper) + # Only pass wrapper if provided, for backward compatibility + if wrapper is not None: + popped = await self.underlying_session.pop_item(wrapper) + else: + popped = await self.underlying_session.pop_item() if popped: self._compaction_candidate_items = None self._session_items = None @@ -319,7 +331,11 @@ async def clear_session( Args: wrapper: Optional RunContextWrapper for accessing context data during clear """ - await self.underlying_session.clear_session(wrapper) + # Only pass wrapper if provided, for backward compatibility + if wrapper is not None: + await self.underlying_session.clear_session(wrapper) + else: + await self.underlying_session.clear_session() self._compaction_candidate_items = [] self._session_items = [] self._deferred_response_id = None