feat: Option to use shallow-copy for session in InMemorySessionService

PiperOrigin-RevId: 893606046

Author: google-genai-bot

src/google/adk/features/_feature_registry.py

@@ -53,6 +53,7 @@ class FeatureName(str, Enum):
   TOOL_CONFIRMATION = "TOOL_CONFIRMATION"
   PLUGGABLE_AUTH = "PLUGGABLE_AUTH"
   SNAKE_CASE_SKILL_NAME = "SNAKE_CASE_SKILL_NAME"
+  IN_MEMORY_SESSION_SERVICE_LIGHT_COPY = "IN_MEMORY_SESSION_SERVICE_LIGHT_COPY"
 
 
 class FeatureStage(Enum):
@@ -166,6 +167,9 @@ class FeatureConfig:
     FeatureName.SNAKE_CASE_SKILL_NAME: FeatureConfig(
         FeatureStage.EXPERIMENTAL, default_on=False
     ),
+    FeatureName.IN_MEMORY_SESSION_SERVICE_LIGHT_COPY: FeatureConfig(
+        FeatureStage.WIP, default_on=False
+    ),
 }
 
 # Track which experimental features have already warned (warn only once)

src/google/adk/sessions/in_memory_session_service.py

@@ -25,6 +25,8 @@
 from . import _session_util
 from ..errors.already_exists_error import AlreadyExistsError
 from ..events.event import Event
+from ..features import FeatureName
+from ..features import is_feature_enabled
 from .base_session_service import BaseSessionService
 from .base_session_service import GetSessionConfig
 from .base_session_service import ListSessionsResponse
@@ -34,6 +36,28 @@
 logger = logging.getLogger('google_adk.' + __name__)
 
 
+def _light_copy(session: Session) -> Session:
+  """Returns a light copy of the session.
+
+  Main difference between this and true shallow-copy is that container fields
+  (e.g., events and state) are also shallow-copied. What this means is appending
+  to events/state of the copied session won't affect the original while avoiding
+  the potentially expensive cost of a full/recursive deep-copy of all events and
+  state.
+  """
+  copied_session = session.model_copy(deep=False)
+  copied_session.events = copy.copy(session.events)
+  copied_session.state = copy.copy(session.state)
+  return copied_session
+
+
+def _copy_session(session: Session) -> Session:
+  if is_feature_enabled(FeatureName.IN_MEMORY_SESSION_SERVICE_LIGHT_COPY):
+    return _light_copy(session)
+  else:
+    return copy.deepcopy(session)
+
+
 class InMemorySessionService(BaseSessionService):
   """An in-memory implementation of the session service.
 
@@ -124,7 +148,7 @@ def _create_session_impl(
       self.sessions[app_name][user_id] = {}
     self.sessions[app_name][user_id][session_id] = session
 
-    copied_session = copy.deepcopy(session)
+    copied_session = _copy_session(session)
     return self._merge_state(app_name, user_id, copied_session)
 
   @override
@@ -175,7 +199,7 @@ def _get_session_impl(
       return None
 
     session = self.sessions[app_name][user_id].get(session_id)
-    copied_session = copy.deepcopy(session)
+    copied_session = _copy_session(session)
 
     if config:
       if config.num_recent_events is not None:
@@ -247,13 +271,13 @@ def _list_sessions_impl(
     if user_id is None:
       for uid in list(self.sessions[app_name].keys()):
         for session in list(self.sessions[app_name][uid].values()):
-          copied_session = copy.deepcopy(session)
+          copied_session = _copy_session(session)
           copied_session.events = []
           copied_session = self._merge_state(app_name, uid, copied_session)
           sessions_without_events.append(copied_session)
     else:
       for session in list(self.sessions[app_name][user_id].values()):
-        copied_session = copy.deepcopy(session)
+        copied_session = _copy_session(session)
         copied_session.events = []
         copied_session = self._merge_state(app_name, user_id, copied_session)
         sessions_without_events.append(copied_session)

tests/unittests/sessions/test_session_service.py

@@ -23,6 +23,8 @@
 from google.adk.errors.already_exists_error import AlreadyExistsError
 from google.adk.events.event import Event
 from google.adk.events.event_actions import EventActions
+from google.adk.features import FeatureName
+from google.adk.features import override_feature_enabled
 from google.adk.sessions import database_session_service
 from google.adk.sessions.base_session_service import GetSessionConfig
 from google.adk.sessions.database_session_service import DatabaseSessionService
@@ -35,6 +37,7 @@
 
 class SessionServiceType(enum.Enum):
   IN_MEMORY = 'IN_MEMORY'
+  IN_MEMORY_WITH_LIGHT_COPY_ENABLED = 'IN_MEMORY_WITH_LIGHT_COPY_ENABLED'
   DATABASE = 'DATABASE'
   SQLITE = 'SQLITE'
 
@@ -48,22 +51,33 @@ def get_session_service(
     return DatabaseSessionService('sqlite+aiosqlite:///:memory:')
   if service_type == SessionServiceType.SQLITE:
     return SqliteSessionService(str(tmp_path / 'sqlite.db'))
+  if service_type == SessionServiceType.IN_MEMORY_WITH_LIGHT_COPY_ENABLED:
+    return InMemorySessionService()
   return InMemorySessionService()
 
 
 @pytest.fixture(
     params=[
         SessionServiceType.IN_MEMORY,
+        SessionServiceType.IN_MEMORY_WITH_LIGHT_COPY_ENABLED,
         SessionServiceType.DATABASE,
         SessionServiceType.SQLITE,
     ]
 )
 async def session_service(request, tmp_path):
   """Provides a session service and closes database backends on teardown."""
+  if request.param == SessionServiceType.IN_MEMORY_WITH_LIGHT_COPY_ENABLED:
+    override_feature_enabled(
+        FeatureName.IN_MEMORY_SESSION_SERVICE_LIGHT_COPY, True
+    )
   service = get_session_service(request.param, tmp_path)
   yield service
   if isinstance(service, DatabaseSessionService):
     await service.close()
+  if request.param == SessionServiceType.IN_MEMORY_WITH_LIGHT_COPY_ENABLED:
+    override_feature_enabled(
+        FeatureName.IN_MEMORY_SESSION_SERVICE_LIGHT_COPY, False
+    )
 
 
 def test_database_session_service_enables_pool_pre_ping_by_default():