diff --git a/src/google/adk/sessions/in_memory_session_service.py b/src/google/adk/sessions/in_memory_session_service.py index 8c6bfafedf..bd48ff22fe 100644 --- a/src/google/adk/sessions/in_memory_session_service.py +++ b/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: diff --git a/tests/unittests/sessions/test_session_service.py b/tests/unittests/sessions/test_session_service.py index 30dadf4517..b0b8322d06 100644 --- a/tests/unittests/sessions/test_session_service.py +++ b/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."""