Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
23 changes: 20 additions & 3 deletions src/google/adk/sessions/vertex_ai_session_service.py
Original file line number Diff line number Diff line change
Expand Up @@ -78,6 +78,21 @@ def _set_internal_custom_metadata(
}


def _drop_vertex_unsupported_part_fields(content_dict: dict[str, Any]) -> None:
"""Drops Part fields the Vertex AI Agent Engine Sessions API rejects.

``part_metadata`` is a Gemini Developer API-only field (the model path guards
it in ``genai`` ``_Part_to_vertex``); the Agent Engine Sessions API does not
accept it and fails ``appendEvent`` with ``400 INVALID_ARGUMENT`` ("Unknown
name \"part_metadata\" at 'event.content.parts[0]'"). Mutates the serialized
content dict in place; tolerant of either field-name or alias serialization.
"""
for part in content_dict.get('parts') or []:
if isinstance(part, dict):
part.pop('part_metadata', None)
part.pop('partMetadata', None)


class VertexAiSessionService(BaseSessionService):
"""Connects to the Vertex AI Agent Engine Session Service using Agent Engine SDK.

Expand Down Expand Up @@ -339,9 +354,9 @@ async def append_event(self, session: Session, event: Event) -> Event:
# Build config (Monolithic approach)
config = {}
if event.content:
config['content'] = event.content.model_dump(
exclude_none=True, mode='json'
)
content_dict = event.content.model_dump(exclude_none=True, mode='json')
_drop_vertex_unsupported_part_fields(content_dict)
config['content'] = content_dict
if event.actions:
config['actions'] = {
'skip_summarization': event.actions.skip_summarization,
Expand Down Expand Up @@ -406,6 +421,8 @@ async def append_event(self, session: Session, event: Event) -> Event:
mode='json',
by_alias=True,
)
if isinstance(config['raw_event'].get('content'), dict):
_drop_vertex_unsupported_part_fields(config['raw_event']['content'])

# Retry without raw_event if client side validation fails for older SDK
# versions.
Expand Down
80 changes: 80 additions & 0 deletions tests/unittests/sessions/test_vertex_ai_session_service.py
Original file line number Diff line number Diff line change
Expand Up @@ -1088,6 +1088,86 @@ async def test_append_event():
assert retrieved_session.events[1] == event_to_append


@pytest.mark.asyncio
@pytest.mark.usefixtures('mock_get_api_client')
async def test_append_event_strips_unsupported_part_metadata(
mock_api_client_instance: MockAsyncClient,
) -> None:
"""part_metadata must not reach the Sessions API (#6014).

``Part.part_metadata`` is a Gemini Developer API-only field; the Vertex AI
Agent Engine Sessions ``appendEvent`` API rejects it with 400 INVALID_ARGUMENT
("Unknown name \"part_metadata\""). It must be dropped from both the
``content`` and ``raw_event`` payloads, while the part text is preserved.
"""
session_service = mock_vertex_ai_session_service()
session = await session_service.get_session(
app_name='123', user_id='user', session_id='1'
)
event_to_append = Event(
invocation_id='inv_part_metadata',
author='user',
timestamp=1734005533.0,
content=genai_types.Content(
parts=[
genai_types.Part(
text='hello', part_metadata={'source': 'portal'}
),
genai_types.Part(text='world', part_metadata={'n': 1}),
],
),
)

await session_service.append_event(session, event_to_append)

appended = mock_api_client_instance.event_dict['1'][0][-1]
for part in appended['content']['parts']:
assert 'part_metadata' not in part
assert 'partMetadata' not in part
for part in appended['raw_event']['content']['parts']:
assert 'part_metadata' not in part
assert 'partMetadata' not in part
assert [p['text'] for p in appended['content']['parts']] == ['hello', 'world']


@pytest.mark.asyncio
@pytest.mark.usefixtures('mock_get_api_client')
async def test_append_event_with_part_metadata_round_trips(
mock_api_client_instance: MockAsyncClient,
) -> None:
"""Reconstruction side of #6014: an event carrying part_metadata appends and
reads back without error. part_metadata is dropped (unsupported on Vertex),
but the session round-trips and the part text is preserved.
"""
session_service = mock_vertex_ai_session_service()
session = await session_service.get_session(
app_name='123', user_id='user', session_id='1'
)
event_to_append = Event(
invocation_id='inv_part_metadata_rt',
author='user',
timestamp=1734005533.0,
content=genai_types.Content(
role='user',
parts=[
genai_types.Part(text='hello', part_metadata={'source': 'portal'})
],
),
)

await session_service.append_event(session, event_to_append)
retrieved = await session_service.get_session(
app_name='123', user_id='user', session_id='1'
)

appended = next(
e for e in retrieved.events if e.invocation_id == 'inv_part_metadata_rt'
)
assert appended.content is not None
assert appended.content.parts[0].text == 'hello'
assert appended.content.parts[0].part_metadata is None


@pytest.mark.asyncio
@pytest.mark.usefixtures('mock_get_api_client')
async def test_append_event_with_compaction():
Expand Down