_compute_state_delta_for_rewind nullifies initial session state keys not present in event state_deltas

[mkbctrl]

Summary

_compute_state_delta_for_rewind nullifies session state keys that were set via create_session(state={...}) but never appear in any event's state_delta. After rewind_async(), these keys are set to None in the session state.

Root Cause

_compute_state_delta_for_rewind (runners.py L615-647) rebuilds state_at_rewind_point by replaying only event state_delta entries:

state_at_rewind_point: dict[str, Any] = {}
for i in range(rewind_event_index):
    if session.events[i].actions.state_delta:
        for k, v in session.events[i].actions.state_delta.items():
            if k.startswith('app:') or k.startswith('user:'):
                continue
            # ...
            state_at_rewind_point[k] = v

Then it nullifies any key in current_state that isn't in state_at_rewind_point:

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  # <-- nullified

Keys set via create_session(state={"my_key": "my_value"}) are present in session.state (via current_state) but never appear in any event's state_delta — they were set as initial state, not through events. The algorithm sees them in current_state but not in state_at_rewind_point, and emits my_key: None.

Only app: and user: prefixed keys are protected from this nullification. Any other initial state key is affected.

Reproduction

from google.adk.sessions import InMemorySessionService
from google.adk.runners import Runner
from google.adk.agents import LlmAgent
from google.adk.events import Event, EventActions

svc = InMemorySessionService()

# Create session with initial state
session = await svc.create_session(
    app_name="test_app",
    user_id="user",
    state={"my_key": "important_value"},  # set via create_session
)

# Simulate an invocation by appending an event
ev = Event(
    invocation_id="inv_A",
    author="model",
    actions=EventActions(state_delta={"other_key": "foo"}),
)
await svc.append_event(session, ev)

# Simulate a second invocation
ev2 = Event(
    invocation_id="inv_B",
    author="model",
    actions=EventActions(state_delta={"other_key": "bar"}),
)
await svc.append_event(session, ev2)

# Rewind before inv_B
agent = LlmAgent(name="test_app", model="gemini-2.5-flash")
runner = Runner(agent=agent, session_service=svc)
await runner.rewind_async(
    user_id="user",
    session_id=session.id,
    rewind_before_invocation_id="inv_B",
)

# Re-fetch session
rewound = await svc.get_session(
    app_name="test_app", user_id="user", session_id=session.id
)

print(rewound.state.get("my_key"))    # None — was "important_value"
print(rewound.state.get("other_key")) # "foo" — correctly restored

Expected Behavior

After rewind_async(), my_key should still be "important_value" — it was set before any invocation and should be unaffected by rewinding invocations.

Actual Behavior

my_key is None after rewind because _compute_state_delta_for_rewind doesn't account for initial session state.

Impact

This affects any integration that sets initial session state via create_session(state={...}). In our case, ag-ui-adk sets _ag_ui_thread_id and _ag_ui_user_id as initial state — after rewind, session lookup by these keys fails because they're nullified.

Suggested Fix

Include initial session state when building state_at_rewind_point. For example, seed it with the session's initial state before replaying event deltas:

async def _compute_state_delta_for_rewind(
    self, session: Session, rewind_event_index: int
) -> dict[str, Any]:
    # Seed with initial state (keys set via create_session)
    state_at_rewind_point: dict[str, Any] = dict(session.initial_state or {})
    for i in range(rewind_event_index):
        # ... existing replay logic

If initial_state is not persisted on the Session object, an alternative is to treat keys present in current_state but absent from all event state_delta entries (across the entire session, not just up to the rewind point) as initial state that should be preserved.

Workaround

After rewind_async(), append a synthetic event that re-sets the nullified keys:

Event(
    invocation_id="",  # empty — skipped by checkpoint enumeration
    author="system",
    actions=EventActions(state_delta=saved_initial_state),
)

Environment

• google-adk: 1.26.x
• ag-ui-adk: 0.5.1
• Python 3.13