From 49d4441ebd4b715919c52ea15319df0abdff9609 Mon Sep 17 00:00:00 2001 From: Google Team Member Date: Mon, 22 Jun 2026 12:04:01 -0700 Subject: [PATCH] feat: add support for session TTL and expiration in Vertex AI session service PiperOrigin-RevId: 936180164 --- src/google/adk/cli/cli_deploy.py | 17 +++++--- .../adk/sessions/vertex_ai_session_service.py | 6 +++ .../test_vertex_ai_session_service.py | 40 +++++++++++++++++++ 3 files changed, 58 insertions(+), 5 deletions(-) diff --git a/src/google/adk/cli/cli_deploy.py b/src/google/adk/cli/cli_deploy.py index 4e643d11ca..e7fed4a4ef 100644 --- a/src/google/adk/cli/cli_deploy.py +++ b/src/google/adk/cli/cli_deploy.py @@ -147,6 +147,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', @@ -215,19 +217,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', diff --git a/src/google/adk/sessions/vertex_ai_session_service.py b/src/google/adk/sessions/vertex_ai_session_service.py index 6b54dcc140..154555371d 100644 --- a/src/google/adk/sessions/vertex_ai_session_service.py +++ b/src/google/adk/sessions/vertex_ai_session_service.py @@ -159,12 +159,18 @@ 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. + Returns: The created session. """ + if kwargs.get('ttl') is not None and kwargs.get('expire_time') is not None: + 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 {} diff --git a/tests/unittests/sessions/test_vertex_ai_session_service.py b/tests/unittests/sessions/test_vertex_ai_session_service.py index 88d0191814..2aea612afc 100644 --- a/tests/unittests/sessions/test_vertex_ai_session_service.py +++ b/tests/unittests/sessions/test_vertex_ai_session_service.py @@ -1008,6 +1008,46 @@ 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_create_session_with_ttl_none_and_expire_time_none_does_not_raise( + mock_api_client_instance, +): + session_service = mock_vertex_ai_session_service() + # None means "not set"; passing both as None must not raise. + await session_service.create_session( + app_name='123', user_id='user', ttl=None, expire_time=None + ) + + @pytest.mark.asyncio @pytest.mark.usefixtures('mock_get_api_client') async def test_append_event():