github-featured/google--adk-python
2
0
Fork 0
Code Issues Pull Requests Actions 8 Deployments Models Agent Notes Packages Projects Releases Wiki Activity

fix: fix rewind to preserve initial session state

Browse Source 
The rewind logic is updated to ensure that state keys set during session creation are not nullified when rewinding. Previously, any key not present in the state at the rewind point was removed. Now, only keys that have appeared in any event's state delta are considered for nullification during a rewind, preventing the removal of initial session state

Close #4933

PiperOrigin-RevId: 905322038
This commit is contained in:
Google Team Member
2026-04-24 19:00:07 -07:00
committed by Copybara-Service
parent c65dd5580f
commit af1b00a12b
8 changed files with 65 additions and 345 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/google/adk/runners.py
+12 -37
View File
@@ -646,12 +646,6 @@ class Runner:
session_id=session_id,
get_session_config=run_config.get_session_config,
)
if not rewind_before_invocation_id:
# Guard against matching the synthetic initial-state event that is
# appended by `create_session`; that event has an empty invocation_id by
# design and is not a valid rewind target.
raise ValueError('rewind_before_invocation_id must be non-empty.')
rewind_event_index = -1
for i, event in enumerate(session.events):
if event.invocation_id == rewind_before_invocation_id:
@@ -692,34 +686,16 @@ class Runner:
self, session: Session, rewind_event_index: int
) -> dict[str, Any]:
"""Computes the state delta to reverse changes."""
# State at the rewind point is reconstructed entirely from the event
# stream. Session-scoped initial state from `create_session` is captured
# as a synthetic event by `BaseSessionService._record_initial_state_event`,
# so walking events naturally restores initial values even when a later
# event overwrote them.
state_at_rewind_point: dict[str, Any] = {}
all_event_keys: set[str] = set()
for event in session.events[:rewind_event_index]:
if not event.actions.state_delta:
continue
for k, v in event.actions.state_delta.items():
if k.startswith('app:') or k.startswith('user:'):
continue
all_event_keys.add(k)
if v is None:
state_at_rewind_point.pop(k, None)
else:
state_at_rewind_point[k] = v
# Collect any other keys touched by events after the rewind point so we
# know which keys were ever event-sourced.
for event in session.events[rewind_event_index:]:
if not event.actions.state_delta:
continue
for k in event.actions.state_delta:
if not k.startswith('app:') and not k.startswith('user:'):
all_event_keys.add(k)
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
if v is None:
state_at_rewind_point.pop(k, None)
else:
state_at_rewind_point[k] = v
current_state = session.state
rewind_state_delta = {}
@@ -730,13 +706,12 @@ class Runner:
rewind_state_delta[key] = value_at_rewind
# 2. Set keys to None in rewind_state_delta if they are in current_state
# but not in state_at_rewind_point. Only nullify keys that were
# introduced or modified through events; keys set outside the event
# stream are preserved.
# but not in state_at_rewind_point. These keys were added after the
# rewind point and need to be removed.
for key in current_state:
if key.startswith('app:') or key.startswith('user:'):
continue
if key not in state_at_rewind_point and key in all_event_keys:
if key not in state_at_rewind_point:
rewind_state_delta[key] = None
return rewind_state_delta
src/google/adk/sessions/base_session_service.py
-35
View File
@@ -18,12 +18,10 @@ import abc
from typing import Any
from typing import Optional
from google.adk.platform import time as platform_time
from pydantic import BaseModel
from pydantic import Field
from ..events.event import Event
from ..events.event_actions import EventActions
from .session import Session
from .state import State
@@ -162,36 +160,3 @@ class BaseSessionService(abc.ABC):
return
for key, value in event.actions.state_delta.items():
session.state.update({key: value})
async def _record_initial_state_event(
self, session: Session, state: Optional[dict[str, Any]]
) -> None:
"""Appends a synthetic event carrying the initial non-temp session state.
Subclasses call this from `create_session` so that initial state flows
through `append_event` (the single state-merging path) and so that
`rewind_async` can restore session-scoped initial values for keys later
overwritten or introduced by subsequent events.
Args:
session: The newly created session to attach the event to.
state: The initial state dict supplied to `create_session`. Temp-prefixed
keys are dropped because temp state is ephemeral and never persisted.
"""
if not state:
return
state_delta = {
k: v for k, v in state.items() if not k.startswith(State.TEMP_PREFIX)
}
if not state_delta:
return
# Round to microseconds so the timestamp roundtrips exactly through
# storage backends that persist timestamps as datetime (microsecond
# precision) — keeps in-memory and reloaded events comparable.
timestamp = round(platform_time.get_time(), 6)
initial_event = Event(
author='user',
timestamp=timestamp,
actions=EventActions(state_delta=dict(state_delta)),
)
await self.append_event(session=session, event=initial_event)
src/google/adk/sessions/database_session_service.py
+21 -12
View File
@@ -417,13 +417,11 @@ class DatabaseSessionService(BaseSessionService):
state: Optional[dict[str, Any]] = None,
session_id: Optional[str] = None,
) -> Session:
# 1. Ensure app/user state rows exist (append_event requires them) and
# insert an empty session row.
# 2. Build the in-memory session reflecting any pre-existing app/user
# state.
# 3. Apply the caller-supplied initial state through the synthetic event
# in `_record_initial_state_event` so all state writes share a single
# code path.
# 1. Populate states.
# 2. Build storage session object
# 3. Add the object to the table
# 4. Build the session object with generated id
# 5. Return the session
await self._prepare_tables()
schema = self._get_schema_classes()
async with self._rollback_on_exception_session() as sql_session:
@@ -434,7 +432,6 @@ class DatabaseSessionService(BaseSessionService):
f"Session with id {session_id} already exists."
)
# Get or create state rows, handling concurrent insert races.
# `append_event` requires the app/user state rows to exist.
storage_app_state = await _get_or_create_state(
sql_session=sql_session,
state_model=schema.StorageAppState,
@@ -448,6 +445,19 @@ class DatabaseSessionService(BaseSessionService):
defaults={"app_name": app_name, "user_id": user_id, "state": {}},
)
# Extract state deltas
state_deltas = _session_util.extract_state_delta(state)
app_state_delta = state_deltas["app"]
user_state_delta = state_deltas["user"]
session_state = state_deltas["session"]
# Apply state delta
if app_state_delta:
storage_app_state.state = storage_app_state.state | app_state_delta
if user_state_delta:
storage_user_state.state = storage_user_state.state | user_state_delta
# Store the session
now = datetime.fromtimestamp(platform_time.get_time(), tz=timezone.utc)
is_sqlite = self.db_engine.dialect.name == _SQLITE_DIALECT
is_postgresql = self.db_engine.dialect.name == _POSTGRESQL_DIALECT
@@ -458,21 +468,20 @@ class DatabaseSessionService(BaseSessionService):
app_name=app_name,
user_id=user_id,
id=session_id,
state={},
state=session_state,
create_time=now,
update_time=now,
)
sql_session.add(storage_session)
await sql_session.commit()
# Merge states for response
merged_state = _merge_state(
storage_app_state.state, storage_user_state.state, {}
storage_app_state.state, storage_user_state.state, session_state
)
session = storage_session.to_session(
state=merged_state, is_sqlite=is_sqlite
)
await self._record_initial_state_event(session, state)
return session
@override
src/google/adk/sessions/in_memory_session_service.py
+2 -8
View File
@@ -83,18 +83,12 @@ class InMemorySessionService(BaseSessionService):
state: Optional[dict[str, Any]] = None,
session_id: Optional[str] = None,
) -> Session:
# Initial state flows through `_record_initial_state_event` ->
# `append_event` so the in-memory dicts and the event stream are written
# exactly once. The deprecated `create_session_sync` keeps the legacy
# direct-write path because it cannot await `append_event`.
session = self._create_session_impl(
return self._create_session_impl(
app_name=app_name,
user_id=user_id,
state=None,
state=state,
session_id=session_id,
)
await self._record_initial_state_event(session, state)
return session
def create_session_sync(
self,
src/google/adk/sessions/sqlite_session_service.py
+25 -13
View File
@@ -179,9 +179,25 @@ class SqliteSessionService(BaseSessionService):
f"Session with id {session_id} already exists."
)
# Insert the session row with empty per-session state. Initial state
# (including app:/user:-prefixed keys) is applied through the synthetic
# event below so that all state writes go through `append_event`.
# Extract state deltas
state_deltas = _session_util.extract_state_delta(state)
app_state_delta = state_deltas["app"]
user_state_delta = state_deltas["user"]
session_state = state_deltas["session"]
# Apply state delta and update/insert states atomically
if app_state_delta:
await self._upsert_app_state(db, app_name, app_state_delta, now)
if user_state_delta:
await self._upsert_user_state(
db, app_name, user_id, user_state_delta, now
)
# Fetch current state after upserts
storage_app_state = await self._get_app_state(db, app_name)
storage_user_state = await self._get_user_state(db, app_name, user_id)
# Store the session
await db.execute(
"""
INSERT INTO sessions (app_name, user_id, id, state, create_time, update_time)
@@ -191,19 +207,18 @@ class SqliteSessionService(BaseSessionService):
app_name,
user_id,
session_id,
json.dumps({}),
json.dumps(session_state),
now,
now,
),
)
await db.commit()
# Reflect already-persisted app/user state so subsequent appends start
# from the correct merged view.
storage_app_state = await self._get_app_state(db, app_name)
storage_user_state = await self._get_user_state(db, app_name, user_id)
merged_state = _merge_state(storage_app_state, storage_user_state, {})
session = Session(
# Merge states for response
merged_state = _merge_state(
storage_app_state, storage_user_state, session_state
)
return Session(
app_name=app_name,
user_id=user_id,
id=session_id,
@@ -212,9 +227,6 @@ class SqliteSessionService(BaseSessionService):
last_update_time=now,
)
await self._record_initial_state_event(session, state)
return session
@override
async def get_session(
self,
src/google/adk/sessions/vertex_ai_session_service.py
+4 -21
View File
@@ -125,13 +125,10 @@ class VertexAiSessionService(BaseSessionService):
"""
reasoning_engine_id = self._get_reasoning_engine_id(app_name)
# Initial state is persisted exclusively through the synthetic event
# below (which is sent via `events.append`); avoid passing it as
# `session_state` here so the same data is not written to the backend
# twice.
config = dict(kwargs)
config = {'session_state': state} if state else {}
if session_id:
config['session_id'] = session_id
config.update(kwargs)
async with self._get_api_client() as api_client:
api_response = await api_client.agent_engines.sessions.create(
name=f'reasoningEngines/{reasoning_engine_id}',
@@ -146,11 +143,9 @@ class VertexAiSessionService(BaseSessionService):
app_name=app_name,
user_id=user_id,
id=session_id,
state={},
state=getattr(get_session_response, 'session_state', None) or {},
last_update_time=get_session_response.update_time.timestamp(),
)
await self._record_initial_state_event(session, state)
return session
@override
@@ -218,21 +213,9 @@ class VertexAiSessionService(BaseSessionService):
# to discard events written milliseconds after the session resource was
# updated. Clock skew between those writes can otherwise drop tool_result
# events and permanently break the replayed conversation.
#
# Apply each event's state_delta as we go so callers see the same state
# whether or not the backend mirrors it onto the session_state field
# (e.g. Vertex stores initial state via the synthetic create_session
# event rather than the session_state field).
if events_iterator is not None:
async for event in events_iterator:
adk_event = _from_api_event(event)
session.events.append(adk_event)
if adk_event.actions and adk_event.actions.state_delta:
for key, value in adk_event.actions.state_delta.items():
if value is None:
session.state.pop(key, None)
else:
session.state[key] = value
session.events.append(_from_api_event(event))
if config:
# Filter events based on num_recent_events.
tests/unittests/runners/test_runner_rewind.py
-199
View File
@@ -280,205 +280,6 @@ class TestRunnerRewind:
filename="f2",
) == types.Part.from_text(text="f2v0")
@pytest.mark.asyncio
async def test_rewind_preserves_initial_session_state(self):
"""Tests that rewind preserves state keys set via create_session."""
runner = self.runner
user_id = "test_user"
session_id = "test_session"
# Create session with initial state (not from events).
session = await runner.session_service.create_session(
app_name=runner.app_name,
user_id=user_id,
session_id=session_id,
state={"tenant_id": "t1", "user_pref": "dark"},
)
# invocation1: sets k1 via event
event1 = Event(
invocation_id="invocation1",
author="agent",
content=types.Content(parts=[types.Part.from_text(text="event1")]),
actions=EventActions(state_delta={"k1": "v1"}),
)
await runner.session_service.append_event(session=session, event=event1)
# invocation2: sets k2 via event
event2 = Event(
invocation_id="invocation2",
author="agent",
content=types.Content(parts=[types.Part.from_text(text="event2")]),
actions=EventActions(state_delta={"k2": "v2"}),
)
await runner.session_service.append_event(session=session, event=event2)
session = await runner.session_service.get_session(
app_name=runner.app_name, user_id=user_id, session_id=session_id
)
assert session.state == {
"tenant_id": "t1",
"user_pref": "dark",
"k1": "v1",
"k2": "v2",
}
# Rewind before invocation2
await runner.rewind_async(
user_id=user_id,
session_id=session_id,
rewind_before_invocation_id="invocation2",
)
session = await runner.session_service.get_session(
app_name=runner.app_name, user_id=user_id, session_id=session_id
)
# Initial state keys must be preserved (not nullified).
assert session.state["tenant_id"] == "t1"
assert session.state["user_pref"] == "dark"
# k1 was set in invocation1 (before rewind point), should be preserved.
assert session.state["k1"] == "v1"
# k2 was set in invocation2 (after rewind point), should be nullified.
assert not session.state["k2"]
@pytest.mark.asyncio
async def test_rewind_before_first_invocation_preserves_initial_state(self):
"""Tests rewind to before first invocation preserves initial state."""
runner = self.runner
user_id = "test_user"
session_id = "test_session"
session = await runner.session_service.create_session(
app_name=runner.app_name,
user_id=user_id,
session_id=session_id,
state={"config_key": "config_val"},
)
# Single invocation that adds event state
event1 = Event(
invocation_id="invocation1",
author="agent",
content=types.Content(parts=[types.Part.from_text(text="event1")]),
actions=EventActions(state_delta={"k1": "v1"}),
)
await runner.session_service.append_event(session=session, event=event1)
# Rewind before the only invocation
await runner.rewind_async(
user_id=user_id,
session_id=session_id,
rewind_before_invocation_id="invocation1",
)
session = await runner.session_service.get_session(
app_name=runner.app_name, user_id=user_id, session_id=session_id
)
# Initial state must survive.
assert session.state["config_key"] == "config_val"
# Event-sourced key must be nullified.
assert not session.state["k1"]
@pytest.mark.asyncio
async def test_rewind_restores_initial_value_when_event_overwrote_it(self):
"""Tests rewind restores initial value for keys overwritten by events.
When a key is set in the initial session state and later modified by an
event, rewinding to before that event must restore the initial value
rather than nullifying the key.
"""
runner = self.runner
user_id = "test_user"
session_id = "test_session"
# Create session with an initial state value for tenant_id.
session = await runner.session_service.create_session(
app_name=runner.app_name,
user_id=user_id,
session_id=session_id,
state={"tenant_id": "t_initial"},
)
# invocation1: overwrites tenant_id and introduces a brand-new key.
event1 = Event(
invocation_id="invocation1",
author="agent",
content=types.Content(parts=[types.Part.from_text(text="event1")]),
actions=EventActions(
state_delta={"tenant_id": "t_event", "k_new": "v_new"}
),
)
await runner.session_service.append_event(session=session, event=event1)
session = await runner.session_service.get_session(
app_name=runner.app_name, user_id=user_id, session_id=session_id
)
assert session.state == {"tenant_id": "t_event", "k_new": "v_new"}
# Rewind before invocation1.
await runner.rewind_async(
user_id=user_id,
session_id=session_id,
rewind_before_invocation_id="invocation1",
)
session = await runner.session_service.get_session(
app_name=runner.app_name, user_id=user_id, session_id=session_id
)
# tenant_id must be restored to its initial value, not nullified.
assert session.state["tenant_id"] == "t_initial"
# k_new was introduced by the rewound event and must be nullified.
assert not session.state.get("k_new")
@pytest.mark.asyncio
async def test_rewind_with_no_initial_session_state(self):
"""Rewind works correctly for sessions created without initial state."""
runner = self.runner
user_id = "test_user"
session_id = "test_session"
# Create session with no initial state.
session = await runner.session_service.create_session(
app_name=runner.app_name,
user_id=user_id,
session_id=session_id,
)
event1 = Event(
invocation_id="invocation1",
author="agent",
content=types.Content(parts=[types.Part.from_text(text="event1")]),
actions=EventActions(state_delta={"k1": "v1"}),
)
await runner.session_service.append_event(session=session, event=event1)
event2 = Event(
invocation_id="invocation2",
author="agent",
content=types.Content(parts=[types.Part.from_text(text="event2")]),
actions=EventActions(state_delta={"k2": "v2"}),
)
await runner.session_service.append_event(session=session, event=event2)
session = await runner.session_service.get_session(
app_name=runner.app_name, user_id=user_id, session_id=session_id
)
assert session.state == {"k1": "v1", "k2": "v2"}
await runner.rewind_async(
user_id=user_id,
session_id=session_id,
rewind_before_invocation_id="invocation2",
)
session = await runner.session_service.get_session(
app_name=runner.app_name, user_id=user_id, session_id=session_id
)
# k1 set before the rewind point survives; k2 from the rewound event is
# nullified.
assert session.state["k1"] == "v1"
assert not session.state.get("k2")
class TestRunnerRewindNoFileData:
"""Tests that rewind works with artifact services that reject file_data."""
tests/unittests/sessions/test_vertex_ai_session_service.py
+1 -20
View File
@@ -921,29 +921,10 @@ async def test_create_session():
assert session.user_id == 'user'
assert session.last_update_time is not None
# Session-scoped initial state is recorded as a synthetic event so that
# `rewind_async` can restore those values for keys later overwritten or
# introduced by subsequent events.
assert len(session.events) == 1
assert session.events[0].author == 'user'
assert not session.events[0].invocation_id
assert session.events[0].actions.state_delta == state
session_id = session.id
got_session = await session_service.get_session(
assert session == await session_service.get_session(
app_name='123', user_id='user', session_id=session_id
)
assert got_session.id == session.id
assert got_session.app_name == session.app_name
assert got_session.user_id == session.user_id
assert got_session.state == session.state
# The retrieved event has a server-assigned id and the session's
# last_update_time advances when an event is appended; otherwise the
# round-tripped session matches the synthetic initial-state event.
assert len(got_session.events) == 1
assert got_session.events[0].author == 'user'
assert not got_session.events[0].invocation_id
assert got_session.events[0].actions.state_delta == state
@pytest.mark.asyncio