From 3ab365eb8a3c30512441d76fd299422debfaa38b Mon Sep 17 00:00:00 2001 From: Google Team Member Date: Mon, 18 May 2026 02:45:34 -0700 Subject: [PATCH] feat: add support for session TTL and expiration in Vertex AI session service PiperOrigin-RevId: 917111472 --- src/google/adk/cli/cli_deploy.py | 17 +++++++---- .../adk/sessions/base_session_service.py | 3 ++ .../adk/sessions/database_session_service.py | 1 + .../adk/sessions/in_memory_session_service.py | 5 ++++ .../adk/sessions/sqlite_session_service.py | 1 + .../adk/sessions/vertex_ai_session_service.py | 11 ++++++-- .../test_vertex_ai_session_service.py | 28 +++++++++++++++++++ 7 files changed, 59 insertions(+), 7 deletions(-) diff --git a/src/google/adk/cli/cli_deploy.py b/src/google/adk/cli/cli_deploy.py index 4350ae17ce..2dd6ccf1bf 100644 --- a/src/google/adk/cli/cli_deploy.py +++ b/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', diff --git a/src/google/adk/sessions/base_session_service.py b/src/google/adk/sessions/base_session_service.py index 324abe230b..63a0551f38 100644 --- a/src/google/adk/sessions/base_session_service.py +++ b/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. diff --git a/src/google/adk/sessions/database_session_service.py b/src/google/adk/sessions/database_session_service.py index d033f1f234..55288b1561 100644 --- a/src/google/adk/sessions/database_session_service.py +++ b/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 diff --git a/src/google/adk/sessions/in_memory_session_service.py b/src/google/adk/sessions/in_memory_session_service.py index b8f6cfab46..f1387eebe6 100644 --- a/src/google/adk/sessions/in_memory_session_service.py +++ b/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,6 +99,7 @@ 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( @@ -104,6 +107,7 @@ def create_session_sync( 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 diff --git a/src/google/adk/sessions/sqlite_session_service.py b/src/google/adk/sessions/sqlite_session_service.py index 427bc3e73e..a69496866f 100644 --- a/src/google/adk/sessions/sqlite_session_service.py +++ b/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() diff --git a/src/google/adk/sessions/vertex_ai_session_service.py b/src/google/adk/sessions/vertex_ai_session_service.py index 8c1fdc134e..76bc66e1e5 100644 --- a/src/google/adk/sessions/vertex_ai_session_service.py +++ b/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 {} diff --git a/tests/unittests/sessions/test_vertex_ai_session_service.py b/tests/unittests/sessions/test_vertex_ai_session_service.py index c5c9996ef5..01e984331a 100644 --- a/tests/unittests/sessions/test_vertex_ai_session_service.py +++ b/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():