diff --git a/backend/druks/core/tasks.py b/backend/druks/core/tasks.py index cc74bc2e..d8289d56 100644 --- a/backend/druks/core/tasks.py +++ b/backend/druks/core/tasks.py @@ -4,10 +4,12 @@ from druks.harnesses.datastructures import RotationResult from druks.harnesses.directory import refresh_added_catalogs from druks.harnesses.providers import get_provider, get_providers +from druks.models import Base from druks.sandbox import gate from druks.sandbox.client import sandbox_client from druks.sandbox.models import SandboxIdentity from druks.secrets.models import VaultSecret +from druks.usage.models import UsageScrape from druks.workflows import task logger = logging.getLogger(__name__) @@ -27,6 +29,13 @@ async def refresh_tokens() -> None: await _refresh() +@task(every="*/5 * * * *") +async def refresh_usage() -> None: + for subscription in await VaultSecret.list_subscriptions(): + if await UsageScrape.is_due(subscription, now=Base.utc_now()): + await get_provider(subscription.audience_name).poll_usage(subscription) + + @task(every="0 * * * *") async def release_orphan_boxes() -> None: await _release_orphan_boxes() diff --git a/backend/druks/usage/models.py b/backend/druks/usage/models.py index bcb8fd76..d3460bf1 100644 --- a/backend/druks/usage/models.py +++ b/backend/druks/usage/models.py @@ -1,5 +1,6 @@ from datetime import datetime, timedelta -from typing import Any +from itertools import pairwise +from typing import TYPE_CHECKING, Any from sqlalchemy import ForeignKey, Index, delete, select from sqlalchemy.dialects.postgresql import JSONB @@ -7,6 +8,11 @@ from druks.db import Base, db_session +if TYPE_CHECKING: + from druks.secrets.models import VaultSecret + +_POLL_INTERVAL_MINUTES = (5, 10, 20, 40, 60) + class UsageScrape(Base): __tablename__ = "usage_scrapes" @@ -31,8 +37,7 @@ class UsageScrape(Base): # Subscription tier when the CLI surfaces it (e.g. ``pro``, ``max``, # ``plus``). Display-only. plan_tier: Mapped[str | None] - # Five-hour rolling window. Claude exposes this directly; Codex - # doesn't have a 5h concept yet so it stays null for the codex row. + # The provider's five-hour rolling window. five_hour_percent_left: Mapped[int | None] five_hour_resets_at: Mapped[datetime | None] # Weekly windows in provider order, including separately metered models. @@ -43,6 +48,52 @@ class UsageScrape(Base): # a quota bar that never moves. unlimited: Mapped[bool] = mapped_column(default=False) + @classmethod + async def is_due(cls, subscription: "VaultSecret", *, now: datetime) -> bool: + """Scrape history, completed calls, and window resets determine when a poll is due.""" + # Cycle: durable.models loads harnesses, whose providers load UsageScrape. + from druks.durable.models import AgentCall + + stmt = ( + select(cls) + .where( + cls.provider == subscription.audience_name, + cls.account_id == subscription.account_id, + ) + .order_by(cls.scraped_at.desc(), cls.id.desc()) + .limit(len(_POLL_INTERVAL_MINUTES)) + ) + rows = list(await db_session().scalars(stmt)) + + if not rows: + return True + latest_scrape = rows[0] + finished_call = select(AgentCall.id).where( + AgentCall.subscription_id == subscription.id, + AgentCall.finished_at > latest_scrape.scraped_at, + ) + + if await db_session().scalar(select(finished_call.exists())): + return True + exhausted_reset = latest_scrape.soonest_reset_after( + latest_scrape.scraped_at, exhausted_only=True + ) + + if exhausted_reset: + return now >= exhausted_reset + reset = latest_scrape.soonest_reset_after(latest_scrape.scraped_at) + + if reset and now >= reset: + return True + unchanged_pairs = 0 + + for newer_scrape, older_scrape in pairwise(rows): + if newer_scrape.quota != older_scrape.quota: + break + unchanged_pairs += 1 + interval = timedelta(minutes=_POLL_INTERVAL_MINUTES[unchanged_pairs]) + return now >= latest_scrape.scraped_at + interval + @classmethod async def latest_for(cls, provider_id: str, account_id: str) -> "UsageScrape | None": stmt = ( @@ -69,20 +120,34 @@ async def history_for( ) return list((await db_session().execute(stmt)).scalars()) + @property + def quota(self) -> tuple[str | None, int | None, list[int | None]]: + return ( + self.error, + self.five_hour_percent_left, + [week["percent_left"] for week in self.weeks], + ) + def binding_week(self) -> dict[str, Any] | None: """The window closest to exhaustion — whichever stops work first.""" reported_windows = [week for week in self.weeks if week["percent_left"] is not None] if reported_windows: return min(reported_windows, key=lambda week: week["percent_left"]) - def soonest_reset_after(self, now: datetime) -> datetime | None: + def soonest_reset_after( + self, after: datetime, *, exhausted_only: bool = False + ) -> datetime | None: resets = [] - if self.five_hour_resets_at and self.five_hour_resets_at > now: + if ( + self.five_hour_resets_at + and self.five_hour_resets_at > after + and (not exhausted_only or self.five_hour_percent_left == 0) + ): resets.append(self.five_hour_resets_at) for week in self.weeks: - if week["resets_at"]: + if week["resets_at"] and (not exhausted_only or week["percent_left"] == 0): reset = datetime.fromisoformat(week["resets_at"]) - if reset > now: + if reset > after: resets.append(reset) if resets: return min(resets) diff --git a/backend/tests/test_usage_schedule.py b/backend/tests/test_usage_schedule.py new file mode 100644 index 00000000..71a993be --- /dev/null +++ b/backend/tests/test_usage_schedule.py @@ -0,0 +1,233 @@ +from datetime import UTC, datetime, timedelta +from typing import Any + +import pytest +from conftest import ( + connect_anthropic_subscription, + connect_provider, + seed_note_run, +) +from druks.core.tasks import refresh_usage +from druks.harnesses.datastructures import ParsedMetric, ParsedUsage +from druks.harnesses.providers import AnthropicProvider, OpenAiProvider +from druks.models import Base +from druks.secrets.models import VaultSecret +from druks.testing import seed_call +from druks.usage.models import UsageScrape + +NOW = datetime(2026, 9, 8, 12, tzinfo=UTC) + + +@pytest.fixture +async def subscription(druks_db) -> VaultSecret: + return await connect_anthropic_subscription("op@example.com") + + +async def _scrape( + subscription: VaultSecret, + at: datetime, + *, + five: int | None = 50, + reset: datetime | None = None, + weeks: list[dict[str, Any]] | None = None, + error: str | None = None, +) -> UsageScrape: + row = UsageScrape( + provider=subscription.audience_name, + account_id=subscription.account_id, + scraped_at=at, + five_hour_percent_left=five, + five_hour_resets_at=reset, + weeks=weeks or [], + parse_ok=not error, + error=error, + ) + await row.save() + return row + + +async def test_first_scrape_ignores_other_subscriptions(subscription) -> None: + """A provider or account with history does not delay a new subscription.""" + other_account = await connect_anthropic_subscription("other@example.com") + other_provider = await connect_provider(OpenAiProvider, {}) + await _scrape(other_account, NOW) + await _scrape(other_provider, NOW) + + assert await UsageScrape.is_due(subscription, now=NOW) + + +@pytest.mark.parametrize("count,minutes", [(1, 5), (2, 10), (3, 20), (4, 40), (5, 60), (6, 60)]) +@pytest.mark.parametrize("error", [None, "auth_required"]) +async def test_unchanged_scrapes_double_interval_to_one_hour( + subscription, count, minutes, error +) -> None: + """Successful scrapes and unchanged errors use the same bounded interval.""" + for index in range(count): + await _scrape( + subscription, + NOW - timedelta(minutes=count - index - 1), + five=None if error else 50, + error=error, + ) + due_at = NOW + timedelta(minutes=minutes) + + assert not await UsageScrape.is_due(subscription, now=due_at - timedelta(seconds=1)) + assert await UsageScrape.is_due(subscription, now=due_at) + + +@pytest.mark.parametrize( + "changes", + [ + {"five": 49}, + {"five": None}, + {"error": "timeout"}, + {"weeks": [{"model": "Fable", "percent_left": 49, "resets_at": None}]}, + {"weeks": []}, + ], +) +async def test_changed_values_restart_the_interval(subscription, changes) -> None: + """A changed quota or error starts a new five-minute interval.""" + values = {"weeks": [{"model": "Fable", "percent_left": 50, "resets_at": None}]} + + for index in range(5): + await _scrape(subscription, NOW - timedelta(minutes=5 - index), **values) + await _scrape(subscription, NOW, **(values | changes)) + + assert not await UsageScrape.is_due(subscription, now=NOW + timedelta(minutes=4)) + assert await UsageScrape.is_due(subscription, now=NOW + timedelta(minutes=5)) + + +@pytest.mark.parametrize("status", ["succeeded", "failed"]) +async def test_finished_call_polls_on_the_next_tick(subscription, druks_db, status) -> None: + """A call billed to this subscription bypasses the idle delay.""" + await _scrape(subscription, NOW) + run = await seed_note_run(druks_db) + call = await seed_call( + druks_db, run, "summarize", subscription_id=subscription.id, status=status + ) + call.finished_at = NOW + timedelta(seconds=1) + await druks_db.flush() + + assert await UsageScrape.is_due(subscription, now=NOW + timedelta(minutes=1)) + + +@pytest.mark.parametrize("source", ["other_account", "other_provider", "api_key", "running", "old"]) +async def test_unrelated_or_unfinished_calls_do_not_bypass_delay( + subscription, druks_db, source +) -> None: + """Only a later completion on this subscription makes its scrape due.""" + await _scrape(subscription, NOW) + charged_subscription = subscription + + if source == "other_account": + charged_subscription = await connect_anthropic_subscription("other@example.com") + elif source == "other_provider": + charged_subscription = await connect_provider(OpenAiProvider, {}) + run = await seed_note_run(druks_db) + call = await seed_call( + druks_db, + run, + "summarize", + subscription_id=None if source == "api_key" else charged_subscription.id, + ) + call.finished_at = NOW + timedelta(seconds=1) + + if source == "running": + call.finished_at = None + elif source == "old": + call.finished_at = NOW + await druks_db.flush() + + assert not await UsageScrape.is_due(subscription, now=NOW + timedelta(minutes=1)) + + +@pytest.mark.parametrize("window", ["five_hour", "weekly"]) +async def test_finished_call_bypasses_exhausted_window(subscription, druks_db, window) -> None: + """A completed call makes an exhausted snapshot due before its reset.""" + reset = NOW + timedelta(hours=2) + values = {"five": 0, "reset": reset} + + if window == "weekly": + values = {"weeks": [{"model": "Fable", "percent_left": 0, "resets_at": reset.isoformat()}]} + await _scrape(subscription, NOW, **values) + run = await seed_note_run(druks_db) + call = await seed_call(druks_db, run, "summarize", subscription_id=subscription.id) + call.finished_at = NOW + timedelta(seconds=1) + await druks_db.flush() + + assert await UsageScrape.is_due(subscription, now=NOW + timedelta(minutes=1)) + + +async def test_soonest_exhausted_reset_controls_polling(subscription) -> None: + """A nonempty window's earlier reset does not end the exhaustion delay.""" + await _scrape( + subscription, + NOW, + five=0, + reset=NOW + timedelta(hours=2), + weeks=[ + { + "model": None, + "percent_left": 50, + "resets_at": (NOW + timedelta(minutes=5)).isoformat(), + }, + { + "model": "Fable", + "percent_left": 0, + "resets_at": (NOW + timedelta(hours=1)).isoformat(), + }, + ], + ) + + assert not await UsageScrape.is_due(subscription, now=NOW + timedelta(minutes=59)) + assert await UsageScrape.is_due(subscription, now=NOW + timedelta(hours=1)) + + +@pytest.mark.parametrize("reset", [None, NOW - timedelta(seconds=1), NOW]) +async def test_zero_without_a_future_reset_uses_normal_interval(subscription, reset) -> None: + """An unknown or expired reset cannot block polling indefinitely.""" + await _scrape(subscription, NOW, five=0, reset=reset) + + assert not await UsageScrape.is_due(subscription, now=NOW + timedelta(minutes=4)) + assert await UsageScrape.is_due(subscription, now=NOW + timedelta(minutes=5)) + + +@pytest.mark.parametrize("error", [None, "timeout"]) +async def test_eight_idle_hours_need_at_most_twelve_polls(subscription, error) -> None: + """Five-minute ticks respect the hourly cap throughout an idle stretch.""" + polls = 0 + + for minute in range(0, 8 * 60 + 1, 5): + now = NOW + timedelta(minutes=minute) + + if await UsageScrape.is_due(subscription, now=now): + await _scrape(subscription, now, five=None if error else 50, error=error) + polls += 1 + + assert 8 <= polls <= 12 + + +async def test_task_polls_only_due_subscriptions(subscription, monkeypatch) -> None: + """The task polls only due subscriptions across providers.""" + idle = await connect_anthropic_subscription("idle@example.com") + exhausted = await connect_anthropic_subscription("exhausted@example.com") + revoked = await connect_anthropic_subscription("revoked@example.com") + revoked.revoked_at = NOW + openai = await connect_provider(OpenAiProvider, {}) + await _scrape(idle, NOW - timedelta(minutes=1)) + await _scrape(exhausted, NOW, five=0, reset=NOW + timedelta(hours=1)) + expected = [subscription.id, openai.id] + fetched = [] + + async def fetch_usage(subscription, *, now=None): + fetched.append(subscription.id) + return ParsedUsage(ok=True, five_hour=ParsedMetric(percent_left=50, resets_at=None)) + + monkeypatch.setattr(Base, "utc_now", lambda: NOW) + monkeypatch.setattr(AnthropicProvider, "fetch_usage", fetch_usage) + monkeypatch.setattr(OpenAiProvider, "fetch_usage", fetch_usage) + + await refresh_usage._function() + + assert fetched == expected + assert (await UsageScrape.latest_for("openai", openai.account_id)).scraped_at == NOW diff --git a/docs/configuration.md b/docs/configuration.md index 5215b278..9fbfaa53 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -367,6 +367,11 @@ model catalog timestamp. The weekly quota stays in the provider row. Anthropic and OpenAI fetch separate model lists. Added providers use the cached Models.dev directory. +Druks polls subscription usage every five minutes, with intervals up to one hour +while values stay unchanged. An exhausted window waits for its reset unless an +agent call finishes on the subscription. Manual refresh keeps a 60-second +minimum between polls. + The `claude` and `codex` CLIs run on their own vendor's subscription or key. `opencode` and `pi` run on an API key only, for Anthropic or OpenAI. A key for a Models.dev provider stores, but an agent on that provider refuses to run: no