From 01af10ab5c127e9fb6d3fc100c9f56b2c2490444 Mon Sep 17 00:00:00 2001 From: Yuting Lin Date: Tue, 24 Mar 2026 13:00:56 +0000 Subject: [PATCH 1/4] fix: include rewind_before_invocation_id in VertexAiSessionService.append_event (#4930) VertexAiSessionService.append_event() was not including rewind_before_invocation_id in the actions dict sent to the Vertex AI API. This caused rewind events to lose their target invocation ID when persisted, making session rewind non-functional for Agent Engine deployments. Added the missing field to the actions dict and a round-trip test. Fixes #4930 --- .../adk/sessions/vertex_ai_session_service.py | 3 +++ .../test_vertex_ai_session_service.py | 27 +++++++++++++++++++ 2 files changed, 30 insertions(+) diff --git a/src/google/adk/sessions/vertex_ai_session_service.py b/src/google/adk/sessions/vertex_ai_session_service.py index 9e5c9bb2ec..aff0b01277 100644 --- a/src/google/adk/sessions/vertex_ai_session_service.py +++ b/src/google/adk/sessions/vertex_ai_session_service.py @@ -282,6 +282,9 @@ async def append_event(self, session: Session, event: Event) -> Event: k: json.loads(v.model_dump_json(exclude_none=True, by_alias=True)) for k, v in event.actions.requested_auth_configs.items() }, + 'rewind_before_invocation_id': ( + event.actions.rewind_before_invocation_id + ), # TODO: add requested_tool_confirmations, agent_state once # they are available in the API. # Note: compaction is stored via event_metadata.custom_metadata. diff --git a/tests/unittests/sessions/test_vertex_ai_session_service.py b/tests/unittests/sessions/test_vertex_ai_session_service.py index 20fdbe3c6d..2282aa1799 100644 --- a/tests/unittests/sessions/test_vertex_ai_session_service.py +++ b/tests/unittests/sessions/test_vertex_ai_session_service.py @@ -829,6 +829,33 @@ 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_with_rewind(): + """rewind_before_invocation_id round-trips through append_event and get_session.""" + 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='rewind_invocation', + author='model', + timestamp=1734005533.0, + actions=EventActions( + rewind_before_invocation_id='target_invocation', + ), + ) + + await session_service.append_event(session, event_to_append) + + retrieved_session = await session_service.get_session( + app_name='123', user_id='user', session_id='1' + ) + + appended_event = retrieved_session.events[-1] + assert appended_event.actions.rewind_before_invocation_id == 'target_invocation' + + @pytest.mark.asyncio @pytest.mark.usefixtures('mock_get_api_client') async def test_append_event_with_compaction(): From a91252c5ab350e98e1c216b54ab6e70e1613df26 Mon Sep 17 00:00:00 2001 From: Yuting Lin Date: Tue, 24 Mar 2026 13:44:49 +0000 Subject: [PATCH 2/4] style: apply pyink formatting to test file --- tests/unittests/sessions/test_vertex_ai_session_service.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/unittests/sessions/test_vertex_ai_session_service.py b/tests/unittests/sessions/test_vertex_ai_session_service.py index 2282aa1799..d65ea3f0ab 100644 --- a/tests/unittests/sessions/test_vertex_ai_session_service.py +++ b/tests/unittests/sessions/test_vertex_ai_session_service.py @@ -853,7 +853,9 @@ async def test_append_event_with_rewind(): ) appended_event = retrieved_session.events[-1] - assert appended_event.actions.rewind_before_invocation_id == 'target_invocation' + assert ( + appended_event.actions.rewind_before_invocation_id == 'target_invocation' + ) @pytest.mark.asyncio From 26bb3865f305b4d20b03902e9af12087a85cac73 Mon Sep 17 00:00:00 2001 From: Yuting Lin Date: Tue, 24 Mar 2026 13:52:38 +0000 Subject: [PATCH 3/4] ci: retrigger CI checks From 7c2e4ce2b076fae0445ba85e21654813378a1a21 Mon Sep 17 00:00:00 2001 From: Yuting Lin Date: Wed, 1 Apr 2026 15:47:23 +0000 Subject: [PATCH 4/4] fix: store rewind_before_invocation_id in custom_metadata to avoid Pydantic validation error The Vertex AI SDK's Pydantic model does not yet support rewind_before_invocation_id in EventActions, causing 'Extra inputs are not permitted' errors. Store it in custom_metadata following the same pattern used for compaction. Addresses feedback from lucasbarzotto-axonify on PR #4978. --- .../adk/sessions/vertex_ai_session_service.py | 34 +++++++++++++++---- 1 file changed, 28 insertions(+), 6 deletions(-) diff --git a/src/google/adk/sessions/vertex_ai_session_service.py b/src/google/adk/sessions/vertex_ai_session_service.py index aff0b01277..a3efeee60b 100644 --- a/src/google/adk/sessions/vertex_ai_session_service.py +++ b/src/google/adk/sessions/vertex_ai_session_service.py @@ -43,6 +43,7 @@ logger = logging.getLogger('google_adk.' + __name__) _COMPACTION_CUSTOM_METADATA_KEY = '_compaction' +_REWIND_CUSTOM_METADATA_KEY = '_rewind_before_invocation_id' _USAGE_METADATA_CUSTOM_METADATA_KEY = '_usage_metadata' @@ -282,12 +283,10 @@ async def append_event(self, session: Session, event: Event) -> Event: k: json.loads(v.model_dump_json(exclude_none=True, by_alias=True)) for k, v in event.actions.requested_auth_configs.items() }, - 'rewind_before_invocation_id': ( - event.actions.rewind_before_invocation_id - ), # TODO: add requested_tool_confirmations, agent_state once # they are available in the API. - # Note: compaction is stored via event_metadata.custom_metadata. + # Note: compaction and rewind_before_invocation_id are stored via + # event_metadata.custom_metadata. } if event.error_code: config['error_code'] = event.error_code @@ -323,6 +322,16 @@ async def append_event(self, session: Session, event: Event) -> Event: key=_COMPACTION_CUSTOM_METADATA_KEY, value=compaction_dict, ) + # Store rewind_before_invocation_id in custom_metadata since the Vertex AI + # service does not yet support the field in EventActions. + # TODO: Stop writing to custom_metadata once the Vertex AI service + # supports rewind_before_invocation_id natively in EventActions. + if event.actions and event.actions.rewind_before_invocation_id: + _set_internal_custom_metadata( + metadata_dict, + key=_REWIND_CUSTOM_METADATA_KEY, + value=event.actions.rewind_before_invocation_id, + ) # Store usage_metadata in custom_metadata since the Vertex AI service # does not persist it in EventMetadata. if event.usage_metadata: @@ -408,15 +417,20 @@ def _from_api_event(api_event_obj: vertexai.types.SessionEvent) -> Event: # written before native compaction support store compaction data # in custom_metadata under the compaction metadata key. compaction_data = None + rewind_data = None usage_metadata_data = None if custom_metadata and ( _COMPACTION_CUSTOM_METADATA_KEY in custom_metadata + or _REWIND_CUSTOM_METADATA_KEY in custom_metadata or _USAGE_METADATA_CUSTOM_METADATA_KEY in custom_metadata ): custom_metadata = dict(custom_metadata) # avoid mutating the API response compaction_data = custom_metadata.pop( _COMPACTION_CUSTOM_METADATA_KEY, None ) + rewind_data = custom_metadata.pop( + _REWIND_CUSTOM_METADATA_KEY, None + ) usage_metadata_data = custom_metadata.pop( _USAGE_METADATA_CUSTOM_METADATA_KEY, None ) @@ -434,6 +448,7 @@ def _from_api_event(api_event_obj: vertexai.types.SessionEvent) -> Event: branch = None custom_metadata = None compaction_data = None + rewind_data = None usage_metadata_data = None grounding_metadata = None @@ -445,11 +460,18 @@ def _from_api_event(api_event_obj: vertexai.types.SessionEvent) -> Event: } if compaction_data: renamed_actions_dict['compaction'] = compaction_data + if rewind_data: + renamed_actions_dict['rewind_before_invocation_id'] = rewind_data event_actions = EventActions.model_validate(renamed_actions_dict) else: - if compaction_data: + if compaction_data or rewind_data: event_actions = EventActions( - compaction=EventCompaction.model_validate(compaction_data) + compaction=( + EventCompaction.model_validate(compaction_data) + if compaction_data + else None + ), + rewind_before_invocation_id=rewind_data, ) else: event_actions = EventActions()