fix: Ensure database sessions are always rolled back on errors

Fixes issue#3328

Co-authored-by: Liang Wu <wuliang@google.com>
PiperOrigin-RevId: 863314939

tests/unittests/sessions/test_session_service.py

@@ -643,3 +643,216 @@ async def test_partial_events_are_not_persisted(session_service):
       app_name=app_name, user_id=user_id, session_id=session.id
   )
   assert len(session_got.events) == 0
+
+
+# ---------------------------------------------------------------------------
+# Rollback tests – verify _rollback_on_exception_session explicitly rolls back
+# on errors
+# ---------------------------------------------------------------------------
+class _RollbackSpySession:
+  """Wraps an AsyncSession to spy on rollback() and optionally fail commit()."""
+
+  def __init__(self, real_session, *, fail_commit=False):
+    self._real = real_session
+    self._fail_commit = fail_commit
+    self.rollback_called = False
+
+  async def __aenter__(self):
+    self._real = await self._real.__aenter__()
+    return self
+
+  async def __aexit__(self, *args):
+    return await self._real.__aexit__(*args)
+
+  async def commit(self):
+    if self._fail_commit:
+      raise RuntimeError('simulated commit failure')
+    return await self._real.commit()
+
+  async def rollback(self):
+    self.rollback_called = True
+    return await self._real.rollback()
+
+  def __getattr__(self, name):
+    return getattr(self._real, name)
+
+
+@pytest.mark.asyncio
+async def test_create_session_calls_rollback_on_commit_failure():
+  """Verifies that a commit failure during create_session triggers an explicit
+  rollback() call via _rollback_on_exception_session, not just a close()."""
+  service = DatabaseSessionService('sqlite+aiosqlite:///:memory:')
+  try:
+    # Ensure tables are initialized.
+    await service.create_session(
+        app_name='app', user_id='user', session_id='good'
+    )
+
+    original_factory = service.database_session_factory
+    spy_sessions = []
+
+    def _spy_factory():
+      spy = _RollbackSpySession(original_factory(), fail_commit=True)
+      spy_sessions.append(spy)
+      return spy
+
+    service.database_session_factory = _spy_factory
+
+    with pytest.raises(RuntimeError, match='simulated commit failure'):
+      await service.create_session(
+          app_name='app', user_id='user', session_id='should_fail'
+      )
+
+    # The key assertion: rollback() must have been called explicitly.
+    assert len(spy_sessions) == 1
+    assert spy_sessions[0].rollback_called, (
+        'rollback() was not called – _rollback_on_exception_session is not'
+        ' protecting this path'
+    )
+
+    # Restore and verify the failed session was not persisted.
+    service.database_session_factory = original_factory
+    assert (
+        await service.get_session(
+            app_name='app', user_id='user', session_id='should_fail'
+        )
+        is None
+    )
+  finally:
+    await service.close()
+
+
+@pytest.mark.asyncio
+async def test_append_event_calls_rollback_on_commit_failure():
+  """Verifies that a commit failure during append_event triggers an explicit
+  rollback() call via _rollback_on_exception_session."""
+  service = DatabaseSessionService('sqlite+aiosqlite:///:memory:')
+  try:
+    session = await service.create_session(
+        app_name='app', user_id='user', session_id='s1'
+    )
+
+    # Successfully append one event first.
+    event1 = Event(
+        invocation_id='inv1',
+        author='user',
+        actions=EventActions(state_delta={'key1': 'value1'}),
+    )
+    await service.append_event(session, event1)
+
+    original_factory = service.database_session_factory
+    spy_sessions = []
+
+    def _spy_factory():
+      spy = _RollbackSpySession(original_factory(), fail_commit=True)
+      spy_sessions.append(spy)
+      return spy
+
+    service.database_session_factory = _spy_factory
+
+    event2 = Event(
+        invocation_id='inv2',
+        author='user',
+        actions=EventActions(state_delta={'key2': 'value2'}),
+    )
+    with pytest.raises(RuntimeError, match='simulated commit failure'):
+      await service.append_event(session, event2)
+
+    assert len(spy_sessions) == 1
+    assert spy_sessions[0].rollback_called, (
+        'rollback() was not called – _rollback_on_exception_session is not'
+        ' protecting this path'
+    )
+
+    # Restore and verify only the first event was persisted.
+    service.database_session_factory = original_factory
+    got = await service.get_session(
+        app_name='app', user_id='user', session_id='s1'
+    )
+    assert len(got.events) == 1
+    assert got.events[0].invocation_id == 'inv1'
+  finally:
+    await service.close()
+
+
+@pytest.mark.asyncio
+async def test_delete_session_calls_rollback_on_commit_failure():
+  """Verifies that a commit failure during delete_session triggers an explicit
+  rollback() call via _rollback_on_exception_session."""
+  service = DatabaseSessionService('sqlite+aiosqlite:///:memory:')
+  try:
+    await service.create_session(
+        app_name='app', user_id='user', session_id='s1'
+    )
+
+    original_factory = service.database_session_factory
+    spy_sessions = []
+
+    def _spy_factory():
+      spy = _RollbackSpySession(original_factory(), fail_commit=True)
+      spy_sessions.append(spy)
+      return spy
+
+    service.database_session_factory = _spy_factory
+
+    with pytest.raises(RuntimeError, match='simulated commit failure'):
+      await service.delete_session(
+          app_name='app', user_id='user', session_id='s1'
+      )
+
+    assert len(spy_sessions) == 1
+    assert spy_sessions[0].rollback_called, (
+        'rollback() was not called – _rollback_on_exception_session is not'
+        ' protecting this path'
+    )
+
+    # Restore and verify the session still exists (delete was rolled back).
+    service.database_session_factory = original_factory
+    got = await service.get_session(
+        app_name='app', user_id='user', session_id='s1'
+    )
+    assert got is not None
+  finally:
+    await service.close()
+
+
+@pytest.mark.asyncio
+async def test_service_recovers_after_multiple_failures():
+  """After several consecutive commit failures, every single one must trigger
+  a rollback() call and the service must remain functional afterward."""
+  service = DatabaseSessionService('sqlite+aiosqlite:///:memory:')
+  try:
+    await service.create_session(
+        app_name='app', user_id='user', session_id='seed'
+    )
+
+    original_factory = service.database_session_factory
+    spy_sessions = []
+
+    def _spy_factory():
+      spy = _RollbackSpySession(original_factory(), fail_commit=True)
+      spy_sessions.append(spy)
+      return spy
+
+    service.database_session_factory = _spy_factory
+
+    num_failures = 5
+    for i in range(num_failures):
+      with pytest.raises(RuntimeError, match='simulated commit failure'):
+        await service.create_session(
+            app_name='app', user_id='user', session_id=f'fail_{i}'
+        )
+
+    # Every failure must have triggered a rollback.
+    assert len(spy_sessions) == num_failures
+    for i, spy in enumerate(spy_sessions):
+      assert spy.rollback_called, f'rollback() was not called on failure #{i}'
+
+    # Restore and verify the service is still healthy.
+    service.database_session_factory = original_factory
+    session = await service.create_session(
+        app_name='app', user_id='user', session_id='recovered'
+    )
+    assert session.id == 'recovered'
+  finally:
+    await service.close()