Skip to content

feat: pass RunContextWrapper to Session methods#3051

Open
LDLZM wants to merge 3 commits intoopenai:mainfrom
LDLZM:feat/pass-context-wrapper-to-session-methods
Open

feat: pass RunContextWrapper to Session methods#3051
LDLZM wants to merge 3 commits intoopenai:mainfrom
LDLZM:feat/pass-context-wrapper-to-session-methods

Conversation

@LDLZM
Copy link
Copy Markdown

@LDLZM LDLZM commented Apr 29, 2026

Summary

  • Add optional wrapper: RunContextWrapper[Any] | None = None parameter to Session protocol and SessionABC methods (get_items, add_items, pop_item, clear_session)
  • Update all built-in session implementations to accept the new parameter
  • Thread context_wrapper from run.py / run_loop.py through session_persistence.py call sites

Motivation

Custom session implementations currently cannot access the run context during history operations. Passing RunContextWrapper directly enables custom sessions to manage conversation history and context together, improving multi-turn stateful agent workflows.

Test Plan

  • All 2409 existing tests pass
  • Updated test utilities (SimpleListSession, DummySession, etc.) and fixtures for the new signatures
  • mypy type check passes

Closes #2072

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
@github-actions github-actions Bot added enhancement New feature or request feature:sessions labels Apr 29, 2026
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: 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)
Copy link
Copy Markdown

Choose a reason for hiding this comment

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

P1 Badge 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)
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 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
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: 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()
Copy link
Copy Markdown

Choose a reason for hiding this comment

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

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

Comment on lines 164 to +166
limit: int | None = None,
branch_id: str | None = None,
wrapper: Any = None,
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 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
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: 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)
Copy link
Copy Markdown

Choose a reason for hiding this comment

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

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

Comment on lines +308 to 312
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[:])
Copy link
Copy Markdown

Choose a reason for hiding this comment

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

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

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

enhancement New feature or request feature:core

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Pass RunContextWrapper[T] to Session methods

1 participant