From 0ec01956e86df6ae8e6553c70e410f1f8238ba88 Mon Sep 17 00:00:00 2001 From: George Weale Date: Mon, 17 Nov 2025 11:12:34 -0800 Subject: [PATCH] fix: keep vertex session event history intact Close #3504 Co-authored-by: George Weale PiperOrigin-RevId: 833417574 --- .../adk/sessions/vertex_ai_session_service.py | 11 ++++---- .../test_vertex_ai_session_service.py | 25 +++++++++++++++++++ 2 files changed, 31 insertions(+), 5 deletions(-) diff --git a/src/google/adk/sessions/vertex_ai_session_service.py b/src/google/adk/sessions/vertex_ai_session_service.py index 252a69e0..cce7e99b 100644 --- a/src/google/adk/sessions/vertex_ai_session_service.py +++ b/src/google/adk/sessions/vertex_ai_session_service.py @@ -176,11 +176,12 @@ class VertexAiSessionService(BaseSessionService): state=getattr(get_session_response, 'session_state', None) or {}, last_update_time=update_timestamp, ) - session.events += [ - _from_api_event(event) - async for event in events_iterator - if event.timestamp.timestamp() <= update_timestamp - ] + # Preserve the entire event stream that Vertex returns rather than trying + # to discard events written milliseconds after the session resource was + # updated. Clock skew between those writes can otherwise drop tool_result + # events and permanently break the replayed conversation. + async for event in events_iterator: + session.events.append(_from_api_event(event)) if config: # Filter events based on num_recent_events. diff --git a/tests/unittests/sessions/test_vertex_ai_session_service.py b/tests/unittests/sessions/test_vertex_ai_session_service.py index fa80dc9a..14d2b15b 100644 --- a/tests/unittests/sessions/test_vertex_ai_session_service.py +++ b/tests/unittests/sessions/test_vertex_ai_session_service.py @@ -531,6 +531,31 @@ async def test_get_session_with_after_timestamp_filter(): assert session.events[0].id == '456' +@pytest.mark.asyncio +@pytest.mark.usefixtures('mock_get_api_client') +async def test_get_session_keeps_events_newer_than_update_time( + mock_api_client_instance: MockAsyncClient, +) -> None: + future_event_time = isoparse( + MOCK_SESSION_JSON_1['update_time'] + ) + datetime.timedelta(seconds=1) + event = mock_api_client_instance.event_dict['1'][0][0] + event['timestamp'] = future_event_time.isoformat().replace('+00:00', 'Z') + session_service = mock_vertex_ai_session_service() + + session = await session_service.get_session( + app_name='123', user_id='user', session_id='1' + ) + + assert session is not None + assert len(session.events) == 1 + assert session.events[0].timestamp == future_event_time.timestamp() + assert session.events[0].timestamp > session.last_update_time, ( + 'Event timestamp should exceed session update_time to guard against' + ' filtering.' + ) + + @pytest.mark.asyncio @pytest.mark.usefixtures('mock_get_api_client') async def test_get_session_with_many_events(mock_api_client_instance):