From ab4297f0cdcbed23ee158c38511f6d536fd7a69a Mon Sep 17 00:00:00 2001 From: chelsealong Date: Thu, 3 Sep 2026 09:14:28 +0000 Subject: [PATCH] fix(sessions): defer local state/event mutation until Vertex append succeeds VertexAiSessionService.append_event() called the base class's append_event() (which applies normal state, appends the event, and mutates temp state) before the remote Agent Engine append. If the remote call failed, the exception propagated but the local session was already mutated, so a retry after recovery re-applied the state delta and appended the event a second time. Temp-scoped state is still applied to the in-memory session and stripped from the outgoing event before the remote call, matching DatabaseSessionService/SqliteSessionService/FirestoreSessionService. Normal state and the event are now only applied after the remote append succeeds. Fixes #6998 --- .../adk/sessions/vertex_ai_session_service.py | 14 ++++- .../test_vertex_ai_session_service.py | 58 +++++++++++++++++++ 2 files changed, 70 insertions(+), 2 deletions(-) diff --git a/src/google/adk/sessions/vertex_ai_session_service.py b/src/google/adk/sessions/vertex_ai_session_service.py index d36406d3797..c4df307aa0a 100644 --- a/src/google/adk/sessions/vertex_ai_session_service.py +++ b/src/google/adk/sessions/vertex_ai_session_service.py @@ -390,8 +390,14 @@ async def get_user_state( @override async def append_event(self, session: Session, event: Event) -> Event: - # Update the in-memory session. - await super().append_event(session=session, event=event) + if not event.partial: + # Apply temp-scoped state to the in-memory session and strip it from + # the event before the remote append succeeds. Normal state and the + # event itself are only applied to the session once the remote append + # succeeds, so a failed append leaves the session unchanged and a + # retry does not re-apply state or duplicate the event. + self._apply_temp_state(session, event) + event = self._trim_temp_delta_state(event) _validate_session_id(session.id) reasoning_engine_id = self._get_reasoning_engine_id(session.app_name) @@ -496,6 +502,10 @@ async def _do_append(cfg: dict[str, Any]) -> None: if 'raw_event' in config: del config['raw_event'] await _do_append(config) + + if not event.partial: + self._update_session_state(session, event) + session.events.append(event) return event def _get_reasoning_engine_id(self, app_name: str) -> str: diff --git a/tests/unittests/sessions/test_vertex_ai_session_service.py b/tests/unittests/sessions/test_vertex_ai_session_service.py index b589d8c3685..9b847961dd7 100644 --- a/tests/unittests/sessions/test_vertex_ai_session_service.py +++ b/tests/unittests/sessions/test_vertex_ai_session_service.py @@ -11,6 +11,8 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. +from collections.abc import AsyncIterator +from contextlib import asynccontextmanager import copy import datetime import re @@ -1165,6 +1167,62 @@ async def test_append_event(): assert retrieved_session.events[1] == event_to_append +@pytest.mark.asyncio +async def test_append_event_does_not_mutate_session_on_remote_failure() -> None: + """Regression test for #6998. + + A failed remote append must leave normal state and the event list + untouched (temp state remains, since it is invocation-local), and a + successful retry must apply the delta and append the event exactly once. + """ + append = mock.AsyncMock(side_effect=[RuntimeError('network failure'), None]) + client = types.SimpleNamespace( + agent_engines=types.SimpleNamespace( + sessions=types.SimpleNamespace( + events=types.SimpleNamespace(append=append), + ) + ) + ) + + @asynccontextmanager + async def fake_client() -> AsyncIterator[types.SimpleNamespace]: + yield client + + session_service = mock_vertex_ai_session_service() + session = Session( + id='1', + app_name='123', + user_id='user', + state={'existing': 'value'}, + ) + event = Event( + invocation_id='invocation', + author='model', + actions=EventActions( + state_delta={ + 'normal': 'persisted', + 'temp:scratch': 'ephemeral', + } + ), + ) + + with mock.patch.object(session_service, '_get_api_client', fake_client): + with pytest.raises(RuntimeError): + await session_service.append_event(session, event) + + assert session.state == {'existing': 'value', 'temp:scratch': 'ephemeral'} + assert len(session.events) == 0 + + await session_service.append_event(session, event) + + assert session.state == { + 'existing': 'value', + 'temp:scratch': 'ephemeral', + 'normal': 'persisted', + } + assert len(session.events) == 1 + + @pytest.mark.asyncio @pytest.mark.usefixtures('mock_get_api_client') async def test_append_event_strips_unsupported_part_metadata(