google / google/adk-python

_compute_state_delta_for_rewind nullifies initial session state keys not present in event state_deltas

Abierto
#4,839 5 comentarios 0 reacciones 1 asignado Asignado a @Jacksunwei Ver en GitHub
core needs review
Lenguaje dominante
Python
Estrellas
21.5k
Forks
4k
Merge medio
1 d 14 h
PR fusionados (30 d)
37

Descripción

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

Guía de contribución

Abrir la guía de contribución

Evaluación

Este issue todavía no se ha evaluado.

Recibe los nuevos issues en tu correo

Un resumen breve de issues de GitHub para principiantes.