From 49aa638c5862bbe3e7074a5943584afd14b54e05 Mon Sep 17 00:00:00 2001 From: Johannes du Plessis Date: Tue, 9 Jun 2026 12:53:15 -0700 Subject: [PATCH] fix: schedule usage snapshot runs with explicit empty input (#1470) --- agent/dashboard/usage_snapshot_cron.py | 39 ++++++++++++++++++----- tests/test_usage_snapshot.py | 43 +++++++++++++++++++++++--- 2 files changed, 70 insertions(+), 12 deletions(-) diff --git a/agent/dashboard/usage_snapshot_cron.py b/agent/dashboard/usage_snapshot_cron.py index 646c2653..99684c15 100644 --- a/agent/dashboard/usage_snapshot_cron.py +++ b/agent/dashboard/usage_snapshot_cron.py @@ -17,6 +17,11 @@ logger = logging.getLogger(__name__) _ASSISTANT_ID = "usage_snapshot" _CRON_KIND = "usage_snapshot" +# A StateGraph run with no input raises EmptyInputError at __start__, so every +# scheduled run must carry an explicit empty input. rev marks crons created +# with the correct payload; older revs are reaped and recreated. +_CRON_INPUT: dict[str, Any] = {} +_CRON_REV = 2 _SCHEDULE = "7,17,27,37,47,57 * * * *" # ~every 10 min, staggered off :00 _META_NAMESPACE: list[str] = ["agent_usage", "meta"] _META_CRON_KEY = "cron" @@ -31,10 +36,18 @@ def _cron_id_of(cron: Any) -> str | None: return cron_id if isinstance(cron_id, str) and cron_id else None +def _cron_rev_of(cron: Any) -> int: + metadata = cron.get("metadata") if isinstance(cron, dict) else getattr(cron, "metadata", None) + rev = metadata.get("rev") if isinstance(metadata, dict) else None + return rev if isinstance(rev, int) else 1 + + async def _existing_cron_id() -> str | None: item = await _client().store.get_item(_META_NAMESPACE, _META_CRON_KEY) value = item.get("value") if isinstance(item, dict) else getattr(item, "value", None) - cron_id = value.get("cron_id") if isinstance(value, dict) else None + if not isinstance(value, dict) or value.get("rev") != _CRON_REV: + return None + cron_id = value.get("cron_id") return cron_id if isinstance(cron_id, str) and cron_id else None @@ -55,12 +68,19 @@ async def ensure_usage_snapshot_cron() -> str | None: try: crons = await _client().crons.search(metadata={"kind": _CRON_KIND}, limit=10) - cron_ids = [cid for c in (crons or []) if (cid := _cron_id_of(c))] - if cron_ids: - keep = cron_ids[0] + current = [ + cid for c in (crons or []) if (cid := _cron_id_of(c)) and _cron_rev_of(c) == _CRON_REV + ] + stale = [ + cid for c in (crons or []) if (cid := _cron_id_of(c)) and _cron_rev_of(c) != _CRON_REV + ] + # Stale revs were created without input and fail at __start__; replace them. + await _delete_crons(stale) + if current: + keep = current[0] # search-then-create isn't atomic; concurrent replicas can each # create one. Reap any extras so we don't fire the build N times. - await _delete_crons(cron_ids[1:]) + await _delete_crons(current[1:]) await _store_cron_id(keep) _registered_cron_id = keep return keep @@ -71,7 +91,8 @@ async def ensure_usage_snapshot_cron() -> str | None: cron = await _client().crons.create( _ASSISTANT_ID, schedule=_SCHEDULE, - metadata={"kind": _CRON_KIND}, + input=_CRON_INPUT, + metadata={"kind": _CRON_KIND, "rev": _CRON_REV}, ) except Exception: logger.exception("Failed to create usage snapshot cron") @@ -95,7 +116,9 @@ async def _delete_crons(cron_ids: list[str]) -> None: async def _store_cron_id(cron_id: str) -> None: try: - await _client().store.put_item(_META_NAMESPACE, _META_CRON_KEY, {"cron_id": cron_id}) + await _client().store.put_item( + _META_NAMESPACE, _META_CRON_KEY, {"cron_id": cron_id, "rev": _CRON_REV} + ) except Exception: logger.debug("Could not persist usage snapshot cron id", exc_info=True) @@ -103,7 +126,7 @@ async def _store_cron_id(cron_id: str) -> None: async def trigger_usage_snapshot_build() -> dict[str, Any]: """Fire-and-forget: schedule (not execute) one immediate build run.""" try: - await _client().runs.create(None, _ASSISTANT_ID) + await _client().runs.create(None, _ASSISTANT_ID, input=_CRON_INPUT) except Exception: logger.debug("Could not schedule immediate usage snapshot build", exc_info=True) return {"status": "error"} diff --git a/tests/test_usage_snapshot.py b/tests/test_usage_snapshot.py index c95e972c..669612a8 100644 --- a/tests/test_usage_snapshot.py +++ b/tests/test_usage_snapshot.py @@ -49,7 +49,7 @@ class FakeRuns: self.created: list[tuple] = [] async def create(self, thread, assistant_id, **kwargs): - self.created.append((thread, assistant_id)) + self.created.append((thread, assistant_id, kwargs)) return {"run_id": "run-1"} @@ -114,6 +114,9 @@ async def test_cron_registration_idempotent_and_creates_once(monkeypatch): first = await usage_snapshot_cron.ensure_usage_snapshot_cron() assert first == "cron-new" assert len(client.crons.created) == 1 + created = client.crons.created[0] + assert created["input"] == {} + assert created["metadata"]["rev"] == usage_snapshot_cron._CRON_REV # Second call reads the persisted id, creates nothing new. second = await usage_snapshot_cron.ensure_usage_snapshot_cron() @@ -123,7 +126,10 @@ async def test_cron_registration_idempotent_and_creates_once(monkeypatch): @pytest.mark.asyncio async def test_cron_reuses_existing_search_hit(monkeypatch): - client = FakeClient(crons=FakeCrons(existing=[{"cron_id": "cron-existing"}])) + rev = usage_snapshot_cron._CRON_REV + client = FakeClient( + crons=FakeCrons(existing=[{"cron_id": "cron-existing", "metadata": {"rev": rev}}]) + ) monkeypatch.setattr(usage_snapshot_cron, "_client", lambda: client) cron_id = await usage_snapshot_cron.ensure_usage_snapshot_cron() @@ -133,8 +139,14 @@ async def test_cron_reuses_existing_search_hit(monkeypatch): @pytest.mark.asyncio async def test_cron_reaps_duplicate_search_hits(monkeypatch): + rev = usage_snapshot_cron._CRON_REV client = FakeClient( - crons=FakeCrons(existing=[{"cron_id": "cron-a"}, {"cron_id": "cron-b"}]), + crons=FakeCrons( + existing=[ + {"cron_id": "cron-a", "metadata": {"rev": rev}}, + {"cron_id": "cron-b", "metadata": {"rev": rev}}, + ] + ), ) monkeypatch.setattr(usage_snapshot_cron, "_client", lambda: client) @@ -144,6 +156,29 @@ async def test_cron_reaps_duplicate_search_hits(monkeypatch): assert client.crons.created == [] +@pytest.mark.asyncio +async def test_cron_replaces_stale_rev(monkeypatch): + """Crons created without input (rev 1) fail at __start__ — reap and recreate.""" + client = FakeClient(crons=FakeCrons(existing=[{"cron_id": "cron-old"}])) + monkeypatch.setattr(usage_snapshot_cron, "_client", lambda: client) + + cron_id = await usage_snapshot_cron.ensure_usage_snapshot_cron() + assert cron_id == "cron-new" + assert client.crons.deleted == ["cron-old"] + assert client.crons.created[0]["input"] == {} + + +@pytest.mark.asyncio +async def test_stale_stored_cron_id_is_ignored(monkeypatch): + store = FakeStore() + store.values[(("agent_usage", "meta"), "cron")] = {"cron_id": "cron-old"} + client = FakeClient(store=store) + monkeypatch.setattr(usage_snapshot_cron, "_client", lambda: client) + + cron_id = await usage_snapshot_cron.ensure_usage_snapshot_cron() + assert cron_id == "cron-new" + + @pytest.mark.asyncio async def test_cron_memoizes_and_skips_loopback(monkeypatch): store = FakeStore() @@ -167,7 +202,7 @@ async def test_trigger_build_schedules_threadless_run(monkeypatch): result = await usage_snapshot_cron.trigger_usage_snapshot_build() assert result["status"] == "scheduled" - assert client.runs.created == [(None, "usage_snapshot")] + assert client.runs.created == [(None, "usage_snapshot", {"input": {}})] def test_lifespan_retains_no_background_task():