feat: out-of-process usage-snapshot builder (Phase 1) (#1468)

* feat: out-of-process usage-snapshot builder (Phase 1)

Move all usage-tab compute off the run-serving HTTP process (the #1434 bug
class). Read path is now a pure cache read with a typed computing placeholder
on cold miss; a dedicated usage_snapshot graph + global ~10min cron rebuild
every period's snapshot out-of-process, wrapped in asyncio.timeout and gated by
USAGE_SNAPSHOT_CRON_ENABLED. Lifespan makes one fire-and-forget scheduling call,
never a retained/looping task.

* fix: address PR review on usage-snapshot builder

- admin guard on /admin/usage/rebuild (was any logged-in user)
- cap lifespan loopback calls with asyncio.timeout(5) so a startup hang
  can't block boot
- memoize cron registration so steady-state requests skip the loopback
  check; reap duplicate crons from concurrent-replica races
- thread the computing flag through the usage payload too
This commit is contained in:
Johannes du Plessis 2026-06-09 11:11:06 -07:00 • committed by GitHub
parent f349311d20
commit 798cc6edf1
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
8 changed files with 467 additions and 41 deletions

View file

@ -5,7 +5,7 @@ from __future__ import annotations
import asyncio
import logging
from collections import Counter
from collections.abc import Awaitable, Callable
from collections.abc import Callable
from datetime import UTC, datetime, timedelta
from typing import Any, Literal
@ -25,7 +25,6 @@ _AGENT_SOURCES = frozenset({"dashboard", "github", "slack", "linear"})
_PR_REFRESH_INTERVAL_MS = 10 * 60 * 1000
_MAX_PR_REFRESH_PER_REQUEST = 25
_PR_REFRESH_CONCURRENCY = 5
_CACHE_TTL_MS = 10 * 60 * 1000
_CACHE_SEARCH_LIMIT = 1000
_GITHUB_API = "https://api.github.com"
@ -605,24 +604,45 @@ async def refresh_reviewer_stats_cache(period: str | None = "30d") -> dict[str,
return snapshot
def _empty_usage_snapshot(period: Period) -> dict[str, Any]:
return {"period": period, "users": [], "total_members": 0, "computing": True}
def _empty_reviewer_snapshot(period: Period) -> dict[str, Any]:
return {
"period": period,
"reviewed_prs": 0,
"prs_with_findings": 0,
"findings_recorded": 0,
"surfaced_findings": 0,
"addressed_findings": 0,
"resolved_after_update": 0,
"dismissed_findings": 0,
"unresolved_surfaced_findings": 0,
"resolution_rate": 0.0,
"human_replies": 0,
"severity_counts": {},
"top_categories": [],
"computing": True,
}
async def _cached_snapshot(
namespace: list[str],
period: Period,
refresh: Callable[[str | None], Awaitable[dict[str, Any]]],
*,
schedule_refresh: Callable[[Period], None] | None = None,
empty: Callable[[Period], dict[str, Any]],
) -> tuple[dict[str, Any], int | None]:
"""Pure read: return the precomputed snapshot (any age) or a typed placeholder.
Never computes inline and never schedules a refresh — the scheduled builder
owns refresh cadence (see ``agent/usage_snapshot.py``).
"""
cached = await _get_value(namespace, period)
if cached:
snapshot = cached.get("snapshot")
generated_at_ms = _coerce_int(cached.get("generated_at_ms"))
if isinstance(snapshot, dict):
if not generated_at_ms or _now_ms() - generated_at_ms <= _CACHE_TTL_MS:
return snapshot, generated_at_ms or None
if schedule_refresh is not None:
schedule_refresh(period)
return snapshot, generated_at_ms
return await refresh(period), _now_ms()
return snapshot, _coerce_int(cached.get("generated_at_ms")) or None
return empty(period), None
def _usage_payload_from_snapshot(
@ -677,6 +697,7 @@ def _usage_payload_from_snapshot(
"total_members": _coerce_int(snapshot.get("total_members")),
"current_user_rank": current_user_row["rank"] if current_user_row else None,
"generated_at_ms": generated_at_ms,
"computing": bool(snapshot.get("computing")),
}
@ -686,22 +707,18 @@ async def list_agent_usage_leaderboard(
limit: int,
current_login: str | None,
current_email: str | None,
schedule_usage_refresh: Callable[[Period], None] | None = None,
schedule_reviewer_refresh: Callable[[Period], None] | None = None,
) -> dict[str, Any]:
"""Return cached Open SWE Agent and reviewer usage stats."""
"""Return cached Open SWE Agent and reviewer usage stats (pure cache read)."""
normalized_period = _normalize_period(period)
usage_snapshot, generated_at_ms = await _cached_snapshot(
USAGE_LEADERBOARD_CACHE_NAMESPACE,
normalized_period,
refresh_usage_leaderboard_cache,
schedule_refresh=schedule_usage_refresh,
_empty_usage_snapshot,
)
reviewer_stats, reviewer_generated_at_ms = await _cached_snapshot(
REVIEWER_STATS_CACHE_NAMESPACE,
normalized_period,
refresh_reviewer_stats_cache,
schedule_refresh=schedule_reviewer_refresh,
_empty_reviewer_snapshot,
)
payload = _usage_payload_from_snapshot(
usage_snapshot,

View file

@ -8,16 +8,12 @@ import os
from typing import Any
import httpx
from fastapi import APIRouter, BackgroundTasks, Depends, HTTPException, Request
from fastapi import APIRouter, Depends, HTTPException, Request
from fastapi.responses import RedirectResponse, Response, StreamingResponse
from pydantic import BaseModel
from .admin import is_admin
from .agent_usage import (
list_agent_usage_leaderboard,
refresh_reviewer_stats_cache,
refresh_usage_leaderboard_cache,
)
from .agent_usage import list_agent_usage_leaderboard
from .analyzer_cron import remove_continual_cron
from .enabled_repos import (
list_enabled_review_repos,
@ -97,6 +93,7 @@ from .thread_api import (
send_dashboard_message,
stream_dashboard_thread,
)
from .usage_snapshot_cron import ensure_usage_snapshot_cron, trigger_usage_snapshot_build
from .user_mappings import (
delete_mapping,
get_mapping,
@ -710,25 +707,27 @@ async def api_delete_review_style(
@router.get("/agent-usage-leaderboard")
async def api_agent_usage_leaderboard(
background_tasks: BackgroundTasks,
period: str | None = "30d",
limit: int = 10,
session: dict[str, Any] = _SESSION_DEP,
) -> dict[str, Any]:
await ensure_usage_snapshot_cron()
return await list_agent_usage_leaderboard(
period=period,
limit=limit,
current_login=session["sub"],
current_email=session.get("email"),
schedule_usage_refresh=lambda cache_period: background_tasks.add_task(
refresh_usage_leaderboard_cache, cache_period
),
schedule_reviewer_refresh=lambda cache_period: background_tasks.add_task(
refresh_reviewer_stats_cache, cache_period
),
)
@router.post("/admin/usage/rebuild")
async def api_admin_usage_rebuild(
session: dict[str, Any] = _ADMIN_DEP,
) -> dict[str, Any]:
await ensure_usage_snapshot_cron()
return await trigger_usage_snapshot_build()
@router.get("/schedules")
async def api_list_schedules(
session: dict[str, Any] = _SESSION_DEP,

View file

@ -0,0 +1,110 @@
"""Global cron + on-demand scheduling for the usage-snapshot builder.
One global cron rebuilds every period's snapshot ~every 10 min (staggered off
:00). Registration is idempotent and concrete (mirrors ``analyzer_cron``): we
look for an existing cron tagged ``{"kind": "usage_snapshot"}`` before creating
one, and stash the ``cron_id`` under ``["agent_usage", "meta"] / "cron"``.
"""
from __future__ import annotations
import logging
from typing import Any
from .review_style_jobs import _client
logger = logging.getLogger(__name__)
_ASSISTANT_ID = "usage_snapshot"
_CRON_KIND = "usage_snapshot"
_SCHEDULE = "7,17,27,37,47,57 * * * *" # ~every 10 min, staggered off :00
_META_NAMESPACE: list[str] = ["agent_usage", "meta"]
_META_CRON_KEY = "cron"
# Once we've confirmed a cron is registered, steady-state requests skip the
# loopback check entirely — there is nothing left to re-verify.
_registered_cron_id: str | None = None
def _cron_id_of(cron: Any) -> str | None:
cron_id = cron.get("cron_id") if isinstance(cron, dict) else getattr(cron, "cron_id", None)
return cron_id if isinstance(cron_id, str) and cron_id else None
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
return cron_id if isinstance(cron_id, str) and cron_id else None
async def ensure_usage_snapshot_cron() -> str | None:
"""Idempotently register the global usage-snapshot cron. Returns its id."""
global _registered_cron_id
if _registered_cron_id:
return _registered_cron_id
try:
existing = await _existing_cron_id()
except Exception:
logger.debug("Could not read usage snapshot cron meta", exc_info=True)
existing = None
if existing:
_registered_cron_id = existing
return existing
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]
# 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 _store_cron_id(keep)
_registered_cron_id = keep
return keep
except Exception:
logger.debug("Usage snapshot cron search failed", exc_info=True)
try:
cron = await _client().crons.create(
_ASSISTANT_ID,
schedule=_SCHEDULE,
metadata={"kind": _CRON_KIND},
)
except Exception:
logger.exception("Failed to create usage snapshot cron")
return None
cron_id = _cron_id_of(cron)
if cron_id:
await _store_cron_id(cron_id)
_registered_cron_id = cron_id
return cron_id
return None
async def _delete_crons(cron_ids: list[str]) -> None:
for cron_id in cron_ids:
try:
await _client().crons.delete(cron_id)
except Exception:
logger.debug("Could not delete duplicate usage cron %s", cron_id, exc_info=True)
async def _store_cron_id(cron_id: str) -> None:
try:
await _client().store.put_item(_META_NAMESPACE, _META_CRON_KEY, {"cron_id": cron_id})
except Exception:
logger.debug("Could not persist usage snapshot cron id", exc_info=True)
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)
except Exception:
logger.debug("Could not schedule immediate usage snapshot build", exc_info=True)
return {"status": "error"}
return {"status": "scheduled"}

75
agent/usage_snapshot.py Normal file
View file

@ -0,0 +1,75 @@
"""Out-of-process builder graph for the usage-tab snapshots.
A single pure-python node (no LLM, no sandbox) that rebuilds every period's
usage leaderboard + reviewer-stats snapshot and writes them to the cache
namespaces. Scheduled by a global cron (see ``dashboard/usage_snapshot_cron.py``)
so heavy/looping work never runs on the run-serving HTTP process (the #1434 bug
class). The whole build is bounded by ``asyncio.timeout`` so a wedged build
self-cancels before the next tick.
"""
from __future__ import annotations
import asyncio
import logging
import os
from typing import Any, TypedDict
from langgraph.graph import END, START, StateGraph
from langgraph.graph.state import RunnableConfig
from .dashboard.agent_usage import (
refresh_reviewer_stats_cache,
refresh_usage_leaderboard_cache,
)
logger = logging.getLogger(__name__)
_PERIODS = ("7d", "30d", "all")
_BUILD_TIMEOUT_S = 120
class UsageSnapshotState(TypedDict, total=False):
result: dict[str, Any]
def _cron_enabled() -> bool:
value = os.environ.get("USAGE_SNAPSHOT_CRON_ENABLED", "true").strip().lower()
return value not in {"0", "false", "no", "off"}
async def _build(state: UsageSnapshotState, config: RunnableConfig) -> dict[str, Any]:
if not _cron_enabled():
logger.info("Usage snapshot build skipped: USAGE_SNAPSHOT_CRON_ENABLED is off")
return {"result": {"status": "disabled"}}
built: list[str] = []
failed: list[str] = []
try:
async with asyncio.timeout(_BUILD_TIMEOUT_S):
for period in _PERIODS:
for label, refresh in (
("usage", refresh_usage_leaderboard_cache),
("reviewer", refresh_reviewer_stats_cache),
):
try:
await refresh(period)
built.append(f"{label}:{period}")
except Exception:
logger.exception("Usage snapshot build failed for %s:%s", label, period)
failed.append(f"{label}:{period}")
except TimeoutError:
logger.warning("Usage snapshot build exceeded %ss budget; cancelled", _BUILD_TIMEOUT_S)
return {"result": {"status": "timeout", "built": built, "failed": failed}}
status = "ok" if not failed else "partial"
logger.info("Usage snapshot build %s: built=%s failed=%s", status, built, failed)
return {"result": {"status": status, "built": built, "failed": failed}}
def get_usage_snapshot(config: RunnableConfig | None = None):
builder = StateGraph(UsageSnapshotState)
builder.add_node("build", _build)
builder.add_edge(START, "build")
builder.add_edge("build", END)
return builder.compile().with_config(config or {})

View file

@ -1,5 +1,6 @@
"""Custom FastAPI routes for LangGraph server."""
import asyncio
import hashlib
import hmac
import json
@ -110,9 +111,26 @@ logger = logging.getLogger(__name__)
@asynccontextmanager
async def lifespan(_app: FastAPI) -> AsyncIterator[None]:
from .dashboard.usage_snapshot_cron import (
ensure_usage_snapshot_cron,
trigger_usage_snapshot_build,
)
from .utils.sandbox import validate_sandbox_startup_config
validate_sandbox_startup_config()
# One bounded scheduling call — never a retained/looping task (the exact
# #1434 bug class). These are loopback HTTP calls into a server that is
# still starting, so cap them hard: a hang must not block startup. Any
# failure/timeout is fine — the cron and lazy per-request registration both
# catch up within the cadence.
try:
async with asyncio.timeout(5):
await ensure_usage_snapshot_cron()
await trigger_usage_snapshot_build()
except Exception:
logger.debug("Usage snapshot bootstrap on startup skipped", exc_info=True)
yield

View file

@ -5,7 +5,8 @@
"agent": "agent.server:get_agent",
"reviewer": "agent.reviewer:get_reviewer_agent",
"analyzer": "agent.analyzer:get_analyzer",
"scheduler": "agent.scheduler:get_scheduler"
"scheduler": "agent.scheduler:get_scheduler",
"usage_snapshot": "agent.usage_snapshot:get_usage_snapshot"
},
"dependencies": ["."],
"http": {

View file

@ -38,7 +38,7 @@ class FakeClient:
@pytest.mark.asyncio
async def test_cached_usage_payload_returns_stale_snapshot_and_schedules_refresh(monkeypatch):
async def test_cached_usage_payload_returns_snapshot_without_refresh(monkeypatch):
usage_snapshot = {
"period": "30d",
"total_members": 2,
@ -101,27 +101,54 @@ async def test_cached_usage_payload_returns_stale_snapshot_and_schedules_refresh
}
)
monkeypatch.setattr(agent_usage, "_client", lambda: FakeClient(store=store))
monkeypatch.setattr(agent_usage, "_now_ms", lambda: agent_usage._CACHE_TTL_MS + 2)
usage_refreshes: list[str] = []
reviewer_refreshes: list[str] = []
refreshed: list[str] = []
monkeypatch.setattr(
agent_usage,
"refresh_usage_leaderboard_cache",
lambda *a, **k: refreshed.append("usage"),
)
monkeypatch.setattr(
agent_usage,
"refresh_reviewer_stats_cache",
lambda *a, **k: refreshed.append("reviewer"),
)
payload = await agent_usage.list_agent_usage_leaderboard(
period="30d",
limit=1,
current_login="octo",
current_email="octo@example.com",
schedule_usage_refresh=usage_refreshes.append,
schedule_reviewer_refresh=reviewer_refreshes.append,
)
assert usage_refreshes == ["30d"]
assert reviewer_refreshes == ["30d"]
# Pure read: stale snapshot is served as-is, never refreshed inline.
assert refreshed == []
assert payload["rows"][0]["user"]["email"] == "octo@example.com"
assert payload["total_members"] == 2
assert payload["reviewer_stats"]["surfaced_findings"] == 1
assert store.puts == []
@pytest.mark.asyncio
async def test_cold_cache_returns_computing_placeholder(monkeypatch):
store = FakeStore({})
monkeypatch.setattr(agent_usage, "_client", lambda: FakeClient(store=store))
payload = await agent_usage.list_agent_usage_leaderboard(
period="7d",
limit=10,
current_login="octo",
current_email="octo@example.com",
)
assert payload["rows"] == []
assert payload["total_members"] == 0
assert payload["generated_at_ms"] is None
assert payload["computing"] is True
assert payload["reviewer_stats"]["computing"] is True
assert store.puts == []
@pytest.mark.asyncio
async def test_reviewer_stats_snapshot_counts_surfaced_and_resolved_findings(monkeypatch):
threads = [

View file

@ -0,0 +1,179 @@
from __future__ import annotations
import inspect
import pytest
from agent import usage_snapshot
from agent.dashboard import usage_snapshot_cron
class FakeStore:
def __init__(self) -> None:
self.values: dict[tuple[tuple[str, ...], str], dict] = {}
async def get_item(self, namespace, key):
value = self.values.get((tuple(namespace), key))
return {"value": value} if value is not None else None
async def put_item(self, namespace, key, value):
self.values[(tuple(namespace), key)] = value
class FakeCrons:
def __init__(self, existing: list[dict] | None = None) -> None:
self.existing = existing or []
self.created: list[dict] = []
self.deleted: list[str] = []
async def search(self, **kwargs):
return self.existing
async def create(self, assistant_id, **kwargs):
self.created.append({"assistant_id": assistant_id, **kwargs})
return {"cron_id": "cron-new"}
async def delete(self, cron_id):
self.deleted.append(cron_id)
@pytest.fixture(autouse=True)
def _reset_cron_memo():
usage_snapshot_cron._registered_cron_id = None
yield
usage_snapshot_cron._registered_cron_id = None
class FakeRuns:
def __init__(self) -> None:
self.created: list[tuple] = []
async def create(self, thread, assistant_id, **kwargs):
self.created.append((thread, assistant_id))
return {"run_id": "run-1"}
class FakeClient:
def __init__(self, store=None, crons=None, runs=None):
self.store = store or FakeStore()
self.crons = crons or FakeCrons()
self.runs = runs or FakeRuns()
@pytest.mark.asyncio
async def test_builder_refreshes_all_periods(monkeypatch):
calls: list[tuple[str, str]] = []
async def fake_usage(period):
calls.append(("usage", period))
async def fake_reviewer(period):
calls.append(("reviewer", period))
monkeypatch.setattr(usage_snapshot, "refresh_usage_leaderboard_cache", fake_usage)
monkeypatch.setattr(usage_snapshot, "refresh_reviewer_stats_cache", fake_reviewer)
monkeypatch.setenv("USAGE_SNAPSHOT_CRON_ENABLED", "true")
result = await usage_snapshot._build({}, {})
assert result["result"]["status"] == "ok"
assert ("usage", "7d") in calls
assert ("reviewer", "all") in calls
assert len(calls) == 6
@pytest.mark.asyncio
async def test_builder_kill_switch(monkeypatch):
monkeypatch.setenv("USAGE_SNAPSHOT_CRON_ENABLED", "false")
result = await usage_snapshot._build({}, {})
assert result["result"]["status"] == "disabled"
@pytest.mark.asyncio
async def test_builder_partial_on_failure(monkeypatch):
async def fake_usage(period):
raise RuntimeError("boom")
async def fake_reviewer(period):
return None
monkeypatch.setattr(usage_snapshot, "refresh_usage_leaderboard_cache", fake_usage)
monkeypatch.setattr(usage_snapshot, "refresh_reviewer_stats_cache", fake_reviewer)
monkeypatch.setenv("USAGE_SNAPSHOT_CRON_ENABLED", "true")
result = await usage_snapshot._build({}, {})
assert result["result"]["status"] == "partial"
assert len(result["result"]["failed"]) == 3
@pytest.mark.asyncio
async def test_cron_registration_idempotent_and_creates_once(monkeypatch):
client = FakeClient()
monkeypatch.setattr(usage_snapshot_cron, "_client", lambda: client)
first = await usage_snapshot_cron.ensure_usage_snapshot_cron()
assert first == "cron-new"
assert len(client.crons.created) == 1
# Second call reads the persisted id, creates nothing new.
second = await usage_snapshot_cron.ensure_usage_snapshot_cron()
assert second == "cron-new"
assert len(client.crons.created) == 1
@pytest.mark.asyncio
async def test_cron_reuses_existing_search_hit(monkeypatch):
client = FakeClient(crons=FakeCrons(existing=[{"cron_id": "cron-existing"}]))
monkeypatch.setattr(usage_snapshot_cron, "_client", lambda: client)
cron_id = await usage_snapshot_cron.ensure_usage_snapshot_cron()
assert cron_id == "cron-existing"
assert client.crons.created == []
@pytest.mark.asyncio
async def test_cron_reaps_duplicate_search_hits(monkeypatch):
client = FakeClient(
crons=FakeCrons(existing=[{"cron_id": "cron-a"}, {"cron_id": "cron-b"}]),
)
monkeypatch.setattr(usage_snapshot_cron, "_client", lambda: client)
cron_id = await usage_snapshot_cron.ensure_usage_snapshot_cron()
assert cron_id == "cron-a"
assert client.crons.deleted == ["cron-b"]
assert client.crons.created == []
@pytest.mark.asyncio
async def test_cron_memoizes_and_skips_loopback(monkeypatch):
store = FakeStore()
client = FakeClient(store=store)
monkeypatch.setattr(usage_snapshot_cron, "_client", lambda: client)
await usage_snapshot_cron.ensure_usage_snapshot_cron()
# After registration, a poisoned client proves no further loopback happens.
def _boom():
raise AssertionError("ensure_usage_snapshot_cron should be memoized")
monkeypatch.setattr(usage_snapshot_cron, "_client", _boom)
assert await usage_snapshot_cron.ensure_usage_snapshot_cron() == "cron-new"
@pytest.mark.asyncio
async def test_trigger_build_schedules_threadless_run(monkeypatch):
client = FakeClient()
monkeypatch.setattr(usage_snapshot_cron, "_client", lambda: client)
result = await usage_snapshot_cron.trigger_usage_snapshot_build()
assert result["status"] == "scheduled"
assert client.runs.created == [(None, "usage_snapshot")]
def test_lifespan_retains_no_background_task():
"""Regression guard for the #1434 bug: lifespan must not spawn a looping task."""
from agent import webapp
src = inspect.getsource(webapp.lifespan)
assert "create_task" not in src
assert "while True" not in src