google / google/adk-python

_compute_state_delta_for_rewind nullifies initial session state keys not present in event state_deltas

Open
#4,839 5 comments 0 reactions 1 assignee Assigned to @Jacksunwei View on GitHub
core needs review
Dominant language
Python
Stars
21.5k
Forks
4k
Avg merge
1d 14h
Merged PRs (30d)
37

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:

```python
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`:

```python
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

```python
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:

```python
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:

```python
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

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.