Skip to content

Commit 3804bfb

Browse files
fix(sessions): Accept Vertex session resource names
Normalize fully qualified Vertex session resource names before building Agent Engine session paths. This keeps strict plain session ID validation while allowing Gemini Enterprise and Vertex resource-name callers to query sessions successfully.
1 parent 0c6974c commit 3804bfb

2 files changed

Lines changed: 76 additions & 3 deletions

File tree

src/google/adk/sessions/vertex_ai_session_service.py

Lines changed: 20 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,11 @@
4848
_USAGE_METADATA_CUSTOM_METADATA_KEY = '_usage_metadata'
4949

5050
_SESSION_ID_PATTERN = re.compile(r'^[A-Za-z0-9_-]+$')
51+
_SESSION_RESOURCE_NAME_PATTERN = re.compile(
52+
r'^projects/[^/]+/locations/[^/]+/'
53+
r'(?:collections/[^/]+/engines/[^/]+|reasoningEngines/[^/]+)/'
54+
r'sessions/([^/]+)$'
55+
)
5156

5257

5358
def _validate_session_id(session_id: str) -> None:
@@ -61,6 +66,19 @@ def _validate_session_id(session_id: str) -> None:
6166
)
6267

6368

69+
def _normalize_session_id(session_id: str) -> str:
70+
"""Returns the plain session ID from a session ID or Vertex resource name."""
71+
match = (
72+
_SESSION_RESOURCE_NAME_PATTERN.fullmatch(session_id)
73+
if isinstance(session_id, str)
74+
else None
75+
)
76+
if match:
77+
session_id = match.group(1)
78+
_validate_session_id(session_id)
79+
return session_id
80+
81+
6482
def _quote_filter_literal(value: str) -> str:
6583
"""Quotes filter values so embedded metacharacters stay inside the literal."""
6684
escaped_value = value.replace('\\', '\\\\').replace('"', '\\"')
@@ -178,7 +196,7 @@ async def get_session(
178196
session_id: str,
179197
config: Optional[GetSessionConfig] = None,
180198
) -> Optional[Session]:
181-
_validate_session_id(session_id)
199+
session_id = _normalize_session_id(session_id)
182200
reasoning_engine_id = self._get_reasoning_engine_id(app_name)
183201
session_resource_name = (
184202
f'reasoningEngines/{reasoning_engine_id}/sessions/{session_id}'
@@ -278,7 +296,7 @@ async def list_sessions(
278296
async def delete_session(
279297
self, *, app_name: str, user_id: str, session_id: str
280298
) -> None:
281-
_validate_session_id(session_id)
299+
session_id = _normalize_session_id(session_id)
282300
reasoning_engine_id = self._get_reasoning_engine_id(app_name)
283301
session_resource_name = (
284302
f'reasoningEngines/{reasoning_engine_id}/sessions/{session_id}'

tests/unittests/sessions/test_vertex_ai_session_service.py

Lines changed: 56 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -725,6 +725,53 @@ async def test_get_and_delete_session():
725725
assert str(excinfo.value) == '404 Session not found: 1'
726726

727727

728+
@pytest.mark.asyncio
729+
@pytest.mark.usefixtures('mock_get_api_client')
730+
@pytest.mark.parametrize(
731+
'resource_name',
732+
[
733+
(
734+
'projects/my-project/locations/global/collections/'
735+
'default_collection/engines/my-app/sessions/1'
736+
),
737+
'projects/my-project/locations/us-central1/reasoningEngines/123/sessions/1',
738+
],
739+
)
740+
async def test_get_session_accepts_fully_qualified_resource_name(
741+
resource_name: str,
742+
):
743+
session_service = mock_vertex_ai_session_service()
744+
745+
session = await session_service.get_session(
746+
app_name='123',
747+
user_id='user',
748+
session_id=resource_name,
749+
)
750+
751+
assert session == MOCK_SESSION
752+
753+
754+
@pytest.mark.asyncio
755+
@pytest.mark.usefixtures('mock_get_api_client')
756+
async def test_delete_session_accepts_fully_qualified_resource_name():
757+
session_service = mock_vertex_ai_session_service()
758+
759+
await session_service.delete_session(
760+
app_name='123',
761+
user_id='user',
762+
session_id=(
763+
'projects/my-project/locations/global/collections/default_collection/'
764+
'engines/my-app/sessions/1'
765+
),
766+
)
767+
768+
with pytest.raises(api_core_exceptions.NotFound) as excinfo:
769+
await session_service.get_session(
770+
app_name='123', user_id='user', session_id='1'
771+
)
772+
assert str(excinfo.value) == '404 Session not found: 1'
773+
774+
728775
@pytest.mark.asyncio
729776
@pytest.mark.usefixtures('mock_get_api_client')
730777
async def test_delete_session_rejects_other_users_session():
@@ -753,7 +800,15 @@ async def test_session_id_path_traversal_rejected():
753800
"""Session IDs containing path-traversal characters must be rejected."""
754801
session_service = mock_vertex_ai_session_service()
755802

756-
for bad_id in ['..', '../foo', '..?force=true', 'a/b', '']:
803+
for bad_id in [
804+
'..',
805+
'../foo',
806+
'..?force=true',
807+
'a/b',
808+
'',
809+
'projects/my-project/locations/global/sessions/../foo',
810+
'reasoningEngines/123/sessions/1',
811+
]:
757812
with pytest.raises(ValueError):
758813
await session_service.delete_session(
759814
app_name='123', user_id='user', session_id=bad_id

0 commit comments

Comments
 (0)