test_session_service.py

# Copyright 2026 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# 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.

import asyncio
from contextlib import asynccontextmanager
from datetime import datetime
from datetime import timezone
import enum
import sqlite3
from unittest import mock

from google.adk.errors.already_exists_error import AlreadyExistsError
from google.adk.events.event import Event
from google.adk.events.event_actions import EventActions
from google.adk.features import FeatureName
from google.adk.features import override_feature_enabled
from google.adk.sessions import database_session_service
from google.adk.sessions.base_session_service import GetSessionConfig
from google.adk.sessions.database_session_service import DatabaseSessionService
from google.adk.sessions.in_memory_session_service import InMemorySessionService
from google.adk.sessions.sqlite_session_service import SqliteSessionService
from google.genai import types
import pytest
from sqlalchemy import delete


class SessionServiceType(enum.Enum):
  IN_MEMORY = 'IN_MEMORY'
  IN_MEMORY_WITH_LIGHT_COPY_ENABLED = 'IN_MEMORY_WITH_LIGHT_COPY_ENABLED'
  DATABASE = 'DATABASE'
  SQLITE = 'SQLITE'


def get_session_service(
    service_type: SessionServiceType = SessionServiceType.IN_MEMORY,
    tmp_path=None,
):
  """Creates a session service for testing."""
  if service_type == SessionServiceType.DATABASE:
    return DatabaseSessionService('sqlite+aiosqlite:///:memory:')
  if service_type == SessionServiceType.SQLITE:
    return SqliteSessionService(str(tmp_path / 'sqlite.db'))
  if service_type == SessionServiceType.IN_MEMORY_WITH_LIGHT_COPY_ENABLED:
    return InMemorySessionService()
  return InMemorySessionService()


@pytest.fixture(
    params=[
        SessionServiceType.IN_MEMORY,
        SessionServiceType.IN_MEMORY_WITH_LIGHT_COPY_ENABLED,
        SessionServiceType.DATABASE,
        SessionServiceType.SQLITE,
    ]
)
async def session_service(request, tmp_path):
  """Provides a session service and closes database backends on teardown."""
  if request.param == SessionServiceType.IN_MEMORY_WITH_LIGHT_COPY_ENABLED:
    override_feature_enabled(
        FeatureName.IN_MEMORY_SESSION_SERVICE_LIGHT_COPY, True
    )
  service = get_session_service(request.param, tmp_path)
  yield service
  if isinstance(service, DatabaseSessionService):
    await service.close()
  if request.param == SessionServiceType.IN_MEMORY_WITH_LIGHT_COPY_ENABLED:
    override_feature_enabled(
        FeatureName.IN_MEMORY_SESSION_SERVICE_LIGHT_COPY, False
    )


def test_database_session_service_enables_pool_pre_ping_by_default():
  captured_kwargs = {}

  def fake_create_async_engine(_db_url: str, **kwargs):
    captured_kwargs.update(kwargs)
    fake_engine = mock.Mock()
    fake_engine.dialect.name = 'postgresql'
    fake_engine.sync_engine = mock.Mock()
    return fake_engine

  with mock.patch.object(
      database_session_service,
      'create_async_engine',
      side_effect=fake_create_async_engine,
  ):
    database_session_service.DatabaseSessionService(
        'postgresql+psycopg2://user:pass@localhost:5432/db'
    )

  assert captured_kwargs.get('pool_pre_ping') is True


@pytest.mark.parametrize('dialect_name', ['sqlite', 'postgresql'])
def test_database_session_service_strips_timezone_for_dialect(dialect_name):
  """Verifies that timezone-aware datetimes are converted to naive datetimes
  for SQLite and PostgreSQL to avoid 'can't subtract offset-naive and
  offset-aware datetimes' errors.

  PostgreSQL's default TIMESTAMP type is WITHOUT TIME ZONE, which cannot
  accept timezone-aware datetime objects when using asyncpg. SQLite also
  requires naive datetimes.
  """
  # Simulate the logic in create_session
  is_sqlite = dialect_name == 'sqlite'
  is_postgres = dialect_name == 'postgresql'

  now = datetime.now(timezone.utc)
  assert now.tzinfo is not None  # Starts with timezone

  if is_sqlite or is_postgres:
    now = now.replace(tzinfo=None)

  # Both SQLite and PostgreSQL should have timezone stripped
  assert now.tzinfo is None


def test_database_session_service_preserves_timezone_for_other_dialects():
  """Verifies that timezone info is preserved for dialects that support it."""
  # For dialects like MySQL with explicit timezone support, we don't strip
  dialect_name = 'mysql'
  is_sqlite = dialect_name == 'sqlite'
  is_postgres = dialect_name == 'postgresql'

  now = datetime.now(timezone.utc)
  assert now.tzinfo is not None

  if is_sqlite or is_postgres:
    now = now.replace(tzinfo=None)

  # MySQL should preserve timezone (if the column type supports it)
  assert now.tzinfo is not None


def test_database_session_service_respects_pool_pre_ping_override():
  captured_kwargs = {}

  def fake_create_async_engine(_db_url: str, **kwargs):
    captured_kwargs.update(kwargs)
    fake_engine = mock.Mock()
    fake_engine.dialect.name = 'postgresql'
    fake_engine.sync_engine = mock.Mock()
    return fake_engine

  with mock.patch.object(
      database_session_service,
      'create_async_engine',
      side_effect=fake_create_async_engine,
  ):
    database_session_service.DatabaseSessionService(
        'postgresql+psycopg2://user:pass@localhost:5432/db',
        pool_pre_ping=False,
    )

  assert captured_kwargs.get('pool_pre_ping') is False


def test_database_session_service_creates_read_only_engine_for_spanner():
  captured_binds = []
  fake_engine = mock.Mock()
  fake_engine.dialect.name = 'spanner+spanner'
  fake_engine.sync_engine = mock.Mock()
  read_only_engine = mock.Mock()
  fake_engine.execution_options.return_value = read_only_engine

  def fake_async_sessionmaker(*, bind, expire_on_commit, **kwargs):
    del expire_on_commit
    del kwargs
    captured_binds.append(bind)
    return mock.Mock()

  with (
      mock.patch.object(
          database_session_service,
          'create_async_engine',
          return_value=fake_engine,
      ),
      mock.patch.object(
          database_session_service,
          'async_sessionmaker',
          side_effect=fake_async_sessionmaker,
      ),
  ):
    database_session_service.DatabaseSessionService(
        'spanner+spanner:///projects/test/instances/test/databases/test'
    )

  assert captured_binds == [fake_engine, read_only_engine]
  fake_engine.execution_options.assert_called_once_with(read_only=True)


def test_database_session_service_creates_read_only_engine_for_other_dialects():
  captured_binds = []
  fake_engine = mock.Mock()
  fake_engine.dialect.name = 'postgresql'
  fake_engine.sync_engine = mock.Mock()
  read_only_engine = mock.Mock()
  fake_engine.execution_options.return_value = read_only_engine

  def fake_async_sessionmaker(*, bind, expire_on_commit, **kwargs):
    del expire_on_commit
    del kwargs
    captured_binds.append(bind)
    return mock.Mock()

  with (
      mock.patch.object(
          database_session_service,
          'create_async_engine',
          return_value=fake_engine,
      ),
      mock.patch.object(
          database_session_service,
          'async_sessionmaker',
          side_effect=fake_async_sessionmaker,
      ),
  ):
    database_session_service.DatabaseSessionService(
        'postgresql+psycopg2://user:pass@localhost:5432/db'
    )

  assert captured_binds == [fake_engine, read_only_engine]
  fake_engine.execution_options.assert_called_once_with(read_only=True)


@pytest.mark.asyncio
async def test_sqlite_session_service_accepts_sqlite_urls(
    tmp_path, monkeypatch
):
  monkeypatch.chdir(tmp_path)

  service = SqliteSessionService('sqlite+aiosqlite:///./sessions.db')
  await service.create_session(app_name='app', user_id='user')
  assert (tmp_path / 'sessions.db').exists()

  service = SqliteSessionService('sqlite:///./sessions2.db')
  await service.create_session(app_name='app', user_id='user')
  assert (tmp_path / 'sessions2.db').exists()


@pytest.mark.asyncio
async def test_sqlite_session_service_preserves_uri_query_parameters(
    tmp_path, monkeypatch
):
  monkeypatch.chdir(tmp_path)
  db_path = tmp_path / 'readonly.db'
  with sqlite3.connect(db_path) as conn:
    conn.execute('CREATE TABLE IF NOT EXISTS t (id INTEGER)')
    conn.commit()

  service = SqliteSessionService(f'sqlite+aiosqlite:///{db_path}?mode=ro')
  # `mode=ro` opens the DB read-only; schema creation should fail.
  with pytest.raises(sqlite3.OperationalError, match=r'readonly'):
    await service.create_session(app_name='app', user_id='user')


@pytest.mark.asyncio
async def test_sqlite_session_service_accepts_absolute_sqlite_urls(tmp_path):
  abs_db_path = tmp_path / 'absolute.db'
  abs_url = 'sqlite+aiosqlite:////' + str(abs_db_path).lstrip('/')
  service = SqliteSessionService(abs_url)
  await service.create_session(app_name='app', user_id='user')
  assert abs_db_path.exists()


@pytest.mark.asyncio
async def test_get_empty_session(session_service):
  assert not await session_service.get_session(
      app_name='my_app', user_id='test_user', session_id='123'
  )


@pytest.mark.asyncio
async def test_database_session_service_get_session_uses_read_only_factory():
  service = DatabaseSessionService('sqlite+aiosqlite:///:memory:')
  service._prepare_tables = mock.AsyncMock()

  read_only_session = mock.AsyncMock()
  read_only_session.get = mock.AsyncMock(return_value=None)

  @asynccontextmanager
  async def fake_read_only_session():
    yield read_only_session

  service.database_session_factory = mock.Mock(
      side_effect=AssertionError('write session factory should not be used')
  )
  service._read_only_database_session_factory = mock.Mock(
      return_value=fake_read_only_session()
  )

  session = await service.get_session(
      app_name='my_app', user_id='test_user', session_id='123'
  )

  assert session is None
  service._read_only_database_session_factory.assert_called_once_with()
  service.database_session_factory.assert_not_called()

  await service.close()


@pytest.mark.asyncio
async def test_database_session_service_list_sessions_uses_read_only_factory():
  service = DatabaseSessionService('sqlite+aiosqlite:///:memory:')
  service._prepare_tables = mock.AsyncMock()

  read_only_session = mock.AsyncMock()
  empty_result = mock.Mock()
  empty_result.scalars.return_value.all.return_value = []
  read_only_session.execute = mock.AsyncMock(return_value=empty_result)
  read_only_session.get = mock.AsyncMock(return_value=None)

  @asynccontextmanager
  async def fake_read_only_session():
    yield read_only_session

  service.database_session_factory = mock.Mock(
      side_effect=AssertionError('write session factory should not be used')
  )
  service._read_only_database_session_factory = mock.Mock(
      return_value=fake_read_only_session()
  )

  response = await service.list_sessions(app_name='my_app', user_id='test_user')

  assert response.sessions == []
  service._read_only_database_session_factory.assert_called_once_with()
  service.database_session_factory.assert_not_called()

  await service.close()


@pytest.mark.asyncio
async def test_create_get_session(session_service):
  app_name = 'my_app'
  user_id = 'test_user'
  state = {'key': 'value'}

  session = await session_service.create_session(
      app_name=app_name, user_id=user_id, state=state
  )
  assert session.app_name == app_name
  assert session.user_id == user_id
  assert session.id
  assert session.state == state
  assert (
      session.last_update_time
      <= datetime.now().astimezone(timezone.utc).timestamp()
  )

  got_session = await session_service.get_session(
      app_name=app_name, user_id=user_id, session_id=session.id
  )
  assert got_session == session
  assert (
      got_session.last_update_time
      <= datetime.now().astimezone(timezone.utc).timestamp()
  )

  session_id = session.id
  await session_service.delete_session(
      app_name=app_name, user_id=user_id, session_id=session_id
  )

  assert (
      await session_service.get_session(
          app_name=app_name, user_id=user_id, session_id=session.id
      )
      is None
  )


@pytest.mark.asyncio
async def test_create_and_list_sessions(session_service):
  app_name = 'my_app'
  user_id = 'test_user'

  session_ids = ['session' + str(i) for i in range(5)]
  for session_id in session_ids:
    await session_service.create_session(
        app_name=app_name,
        user_id=user_id,
        session_id=session_id,
        state={'key': 'value' + session_id},
    )

  list_sessions_response = await session_service.list_sessions(
      app_name=app_name, user_id=user_id
  )
  sessions = list_sessions_response.sessions
  assert len(sessions) == len(session_ids)
  assert {s.id for s in sessions} == set(session_ids)
  for session in sessions:
    assert session.state == {'key': 'value' + session.id}


@pytest.mark.asyncio
async def test_list_sessions_all_users(session_service):
  app_name = 'my_app'
  user_id_1 = 'user1'
  user_id_2 = 'user2'

  await session_service.create_session(
      app_name=app_name,
      user_id=user_id_1,
      session_id='session1a',
      state={'key': 'value1a'},
  )
  await session_service.create_session(
      app_name=app_name,
      user_id=user_id_1,
      session_id='session1b',
      state={'key': 'value1b'},
  )
  await session_service.create_session(
      app_name=app_name,
      user_id=user_id_2,
      session_id='session2a',
      state={'key': 'value2a'},
  )

  # List sessions for user1 - should contain merged state
  list_sessions_response_1 = await session_service.list_sessions(
      app_name=app_name, user_id=user_id_1
  )
  sessions_1 = list_sessions_response_1.sessions
  assert len(sessions_1) == 2
  sessions_1_map = {s.id: s for s in sessions_1}
  assert sessions_1_map['session1a'].state == {'key': 'value1a'}
  assert sessions_1_map['session1b'].state == {'key': 'value1b'}

  # List sessions for user2 - should contain merged state
  list_sessions_response_2 = await session_service.list_sessions(
      app_name=app_name, user_id=user_id_2
  )
  sessions_2 = list_sessions_response_2.sessions
  assert len(sessions_2) == 1
  assert sessions_2[0].id == 'session2a'
  assert sessions_2[0].state == {'key': 'value2a'}

  # List sessions for all users - should contain merged state
  list_sessions_response_all = await session_service.list_sessions(
      app_name=app_name, user_id=None
  )
  sessions_all = list_sessions_response_all.sessions
  assert len(sessions_all) == 3
  sessions_all_map = {s.id: s for s in sessions_all}
  assert sessions_all_map['session1a'].state == {'key': 'value1a'}
  assert sessions_all_map['session1b'].state == {'key': 'value1b'}
  assert sessions_all_map['session2a'].state == {'key': 'value2a'}


@pytest.mark.asyncio
async def test_app_state_is_shared_by_all_users_of_app(session_service):
  app_name = 'my_app'
  # User 1 creates a session, establishing app:k1
  session1 = await session_service.create_session(
      app_name=app_name, user_id='u1', session_id='s1', state={'app:k1': 'v1'}
  )
  # User 1 appends an event to session1, establishing app:k2
  event = Event(
      invocation_id='inv1',
      author='user',
      actions=EventActions(state_delta={'app:k2': 'v2'}),
  )
  await session_service.append_event(session=session1, event=event)

  # User 2 creates a new session session2, it should see app:k1 and app:k2
  session2 = await session_service.create_session(
      app_name=app_name, user_id='u2', session_id='s2'
  )
  assert session2.state == {'app:k1': 'v1', 'app:k2': 'v2'}

  # If we get session session1 again, it should also see both
  session1_got = await session_service.get_session(
      app_name=app_name, user_id='u1', session_id='s1'
  )
  assert session1_got.state.get('app:k1') == 'v1'
  assert session1_got.state.get('app:k2') == 'v2'


@pytest.mark.asyncio
async def test_user_state_is_shared_only_by_user_sessions(session_service):
  app_name = 'my_app'
  # User 1 creates a session, establishing user:k1 for user 1
  session1 = await session_service.create_session(
      app_name=app_name, user_id='u1', session_id='s1', state={'user:k1': 'v1'}
  )
  # User 1 appends an event to session1, establishing user:k2 for user 1
  event = Event(
      invocation_id='inv1',
      author='user',
      actions=EventActions(state_delta={'user:k2': 'v2'}),
  )
  await session_service.append_event(session=session1, event=event)

  # Another session for User 1 should see user:k1 and user:k2
  session1b = await session_service.create_session(
      app_name=app_name, user_id='u1', session_id='s1b'
  )
  assert session1b.state == {'user:k1': 'v1', 'user:k2': 'v2'}

  # A session for User 2 should NOT see user:k1 or user:k2
  session2 = await session_service.create_session(
      app_name=app_name, user_id='u2', session_id='s2'
  )
  assert session2.state == {}


@pytest.mark.asyncio
async def test_session_state_is_not_shared(session_service):
  app_name = 'my_app'
  # User 1 creates a session session1, establishing sk1 only for session1
  session1 = await session_service.create_session(
      app_name=app_name, user_id='u1', session_id='s1', state={'sk1': 'v1'}
  )
  # User 1 appends an event to session1, establishing sk2 only for session1
  event = Event(
      invocation_id='inv1',
      author='user',
      actions=EventActions(state_delta={'sk2': 'v2'}),
  )
  await session_service.append_event(session=session1, event=event)

  # Getting session1 should show sk1 and sk2
  session1_got = await session_service.get_session(
      app_name=app_name, user_id='u1', session_id='s1'
  )
  assert session1_got.state.get('sk1') == 'v1'
  assert session1_got.state.get('sk2') == 'v2'

  # Creating another session session1b for User 1 should NOT see sk1 or sk2
  session1b = await session_service.create_session(
      app_name=app_name, user_id='u1', session_id='s1b'
  )
  assert session1b.state == {}


@pytest.mark.asyncio
async def test_temp_state_is_not_persisted_in_state_or_events(session_service):
  app_name = 'my_app'
  user_id = 'u1'
  session = await session_service.create_session(
      app_name=app_name, user_id=user_id, session_id='s1'
  )
  event = Event(
      invocation_id='inv1',
      author='user',
      actions=EventActions(state_delta={'temp:k1': 'v1', 'sk': 'v2'}),
  )
  await session_service.append_event(session=session, event=event)

  # Temp state IS available in the in-memory session (same invocation)
  assert session.state.get('temp:k1') == 'v1'
  assert session.state.get('sk') == 'v2'

  # Check event as stored in session does not contain temp keys in state_delta
  assert 'temp:k1' not in event.actions.state_delta
  assert event.actions.state_delta.get('sk') == 'v2'


@pytest.mark.asyncio
async def test_temp_state_visible_across_sequential_events(session_service):
  """Temp state set by one event should be readable before the next event.

  This simulates a SequentialAgent where agent-1 writes output_key='temp:out'
  and agent-2 needs to read it from session.state within the same invocation.
  """
  app_name = 'my_app'
  user_id = 'u1'
  session = await session_service.create_session(
      app_name=app_name, user_id=user_id, session_id='s_seq'
  )

  # Agent-1 writes temp state
  event1 = Event(
      invocation_id='inv1',
      author='agent1',
      actions=EventActions(state_delta={'temp:output': 'result_from_a1'}),
  )
  await session_service.append_event(session=session, event=event1)

  # Agent-2 should be able to read temp state from the same session object
  assert session.state.get('temp:output') == 'result_from_a1'

  # But the event delta should NOT contain the temp key (not persisted)
  assert 'temp:output' not in event1.actions.state_delta


@pytest.mark.asyncio
async def test_get_session_respects_user_id(session_service):
  app_name = 'my_app'
  # u1 creates session 's1' and adds an event
  session1 = await session_service.create_session(
      app_name=app_name, user_id='u1', session_id='s1'
  )
  event = Event(invocation_id='inv1', author='user')
  await session_service.append_event(session1, event)
  # u2 creates a session with the same session_id 's1'
  await session_service.create_session(
      app_name=app_name, user_id='u2', session_id='s1'
  )
  # Check that getting s1 for u2 returns u2's session (with no events)
  # not u1's session.
  session2_got = await session_service.get_session(
      app_name=app_name, user_id='u2', session_id='s1'
  )
  assert session2_got.user_id == 'u2'
  assert len(session2_got.events) == 0


@pytest.mark.asyncio
async def test_create_session_with_existing_id_raises_error(session_service):
  app_name = 'my_app'
  user_id = 'test_user'
  session_id = 'existing_session'

  # Create the first session
  await session_service.create_session(
      app_name=app_name,
      user_id=user_id,
      session_id=session_id,
  )

  # Attempt to create a session with the same ID
  with pytest.raises(AlreadyExistsError):
    await session_service.create_session(
        app_name=app_name,
        user_id=user_id,
        session_id=session_id,
    )


@pytest.mark.asyncio
async def test_append_event_bytes(session_service):
  app_name = 'my_app'
  user_id = 'user'

  session = await session_service.create_session(
      app_name=app_name, user_id=user_id
  )

  test_content = types.Content(
      role='user',
      parts=[
          types.Part.from_bytes(data=b'test_image_data', mime_type='image/png'),
      ],
  )
  test_grounding_metadata = types.GroundingMetadata(
      search_entry_point=types.SearchEntryPoint(sdk_blob=b'test_sdk_blob')
  )
  event = Event(
      invocation_id='invocation',
      author='user',
      content=test_content,
      grounding_metadata=test_grounding_metadata,
  )
  await session_service.append_event(session=session, event=event)

  assert session.events[0].content == test_content

  session = await session_service.get_session(
      app_name=app_name, user_id=user_id, session_id=session.id
  )
  events = session.events
  assert len(events) == 1
  assert events[0].content == test_content
  assert events[0].grounding_metadata == test_grounding_metadata


@pytest.mark.asyncio
async def test_append_event_complete(session_service):
  app_name = 'my_app'
  user_id = 'user'

  session = await session_service.create_session(
      app_name=app_name, user_id=user_id
  )
  event = Event(
      invocation_id='invocation',
      author='user',
      content=types.Content(role='user', parts=[types.Part(text='test_text')]),
      turn_complete=True,
      partial=False,
      actions=EventActions(
          artifact_delta={
              'file': 0,
          },
          transfer_to_agent='agent',
          escalate=True,
      ),
      long_running_tool_ids={'tool1'},
      error_code='error_code',
      error_message='error_message',
      interrupted=True,
      grounding_metadata=types.GroundingMetadata(
          web_search_queries=['query1'],
      ),
      usage_metadata=types.GenerateContentResponseUsageMetadata(
          prompt_token_count=1, candidates_token_count=1, total_token_count=2
      ),
      citation_metadata=types.CitationMetadata(),
      custom_metadata={'custom_key': 'custom_value'},
      timestamp=1700000000.123,
      input_transcription=types.Transcription(
          text='input transcription',
          finished=True,
      ),
      output_transcription=types.Transcription(
          text='output transcription',
          finished=True,
      ),
  )
  await session_service.append_event(session=session, event=event)

  assert (
      await session_service.get_session(
          app_name=app_name, user_id=user_id, session_id=session.id
      )
      == session
  )


@pytest.mark.asyncio
async def test_session_last_update_time_updates_on_event(session_service):
  app_name = 'my_app'
  user_id = 'user'

  session = await session_service.create_session(
      app_name=app_name, user_id=user_id
  )
  original_update_time = session.last_update_time

  event_timestamp = original_update_time + 10
  event = Event(
      invocation_id='invocation',
      author='user',
      timestamp=event_timestamp,
  )
  await session_service.append_event(session=session, event=event)

  assert session.last_update_time == pytest.approx(event_timestamp, abs=1e-6)

  refreshed_session = await session_service.get_session(
      app_name=app_name, user_id=user_id, session_id=session.id
  )
  assert refreshed_session is not None
  assert refreshed_session.last_update_time == pytest.approx(
      event_timestamp, abs=1e-6
  )
  assert refreshed_session.last_update_time > original_update_time


@pytest.mark.asyncio
async def test_append_event_to_stale_session():
  session_service = get_session_service(
      service_type=SessionServiceType.DATABASE
  )

  async with session_service:
    app_name = 'my_app'
    user_id = 'user'
    current_time = datetime.now().astimezone(timezone.utc).timestamp()

    original_session = await session_service.create_session(
        app_name=app_name, user_id=user_id
    )
    event1 = Event(
        invocation_id='inv1',
        author='user',
        timestamp=current_time + 1,
        actions=EventActions(state_delta={'sk1': 'v1'}),
    )
    await session_service.append_event(original_session, event1)

    updated_session = await session_service.get_session(
        app_name=app_name, user_id=user_id, session_id=original_session.id
    )
    event2 = Event(
        invocation_id='inv2',
        author='user',
        timestamp=current_time + 2,
        actions=EventActions(state_delta={'sk2': 'v2'}),
    )
    await session_service.append_event(updated_session, event2)

    # original_session is now stale
    assert original_session.last_update_time < updated_session.last_update_time
    assert len(original_session.events) == 1
    assert 'sk2' not in original_session.state

    # Appending another event to stale original_session should be rejected.
    event3 = Event(
        invocation_id='inv3',
        author='user',
        timestamp=current_time + 3,
        actions=EventActions(state_delta={'sk3': 'v3'}),
    )
    with pytest.raises(ValueError, match='modified in storage'):
      await session_service.append_event(original_session, event3)

    # If we fetch session from DB, it should only contain the committed events.
    session_final = await session_service.get_session(
        app_name=app_name, user_id=user_id, session_id=original_session.id
    )
    assert len(session_final.events) == 2
    assert session_final.state.get('sk1') == 'v1'
    assert session_final.state.get('sk2') == 'v2'
    assert session_final.state.get('sk3') is None
    assert [e.invocation_id for e in session_final.events] == [
        'inv1',
        'inv2',
    ]


@pytest.mark.asyncio
async def test_append_event_raises_if_app_state_row_missing():
  service = DatabaseSessionService('sqlite+aiosqlite:///:memory:')
  try:
    session = await service.create_session(
        app_name='my_app', user_id='user', session_id='s1'
    )
    schema = service._get_schema_classes()
    async with service.database_session_factory() as sql_session:
      await sql_session.execute(
          delete(schema.StorageAppState).where(
              schema.StorageAppState.app_name == session.app_name
          )
      )
      await sql_session.commit()

    event = Event(
        invocation_id='inv1',
        author='user',
        actions=EventActions(state_delta={'k': 'v'}),
    )
    with pytest.raises(ValueError, match='App state missing'):
      await service.append_event(session, event)
  finally:
    await service.close()


@pytest.mark.asyncio
async def test_append_event_raises_if_user_state_row_missing():
  service = DatabaseSessionService('sqlite+aiosqlite:///:memory:')
  try:
    session = await service.create_session(
        app_name='my_app', user_id='user', session_id='s1'
    )
    schema = service._get_schema_classes()
    async with service.database_session_factory() as sql_session:
      await sql_session.execute(
          delete(schema.StorageUserState).where(
              schema.StorageUserState.app_name == session.app_name,
              schema.StorageUserState.user_id == session.user_id,
          )
      )
      await sql_session.commit()

    event = Event(
        invocation_id='inv1',
        author='user',
        actions=EventActions(state_delta={'k': 'v'}),
    )
    with pytest.raises(ValueError, match='User state missing'):
      await service.append_event(session, event)
  finally:
    await service.close()


@pytest.mark.asyncio
async def test_append_event_concurrent_stale_sessions_reject_stale_writer():
  session_service = get_session_service(
      service_type=SessionServiceType.DATABASE
  )

  async with session_service:
    app_name = 'my_app'
    user_id = 'user'
    session = await session_service.create_session(
        app_name=app_name, user_id=user_id
    )

    iteration_count = 8
    for i in range(iteration_count):
      latest_session = await session_service.get_session(
          app_name=app_name, user_id=user_id, session_id=session.id
      )
      stale_session_1 = latest_session.model_copy(deep=True)
      stale_session_2 = latest_session.model_copy(deep=True)
      base_timestamp = latest_session.last_update_time + 10.0
      event_1 = Event(
          invocation_id=f'inv{i}-1',
          author='user',
          timestamp=base_timestamp + 1.0,
          actions=EventActions(state_delta={f'sk{i}-1': f'v{i}-1'}),
      )
      event_2 = Event(
          invocation_id=f'inv{i}-2',
          author='user',
          timestamp=base_timestamp + 2.0,
          actions=EventActions(state_delta={f'sk{i}-2': f'v{i}-2'}),
      )

      results = await asyncio.gather(
          session_service.append_event(stale_session_1, event_1),
          session_service.append_event(stale_session_2, event_2),
          return_exceptions=True,
      )
      errors = [result for result in results if isinstance(result, Exception)]
      successes = [
          result for result in results if not isinstance(result, Exception)
      ]
      assert len(successes) == 1
      assert len(errors) == 1
      assert isinstance(errors[0], ValueError)
      assert 'modified in storage' in str(errors[0])

    session_final = await session_service.get_session(
        app_name=app_name, user_id=user_id, session_id=session.id
    )

    for i in range(iteration_count):
      event_values = {
          session_final.state.get(f'sk{i}-1'),
          session_final.state.get(f'sk{i}-2'),
      }
      assert event_values & {f'v{i}-1', f'v{i}-2'}
      assert None in event_values
    assert len(session_final.events) == iteration_count


@pytest.mark.asyncio
async def test_append_event_allows_timestamp_drift_for_current_session():
  service = DatabaseSessionService('sqlite+aiosqlite:///:memory:')
  try:
    session = await service.create_session(
        app_name='my_app', user_id='user', session_id='s1'
    )
    event1 = Event(
        invocation_id='inv1',
        author='user',
        timestamp=session.last_update_time + 10,
    )
    await service.append_event(session, event1)

    # Simulate a float round-trip mismatch without changing the persisted
    # session revision.
    session.last_update_time -= 0.0001

    event2 = Event(
        invocation_id='inv2',
        author='user',
        timestamp=event1.timestamp + 10,
    )
    await service.append_event(session, event2)

    refreshed_session = await service.get_session(
        app_name='my_app', user_id='user', session_id=session.id
    )
    assert [event.invocation_id for event in refreshed_session.events] == [
        'inv1',
        'inv2',
    ]
  finally:
    await service.close()


@pytest.mark.asyncio
async def test_append_event_allows_markerless_current_session():
  service = DatabaseSessionService('sqlite+aiosqlite:///:memory:')
  try:
    session = await service.create_session(
        app_name='my_app', user_id='user', session_id='s1'
    )
    event1 = Event(
        invocation_id='inv1',
        author='user',
        timestamp=session.last_update_time + 10,
    )
    await service.append_event(session, event1)

    session._storage_update_marker = None
    session.last_update_time -= 0.0001

    event2 = Event(
        invocation_id='inv2',
        author='user',
        timestamp=event1.timestamp + 10,
    )
    await service.append_event(session, event2)

    refreshed_session = await service.get_session(
        app_name='my_app', user_id='user', session_id=session.id
    )
    assert [event.invocation_id for event in refreshed_session.events] == [
        'inv1',
        'inv2',
    ]
  finally:
    await service.close()


@pytest.mark.asyncio
async def test_get_session_with_config(session_service):
  app_name = 'my_app'
  user_id = 'user'

  num_test_events = 5
  session = await session_service.create_session(
      app_name=app_name, user_id=user_id
  )
  for i in range(1, num_test_events + 1):
    event = Event(author='user', timestamp=i)
    await session_service.append_event(session, event)

  # No config, expect all events to be returned.
  session = await session_service.get_session(
      app_name=app_name, user_id=user_id, session_id=session.id
  )
  events = session.events
  assert len(events) == num_test_events

  # Only expect the most recent 3 events.
  num_recent_events = 3
  config = GetSessionConfig(num_recent_events=num_recent_events)
  session = await session_service.get_session(
      app_name=app_name, user_id=user_id, session_id=session.id, config=config
  )
  events = session.events
  assert len(events) == num_recent_events
  assert events[0].timestamp == num_test_events - num_recent_events + 1

  # Only expect events after timestamp 4.0 (inclusive), i.e., 2 events.
  after_timestamp = 4.0
  config = GetSessionConfig(after_timestamp=after_timestamp)
  session = await session_service.get_session(
      app_name=app_name, user_id=user_id, session_id=session.id, config=config
  )
  events = session.events
  assert len(events) == num_test_events - after_timestamp + 1
  assert events[0].timestamp == after_timestamp

  # Expect no events if none are > after_timestamp.
  way_after_timestamp = num_test_events * 10
  config = GetSessionConfig(after_timestamp=way_after_timestamp)
  session = await session_service.get_session(
      app_name=app_name, user_id=user_id, session_id=session.id, config=config
  )
  assert not session.events

  # Both filters applied, i.e., of 3 most recent events, only 2 are after
  # timestamp 4.0, so expect 2 events.
  config = GetSessionConfig(
      after_timestamp=after_timestamp, num_recent_events=num_recent_events
  )
  session = await session_service.get_session(
      app_name=app_name, user_id=user_id, session_id=session.id, config=config
  )
  events = session.events
  assert len(events) == num_test_events - after_timestamp + 1


@pytest.mark.asyncio
async def test_partial_events_are_not_persisted(session_service):
  app_name = 'my_app'
  user_id = 'user'
  session = await session_service.create_session(
      app_name=app_name, user_id=user_id
  )
  event = Event(author='user', partial=True)
  await session_service.append_event(session, event)

  # Check in-memory session
  assert len(session.events) == 0
  # Check persisted session
  session_got = await session_service.get_session(
      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()


@pytest.mark.asyncio
async def test_concurrent_prepare_tables_no_race_condition():
  """Verifies that concurrent calls to _prepare_tables wait for table creation.
  Reproduces the race condition from
  https://github.com/google/adk-python/issues/4445: when concurrent requests
  arrive at startup, _prepare_tables must not return before tables exist.
  Previously, the early-return guard checked _db_schema_version (set during
  schema detection) instead of _tables_created, so a second request could
  slip through after schema detection but before table creation finished.
  """
  service = DatabaseSessionService('sqlite+aiosqlite:///:memory:')
  try:
    # Tables haven't been created yet.
    assert not service._tables_created
    assert service._db_schema_version is None

    # Launch several concurrent create_session calls, each with a unique
    # app_name to avoid IntegrityError on the shared app_states row.
    # Each will call _prepare_tables internally.  If the race condition
    # exists, some of these will fail because the "sessions" table doesn't
    # exist yet.
    num_concurrent = 5
    results = await asyncio.gather(
        *[
            service.create_session(
                app_name=f'app_{i}', user_id='user', session_id=f'sess_{i}'
            )
            for i in range(num_concurrent)
        ],
        return_exceptions=True,
    )

    # Every call must succeed – no exceptions allowed.
    for i, result in enumerate(results):
      assert not isinstance(result, BaseException), (
          f'Concurrent create_session #{i} raised {result!r}; tables were'
          ' likely not ready due to the _prepare_tables race condition.'
      )

    # All sessions should be retrievable.
    for i in range(num_concurrent):
      session = await service.get_session(
          app_name=f'app_{i}', user_id='user', session_id=f'sess_{i}'
      )
      assert session is not None, f'Session sess_{i} not found after creation.'

    assert service._tables_created
  finally:
    await service.close()


@pytest.mark.asyncio
async def test_prepare_tables_serializes_schema_detection_and_creation():
  """Verifies schema detection and table creation happen atomically under one
  lock, so concurrent callers cannot observe a partially-initialized state.
  After _prepare_tables completes, both _db_schema_version and _tables_created
  must be set.
  """
  service = DatabaseSessionService('sqlite+aiosqlite:///:memory:')
  try:
    assert not service._tables_created
    assert service._db_schema_version is None

    await service._prepare_tables()

    # Both must be set after a single _prepare_tables call.
    assert service._tables_created
    assert service._db_schema_version is not None

    # Verify tables actually exist by performing a real operation.
    session = await service.create_session(
        app_name='app', user_id='user', session_id='s1'
    )
    assert session is not None
    assert session.id == 's1'
  finally:
    await service.close()


@pytest.mark.asyncio
async def test_get_or_create_state_returns_existing_row():
  """_get_or_create_state returns an existing row without inserting."""
  service = DatabaseSessionService('sqlite+aiosqlite:///:memory:')
  try:
    await service._prepare_tables()
    schema = service._get_schema_classes()

    # Pre-create the app_state row.
    async with service.database_session_factory() as sql_session:
      sql_session.add(schema.StorageAppState(app_name='app1', state={'k': 'v'}))
      await sql_session.commit()

    # _get_or_create_state should find and return it.
    async with service.database_session_factory() as sql_session:
      row = await database_session_service._get_or_create_state(
          sql_session=sql_session,
          state_model=schema.StorageAppState,
          primary_key='app1',
          defaults={'app_name': 'app1', 'state': {}},
      )
      assert row.app_name == 'app1'
      assert row.state == {'k': 'v'}
  finally:
    await service.close()


@pytest.mark.asyncio
async def test_get_or_create_state_creates_new_row():
  """_get_or_create_state creates a row when none exists."""
  service = DatabaseSessionService('sqlite+aiosqlite:///:memory:')
  try:
    await service._prepare_tables()
    schema = service._get_schema_classes()

    async with service.database_session_factory() as sql_session:
      row = await database_session_service._get_or_create_state(
          sql_session=sql_session,
          state_model=schema.StorageAppState,
          primary_key='new_app',
          defaults={'app_name': 'new_app', 'state': {}},
      )
      await sql_session.commit()
      assert row.app_name == 'new_app'
      assert row.state == {}

    # Verify the row was actually persisted.
    async with service.database_session_factory() as sql_session:
      persisted = await sql_session.get(schema.StorageAppState, 'new_app')
      assert persisted is not None
  finally:
    await service.close()


@pytest.mark.asyncio
async def test_get_or_create_state_handles_race_condition():
  """_get_or_create_state recovers when a concurrent INSERT wins the race.

  Simulates the race from https://github.com/google/adk-python/issues/4954:
  the initial SELECT returns None (another caller hasn't committed yet), but
  by the time we INSERT, the other caller has committed — so the INSERT fails
  with IntegrityError and we fall back to re-fetching.
  """
  service = DatabaseSessionService('sqlite+aiosqlite:///:memory:')
  try:
    await service._prepare_tables()
    schema = service._get_schema_classes()

    # Pre-create the row to guarantee the INSERT will fail.
    async with service.database_session_factory() as sql_session:
      sql_session.add(schema.StorageAppState(app_name='race_app', state={}))
      await sql_session.commit()

    # Patch session.get to return None on the first call (simulating the
    # race window), then fall through to the real implementation.
    async with service.database_session_factory() as sql_session:
      original_get = sql_session.get
      call_count = 0

      async def patched_get(*args, **kwargs):
        nonlocal call_count
        call_count += 1
        if call_count == 1:
          return None  # Simulate: row not yet visible
        return await original_get(*args, **kwargs)

      sql_session.get = patched_get

      row = await database_session_service._get_or_create_state(
          sql_session=sql_session,
          state_model=schema.StorageAppState,
          primary_key='race_app',
          defaults={'app_name': 'race_app', 'state': {}},
      )
      assert row.app_name == 'race_app'
      # The function should have called get twice: once before the INSERT
      # (patched to return None) and once after the IntegrityError.
      assert call_count == 2
  finally:
    await service.close()


@pytest.mark.asyncio
async def test_create_session_sequential_same_app_name():
  """Sequential create_session calls for the same app_name work correctly.

  The second call reuses the existing app_states row.
  """
  service = DatabaseSessionService('sqlite+aiosqlite:///:memory:')
  try:
    s1 = await service.create_session(
        app_name='shared', user_id='u1', session_id='s1'
    )
    s2 = await service.create_session(
        app_name='shared', user_id='u2', session_id='s2'
    )
    assert s1.app_name == 'shared'
    assert s2.app_name == 'shared'

    got1 = await service.get_session(
        app_name='shared', user_id='u1', session_id='s1'
    )
    got2 = await service.get_session(
        app_name='shared', user_id='u2', session_id='s2'
    )
    assert got1 is not None
    assert got2 is not None
  finally:
    await service.close()


@pytest.mark.asyncio
async def test_prepare_tables_idempotent_after_creation():
  """Calling _prepare_tables multiple times is safe and idempotent.
  After tables are created, subsequent calls should return immediately via
  the fast path without errors.
  """
  service = DatabaseSessionService('sqlite+aiosqlite:///:memory:')
  try:
    await service._prepare_tables()
    assert service._tables_created

    # Call again — should be a no-op via the fast path.
    await service._prepare_tables()
    assert service._tables_created

    # Service should still work.
    session = await service.create_session(
        app_name='app', user_id='user', session_id='s1'
    )
    assert session.id == 's1'
  finally:
    await service.close()


@pytest.mark.asyncio
@pytest.mark.parametrize(
    'state_delta, expect_app_lock, expect_user_lock',
    [
        pytest.param(
            None,
            False,
            False,
            id='no_state_delta',
        ),
        pytest.param(
            {'session_key': 'v'},
            False,
            False,
            id='session_only_delta',
        ),
        pytest.param(
            {'app:key': 'v'},
            True,
            False,
            id='app_delta_only',
        ),
        pytest.param(
            {'user:key': 'v'},
            False,
            True,
            id='user_delta_only',
        ),
        pytest.param(
            {'app:a': '1', 'user:b': '2', 'sk': '3'},
            True,
            True,
            id='all_scopes',
        ),
    ],
)
async def test_append_event_locks_only_scopes_with_deltas(
    state_delta, expect_app_lock, expect_user_lock
):
  """FOR UPDATE should only be requested for state scopes that have deltas."""
  service = DatabaseSessionService('sqlite+aiosqlite:///:memory:')

  lock_requests = []
  original_fn = database_session_service._select_required_state

  async def tracking_fn(**kwargs):
    lock_requests.append({
        'model': kwargs['state_model'].__tablename__,
        'use_row_level_locking': kwargs['use_row_level_locking'],
    })
    return await original_fn(**kwargs)

  try:
    session = await service.create_session(
        app_name='app', user_id='user', session_id='s1'
    )

    database_session_service._select_required_state = tracking_fn
    lock_requests.clear()

    event_kwargs = {'invocation_id': 'inv', 'author': 'user'}
    if state_delta is not None:
      event_kwargs['actions'] = EventActions(state_delta=state_delta)
    event = Event(**event_kwargs)
    await service.append_event(session, event)

    app_req = next(
        (r for r in lock_requests if r['model'] == 'app_states'), None
    )
    user_req = next(
        (r for r in lock_requests if r['model'] == 'user_states'), None
    )

    # SQLite doesn't support row-level locking so use_row_level_locking is
    # always False. The important check is that locking is not requested
    # when there is no delta (it must never be True without a delta).
    if not expect_app_lock:
      assert (
          app_req is None or not app_req['use_row_level_locking']
      ), 'app_states should not be locked without an app: delta'
    if not expect_user_lock:
      assert (
          user_req is None or not user_req['use_row_level_locking']
      ), 'user_states should not be locked without a user: delta'
  finally:
    database_session_service._select_required_state = original_fn
    await service.close()

import json


class MockPIIMaskerTransformer:
  def before_persist_state(self, state):
    return {k: f"{v}_masked" if isinstance(v, str) else v for k, v in state.items()}

  def after_load_state(self, state):
    return {k: v.replace("_masked", "") if isinstance(v, str) and v.endswith("_masked") else v for k, v in state.items()}

  def before_persist_event(self, event: Event) -> Event:
    new_event = event.model_copy() if hasattr(event, "model_copy") else event.copy()
    if new_event.invocation_id:
      new_event.invocation_id += "_masked"
    return new_event

  def after_load_event(self, event: Event) -> Event:
    new_event = event.model_copy() if hasattr(event, "model_copy") else event.copy()
    if new_event.invocation_id and new_event.invocation_id.endswith("_masked"):
      new_event.invocation_id = new_event.invocation_id.replace("_masked", "")
    return new_event

@pytest.mark.asyncio
async def test_session_data_transformer():
  service = DatabaseSessionService('sqlite+aiosqlite:///:memory:', transformer=MockPIIMaskerTransformer())
  try:
    session = await service.create_session(
        app_name='app', user_id='user', session_id='s1', state={'app:secret': 'foo', 'user:pii': 'bar'}
    )
    assert session.state == {'app:secret': 'foo', 'user:pii': 'bar'}

    # Verify persistence has been masked
    async with service.db_engine.connect() as conn:
      from sqlalchemy import text
      result = await conn.execute(text("SELECT state FROM app_states WHERE app_name = 'app'"))
      app_state_json = result.scalar()
      assert "foo_masked" in json.dumps(app_state_json)

    event = Event(invocation_id='inv1', author='user', actions=EventActions(state_delta={'sk1': 'pass'}))
    returned_event = await service.append_event(session, event)

    assert returned_event.invocation_id == 'inv1'
    assert session.state.get('sk1') == 'pass'

    # Check event persistence
    async with service.db_engine.connect() as conn:
      result = await conn.execute(text("SELECT id, state FROM sessions WHERE id = 's1'"))
      row = result.fetchone()
      assert "pass_masked" in json.dumps(row[1])

      result_evt = await conn.execute(text("SELECT event_data FROM events WHERE session_id = 's1' LIMIT 1"))
      evt_payload = result_evt.scalar()
      assert "inv1_masked" in json.dumps(evt_payload)

    # Check retrieval unmasks
    loaded_session = await service.get_session(app_name='app', user_id='user', session_id='s1')
    assert loaded_session.state == {'app:secret': 'foo', 'user:pii': 'bar', 'sk1': 'pass'}
    assert len(loaded_session.events) == 1
    assert loaded_session.events[0].invocation_id == 'inv1'
  finally:
    await service.close()

class ErrorMaskerTransformer:
  def before_persist_state(self, state):
    raise ValueError("Transformer exception test")

  def after_load_state(self, state):
    return state

  def before_persist_event(self, event: Event) -> Event:
    return event

  def after_load_event(self, event: Event) -> Event:
    return event

@pytest.mark.asyncio
async def test_session_data_transformer_handles_exception():
  service = DatabaseSessionService('sqlite+aiosqlite:///:memory:', transformer=ErrorMaskerTransformer())
  try:
    with pytest.raises(ValueError, match="Transformer exception test"):
      await service.create_session(
          app_name='app', user_id='user', session_id='s1', state={'app:secret': 'foo'}
      )
  finally:
    await service.close()