diff --git a/agentplatform/agent_engines/templates/adk.py b/agentplatform/agent_engines/templates/adk.py index 91737e688a..115351c8bc 100644 --- a/agentplatform/agent_engines/templates/adk.py +++ b/agentplatform/agent_engines/templates/adk.py @@ -1616,6 +1616,10 @@ async def async_create_session( will be be generated for the session. state (dict[str, Any]): Optional. The initial state of the session. + ttl (str): + Optional. The time-to-live for the session. + expire_time (str): + Optional. The expiration time for the session. **kwargs (dict[str, Any]): Optional. Additional keyword arguments to pass to the session service. diff --git a/tests/unit/agentplatform/frameworks/test_frameworks_adk.py b/tests/unit/agentplatform/frameworks/test_frameworks_adk.py index 6da7d41409..7e3d5dba6f 100644 --- a/tests/unit/agentplatform/frameworks/test_frameworks_adk.py +++ b/tests/unit/agentplatform/frameworks/test_frameworks_adk.py @@ -41,6 +41,7 @@ ) from google.genai import types import pytest +from google.adk.sessions.base_session_service import BaseSessionService try: @@ -701,6 +702,50 @@ async def test_async_create_session(self, get_project_id_mock: mock.Mock): assert session2["user_id"] == _TEST_USER_ID assert session2["id"] == "test_session_id" + @pytest.mark.asyncio + async def test_async_create_session_with_ttl_kwargs( + self, get_project_id_mock: mock.Mock + ): + mock_session_service = mock.Mock(spec=BaseSessionService) + mock_session_service.create_session = mock.AsyncMock() + app = adk_template.AdkApp( + agent=_TEST_AGENT, + session_service_builder=lambda: mock_session_service, + ) + await app.async_create_session( + user_id=_TEST_USER_ID, + ttl="7200s", + ) + mock_session_service.create_session.assert_called_once_with( + app_name=app._app_name(), + user_id=_TEST_USER_ID, + session_id=None, + state=None, + ttl="7200s", + ) + + @pytest.mark.asyncio + async def test_async_create_session_with_expire_time_kwargs( + self, get_project_id_mock: mock.Mock + ): + mock_session_service = mock.Mock(spec=BaseSessionService) + mock_session_service.create_session = mock.AsyncMock() + app = adk_template.AdkApp( + agent=_TEST_AGENT, + session_service_builder=lambda: mock_session_service, + ) + await app.async_create_session( + user_id=_TEST_USER_ID, + expire_time="2026-03-01T00:00:00Z", + ) + mock_session_service.create_session.assert_called_once_with( + app_name=app._app_name(), + user_id=_TEST_USER_ID, + session_id=None, + state=None, + expire_time="2026-03-01T00:00:00Z", + ) + @pytest.mark.asyncio async def test_async_get_session(self, get_project_id_mock: mock.Mock): app = adk_template.AdkApp(agent=_TEST_AGENT) diff --git a/tests/unit/vertex_adk/test_agent_engine_templates_adk.py b/tests/unit/vertex_adk/test_agent_engine_templates_adk.py index 2e80315044..044852b575 100644 --- a/tests/unit/vertex_adk/test_agent_engine_templates_adk.py +++ b/tests/unit/vertex_adk/test_agent_engine_templates_adk.py @@ -40,6 +40,7 @@ from vertexai.agent_engines.templates import adk as adk_template from google.genai import types import pytest +from google.adk.sessions.base_session_service import BaseSessionService try: @@ -620,6 +621,51 @@ async def test_async_create_session(self, get_project_id_mock: mock.Mock): assert session2["user_id"] == _TEST_USER_ID assert session2["id"] == "test_session_id" + @pytest.mark.asyncio + async def test_async_create_session_with_ttl_kwargs( + self, get_project_id_mock: mock.Mock + ): + mock_session_service = mock.Mock(spec=BaseSessionService) + mock_session_service.create_session = mock.AsyncMock() + app = agent_engines.AdkApp( + agent=_TEST_AGENT, + session_service_builder=lambda: mock_session_service, + ) + await app.async_create_session( + user_id=_TEST_USER_ID, + ttl="7200s", + ) + mock_session_service.create_session.assert_called_once_with( + app_name=app._app_name(), + user_id=_TEST_USER_ID, + session_id=None, + state=None, + ttl="7200s", + ) + + @pytest.mark.asyncio + async def test_async_create_session_with_expire_time_kwargs( + self, get_project_id_mock: mock.Mock + ): + mock_session_service = mock.Mock(spec=BaseSessionService) + mock_session_service.create_session = mock.AsyncMock() + app = agent_engines.AdkApp( + agent=_TEST_AGENT, + session_service_builder=lambda: mock_session_service, + ) + await app.async_create_session( + user_id=_TEST_USER_ID, + expire_time="2026-03-01T00:00:00Z", + ) + mock_session_service.create_session.assert_called_once_with( + app_name=app._app_name(), + user_id=_TEST_USER_ID, + session_id=None, + state=None, + expire_time="2026-03-01T00:00:00Z", + ) + + @pytest.mark.asyncio async def test_async_get_session(self, get_project_id_mock: mock.Mock): app = agent_engines.AdkApp(agent=_TEST_AGENT) diff --git a/vertexai/agent_engines/templates/adk.py b/vertexai/agent_engines/templates/adk.py index 58c3025285..48fcb04ddd 100644 --- a/vertexai/agent_engines/templates/adk.py +++ b/vertexai/agent_engines/templates/adk.py @@ -1576,6 +1576,10 @@ async def async_create_session( will be be generated for the session. state (dict[str, Any]): Optional. The initial state of the session. + ttl (str): + Optional. The time-to-live for the session. + expire_time (str): + Optional. The expiration time for the session. **kwargs (dict[str, Any]): Optional. Additional keyword arguments to pass to the session service.