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

Co-authored-by: George Weale <gweale@google.com>
PiperOrigin-RevId: 905271916

Author: GWeale

src/google/adk/runners.py

@@ -646,6 +646,12 @@ async def rewind_async(
         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:
@@ -686,16 +692,34 @@ async def _compute_state_delta_for_rewind(
       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] = {}
-    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
+    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)
 
     current_state = session.state
     rewind_state_delta = {}
@@ -706,12 +730,13 @@ async def _compute_state_delta_for_rewind(
         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. These keys were added after the
-    #    rewind point and need to be removed.
+    #    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.
     for key in current_state:
       if key.startswith('app:') or key.startswith('user:'):
         continue
-      if key not in state_at_rewind_point:
+      if key not in state_at_rewind_point and key in all_event_keys:
         rewind_state_delta[key] = None
 
     return rewind_state_delta

src/google/adk/sessions/base_session_service.py

@@ -18,10 +18,12 @@
 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
 
@@ -160,3 +162,36 @@ def _update_session_state(self, session: Session, event: Event) -> None:
       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

@@ -417,11 +417,13 @@ async def create_session(
       state: Optional[dict[str, Any]] = None,
       session_id: Optional[str] = None,
   ) -> Session:
-    # 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
+    # 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.
     await self._prepare_tables()
     schema = self._get_schema_classes()
     async with self._rollback_on_exception_session() as sql_session:
@@ -432,6 +434,7 @@ async def create_session(
             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,
@@ -445,19 +448,6 @@ async def create_session(
           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
@@ -468,20 +458,21 @@ async def create_session(
           app_name=app_name,
           user_id=user_id,
           id=session_id,
-          state=session_state,
+          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, session_state
+          storage_app_state.state, storage_user_state.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

@@ -83,12 +83,18 @@ async def create_session(
       state: Optional[dict[str, Any]] = None,
       session_id: Optional[str] = None,
   ) -> Session:
-    return self._create_session_impl(
+    # 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(
         app_name=app_name,
         user_id=user_id,
-        state=state,
+        state=None,
         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

@@ -179,25 +179,9 @@ async def create_session(
               f"Session with id {session_id} already exists."
           )
 
-      # 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
+      # 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`.
       await db.execute(
           """
           INSERT INTO sessions (app_name, user_id, id, state, create_time, update_time)
@@ -207,18 +191,19 @@ async def create_session(
               app_name,
               user_id,
               session_id,
-              json.dumps(session_state),
+              json.dumps({}),
               now,
               now,
           ),
       )
       await db.commit()
 
-      # Merge states for response
-      merged_state = _merge_state(
-          storage_app_state, storage_user_state, session_state
-      )
-      return Session(
+      # 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(
           app_name=app_name,
           user_id=user_id,
           id=session_id,
@@ -227,6 +212,9 @@ async def create_session(
           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

@@ -125,10 +125,13 @@ async def create_session(
     """
     reasoning_engine_id = self._get_reasoning_engine_id(app_name)
 
-    config = {'session_state': state} if state else {}
+    # 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)
     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}',
@@ -143,9 +146,11 @@ async def create_session(
         app_name=app_name,
         user_id=user_id,
         id=session_id,
-        state=getattr(get_session_response, 'session_state', None) or {},
+        state={},
         last_update_time=get_session_response.update_time.timestamp(),
     )
+
+    await self._record_initial_state_event(session, state)
     return session
 
   @override
@@ -213,9 +218,21 @@ async def get_session(
       # 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:
-          session.events.append(_from_api_event(event))
+          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
 
     if config:
       # Filter events based on num_recent_events.