feat: add support for session TTL and expiration in Vertex AI session service

PiperOrigin-RevId: 917111472

Author: google-genai-bot

src/google/adk/cli/cli_deploy.py

@@ -171,6 +171,8 @@ def _ensure_agent_engine_dependency(requirements_txt_path: str) -> None:
                 'user_id': {'type': 'string'},
                 'session_id': {'type': 'string', 'nullable': True},
                 'state': {'type': 'object', 'nullable': True},
+                'ttl': {'type': 'string', 'nullable': True},
+                'expire_time': {'type': 'string', 'nullable': True},
             },
             'required': ['user_id'],
             'type': 'object',
@@ -239,19 +241,24 @@ def _ensure_agent_engine_dependency(requirements_txt_path: str) -> None:
             'Creates a new session.\n\n        Args:\n            user_id'
             ' (str):\n                Required. The ID of the user.\n          '
             '  session_id (str):\n                Optional. The ID of the'
-            ' session. If not provided, an ID\n                will be be'
+            ' session. If not provided, an ID\n                will be'
             ' generated for the session.\n            state (dict[str, Any]):\n'
             '                Optional. The initial state of the session.\n     '
-            '       **kwargs (dict[str, Any]):\n                Optional.'
-            ' Additional keyword arguments to pass to the\n               '
-            ' session service.\n\n        Returns:\n            Session: The'
-            ' newly created session instance.\n        '
+            '       ttl (str):\n                Optional. The time-to-live for'
+            ' the session.\n            expire_time (str):\n               '
+            ' Optional. The expiration time for the session.\n           '
+            ' **kwargs (dict[str, Any]):\n                Optional. Additional'
+            ' keyword arguments to pass to the\n                session'
+            ' service.\n\n        Returns:\n            Session: The newly'
+            ' created session instance.\n        '
         ),
         'parameters': {
             'properties': {
                 'user_id': {'type': 'string'},
                 'session_id': {'type': 'string', 'nullable': True},
                 'state': {'type': 'object', 'nullable': True},
+                'ttl': {'type': 'string', 'nullable': True},
+                'expire_time': {'type': 'string', 'nullable': True},
             },
             'required': ['user_id'],
             'type': 'object',

src/google/adk/sessions/base_session_service.py

@@ -65,6 +65,7 @@ async def create_session(
       user_id: str,
       state: Optional[dict[str, Any]] = None,
       session_id: Optional[str] = None,
+      **kwargs: Any,
   ) -> Session:
     """Creates a new session.
 
@@ -74,6 +75,8 @@ async def create_session(
       state: the initial state of the session.
       session_id: the client-provided id of the session. If not provided, a
         generated ID will be used.
+      **kwargs: Additional keyword arguments to pass to the session creation,
+        such as 'ttl' or 'expire_time'. Supported by some implementations.
 
     Returns:
       session: The newly created session instance.

src/google/adk/sessions/database_session_service.py

@@ -416,6 +416,7 @@ async def create_session(
       user_id: str,
       state: Optional[dict[str, Any]] = None,
       session_id: Optional[str] = None,
+      **kwargs: Any,
   ) -> Session:
     # 1. Populate states.
     # 2. Build storage session object

src/google/adk/sessions/in_memory_session_service.py

@@ -82,12 +82,14 @@ async def create_session(
       user_id: str,
       state: Optional[dict[str, Any]] = None,
       session_id: Optional[str] = None,
+      **kwargs: Any,
   ) -> Session:
     return self._create_session_impl(
         app_name=app_name,
         user_id=user_id,
         state=state,
         session_id=session_id,
+        **kwargs,
     )
 
   def create_session_sync(
@@ -97,13 +99,15 @@ def create_session_sync(
       user_id: str,
       state: Optional[dict[str, Any]] = None,
       session_id: Optional[str] = None,
+      **kwargs: Any,
   ) -> Session:
     logger.warning('Deprecated. Please migrate to the async method.')
     return self._create_session_impl(
         app_name=app_name,
         user_id=user_id,
         state=state,
         session_id=session_id,
+        **kwargs,
     )
 
   def _create_session_impl(
@@ -113,6 +117,7 @@ def _create_session_impl(
       user_id: str,
       state: Optional[dict[str, Any]] = None,
       session_id: Optional[str] = None,
+      **kwargs: Any,
   ) -> Session:
     if session_id and self._get_session_impl(
         app_name=app_name, user_id=user_id, session_id=session_id

src/google/adk/sessions/sqlite_session_service.py

@@ -161,6 +161,7 @@ async def create_session(
       user_id: str,
       state: Optional[dict[str, Any]] = None,
       session_id: Optional[str] = None,
+      **kwargs: Any,
   ) -> Session:
     if session_id:
       session_id = session_id.strip()

src/google/adk/sessions/vertex_ai_session_service.py

@@ -117,12 +117,19 @@ async def create_session(
       state: The initial state of the session.
       session_id: The ID of the session.
       **kwargs: Additional arguments to pass to the session creation. E.g. set
+        ttl='7200s' to set the session time-to-live or
         expire_time='2025-10-01T00:00:00Z' to set the session expiration time.
-        See https://cloud.google.com/vertex-ai/generative-ai/docs/reference/rest/v1beta1/projects.locations.reasoningEngines.sessions
-        for more details.
+        See
+        https://cloud.google.com/vertex-ai/generative-ai/docs/reference/rest/v1beta1/projects.locations.reasoningEngines.sessions
+          for more details.
+
     Returns:
       The created session.
     """
+    if 'ttl' in kwargs and 'expire_time' in kwargs:
+      raise ValueError(
+          "Cannot specify both 'ttl' and 'expire_time' simultaneously."
+      )
     reasoning_engine_id = self._get_reasoning_engine_id(app_name)
 
     config = {'session_state': state} if state else {}

tests/unittests/sessions/test_vertex_ai_session_service.py

@@ -967,6 +967,34 @@ async def test_create_session_with_custom_config(mock_api_client_instance):
   )
 
 
+@pytest.mark.asyncio
+@pytest.mark.usefixtures('mock_get_api_client')
+async def test_create_session_with_ttl(mock_api_client_instance):
+  session_service = mock_vertex_ai_session_service()
+
+  ttl = '7200s'
+  await session_service.create_session(app_name='123', user_id='user', ttl=ttl)
+  assert mock_api_client_instance.last_create_session_config['ttl'] == ttl
+
+
+@pytest.mark.asyncio
+@pytest.mark.usefixtures('mock_get_api_client')
+async def test_create_session_with_ttl_and_expire_time_raises_value_error(
+    mock_api_client_instance,
+):
+  session_service = mock_vertex_ai_session_service()
+  with pytest.raises(
+      ValueError,
+      match="Cannot specify both 'ttl' and 'expire_time' simultaneously.",
+  ):
+    await session_service.create_session(
+        app_name='123',
+        user_id='user',
+        ttl='7200s',
+        expire_time='2025-12-12T12:12:12.123456Z',
+    )
+
+
 @pytest.mark.asyncio
 @pytest.mark.usefixtures('mock_get_api_client')
 async def test_append_event():