Repository navigation
_compute_state_delta_for_rewind nullifies initial session state keys not present in event state_deltas #4839
Description
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] = vThen 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 # <-- nullifiedKeys 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 restoredExpected 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 logicIf 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.xag-ui-adk: 0.5.1- Python 3.13
Activity
fix: preserve initial session state keys during rewind Fixes google#4839
surajksharma07 commentedon Mar 16, 2026
@mkbctrl the rewind logic initializes with an empty state_at_rewind_point and only replays event deltas, so any keys set at session creation are never picked up and end up getting nullified.
That explains the behavior you're seeing.
A possible workaround: narrow the nullification to only keys that have actually appeared in at least one event delta throughout the session.
That way, initial state keys like _ag_ui_thread_id would be left untouched after a rewind rather than getting wiped.
Would you mind testing that on your end and letting us know if it holds up for your use case?
Want to make sure it's solid before we move ahead with the PR.
Tracking this through #4845 in the meantime.
weiguangli-io commentedon Mar 16, 2026
Technical Root Cause
In
src/google/adk/runners.py,_compute_state_delta_for_rewind(line 670) initializesstate_at_rewind_pointas an empty dict (line 674), then only replays event deltas:State keys set at session creation (
create_session(state={...})) never appear in any event delta, so they are absent fromstate_at_rewind_point. The nullification loop at line 696-700 then sets them toNone:Suggested Fix
Seed
state_at_rewind_pointwith the session's initial state before replaying event deltas. The initial state can be captured from the session creation parameters (the state passed tocreate_session). Concretely, change line 674 from:to something like:
This requires surfacing the original creation state on the
Sessionobject (or reconstructing it by diffingsession.stateagainst cumulative event deltas). The workaround mentioned above (only nullifying keys that appeared in at least one event delta) would also work but is less precise — it wouldn't correctly rewind a key that was both set at creation and modified by an event delta.Happy to submit a PR with the fix if that would be helpful.
surajksharma07 commentedon Mar 17, 2026
@weiguangli-io thanks for the thorough breakdown!
You're right — seeding from session.initial_state is definitely the cleaner approach, and it sidesteps that edge case where a key exists in the initial state but gets modified later by an event delta.
We're already tracking this through #4845, so no need for a separate PR, but you're more than welcome to drop any additional context or thoughts over there if you'd like.
@mkbctrl once #4845 lands, would you be able to verify on your end that it covers the behavior you've been seeing?
surajksharma07 commentedon Mar 26, 2026
@Jacksunwei Please have a look into it once.
mkbctrl commentedon Jun 25, 2026
@surajksharma07 I missed the notification, I will verify it the moment it lands for sure :) Thanks for tending to it ❤️