fix: thread run context wrapper through session retry-rewind and compaction paths

Addresses review feedback: a wrapper-aware session that scopes storage
by wrapper.context was getting the wrapper on normal reads/writes but
not on the conversation-retry rewind path or the compaction decorator's
internal history reads/replacements, so those operated on the
unscoped/default store.

- Thread the wrapper through rewind_session_items, the tail-suffix
  rewind, cleanup verification, and popped-item restoration.
- Thread the wrapper through OpenAIResponsesCompactionSession's
  run_compaction, candidate loading, and replace/restore helpers, and
  add it to the run_compaction protocol method.

clear_session/pop_item keep their existing signatures, matching the
get_items/add_items-only scope of this change.

Author: jawwad-ali

src/agents/memory/openai_responses_compaction_session.py

@@ -159,7 +159,12 @@ def _resolve_compaction_mode_for_response(
             return "input"
         return _resolve_compaction_mode(mode, response_id=response_id, store=store)
 
-    async def run_compaction(self, args: OpenAIResponsesCompactionArgs | None = None) -> None:
+    async def run_compaction(
+        self,
+        args: OpenAIResponsesCompactionArgs | None = None,
+        *,
+        wrapper: RunContextWrapper[Any] | None = None,
+    ) -> None:
         """Run compaction using responses.compact API."""
         if args and args.get("response_id"):
             self._response_id = args["response_id"]
@@ -184,7 +189,9 @@ async def run_compaction(self, args: OpenAIResponsesCompactionArgs | None = None
                 "when using previous_response_id compaction."
             )
 
-        compaction_candidate_items, session_items = await self._ensure_compaction_candidates()
+        compaction_candidate_items, session_items = await self._ensure_compaction_candidates(
+            wrapper=wrapper
+        )
 
         force = args.get("force", False) if args else False
         should_compact = force or self.should_trigger_compaction(
@@ -220,10 +227,11 @@ async def run_compaction(self, args: OpenAIResponsesCompactionArgs | None = None
             _normalize_compaction_output_items(compacted.output or [])
         )
 
-        previous_items = await self._get_all_underlying_session_items()
+        previous_items = await self._get_all_underlying_session_items(wrapper=wrapper)
         await self._replace_underlying_session_items(
             output_items=output_items,
             previous_items=previous_items,
+            wrapper=wrapper,
         )
 
         self._compaction_candidate_items = select_compaction_candidate_items(output_items)
@@ -235,7 +243,7 @@ async def run_compaction(self, args: OpenAIResponsesCompactionArgs | None = None
             f"candidates={len(self._compaction_candidate_items)})"
         )
 
-    async def get_items(
+    async def _underlying_get_items(
         self,
         limit: int | None = None,
         *,
@@ -249,37 +257,65 @@ async def get_items(
             return await self.underlying_session.get_items(limit, wrapper=wrapper)
         return await self.underlying_session.get_items(limit)
 
-    async def _get_all_underlying_session_items(self) -> list[TResponseInputItem]:
-        return await self.underlying_session.get_items(limit=_ALL_SESSION_ITEMS_LIMIT)
+    async def _underlying_add_items(
+        self,
+        items: list[TResponseInputItem],
+        *,
+        wrapper: RunContextWrapper[Any] | None = None,
+    ) -> None:
+        if wrapper is not None and session_method_accepts_wrapper(
+            self.underlying_session.add_items
+        ):
+            await self.underlying_session.add_items(items, wrapper=wrapper)
+            return
+        await self.underlying_session.add_items(items)
+
+    async def get_items(
+        self,
+        limit: int | None = None,
+        *,
+        wrapper: RunContextWrapper[Any] | None = None,
+    ) -> list[TResponseInputItem]:
+        return await self._underlying_get_items(limit, wrapper=wrapper)
+
+    async def _get_all_underlying_session_items(
+        self, *, wrapper: RunContextWrapper[Any] | None = None
+    ) -> list[TResponseInputItem]:
+        return await self._underlying_get_items(_ALL_SESSION_ITEMS_LIMIT, wrapper=wrapper)
 
     async def _replace_underlying_session_items(
         self,
         *,
         output_items: list[TResponseInputItem],
         previous_items: list[TResponseInputItem],
+        wrapper: RunContextWrapper[Any] | None = None,
     ) -> None:
         try:
             await self.underlying_session.clear_session()
         except Exception as clear_error:
             await self._restore_underlying_session_items_after_failed_clear(
-                previous_items, clear_error
+                previous_items, clear_error, wrapper=wrapper
             )
             raise
 
         try:
             if output_items:
-                await self.underlying_session.add_items(output_items)
+                await self._underlying_add_items(output_items, wrapper=wrapper)
         except Exception as replacement_error:
-            await self._restore_underlying_session_items(previous_items, replacement_error)
+            await self._restore_underlying_session_items(
+                previous_items, replacement_error, wrapper=wrapper
+            )
             raise
 
     async def _restore_underlying_session_items_after_failed_clear(
         self,
         previous_items: list[TResponseInputItem],
         clear_error: Exception,
+        *,
+        wrapper: RunContextWrapper[Any] | None = None,
     ) -> None:
         try:
-            current_items = await self._get_all_underlying_session_items()
+            current_items = await self._get_all_underlying_session_items(wrapper=wrapper)
         except Exception:
             logger.warning(
                 "Failed to inspect session history after compaction replacement clear failed.",
@@ -291,7 +327,7 @@ async def _restore_underlying_session_items_after_failed_clear(
             return
 
         await self._restore_underlying_session_items(
-            previous_items, clear_error, clear_existing_items=False
+            previous_items, clear_error, clear_existing_items=False, wrapper=wrapper
         )
 
     async def _restore_underlying_session_items(
@@ -300,12 +336,13 @@ async def _restore_underlying_session_items(
         replacement_error: Exception,
         *,
         clear_existing_items: bool = True,
+        wrapper: RunContextWrapper[Any] | None = None,
     ) -> None:
         try:
             if clear_existing_items:
                 await self.underlying_session.clear_session()
             if previous_items:
-                await self.underlying_session.add_items(list(previous_items))
+                await self._underlying_add_items(list(previous_items), wrapper=wrapper)
         except Exception:
             logger.warning(
                 "Failed to restore session history after compaction replacement failed.",
@@ -318,10 +355,18 @@ async def _restore_underlying_session_items(
             replacement_error,
         )
 
-    async def _defer_compaction(self, response_id: str, store: bool | None = None) -> None:
+    async def _defer_compaction(
+        self,
+        response_id: str,
+        store: bool | None = None,
+        *,
+        wrapper: RunContextWrapper[Any] | None = None,
+    ) -> None:
         if self._deferred_response_id is not None:
             return
-        compaction_candidate_items, session_items = await self._ensure_compaction_candidates()
+        compaction_candidate_items, session_items = await self._ensure_compaction_candidates(
+            wrapper=wrapper
+        )
         resolved_mode = self._resolve_compaction_mode_for_response(
             response_id=response_id,
             store=store,
@@ -350,12 +395,7 @@ async def add_items(
         *,
         wrapper: RunContextWrapper[Any] | None = None,
     ) -> None:
-        if wrapper is not None and session_method_accepts_wrapper(
-            self.underlying_session.add_items
-        ):
-            await self.underlying_session.add_items(items, wrapper=wrapper)
-        else:
-            await self.underlying_session.add_items(items)
+        await self._underlying_add_items(items, wrapper=wrapper)
         if self._compaction_candidate_items is not None:
             new_items = _normalize_compaction_session_items(items)
             new_candidates = select_compaction_candidate_items(new_items)
@@ -379,12 +419,16 @@ async def clear_session(self) -> None:
 
     async def _ensure_compaction_candidates(
         self,
+        *,
+        wrapper: RunContextWrapper[Any] | None = None,
     ) -> tuple[list[TResponseInputItem], list[TResponseInputItem]]:
         """Lazy-load and cache compaction candidates."""
         if self._compaction_candidate_items is not None and self._session_items is not None:
             return (self._compaction_candidate_items[:], self._session_items[:])
 
-        history = _normalize_compaction_session_items(await self.underlying_session.get_items())
+        history = _normalize_compaction_session_items(
+            await self._underlying_get_items(wrapper=wrapper)
+        )
         candidates = select_compaction_candidate_items(history)
         self._compaction_candidate_items = candidates
         self._session_items = history

src/agents/memory/session.py

@@ -162,8 +162,19 @@ class OpenAIResponsesCompactionArgs(TypedDict, total=False):
 class OpenAIResponsesCompactionAwareSession(Session, Protocol):
     """Protocol for session implementations that support responses compaction."""
 
-    async def run_compaction(self, args: OpenAIResponsesCompactionArgs | None = None) -> None:
-        """Run the compaction process for the session."""
+    async def run_compaction(
+        self,
+        args: OpenAIResponsesCompactionArgs | None = None,
+        *,
+        wrapper: RunContextWrapper[Any] | None = None,
+    ) -> None:
+        """Run the compaction process for the session.
+
+        Args:
+            args: Optional compaction arguments.
+            wrapper: Optional run context wrapper for the current run. Implementations
+                may ignore this parameter.
+        """
         ...
 
 

src/agents/run_internal/run_loop.py

@@ -1463,7 +1463,9 @@ def _tool_search_fingerprint(raw_item: Any) -> str:
 
     async def rewind_model_request() -> 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
+        )
         if server_conversation_tracker is not None:
             server_conversation_tracker.rewind_input(filtered.input)
 
@@ -1887,7 +1889,9 @@ async def get_new_response(
 
     async def rewind_model_request() -> 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
+        )
         if server_conversation_tracker is not None:
             server_conversation_tracker.rewind_input(filtered.input)
 

src/agents/run_internal/session_persistence.py

@@ -398,7 +398,10 @@ 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 session_method_accepts_wrapper(defer_compaction):
+                    result = defer_compaction(response_id, store=store, wrapper=wrapper)
+                else:
+                    result = defer_compaction(response_id, store=store)
                 if inspect.isawaitable(result):
                     await result
             logger.debug(
@@ -424,7 +427,10 @@ async def save_result_to_session(
         }
         if store is not None:
             compaction_args["store"] = store
-        await session.run_compaction(compaction_args)
+        if session_method_accepts_wrapper(session.run_compaction):
+            await session.run_compaction(compaction_args, wrapper=wrapper)
+        else:
+            await session.run_compaction(compaction_args)
 
     return saved_run_items_count
 
@@ -459,6 +465,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
@@ -499,6 +507,7 @@ async def rewind_session_items(
             "Skipping session rewind because the current tail does not match the retry-owned suffix"
         ),
         pop_failure_warning="Failed to rewind session item: %s",
+        wrapper=wrapper,
     )
     if not rewound:
         return
@@ -507,13 +516,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:
         logger.debug("Failed to peek session items while rewinding: %s", exc)
         return
@@ -526,7 +536,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:
         logger.debug("Failed to inspect session tail while stripping stray items: %s", exc)
         return
@@ -554,6 +564,7 @@ async def rewind_session_items(
             "retry-owned conversation items"
         ),
         pop_failure_warning="Failed to strip stray session item: %s",
+        wrapper=wrapper,
     )
 
 
@@ -563,6 +574,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
@@ -575,7 +587,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:
             logger.debug("Failed to verify session cleanup (attempt %d): %s", attempt + 1, exc)
             await asyncio.sleep(0.1 * (attempt + 1))
@@ -700,13 +712,16 @@ async def _rewind_session_tail_suffix(
     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:
         logger.warning(pop_failure_warning, exc)
         return False
@@ -734,12 +749,12 @@ async def _rewind_session_tail_suffix(
             if inspect.isawaitable(result):
                 result = await result
         except Exception as exc:
-            await _restore_popped_session_items(session, popped_items)
+            await _restore_popped_session_items(session, popped_items, wrapper=wrapper)
             logger.warning(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
 
@@ -748,15 +763,18 @@ 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
 
     return True
 
 
 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:
@@ -767,9 +785,7 @@ async def _restore_popped_session_items(
         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:
         logger.warning("Failed to restore session items after a rewind mismatch: %s", exc)