Skip to content

Commit 7232086

Browse files
vertex-sdk-botcopybara-github
authored andcommitted
feat: add support for session TTL and expiration in Vertex AI session service
PiperOrigin-RevId: 917111472
1 parent 1585602 commit 7232086

3 files changed

Lines changed: 54 additions & 0 deletions

File tree

agentplatform/agent_engines/templates/adk.py

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1568,6 +1568,10 @@ async def async_create_session(
15681568
will be be generated for the session.
15691569
state (dict[str, Any]):
15701570
Optional. The initial state of the session.
1571+
ttl (str):
1572+
Optional. The time-to-live for the session.
1573+
expire_time (str):
1574+
Optional. The expiration time for the session.
15711575
**kwargs (dict[str, Any]):
15721576
Optional. Additional keyword arguments to pass to the
15731577
session service.

tests/unit/vertex_adk/test_agent_engine_templates_adk.py

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,7 @@
4040
from vertexai.agent_engines.templates import adk as adk_template
4141
from google.genai import types
4242
import pytest
43+
from google.adk.sessions.base_session_service import BaseSessionService
4344

4445

4546
try:
@@ -554,6 +555,51 @@ async def test_async_create_session(self, get_project_id_mock: mock.Mock):
554555
assert session2["user_id"] == _TEST_USER_ID
555556
assert session2["id"] == "test_session_id"
556557

558+
@pytest.mark.asyncio
559+
async def test_async_create_session_with_ttl_kwargs(
560+
self, get_project_id_mock: mock.Mock
561+
):
562+
mock_session_service = mock.Mock(spec=BaseSessionService)
563+
mock_session_service.create_session = mock.AsyncMock()
564+
app = agent_engines.AdkApp(
565+
agent=_TEST_AGENT,
566+
session_service_builder=lambda: mock_session_service,
567+
)
568+
await app.async_create_session(
569+
user_id=_TEST_USER_ID,
570+
ttl="7200s",
571+
)
572+
mock_session_service.create_session.assert_called_once_with(
573+
app_name=app._app_name(),
574+
user_id=_TEST_USER_ID,
575+
session_id=None,
576+
state=None,
577+
ttl="7200s",
578+
)
579+
580+
@pytest.mark.asyncio
581+
async def test_async_create_session_with_expire_time_kwargs(
582+
self, get_project_id_mock: mock.Mock
583+
):
584+
mock_session_service = mock.Mock(spec=BaseSessionService)
585+
mock_session_service.create_session = mock.AsyncMock()
586+
app = agent_engines.AdkApp(
587+
agent=_TEST_AGENT,
588+
session_service_builder=lambda: mock_session_service,
589+
)
590+
await app.async_create_session(
591+
user_id=_TEST_USER_ID,
592+
expire_time="2026-03-01T00:00:00Z",
593+
)
594+
mock_session_service.create_session.assert_called_once_with(
595+
app_name=app._app_name(),
596+
user_id=_TEST_USER_ID,
597+
session_id=None,
598+
state=None,
599+
expire_time="2026-03-01T00:00:00Z",
600+
)
601+
602+
557603
@pytest.mark.asyncio
558604
async def test_async_get_session(self, get_project_id_mock: mock.Mock):
559605
app = agent_engines.AdkApp(agent=_TEST_AGENT)

vertexai/agent_engines/templates/adk.py

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1568,6 +1568,10 @@ async def async_create_session(
15681568
will be be generated for the session.
15691569
state (dict[str, Any]):
15701570
Optional. The initial state of the session.
1571+
ttl (str):
1572+
Optional. The time-to-live for the session.
1573+
expire_time (str):
1574+
Optional. The expiration time for the session.
15711575
**kwargs (dict[str, Any]):
15721576
Optional. Additional keyword arguments to pass to the
15731577
session service.

0 commit comments

Comments
 (0)