From af1b00a12b8dd6eee844cc28df7bcd4838e22c1a Mon Sep 17 00:00:00 2001 From: Google Team Member Date: Fri, 24 Apr 2026 19:00:07 -0700 Subject: [PATCH] fix: fix rewind to preserve initial session state 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 --- src/google/adk/runners.py | 49 ++--- .../adk/sessions/base_session_service.py | 35 --- .../adk/sessions/database_session_service.py | 33 +-- .../adk/sessions/in_memory_session_service.py | 10 +- .../adk/sessions/sqlite_session_service.py | 38 ++-- .../adk/sessions/vertex_ai_session_service.py | 25 +-- tests/unittests/runners/test_runner_rewind.py | 199 ------------------ .../test_vertex_ai_session_service.py | 21 +- 8 files changed, 65 insertions(+), 345 deletions(-) diff --git a/src/google/adk/runners.py b/src/google/adk/runners.py index e2393ab6..0f36d638 100644 --- a/src/google/adk/runners.py +++ b/src/google/adk/runners.py @@ -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 diff --git a/src/google/adk/sessions/base_session_service.py b/src/google/adk/sessions/base_session_service.py index c7468bd0..af94bb9e 100644 --- a/src/google/adk/sessions/base_session_service.py +++ b/src/google/adk/sessions/base_session_service.py @@ -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) diff --git a/src/google/adk/sessions/database_session_service.py b/src/google/adk/sessions/database_session_service.py index 16929e9d..d033f1f2 100644 --- a/src/google/adk/sessions/database_session_service.py +++ b/src/google/adk/sessions/database_session_service.py @@ -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 diff --git a/src/google/adk/sessions/in_memory_session_service.py b/src/google/adk/sessions/in_memory_session_service.py index 27daa00f..1bef5160 100644 --- a/src/google/adk/sessions/in_memory_session_service.py +++ b/src/google/adk/sessions/in_memory_session_service.py @@ -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, diff --git a/src/google/adk/sessions/sqlite_session_service.py b/src/google/adk/sessions/sqlite_session_service.py index c00ca197..427bc3e7 100644 --- a/src/google/adk/sessions/sqlite_session_service.py +++ b/src/google/adk/sessions/sqlite_session_service.py @@ -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, diff --git a/src/google/adk/sessions/vertex_ai_session_service.py b/src/google/adk/sessions/vertex_ai_session_service.py index c73a9135..8c1fdc13 100644 --- a/src/google/adk/sessions/vertex_ai_session_service.py +++ b/src/google/adk/sessions/vertex_ai_session_service.py @@ -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. diff --git a/tests/unittests/runners/test_runner_rewind.py b/tests/unittests/runners/test_runner_rewind.py index f95c3bd7..22ba06f4 100644 --- a/tests/unittests/runners/test_runner_rewind.py +++ b/tests/unittests/runners/test_runner_rewind.py @@ -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.""" diff --git a/tests/unittests/sessions/test_vertex_ai_session_service.py b/tests/unittests/sessions/test_vertex_ai_session_service.py index 229b0445..0b265c3c 100644 --- a/tests/unittests/sessions/test_vertex_ai_session_service.py +++ b/tests/unittests/sessions/test_vertex_ai_session_service.py @@ -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