Skip to content

Commit 21b9dcf

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 21b9dcf

2 files changed

Lines changed: 28 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: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -554,6 +554,30 @@ async def test_async_create_session(self, get_project_id_mock: mock.Mock):
554554
assert session2["user_id"] == _TEST_USER_ID
555555
assert session2["id"] == "test_session_id"
556556

557+
@pytest.mark.asyncio
558+
async def test_async_create_session_with_kwargs(
559+
self, get_project_id_mock: mock.Mock
560+
):
561+
mock_session_service = mock.AsyncMock()
562+
app = agent_engines.AdkApp(
563+
agent=_TEST_AGENT,
564+
session_service_builder=lambda: mock_session_service,
565+
)
566+
await app.async_create_session(
567+
user_id=_TEST_USER_ID,
568+
ttl="7200s",
569+
expire_time="2026-03-01T00:00:00Z",
570+
)
571+
mock_session_service.create_session.assert_called_once_with(
572+
app_name=app._app_name(),
573+
user_id=_TEST_USER_ID,
574+
session_id=None,
575+
state=None,
576+
ttl="7200s",
577+
expire_time="2026-03-01T00:00:00Z",
578+
)
579+
580+
557581
@pytest.mark.asyncio
558582
async def test_async_get_session(self, get_project_id_mock: mock.Mock):
559583
app = agent_engines.AdkApp(agent=_TEST_AGENT)

0 commit comments

Comments
 (0)