fix(sessions): stop returned sessions aliasing scoped state

src/google/adk/sessions/in_memory_session_service.py

@@ -59,6 +59,13 @@ def _copy_session(session: Session) -> Session:
     return copy.deepcopy(session)
 
 
+def _copy_state(state: dict[str, Any]) -> dict[str, Any]:
+  """Copies state as deeply as the session service is configured to copy."""
+  if is_feature_enabled(FeatureName.IN_MEMORY_SESSION_SERVICE_LIGHT_COPY):
+    return dict(state)
+  return copy.deepcopy(state)
+
+
 class InMemorySessionService(BaseSessionService):
   """An in-memory implementation of the session service.
 
@@ -225,10 +232,8 @@ def _merge_state(
     """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():
-        copied_session.state[State.APP_PREFIX + key] = self.app_state[app_name][
-            key
-        ]
+      for key, value in _copy_state(self.app_state[app_name]).items():
+        copied_session.state[State.APP_PREFIX + key] = value
 
     if (
         app_name not in self.user_state
@@ -237,10 +242,8 @@ def _merge_state(
       return copied_session
 
     # Merge session state with user state.
-    for key in self.user_state[app_name][user_id].keys():
-      copied_session.state[State.USER_PREFIX + key] = self.user_state[app_name][
-          user_id
-      ][key]
+    for key, value in _copy_state(self.user_state[app_name][user_id]).items():
+      copied_session.state[State.USER_PREFIX + key] = value
     return copied_session
 
   @override
@@ -319,11 +322,7 @@ async def get_user_state(
       self, *, app_name: str, user_id: str
   ) -> dict[str, Any]:
     user_state = self.user_state.get(app_name, {}).get(user_id, {})
-    # Copy as deeply as _copy_session copies a session's own state, so user
-    # state is no more reachable through the result than session state is.
-    if is_feature_enabled(FeatureName.IN_MEMORY_SESSION_SERVICE_LIGHT_COPY):
-      return dict(user_state)
-    return copy.deepcopy(user_state)
+    return _copy_state(user_state)
 
   @override
   async def append_event(self, session: Session, event: Event) -> Event:

tests/unittests/sessions/test_session_service.py

@@ -2840,6 +2840,55 @@ async def test_get_user_state_copies_to_session_state_depth(light_copy):
     )
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize('light_copy', [False, True])
+@pytest.mark.parametrize('session_source', ['create', 'get', 'list'])
+async def test_returned_session_scoped_state_uses_configured_copy_depth(
+    light_copy, session_source
+):
+  """Returned sessions copy nested scoped state to the configured depth."""
+  override_feature_enabled(
+      FeatureName.IN_MEMORY_SESSION_SERVICE_LIGHT_COPY, light_copy
+  )
+  try:
+    service = InMemorySessionService()
+    created = await service.create_session(
+        app_name='my_app',
+        user_id='u1',
+        session_id='s1',
+        state={
+            'app:config': {'theme': 'light'},
+            'user:profile': {'name': 'Alice'},
+        },
+    )
+
+    if session_source == 'create':
+      returned = created
+    elif session_source == 'get':
+      returned = await service.get_session(
+          app_name='my_app', user_id='u1', session_id='s1'
+      )
+    else:
+      returned = (
+          await service.list_sessions(app_name='my_app', user_id='u1')
+      ).sessions[0]
+
+    returned.state['app:config']['theme'] = 'dark'
+    returned.state['user:profile']['name'] = 'Mallory'
+    later = await service.create_session(
+        app_name='my_app', user_id='u1', session_id='s2'
+    )
+
+    expected_theme = 'dark' if light_copy else 'light'
+    expected_name = 'Mallory' if light_copy else 'Alice'
+    assert later.state['app:config']['theme'] == expected_theme
+    assert later.state['user:profile']['name'] == expected_name
+  finally:
+    override_feature_enabled(
+        FeatureName.IN_MEMORY_SESSION_SERVICE_LIGHT_COPY, False
+    )
+
+
 @pytest.mark.asyncio
 async def test_vertex_ai_session_service_raises_not_implemented_for_get_user_state():
   """Verifies VertexAiSessionService raises NotImplementedError."""