chore: Refactor and fix state management in the session service

Also refactoring the test cases to focus on the expected behaviors

PiperOrigin-RevId: 820734484

Author: DeanChensj

src/google/adk/sessions/_session_util.py

@@ -19,6 +19,8 @@
 from typing import Type
 from typing import TypeVar
 
+from .state import State
+
 M = TypeVar("M")
 
 
@@ -29,3 +31,19 @@ def decode_model(
   if data is None:
     return None
   return model_cls.model_validate(data)
+
+
+def extract_state_delta(
+    state: dict[str, Any],
+) -> dict[str, dict[str, Any]]:
+  """Extracts app, user, and session state deltas from a state dictionary."""
+  deltas = {"app": {}, "user": {}, "session": {}}
+  if state:
+    for key in state.keys():
+      if key.startswith(State.APP_PREFIX):
+        deltas["app"][key.removeprefix(State.APP_PREFIX)] = state[key]
+      elif key.startswith(State.USER_PREFIX):
+        deltas["user"][key.removeprefix(State.USER_PREFIX)] = state[key]
+      elif not key.startswith(State.TEMP_PREFIX):
+        deltas["session"][key] = state[key]
+  return deltas

src/google/adk/sessions/database_session_service.py

@@ -465,40 +465,31 @@ async def create_session(
     # 5. Return the session
 
     with self.database_session_factory() as sql_session:
-
       # Fetch app and user states from storage
       storage_app_state = sql_session.get(StorageAppState, (app_name))
-      storage_user_state = sql_session.get(
-          StorageUserState, (app_name, user_id)
-      )
-
-      app_state = storage_app_state.state if storage_app_state else {}
-      user_state = storage_user_state.state if storage_user_state else {}
-
-      # Create state tables if not exist
       if not storage_app_state:
         storage_app_state = StorageAppState(app_name=app_name, state={})
         sql_session.add(storage_app_state)
+      storage_user_state = sql_session.get(
+          StorageUserState, (app_name, user_id)
+      )
       if not storage_user_state:
         storage_user_state = StorageUserState(
             app_name=app_name, user_id=user_id, state={}
         )
         sql_session.add(storage_user_state)
 
       # Extract state deltas
-      app_state_delta, user_state_delta, session_state = _extract_state_delta(
-          state
-      )
+      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
-      app_state.update(app_state_delta)
-      user_state.update(user_state_delta)
-
-      # Store app and user state
       if app_state_delta:
-        storage_app_state.state = app_state
+        storage_app_state.state = storage_app_state.state | app_state_delta
       if user_state_delta:
-        storage_user_state.state = user_state
+        storage_user_state.state = storage_user_state.state | user_state_delta
 
       # Store the session
       storage_session = StorageSession(
@@ -513,7 +504,9 @@ async def create_session(
       sql_session.refresh(storage_session)
 
       # Merge states for response
-      merged_state = _merge_state(app_state, user_state, session_state)
+      merged_state = _merge_state(
+          storage_app_state.state, storage_user_state.state, session_state
+      )
       session = storage_session.to_session(state=merged_state)
     return session
 
@@ -536,19 +529,18 @@ async def get_session(
       if storage_session is None:
         return None
 
+      query = sql_session.query(StorageEvent).filter(
+          StorageEvent.app_name == app_name,
+          StorageEvent.user_id == user_id,
+          StorageEvent.session_id == storage_session.id,
+      )
+
       if config and config.after_timestamp:
         after_dt = datetime.fromtimestamp(config.after_timestamp)
-        timestamp_filter = StorageEvent.timestamp >= after_dt
-      else:
-        timestamp_filter = True
+        query = query.filter(StorageEvent.timestamp >= after_dt)
 
       storage_events = (
-          sql_session.query(StorageEvent)
-          .filter(StorageEvent.app_name == app_name)
-          .filter(StorageEvent.session_id == storage_session.id)
-          .filter(StorageEvent.user_id == user_id)
-          .filter(timestamp_filter)
-          .order_by(StorageEvent.timestamp.desc())
+          query.order_by(StorageEvent.timestamp.desc())
           .limit(
               config.num_recent_events
               if config and config.num_recent_events
@@ -660,30 +652,21 @@ async def append_event(self, session: Session, event: Event) -> Event:
           StorageUserState, (session.app_name, session.user_id)
       )
 
-      app_state = storage_app_state.state if storage_app_state else {}
-      user_state = storage_user_state.state if storage_user_state else {}
-      session_state = storage_session.state
-
       # Extract state delta
-      app_state_delta = {}
-      user_state_delta = {}
-      session_state_delta = {}
-      if event.actions:
-        if event.actions.state_delta:
-          app_state_delta, user_state_delta, session_state_delta = (
-              _extract_state_delta(event.actions.state_delta)
-          )
-
-      # Merge state and update storage
-      if app_state_delta:
-        app_state.update(app_state_delta)
-        storage_app_state.state = app_state
-      if user_state_delta:
-        user_state.update(user_state_delta)
-        storage_user_state.state = user_state
-      if session_state_delta:
-        session_state.update(session_state_delta)
-        storage_session.state = session_state
+      if event.actions and event.actions.state_delta:
+        state_deltas = _session_util.extract_state_delta(
+            event.actions.state_delta
+        )
+        app_state_delta = state_deltas["app"]
+        user_state_delta = state_deltas["user"]
+        session_state_delta = state_deltas["session"]
+        # Merge state and update storage
+        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
+        if session_state_delta:
+          storage_session.state = storage_session.state | session_state_delta
 
       sql_session.add(StorageEvent.from_event(session, event))
 
@@ -698,21 +681,6 @@ async def append_event(self, session: Session, event: Event) -> Event:
     return event
 
 
-def _extract_state_delta(state: dict[str, Any]):
-  app_state_delta = {}
-  user_state_delta = {}
-  session_state_delta = {}
-  if state:
-    for key in state.keys():
-      if key.startswith(State.APP_PREFIX):
-        app_state_delta[key.removeprefix(State.APP_PREFIX)] = state[key]
-      elif key.startswith(State.USER_PREFIX):
-        user_state_delta[key.removeprefix(State.USER_PREFIX)] = state[key]
-      elif not key.startswith(State.TEMP_PREFIX):
-        session_state_delta[key] = state[key]
-  return app_state_delta, user_state_delta, session_state_delta
-
-
 def _merge_state(app_state, user_state, session_state):
   # Merge states for response
   merged_state = copy.deepcopy(session_state)

src/google/adk/sessions/in_memory_session_service.py

@@ -22,6 +22,7 @@
 
 from typing_extensions import override
 
+from . import _session_util
 from ..events.event import Event
 from .base_session_service import BaseSessionService
 from .base_session_service import GetSessionConfig
@@ -88,6 +89,17 @@ def _create_session_impl(
       state: Optional[dict[str, Any]] = None,
       session_id: Optional[str] = None,
   ) -> Session:
+    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']
+    if app_state_delta:
+      self.app_state.setdefault(app_name, {}).update(app_state_delta)
+    if user_state_delta:
+      self.user_state.setdefault(app_name, {}).setdefault(user_id, {}).update(
+          user_state_delta
+      )
+
     session_id = (
         session_id.strip()
         if session_id and session_id.strip()
@@ -97,7 +109,7 @@ def _create_session_impl(
         app_name=app_name,
         user_id=user_id,
         id=session_id,
-        state=state or {},
+        state=session_state or {},
         last_update_time=time.time(),
     )
 
@@ -174,11 +186,13 @@ def _get_session_impl(
         if i >= 0:
           copied_session.events = copied_session.events[i + 1 :]
 
+    # Return a copy of the session object with merged state.
     return self._merge_state(app_name, user_id, copied_session)
 
   def _merge_state(
       self, app_name: str, user_id: str, copied_session: Session
   ) -> Session:
+    """Merges app and user state into session state."""
     # Merge app state
     if app_name in self.app_state:
       for key in self.app_state[app_name].keys():
@@ -269,11 +283,9 @@ def _delete_session_impl(
 
   @override
   async def append_event(self, session: Session, event: Event) -> Event:
-    # Update the in-memory session.
-    await super().append_event(session=session, event=event)
-    session.last_update_time = event.timestamp
+    if event.partial:
+      return event
 
-    # Update the storage session
     app_name = session.app_name
     user_id = session.user_id
     session_id = session.id
@@ -293,21 +305,29 @@ def _warning(message: str) -> None:
       _warning(f'session_id {session_id} not in sessions[app_name][user_id]')
       return event
 
-    if event.actions and event.actions.state_delta:
-      for key in event.actions.state_delta:
-        if key.startswith(State.APP_PREFIX):
-          self.app_state.setdefault(app_name, {})[
-              key.removeprefix(State.APP_PREFIX)
-          ] = event.actions.state_delta[key]
-
-        if key.startswith(State.USER_PREFIX):
-          self.user_state.setdefault(app_name, {}).setdefault(user_id, {})[
-              key.removeprefix(State.USER_PREFIX)
-          ] = event.actions.state_delta[key]
+    # Update the in-memory session.
+    await super().append_event(session=session, event=event)
+    session.last_update_time = event.timestamp
 
+    # Update the storage session
     storage_session = self.sessions[app_name][user_id].get(session_id)
-    await super().append_event(session=storage_session, event=event)
-
+    storage_session.events.append(event)
     storage_session.last_update_time = event.timestamp
 
+    if event.actions and event.actions.state_delta:
+      state_deltas = _session_util.extract_state_delta(
+          event.actions.state_delta
+      )
+      app_state_delta = state_deltas['app']
+      user_state_delta = state_deltas['user']
+      session_state_delta = state_deltas['session']
+      if app_state_delta:
+        self.app_state.setdefault(app_name, {}).update(app_state_delta)
+      if user_state_delta:
+        self.user_state.setdefault(app_name, {}).setdefault(user_id, {}).update(
+            user_state_delta
+        )
+      if session_state_delta:
+        storage_session.state.update(session_state_delta)
+
     return event