Skip to content

fix(sessions): defer local state/event mutation until Vertex append succeeds - #6999

Open
chelsealong wants to merge 1 commit into
google:mainfrom
chelsealong:fix-6998-vertex-session-append-order
Open

fix(sessions): defer local state/event mutation until Vertex append succeeds#6999
chelsealong wants to merge 1 commit into
google:mainfrom
chelsealong:fix-6998-vertex-session-append-order

Conversation

@chelsealong

Copy link
Copy Markdown
Contributor

Fixes #6998

Problem

VertexAiSessionService.append_event() called super().append_event()
— which applies the normal state delta, mutates temp-scoped state, and
appends the event to session.eventsbefore making the remote
Agent Engine append call. If the remote call raised, the exception
propagated to the caller, but the local session had already been
mutated. A retry of the same event after the remote service recovered
then re-applied the state delta and appended the event a second time.

DatabaseSessionService, SqliteSessionService, and
FirestoreSessionService all apply temp-scoped state (and strip it
from the event) before persisting, but defer the normal state update
and the event append until after storage succeeds. VertexAiSessionService
did not follow this pattern.

Fix

In VertexAiSessionService.append_event():

  • Apply temp-scoped state to the in-memory session and strip it from
    the outgoing event before the remote call (unchanged data still
    goes out over the wire, temp state remains available in-memory for
    the rest of the invocation).
  • Only apply the normal state delta and append the event to
    session.events after the remote append call succeeds.

This mirrors the ordering already used by the other three
BaseSessionService implementations.

Testing plan

Added test_append_event_does_not_mutate_session_on_remote_failure,
adapted directly from the issue's minimal reproduction: a remote
append call that fails once then succeeds.

Confirmed the test fails without the fix (stashed the source change,
kept the test):

$ python -m pytest tests/unittests/sessions/test_vertex_ai_session_service.py -k test_append_event_does_not_mutate_session_on_remote_failure -v
...
>       assert session.state == {'existing': 'value', 'temp:scratch': 'ephemeral'}
E       AssertionError: assert {'existing': ...: 'persisted'} == {'existing': ...: 'ephemeral'}
E         Left contains 1 more item:
E         {'normal': 'persisted'}
======================= 1 failed, 46 deselected in 2.52s =======================

And passes with the fix applied:

$ python -m pytest tests/unittests/sessions/test_vertex_ai_session_service.py -v
...
tests/unittests/sessions/test_vertex_ai_session_service.py::test_append_event_does_not_mutate_session_on_remote_failure PASSED
...
======================== 47 passed, 1 warning in 2.51s =========================

Full sessions unit test suite also passes:

$ python -m pytest tests/unittests/sessions/ -q
...
453 passed, 2 xfailed, 9 warnings in 8.32s

isort and pyink both report the changed files as clean.

AI assistance disclosure

This change was authored with the assistance of an AI coding agent
(Claude), including the fix implementation, regression test, and this
PR description. All changes were reviewed and verified locally before
submission (tests run and shown to fail without the fix / pass with
it, lint run).

…ucceeds

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 google#6998
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

VertexAiSessionService mutates local session before append_event succeeds

2 participants