From 23f57daacd65161f4a2aacd7f4e9e179ee645081 Mon Sep 17 00:00:00 2001 From: Kazuhiro Sera Date: Wed, 5 Aug 2026 21:01:56 +0900 Subject: [PATCH 1/3] feat: pass run context to custom sessions --- docs/sessions/index.md | 52 +++ .../extensions/memory/encrypt_session.py | 67 ++- src/agents/memory/session.py | 67 ++- src/agents/run.py | 28 +- .../run_internal/agent_runner_helpers.py | 5 + src/agents/run_internal/run_loop.py | 26 +- .../run_internal/session_persistence.py | 136 ++++-- .../extensions/memory/test_encrypt_session.py | 71 +++- tests/memory/test_session_context_wrapper.py | 386 ++++++++++++++++++ tests/test_agent_runner_streamed.py | 2 + 10 files changed, 792 insertions(+), 48 deletions(-) create mode 100644 tests/memory/test_session_context_wrapper.py diff --git a/docs/sessions/index.md b/docs/sessions/index.md index 95f66172d0..ba4190cae0 100644 --- a/docs/sessions/index.md +++ b/docs/sessions/index.md @@ -684,6 +684,58 @@ result = await Runner.run( ) ``` +### Accessing the run context in a custom session + +A custom session can opt into the active [`RunContextWrapper`][agents.run_context.RunContextWrapper] by adding an optional keyword-only `wrapper` parameter to all four session methods. Existing session implementations do not need to change: the runner only passes this keyword when the complete session implementation accepts it. + +```python +from typing import Any + +from agents import RunContextWrapper +from agents.items import TResponseInputItem +from agents.memory.session import SessionABC + + +class TenantSession(SessionABC): + async def get_items( + self, + limit: int | None = None, + *, + wrapper: RunContextWrapper[Any] | None = None, + ) -> list[TResponseInputItem]: + tenant = wrapper.context.tenant_id if wrapper is not None else "default" + return await self.store.read(tenant=tenant, limit=limit) + + async def add_items( + self, + items: list[TResponseInputItem], + *, + wrapper: RunContextWrapper[Any] | None = None, + ) -> None: + tenant = wrapper.context.tenant_id if wrapper is not None else "default" + await self.store.append(tenant=tenant, items=items) + + async def pop_item( + self, + *, + wrapper: RunContextWrapper[Any] | None = None, + ) -> TResponseInputItem | None: + tenant = wrapper.context.tenant_id if wrapper is not None else "default" + return await self.store.pop(tenant=tenant) + + async def clear_session( + self, + *, + wrapper: RunContextWrapper[Any] | None = None, + ) -> None: + tenant = wrapper.context.tenant_id if wrapper is not None else "default" + await self.store.clear(tenant=tenant) +``` + +When context selects a storage scope, accept `wrapper` consistently on `get_items`, `add_items`, `pop_item`, and `clear_session`. The opt-in is session-level so history loading, persistence, and retry rollback cannot address different stores. Treat the wrapper as run-local state; persist values from `wrapper.context` only when that is an intentional application-level durability decision. + +`EncryptedSession` forwards the wrapper to an underlying session that supports this opt-in. `OpenAIResponsesCompactionSession` does not currently forward it; use the context-aware session directly when compaction state would also need context-specific isolation. + ## Community session implementations The community has developed additional session implementations: diff --git a/src/agents/extensions/memory/encrypt_session.py b/src/agents/extensions/memory/encrypt_session.py index 19ba7a5683..8b2eb18226 100644 --- a/src/agents/extensions/memory/encrypt_session.py +++ b/src/agents/extensions/memory/encrypt_session.py @@ -37,8 +37,9 @@ from typing_extensions import TypedDict from ...items import TResponseInputItem -from ...memory.session import SessionABC +from ...memory.session import SessionABC, _call_session_method, _get_session_wrapper from ...memory.session_settings import SessionSettings, resolve_session_limit +from ...run_context import RunContextWrapper class EncryptedEnvelope(TypedDict): @@ -180,12 +181,25 @@ def _unwrap_valid_items( valid_items.append(item) return valid_items - 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]: + wrapper = _get_session_wrapper(self.underlying_session, wrapper) effective_limit = resolve_session_limit(limit, self.session_settings) if effective_limit is not None and effective_limit > 0: window = effective_limit while True: - encrypted_items = await self.underlying_session.get_items(window) + encrypted_items = cast( + list[TResponseInputItem], + await _call_session_method( + self.underlying_session.get_items, + window, + wrapper=wrapper, + ), + ) valid_items = self._unwrap_valid_items(encrypted_items) if len(valid_items) >= effective_limit: return valid_items[-effective_limit:] @@ -193,21 +207,54 @@ async def get_items(self, limit: int | None = None) -> list[TResponseInputItem]: return valid_items window *= 2 - encrypted_items = await self.underlying_session.get_items(limit) + encrypted_items = cast( + list[TResponseInputItem], + await _call_session_method( + self.underlying_session.get_items, + limit, + wrapper=wrapper, + ), + ) return self._unwrap_valid_items(encrypted_items) - async def add_items(self, items: list[TResponseInputItem]) -> None: + async def add_items( + self, + items: list[TResponseInputItem], + *, + wrapper: RunContextWrapper[Any] | None = None, + ) -> None: + wrapper = _get_session_wrapper(self.underlying_session, wrapper) wrapped: list[EncryptedEnvelope] = [self._wrap(it) for it in items] - await self.underlying_session.add_items(cast(list[TResponseInputItem], wrapped)) + await _call_session_method( + self.underlying_session.add_items, + cast(list[TResponseInputItem], wrapped), + wrapper=wrapper, + ) - async def pop_item(self) -> TResponseInputItem | None: + async def pop_item( + self, + *, + wrapper: RunContextWrapper[Any] | None = None, + ) -> TResponseInputItem | None: + wrapper = _get_session_wrapper(self.underlying_session, wrapper) while True: - enc = await self.underlying_session.pop_item() + enc = await _call_session_method( + self.underlying_session.pop_item, + wrapper=wrapper, + ) if not enc: return None item = self._unwrap(enc) if item is not None: return item - async def clear_session(self) -> None: - await self.underlying_session.clear_session() + async def clear_session( + self, + *, + wrapper: RunContextWrapper[Any] | None = None, + ) -> None: + wrapper = _get_session_wrapper(self.underlying_session, wrapper) + await _call_session_method( + self.underlying_session.clear_session, + wrapper=wrapper, + ) diff --git a/src/agents/memory/session.py b/src/agents/memory/session.py index 1781b7ac9f..3f9b2ddf64 100644 --- a/src/agents/memory/session.py +++ b/src/agents/memory/session.py @@ -1,12 +1,14 @@ from __future__ import annotations +import inspect from abc import ABC, abstractmethod -from typing import TYPE_CHECKING, Literal, Protocol, TypeGuard, runtime_checkable +from typing import TYPE_CHECKING, Any, Literal, Protocol, TypeGuard, runtime_checkable from typing_extensions import TypedDict if TYPE_CHECKING: from ..items import TResponseInputItem + from ..run_context import RunContextWrapper from .session_settings import SessionSettings @@ -148,3 +150,66 @@ def is_openai_responses_compaction_aware_session( except Exception: return False return callable(run_compaction) + + +def _session_method_accepts_wrapper(method: Any) -> bool: + """Return whether a session method opts into receiving ``wrapper``. + + The public ``Session`` protocol keeps its released signatures so existing structural + implementations remain type-compatible. Custom sessions can opt in by adding a keyword + parameter named ``wrapper`` or by accepting arbitrary keyword arguments. + """ + try: + parameters = inspect.signature(method).parameters.values() + except (TypeError, ValueError): + return False + + return any( + parameter.kind is inspect.Parameter.VAR_KEYWORD + or ( + parameter.name == "wrapper" + and parameter.kind + in (inspect.Parameter.POSITIONAL_OR_KEYWORD, inspect.Parameter.KEYWORD_ONLY) + ) + for parameter in parameters + ) + + +def _session_accepts_wrapper(session: Any) -> bool: + """Return whether every history operation accepts ``wrapper``.""" + try: + methods = ( + session.get_items, + session.add_items, + session.pop_item, + session.clear_session, + ) + except Exception: + return False + return all(_session_method_accepts_wrapper(method) for method in methods) + + +def _get_session_wrapper( + session: Any, + wrapper: RunContextWrapper[Any] | None, +) -> RunContextWrapper[Any] | None: + """Return ``wrapper`` only for sessions with a complete context-aware contract.""" + if wrapper is None or not _session_accepts_wrapper(session): + return None + return wrapper + + +async def _call_session_method( + method: Any, + /, + *args: Any, + wrapper: RunContextWrapper[Any] | None = None, + **kwargs: Any, +) -> Any: + """Call a session method with its legacy shape unless it opts into ``wrapper``.""" + if wrapper is not None and _session_method_accepts_wrapper(method): + kwargs["wrapper"] = wrapper + result = method(*args, **kwargs) + if inspect.isawaitable(result): + return await result + return result diff --git a/src/agents/run.py b/src/agents/run.py index 00028cf406..e95bb0d7f7 100644 --- a/src/agents/run.py +++ b/src/agents/run.py @@ -108,6 +108,7 @@ NextStepRunAgain, ) from .run_internal.session_persistence import ( + _session_get_items, persist_session_items_for_guardrail_trip, prepare_input_with_session, reconcile_nested_history_owned_session_item_refs, @@ -525,6 +526,9 @@ async def run( previous_response_id=previous_response_id, auto_previous_response_id=auto_previous_response_id, ) + context_wrapper = ensure_context_wrapper(context) + context = context_wrapper.context + set_agent_tool_state_scope(context_wrapper, None) server_manages_conversation = ( conversation_id is not None @@ -540,6 +544,7 @@ async def run( run_config.session_settings, include_history_in_prepared_input=False, preserve_dropped_new_items=True, + wrapper=context_wrapper, ) original_input_for_state = raw_input session_input_items_for_persistence = [] @@ -552,6 +557,7 @@ async def run( session, run_config.session_input_callback, run_config.session_settings, + wrapper=context_wrapper, ) original_input_for_state = prepared_input @@ -588,7 +594,10 @@ async def run( session_input_items: list[TResponseInputItem] | None = None if session is not None: try: - session_input_items = await session.get_items() + session_input_items = await _session_get_items( + session, + wrapper=context_wrapper, + ) except Exception: session_input_items = None server_conversation_tracker.hydrate_from_state( @@ -646,8 +655,6 @@ async def run( generated_items = [] session_items = [] model_responses = [] - context_wrapper = ensure_context_wrapper(context) - set_agent_tool_state_scope(context_wrapper, None) run_state = RunState( context=context_wrapper, original_input=original_input, @@ -782,6 +789,7 @@ def _finalize_result(result: RunResult) -> RunResult: [], run_state, store=store_setting, + wrapper=context_wrapper, ) session_input_items_for_persistence = [] except BaseException: @@ -825,6 +833,7 @@ def _finalize_result(result: RunResult) -> RunResult: original_user_input, run_state, store=store_setting, + wrapper=context_wrapper, ) ) raise @@ -875,6 +884,7 @@ def _finalize_result(result: RunResult) -> RunResult: [], run_state, store=store_setting, + wrapper=context_wrapper, ) session_input_items_for_persistence = [] if run_state is not None and run_state._current_step is not None: @@ -944,6 +954,7 @@ def _finalize_result(result: RunResult) -> RunResult: run_state._reasoning_item_id_policy ), store=store_setting, + wrapper=context_wrapper, ) ) @@ -1057,6 +1068,7 @@ def _finalize_result(result: RunResult) -> RunResult: run_state, response_id=turn_result.model_response.response_id, store=store_setting, + wrapper=context_wrapper, ) result._original_input = copy_input_items(original_input) return _finalize_result(result) @@ -1191,6 +1203,7 @@ def _finalize_result(result: RunResult) -> RunResult: response_id=None, reasoning_item_id_policy=resolved_reasoning_item_id_policy, store=store_setting, + wrapper=context_wrapper, ) result._original_input = copy_input_items(original_input) return _finalize_result(result) @@ -1245,6 +1258,7 @@ def _finalize_result(result: RunResult) -> RunResult: original_user_input, run_state, store=store_setting, + wrapper=context_wrapper, ) ) raise @@ -1302,6 +1316,7 @@ def _finalize_result(result: RunResult) -> RunResult: original_user_input, run_state, store=store_setting, + wrapper=context_wrapper, ) ) raise @@ -1431,6 +1446,7 @@ def _finalize_result(result: RunResult) -> RunResult: run_state._reasoning_item_id_policy ), store=store_setting, + wrapper=context_wrapper, ) run_state._current_turn_persisted_item_count += saved_count else: @@ -1441,6 +1457,7 @@ def _finalize_result(result: RunResult) -> RunResult: run_state, response_id=turn_result.model_response.response_id, store=store_setting, + wrapper=context_wrapper, ) # After the first resumed turn, treat subsequent turns as fresh @@ -1467,6 +1484,7 @@ def _finalize_result(result: RunResult) -> RunResult: items=_retained_items_for_blocked_output(items_to_save_turn), response_id=turn_result.model_response.response_id, store=store_setting, + wrapper=context_wrapper, ) raise except (Exception, asyncio.CancelledError): @@ -1480,6 +1498,7 @@ def _finalize_result(result: RunResult) -> RunResult: items=items_to_save_turn, response_id=turn_result.model_response.response_id, store=store_setting, + wrapper=context_wrapper, ) raise @@ -1491,6 +1510,7 @@ def _finalize_result(result: RunResult) -> RunResult: items=items_to_save_turn, response_id=turn_result.model_response.response_id, store=store_setting, + wrapper=context_wrapper, ) # Ensure starting_input is not None and not RunState @@ -1539,6 +1559,7 @@ def _finalize_result(result: RunResult) -> RunResult: run_state, response_id=turn_result.model_response.response_id, store=store_setting, + wrapper=context_wrapper, ) append_model_response_if_new( model_responses, turn_result.model_response @@ -1600,6 +1621,7 @@ def _finalize_result(result: RunResult) -> RunResult: items=session_items_for_turn(turn_result), response_id=turn_result.model_response.response_id, store=store_setting, + wrapper=context_wrapper, ) continue else: diff --git a/src/agents/run_internal/agent_runner_helpers.py b/src/agents/run_internal/agent_runner_helpers.py index d380ebe649..348908b79f 100644 --- a/src/agents/run_internal/agent_runner_helpers.py +++ b/src/agents/run_internal/agent_runner_helpers.py @@ -474,6 +474,7 @@ async def save_turn_items_if_needed( items: list[RunItem], response_id: str | None, store: bool | None = None, + wrapper: RunContextWrapper[Any] | None = None, ) -> None: """Persist turn items when persistence is enabled and guardrails allow it.""" if not session_persistence_enabled: @@ -489,6 +490,7 @@ async def save_turn_items_if_needed( run_state, response_id=response_id, store=store, + wrapper=wrapper, ) @@ -501,6 +503,7 @@ async def save_final_turn_items_after_guardrails( items: list[RunItem], response_id: str | None, store: bool | None = None, + wrapper: RunContextWrapper[Any] | None = None, ) -> None: """Persist deferred final-turn items without skipping a partially persisted resumed turn.""" if not session_persistence_enabled or not items: @@ -515,6 +518,7 @@ async def save_final_turn_items_after_guardrails( response_id=response_id, reasoning_item_id_policy=run_state._reasoning_item_id_policy, store=store, + wrapper=wrapper, ) return await save_result_to_session( @@ -524,6 +528,7 @@ async def save_final_turn_items_after_guardrails( run_state, response_id=response_id, store=store, + wrapper=wrapper, ) diff --git a/src/agents/run_internal/run_loop.py b/src/agents/run_internal/run_loop.py index 643238d914..66924bc5fb 100644 --- a/src/agents/run_internal/run_loop.py +++ b/src/agents/run_internal/run_loop.py @@ -157,6 +157,7 @@ ToolRunShellCall, ) from .session_persistence import ( + _session_get_items, persist_session_items_for_guardrail_trip, prepare_input_with_session, reconcile_nested_history_owned_session_item_refs, @@ -367,6 +368,7 @@ async def _save_resumed_stream_items( response_id=response_id, reasoning_item_id_policy=streamed_result._reasoning_item_id_policy, store=store, + wrapper=streamed_result.context_wrapper, ) if run_state is not None: run_state._current_turn_persisted_item_count = ( @@ -398,6 +400,7 @@ async def _save_stream_items( run_state, response_id=response_id, store=store, + wrapper=streamed_result.context_wrapper, ) if update_persisted_count and streamed_result._state is not None: streamed_result._current_turn_persisted_item_count = ( @@ -728,7 +731,10 @@ def _sync_conversation_tracking_from_tracker() -> None: session_items: list[TResponseInputItem] | None = None if session is not None: try: - session_items = await session.get_items() + session_items = await _session_get_items( + session, + wrapper=context_wrapper, + ) except Exception: session_items = None server_conversation_tracker.hydrate_from_state( @@ -768,6 +774,7 @@ def _sync_conversation_tracking_from_tracker() -> None: run_config.session_settings, include_history_in_prepared_input=not server_manages_conversation, preserve_dropped_new_items=True, + wrapper=context_wrapper, ) streamed_result.input = prepared_input streamed_result._original_input = copy_input_items(prepared_input) @@ -871,6 +878,7 @@ async def _save_stream_items_without_count( store=current_agent.model_settings.resolve( run_config.model_settings ).store, + wrapper=context_wrapper, ) ) raise InputGuardrailTripwireTriggered(result) @@ -1180,6 +1188,7 @@ async def _save_stream_items_without_count( store=current_agent.model_settings.resolve( run_config.model_settings ).store, + wrapper=context_wrapper, ) ) raise InputGuardrailTripwireTriggered(result) @@ -1648,7 +1657,13 @@ def _tool_search_fingerprint(raw_item: Any) -> str: ) ] if input_items_to_save: - await save_result_to_session(session, input_items_to_save, [], streamed_result._state) + await save_result_to_session( + session, + input_items_to_save, + [], + streamed_result._state, + wrapper=context_wrapper, + ) previous_response_id = ( server_conversation_tracker.previous_response_id @@ -2121,7 +2136,12 @@ async def get_new_response( async def rewind_model_request() -> None: if server_conversation_tracker is not None: items_to_rewind = session_items_to_rewind if session_items_to_rewind is not None else [] - await rewind_session_items(session, items_to_rewind, server_conversation_tracker) + await rewind_session_items( + session, + items_to_rewind, + server_conversation_tracker, + wrapper=context_wrapper, + ) server_conversation_tracker.rewind_input(filtered.input) with model_run_context(tool_use_tracker): diff --git a/src/agents/run_internal/session_persistence.py b/src/agents/run_internal/session_persistence.py index 3f44e7d2d7..b9dc6449ad 100644 --- a/src/agents/run_internal/session_persistence.py +++ b/src/agents/run_internal/session_persistence.py @@ -29,6 +29,8 @@ is_openai_responses_compaction_aware_session, ) from ..memory.openai_conversations_session import OpenAIConversationsSession +from ..memory.session import _call_session_method, _get_session_wrapper +from ..run_context import RunContextWrapper from ..run_state import RunState from .items import ( NestedHistoryOwnedItem, @@ -65,6 +67,48 @@ ] +_SESSION_LIMIT_UNSET = object() + + +async def _session_get_items( + session: Session, + limit: int | None | object = _SESSION_LIMIT_UNSET, + *, + wrapper: RunContextWrapper[Any] | None = None, +) -> list[TResponseInputItem]: + """Read session items while preserving the legacy method call shape.""" + 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) + + +async def _session_add_items( + session: Session, + items: list[TResponseInputItem], + *, + wrapper: RunContextWrapper[Any] | None = None, +) -> None: + """Append session items while preserving the legacy method call shape.""" + wrapper = _get_session_wrapper(session, wrapper) + await _call_session_method(session.add_items, items, wrapper=wrapper) + + +async def _session_pop_item( + session: Session, + *, + wrapper: RunContextWrapper[Any] | None = None, +) -> TResponseInputItem | None: + """Pop a session item while preserving the legacy method call shape.""" + wrapper = _get_session_wrapper(session, wrapper) + return cast( + TResponseInputItem | None, + await _call_session_method(session.pop_item, wrapper=wrapper), + ) + + def resolve_nested_history_owned_session_item_refs( session_items: Sequence[RunItem], current_input: str | Sequence[TResponseInputItem], @@ -162,6 +206,7 @@ async def prepare_input_with_session( *, include_history_in_prepared_input: bool = True, preserve_dropped_new_items: bool = False, + wrapper: RunContextWrapper[Any] | None = None, ) -> tuple[str | list[TResponseInputItem], list[TResponseInputItem]]: """Prepare model input from session history plus the new turn input. @@ -186,9 +231,13 @@ async def prepare_input_with_session( resolved_settings = resolved_settings.resolve(session_settings) if resolved_settings.limit is not None: - history = await session.get_items(limit=resolved_settings.limit) + history = await _session_get_items( + session, + limit=resolved_settings.limit, + wrapper=wrapper, + ) else: - history = await session.get_items() + history = await _session_get_items(session, wrapper=wrapper) is_openai_conversation_session = isinstance(session, OpenAIConversationsSession) converted_history = [ strip_internal_input_item_metadata(ensure_input_item_format(item)) for item in history @@ -307,6 +356,7 @@ async def persist_session_items_for_guardrail_trip( original_user_input: str | list[TResponseInputItem] | None, run_state: RunState | None, store: bool | None = None, + wrapper: RunContextWrapper[Any] | None = None, ) -> list[TResponseInputItem] | None: """ Persist input items when a guardrail tripwire is triggered. @@ -321,7 +371,14 @@ async def persist_session_items_for_guardrail_trip( input_items_for_save: list[TResponseInputItem] = ( updated_session_input_items if updated_session_input_items is not None else [] ) - await save_result_to_session(session, input_items_for_save, [], run_state, store=store) + await save_result_to_session( + session, + input_items_for_save, + [], + run_state, + store=store, + wrapper=wrapper, + ) return updated_session_input_items @@ -366,6 +423,7 @@ async def save_result_to_session( response_id: str | None = None, reasoning_item_id_policy: ReasoningItemIdPolicy | None = None, store: bool | None = None, + wrapper: RunContextWrapper[Any] | None = None, ) -> int: """ Persist a turn to the session store, keeping track of what was already saved so retries @@ -379,6 +437,8 @@ async def save_result_to_session( if session is None: return 0 + wrapper = _get_session_wrapper(session, wrapper) + new_run_items: list[RunItem] if already_persisted >= len(new_items): new_run_items = [] @@ -459,7 +519,7 @@ async def save_result_to_session( run_state._current_turn_persisted_item_count = already_persisted + saved_run_items_count return saved_run_items_count - await session.add_items(items_to_save) + await _session_add_items(session, items_to_save, wrapper=wrapper) if run_state: run_state._current_turn_persisted_item_count = already_persisted + saved_run_items_count @@ -471,9 +531,12 @@ async def save_result_to_session( if has_local_tool_outputs: defer_compaction = getattr(session, "_defer_compaction", None) if callable(defer_compaction): - result = defer_compaction(response_id, store=store) - if inspect.isawaitable(result): - await result + await _call_session_method( + defer_compaction, + response_id, + store=store, + wrapper=wrapper, + ) logger.debug( "skip: deferring compaction for response %s due to local tool outputs", response_id, @@ -497,7 +560,11 @@ async def save_result_to_session( } if store is not None: compaction_args["store"] = store - await session.run_compaction(compaction_args) + await _call_session_method( + session.run_compaction, + compaction_args, + wrapper=wrapper, + ) return saved_run_items_count @@ -510,6 +577,7 @@ async def save_resumed_turn_items( response_id: str | None, reasoning_item_id_policy: ReasoningItemIdPolicy | None = None, store: bool | None = None, + wrapper: RunContextWrapper[Any] | None = None, ) -> int: """Persist resumed turn items and return the updated persisted count.""" if session is None or not items: @@ -522,6 +590,7 @@ async def save_resumed_turn_items( response_id=response_id, reasoning_item_id_policy=reasoning_item_id_policy, store=store, + wrapper=wrapper, ) return persisted_count + saved_count @@ -530,6 +599,8 @@ async def rewind_session_items( session: Session | None, items: Sequence[TResponseInputItem], server_tracker: OpenAIServerConversationTracker | None = None, + *, + wrapper: RunContextWrapper[Any] | None = None, ) -> None: """ Best-effort helper to roll back items recently persisted to a session when a conversation @@ -538,8 +609,7 @@ async def rewind_session_items( if session is None or not items: return - pop_item = getattr(session, "pop_item", None) - if not callable(pop_item): + if not callable(getattr(session, "pop_item", None)): return ignore_ids_for_matching = _ignore_ids_for_matching(session) @@ -564,13 +634,13 @@ async def rewind_session_items( snapshot_serializations = target_serializations.copy() rewound = await _rewind_session_tail_suffix( session=session, - pop_item=pop_item, expected_serializations=target_serializations, ignore_ids_for_matching=ignore_ids_for_matching, mismatch_warning=( "Skipping session rewind because the current tail does not match the retry-owned suffix" ), pop_failure_warning="Failed to rewind session item", + wrapper=wrapper, ) if not rewound: return @@ -579,13 +649,14 @@ async def rewind_session_items( session, snapshot_serializations, ignore_ids_for_matching=ignore_ids_for_matching, + wrapper=wrapper, ) if session is None or server_tracker is None: return try: - latest_items = await session.get_items(limit=1) + latest_items = await _session_get_items(session, limit=1, wrapper=wrapper) except Exception as exc: log_model_and_tool_action_debug(logger, "Failed to peek session items while rewinding", exc) return @@ -598,7 +669,7 @@ async def rewind_session_items( return try: - session_items = await session.get_items() + session_items = await _session_get_items(session, wrapper=wrapper) except Exception as exc: log_model_and_tool_action_debug( logger, "Failed to inspect session tail while stripping stray items", exc @@ -620,7 +691,6 @@ async def rewind_session_items( ) await _rewind_session_tail_suffix( session=session, - pop_item=pop_item, expected_serializations=stray_serializations, ignore_ids_for_matching=ignore_ids_for_matching, mismatch_warning=( @@ -628,6 +698,7 @@ async def rewind_session_items( "retry-owned conversation items" ), pop_failure_warning="Failed to strip stray session item", + wrapper=wrapper, ) @@ -637,6 +708,7 @@ async def wait_for_session_cleanup( *, max_attempts: int = 5, ignore_ids_for_matching: bool = False, + wrapper: RunContextWrapper[Any] | None = None, ) -> None: """ Confirm that rewound items are no longer present in the session tail so the store stays @@ -649,7 +721,7 @@ async def wait_for_session_cleanup( for attempt in range(max_attempts): try: - tail_items = await session.get_items(limit=window) + tail_items = await _session_get_items(session, limit=window, wrapper=wrapper) except Exception as exc: log_model_and_tool_action_debug( logger, f"Failed to verify session cleanup (attempt {attempt + 1})", exc @@ -771,18 +843,22 @@ def _fingerprint_or_repr(item: TResponseInputItem, *, ignore_ids_for_matching: b async def _rewind_session_tail_suffix( *, session: Session, - pop_item: Any, expected_serializations: Sequence[str], ignore_ids_for_matching: bool, mismatch_warning: str, pop_failure_warning: str, + wrapper: RunContextWrapper[Any] | None = None, ) -> bool: """Remove an exact serialized suffix from the session tail, aborting when the tail diverges.""" if not expected_serializations: return True try: - tail_items = await session.get_items(limit=len(expected_serializations)) + tail_items = await _session_get_items( + session, + limit=len(expected_serializations), + wrapper=wrapper, + ) except Exception as exc: log_model_and_tool_action_warning(logger, pop_failure_warning, exc) return False @@ -806,16 +882,14 @@ async def _rewind_session_tail_suffix( popped_items: list[TResponseInputItem] = [] for expected in reversed(expected_serializations): try: - result = pop_item() - if inspect.isawaitable(result): - result = await result + result = await _session_pop_item(session, wrapper=wrapper) except Exception as exc: - await _restore_popped_session_items(session, popped_items) + await _restore_popped_session_items(session, popped_items, wrapper=wrapper) log_model_and_tool_action_warning(logger, pop_failure_warning, exc) return False if result is None: - await _restore_popped_session_items(session, popped_items) + await _restore_popped_session_items(session, popped_items, wrapper=wrapper) logger.warning(mismatch_warning) return False @@ -824,7 +898,7 @@ async def _rewind_session_tail_suffix( result, ignore_ids_for_matching=ignore_ids_for_matching ) if popped_serialized != expected: - await _restore_popped_session_items(session, popped_items) + await _restore_popped_session_items(session, popped_items, wrapper=wrapper) logger.warning(mismatch_warning) return False @@ -832,20 +906,24 @@ async def _rewind_session_tail_suffix( async def _restore_popped_session_items( - session: Session, popped_items: Sequence[TResponseInputItem] + session: Session, + popped_items: Sequence[TResponseInputItem], + *, + wrapper: RunContextWrapper[Any] | None = None, ) -> None: """Best-effort restoration for items popped during a failed rewind attempt.""" if not popped_items: return - add_items = getattr(session, "add_items", None) - if not callable(add_items): + if not callable(getattr(session, "add_items", None)): return try: - result = add_items(list(reversed(popped_items))) - if inspect.isawaitable(result): - await result + await _session_add_items( + session, + list(reversed(popped_items)), + wrapper=wrapper, + ) except Exception as exc: log_model_and_tool_action_warning( logger, "Failed to restore session items after a rewind mismatch", exc diff --git a/tests/extensions/memory/test_encrypt_session.py b/tests/extensions/memory/test_encrypt_session.py index 71d2bd13b6..fb6da900dd 100644 --- a/tests/extensions/memory/test_encrypt_session.py +++ b/tests/extensions/memory/test_encrypt_session.py @@ -2,7 +2,7 @@ import tempfile from pathlib import Path -from typing import cast +from typing import Any, cast import pytest @@ -10,7 +10,14 @@ from cryptography.fernet import Fernet -from agents import Agent, Runner, SessionSettings, SQLiteSession, TResponseInputItem +from agents import ( + Agent, + RunContextWrapper, + Runner, + SessionSettings, + SQLiteSession, + TResponseInputItem, +) from agents.extensions.memory.encrypt_session import EncryptedSession from tests.fake_model import FakeModel from tests.test_responses import get_text_message @@ -161,6 +168,66 @@ async def test_encrypted_session_clear(encryption_key: str, underlying_session: underlying_session.close() +async def test_encrypted_session_forwards_wrapper_to_all_underlying_operations( + encryption_key: str, +): + class ContextAwareUnderlying: + def __init__(self) -> None: + self.session_id = "test_session" + self.session_settings = None + self.items: list[TResponseInputItem] = [] + self.wrappers: list[RunContextWrapper[Any] | None] = [] + + async def get_items( + self, + limit: int | None = None, + *, + wrapper: RunContextWrapper[Any] | None = None, + ) -> list[TResponseInputItem]: + self.wrappers.append(wrapper) + return list(self.items if limit is None else self.items[-limit:]) + + async def add_items( + self, + items: list[TResponseInputItem], + *, + wrapper: RunContextWrapper[Any] | None = None, + ) -> None: + self.wrappers.append(wrapper) + self.items.extend(items) + + async def pop_item( + self, + *, + wrapper: RunContextWrapper[Any] | None = None, + ) -> TResponseInputItem | None: + self.wrappers.append(wrapper) + return self.items.pop() if self.items else None + + async def clear_session( + self, + *, + wrapper: RunContextWrapper[Any] | None = None, + ) -> None: + self.wrappers.append(wrapper) + self.items.clear() + + underlying = ContextAwareUnderlying() + session = EncryptedSession( + session_id="test_session", + underlying_session=cast(Any, underlying), + encryption_key=encryption_key, + ) + wrapper = RunContextWrapper(context={"tenant": "a"}) + + await session.add_items([{"role": "user", "content": "hello"}], wrapper=wrapper) + assert await session.get_items(wrapper=wrapper) == [{"role": "user", "content": "hello"}] + assert await session.pop_item(wrapper=wrapper) == {"role": "user", "content": "hello"} + await session.clear_session(wrapper=wrapper) + + assert underlying.wrappers == [wrapper, wrapper, wrapper, wrapper] + + async def test_encrypted_session_ttl_expiration( encryption_key: str, underlying_session: SQLiteSession, set_fernet_time ): diff --git a/tests/memory/test_session_context_wrapper.py b/tests/memory/test_session_context_wrapper.py new file mode 100644 index 0000000000..b3f0b06e7d --- /dev/null +++ b/tests/memory/test_session_context_wrapper.py @@ -0,0 +1,386 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, cast + +import pytest + +from agents import Agent, RunContextWrapper, Runner, SessionSettings, TResponseInputItem +from agents.exceptions import InputGuardrailTripwireTriggered +from agents.guardrail import GuardrailFunctionOutput, InputGuardrail +from agents.memory import OpenAIResponsesCompactionSession +from agents.memory.session import _session_accepts_wrapper, _session_method_accepts_wrapper +from agents.run_internal.session_persistence import rewind_session_items +from agents.tool import function_tool +from tests.fake_model import FakeModel +from tests.test_responses import get_function_tool_call, get_text_message + + +@dataclass +class TenantContext: + tenant_id: str + + +class ContextAwareSession: + def __init__(self) -> None: + self.session_id = "context-aware" + self.session_settings: SessionSettings | None = None + self.items_by_scope: dict[str, list[TResponseInputItem]] = {"default": []} + self.calls: list[tuple[str, RunContextWrapper[Any] | None]] = [] + + def _scope(self, wrapper: RunContextWrapper[Any] | None) -> str: + if wrapper is None: + return "default" + context = cast(TenantContext, wrapper.context) + return context.tenant_id + + async def get_items( + self, + limit: int | None = None, + *, + wrapper: RunContextWrapper[Any] | None = None, + ) -> list[TResponseInputItem]: + self.calls.append(("get_items", wrapper)) + items = self.items_by_scope.setdefault(self._scope(wrapper), []) + if limit is None: + return list(items) + return list(items[-limit:]) + + async def add_items( + self, + items: list[TResponseInputItem], + *, + wrapper: RunContextWrapper[Any] | None = None, + ) -> None: + self.calls.append(("add_items", wrapper)) + self.items_by_scope.setdefault(self._scope(wrapper), []).extend(items) + + async def pop_item( + self, + *, + wrapper: RunContextWrapper[Any] | None = None, + ) -> TResponseInputItem | None: + self.calls.append(("pop_item", wrapper)) + items = self.items_by_scope.setdefault(self._scope(wrapper), []) + return items.pop() if items else None + + async def clear_session( + self, + *, + wrapper: RunContextWrapper[Any] | None = None, + ) -> None: + self.calls.append(("clear_session", wrapper)) + self.items_by_scope[self._scope(wrapper)] = [] + + +class LegacySession: + def __init__(self) -> None: + self.session_id = "legacy" + self.session_settings = None + self.items: list[TResponseInputItem] = [] + self.get_calls = 0 + + async def get_items(self) -> list[TResponseInputItem]: + self.get_calls += 1 + return list(self.items) + + async def add_items(self, items: list[TResponseInputItem]) -> None: + self.items.extend(items) + + async def pop_item(self) -> TResponseInputItem | None: + return self.items.pop() if self.items else None + + async def clear_session(self) -> None: + self.items.clear() + + +@pytest.mark.parametrize("streamed", [False, True]) +@pytest.mark.asyncio +async def test_runner_passes_same_wrapper_to_context_aware_session(streamed: bool) -> None: + session = ContextAwareSession() + model = FakeModel(initial_output=[get_text_message("ok")]) + agent = Agent(name="test", model=model) + context = TenantContext(tenant_id="tenant-a") + + if streamed: + result: Any = Runner.run_streamed(agent, "hello", context=context, session=session) + async for _ in result.stream_events(): + pass + else: + result = await Runner.run(agent, "hello", context=context, session=session) + + assert result.final_output == "ok" + assert [name for name, _ in session.calls] == [ + "get_items", + "add_items", + "add_items", + ] + assert all(wrapper is result.context_wrapper for _, wrapper in session.calls) + assert session.items_by_scope["default"] == [] + assert len(session.items_by_scope["tenant-a"]) == 2 + + +@pytest.mark.asyncio +async def test_runner_preserves_legacy_session_call_shapes() -> None: + session = LegacySession() + model = FakeModel(initial_output=[get_text_message("ok")]) + agent = Agent(name="test", model=model) + + result = await Runner.run( + agent, + "hello", + context=TenantContext(tenant_id="tenant-a"), + session=cast(Any, session), + ) + + assert result.final_output == "ok" + assert session.get_calls == 1 + assert len(session.items) == 2 + + +@pytest.mark.asyncio +async def test_runner_does_not_partially_enable_context_aware_session() -> None: + class PartialSession: + def __init__(self) -> None: + self.session_id = "partial" + self.session_settings: SessionSettings | None = None + self.items: list[TResponseInputItem] = [] + self.wrappers: list[RunContextWrapper[Any] | None] = [] + + async def get_items( + self, + limit: int | None = None, + *, + wrapper: RunContextWrapper[Any] | None = None, + ) -> list[TResponseInputItem]: + self.wrappers.append(wrapper) + return list(self.items if limit is None else self.items[-limit:]) + + async def add_items( + self, + items: list[TResponseInputItem], + *, + wrapper: RunContextWrapper[Any] | None = None, + ) -> None: + self.wrappers.append(wrapper) + self.items.extend(items) + + async def pop_item(self) -> TResponseInputItem | None: + return None + + async def clear_session(self) -> None: + pass + + session = PartialSession() + model = FakeModel(initial_output=[get_text_message("ok")]) + + result = await Runner.run( + Agent(name="test", model=model), + "hello", + context=TenantContext(tenant_id="tenant-a"), + session=session, + ) + + assert result.final_output == "ok" + assert session.wrappers == [None, None, None] + assert len(session.items) == 2 + + +@pytest.mark.asyncio +async def test_retry_rewind_uses_same_context_scope_for_reads_pops_and_cleanup() -> None: + session = ContextAwareSession() + wrapper = RunContextWrapper(context=TenantContext(tenant_id="tenant-a")) + items: list[TResponseInputItem] = [ + {"role": "user", "content": "hello"}, + {"role": "assistant", "content": "hi"}, + ] + session.items_by_scope["default"] = [{"role": "user", "content": "keep"}] + await session.add_items(items, wrapper=wrapper) + session.calls.clear() + + await rewind_session_items(session, items, wrapper=wrapper) + + assert session.items_by_scope["tenant-a"] == [] + assert session.items_by_scope["default"] == [{"role": "user", "content": "keep"}] + assert [name for name, _ in session.calls] == [ + "get_items", + "pop_item", + "pop_item", + "get_items", + ] + assert all(call_wrapper is wrapper for _, call_wrapper in session.calls) + + +@pytest.mark.asyncio +async def test_retry_rewind_restores_partial_pops_in_the_same_context_scope() -> None: + class FailingSecondPopSession(ContextAwareSession): + def __init__(self) -> None: + super().__init__() + self.pop_count = 0 + + async def pop_item( + self, + *, + wrapper: RunContextWrapper[Any] | None = None, + ) -> TResponseInputItem | None: + self.calls.append(("pop_item", wrapper)) + self.pop_count += 1 + if self.pop_count == 2: + raise RuntimeError("pop failed") + items = self.items_by_scope.setdefault(self._scope(wrapper), []) + return items.pop() if items else None + + session = FailingSecondPopSession() + wrapper = RunContextWrapper(context=TenantContext(tenant_id="tenant-a")) + items: list[TResponseInputItem] = [ + {"role": "user", "content": "hello"}, + {"role": "assistant", "content": "hi"}, + ] + session.items_by_scope["tenant-a"] = list(items) + session.items_by_scope["default"] = [{"role": "user", "content": "keep"}] + + await rewind_session_items(session, items, wrapper=wrapper) + + assert session.items_by_scope["tenant-a"] == items + assert session.items_by_scope["default"] == [{"role": "user", "content": "keep"}] + assert [name for name, _ in session.calls] == [ + "get_items", + "pop_item", + "pop_item", + "add_items", + ] + assert all(call_wrapper is wrapper for _, call_wrapper in session.calls) + + +@pytest.mark.parametrize("streamed", [False, True]) +@pytest.mark.asyncio +async def test_input_guardrail_persists_in_the_context_scope(streamed: bool) -> None: + def guardrail_function( + _context: RunContextWrapper[Any], _agent: Agent[Any], _input: Any + ) -> GuardrailFunctionOutput: + return GuardrailFunctionOutput(output_info=None, tripwire_triggered=True) + + session = ContextAwareSession() + session.items_by_scope["default"] = [{"role": "user", "content": "keep"}] + context = TenantContext(tenant_id="tenant-a") + agent = Agent( + name="test", + model=FakeModel(initial_output=[get_text_message("not persisted")]), + input_guardrails=[InputGuardrail(guardrail_function=guardrail_function)], + ) + + with pytest.raises(InputGuardrailTripwireTriggered): + if streamed: + result = Runner.run_streamed(agent, "hello", context=context, session=session) + async for _ in result.stream_events(): + pass + else: + await Runner.run(agent, "hello", context=context, session=session) + + assert session.items_by_scope["default"] == [{"role": "user", "content": "keep"}] + assert session.items_by_scope["tenant-a"] == [{"role": "user", "content": "hello"}] + assert all(wrapper is not None and wrapper.context is context for _, wrapper in session.calls) + + +@pytest.mark.parametrize("streamed", [False, True]) +@pytest.mark.asyncio +async def test_resumed_run_persists_in_the_context_scope(streamed: bool) -> None: + async def test_tool() -> str: + return "tool result" + + tool = function_tool(test_tool, name_override="test_tool", needs_approval=True) + model = FakeModel() + model.add_multiple_turn_outputs( + [ + [get_function_tool_call("test_tool", "{}", call_id="call-resume")], + [get_text_message("done")], + ] + ) + agent = Agent(name="test", model=model, tools=[tool]) + session = ContextAwareSession() + session.items_by_scope["default"] = [{"role": "user", "content": "keep"}] + context = TenantContext(tenant_id="tenant-a") + + if streamed: + first: Any = Runner.run_streamed(agent, "hello", context=context, session=session) + async for _ in first.stream_events(): + pass + else: + first = await Runner.run(agent, "hello", context=context, session=session) + + assert len(first.interruptions) == 1 + state = first.to_state() + state.approve(first.interruptions[0]) + session.calls.clear() + + if streamed: + resumed: Any = Runner.run_streamed(agent, state, session=session) + async for _ in resumed.stream_events(): + pass + else: + resumed = await Runner.run(agent, state, session=session) + + assert resumed.final_output == "done" + assert session.items_by_scope["default"] == [{"role": "user", "content": "keep"}] + assert all(wrapper is resumed.context_wrapper for _, wrapper in session.calls) + assert any( + isinstance(item, dict) + and item.get("type") == "function_call_output" + and item.get("call_id") == "call-resume" + for item in session.items_by_scope["tenant-a"] + ) + + +@pytest.mark.asyncio +async def test_compaction_session_keeps_context_aware_underlying_on_legacy_scope() -> None: + underlying = ContextAwareSession() + underlying.items_by_scope["default"] = [{"role": "user", "content": "existing"}] + session = OpenAIResponsesCompactionSession( + session_id="compaction", + underlying_session=underlying, + should_trigger_compaction=lambda _: False, + ) + + result = await Runner.run( + Agent(name="test", model=FakeModel(initial_output=[get_text_message("done")])), + "hello", + context=TenantContext(tenant_id="tenant-a"), + session=session, + ) + + assert result.final_output == "done" + assert not _session_accepts_wrapper(session) + assert "tenant-a" not in underlying.items_by_scope + assert len(underlying.items_by_scope["default"]) == 3 + assert underlying.calls + assert all(wrapper is None for _, wrapper in underlying.calls) + + +def test_session_wrapper_method_requires_named_keyword_or_kwargs() -> None: + class Methods: + async def legacy(self) -> None: + pass + + async def positional_only(self, wrapper: Any, /) -> None: + pass + + async def keyword(self, *, wrapper: Any = None) -> None: + pass + + async def kwargs(self, **kwargs: Any) -> None: + pass + + methods = Methods() + + assert not _session_method_accepts_wrapper(methods.legacy) + assert not _session_method_accepts_wrapper(methods.positional_only) + assert _session_method_accepts_wrapper(methods.keyword) + assert _session_method_accepts_wrapper(methods.kwargs) + + +def test_session_wrapper_opt_in_requires_all_history_operations() -> None: + session = ContextAwareSession() + assert _session_accepts_wrapper(session) + + cast(Any, session).clear_session = LegacySession().clear_session + assert not _session_accepts_wrapper(session) diff --git a/tests/test_agent_runner_streamed.py b/tests/test_agent_runner_streamed.py index a908eced94..2a3c605817 100644 --- a/tests/test_agent_runner_streamed.py +++ b/tests/test_agent_runner_streamed.py @@ -2615,6 +2615,7 @@ async def save_wrapper( response_id: str | None, reasoning_item_id_policy: str | None = None, store: bool | None = None, + wrapper: RunContextWrapper[Any] | None = None, ) -> int: observed_counts.append(persisted_count) result = await real_save_resumed( @@ -2624,6 +2625,7 @@ async def save_wrapper( response_id=response_id, reasoning_item_id_policy=reasoning_item_id_policy, store=store, + wrapper=wrapper, ) return int(result) From ab0b7bf78343b12d3bed8097d87be8e176fb9ac8 Mon Sep 17 00:00:00 2001 From: Kazuhiro Sera Date: Wed, 5 Aug 2026 21:06:29 +0900 Subject: [PATCH 2/3] docs: remove session wrapper documentation --- docs/sessions/index.md | 52 ------------------------------------------ 1 file changed, 52 deletions(-) diff --git a/docs/sessions/index.md b/docs/sessions/index.md index ba4190cae0..95f66172d0 100644 --- a/docs/sessions/index.md +++ b/docs/sessions/index.md @@ -684,58 +684,6 @@ result = await Runner.run( ) ``` -### Accessing the run context in a custom session - -A custom session can opt into the active [`RunContextWrapper`][agents.run_context.RunContextWrapper] by adding an optional keyword-only `wrapper` parameter to all four session methods. Existing session implementations do not need to change: the runner only passes this keyword when the complete session implementation accepts it. - -```python -from typing import Any - -from agents import RunContextWrapper -from agents.items import TResponseInputItem -from agents.memory.session import SessionABC - - -class TenantSession(SessionABC): - async def get_items( - self, - limit: int | None = None, - *, - wrapper: RunContextWrapper[Any] | None = None, - ) -> list[TResponseInputItem]: - tenant = wrapper.context.tenant_id if wrapper is not None else "default" - return await self.store.read(tenant=tenant, limit=limit) - - async def add_items( - self, - items: list[TResponseInputItem], - *, - wrapper: RunContextWrapper[Any] | None = None, - ) -> None: - tenant = wrapper.context.tenant_id if wrapper is not None else "default" - await self.store.append(tenant=tenant, items=items) - - async def pop_item( - self, - *, - wrapper: RunContextWrapper[Any] | None = None, - ) -> TResponseInputItem | None: - tenant = wrapper.context.tenant_id if wrapper is not None else "default" - return await self.store.pop(tenant=tenant) - - async def clear_session( - self, - *, - wrapper: RunContextWrapper[Any] | None = None, - ) -> None: - tenant = wrapper.context.tenant_id if wrapper is not None else "default" - await self.store.clear(tenant=tenant) -``` - -When context selects a storage scope, accept `wrapper` consistently on `get_items`, `add_items`, `pop_item`, and `clear_session`. The opt-in is session-level so history loading, persistence, and retry rollback cannot address different stores. Treat the wrapper as run-local state; persist values from `wrapper.context` only when that is an intentional application-level durability decision. - -`EncryptedSession` forwards the wrapper to an underlying session that supports this opt-in. `OpenAIResponsesCompactionSession` does not currently forward it; use the context-aware session directly when compaction state would also need context-specific isolation. - ## Community session implementations The community has developed additional session implementations: From 3874fa473b45af04f2c42fc320cf0b45e700d884 Mon Sep 17 00:00:00 2001 From: Kazuhiro Sera Date: Wed, 5 Aug 2026 22:15:20 +0900 Subject: [PATCH 3/3] fix: require explicit session wrapper opt-in --- src/agents/memory/session.py | 15 ++-- tests/memory/test_session_context_wrapper.py | 77 +++++++++++++++++++- 2 files changed, 81 insertions(+), 11 deletions(-) diff --git a/src/agents/memory/session.py b/src/agents/memory/session.py index 3f9b2ddf64..26690c2c71 100644 --- a/src/agents/memory/session.py +++ b/src/agents/memory/session.py @@ -156,21 +156,18 @@ def _session_method_accepts_wrapper(method: Any) -> bool: """Return whether a session method opts into receiving ``wrapper``. The public ``Session`` protocol keeps its released signatures so existing structural - implementations remain type-compatible. Custom sessions can opt in by adding a keyword - parameter named ``wrapper`` or by accepting arbitrary keyword arguments. + implementations remain type-compatible. Custom sessions can opt in by adding a ``wrapper`` + parameter that can be passed by keyword. """ try: parameters = inspect.signature(method).parameters.values() - except (TypeError, ValueError): + except Exception: return False return any( - parameter.kind is inspect.Parameter.VAR_KEYWORD - or ( - parameter.name == "wrapper" - and parameter.kind - in (inspect.Parameter.POSITIONAL_OR_KEYWORD, inspect.Parameter.KEYWORD_ONLY) - ) + parameter.name == "wrapper" + and parameter.kind + in (inspect.Parameter.POSITIONAL_OR_KEYWORD, inspect.Parameter.KEYWORD_ONLY) for parameter in parameters ) diff --git a/tests/memory/test_session_context_wrapper.py b/tests/memory/test_session_context_wrapper.py index b3f0b06e7d..0a7cf65f24 100644 --- a/tests/memory/test_session_context_wrapper.py +++ b/tests/memory/test_session_context_wrapper.py @@ -94,6 +94,44 @@ async def clear_session(self) -> None: self.items.clear() +class LegacyKwargsSession: + def __init__(self) -> None: + self.session_id = "legacy-kwargs" + self.session_settings: SessionSettings | None = None + self.items: list[TResponseInputItem] = [] + self.kwargs_calls: list[dict[str, Any]] = [] + + async def get_items(self, limit: int | None = None, **kwargs: Any) -> list[TResponseInputItem]: + self.kwargs_calls.append(kwargs) + if limit is None: + return list(self.items) + return list(self.items[-limit:]) + + async def add_items(self, items: list[TResponseInputItem], **kwargs: Any) -> None: + self.kwargs_calls.append(kwargs) + self.items.extend(items) + + async def pop_item(self, **kwargs: Any) -> TResponseInputItem | None: + self.kwargs_calls.append(kwargs) + return self.items.pop() if self.items else None + + async def clear_session(self, **kwargs: Any) -> None: + self.kwargs_calls.append(kwargs) + self.items.clear() + + +class UninspectableAsyncMethod: + def __init__(self, method: Any) -> None: + self.method = method + + @property + def __signature__(self) -> Any: + raise RuntimeError("signature unavailable") + + async def __call__(self, *args: Any, **kwargs: Any) -> Any: + return await self.method(*args, **kwargs) + + @pytest.mark.parametrize("streamed", [False, True]) @pytest.mark.asyncio async def test_runner_passes_same_wrapper_to_context_aware_session(streamed: bool) -> None: @@ -138,6 +176,41 @@ async def test_runner_preserves_legacy_session_call_shapes() -> None: assert len(session.items) == 2 +@pytest.mark.asyncio +async def test_runner_does_not_treat_legacy_kwargs_as_wrapper_opt_in() -> None: + session = LegacyKwargsSession() + model = FakeModel(initial_output=[get_text_message("ok")]) + + result = await Runner.run( + Agent(name="test", model=model), + "hello", + context=TenantContext(tenant_id="tenant-a"), + session=session, + ) + + assert result.final_output == "ok" + assert session.kwargs_calls == [{}, {}, {}] + assert len(session.items) == 2 + + +@pytest.mark.asyncio +async def test_runner_preserves_legacy_calls_when_signature_inspection_fails() -> None: + session = cast(Any, LegacySession()) + session.get_items = UninspectableAsyncMethod(session.get_items) + model = FakeModel(initial_output=[get_text_message("ok")]) + + result = await Runner.run( + Agent(name="test", model=model), + "hello", + context=TenantContext(tenant_id="tenant-a"), + session=session, + ) + + assert result.final_output == "ok" + assert session.get_calls == 1 + assert len(session.items) == 2 + + @pytest.mark.asyncio async def test_runner_does_not_partially_enable_context_aware_session() -> None: class PartialSession: @@ -356,7 +429,7 @@ async def test_compaction_session_keeps_context_aware_underlying_on_legacy_scope assert all(wrapper is None for _, wrapper in underlying.calls) -def test_session_wrapper_method_requires_named_keyword_or_kwargs() -> None: +def test_session_wrapper_method_requires_named_wrapper_parameter() -> None: class Methods: async def legacy(self) -> None: pass @@ -375,7 +448,7 @@ async def kwargs(self, **kwargs: Any) -> None: assert not _session_method_accepts_wrapper(methods.legacy) assert not _session_method_accepts_wrapper(methods.positional_only) assert _session_method_accepts_wrapper(methods.keyword) - assert _session_method_accepts_wrapper(methods.kwargs) + assert not _session_method_accepts_wrapper(methods.kwargs) def test_session_wrapper_opt_in_requires_all_history_operations() -> None: