_compute_state_delta_for_rewind nullifies initial session state keys not present in event state_deltas
- Lingua principale
- Python
- Stelle
- 21.5k
- Fork
- 4k
- Merge medio
- 1g 14h
- PR unite (30g)
- 37
Descrizione
## 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
Guida per i contributori
Apri la guida per i contributori
Valutazione
Questa issue non è ancora stata valutata.