fix(graph): refuse change-tier without target capacity; make the drain fail loudly

[jfrench9]

Summary

change-tier is the designated upsell route from Standard to Large/XLarge, and those tiers are provisioned on request (their writer ASGs sit at desired 0 by policy). The operation did not check that first: it moved the customer's Stripe price, set the subscription upgrading, and enqueued a migration whose worker raised the target ASG's desired capacity itself and then waited at most 300 s for a cold instance to boot, register, and claim the volume — a window a cold boot does not reliably fit, ending in a rollback that has only ever run in a test. Separately, the migration's drain step polled admin routes the graph API does not serve, timed out after 120 s, and detached the volume anyway.

This PR makes change-tier apply the same refuse-the-sale rule checkout already applies — a healthy writer on the target tier with a free slot, or a 409 before anything moves — and makes the drain refuse rather than time out into a detach.

Changes

• operations/graph/capacity.py (new) — tier_capacity_status(tier): ready when the tier has a healthy writer with a free slot; scalable and any lookup failure read as at_capacity (fail closed). Lifted verbatim from the checkout router so both sale paths share one definition.
• routers/billing/checkout.py — imports the shared helper under its existing private name; behaviour unchanged.
• operations/graph/commands/tier.py — after tier validation and before any write: tier_capacity_status(new_tier) != "ready" → 409 "No capacity is currently available on the '' tier. Request access from the tier picker or contact support." Nothing is committed, Stripe is not called, no task is enqueued.
• operations/graph/tasks/graph_tier_upgrade.py — _drain_instance now has three outcomes instead of one:
  • graph API unreachable → treated as drained (the maintenance-window procedure stops the container before a tier change, and every write reaches the volume through that API);
  • graph API up but 404 on /admin/drain, or any error status → DrainRefusedError before the detach;
  • drain timeout with connections still open → DrainRefusedError before the detach.
    The existing except Exception in run() catches it, rolls back the registry state (nothing has been detached), and re-raises. Module docstring updated to say so. Reviewers: the "unreachable counts as drained" branch is the judgment call — it is what lets the documented stop-the-container procedure work without a real drain endpoint.
• Tests — tests/operations/graph/test_capacity.py (new, 5); test_tier_command.py gains an autouse "capacity ready" fixture and a 409 test asserting no commit / no Stripe / no enqueue; test_graph_tier_upgrade.py gains three drain tests (unreachable, 404, timeout) and gives the existing happy-path stub a real status code.

Breaking Changes

None. The only observable change is a new 409 from POST /v1/graphs/{g}/operations/change-tier when the target tier has no capacity; the operation envelope and request model are unchanged, so nothing in either SDK tier moves.

Testing

• just test — 14,185 passed, 42 skipped.
• just test-code — clean (ruff, format, basedpyright, cfn-lint); also re-run by the pre-commit hook on both commits.
• Not exercised against a live fleet; the change-tier walk on a throwaway graph is the next step and is documented separately.

Certification

• I have the right to submit this work under the Apache 2.0 license, and do so. Where any part of it is owned by my employer, I have their permission.

[jfrench9]

@claude please review this PR

[claude]

Claude finished @jfrench9's task in 2m 39s —— View job

---

Reviewing PR #1305

• Gather PR metadata, diff, and existing review comments
• Review changed files against CLAUDE.md conventions
• Check test coverage
• Post review feedback

PR Summary

Title: fix(graph): refuse change-tier without target capacity; make the drain fail loudly
Author: jfrench9 | Branch: bugfix/change-tier-capacity-gate → main
Changes: +237 / -28 across 7 files

Adds a shared tier_capacity_status() helper (lifted verbatim from checkout's existing private helper) and gates change-tier on it before any DB/Stripe writes, returning 409 on at_capacity/scalable/lookup-failure. Also hardens _drain_instance in the migration worker so a graph API that's up-but-unresponsive-to-drain (404, error status, or timeout with open connections) raises DrainRefusedError before the volume is ever detached, instead of logging a warning and detaching anyway.

Existing Review Feedback

No prior human or bot review comments/reviews on this PR — this is the first pass.

My Review

I traced both changed control-flow paths end-to-end against the tests and the callers.

operations/graph/commands/tier.py — the capacity gate sits correctly after tier/subscription validation and before db.commit(), Stripe, and enqueue_task. Confirmed via test_target_tier_without_capacity_raises_409: subscription/graph state untouched, db.commit/payment provider/worker enqueue all unasserted-called. new_tier is already validated against BillingConfig.get_subscription_plan() before reaching the gate, and its values (ladybug-standard/-large/-xlarge) line up with GraphTier enum values in config/graph_tier.py, so the lazy GraphTier(tier) construction inside tier_capacity_status won't spuriously fail-closed on a legitimate tier.

operations/graph/capacity.py — straight lift from checkout's prior private helper; checkout.py's behavior is unchanged (confirmed the diff only moves the function and re-imports it under its old private name; env import is still used elsewhere in that file for BILLING_ENABLED, so nothing goes dead). check_tier_capacity in allocation_manager.py is an existing, separately-tested method — this PR doesn't touch its logic, just calls it from a second place.

operations/graph/tasks/graph_tier_upgrade.py — DrainRefusedError is a plain RuntimeError subclass, so it's caught by the existing broad except Exception in execute(), which runs _rollback(). Since _drain_instance now raises before volume_detached is set True, _rollback correctly skips the reattach branch (if volume_detached and ...) — matches the commit message's "nothing has been detached" claim.

One gap worth a look: _drain_instance still has its original early-return for a missing private_ip (line 332-334 — if not private_ip: logger.warning(...); return), which silently counts as "drained" without ever calling the graph API. That's the same shape of bug this PR is fixing elsewhere (skip real confirmation → proceed to detach), just triggered by an empty registry field instead of a timeout/404. Given the module docstring's new claim that "the volume is never detached under possible writes," it seems like this branch should also raise DrainRefusedError rather than return silently — worth confirming whether an empty private_ip is truly unreachable-equivalent (e.g., always paired with a stopped container) or whether it's a distinct condition that deserves the same refusal treatment.

The "unreachable graph API counts as drained" branch is the one the PR description explicitly flags as a judgment call, and I think it's defensible given the stated invariant (maintenance-window procedure stops the container before any tier change), but it is worth being explicit that this treats "API unreachable" and "container stopped" as equivalent — a transient network/security-group issue during a live migration would also read as "drained" under this logic. If the maintenance-window precondition is enforced elsewhere (a state check, a runbook step), that's fine as documented; if it's only convention, this is a soft spot.

Tests — read through all three new/changed test files rather than trusting green. test_capacity.py's 5 cases correctly cover ready/scalable/at_capacity/lookup-exception/invalid-tier-string. test_tier_command.py's new autouse fixture patches tier_capacity_status at its lazy-import source (robosystems.operations.graph.capacity.tier_capacity_status), which is correct given tier.py imports it inside the function body — confirmed the patch target actually matches the import site. test_graph_tier_upgrade.py's three new drain tests each assert client.get.assert_not_called()/assert_called() appropriately to distinguish "refused before polling" from "refused after polling," and the pre-existing happy-path stub now returns a real status_code=202 instead of an unconfigured AsyncMock(), which is a good catch — the old stub would've silently returned a MagicMock for response.status_code that happens to not equal 404 or >=400 by luck of mock comparison, not by design.

Issues (should fix before merge)

None blocking.

Suggestions (non-blocking)

1. graph_tier_upgrade.py:332-334 — consider whether the missing-private_ip branch should also raise DrainRefusedError instead of silently returning, for the reasons above. If there's a reason it's safe (e.g., this state genuinely can't occur without the container already being stopped), a one-line comment would help the next reader who re-derives the same question.
2. Minor: the response.status_code == 404 check and the following >= 400 check overlap (404 satisfies both), which is intentional for a more specific message — fine as-is, just flagging in case it wasn't deliberate.

Questions

• Is the "maintenance-window stops the container first" precondition enforced anywhere in code/ops tooling, or is it purely procedural? That determines how much weight the unreachable-counts-as-drained branch is carrying.

Verdict

APPROVE — the core fix (refuse-the-sale gate + fail-loud drain) is correct, well-tested, and matches the PR description's claims. The one suggestion above (missing-private_ip branch) is worth a follow-up but doesn't block this PR, since it's a narrowing of a pre-existing gap rather than a regression introduced here.

[jfrench9]

Follow-up in ace5b80: (1) no private IP now raises DrainRefusedError instead of returning — nothing could be asked, so nothing is detached; (2) "unreachable" is confirmed across three attempts (DRAIN_UNREACHABLE_CONFIRMATIONS) before it is read as a stopped container, with a test that a first-attempt blip followed by a live API still refuses. On the question: the maintenance-window precondition is procedural (the runbook), but the code enforces its effect — graph_api mounts no drain endpoint, so an API that is up always hits the 404 branch and refuses; a stopped container is the only path through, and now a brief blip cannot impersonate one. The 404-then->= 400 overlap is deliberate for the more specific message.