diff --git a/src/google/adk/sessions/_rewind_utils.py b/src/google/adk/sessions/_rewind_utils.py index 233b9e4fe9..7c1394bc1c 100644 --- a/src/google/adk/sessions/_rewind_utils.py +++ b/src/google/adk/sessions/_rewind_utils.py @@ -30,6 +30,7 @@ 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 @@ -37,36 +38,58 @@ 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 diff --git a/tests/unittests/runners/test_runner_rewind.py b/tests/unittests/runners/test_runner_rewind.py index a9fe5d3bb5..721b3692fd 100644 --- a/tests/unittests/runners/test_runner_rewind.py +++ b/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.""" diff --git a/tests/unittests/sessions/test_rewind_utils.py b/tests/unittests/sessions/test_rewind_utils.py index 0031866b1d..07275a7750 100644 --- a/tests/unittests/sessions/test_rewind_utils.py +++ b/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=[])