diff --git a/src/langbot/pkg/entity/persistence/agent_run.py b/src/langbot/pkg/entity/persistence/agent_run.py index 7d5dc0a8c..af1f9f6b6 100644 --- a/src/langbot/pkg/entity/persistence/agent_run.py +++ b/src/langbot/pkg/entity/persistence/agent_run.py @@ -116,6 +116,7 @@ class AgentRun(Base): ), sqlalchemy.Index('ix_agent_run_runner_status', 'runner_id', 'status'), sqlalchemy.Index('ix_agent_run_queue_claim', 'queue_name', 'status', 'priority', 'id'), + sqlalchemy.Index('ix_agent_run_workspace_status_lease', 'workspace_id', 'status', 'execution_lease_expires_at'), ) diff --git a/src/langbot/pkg/entity/persistence/monitoring.py b/src/langbot/pkg/entity/persistence/monitoring.py index ecf032986..e1a92159e 100644 --- a/src/langbot/pkg/entity/persistence/monitoring.py +++ b/src/langbot/pkg/entity/persistence/monitoring.py @@ -40,6 +40,12 @@ class MonitoringMessage(Base): __table_args__ = ( sqlalchemy.Index('ix_monitoring_messages_workspace_timestamp', 'workspace_uuid', 'timestamp'), sqlalchemy.Index('ix_monitoring_messages_workspace_session', 'workspace_uuid', 'session_id'), + sqlalchemy.Index( + 'ix_monitoring_messages_workspace_status_lease', + 'workspace_uuid', + 'status', + 'execution_lease_expires_at', + ), ) diff --git a/src/langbot/pkg/persistence/alembic/versions/0039_execution_lease_indexes.py b/src/langbot/pkg/persistence/alembic/versions/0039_execution_lease_indexes.py new file mode 100644 index 000000000..4d66427e4 --- /dev/null +++ b/src/langbot/pkg/persistence/alembic/versions/0039_execution_lease_indexes.py @@ -0,0 +1,29 @@ +"""Bound host lease maintenance to active executions within each Workspace.""" + +from alembic import op +import sqlalchemy as sa + +revision = '0039_execution_lease_indexes' +down_revision = '0038_llm_call_details' +branch_labels = None +depends_on = None + + +def upgrade(): + inspector = sa.inspect(op.get_bind()) + for table, workspace_column in ( + ('agent_run', 'workspace_id'), + ('monitoring_messages', 'workspace_uuid'), + ): + index = f'ix_{table}_workspace_status_lease' + # Fresh installations create these indexes from ORM metadata before migrations. + if index not in {item['name'] for item in inspector.get_indexes(table)}: + op.create_index(index, table, [workspace_column, 'status', 'execution_lease_expires_at']) + + +def downgrade(): + inspector = sa.inspect(op.get_bind()) + for table in ('agent_run', 'monitoring_messages'): + index = f'ix_{table}_workspace_status_lease' + if index in {item['name'] for item in inspector.get_indexes(table)}: + op.drop_index(index, table_name=table) diff --git a/tests/unit_tests/persistence/test_execution_lease_indexes.py b/tests/unit_tests/persistence/test_execution_lease_indexes.py new file mode 100644 index 000000000..468871087 --- /dev/null +++ b/tests/unit_tests/persistence/test_execution_lease_indexes.py @@ -0,0 +1,134 @@ +"""Lease maintenance must not scan a Workspace's completed execution history.""" + +import datetime +import importlib + +import pytest +import sqlalchemy as sa +from alembic.migration import MigrationContext +from alembic.operations import Operations +from sqlalchemy.ext.asyncio import create_async_engine + +from langbot.pkg.entity.persistence import user, workspace # noqa: F401 +from langbot.pkg.entity.persistence.agent_run import AgentRun +from langbot.pkg.entity.persistence.base import Base +from langbot.pkg.entity.persistence.monitoring import MonitoringMessage +from langbot.pkg.utils.execution_lease import reconcile_execution_leases + + +MIGRATION = importlib.import_module('langbot.pkg.persistence.alembic.versions.0039_execution_lease_indexes') +LEASE_INDEXES = { + 'agent_run': ('ix_agent_run_workspace_status_lease', ['workspace_id', 'status', 'execution_lease_expires_at']), + 'monitoring_messages': ( + 'ix_monitoring_messages_workspace_status_lease', + ['workspace_uuid', 'status', 'execution_lease_expires_at'], + ), +} + + +def _create_schema(connection, source): + Base.metadata.create_all(connection) + if source == 'migration': + for table, (index, _) in LEASE_INDEXES.items(): + next(item for item in Base.metadata.tables[table].indexes if item.name == index).drop(connection) + with Operations.context(MigrationContext.configure(connection)): + MIGRATION.upgrade() + + +@pytest.mark.parametrize('source', ['metadata', 'migration']) +def test_lease_index_migration_preserves_data_and_existing_indexes(source): + engine = sa.create_engine('sqlite:///:memory:') + try: + with engine.begin() as connection: + _create_schema(connection, source) + connection.execute(sa.insert(AgentRun).values(run_id='history', runner_id='runner', status='completed')) + original_indexes = { + table: {index['name'] for index in sa.inspect(connection).get_indexes(table)} for table in LEASE_INDEXES + } + with Operations.context(MigrationContext.configure(connection)): + MIGRATION.upgrade() + for table, (index, columns) in LEASE_INDEXES.items(): + indexes = {item['name']: item['column_names'] for item in sa.inspect(connection).get_indexes(table)} + assert indexes[index] == columns + assert original_indexes[table] <= indexes.keys() + + MIGRATION.downgrade() + for table, (index, _) in LEASE_INDEXES.items(): + remaining = {item['name'] for item in sa.inspect(connection).get_indexes(table)} + assert remaining == original_indexes[table] - {index} + assert connection.scalar(sa.select(AgentRun.status)) == 'completed' + + MIGRATION.upgrade() + assert connection.scalar(sa.select(AgentRun.status)) == 'completed' + finally: + engine.dispose() + + +@pytest.mark.asyncio +@pytest.mark.parametrize('source', ['metadata', 'migration']) +@pytest.mark.parametrize('model', [AgentRun, MonitoringMessage]) +async def test_lease_queries_use_workspace_status_expiry_index(model, source): + engine = create_async_engine('sqlite+aiosqlite:///:memory:') + now = datetime.datetime(2026, 10, 17) + statements = [] + + def capture_statement(connection, cursor, statement, parameters, context, executemany): + if statement.startswith(('SELECT ', 'UPDATE ')) and model.__tablename__ in statement: + statements.append((statement, parameters)) + + try: + async with engine.begin() as connection: + await connection.run_sync(lambda sync: _create_schema(sync, source)) + common = { + 'execution_owner_id': 'live-host', + 'execution_lease_expires_at': now - datetime.timedelta(seconds=1), + } + if model is AgentRun: + rows = [ + dict( + common, + run_id=str(i), + runner_id='runner', + workspace_id='workspace', + status='completed', + created_at=now, + updated_at=now, + ) + for i in range(2000) + ] + else: + rows = [ + dict( + common, + id=str(i), + workspace_uuid='workspace', + timestamp=now, + bot_id='bot', + bot_name='bot', + pipeline_id='pipeline', + pipeline_name='pipeline', + message_content='Historical message ' * 64, + session_id='session', + status='success', + level='info', + ) + for i in range(2000) + ] + await connection.execute(sa.insert(model), rows) + + sa.event.listen(engine.sync_engine, 'before_cursor_execute', capture_statement) + try: + assert await reconcile_execution_leases(engine, 'workspace', now=now, owner_id='live-host') == 0 + finally: + sa.event.remove(engine.sync_engine, 'before_cursor_execute', capture_statement) + + assert len(statements) == 2 + async with engine.connect() as connection: + for statement, parameters in statements: + plan = (await connection.exec_driver_sql('EXPLAIN QUERY PLAN ' + statement, parameters)).all() + details = ' '.join(row[3] for row in plan) + assert LEASE_INDEXES[model.__tablename__][0] in details, details + assert 'execution_lease_expires_at' in details, details + assert 'USE TEMP B-TREE' not in details, details + finally: + await engine.dispose()