fix(sessions): accept a full session resource name again on v1
The session id validation ported in #6809 requires a bare id, but Agent Engine passes the full projects/.../sessions/{id} resource name, so get_session, delete_session and create_session now reject it. Upstream hit the same thing and added _extract_short_session_id 13 days after the commit that was ported; v1 took the check without the follow-up. Ports that normalizer verbatim and calls it before validation at the three sites upstream patched. Ids that are not session resource names are unaffected.
This commit is contained in:
@@ -50,6 +50,29 @@ _USAGE_METADATA_CUSTOM_METADATA_KEY = '_usage_metadata'
|
||||
_SESSION_ID_PATTERN = re.compile(r'^[A-Za-z0-9_-]+$')
|
||||
|
||||
|
||||
def _extract_short_session_id(
|
||||
session_id: str, expected_engine_id: str | None = None
|
||||
) -> str:
|
||||
"""Extracts the short session ID if a full resource name is provided."""
|
||||
if isinstance(session_id, str) and '/' in session_id:
|
||||
parts = session_id.split('/')
|
||||
if len(parts) >= 2 and parts[-2] == 'sessions':
|
||||
if (
|
||||
len(parts) >= 4
|
||||
and parts[-4] == 'reasoningEngines'
|
||||
and expected_engine_id
|
||||
):
|
||||
passed_engine_id = parts[-3]
|
||||
if passed_engine_id != expected_engine_id:
|
||||
raise ValueError(
|
||||
'Session resource name mismatch: session belongs to '
|
||||
f'reasoningEngine {passed_engine_id!r}, but service is '
|
||||
f'configured for {expected_engine_id!r}.'
|
||||
)
|
||||
return parts[-1]
|
||||
return session_id
|
||||
|
||||
|
||||
def _validate_session_id(session_id: str) -> None:
|
||||
"""Rejects session IDs that could escape the URL path segment."""
|
||||
if not isinstance(session_id, str) or not _SESSION_ID_PATTERN.fullmatch(
|
||||
@@ -140,6 +163,9 @@ class VertexAiSessionService(BaseSessionService):
|
||||
|
||||
config = {'session_state': state} if state else {}
|
||||
if session_id:
|
||||
session_id = _extract_short_session_id(
|
||||
session_id, expected_engine_id=reasoning_engine_id
|
||||
)
|
||||
_validate_session_id(session_id)
|
||||
config['session_id'] = session_id
|
||||
config.update(kwargs)
|
||||
@@ -171,8 +197,11 @@ class VertexAiSessionService(BaseSessionService):
|
||||
session_id: str,
|
||||
config: Optional[GetSessionConfig] = None,
|
||||
) -> Optional[Session]:
|
||||
_validate_session_id(session_id)
|
||||
reasoning_engine_id = self._get_reasoning_engine_id(app_name)
|
||||
session_id = _extract_short_session_id(
|
||||
session_id, expected_engine_id=reasoning_engine_id
|
||||
)
|
||||
_validate_session_id(session_id)
|
||||
session_resource_name = (
|
||||
f'reasoningEngines/{reasoning_engine_id}/sessions/{session_id}'
|
||||
)
|
||||
@@ -271,8 +300,11 @@ class VertexAiSessionService(BaseSessionService):
|
||||
async def delete_session(
|
||||
self, *, app_name: str, user_id: str, session_id: str
|
||||
) -> None:
|
||||
_validate_session_id(session_id)
|
||||
reasoning_engine_id = self._get_reasoning_engine_id(app_name)
|
||||
session_id = _extract_short_session_id(
|
||||
session_id, expected_engine_id=reasoning_engine_id
|
||||
)
|
||||
_validate_session_id(session_id)
|
||||
session_resource_name = (
|
||||
f'reasoningEngines/{reasoning_engine_id}/sessions/{session_id}'
|
||||
)
|
||||
|
||||
@@ -31,6 +31,8 @@ from google.adk.events.event_actions import EventCompaction
|
||||
from google.adk.models.cache_metadata import CacheMetadata
|
||||
from google.adk.sessions.base_session_service import GetSessionConfig
|
||||
from google.adk.sessions.session import Session
|
||||
from google.adk.sessions.vertex_ai_session_service import _extract_short_session_id
|
||||
from google.adk.sessions.vertex_ai_session_service import _validate_session_id
|
||||
from google.adk.sessions.vertex_ai_session_service import VertexAiSessionService
|
||||
from google.api_core import exceptions as api_core_exceptions
|
||||
from google.genai import types as genai_types
|
||||
@@ -764,6 +766,70 @@ async def test_session_id_path_traversal_rejected():
|
||||
)
|
||||
|
||||
|
||||
def test_extract_short_session_id_short_id():
|
||||
assert _extract_short_session_id('123') == '123'
|
||||
assert _extract_short_session_id('session-123_abc') == 'session-123_abc'
|
||||
|
||||
|
||||
def test_extract_short_session_id_strips_full_resource_name():
|
||||
resource_name = (
|
||||
'projects/test-project/locations/test-location/'
|
||||
'reasoningEngines/123/sessions/3'
|
||||
)
|
||||
assert _extract_short_session_id(resource_name) == '3'
|
||||
assert (
|
||||
_extract_short_session_id(resource_name, expected_engine_id='123') == '3'
|
||||
)
|
||||
|
||||
|
||||
def test_extract_short_session_id_mismatch():
|
||||
resource_name = (
|
||||
'projects/test-project/locations/test-location/'
|
||||
'reasoningEngines/wrong/sessions/3'
|
||||
)
|
||||
with pytest.raises(ValueError, match='Session resource name mismatch'):
|
||||
_extract_short_session_id(resource_name, expected_engine_id='123')
|
||||
|
||||
|
||||
def test_validate_session_id_rejects_invalid_chars():
|
||||
with pytest.raises(ValueError, match='Invalid session_id'):
|
||||
_validate_session_id('invalid@id')
|
||||
with pytest.raises(ValueError, match='Invalid session_id'):
|
||||
_validate_session_id('invalid/id')
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.usefixtures('mock_get_api_client')
|
||||
async def test_get_session_accepts_a_full_resource_name():
|
||||
"""Agent Engine passes the full resource name, not the short id."""
|
||||
session_service = mock_vertex_ai_session_service()
|
||||
resource_name = (
|
||||
'projects/test-project/locations/test-location/'
|
||||
'reasoningEngines/123/sessions/1'
|
||||
)
|
||||
|
||||
session = await session_service.get_session(
|
||||
app_name='123', user_id='user', session_id=resource_name
|
||||
)
|
||||
|
||||
assert session.id == '1'
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.usefixtures('mock_get_api_client')
|
||||
async def test_get_session_rejects_a_resource_name_for_another_engine():
|
||||
session_service = mock_vertex_ai_session_service()
|
||||
resource_name = (
|
||||
'projects/test-project/locations/test-location/'
|
||||
'reasoningEngines/456/sessions/1'
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match='Session resource name mismatch'):
|
||||
await session_service.get_session(
|
||||
app_name='123', user_id='user', session_id=resource_name
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.usefixtures('mock_get_api_client')
|
||||
async def test_get_session_with_page_token():
|
||||
|
||||
Reference in New Issue
Block a user