Rewind nullifies session state keys set via create_session initial state

[lucasbarzotto-axonify]

🔴 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 events

This 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