Repository navigation
Rewind nullifies session state keys set via create_session initial state #4933
Description
🔴 Required Information
Is your feature request related to a specific problem?
When using Runner.rewind_async, the rewind logic (_compute_state_delta_for_rewind in runners.py:686–718) reconstructs the target state by replaying state_delta entries from the event stream. However, state keys set during create_session(state={...}) are never recorded as events — they are written directly to session.state. This means the rewind algorithm has no knowledge of these keys.
As a result, after a rewind:
- Keys that were set at session creation but never modified by any event are treated as "added after the rewind point" and get nullified (set to
None) in the rewind state delta. - The session loses critical initial state (e.g. tenant context, user preferences, configuration injected by middleware at session creation time).
This makes rewind_async unsafe for any application that relies on initial session state set via create_session.
Describe the Solution You'd Like
The rewind algorithm should preserve state keys that were never part of the event stream.
Option A — Skip keys absent from the event stream during rewind:
In _compute_state_delta_for_rewind, when diffing current state against the reconstructed state at the rewind point (runners.py:710–717), only null out keys that appeared in some event's state_delta. Keys in session.state that were never touched by any event should be left untouched:
# Current behavior (runners.py:710-717):
for key in current_state:
if key.startswith("app:") or key.startswith("user:"):
continue
if key not in state_at_rewind_point:
rewind_state_delta[key] = None # <-- nullifies initial state keys
# Proposed: also track which keys ever appeared in any event
all_event_keys = set()
for event in session.events:
if event.actions and event.actions.state_delta:
all_event_keys.update(event.actions.state_delta.keys())
for key in current_state:
if key.startswith("app:") or key.startswith("user:"):
continue
if key not in state_at_rewind_point and key in all_event_keys:
rewind_state_delta[key] = None # only null keys that came from eventsThis is backward-compatible, requires no API changes, and naturally handles any key set outside the event stream (whether from create_session, middleware, or direct state manipulation).
Option B — Allow specifying protected state keys via rewind_async:
Add an optional protected_state_keys parameter to rewind_async (and _compute_state_delta_for_rewind) that accepts a set of key names to never include in the rewind delta:
await runner.rewind_async(
user_id="user-1",
session_id="session-1",
rewind_before_invocation_id="inv-3",
protected_state_keys={"tenant_id", "user_id", "subdomain"},
)This gives applications explicit control over which keys are preserved, and works even for keys that do appear in the event stream but should not be rewound (e.g. counters, audit fields).
Both options can coexist — Option A as the default safe behavior, Option B for additional explicit control.
Impact on your work
This is critical for our production application. We inject authentication and tenant context (user ID, tenant ID, subdomain, etc.) into session state via middleware during create_session. After a rewind, these keys are nullified, which breaks all subsequent agent invocations that depend on this context. We currently work around this with a custom rewind implementation that maintains a blocklist of protected keys.
Willingness to contribute
No
🟡 Recommended Information
Describe Alternatives You've Considered
1. Using app: or user: prefixed keys:
The rewind code already skips keys with app: and user: prefixes. We could rename our initial state keys to use these prefixes. However, these prefixes have specific ADK semantics (app-scoped and user-scoped state sharing across sessions) that don't match our use case — we need session-scoped keys that simply aren't part of the event stream.
2. Re-injecting initial state after rewind:
Call middleware again after every rewind to re-set the nullified keys. This adds complexity, couples rewind callers to middleware internals, and creates a race condition window where the session state is in an invalid state.
Additional Context
- ADK version: 1.27.1
Activity
surajksharma07 commentedon Mar 23, 2026
@lucasbarzotto-axonify Reproduced this on v1.27.2! _compute_state_delta_for_rewind rebuilds the target state purely from the event stream, so any keys written directly at session creation are completely invisible to it — and end up getting nullified as a result.
Your Option A looks like the right call here.
Worth trying: collect all_event_keys while iterating through the events, then in the nullification loop change if key not in state_at_rewind_point to if key not in state_at_rewind_point and key in all_event_keys.
That way only keys that have actually appeared in the event stream are eligible for nullification.
Could you give that a try on your end and let us know if it holds up in your production scenario?
lucasbarzotto-axonify commentedon Mar 24, 2026
Hi @surajksharma07, thanks for confirming the reproduction.
I implemented Option A locally and verified it resolves the primary case. Keys set via
create_sessionthat are never touched by events are correctly preserved across rewind.However, I realized that if a key is set at
create_sessionand subsequently modified by an event, rewinding to before that event still nullifies it. The key is present inall_event_keysbut absent fromstate_at_rewind_point, so Option A treats it as eligible for nullification. This is not a regression from the current behavior, but it remains unresolved.Some ideas that could be more robust (but any other approaches that could solve this are definitely welcome):
protected_state_keysparameter onrewind_async). It gives applications explicit control over which keys are preserved, regardless of their presence in the event stream.surajksharma07 commentedon Mar 25, 2026
@lucasbarzotto-axonify you're right that Option A can't protect a key once it's been touched by an event, because all_event_keys makes it eligible for nullification regardless of where it originated.
Of your two suggestions, the synthetic initial event approach seems like the cleaner long-term fix — if create_session emits an event recording the initial state as its state_delta, the rewind algorithm handles it naturally without any API changes or special-casing.
Option B still makes sense as an escape hatch for cases where you need explicit control, but the synthetic event approach would solve the root problem architecturally.
Worth raising that as the recommended direction in a PR.
Team is looking more into it.
surajksharma07 commentedon Mar 25, 2026
@Jacksunwei Please have a look into it once.
fix: fix rewind to preserve initial session state
fix: fix rewind to preserve initial session state
lucasbarzotto-axonify commentedon May 31, 2026
Hi @Jacksunwei @GWeale
I see this has been marked as closed but it looks like the fix has been reverted by this commit af1b00a#diff-9d44028a5e97ca1fe876261e7a0c26660b0d2ad18fa99e44e3814aec8bd61e8b?
I don't see this fix in the new ADK versions and I still can reproduce the bug in version 1.34.0
Should we reopen this issue?
fix: fix rewind to preserve initial session state
fix: fix rewind to preserve initial session state
fix: keep create_session state when rewinding a session