diff --git a/openhands/automation/scheduler.py b/openhands/automation/scheduler.py index c294586..a9709d7 100644 --- a/openhands/automation/scheduler.py +++ b/openhands/automation/scheduler.py @@ -11,6 +11,7 @@ import asyncio import logging from datetime import datetime, timedelta +from uuid import UUID from zoneinfo import ZoneInfoNotFoundError from croniter import CroniterBadDateError, CroniterBadTypeRangeError, CroniterError @@ -18,7 +19,11 @@ from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker from openhands.automation.db import using_sqlite -from openhands.automation.models import Automation, AutomationRun +from openhands.automation.models import ( + Automation, + AutomationRun, + AutomationRunStatus, +) from openhands.automation.telemetry import capture_automation_event from openhands.automation.utils import get_next_fire_time, is_automation_due, utcnow from openhands.automation.utils.run import create_pending_run @@ -38,6 +43,11 @@ ZoneInfoNotFoundError, ) +_IN_FLIGHT_RUN_STATUSES = ( + AutomationRunStatus.PENDING, + AutomationRunStatus.RUNNING, +) + def _schedule_log_extra(automation: Automation) -> dict[str, str | None]: trigger = automation.trigger or {} @@ -142,6 +152,24 @@ async def _fetch_enabled_automations( return list(result.scalars().all()) +async def _get_in_flight_run_statuses( + session: AsyncSession, + automations: list[Automation], +) -> dict[UUID, AutomationRunStatus]: + """Return one active run status for each automation that has one.""" + automation_ids = [automation.id for automation in automations] + if not automation_ids: + return {} + + result = await session.execute( + select(AutomationRun.automation_id, AutomationRun.status).where( + AutomationRun.automation_id.in_(automation_ids), + AutomationRun.status.in_(_IN_FLIGHT_RUN_STATUSES), + ) + ) + return {automation_id: status for automation_id, status in result} + + async def poll_and_schedule( session_factory: async_sessionmaker[AsyncSession], batch_size: int = DEFAULT_BATCH_SIZE, @@ -152,8 +180,9 @@ async def poll_and_schedule( Fetches enabled automations (using FOR UPDATE SKIP LOCKED on PostgreSQL for multi-worker safety), updates last_polled_at for ALL fetched automations (to ensure fair batch rotation), filters to those that are due, and creates - PENDING runs. All within a single transaction so row locks are held throughout - and no schedules can be lost or duplicated. + PENDING runs only when the automation has no in-flight run. All within a single + transaction so row locks are held throughout and no schedules can be lost or + duplicated by concurrent scheduler workers. Note: SQLite deployments skip row locking (single-process mode assumed). @@ -189,9 +218,23 @@ async def poll_and_schedule( automation.last_polled_at = now due_automations = [a for a in automations if _is_automation_due_safely(a, now)] + in_flight_statuses = await _get_in_flight_run_statuses(session, due_automations) for automation in due_automations: try: + in_flight_status = in_flight_statuses.get(automation.id) + if in_flight_status is not None: + logger.info( + "Skipping cron run for active automation", + extra={ + "automation_id": str(automation.id), + "existing_run_status": in_flight_status.value, + "trigger_source": "cron", + "skip_reason": "automation_run_in_flight", + }, + ) + continue + run = await create_pending_run(session, automation) created_runs.append(run) schedule_properties = { @@ -243,8 +286,9 @@ async def scheduler_loop( ) -> None: """Main scheduler loop that polls for due automations. - For each due automation, creates a PENDING run in the automation_runs table. - The dispatcher (separate process) picks up PENDING runs and executes them. + For each due automation without a PENDING or RUNNING run, creates a PENDING run + in the automation_runs table. The dispatcher picks up PENDING runs and executes + them. Args: session_factory: SQLAlchemy async session factory diff --git a/tests/test_scheduler.py b/tests/test_scheduler.py index 4ee6924..140b667 100644 --- a/tests/test_scheduler.py +++ b/tests/test_scheduler.py @@ -1,13 +1,21 @@ """Tests for the scheduler module.""" import asyncio +import logging import uuid from datetime import UTC, datetime, timedelta import pytest -from sqlalchemy import func, select - -from openhands.automation.models import Automation, AutomationRun, AutomationRunStatus +from sqlalchemy import event, func, select +from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine + +from openhands.automation.db import set_sqlite_mode, using_sqlite +from openhands.automation.models import ( + Automation, + AutomationRun, + AutomationRunStatus, + Base, +) from openhands.automation.scheduler import ( POLL_INTERVAL_SECONDS, poll_and_schedule, @@ -34,6 +42,63 @@ def _utc(*args: int) -> datetime: return datetime(*args, tzinfo=UTC) +@pytest.fixture +async def sqlite_session_factory(): + """Create an isolated in-memory SQLite session factory.""" + previous_sqlite_mode = using_sqlite() + sqlalchemy_logger = logging.getLogger("sqlalchemy.engine") + previous_sqlalchemy_level = sqlalchemy_logger.level + set_sqlite_mode(True) + sqlalchemy_logger.setLevel(logging.WARNING) + engine = create_async_engine("sqlite+aiosqlite:///:memory:") + try: + async with engine.begin() as connection: + await connection.run_sync(Base.metadata.create_all) + yield async_sessionmaker( + engine, + class_=AsyncSession, + expire_on_commit=False, + ) + finally: + await engine.dispose() + set_sqlite_mode(previous_sqlite_mode) + sqlalchemy_logger.setLevel(previous_sqlalchemy_level) + + +@pytest.fixture +def scheduler_telemetry_events(monkeypatch): + """Capture scheduler telemetry without invoking the telemetry backend.""" + event_names: list[str] = [] + + async def record_event(event_name, *args, **kwargs): + event_names.append(event_name) + + monkeypatch.setattr( + "openhands.automation.scheduler.capture_automation_event", + record_event, + ) + return event_names + + +def _due_automation(name: str, now: datetime) -> Automation: + return Automation( + user_id=TEST_USER_ID, + org_id=TEST_ORG_ID, + name=name, + trigger={"type": "cron", "schedule": "* * * * *", "timezone": "UTC"}, + tarball_path="s3://bucket/code.tar.gz", + entrypoint="uv run main.py", + enabled=True, + last_triggered_at=now - timedelta(minutes=2), + created_at=now - timedelta(minutes=5), + ) + + +def _as_utc(value: datetime) -> datetime: + """Normalize SQLite timestamps for deterministic comparisons.""" + return value.replace(tzinfo=UTC) + + class TestGetNextFireTime: """Tests for get_next_fire_time function.""" @@ -361,6 +426,223 @@ def test_automation_not_due_with_timezone_before_schedule(self): class TestPollAndSchedule: """Tests for poll_and_schedule function (atomic poll + run creation).""" + @pytest.mark.parametrize( + "active_status", + [AutomationRunStatus.PENDING, AutomationRunStatus.RUNNING], + ) + async def test_poll_skips_due_automation_with_active_run( + self, + sqlite_session_factory, + active_status, + scheduler_telemetry_events, + caplog, + ): + """An in-flight run prevents a second cron-created run.""" + now = _utc(2026, 7, 30, 12, 0, 30) + original_last_triggered_at = now - timedelta(minutes=2) + async with sqlite_session_factory() as session: + automation = _due_automation("Blocked", now) + session.add(automation) + await session.flush() + session.add( + AutomationRun( + automation_id=automation.id, + status=active_status, + ) + ) + await session.commit() + automation_id = automation.id + + with caplog.at_level(logging.INFO, logger="automation.scheduler"): + created_runs = await poll_and_schedule(sqlite_session_factory, now=now) + + assert created_runs == [] + assert scheduler_telemetry_events == [] + async with sqlite_session_factory() as session: + run_count = await session.scalar( + select(func.count()) + .select_from(AutomationRun) + .where(AutomationRun.automation_id == automation_id) + ) + assert run_count == 1 + updated = await session.get(Automation, automation_id) + assert updated is not None + assert _as_utc(updated.last_polled_at) == now + assert _as_utc(updated.last_triggered_at) == original_last_triggered_at + + skip_record = next( + record + for record in caplog.records + if record.getMessage() == "Skipping cron run for active automation" + ) + assert skip_record.automation_id == str(automation_id) + assert skip_record.existing_run_status == active_status.value + assert skip_record.trigger_source == "cron" + assert skip_record.skip_reason == "automation_run_in_flight" + + @pytest.mark.parametrize( + "terminal_status", + [ + AutomationRunStatus.COMPLETED, + AutomationRunStatus.FAILED, + AutomationRunStatus.CANCELLED, + AutomationRunStatus.SKIPPED, + ], + ) + async def test_poll_overlap_guard_allows_terminal_run( + self, sqlite_session_factory, terminal_status, scheduler_telemetry_events + ): + """Terminal history does not block a new cron run.""" + now = _utc(2026, 7, 30, 12, 0, 30) + async with sqlite_session_factory() as session: + automation = _due_automation("Terminal history", now) + session.add(automation) + await session.flush() + session.add( + AutomationRun(automation_id=automation.id, status=terminal_status) + ) + await session.commit() + automation_id = automation.id + + created_runs = await poll_and_schedule(sqlite_session_factory, now=now) + + assert [run.automation_id for run in created_runs] == [automation_id] + assert scheduler_telemetry_events == [ + "automation_run_scheduled", + "automation_run_created", + ] + + async def test_poll_overlap_guard_allows_automation_without_run( + self, sqlite_session_factory, scheduler_telemetry_events + ): + """A due automation without run history is scheduled normally.""" + now = _utc(2026, 7, 30, 12, 0, 30) + async with sqlite_session_factory() as session: + automation = _due_automation("Free", now) + session.add(automation) + await session.commit() + automation_id = automation.id + + created_runs = await poll_and_schedule(sqlite_session_factory, now=now) + + assert [run.automation_id for run in created_runs] == [automation_id] + assert scheduler_telemetry_events == [ + "automation_run_scheduled", + "automation_run_created", + ] + + async def test_poll_overlap_guard_is_per_automation( + self, sqlite_session_factory, scheduler_telemetry_events + ): + """One batch query isolates active runs by automation.""" + now = _utc(2026, 7, 30, 12, 0, 30) + async with sqlite_session_factory() as session: + blocked = _due_automation("Blocked", now) + free = _due_automation("Free", now) + session.add_all([blocked, free]) + await session.flush() + session.add( + AutomationRun( + automation_id=blocked.id, + status=AutomationRunStatus.PENDING, + ) + ) + await session.commit() + blocked_id = blocked.id + free_id = free.id + + active_run_queries: list[str] = [] + + def capture_active_run_query( + connection, cursor, statement, parameters, context, executemany + ): + if "FROM automation_runs" in statement: + active_run_queries.append(statement) + + engine = sqlite_session_factory.kw["bind"] + event.listen( + engine.sync_engine, "before_cursor_execute", capture_active_run_query + ) + try: + created_runs = await poll_and_schedule(sqlite_session_factory, now=now) + finally: + event.remove( + engine.sync_engine, + "before_cursor_execute", + capture_active_run_query, + ) + + assert [run.automation_id for run in created_runs] == [free_id] + assert len(active_run_queries) == 1 + assert scheduler_telemetry_events == [ + "automation_run_scheduled", + "automation_run_created", + ] + async with sqlite_session_factory() as session: + blocked_run_count = await session.scalar( + select(func.count()) + .select_from(AutomationRun) + .where(AutomationRun.automation_id == blocked_id) + ) + assert blocked_run_count == 1 + + async def test_poll_overlap_guard_preserves_batch_rotation( + self, sqlite_session_factory, scheduler_telemetry_events + ): + """A skipped automation rotates behind an eligible automation.""" + now = _utc(2026, 7, 30, 12, 0, 30) + async with sqlite_session_factory() as session: + blocked = _due_automation("Blocked first", now) + blocked.last_polled_at = now - timedelta(hours=2) + free = _due_automation("Free second", now) + free.last_polled_at = now - timedelta(hours=1) + session.add_all([blocked, free]) + await session.flush() + session.add( + AutomationRun( + automation_id=blocked.id, + status=AutomationRunStatus.RUNNING, + ) + ) + await session.commit() + free_id = free.id + + first_runs = await poll_and_schedule( + sqlite_session_factory, batch_size=1, now=now + ) + second_runs = await poll_and_schedule( + sqlite_session_factory, batch_size=1, now=now + ) + + assert first_runs == [] + assert [run.automation_id for run in second_runs] == [free_id] + assert scheduler_telemetry_events == [ + "automation_run_scheduled", + "automation_run_created", + ] + + async def test_poll_overlap_guard_leaves_not_due_automation_untriggered( + self, sqlite_session_factory, scheduler_telemetry_events + ): + """A polled automation that is not due remains untriggered.""" + now = _utc(2026, 7, 30, 12, 0, 30) + async with sqlite_session_factory() as session: + automation = _due_automation("Not due", now) + automation.last_triggered_at = now + session.add(automation) + await session.commit() + automation_id = automation.id + + created_runs = await poll_and_schedule(sqlite_session_factory, now=now) + + assert created_runs == [] + assert scheduler_telemetry_events == [] + async with sqlite_session_factory() as session: + updated = await session.get(Automation, automation_id) + assert updated is not None + assert _as_utc(updated.last_polled_at) == now + assert _as_utc(updated.last_triggered_at) == now + async def test_poll_creates_runs_for_due_automations(self, async_session_factory): """Creates pending runs for automations that are due.""" async with async_session_factory() as session: