feat: #2072 pass RunContextWrapper to Session get_items/add_items#3591
feat: #2072 pass RunContextWrapper to Session get_items/add_items#3591jawwad-ali wants to merge 1 commit into
Conversation
Sessions can now opt into receiving the current run's RunContextWrapper by accepting a keyword-only wrapper parameter on get_items/add_items. The runner inspects the session's signature and only forwards the wrapper to implementations that accept it, so existing custom sessions keep working unchanged. Wrapper-style sessions (EncryptedSession, OpenAIResponsesCompactionSession) forward the wrapper to opted-in underlying sessions. This revives the approach from openai#2690, which addressed maintainer review feedback before going stale, rebased onto the current runner internals and extended to MongoDBSession and the guardrail-trip persistence path. Closes openai#2072
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: cdf4f0317f
ℹ️ 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".
| input_items_to_save, | ||
| [], | ||
| streamed_result._state, | ||
| wrapper=context_wrapper, |
There was a problem hiding this comment.
Thread wrappers through retry rewind
When a wrapper-aware session scopes storage by wrapper.context (as the new docs example suggests), this write persists streamed input into that scoped store, but rewind_model_request just below still calls rewind_session_items(session, ...) without the wrapper. On a transient stream/model retry after _stream_input_persisted is set, the rewind reads/pops the unscoped/default tail instead of the items just written, so the retry can append the same user input again and leave duplicate session history. Please pass the same wrapper through the rewind/cleanup path as well.
Useful? React with 👍 / 👎.
| if wrapper is not None and session_method_accepts_wrapper( | ||
| self.underlying_session.add_items | ||
| ): | ||
| await self.underlying_session.add_items(items, wrapper=wrapper) |
There was a problem hiding this comment.
Thread wrappers through compaction internals
This forwards wrapper to a context-scoped underlying session for normal writes, but run_compaction() later reads candidates and replaces the underlying store via _ensure_compaction_candidates(), _get_all_underlying_session_items(), and _replace_underlying_session_items() without any wrapper. When OpenAIResponsesCompactionSession wraps a session that uses wrapper.context to choose a user/tenant, compaction decisions are made against the unscoped/default history and the replacement can clear/restore the wrong scope instead of the history that was just written. The compaction path needs to carry the current wrapper (or avoid advertising/forwarding wrapper support for this decorator).
Useful? React with 👍 / 👎.
Summary
Sessions can now opt into receiving the current run's
RunContextWrapper[T]by accepting a keyword-onlywrapperparameter onget_items/add_items, so a custom session can manage conversation history and run context together (per-user storage scoping, context-derived metadata, etc.).Design points:
session_method_accepts_wrapper, following the existingis_openai_responses_compaction_aware_sessioncapability pattern) and only forwards the wrapper to implementations that accept it. Sessions with pre-existing signatures never receive the kwarg — covered by dedicated regression tests.*, wrapper: RunContextWrapper[Any] | None = Noneacross theSessionprotocol,SessionABC, all 10 built-in implementations, andexamples/memory/file_session.py.EncryptedSessionandOpenAIResponsesCompactionSessionpass the wrapper through to opted-in underlying sessions (and never to legacy ones).session_input_callbackmerge path andSessionSettingslimits), turn persistence forrun/run_sync/run_streamed, resumed-turn persistence, and guardrail-trip input persistence.This revives #2690 by @HuxleyHu98, which had addressed the review feedback before timing/staleness closed it. The implementation here incorporates that review feedback, is rebased onto the current
run_internallayout, and extends coverage toMongoDBSession(added since), the guardrail-trip persistence path, and a larger test suite.Test plan
tests/memory/test_session_context_wrapper.pywith 26 tests: wrapper delivery andwrapper.contextidentity acrossRunner.run/run_sync/run_streamed; legacy-session (old signature) backward compatibility through the runner;**kwargs-style sessions; thesession_input_callbackmerge path; theSessionSettings(limit=...)path; guardrail-trip input persistence;EncryptedSessionandOpenAIResponsesCompactionSessionforwarding (including legacy underlying sessions); a signature audit asserting every built-in session exposes the keyword-only parameter; and unit tests for the capability check.make format,make lintpass;make typecheck(mypy + pyright) reports no new errors relative tomain(verified by diffing full mypy output against a cleanmainworktree); full test suite run locally with only pre-existing environment-related failures that reproduce identically on cleanmain.Issue number
Closes #2072
Notes
get_items/add_itemsreceive the wrapper (the runner-driven history read/write paths).pop_item/clear_sessionand internal rewind/compaction-maintenance reads keep their current signatures and behavior.SessionABCwith the old signatures will see mypy/pyright[override]hints until they add the optional kwarg; runtime behavior is unchanged either way. Sessions implementing theSessionprotocol structurally are unaffected at runtime; the runner treats them as not opted in.docs/sessions/index.mdgains a short "Accessing the run context in custom sessions" section (English only; translated docs are generated).Checks
make lintandmake format