feat: pass RunContextWrapper to Session methods#3051
feat: pass RunContextWrapper to Session methods#3051LDLZM wants to merge 3 commits intoopenai:mainfrom
Conversation
Add optional 'wrapper' parameter to Session protocol methods (get_items, add_items, pop_item, clear_session) so custom session implementations can access the run context during history operations. - Update Session protocol and SessionABC in session.py - Update all built-in implementations (SQLiteSession, OpenAIConversationsSession, OpenAIResponsesCompactionSession) - Thread context_wrapper through call sites in session_persistence.py, run.py, and run_loop.py - Update test utilities and fixtures for the new signatures Closes openai#2072
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 5a8fbbec47
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
|
|
||
| if resolved_settings.limit is not None: | ||
| history = await session.get_items(limit=resolved_settings.limit) | ||
| history = await session.get_items(limit=resolved_settings.limit, wrapper=wrapper) |
There was a problem hiding this comment.
Stop breaking existing sessions with
wrapper kwarg
prepare_input_with_session now unconditionally calls session.get_items(..., wrapper=wrapper), but several in-repo SessionABC implementations (e.g. extensions/memory/{redis,sqlalchemy,mongodb,dapr,async_sqlite,encrypt}_session.py and AdvancedSQLiteSession.get_items) still do not accept a wrapper keyword. Using any of those sessions will now fail at runtime with TypeError: ... got an unexpected keyword argument 'wrapper' before the first model call. This change needs to keep compatibility for existing session implementations or update them in the same commit.
Useful? React with 👍 / 👎.
| return saved_run_items_count | ||
|
|
||
| await session.add_items(items_to_save) | ||
| await session.add_items(items_to_save, wrapper=getattr(run_state, "_context", None) if run_state else None) |
There was a problem hiding this comment.
Pass run context when saving without RunState
This write path derives wrapper only from run_state, so in the default non-resumed run flow (where run_state is None) session.add_items always receives wrapper=None even though a context_wrapper was already resolved earlier in the run. Custom sessions that use RunContextWrapper for tenant/user routing will read history with context but persist new items without it, leading to misrouted or rejected writes. The persistence API should accept and forward the active wrapper directly instead of depending on run_state presence.
Useful? React with 👍 / 👎.
… param - Update all 7 extension session implementations to accept wrapper (RedisSession, SQLAlchemySession, MongoDBSession, DaprSession, AsyncSQLiteSession, EncryptedSession, AdvancedSQLiteSession) - Add wrapper parameter directly to save_result_to_session instead of deriving from run_state - Preserve backward compatibility with default None
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 0383b9d12e
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
|
|
||
| history = _normalize_compaction_session_items(await self.underlying_session.get_items()) | ||
| history = _normalize_compaction_session_items( | ||
| await self.underlying_session.get_items() |
There was a problem hiding this comment.
Thread wrapper through compaction candidate loading
OpenAIResponsesCompactionSession now advertises wrapper-aware session methods, but _ensure_compaction_candidates still reads history via underlying_session.get_items() without forwarding any wrapper. For underlying sessions that scope storage by RunContextWrapper (tenant/user routing), compaction decisions will be computed from the wrong context (or fail when wrapper is required), and subsequent compaction can rewrite incorrect history. This path needs to accept and propagate the active wrapper consistently.
Useful? React with 👍 / 👎.
| limit: int | None = None, | ||
| branch_id: str | None = None, | ||
| wrapper: Any = None, |
There was a problem hiding this comment.
Preserve Session positional order in AdvancedSQLite get_items
The new Session.get_items contract adds wrapper as the second optional argument, but this override keeps branch_id in that slot and suppresses the incompatibility with # type: ignore[override]. Any caller using the protocol positionally (e.g. session.get_items(limit, wrapper)) will pass the wrapper object as branch_id, leading to incorrect branch lookup and broken context propagation for this built-in session implementation.
Useful? React with 👍 / 👎.
- Pass wrapper through _ensure_compaction_candidates, _defer_compaction, and run_compaction in OpenAIResponsesCompactionSession - Update session_persistence.py to forward wrapper to compaction calls - Move branch_id to keyword-only in AdvancedSQLiteSession.get_items to avoid positional conflict with new wrapper parameter
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 165e7ab0f7
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| if store is not None: | ||
| compaction_args["store"] = store | ||
| await session.run_compaction(compaction_args) | ||
| await session.run_compaction(compaction_args, wrapper=wrapper) |
There was a problem hiding this comment.
Keep compaction callbacks compatible with existing sessions
save_result_to_session now always invokes session.run_compaction(..., wrapper=wrapper), but the public OpenAIResponsesCompactionAwareSession contract still defines run_compaction(args=None) with no wrapper parameter. Any third-party compaction-aware session that implemented the documented signature will now crash with TypeError: unexpected keyword argument 'wrapper' as soon as compaction is triggered, even when wrapper is None.
Useful? React with 👍 / 👎.
| 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[:]) |
There was a problem hiding this comment.
Key compaction candidate cache by wrapper context
_ensure_compaction_candidates accepts a wrapper but still returns cached candidates/session items without checking which wrapper populated that cache. If one OpenAIResponsesCompactionSession instance is reused across different RunContextWrappers (the new feature this commit enables), the second context can reuse the first context’s cached history, causing compaction decisions and compacted inputs to be computed from the wrong tenant/user data.
Useful? React with 👍 / 👎.
Summary
wrapper: RunContextWrapper[Any] | None = Noneparameter toSessionprotocol andSessionABCmethods (get_items,add_items,pop_item,clear_session)context_wrapperfromrun.py/run_loop.pythroughsession_persistence.pycall sitesMotivation
Custom session implementations currently cannot access the run context during history operations. Passing
RunContextWrapperdirectly enables custom sessions to manage conversation history and context together, improving multi-turn stateful agent workflows.Test Plan
SimpleListSession,DummySession, etc.) and fixtures for the new signaturesmypytype check passesCloses #2072