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(