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