From b8a7649ac26928c0db34c3d40a84b71350eb5db6 Mon Sep 17 00:00:00 2001 From: Hyu Date: Sat, 10 Oct 2026 21:30:04 +0800 Subject: [PATCH] fix(monitoring): expire interrupted host executions using liveness leases --- .../pkg/agent/runner/run_ledger_store.py | 2 + .../pkg/api/http/service/monitoring.py | 3 + src/langbot/pkg/core/app.py | 6 ++ .../pkg/entity/persistence/agent_run.py | 4 + .../pkg/entity/persistence/monitoring.py | 4 + .../versions/0036_execution_host_leases.py | 24 +++++ src/langbot/pkg/utils/execution_lease.py | 101 ++++++++++++++++++ tests/unit_tests/test_execution_lease.py | 54 ++++++++++ 8 files changed, 198 insertions(+) create mode 100644 src/langbot/pkg/persistence/alembic/versions/0036_execution_host_leases.py create mode 100644 src/langbot/pkg/utils/execution_lease.py create mode 100644 tests/unit_tests/test_execution_lease.py diff --git a/src/langbot/pkg/agent/runner/run_ledger_store.py b/src/langbot/pkg/agent/runner/run_ledger_store.py index e34a632c6..f31aff510 100644 --- a/src/langbot/pkg/agent/runner/run_ledger_store.py +++ b/src/langbot/pkg/agent/runner/run_ledger_store.py @@ -13,6 +13,7 @@ from sqlalchemy.orm import sessionmaker from ...persistence.datetime_utils import as_naive_utc from ...utils.inflight import inflight_hub +from ...utils.execution_lease import host_execution_lease from ...entity.persistence.agent_run import AgentRun, AgentRunEvent, AgentRuntime @@ -137,6 +138,7 @@ class RunLedgerStore: thread_id=thread_id, workspace_id=workspace_id, bot_id=bot_id, + **(host_execution_lease() if status == 'running' and queue_name is None else {}), status=status, queue_name=queue_name, priority=priority, diff --git a/src/langbot/pkg/api/http/service/monitoring.py b/src/langbot/pkg/api/http/service/monitoring.py index d4ec46591..fcd98d2aa 100644 --- a/src/langbot/pkg/api/http/service/monitoring.py +++ b/src/langbot/pkg/api/http/service/monitoring.py @@ -555,7 +555,10 @@ class MonitoringService(ExecutionDetailsMixin): run_id = run_id or trace.current_run() or None message_id = str(uuid.uuid4()) message_content = self._sanitize_message_content(message_content) + from ....utils.execution_lease import host_execution_lease + message_data = { + **(host_execution_lease() if status == 'pending' else {}), 'id': message_id, 'workspace_uuid': workspace_uuid, 'timestamp': datetime.datetime.now(datetime.timezone.utc).replace(tzinfo=None), diff --git a/src/langbot/pkg/core/app.py b/src/langbot/pkg/core/app.py index cf3ecd13a..3a066714f 100644 --- a/src/langbot/pkg/core/app.py +++ b/src/langbot/pkg/core/app.py @@ -358,6 +358,12 @@ class Application: async def run(self): self.event_loop_monitor.start() + from ..utils.execution_lease import maintain_execution_leases + self.task_mgr.create_task( + maintain_execution_leases(self), + name='execution-lease-maintenance', + scopes=[core_entities.LifecycleControlScope.APPLICATION], + ) try: if ( self.directory_projection_service is not None diff --git a/src/langbot/pkg/entity/persistence/agent_run.py b/src/langbot/pkg/entity/persistence/agent_run.py index 7d309110d..7d5dc0a8c 100644 --- a/src/langbot/pkg/entity/persistence/agent_run.py +++ b/src/langbot/pkg/entity/persistence/agent_run.py @@ -44,6 +44,10 @@ class AgentRun(Base): bot_id = sqlalchemy.Column(sqlalchemy.String(255), nullable=True, index=True) """Bot UUID this run belongs to.""" + execution_owner_id = sqlalchemy.Column(sqlalchemy.String(36), nullable=True, index=True) + execution_lease_expires_at = sqlalchemy.Column(sqlalchemy.DateTime, nullable=True, index=True) + """Host process lease, independent of execution duration and remote queue claims.""" + status = sqlalchemy.Column(sqlalchemy.String(50), nullable=False, index=True) """Run lifecycle status.""" diff --git a/src/langbot/pkg/entity/persistence/monitoring.py b/src/langbot/pkg/entity/persistence/monitoring.py index 68d68d867..9b07d4f5a 100644 --- a/src/langbot/pkg/entity/persistence/monitoring.py +++ b/src/langbot/pkg/entity/persistence/monitoring.py @@ -21,6 +21,10 @@ class MonitoringMessage(Base): pipeline_name = sqlalchemy.Column(sqlalchemy.String(255), nullable=False) message_content = sqlalchemy.Column(sqlalchemy.Text, nullable=False) session_id = sqlalchemy.Column(sqlalchemy.String(255), nullable=False, index=True) + execution_owner_id = sqlalchemy.Column(sqlalchemy.String(36), nullable=True, index=True) + execution_lease_expires_at = sqlalchemy.Column(sqlalchemy.DateTime, nullable=True, index=True) + """Host process lease, independent of execution duration and remote queue claims.""" + status = sqlalchemy.Column(sqlalchemy.String(50), nullable=False) # success, error, pending level = sqlalchemy.Column(sqlalchemy.String(50), nullable=False) # info, warning, error, debug platform = sqlalchemy.Column(sqlalchemy.String(255), nullable=True) diff --git a/src/langbot/pkg/persistence/alembic/versions/0036_execution_host_leases.py b/src/langbot/pkg/persistence/alembic/versions/0036_execution_host_leases.py new file mode 100644 index 000000000..35052f4ab --- /dev/null +++ b/src/langbot/pkg/persistence/alembic/versions/0036_execution_host_leases.py @@ -0,0 +1,24 @@ +"""Track the liveness of hosts owning direct executions.""" + +from alembic import op +import sqlalchemy as sa + +revision = '0036_execution_host_leases' +down_revision = '0035_pipeline_received_only' +branch_labels = None +depends_on = None + + +def upgrade(): + for table in ('agent_run', 'monitoring_messages'): + op.add_column(table, sa.Column('execution_owner_id', sa.String(36), nullable=True)) + op.add_column(table, sa.Column('execution_lease_expires_at', sa.DateTime(), nullable=True)) + for column in ('execution_owner_id', 'execution_lease_expires_at'): + op.create_index(f'ix_{table}_{column}', table, [column]) + + +def downgrade(): + for table in ('agent_run', 'monitoring_messages'): + for column in ('execution_owner_id', 'execution_lease_expires_at'): + op.drop_index(f'ix_{table}_{column}', table_name=table) + op.drop_column(table, column) diff --git a/src/langbot/pkg/utils/execution_lease.py b/src/langbot/pkg/utils/execution_lease.py new file mode 100644 index 000000000..d25e8f001 --- /dev/null +++ b/src/langbot/pkg/utils/execution_lease.py @@ -0,0 +1,101 @@ +"""Host liveness leases; no age-based timeout for live executions.""" + +import asyncio +import datetime +import uuid + +import sqlalchemy as sa + +from ..entity.persistence.agent_run import AgentRun +from ..entity.persistence.monitoring import MonitoringMessage, MonitoringError +from .inflight import inflight_hub + +OWNER_ID = str(uuid.uuid4()) +LEASE_SECONDS = 120 +INTERRUPTED_REASON = 'Execution interrupted: the LangBot host stopped renewing its lease.' + + +def host_execution_lease(): + return { + 'execution_owner_id': OWNER_ID, + 'execution_lease_expires_at': datetime.datetime.now(datetime.timezone.utc).replace(tzinfo=None) + + datetime.timedelta(seconds=LEASE_SECONDS), + } + + +async def reconcile_execution_leases(engine, workspace_uuid, *, now=None, owner_id=OWNER_ID): + """Renew our live leases, then conditionally finish bounded expired batches. + + Never adopt a lease that already expired: another host may be reaping it. + Remote queued/claimed runs and pre-upgrade rows without ownership are excluded. + """ + now = now or datetime.datetime.now(datetime.timezone.utc).replace(tzinfo=None) + changed = 0 + workspaces = set() + for model, statuses, scope, identity in ( + (AgentRun, ['running'], AgentRun.workspace_id, AgentRun.run_id), + (MonitoringMessage, ['pending'], MonitoringMessage.workspace_uuid, MonitoringMessage.id), + ): + conditions = [model.status.in_(statuses), scope == workspace_uuid] + if model is AgentRun: + conditions.extend([AgentRun.queue_name.is_(None), AgentRun.claimed_by_runtime_id.is_(None)]) + async with engine.begin() as connection: + if engine.dialect.name == 'postgresql': + await connection.execute(sa.text("SELECT set_config('langbot.workspace_uuid', :workspace, true)"), + {'workspace': workspace_uuid}) + await connection.execute(sa.update(model).where( + *conditions, model.execution_owner_id == owner_id, + model.execution_lease_expires_at > now, + ).values(execution_lease_expires_at=now + datetime.timedelta(seconds=LEASE_SECONDS))) + expired = [*conditions, model.execution_lease_expires_at <= now] + rows = (await connection.execute(sa.select(identity, scope).where(*expired) + .order_by(model.execution_lease_expires_at).limit(500))).all() + if not rows: + continue + values = {'status': 'failed', 'status_reason': INTERRUPTED_REASON, + 'finished_at': now, 'updated_at': now} if model is AgentRun else { + 'status': 'error', 'level': 'error', + } + if model is MonitoringMessage: + messages = (await connection.execute(sa.select(model).where( + model.id.in_([row[0] for row in rows]), + ))).mappings().all() + for message in messages: + result = await connection.execute(sa.update(model).where( + *expired, model.id == message['id'], + ).values(**values)) + changed += result.rowcount + if result.rowcount: + await connection.execute(sa.insert(MonitoringError).values( + id=str(uuid.uuid4()), timestamp=now, + error_type='ExecutionInterrupted', error_message=INTERRUPTED_REASON, + message_id=message['id'], + **{key: message[key] for key in ( + 'workspace_uuid', 'bot_id', 'bot_name', 'pipeline_id', + 'pipeline_name', 'session_id', + )}, + )) + else: + result = await connection.execute(sa.update(model).where( + *expired, identity.in_([row[0] for row in rows]), + ).values(**values)) + changed += result.rowcount + workspaces.update(row[1] for row in rows) + for workspace in workspaces: + inflight_hub.notify(workspace) + return changed + + +async def maintain_execution_leases(ap): + while True: + try: + count = 0 + for binding in await ap.workspace_service.list_active_execution_bindings(): + count += await reconcile_execution_leases( + ap.persistence_mgr.get_db_engine(), binding.workspace_uuid, + ) + if count: + ap.logger.info(f'Marked {count} interrupted executions as failed after host lease expiry') + except Exception: + ap.logger.exception('Execution lease maintenance failed') + await asyncio.sleep(20) diff --git a/tests/unit_tests/test_execution_lease.py b/tests/unit_tests/test_execution_lease.py new file mode 100644 index 000000000..2a851eec7 --- /dev/null +++ b/tests/unit_tests/test_execution_lease.py @@ -0,0 +1,54 @@ +import datetime + +import pytest +import sqlalchemy as sa +from sqlalchemy.ext.asyncio import create_async_engine + +from langbot.pkg.entity.persistence.base import Base +from langbot.pkg.entity.persistence.agent_run import AgentRun +from langbot.pkg.entity.persistence.monitoring import MonitoringMessage, MonitoringError +from langbot.pkg.entity.persistence import workspace, user # noqa: F401 +from langbot.pkg.utils.execution_lease import reconcile_execution_leases + + +@pytest.mark.asyncio +async def test_host_lease_reaping_preserves_live_and_external_work(): + engine = create_async_engine('sqlite+aiosqlite:///:memory:') + now = datetime.datetime(2026, 10, 10) + async with engine.begin() as conn: + await conn.run_sync(Base.metadata.create_all) + for name, owner, expiry, status, queue in ( + ('dead', 'dead-host', -1, 'running', None), + ('other-tenant', 'dead-host', -1, 'running', None), + ('ours', 'live-host', 30, 'running', None), + ('other', 'other-host', 30, 'running', None), + ('done', 'dead-host', -1, 'completed', None), + ('legacy', None, None, 'running', None), + ('remote', 'dead-host', -1, 'running', 'remote'), + ): + await conn.execute(sa.insert(AgentRun).values( + run_id=name, runner_id='runner', workspace_id='other-workspace' if name == 'other-tenant' else 'workspace', status=status, + queue_name=queue, execution_owner_id=owner, + execution_lease_expires_at=now+datetime.timedelta(seconds=expiry) if expiry else None, + )) + await conn.execute(sa.insert(MonitoringMessage).values( + id='pipeline', workspace_uuid='workspace', timestamp=now, + bot_id='bot', bot_name='bot', pipeline_id='pipeline', pipeline_name='pipeline', + message_content='hello', session_id='session', status='pending', level='info', + execution_owner_id='dead-host', execution_lease_expires_at=now-datetime.timedelta(seconds=1), + )) + assert await reconcile_execution_leases(engine, 'workspace', now=now, owner_id='live-host') == 2 + assert await reconcile_execution_leases(engine, 'workspace', now=now, owner_id='live-host') == 0 + async with engine.connect() as conn: + rows = {r['run_id']: r for r in (await conn.execute(sa.select(AgentRun))).mappings()} + assert rows['dead']['status'] == 'failed' + assert rows['dead']['status_reason'] + assert rows['dead']['finished_at'] == now + assert rows['ours']['execution_lease_expires_at'] == now+datetime.timedelta(seconds=120) + for key in ('ours', 'other', 'legacy', 'remote', 'other-tenant'): + assert rows[key]['status'] == 'running' + assert rows['done']['status'] == 'completed' + assert (await conn.execute(sa.select(MonitoringMessage.status))).scalar_one() == 'error' + assert (await conn.execute(sa.select(sa.func.count()).select_from(MonitoringError))).scalar_one() == 1 + await engine.dispose() +