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 4 commits 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".
…action 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.
There was a problem hiding this comment.
💡 Codex Review
Fresh evidence after the update is that the tail is now read with wrapper above, but the actual removal still calls the wrapperless pop_item(). For a session that scopes storage by wrapper.context as the new docs describe, a transient model/stream retry can verify the scoped tail and then pop from the unscoped/default backing store, leaving the just-persisted scoped input in place so the retry appends it again. The pop path needs the same context as the read/write path, or wrapper-scoped sessions cannot be used safely with retry rewind.
ℹ️ 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".
…tionSession The compaction decorator rewrites history by clearing and replacing the underlying store during compaction. clear_session is intentionally not part of the get_items/add_items wrapper contract, so scoping only the reads and adds (but not the clear) would let compaction read/write a context-scoped store while clearing the default scope. Make the decorator consistent instead: it accepts the keyword-only wrapper to stay protocol-compatible but does not forward it to the underlying session, so all of its history operations use the same (default) scope. Document that wrapping a context-scoped session in the compaction decorator is unsupported for run-context scoping, while transparent proxies like EncryptedSession still forward the wrapper.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 327ab24cdc
ℹ️ 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".
The retry-rewind path removes items via pop_item, which is outside the get_items/add_items run-context wrapper contract. Scoping only the tail reads (but not the pop) would let a wrapper-scoped session verify the scoped tail and then pop from the default scope, leaving the scoped items in place. Keep the rewind path unscoped so its reads and pops use the same (default) scope, mirroring the OpenAIResponsesCompactionSession decision. Documented that retry rewind and compaction run against the default scope for wrapper-scoped sessions.
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