fix: Ensure database sessions are always rolled back on errors

Fixes issue#3328

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

Author: wuliang229

src/google/adk/sessions/database_session_service.py

@@ -14,11 +14,13 @@
 from __future__ import annotations
 
 import asyncio
+from contextlib import asynccontextmanager
 import copy
 from datetime import datetime
 from datetime import timezone
 import logging
 from typing import Any
+from typing import AsyncIterator
 from typing import Optional
 
 from sqlalchemy import delete
@@ -156,6 +158,23 @@ def __init__(self, db_url: str, **kwargs: Any):
   def _get_schema_classes(self) -> _SchemaClasses:
     return _SchemaClasses(self._db_schema_version)
 
+  @asynccontextmanager
+  async def _rollback_on_exception_session(
+      self,
+  ) -> AsyncIterator[DatabaseSessionFactory]:
+    """Yields a database session with guaranteed rollback on errors.
+
+    On normal exit the caller is responsible for committing; on any exception
+    the transaction is explicitly rolled back before the error propagates,
+    preventing connection-pool exhaustion from lingering invalid transactions.
+    """
+    async with self.database_session_factory() as sql_session:
+      try:
+        yield sql_session
+      except BaseException:
+        await sql_session.rollback()
+        raise
+
   async def _prepare_tables(self):
     """Ensure database tables are ready for use.
 
@@ -204,7 +223,7 @@ async def _prepare_tables(self):
         self._tables_created = True
 
         if self._db_schema_version == _schema_check_utils.LATEST_SCHEMA_VERSION:
-          async with self.database_session_factory() as sql_session:
+          async with self._rollback_on_exception_session() as sql_session:
             # Check if schema version is set, if not, set it to the latest
             # version
             stmt = select(StorageMetadata).where(
@@ -236,7 +255,7 @@ async def create_session(
     # 5. Return the session
     await self._prepare_tables()
     schema = self._get_schema_classes()
-    async with self.database_session_factory() as sql_session:
+    async with self._rollback_on_exception_session() as sql_session:
       if session_id and await sql_session.get(
           schema.StorageSession, (app_name, user_id, session_id)
       ):
@@ -313,7 +332,7 @@ async def get_session(
     # 2. Get all the events based on session id and filtering config
     # 3. Convert and return the session
     schema = self._get_schema_classes()
-    async with self.database_session_factory() as sql_session:
+    async with self._rollback_on_exception_session() as sql_session:
       storage_session = await sql_session.get(
           schema.StorageSession, (app_name, user_id, session_id)
       )
@@ -368,7 +387,7 @@ async def list_sessions(
   ) -> ListSessionsResponse:
     await self._prepare_tables()
     schema = self._get_schema_classes()
-    async with self.database_session_factory() as sql_session:
+    async with self._rollback_on_exception_session() as sql_session:
       stmt = select(schema.StorageSession).filter(
           schema.StorageSession.app_name == app_name
       )
@@ -418,7 +437,7 @@ async def delete_session(
   ) -> None:
     await self._prepare_tables()
     schema = self._get_schema_classes()
-    async with self.database_session_factory() as sql_session:
+    async with self._rollback_on_exception_session() as sql_session:
       stmt = delete(schema.StorageSession).where(
           schema.StorageSession.app_name == app_name,
           schema.StorageSession.user_id == user_id,
@@ -440,7 +459,7 @@ async def append_event(self, session: Session, event: Event) -> Event:
     # 2. Update session attributes based on event config
     # 3. Store event to table
     schema = self._get_schema_classes()
-    async with self.database_session_factory() as sql_session:
+    async with self._rollback_on_exception_session() as sql_session:
       storage_session = await sql_session.get(
           schema.StorageSession, (session.app_name, session.user_id, session.id)
       )

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()