fix(quotas): admit append sweeps only against the locked plan (#1092)

Co-authored-by: Cursor Agent <cursoragent@cursor.com>

Author: RishiDesai

oddish/src/oddish/core/endpoints/sweep.py

@@ -417,7 +417,9 @@ async def create_task_sweep_core(
     of the *raw* client submission so an honest retry is not spuriously rejected;
     when omitted it is computed from ``submission`` as received here.
     """
+    from oddish.config import QuotaMode
     from oddish.core.quota_admission import admit_trials
+    from oddish.core.quotas import acquire_quota_locks
     from oddish.core.sweeps import (
         build_task_submission_from_sweep,
         build_trial_specs_from_sweep,
@@ -525,17 +527,19 @@ async def create_task_sweep_core(
     if submission.append_to_task:
         # Lock order must stay ``quota advisory → task row`` to match
         # ``cancel_trials_if_quota_reached`` (quota first, then task FOR UPDATE).
-        # Taking the task row before ``admit_trials`` inverted that order and
+        # Taking the task row before the quota lock inverted that order and
         # deadlocked under concurrent over-quota cancellation.
         #
-        # Also do not take the org quota lock across reconcile: that window can
-        # take tens of seconds on large tasks and starves every concurrent
-        # submit / enforcement worker (opaque 500s on /tasks/sweep).
+        # Admit only against the locked plan: an unlocked estimate can still
+        # ``QuotaExceeded`` (402) after a concurrent append already filled the
+        # deficit, even when this request would insert fewer trials or none.
+        # Hold the quota advisory only across the short locked plan + admit +
+        # insert — not across the earlier experiment/setup work.
         task = await get_task_for_org_core(
             session, task_id=submission.task_id, org_id=org_id
         )
         # Read-only intent from the unlocked snapshot. Applied under FOR UPDATE
-        # after admission (idempotent flips).
+        # after the quota lock (idempotent flips).
         want_run_analysis = bool(task.run_analysis or submission.run_analysis)
         want_run_probe = bool(task.run_probe or submission.run_probe)
 
@@ -582,20 +586,10 @@ async def create_task_sweep_core(
         target_experiment_id = new_experiment_id or (
             primary_experiment.id if primary_experiment else None
         )
-        # Unlocked estimate for quota admission only. Authoritative plan is
-        # rebuilt under the task row lock below so concurrent appends cannot
-        # both observe the same deficit and overshoot declarative N.
-        planned_trials, _ = await _plan_append_trials(
-            session,
-            task=task,
-            submission=submission,
-            target_experiment_id=target_experiment_id,
-            default_environment=effective_default_env,
-            allowed_environments=allowed_environments,
-        )
-        # Quota advisory lock is acquired here (and held through commit).
-        await admit_trials(session, org_id, billed_user_id, count=len(planned_trials))
-        # Task row lock only after quota — same order as enforcement.
+        # Quota advisory before task FOR UPDATE (ENFORCE only; SHADOW/OFF do
+        # not take these locks on admit or cancel either).
+        if settings.quota_mode == QuotaMode.ENFORCE and org_id is not None:
+            await acquire_quota_locks(session, org_id, billed_user_id)
         await session.refresh(task, with_for_update=True)
         # Allow flipping task.run_analysis from False to True on append.
         # ``run_analysis`` runs at trial-completion time, so updating the
@@ -625,17 +619,8 @@ async def create_task_sweep_core(
             default_environment=effective_default_env,
             allowed_environments=allowed_environments,
         )
-        if len(trials) > len(planned_trials):
-            # Rare: deficit grew while we waited (e.g. concurrent failures).
-            # Re-check the *full* final count — admit_trials does not accumulate
-            # the earlier estimate (it only adds ``count`` to current inflight),
-            # so a delta-only top-up would undercount headroom.
-            await admit_trials(
-                session,
-                org_id,
-                billed_user_id,
-                count=len(trials),
-            )
+        # Authoritative count only (no-op when the locked plan is empty).
+        await admit_trials(session, org_id, billed_user_id, count=len(trials))
 
         append_submission = submission.model_copy(
             update={

oddish/tests/test_quota_lock_try_acquire.py

@@ -99,45 +99,36 @@ async def after_check():
 
 
 def test_append_sweep_preserves_quota_then_task_lock_order():
-    """Regression: task FOR UPDATE before admit inverted lock order vs cancel.
+    """Regression: lock order and admit-against-locked-plan invariants.
 
     Enforcement takes quota advisory first, then task FOR UPDATE. Append must
-    do the same (admit_trials → task FOR UPDATE) or concurrent cancel deadlocks.
-    Also: no early acquire_quota_locks across reconcile (starves /tasks/sweep).
-    Authoritative reconcile must re-run after the task lock so concurrent
-    appends cannot both insert against the same unlocked deficit.
+    do the same or concurrent cancel deadlocks. Admit must use the locked plan
+    only — an unlocked estimate can 402 after a concurrent append already
+    filled the deficit.
     """
     from oddish.core.endpoints import sweep as sweep_mod
 
     source = Path(sweep_mod.__file__).read_text(encoding="utf-8")
-    assert "await acquire_quota_locks" not in source
-    assert "from oddish.core.quotas import acquire_quota_locks" not in source
+    assert "from oddish.core.quotas import acquire_quota_locks" in source
     assert "async def _plan_append_trials(" in source
 
     append_start = source.index("if submission.append_to_task:")
     # Limit to the append branch (create mode follows in the same function).
     append_end = source.index("\n    # Create mode", append_start)
     append_body = source[append_start:append_end]
-    admit_at = append_body.index("await admit_trials(")
+
+    acquire_at = append_body.index("await acquire_quota_locks(")
     for_update_at = append_body.index("with_for_update=True")
-    assert for_update_at > admit_at, (
-        "append path must FOR UPDATE the task only after admit_trials "
-        f"(admit@{admit_at}, for_update@{for_update_at})"
-    )
-    plan_calls = [
-        i
-        for i in range(len(append_body))
-        if append_body.startswith("await _plan_append_trials(", i)
-    ]
-    assert len(plan_calls) >= 2, (
-        "append must plan once for admit sizing and again under the task lock"
+    plan_at = append_body.index("await _plan_append_trials(")
+    admit_at = append_body.index("await admit_trials(")
+
+    assert acquire_at < for_update_at < plan_at < admit_at, (
+        "expected acquire_quota_locks → FOR UPDATE → plan → admit_trials; "
+        f"got acquire@{acquire_at} for_update@{for_update_at} "
+        f"plan@{plan_at} admit@{admit_at}"
     )
-    assert plan_calls[0] < admit_at < for_update_at < plan_calls[1], (
-        "expected unlocked plan → admit → FOR UPDATE → locked re-plan; "
-        f"got plan@{plan_calls} admit@{admit_at} for_update@{for_update_at}"
+    assert append_body.count("await _plan_append_trials(") == 1, (
+        "append must plan once under the task lock, not against an unlocked estimate"
     )
-    # Top-up after a larger locked plan must re-admit the full final count
-    # (not a delta): admit_trials adds ``count`` to live inflight only.
-    top_up = append_body[plan_calls[1] :]
-    assert "count=len(trials)" in top_up
-    assert "count=len(trials) - len(planned_trials)" not in top_up
+    assert "planned_trials" not in append_body
+    assert "count=len(trials)" in append_body[admit_at:]