From cddc380cdc4c60c448ada54fc3a9cf9262749c23 Mon Sep 17 00:00:00 2001 From: galleonli Date: Thu, 8 Oct 2026 03:19:42 -0400 Subject: [PATCH] fix(sessions): refresh stale callers on duplicate event delivery --- .../adk/sessions/in_memory_session_service.py | 11 ++- .../sessions/test_session_service.py | 93 +++++++++++++++++++ tests/unittests/test_runners.py | 48 ++++++++++ 3 files changed, 151 insertions(+), 1 deletion(-) diff --git a/src/google/adk/sessions/in_memory_session_service.py b/src/google/adk/sessions/in_memory_session_service.py index a4f04b8c535..b29c26b34d5 100644 --- a/src/google/adk/sessions/in_memory_session_service.py +++ b/src/google/adk/sessions/in_memory_session_service.py @@ -342,7 +342,7 @@ async def append_event(self, session: Session, event: Event) -> Event: raise SessionNotFoundError(f'Session {session_id} not found.') # Fetch the canonical storage session early so we can drop a re-delivered - # event before modifying any state. The same event can be delivered more + # event before modifying stored state. The same event can be delivered more # than once when the orchestrator broadcasts a shared-state delta to # several concurrent session references; deduplicating here prevents # double-application of state updates and duplicate entries in event lists. @@ -352,6 +352,15 @@ async def append_event(self, session: Session, event: Event) -> Event: # stamp a fixed uuid) are not equal and are kept. storage_session = self.sessions[app_name][user_id][session_id] if any(e == event for e in storage_session.events if e.id == event.id): + if not any(e == event for e in session.events if e.id == event.id): + # The broadcast may still be new to this caller. Refresh its snapshot + # instead of replaying an old delta that could undo newer state updates. + refreshed = self._merge_state( + app_name, user_id, _copy_session(storage_session) + ) + session.state.update(refreshed.state) + session.events[:] = refreshed.events + session.last_update_time = refreshed.last_update_time return event # Update the in-memory session. diff --git a/tests/unittests/sessions/test_session_service.py b/tests/unittests/sessions/test_session_service.py index 9d5a673fe34..7c39df92a98 100644 --- a/tests/unittests/sessions/test_session_service.py +++ b/tests/unittests/sessions/test_session_service.py @@ -3895,6 +3895,99 @@ async def test_append_different_events_not_deduplicated(session_service): assert [e.author for e in retrieved.events] == ['user', 'agent'] +@pytest.mark.asyncio +@pytest.mark.parametrize('light_copy', [False, True]) +@pytest.mark.parametrize('copy_event', [False, True]) +async def test_duplicate_event_refreshes_another_session_reference( + light_copy, copy_event +): + """A broadcast updates each caller without duplicating the stored event.""" + override_feature_enabled( + FeatureName.IN_MEMORY_SESSION_SERVICE_LIGHT_COPY, light_copy + ) + try: + service = InMemorySessionService() + session_args = dict(app_name='app', user_id='user', session_id='session') + first = await service.create_session(**session_args, state={'counter': 0}) + second = await service.get_session(**session_args) + second.state['temp:local'] = 'keep' + state_reference = second.state + events_reference = second.events + event = Event( + author='agent', + timestamp=1.0, + actions=EventActions( + state_delta={ + 'counter': 1, + 'app:settings': {'enabled': True}, + 'user:profile': {'name': 'Alice'}, + } + ), + ) + + await service.append_event(session=first, event=event) + delivered = event.model_copy(deep=True) if copy_event else event + await service.append_event(session=second, event=delivered) + await service.append_event(session=second, event=delivered) + + stored = await service.get_session(**session_args) + assert second.state is state_reference + assert second.events is events_reference + assert second.state == {**stored.state, 'temp:local': 'keep'} + assert second.events == stored.events == [event] + assert second.last_update_time == stored.last_update_time == 1.0 + assert 'temp:local' not in stored.state + + # The refresh must honor the same copy depth as get_session. + second.state['app:settings']['enabled'] = False + second.state['user:profile']['name'] = 'Bob' + stored = await service.get_session(**session_args) + assert stored.state['app:settings']['enabled'] is (not light_copy) + assert stored.state['user:profile']['name'] == ( + 'Bob' if light_copy else 'Alice' + ) + finally: + override_feature_enabled( + FeatureName.IN_MEMORY_SESSION_SERVICE_LIGHT_COPY, False + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize('light_copy', [False, True]) +async def test_duplicate_event_refresh_preserves_newer_updates(light_copy): + """Delivering a missed older event must not roll back a newer update.""" + override_feature_enabled( + FeatureName.IN_MEMORY_SESSION_SERVICE_LIGHT_COPY, light_copy + ) + try: + service = InMemorySessionService() + session_args = dict(app_name='app', user_id='user', session_id='session') + first = await service.create_session(**session_args, state={'counter': 0}) + second = await service.get_session(**session_args) + older = Event( + author='agent', + timestamp=1.0, + actions=EventActions(state_delta={'counter': 1, 'result': 'ready'}), + ) + newer = Event( + author='agent', + timestamp=2.0, + actions=EventActions(state_delta={'counter': 2}), + ) + await service.append_event(session=first, event=older) + await service.append_event(session=second, event=newer) + await service.append_event(session=second, event=older) + + stored = await service.get_session(**session_args) + assert second.state == stored.state == {'counter': 2, 'result': 'ready'} + assert second.events == stored.events == [older, newer] + assert second.last_update_time == stored.last_update_time == 2.0 + finally: + override_feature_enabled( + FeatureName.IN_MEMORY_SESSION_SERVICE_LIGHT_COPY, False + ) + + @pytest.mark.asyncio @pytest.mark.parametrize( 'service_type', diff --git a/tests/unittests/test_runners.py b/tests/unittests/test_runners.py index a1abbc75cdd..bd97d4bd80d 100644 --- a/tests/unittests/test_runners.py +++ b/tests/unittests/test_runners.py @@ -70,6 +70,54 @@ TEST_SESSION_ID = "test_session" +async def test_concurrent_invocations_receive_a_shared_state_event(): + """Both invocation contexts see a broadcast that is stored only once.""" + barrier = asyncio.Barrier(2) + observed_counters = [] + broadcast = Event( + author="broadcast_agent", + actions=EventActions(state_delta={"counter": 1}), + ) + + class BroadcastAgent(BaseAgent): + + async def _run_async_impl( + self, ctx: InvocationContext + ) -> AsyncGenerator[Event, None]: + # Both callers must load their session before either receives the event. + await barrier.wait() + yield broadcast.model_copy(deep=True) + observed_counters.append(ctx.session.state["counter"]) + + session_service = InMemorySessionService() + session_args = dict( + app_name=TEST_APP_ID, user_id=TEST_USER_ID, session_id=TEST_SESSION_ID + ) + await session_service.create_session(**session_args, state={"counter": 0}) + runner = Runner( + app_name=TEST_APP_ID, + agent=BroadcastAgent(name="broadcast_agent"), + session_service=session_service, + ) + + async def invoke(): + async for _ in runner.run_async( + user_id=TEST_USER_ID, + session_id=TEST_SESSION_ID, + new_message=types.Content(role="user", parts=[types.Part(text="read")]), + ): + pass + + try: + await asyncio.wait_for(asyncio.gather(invoke(), invoke()), timeout=5) + stored = await session_service.get_session(**session_args) + assert observed_counters == [1, 1] + assert stored.state["counter"] == 1 + assert sum(e.id == broadcast.id for e in stored.events) == 1 + finally: + await runner.close() + + class MockAgent(BaseAgent): """Mock agent for unit testing."""