From e3f4af0d97592c488674643b597ed2fc2c1eebd9 Mon Sep 17 00:00:00 2001 From: olegbrok Date: Tue, 4 Aug 2026 05:57:44 -0700 Subject: [PATCH] fix(scheduler): atomically claim schedule fires --- src/pinky_daemon/agent_registry.py | 25 ++++-- src/pinky_daemon/scheduler.py | 14 +++- tests/test_scheduler.py | 122 ++++++++++++++++++++++++++++- 3 files changed, 152 insertions(+), 9 deletions(-) diff --git a/src/pinky_daemon/agent_registry.py b/src/pinky_daemon/agent_registry.py index 1ea8ecf5..93b75e15 100644 --- a/src/pinky_daemon/agent_registry.py +++ b/src/pinky_daemon/agent_registry.py @@ -3229,14 +3229,27 @@ def toggle_schedule(self, schedule_id: int, enabled: bool) -> bool: self._db.commit() return cursor.rowcount > 0 - def update_schedule_last_run(self, schedule_id: int, timestamp: float = 0.0) -> None: - """Record when the scheduler decided to fire a schedule.""" + def update_schedule_last_run( + self, + schedule_id: int, + timestamp: float = 0.0, + *, + expected_last_run: float | None = None, + ) -> bool: + """Record a fire, optionally only if ``last_run`` is unchanged.""" ts = timestamp or time.time() - self._db.execute( - "UPDATE agent_schedules SET last_run=? WHERE id=?", - (ts, schedule_id), - ) + if expected_last_run is None: + cursor = self._db.execute( + "UPDATE agent_schedules SET last_run=? WHERE id=?", + (ts, schedule_id), + ) + else: + cursor = self._db.execute( + "UPDATE agent_schedules SET last_run=? WHERE id=? AND last_run=?", + (ts, schedule_id, expected_last_run), + ) self._db.commit() + return cursor.rowcount > 0 def update_schedule_last_delivered( self, schedule_id: int, timestamp: float = 0.0 diff --git a/src/pinky_daemon/scheduler.py b/src/pinky_daemon/scheduler.py index 4be79023..815f7b2c 100644 --- a/src/pinky_daemon/scheduler.py +++ b/src/pinky_daemon/scheduler.py @@ -383,13 +383,23 @@ async def _check_schedules(self, now: float) -> None: continue if cron_matches(schedule.cron, dt): + claimed = self._registry.update_schedule_last_run( + schedule.id, + now, + expected_last_run=schedule.last_run, + ) + if not claimed: + _log( + f"scheduler: lost last_run claim race for " + f"#{schedule.id} — skipping fire" + ) + continue _log(f"scheduler: firing schedule '{schedule.name}' for agent '{schedule.agent_name}' (direct_send={schedule.direct_send}, one_shot={schedule.one_shot})") if self._activity: try: self._activity.log(schedule.agent_name, "schedule_fired", f"Schedule '{schedule.name}' fired") except Exception: pass - self._registry.update_schedule_last_run(schedule.id, now) # Carry the exact fire identity on this queued snapshot. A # later minute can advance the DB row while this cohort still # waits behind a long turn, so failure paths must never reread @@ -572,7 +582,7 @@ async def _deliver_schedule( _log( f"scheduler: delivery confirmed for schedule " f"'{schedule.name}' (#{schedule.id}) for agent " - f"'{schedule.agent_name}'" + f"'{schedule.agent_name}' (fired_at={schedule.last_run})" ) return diff --git a/tests/test_scheduler.py b/tests/test_scheduler.py index e070edae..ecb900b0 100644 --- a/tests/test_scheduler.py +++ b/tests/test_scheduler.py @@ -357,12 +357,35 @@ def test_update_last_run(self, registry): assert s.last_delivered == 0.0 now = time.time() - registry.update_schedule_last_run(s.id, now) + assert registry.update_schedule_last_run(s.id, now) is True schedules = registry.get_schedules("oleg") assert schedules[0].last_run == pytest.approx(now, abs=0.1) assert schedules[0].last_delivered == 0.0 + def test_update_last_run_compare_and_swap(self, registry): + registry.register("oleg") + schedule = registry.add_schedule("oleg", "0 8 * * *") + contender = AgentRegistry(db_path=registry._db_path) + try: + contender_snapshot = contender.get_schedules("oleg")[0] + + assert registry.update_schedule_last_run( + schedule.id, + 100.0, + expected_last_run=schedule.last_run, + ) is True + assert contender.update_schedule_last_run( + schedule.id, + 200.0, + expected_last_run=contender_snapshot.last_run, + ) is False + + stored = registry.get_schedules("oleg")[0] + assert stored.last_run == 100.0 + finally: + contender.close() + def test_update_last_delivered_is_distinct_from_last_run(self, registry): registry.register("oleg") schedule = registry.add_schedule("oleg", "0 8 * * *") @@ -897,6 +920,103 @@ def confirm_delivery(): "schedule_delivered", ] + @pytest.mark.asyncio + async def test_primary_confirm_log_includes_fire_identity( + self, registry, capsys + ): + registry.register("oleg") + schedule = registry.add_schedule( + "oleg", "* * * * *", name="forensic", prompt="run once" + ) + fired_at = 1_800_000_000.0 + assert registry.update_schedule_last_run( + schedule.id, + fired_at, + expected_last_run=schedule.last_run, + ) is True + schedule.last_run = fired_at + + async def confirmed(agent_name, session_id, prompt): + del agent_name, session_id, prompt + return True + + scheduler = AgentScheduler(registry, wake_callback=confirmed) + await scheduler._deliver_schedule(schedule) + + assert ( + f"scheduler: delivery confirmed for schedule 'forensic' " + f"(#{schedule.id}) for agent 'oleg' (fired_at={fired_at})" + ) in capsys.readouterr().err + + @pytest.mark.asyncio + async def test_lost_last_run_claim_race_skips_fire( + self, registry, monkeypatch, capsys + ): + registry.register("oleg") + schedule = registry.add_schedule( + "oleg", "* * * * *", name="claimed", prompt="run once" + ) + wake_calls: list[str] = [] + events: list[str] = [] + real_update = registry.update_schedule_last_run + competing_claim_made = False + + def lose_to_competing_claim( + schedule_id, + timestamp=0.0, + *, + expected_last_run=None, + ): + nonlocal competing_claim_made + if expected_last_run is not None and not competing_claim_made: + competing_claim_made = True + assert real_update( + schedule_id, + timestamp, + expected_last_run=expected_last_run, + ) is True + return real_update( + schedule_id, + timestamp, + expected_last_run=expected_last_run, + ) + + class Activity: + def log(self, agent_name, event_type, summary): + del agent_name, summary + events.append(event_type) + + async def wake_cb(agent_name, session_id, prompt): + del agent_name, session_id + wake_calls.append(prompt) + return True + + monkeypatch.setattr( + registry, + "update_schedule_last_run", + lose_to_competing_claim, + ) + scheduler = AgentScheduler( + registry, wake_callback=wake_cb, activity=Activity() + ) + fired_at = 1_800_000_000.0 + + await scheduler._check_schedules(fired_at) + + stored = registry.get_schedules("oleg")[0] + assert competing_claim_made is True + assert stored.last_run == fired_at + assert stored.last_delivered == 0.0 + assert wake_calls == [] + assert events == [] + assert scheduler._schedule_delivery_tasks == set() + error_log = capsys.readouterr().err + assert ( + f"scheduler: lost last_run claim race for " + f"#{schedule.id} — skipping fire" + ) in error_log + assert "scheduler: firing schedule" not in error_log + @pytest.mark.asyncio async def test_unconfirmed_delivery_stays_fired_but_undelivered( self, registry, capsys