DRU-1065 -- A provider or account switch wakes a call waiting out a usage limit

backend/druks/agents.py

@@ -10,6 +10,8 @@
 
 from dbos import DBOS, StepOptions
 from pydantic import BaseModel, ConfigDict, ValidationError
+from sqlalchemy import select
+from sqlalchemy.ext.asyncio import AsyncSession
 
 from druks.apps.exceptions import AppBotError
 from druks.apps.loader import resolve_app
@@ -22,8 +24,9 @@
 from druks.db import db_session
 from druks.durable.activity import set_run_phase
 from druks.durable.engine import _step_engine, step_session
-from druks.durable.exceptions import WorkflowError
-from druks.durable.models import AgentCall, Artifact
+from druks.durable.enums import RunState
+from druks.durable.exceptions import WaitTimeout, WorkflowError
+from druks.durable.models import AgentCall, Artifact, Run
 from druks.exceptions import ObjectNotFound
 from druks.files.datastructures import File
 from druks.files.exceptions import FileUnavailableError
@@ -55,6 +58,8 @@
 _QUOTA_FALLBACK_WAIT_SECONDS = 30 * 60
 _QUOTA_MAX_WAIT_SECONDS = 6 * 60 * 60
 _REAP_BEFORE_WAIT_SECONDS = 120
+# The topic a run waiting out a quota window listens on for a changed provider or account.
+_QUOTA_TOPIC = "agents.quota"
 
 _in_after_agent: ContextVar[bool] = ContextVar("_in_after_agent", default=False)
 
@@ -232,26 +237,40 @@ async def _invoke() -> Any:
                     # re-runs the whole enclosing step from attempt one anyway.
                     await asyncio.sleep(delay * random.uniform(0.9, 1.1))
 
-        async def _quota_retry_wait() -> float:
-            # Runs as its own step: the body does no IO, and replay reuses the
-            # recorded wait instead of re-reading the scrape.
+        async def _quota_retry_wait() -> dict[str, Any]:
+            # Runs as its own step: replay reuses the recorded wait instead of
+            # re-reading the scrape.
             async with step_session() as session:
-                # The scrape belongs to the subscription account.
                 config = await get_config(session, self.id, workflow.account_id)
                 provider_id = config.model.partition("/")[0]
                 scrape = await UsageScrape.latest_for(
                     session, provider_id, config.charged_account_id
                 )
-                if scrape:
-                    now = datetime.now(UTC)
-                    reset = scrape.soonest_reset_after(now)
-                    if reset:
-                        return (reset - now).total_seconds()
-                # No scrape yet, or nothing ahead of now: a blind half hour.
-                return _QUOTA_FALLBACK_WAIT_SECONDS
+            now = datetime.now(UTC)
+            reset = scrape.soonest_reset_after(now) if scrape else None
+            # No scrape yet, or nothing ahead of now: a blind half hour.
+            seconds = (reset - now).total_seconds() if reset else _QUOTA_FALLBACK_WAIT_SECONDS
+            # Jitter, so one window's reset doesn't wake every run at once.
+            wait = seconds + random.uniform(60, 120)
+            return {
+                "quota_holder": config.quota_holder,
+                "seconds": seconds,
+                "wait": wait,
+                "ends_at": (now + timedelta(seconds=wait)).isoformat(),
+            }
+
+        async def _quota_wait_left(quota_holder: str, ends_at: str) -> float:
+            # Another holder retries now; its own limit, if any, fails the call and
+            # earns that holder its own wait.
+            async with step_session() as session:
+                config = await get_config(session, self.id, workflow.account_id)
+            if config.quota_holder != quota_holder:
+                return 0
+            return (datetime.fromisoformat(ends_at) - datetime.now(UTC)).total_seconds()
 
         transient_retry_index = 0
-        quota_retried = False
+        # The quota holders whose window this call has waited out once already.
+        quota_waited: set[str] = set()
         while True:
             try:
                 result = await DBOS.run_step_async(
@@ -268,28 +287,51 @@ async def _quota_retry_wait() -> float:
                     # Jitter, so a provider outage doesn't wake every run at once.
                     wait_seconds = delay * random.uniform(0.9, 1.1)
                 elif error.retry is Retry.QUOTA:
-                    # One shot: sleep until the provider's window turns, read
-                    # off the latest usage scrape.
-                    if quota_retried:
-                        raise
-                    quota_retried = True
-                    wait_seconds = await DBOS.run_step_async(
+                    # Wait until the provider's window turns, read off the latest
+                    # usage scrape, once per quota holder.
+                    quota = await DBOS.run_step_async(
                         StepOptions(name=f"{workflow.kind}.agent.{self.id}.retry_wait"),
                         _quota_retry_wait,
                     )
-                    # A far-off reset (an exhausted weekly window) isn't worth
-                    # sleeping into — fail now.
-                    if wait_seconds > _QUOTA_MAX_WAIT_SECONDS:
+                    # A quota holder waits once, and a far-off reset (an exhausted
+                    # weekly window) isn't worth sleeping into — fail now.
+                    if (
+                        quota["quota_holder"] in quota_waited
+                        or quota["seconds"] > _QUOTA_MAX_WAIT_SECONDS
+                    ):
                         raise
-                    wait_seconds += random.uniform(60, 120)
-                    resumes_at = datetime.now(UTC) + timedelta(seconds=wait_seconds)
+                    quota_waited.add(quota["quota_holder"])
+                    resumes_at = datetime.fromisoformat(quota["ends_at"])
                     await set_run_phase("waiting_for_quota", ends_at=resumes_at)
                     if workflow.conversation_id:
                         await _post_to_chat(
                             workflow.workflow_id,
                             "chat.report_quota_wait",
                             functools.partial(report_quota_wait, resumes_at=resumes_at),
                         )
+                    wait_seconds = quota["wait"]
+                    if wait_seconds >= _REAP_BEFORE_WAIT_SECONDS:
+                        await workflow._reap_run()
+                    # Each wait's length comes from a step, so replay takes the same turns.
+                    while wait_seconds > 0:
+                        try:
+                            await workflow.wait(
+                                _QUOTA_TOPIC, label="waiting_for_quota", ttl_seconds=wait_seconds
+                            )
+                        except WaitTimeout:
+                            break
+                        # A changed provider or account woke the wait.
+                        wait_seconds = await DBOS.run_step_async(
+                            StepOptions(name=f"{workflow.kind}.agent.{self.id}.quota_check"),
+                            functools.partial(
+                                _quota_wait_left, quota["quota_holder"], quota["ends_at"]
+                            ),
+                        )
+                    if workflow.conversation_id:
+                        await _post_to_chat(
+                            workflow.workflow_id, "chat.report_quota_resume", report_quota_resume
+                        )
+                    continue
                 else:
                     # Permanent or unclassified — never retried.
                     raise
@@ -300,16 +342,21 @@ async def _quota_retry_wait() -> float:
                 # Durable: the end time checkpoints now, so a restart resumes
                 # with the remaining wait, and replay skips it entirely.
                 await DBOS.sleep_async(wait_seconds)
-                if error.retry is Retry.QUOTA and workflow.conversation_id:
-                    await _post_to_chat(
-                        workflow.workflow_id, "chat.report_quota_resume", report_quota_resume
-                    )
             else:
                 # Body-level only, on both passes: in-step calls hit the early return
                 # live and are skipped with their step on replay.
                 workflow.journal.add(result)
                 return result
 
+    @classmethod
+    async def wake_quota_waits(cls, session: AsyncSession) -> None:
+        """Wake every run waiting out a quota window, so it reads its agent's provider and
+        account again. Call it after the change commits."""
+        running = await session.scalars(select(Run.id).where(Run.state == RunState.RUNNING.value))
+        for run_id in running:
+            if await DBOS.get_event_async(run_id, "waiting_on", timeout_seconds=0) == _QUOTA_TOPIC:
+                await DBOS.send_async(run_id, {}, topic=_QUOTA_TOPIC)
+
     async def _run(
         self,
         *,

backend/druks/chat/constants.py

@@ -40,8 +40,7 @@
     "what it waits for and where to answer it.]"
 )
 QUOTA_RESUMED_MESSAGE = (
-    "[Internal: Run {run} continues now that the provider's usage limit reset. Tell the "
-    "person in one short line.]"
+    "[Internal: Run {run} continues after its usage-limit wait. Tell the person in one short line.]"
 )
 FAILURE_MESSAGE = (
     "[Internal: Run {run} failed: {failure}. Tell the person in one short line, with no "

backend/druks/chat/service.py

@@ -747,7 +747,7 @@ async def report_quota_wait(session: AsyncSession, run: Run, *, resumes_at: date
 
 
 async def report_quota_resume(session: AsyncSession, run: Run) -> str | None:
-    """Tell the chat that started a run that it continues after the usage limit reset.
+    """Tell the chat that started a run that it continues after its usage-limit wait.
     Returns the conversation to deliver."""
     return await report_outcome(session, run, QUOTA_RESUMED_MESSAGE.format(run=run.id))
 

backend/druks/harnesses/config.py

@@ -70,6 +70,12 @@ def secrets_hash(self) -> str:
     def charged_account_id(self) -> str | None:
         return self.subscription.account_id if self.subscription else None
 
+    @property
+    def quota_holder(self) -> str:
+        """Whom a usage limit binds: the model's provider and the account signed in to it."""
+        email = self.subscription.identity.get("email") if self.subscription else None
+        return f"{self.model}|{self.billing}|{email}"
+
 
 async def check_config(
     session: AsyncSession, harness_name: str, model: str, billing: str

backend/druks/harnesses/routes.py

@@ -7,6 +7,7 @@
 from druks.accounts.dependencies import current_session_account, current_session_or_setup
 from druks.accounts.models import Account
 from druks.accounts.schemas import AccountResponse
+from druks.agents import Agent
 from druks.api.dependencies import SessionDep
 from druks.secrets.datastructures import Audience
 from druks.secrets.enums import SecretKind
@@ -169,7 +170,7 @@ async def sign_in_with_token(
         credentials, expires_at = provider.parse_token(token)
     except ConnectError as error:
         raise HTTPException(status_code=422, detail=str(error)) from error
-    return await VaultSecret.store(
+    subscription = await VaultSecret.store(
         session,
         SecretKind.SUBSCRIPTION,
         Audience.provider(provider.id),
@@ -179,6 +180,10 @@ async def sign_in_with_token(
         identity={"email": None},
         expires_at=expires_at,
     )
+    # A woken run reads the subscription, so it commits first.
+    await session.commit()
+    await Agent.wake_quota_waits(session)
+    return subscription
 
 
 async def _store_connection(
@@ -214,6 +219,7 @@ async def _store_connection(
     # across an await, and after the commit the reply would need another read.
     response = AccountResponse.model_validate(resolved)
     await session.commit()
+    await Agent.wake_quota_waits(session)
     try:
         # The flow is already spent, so a failed refresh only logs.
         await provider.refresh_catalog(session)

backend/druks/user_settings/routes.py

@@ -3,6 +3,7 @@
 
 from druks.accounts.dependencies import current_account, current_session_account
 from druks.accounts.models import Account
+from druks.agents import Agent
 from druks.api.dependencies import SessionDep
 from druks.apps.loader import get_app, iter_apps
 from druks.apps.registry import agents, bots, channels, workflows
@@ -108,6 +109,9 @@ async def update_settings(session: SessionDep, body: UpdateSettingsRequest) -> I
         await settings.update(**fields)
         if any(field in fields for field in _EXECUTION_DEFAULTS):
             await check_agent_configs(session, settings)
+            # A woken run reads the settings, so they commit first.
+            await session.commit()
+            await Agent.wake_quota_waits(session)
     return settings
 
 
@@ -217,4 +221,9 @@ async def update_app_settings(
         # the just-written overrides off this request's session.
         await apply_schedules(session)
 
+    if body.agent_harnesses or body.agent_models or body.agent_billings:
+        # A woken run reads the overrides, so they commit first.
+        await session.commit()
+        await Agent.wake_quota_waits(session)
+
     return await get_app_settings(session)