# 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.adk.sessions.vertex_ai_session_service import VertexAiSessionService
from google.genai import types
import pytest
from sqlalchemy import delete
from sqlalchemy import text
from sqlalchemy.ext.asyncio import create_async_engine
from sqlalchemy.pool import StaticPool


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', 'mysql', 'mariadb']
)
def test_database_session_service_uses_naive_datetime_for_dialect(dialect_name):
  """Verifies dialects that store DATETIME WITHOUT TIME ZONE are treated as naive.

  SQLite, PostgreSQL, MySQL, and MariaDB all store DATETIME/TIMESTAMP WITHOUT
  TIME ZONE, so create_session must strip tzinfo before storing. Otherwise the
  marker produced by create_session (with +00:00) mismatches the marker read
  back from storage (without +00:00), triggering a false stale-writer error on
  the first append_event after create_session.

  This exercises the production decision (_uses_naive_datetime) directly rather
  than re-implementing the strip logic, so it actually guards create_session.
  """
  fake_engine = mock.Mock()
  fake_engine.dialect.name = dialect_name
  fake_engine.sync_engine = mock.Mock()

  service = DatabaseSessionService(db_engine=fake_engine)

  assert service._uses_naive_datetime() is True


def test_database_session_service_keeps_timezone_for_spanner():
  """Spanner's TIMESTAMP is timezone-aware, so tzinfo must be preserved."""
  fake_engine = mock.Mock()
  fake_engine.dialect.name = 'spanner+spanner'
  fake_engine.sync_engine = mock.Mock()

  service = DatabaseSessionService(db_engine=fake_engine)

  assert service._uses_naive_datetime() is False


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():
  async with DatabaseSessionService('sqlite+aiosqlite:///:memory:') as service:
    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_append_event_when_session_is_same_ref_as_storage_session():
  """Tests that appending an event to a session only appends it once if the user-passed session and the underlying storage session are the same object."""
  service = InMemorySessionService()
  app_name = 'my_app'
  user_id = 'test_user'

  # Create a session
  session = await service.create_session(app_name=app_name, user_id=user_id)

  # Get the actual storage event object from the underlying storage
  storage_session = service.sessions[app_name][user_id][session.id]

  # Append the event to the storage session directly
  event = Event(invocation_id='inv1', author='user')
  await service.append_event(session=storage_session, event=event)

  # Verify that the storage session has only one event
  final_session = await service.get_session(
      app_name=app_name, user_id=user_id, session_id=session.id
  )
  assert len(final_session.events) == 1


@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

  # Explicitly requesting zero recent events should return no event history.
  config = GetSessionConfig(num_recent_events=0)
  session = await session_service.get_session(
      app_name=app_name, user_id=user_id, session_id=session.id, config=config
  )
  assert not session.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()


class _CommitOrderSpySession:
  """SQLAlchemy session spy that marks when commit() has completed."""

  def __init__(self, real_session, on_committed):
    self._real = real_session
    self._on_committed = on_committed

  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):
    result = await self._real.commit()
    self._on_committed()
    return result

  def __getattr__(self, name):
    return getattr(self._real, name)


@pytest.mark.asyncio
async def test_append_event_reads_storage_revision_before_commit():
  """append_event captures session revision before commit completes."""
  service = DatabaseSessionService('sqlite+aiosqlite:///:memory:')
  await service.prepare_tables()
  schema = service._get_schema_classes()
  original_get_update_timestamp = schema.StorageSession.get_update_timestamp
  original_get_update_marker = schema.StorageSession.get_update_marker
  revision_read_state = {'committed': False, 'post_commit_reads': 0}

  def _track_revision_read(original):
    def wrapper(self, *args, **kwargs):
      if revision_read_state['committed']:
        revision_read_state['post_commit_reads'] += 1
      return original(self, *args, **kwargs)

    return wrapper

  schema.StorageSession.get_update_timestamp = _track_revision_read(
      original_get_update_timestamp
  )
  schema.StorageSession.get_update_marker = _track_revision_read(
      original_get_update_marker
  )

  try:
    session = await service.create_session(
        app_name='app', user_id='user', session_id='s1'
    )
    event_timestamp = session.last_update_time + 10
    event = Event(
        invocation_id='inv1',
        author='user',
        timestamp=event_timestamp,
    )

    original_factory = service.database_session_factory

    def _spy_factory():
      return _CommitOrderSpySession(
          original_factory(),
          on_committed=lambda: revision_read_state.update({'committed': True}),
      )

    service.database_session_factory = _spy_factory

    await service.append_event(session, event)

    assert revision_read_state['post_commit_reads'] == 0
    assert session.last_update_time == pytest.approx(event_timestamp, abs=1e-6)
    assert session._storage_update_marker is not None
  finally:
    schema.StorageSession.get_update_timestamp = original_get_update_timestamp
    schema.StorageSession.get_update_marker = original_get_update_marker

    await service.close()


@pytest.mark.asyncio
async def test_create_session_reads_storage_revision_before_commit():
  """create_session captures session revision before commit completes."""
  service = DatabaseSessionService('sqlite+aiosqlite:///:memory:')
  await service.prepare_tables()
  schema = service._get_schema_classes()
  original_get_update_timestamp = schema.StorageSession.get_update_timestamp
  original_get_update_marker = schema.StorageSession.get_update_marker
  revision_read_state = {'committed': False, 'post_commit_reads': 0}

  def _track_revision_read(original):
    def wrapper(self, *args, **kwargs):
      if revision_read_state['committed']:
        revision_read_state['post_commit_reads'] += 1
      return original(self, *args, **kwargs)

    return wrapper

  schema.StorageSession.get_update_timestamp = _track_revision_read(
      original_get_update_timestamp
  )
  schema.StorageSession.get_update_marker = _track_revision_read(
      original_get_update_marker
  )

  try:
    original_factory = service.database_session_factory

    def _spy_factory():
      return _CommitOrderSpySession(
          original_factory(),
          on_committed=lambda: revision_read_state.update({'committed': True}),
      )

    service.database_session_factory = _spy_factory

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

    assert revision_read_state['post_commit_reads'] == 0
    assert session.last_update_time is not None
    assert session._storage_update_marker is not None
  finally:
    schema.StorageSession.get_update_timestamp = original_get_update_timestamp
    schema.StorageSession.get_update_marker = original_get_update_marker

    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 where 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:
  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
async def test_public_prepare_tables_eager_initialization():
  """Calling the public prepare_tables() eagerly initializes tables so that
  the first real database operation does not pay the setup cost.
  """
  async with DatabaseSessionService('sqlite+aiosqlite:///:memory:') as service:
    # Before calling prepare_tables, tables are not created.
    assert not service._tables_created
    assert service._db_schema_version is None

    # Eagerly prepare tables via the public API.
    await service.prepare_tables()

    # Tables should now be ready.
    assert service._tables_created
    assert service._db_schema_version is not None

    # Subsequent operations should work without any additional setup cost.
    session = await service.create_session(
        app_name='app', user_id='user', session_id='s1'
    )
    assert session.id == 's1'


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


@pytest.mark.asyncio
async def test_get_user_state_returns_empty_dict_when_no_state_exists(
    session_service,
):
  """Verifies get_user_state returns empty dict when no state exists."""
  state = await session_service.get_user_state(app_name='my_app', user_id='u1')
  assert not state


@pytest.mark.asyncio
async def test_get_user_state_returns_state_written_via_append_event(
    session_service,
):
  """Verifies get_user_state returns state written via append_event."""
  session = await session_service.create_session(
      app_name='my_app', user_id='u1'
  )
  await session_service.append_event(
      session,
      Event(
          author='system',
          actions=EventActions(
              state_delta={'user:profile': {'name': 'Alice'}, 'session_key': 1}
          ),
      ),
  )

  state = await session_service.get_user_state(app_name='my_app', user_id='u1')

  assert state == {'profile': {'name': 'Alice'}}
  assert 'session_key' not in state


@pytest.mark.asyncio
async def test_get_user_state_is_not_visible_across_users(session_service):
  """Verifies user state is isolated between users."""
  session = await session_service.create_session(
      app_name='my_app', user_id='u1'
  )
  await session_service.append_event(
      session,
      Event(
          author='system',
          actions=EventActions(state_delta={'user:secret': 'only-for-u1'}),
      ),
  )

  other_state = await session_service.get_user_state(
      app_name='my_app', user_id='u2'
  )
  assert not other_state


@pytest.mark.asyncio
async def test_get_user_state_is_not_visible_across_apps(session_service):
  """Verifies user state is isolated between apps."""
  session = await session_service.create_session(
      app_name='my_app', user_id='u1'
  )
  await session_service.append_event(
      session,
      Event(
          author='system',
          actions=EventActions(state_delta={'user:data': 'only-app-a'}),
      ),
  )

  other_state = await session_service.get_user_state(
      app_name='other_app', user_id='u1'
  )
  assert not other_state


@pytest.mark.asyncio
async def test_get_user_state_available_before_session_is_created(
    session_service,
):
  """Verifies user state can be retrieved before a session is created."""
  first_session = await session_service.create_session(
      app_name='my_app', user_id='u1'
  )
  await session_service.append_event(
      first_session,
      Event(
          author='system',
          actions=EventActions(state_delta={'user:ctx': {'v': 1}}),
      ),
  )

  state = await session_service.get_user_state(app_name='my_app', user_id='u1')
  assert state == {'ctx': {'v': 1}}


@pytest.mark.asyncio
async def test_get_user_state_reflects_latest_write(session_service):
  """Verifies get_user_state returns the latest state."""
  session = await session_service.create_session(
      app_name='my_app', user_id='u1'
  )
  await session_service.append_event(
      session,
      Event(
          author='system',
          actions=EventActions(state_delta={'user:counter': 1}),
      ),
  )
  await session_service.append_event(
      session,
      Event(
          author='system',
          actions=EventActions(state_delta={'user:counter': 2}),
      ),
  )

  state = await session_service.get_user_state(app_name='my_app', user_id='u1')
  assert state['counter'] == 2


@pytest.mark.asyncio
async def test_vertex_ai_session_service_raises_not_implemented_for_get_user_state():
  """Verifies VertexAiSessionService raises NotImplementedError."""
  service = VertexAiSessionService(project='proj', location='us-central1')
  with pytest.raises(NotImplementedError):
    await service.get_user_state(app_name='my_app', user_id='u1')


def test_database_session_service_visible_in_module_namespace():
  """DatabaseSessionService must be in dir() so Sphinx autodoc renders it.

  It is imported lazily via module __getattr__, so without an explicit
  __dir__ it drops out of the generated API reference (issue #4331).
  """
  import google.adk.sessions as sessions_module

  assert 'DatabaseSessionService' in dir(sessions_module)
  assert sessions_module.DatabaseSessionService is DatabaseSessionService


@pytest.mark.asyncio
async def test_database_session_service_with_db_url():
  """Test DatabaseSessionService initialization with db_url."""
  # Test db_url as positional argument
  service = DatabaseSessionService('sqlite+aiosqlite:///:memory:')
  app_name = 'test_app'
  user_id = 'test_user'

  # Create and retrieve a session
  session = await service.create_session(
      app_name=app_name, user_id=user_id, state={'key': 'value'}
  )
  assert session.app_name == app_name
  assert session.user_id == user_id
  assert session.state == {'key': 'value'}

  # Let's check that we can retrieve it
  retrieved = await service.get_session(
      app_name=app_name, user_id=user_id, session_id=session.id
  )
  assert retrieved == session

  # test db_url as keyword argument
  service2 = DatabaseSessionService(db_url='sqlite+aiosqlite:///:memory:')
  session2 = await service2.create_session(
      app_name=app_name, user_id=user_id, state={'key': 'value2'}
  )
  assert session2.state == {'key': 'value2'}


@pytest.mark.asyncio
async def test_database_session_service_with_db_engine():
  """Test DatabaseSessionService initialization with db_engine."""
  # Create an engine manually with StaticPool to avoid flakes
  engine = create_async_engine(
      'sqlite+aiosqlite:///:memory:',
      poolclass=StaticPool,
      connect_args={'check_same_thread': False},
  )

  # Create service with db_engine
  service = DatabaseSessionService(db_engine=engine)
  app_name = 'test_app'
  user_id = 'test_user'

  # Create and retrieve a session
  session = await service.create_session(
      app_name=app_name, user_id=user_id, state={'key': 'value'}
  )
  assert session.app_name == app_name
  assert session.user_id == user_id
  assert session.state == {'key': 'value'}

  # Let's check that we can retrieve it
  retrieved = await service.get_session(
      app_name=app_name, user_id=user_id, session_id=session.id
  )
  assert retrieved == session


@pytest.mark.asyncio
async def test_database_session_service_caller_owned_engine_not_disposed_on_close():
  """Verifies that a caller-owned engine is not disposed when the service is closed."""
  engine = create_async_engine(
      'sqlite+aiosqlite:///:memory:',
      poolclass=StaticPool,
      connect_args={'check_same_thread': False},
  )

  service = DatabaseSessionService(db_engine=engine)

  # Use the service
  session = await service.create_session(app_name='app', user_id='user')
  assert session is not None

  # Close the service
  await service.close()

  # Verify engine is still usable by running a query
  async with engine.connect() as conn:
    result = await conn.execute(text('SELECT 1;'))
    assert result.scalar() == 1


@pytest.mark.asyncio
async def test_database_session_service_requires_one_argument():
  """Test that DatabaseSessionService requires exactly one of db_url or db_engine."""
  # Neither argument provided
  with pytest.raises(
      ValueError,
      match="Exactly one of 'db_url' or 'db_engine' must be provided",
  ):
    DatabaseSessionService()

  # Both arguments provided
  engine = create_async_engine('sqlite+aiosqlite:///:memory:')
  with pytest.raises(
      ValueError,
      match="Exactly one of 'db_url' or 'db_engine' must be provided",
  ):
    DatabaseSessionService(
        db_url='sqlite+aiosqlite:///:memory:', db_engine=engine
    )