Skip to content

feat: #2072 pass RunContextWrapper to Session get_items/add_items#3591

Open
jawwad-ali wants to merge 1 commit into
openai:mainfrom
jawwad-ali:feat/2072-session-context-wrapper
Open

feat: #2072 pass RunContextWrapper to Session get_items/add_items#3591
jawwad-ali wants to merge 1 commit into
openai:mainfrom
jawwad-ali:feat/2072-session-context-wrapper

Conversation

@jawwad-ali
Copy link
Copy Markdown
Contributor

Summary

Sessions can now opt into receiving the current run's RunContextWrapper[T] by accepting a keyword-only wrapper parameter on get_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:

  • Zero breaking changes for existing sessions. The runner checks the session's method signature (session_method_accepts_wrapper, following the existing is_openai_responses_compaction_aware_session capability 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.
  • Keyword-only, placed last, consistent everywhere (per the review feedback on feat: propagate RunContextWrapper to session history read/write paths #2690): the parameter is *, wrapper: RunContextWrapper[Any] | None = None across the Session protocol, SessionABC, all 10 built-in implementations, and examples/memory/file_session.py.
  • Wrapper-style sessions forward it. EncryptedSession and OpenAIResponsesCompactionSession pass the wrapper through to opted-in underlying sessions (and never to legacy ones).
  • The wrapper is threaded through every runner-driven session read/write: input preparation (including the session_input_callback merge path and SessionSettings limits), turn persistence for run/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_internal layout, and extends coverage to MongoDBSession (added since), the guardrail-trip persistence path, and a larger test suite.

Test plan

  • New tests/memory/test_session_context_wrapper.py with 26 tests: wrapper delivery and wrapper.context identity across Runner.run / run_sync / run_streamed; legacy-session (old signature) backward compatibility through the runner; **kwargs-style sessions; the session_input_callback merge path; the SessionSettings(limit=...) path; guardrail-trip input persistence; EncryptedSession and OpenAIResponsesCompactionSession forwarding (including legacy underlying sessions); a signature audit asserting every built-in session exposes the keyword-only parameter; and unit tests for the capability check.
  • Existing fake sessions in the test suite were updated to the new signature; mock-based call assertions now expect the forwarded kwarg.
  • make format, make lint pass; make typecheck (mypy + pyright) reports no new errors relative to main (verified by diffing full mypy output against a clean main worktree); full test suite run locally with only pre-existing environment-related failures that reproduce identically on clean main.

Issue number

Closes #2072

Notes

  • Scope intentionally matches feat: propagate RunContextWrapper to session history read/write paths #2690: only get_items/add_items receive the wrapper (the runner-driven history read/write paths). pop_item/clear_session and internal rewind/compaction-maintenance reads keep their current signatures and behavior.
  • Static-typing note: third-party sessions subclassing SessionABC with the old signatures will see mypy/pyright [override] hints until they add the optional kwarg; runtime behavior is unchanged either way. Sessions implementing the Session protocol structurally are unaffected at runtime; the runner treats them as not opted in.
  • docs/sessions/index.md gains a short "Accessing the run context in custom sessions" section (English only; translated docs are generated).

Checks

  • I've added new tests (if relevant)
  • I've added/updated the relevant documentation
  • I've run make lint and make format
  • I've made sure tests pass

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
Copy link
Copy Markdown

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 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,
Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge 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 👍 / 👎.

Comment on lines +353 to +356
if wrapper is not None and session_method_accepts_wrapper(
self.underlying_session.add_items
):
await self.underlying_session.add_items(items, wrapper=wrapper)
Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge 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 👍 / 👎.

@seratch seratch added duplicate This issue or pull request already exists feature:sessions labels Jun 7, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

duplicate This issue or pull request already exists feature:sessions

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Pass RunContextWrapper[T] to Session methods

2 participants