fix(sessions): keep create_session state when rewinding a session

src/google/adk/sessions/_rewind_utils.py

@@ -30,43 +30,66 @@
 from ..platform import uuid as platform_uuid
 from ..sessions.base_session_service import BaseSessionService
 from ..sessions.session import Session
+from ..sessions.state import State
 
 if TYPE_CHECKING:
   from ..artifacts.base_artifact_service import BaseArtifactService
 
 logger = logging.getLogger("google_adk." + __name__)
 
 
+def _is_session_scoped_key(key: str) -> bool:
+  """Returns whether a state key is scoped to the session being rewound."""
+  return not key.startswith((State.APP_PREFIX, State.USER_PREFIX))
+
+
 async def compute_state_delta_for_rewind(
     session: Session, rewind_event_index: int
 ) -> dict[str, Any]:
-  """Computes the state delta to reverse changes."""
+  """Computes the state delta that reverses changes made after a rewind point.
+
+  Only session-scoped keys that an event at or after the rewind point changed
+  are reverted. Every other key keeps its current value.
+
+  Args:
+    session: The session to rewind, with the events it was loaded with.
+    rewind_event_index: Index in `session.events` of the first event to undo.
+
+  Returns:
+    A state delta that restores each changed key to its value at the rewind
+    point, or sets it to None when no earlier event recorded a value for it.
+  """
   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
+  for event in session.events[:rewind_event_index]:
+    for key, value in (event.actions.state_delta or {}).items():
+      if not _is_session_scoped_key(key):
+        continue
+      if value is None:
+        state_at_rewind_point.pop(key, None)
+      else:
+        state_at_rewind_point[key] = value
+
+  # Initial state passed to `create_session`, and state from events that were
+  # not loaded, never appears in the event stream being replayed. Deriving the
+  # delta from every key in `session.state` would clear that state, so only
+  # keys that the undone events wrote are considered.
+  keys_changed_after_rewind_point: dict[str, None] = {}
+  for event in session.events[rewind_event_index:]:
+    for key in event.actions.state_delta or {}:
+      if _is_session_scoped_key(key):
+        keys_changed_after_rewind_point[key] = None
 
   current_state = session.state
-  rewind_state_delta = {}
-
-  # 1. Add/update keys in rewind_state_delta to match state_at_rewind_point.
-  for key, value_at_rewind in state_at_rewind_point.items():
-    if key not in current_state or current_state[key] != value_at_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.
-  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: dict[str, Any] = {}
+  for key in keys_changed_after_rewind_point:
+    if key in state_at_rewind_point:
+      value_at_rewind = state_at_rewind_point[key]
+      if key not in current_state or current_state[key] != value_at_rewind:
+        rewind_state_delta[key] = value_at_rewind
+    elif key in current_state:
+      # No replayed event recorded a value before the rewind point. If the key
+      # came from initial state and an undone event overwrote it, the original
+      # value is not stored anywhere, so the key is cleared.
       rewind_state_delta[key] = None
 
   return rewind_state_delta

tests/unittests/runners/test_runner_rewind.py

@@ -19,11 +19,13 @@
 from typing import Union
 
 from google.adk.agents.base_agent import BaseAgent
+from google.adk.agents.run_config import RunConfig
 from google.adk.artifacts.base_artifact_service import ensure_part
 from google.adk.artifacts.in_memory_artifact_service import InMemoryArtifactService
 from google.adk.events.event import Event
 from google.adk.events.event import EventActions
 from google.adk.runners import Runner
+from google.adk.sessions.base_session_service import GetSessionConfig
 from google.adk.sessions.in_memory_session_service import InMemorySessionService
 from google.adk.sessions.session import Session
 from google.genai import types
@@ -281,6 +283,128 @@ async def test_rewind_async_not_first_invocation(self):
         filename="f2",
     ) == types.Part.from_text(text="f2v0")
 
+  @pytest.mark.asyncio
+  async def test_rewind_async_preserves_initial_session_state(self):
+    """Rewind keeps state passed to create_session and undoes later events.
+
+    Setup: session created with initial state; invocation1 sets k1 and
+      invocation2 sets k2.
+    Act: rewind before invocation2.
+    Assert: initial state and k1 are kept, k2 is cleared.
+    """
+    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={"tenant_id": "t1", "user_pref": "dark"},
+    )
+    for invocation_id, state_delta in (
+        ("invocation1", {"k1": "v1"}),
+        ("invocation2", {"k2": "v2"}),
+    ):
+      await runner.session_service.append_event(
+          session=session,
+          event=Event(
+              invocation_id=invocation_id,
+              author="agent",
+              actions=EventActions(state_delta=state_delta),
+          ),
+      )
+
+    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
+    )
+    assert session.state == {
+        "tenant_id": "t1",
+        "user_pref": "dark",
+        "k1": "v1",
+        "k2": None,
+    }
+
+  @pytest.mark.asyncio
+  async def test_rewind_async_before_first_invocation_preserves_initial_state(
+      self,
+  ):
+    """Rewinding before the first invocation keeps create_session 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"},
+    )
+    await runner.session_service.append_event(
+        session=session,
+        event=Event(
+            invocation_id="invocation1",
+            author="agent",
+            actions=EventActions(state_delta={"k1": "v1"}),
+        ),
+    )
+
+    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
+    )
+    assert session.state == {"config_key": "config_val", "k1": None}
+
+  @pytest.mark.asyncio
+  async def test_rewind_async_with_recent_events_config_keeps_earlier_state(
+      self,
+  ):
+    """Rewind loading only recent events keeps state from unloaded events.
+
+    Setup: invocations 1-3 each set their own key.
+    Act: rewind before invocation3 with a session config that loads only the
+      most recent event.
+    Assert: keys from invocations 1 and 2 are kept, k3 is cleared.
+    """
+    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
+    )
+    for i in range(1, 4):
+      await runner.session_service.append_event(
+          session=session,
+          event=Event(
+              invocation_id=f"invocation{i}",
+              author="agent",
+              actions=EventActions(state_delta={f"k{i}": f"v{i}"}),
+          ),
+      )
+
+    await runner.rewind_async(
+        user_id=user_id,
+        session_id=session_id,
+        rewind_before_invocation_id="invocation3",
+        run_config=RunConfig(
+            get_session_config=GetSessionConfig(num_recent_events=1)
+        ),
+    )
+
+    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", "k3": None}
+
 
 class TestRunnerRewindNoFileData:
   """Tests that rewind works with artifact services that reject file_data."""

tests/unittests/sessions/test_rewind_utils.py

@@ -50,6 +50,30 @@ async def test_compute_state_delta_reverts_state_to_rewind_point():
   assert "app:stay" not in delta
 
 
+async def test_compute_state_delta_leaves_state_without_events_untouched():
+  """Keys that no undone event wrote are left out of the rewind delta."""
+  session = Session(
+      id="s1",
+      app_name="app",
+      user_id="u1",
+      state={"tenant_id": "t1", "k1": "v1", "k2": "v2"},
+      events=[
+          Event(
+              invocation_id="inv1",
+              actions=EventActions(state_delta={"k1": "v1"}),
+          ),
+          Event(
+              invocation_id="inv2",
+              actions=EventActions(state_delta={"k2": "v2"}),
+          ),
+      ],
+  )
+
+  delta = await _rewind_utils.compute_state_delta_for_rewind(session, 1)
+
+  assert delta == {"k2": None}
+
+
 async def test_compute_artifact_delta_returns_empty_when_no_artifact_service():
   """Without an artifact service, artifact delta computation returns empty dict."""
   session = Session(id="s1", app_name="app", user_id="u1", events=[])