From 5c4abdfd5302e90f2b34635718ab221e6fae580b Mon Sep 17 00:00:00 2001 From: "Joseph T. French" Date: Fri, 28 Aug 2026 15:33:25 -0500 Subject: [PATCH 1/3] fix(graph): refuse change-tier without target capacity; make the drain fail loudly MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit change-tier could move a customer's Stripe price and then start a migration whose worker raised the target ASG's desired capacity itself and waited five minutes for a cold instance to claim the volume. It now applies the same refuse-the-sale rule as checkout — a healthy writer on the target tier with a free slot, or 409 before anything moves — through one shared capacity helper in the ops layer. The drain step no longer times out into a detach. An unreachable graph API counts as drained (the maintenance-window procedure stops the container first, and every write goes through that API); a graph API without a drain endpoint, an error response, or a timeout with connections still open now raise DrainRefusedError before the volume is touched. --- robosystems/operations/graph/capacity.py | 35 ++++++++++ robosystems/operations/graph/commands/tier.py | 15 +++++ .../graph/tasks/graph_tier_upgrade.py | 45 ++++++++++--- robosystems/routers/billing/checkout.py | 22 ++----- .../graph/tasks/test_graph_tier_upgrade.py | 64 ++++++++++++++++++- tests/operations/graph/test_capacity.py | 36 +++++++++++ tests/operations/graph/test_tier_command.py | 48 ++++++++++++++ 7 files changed, 239 insertions(+), 26 deletions(-) create mode 100644 robosystems/operations/graph/capacity.py create mode 100644 tests/operations/graph/test_capacity.py diff --git a/robosystems/operations/graph/capacity.py b/robosystems/operations/graph/capacity.py new file mode 100644 index 000000000..75e09020d --- /dev/null +++ b/robosystems/operations/graph/capacity.py @@ -0,0 +1,35 @@ +"""Tier capacity as the sale paths see it. + +Both places the platform commits money against a writer slot — checkout and +``change-tier`` — ask the same question first: is there a healthy writer on +the target tier with a free slot *right now*? Nothing on either path raises +desired capacity on purpose (the high tiers are provisioned on request), so +``scalable`` is not good enough, and any failure to determine capacity reads +as none: refuse the sale rather than collect against a slot that may not +exist. +""" + +from __future__ import annotations + +from robosystems.config import env +from robosystems.logger import logger + + +async def tier_capacity_status(tier: str) -> str: + """``ready`` when a writer for ``tier`` has a free slot; otherwise + ``at_capacity``. ``scalable`` (no slot, ASG below max) counts as + ``at_capacity`` because nothing on the sale paths raises desired capacity. + Any failure to determine capacity reads as ``at_capacity``. + """ + try: + from robosystems.middleware.graph.allocation_manager import ( + LadybugAllocationManager, + ) + from robosystems.middleware.graph.types import GraphTier + + manager = LadybugAllocationManager(environment=env.ENVIRONMENT) + status_value = await manager.check_tier_capacity(GraphTier(tier)) + except Exception as e: + logger.warning(f"Could not determine capacity for tier {tier}: {e}") + return "at_capacity" + return "ready" if status_value == "ready" else "at_capacity" diff --git a/robosystems/operations/graph/commands/tier.py b/robosystems/operations/graph/commands/tier.py index 15718d3cb..1b5519cd7 100644 --- a/robosystems/operations/graph/commands/tier.py +++ b/robosystems/operations/graph/commands/tier.py @@ -37,6 +37,7 @@ async def change_graph_tier_cmd( from robosystems.models.core import OrgRole, OrgUser from robosystems.models.core.graph import Graph from robosystems.models.core.graph.graph_credits import GraphCredits + from robosystems.operations.graph.capacity import tier_capacity_status from robosystems.operations.graph.tier_validation import ( validate_storage_capacity, validate_subgraph_count, @@ -93,6 +94,20 @@ async def change_graph_tier_cmd( if new_tier == old_tier: raise HTTPException(status_code=400, detail="Already on this tier") + # Refuse-the-sale rule, the same one checkout applies: the migration needs + # a healthy writer on the target tier with a free slot *now*. The worker + # can raise the ASG's desired capacity itself, but a cold boot does not fit + # its reattach window, and the high tiers are provisioned on request by + # policy — so refuse here, before Stripe has moved the customer's price. + if await tier_capacity_status(new_tier) != "ready": + raise HTTPException( + status_code=409, + detail=( + f"No capacity is currently available on the '{new_tier}' tier. " + "Request access from the tier picker or contact support." + ), + ) + new_price_cents = plan_config["base_price_cents"] old_price_cents = subscription.base_price_cents is_upgrade = new_price_cents > old_price_cents diff --git a/robosystems/operations/graph/tasks/graph_tier_upgrade.py b/robosystems/operations/graph/tasks/graph_tier_upgrade.py index 9163351a7..b999e6887 100644 --- a/robosystems/operations/graph/tasks/graph_tier_upgrade.py +++ b/robosystems/operations/graph/tasks/graph_tier_upgrade.py @@ -5,7 +5,8 @@ 1. mark the graph ``migrating`` and retag the volume with the new tier 2. snapshot the volume via the Volume Manager Lambda (the rollback point) -3. drain connections on the old instance +3. drain connections on the old instance — or refuse, if it cannot be + confirmed drained; the volume is never detached under possible writes 4. detach the volume — it then sits ``available`` carrying the new tier tag, which is what makes the target tier's ASG claim it 5. ensure the target ASG has capacity @@ -41,6 +42,11 @@ GRAPH_API_PORT = 8001 +class DrainRefusedError(RuntimeError): + """The old instance could not be confirmed drained, so the volume must + not be detached. Nothing has moved yet when this is raised.""" + + def _get_dynamodb_resource(): """DynamoDB resource, pointed at LocalStack in development.""" region = env.AWS_REGION @@ -312,10 +318,16 @@ async def execute(self) -> dict[str, Any]: raise async def _drain_instance(self, private_ip: str) -> None: - """Ask the old instance to close its connections before the detach. - - Bounded by ``DRAIN_TIMEOUT_SECONDS``; the migration proceeds either way, so - a slow drain delays rather than blocks the move. + """Confirm the old instance has stopped serving before the detach. + + Two outcomes count as drained: the instance reports zero active + connections, or its graph API is not listening at all — the + maintenance-window procedure stops the container before a tier change + precisely so nothing can be writing when the volume detaches, and every + write reaches the volume through that API. A graph API that is up but + has no drain endpoint is neither, and so is a drain that times out with + connections still open: raise, because detaching under live writes is the + one thing this step exists to prevent. """ if not private_ip: logger.warning("No private IP for drain, skipping") @@ -325,11 +337,25 @@ async def _drain_instance(self, private_ip: str) -> None: async with httpx.AsyncClient(timeout=10) as client: try: - await client.post(f"{base_url}/admin/drain") + response = await client.post(f"{base_url}/admin/drain") except httpx.HTTPError as e: - logger.warning(f"Drain request failed (may be expected): {e}") + logger.warning( + f"Graph API on {private_ip} is not reachable ({e}); treating the " + "instance as drained — nothing can be writing through it" + ) return + if response.status_code == 404: + raise DrainRefusedError( + f"Graph API on {private_ip} has no drain endpoint; stop the graph " + "container (maintenance window) and retry the tier change" + ) + if response.status_code >= 400: + raise DrainRefusedError( + f"Drain request to {private_ip} failed with HTTP " + f"{response.status_code}; not detaching" + ) + # Poll for connections to reach 0 elapsed = 0 while elapsed < DRAIN_TIMEOUT_SECONDS: @@ -346,7 +372,10 @@ async def _drain_instance(self, private_ip: str) -> None: await asyncio.sleep(DRAIN_POLL_INTERVAL_SECONDS) elapsed += DRAIN_POLL_INTERVAL_SECONDS - logger.warning("Drain timeout reached, proceeding with detach") + raise DrainRefusedError( + f"Drain of {private_ip} timed out after {DRAIN_TIMEOUT_SECONDS}s with " + "connections still open; not detaching" + ) async def _ensure_asg_capacity(self, asg_client: Any, tier: str) -> None: """Ensure the target tier ASG has capacity for a new instance.""" diff --git a/robosystems/routers/billing/checkout.py b/robosystems/routers/billing/checkout.py index 4adbfe8db..558c94d1c 100644 --- a/robosystems/routers/billing/checkout.py +++ b/robosystems/routers/billing/checkout.py @@ -23,23 +23,11 @@ logger = get_logger(__name__) -async def _tier_capacity_status(plan_name: str) -> str: - """``ready`` when a writer for the tier has a free slot; otherwise - ``at_capacity``. ``scalable`` (no slot, ASG below max) counts as - ``at_capacity`` because nothing on the create path raises desired capacity. - Any failure to determine capacity reads as ``at_capacity``. - """ - try: - from ...middleware.graph.allocation_manager import LadybugAllocationManager - from ...middleware.graph.types import GraphTier - - manager = LadybugAllocationManager(environment=env.ENVIRONMENT) - status_value = await manager.check_tier_capacity(GraphTier(plan_name)) - except Exception as e: - logger.warning(f"Could not determine capacity for tier {plan_name}: {e}") - return "at_capacity" - return "ready" if status_value == "ready" else "at_capacity" - +# The same refuse-the-sale rule guards change-tier; the helper lives in the +# ops layer so both paths read one definition of "capacity". +from ...operations.graph.capacity import ( # noqa: E402 + tier_capacity_status as _tier_capacity_status, +) router = APIRouter(prefix="/billing", tags=["Billing"]) diff --git a/tests/operations/graph/tasks/test_graph_tier_upgrade.py b/tests/operations/graph/tasks/test_graph_tier_upgrade.py index e7e2c1d0f..da92d67ba 100644 --- a/tests/operations/graph/tasks/test_graph_tier_upgrade.py +++ b/tests/operations/graph/tasks/test_graph_tier_upgrade.py @@ -9,6 +9,7 @@ import json from unittest.mock import AsyncMock, MagicMock, patch +import httpx import pytest TASK_MODULE = "robosystems.operations.graph.tasks.graph_tier_upgrade" @@ -601,7 +602,7 @@ async def test_drain_succeeds_when_connections_reach_zero(self): with patch("httpx.AsyncClient") as MockClient: client = AsyncMock() - client.post = AsyncMock() + client.post = AsyncMock(return_value=MagicMock(status_code=202)) client.get = AsyncMock(return_value=mock_response) MockClient.return_value.__aenter__ = AsyncMock(return_value=client) MockClient.return_value.__aexit__ = AsyncMock(return_value=False) @@ -611,6 +612,67 @@ async def test_drain_succeeds_when_connections_reach_zero(self): client.post.assert_called_once() client.get.assert_called() + async def test_drain_treats_unreachable_api_as_drained(self): + # The maintenance-window procedure stops the graph container first; + # nothing can write through an API that is not listening. + task = _make_task() + + with patch("httpx.AsyncClient") as MockClient: + client = AsyncMock() + client.post = AsyncMock(side_effect=httpx.ConnectError("refused")) + client.get = AsyncMock() + MockClient.return_value.__aenter__ = AsyncMock(return_value=client) + MockClient.return_value.__aexit__ = AsyncMock(return_value=False) + + await task._drain_instance("10.0.1.5") + + client.get.assert_not_called() + + async def test_drain_refuses_when_api_has_no_drain_endpoint(self): + from robosystems.operations.graph.tasks.graph_tier_upgrade import ( + DrainRefusedError, + ) + + task = _make_task() + post_response = MagicMock() + post_response.status_code = 404 + + with patch("httpx.AsyncClient") as MockClient: + client = AsyncMock() + client.post = AsyncMock(return_value=post_response) + client.get = AsyncMock() + MockClient.return_value.__aenter__ = AsyncMock(return_value=client) + MockClient.return_value.__aexit__ = AsyncMock(return_value=False) + + with pytest.raises(DrainRefusedError, match="no drain endpoint"): + await task._drain_instance("10.0.1.5") + + client.get.assert_not_called() + + async def test_drain_refuses_on_timeout_with_connections_open(self): + from robosystems.operations.graph.tasks.graph_tier_upgrade import ( + DrainRefusedError, + ) + + task = _make_task() + post_response = MagicMock() + post_response.status_code = 202 + poll_response = MagicMock() + poll_response.status_code = 200 + poll_response.json.return_value = {"active_connections": 3} + + with patch("httpx.AsyncClient") as MockClient: + client = AsyncMock() + client.post = AsyncMock(return_value=post_response) + client.get = AsyncMock(return_value=poll_response) + MockClient.return_value.__aenter__ = AsyncMock(return_value=client) + MockClient.return_value.__aexit__ = AsyncMock(return_value=False) + + with pytest.raises(DrainRefusedError, match="still open"): + await task._drain_instance("10.0.1.5") + + client.get.assert_called() + @pytest.mark.unit @pytest.mark.asyncio diff --git a/tests/operations/graph/test_capacity.py b/tests/operations/graph/test_capacity.py new file mode 100644 index 000000000..ebfdcbf16 --- /dev/null +++ b/tests/operations/graph/test_capacity.py @@ -0,0 +1,36 @@ +"""Tests for the shared refuse-the-sale capacity rule.""" + +from unittest.mock import AsyncMock, patch + +import pytest + +_MANAGER = "robosystems.middleware.graph.allocation_manager.LadybugAllocationManager" + + +@pytest.mark.unit +@pytest.mark.asyncio +class TestTierCapacityStatus: + async def _status(self, tier="ladybug-large", *, value=None, raises=None): + from robosystems.operations.graph.capacity import tier_capacity_status + + with patch(_MANAGER) as manager_cls: + manager_cls.return_value.check_tier_capacity = AsyncMock( + side_effect=raises, return_value=value + ) + return await tier_capacity_status(tier) + + async def test_ready_passes_through(self): + assert await self._status(value="ready") == "ready" + + async def test_scalable_is_not_ready(self): + # Nothing on a sale path raises desired capacity; headroom is not a slot. + assert await self._status(value="scalable") == "at_capacity" + + async def test_at_capacity_passes_through(self): + assert await self._status(value="at_capacity") == "at_capacity" + + async def test_lookup_failure_reads_as_no_capacity(self): + assert await self._status(raises=RuntimeError("dynamodb down")) == "at_capacity" + + async def test_unknown_tier_reads_as_no_capacity(self): + assert await self._status(tier="not-a-tier", value="ready") == "at_capacity" diff --git a/tests/operations/graph/test_tier_command.py b/tests/operations/graph/test_tier_command.py index d5ffd4737..4525e8755 100644 --- a/tests/operations/graph/test_tier_command.py +++ b/tests/operations/graph/test_tier_command.py @@ -71,6 +71,14 @@ def _make_graph_credits(allocation=8000): GRAPH_ID = "kg1a2b3c4d5e6f78" +_CAPACITY = "robosystems.operations.graph.capacity.tier_capacity_status" + + +@pytest.fixture(autouse=True) +def target_tier_has_capacity(): + """The refuse-the-sale gate passes unless a test says otherwise.""" + with patch(_CAPACITY, new_callable=AsyncMock, return_value="ready") as gate: + yield gate @pytest.mark.asyncio @@ -230,6 +238,46 @@ async def test_same_tier_raises_400(self): assert exc.value.status_code == 400 assert "Already on this tier" in exc.value.detail + async def test_target_tier_without_capacity_raises_409( + self, target_tier_has_capacity + ): + """The refuse-the-sale rule fires before billing or Stripe move.""" + target_tier_has_capacity.return_value = "at_capacity" + membership = _make_membership() + subscription = _make_subscription() + graph = _make_graph() + db = MagicMock() + with ( + patch(f"{_CORE}.OrgUser") as mock_org_user, + patch(f"{_CORE}.OrgRole") as mock_org_role, + patch(_GRAPH) as mock_graph_cls, + patch(f"{_TIER}.BillingSubscription") as mock_sub_cls, + patch(f"{_TIER}.BillingConfig") as mock_billing_cfg, + patch(_PAYMENT) as mock_provider_factory, + patch(_WORKER, new_callable=AsyncMock) as mock_enqueue, + ): + mock_org_user.get_user_orgs.return_value = [membership] + mock_org_role.OWNER = "OWNER" + membership.role = mock_org_role.OWNER + mock_graph_cls.get_by_id.return_value = graph + mock_sub_cls.get_by_resource.return_value = subscription + mock_billing_cfg.get_subscription_plan.return_value = { + "name": "ladybug-large", + "base_price_cents": 24900, + } + with pytest.raises(HTTPException) as exc: + await self._call(new_tier="ladybug-large", db=db) + + assert exc.value.status_code == 409 + assert "ladybug-large" in exc.value.detail + target_tier_has_capacity.assert_awaited_once_with("ladybug-large") + # Nothing moved: no commit, no Stripe, no worker task. + assert subscription.status == "active" + assert graph.graph_tier == "ladybug-standard" + db.commit.assert_not_called() + mock_provider_factory.assert_not_called() + mock_enqueue.assert_not_called() + async def test_successful_upgrade_returns_operation_id(self): """Happy-path upgrade returns the operation_id string.""" membership = _make_membership() From 23751cc1cfb10d9f4ed57817a0e4e9f3eded1075 Mon Sep 17 00:00:00 2001 From: "Joseph T. French" Date: Fri, 28 Aug 2026 15:34:22 -0500 Subject: [PATCH 2/3] refactor(billing): import the shared capacity helper with the other imports --- robosystems/routers/billing/checkout.py | 10 +++------- 1 file changed, 3 insertions(+), 7 deletions(-) diff --git a/robosystems/routers/billing/checkout.py b/robosystems/routers/billing/checkout.py index 558c94d1c..cfdbba8d8 100644 --- a/robosystems/routers/billing/checkout.py +++ b/robosystems/routers/billing/checkout.py @@ -18,17 +18,13 @@ from ...models.api.common import AUTHENTICATED_ERROR_RESPONSES, RESOURCE_ERROR_RESPONSES from ...models.core import User from ...models.core.billing import BillingCustomer, BillingSubscription +from ...operations.graph.capacity import ( + tier_capacity_status as _tier_capacity_status, +) from ...operations.providers.payment_provider import get_payment_provider logger = get_logger(__name__) - -# The same refuse-the-sale rule guards change-tier; the helper lives in the -# ops layer so both paths read one definition of "capacity". -from ...operations.graph.capacity import ( # noqa: E402 - tier_capacity_status as _tier_capacity_status, -) - router = APIRouter(prefix="/billing", tags=["Billing"]) From ace5b80ad4e2807795437443c15ebf28f022a2ee Mon Sep 17 00:00:00 2001 From: "Joseph T. French" Date: Fri, 28 Aug 2026 17:37:07 -0500 Subject: [PATCH 3/3] fix(graph): drain refuses without a private IP; confirm unreachable before trusting it MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Review follow-ups on #1305. A source instance with no private IP in the registry can no longer be taken as drained — nothing could be asked, so nothing is detached. And "unreachable" is now confirmed across three attempts before it is read as a stopped container, so a network blip during the drain does not pass for the maintenance window. --- .../graph/tasks/graph_tier_upgrade.py | 45 +++++++++++++------ .../graph/tasks/test_graph_tier_upgrade.py | 34 +++++++++++++- 2 files changed, 63 insertions(+), 16 deletions(-) diff --git a/robosystems/operations/graph/tasks/graph_tier_upgrade.py b/robosystems/operations/graph/tasks/graph_tier_upgrade.py index b999e6887..4414bde76 100644 --- a/robosystems/operations/graph/tasks/graph_tier_upgrade.py +++ b/robosystems/operations/graph/tasks/graph_tier_upgrade.py @@ -37,6 +37,9 @@ DRAIN_TIMEOUT_SECONDS = 120 DRAIN_POLL_INTERVAL_SECONDS = 5 +# Consecutive failed attempts before an unreachable graph API is taken to mean +# the container is stopped rather than the network blinked. +DRAIN_UNREACHABLE_CONFIRMATIONS = 3 REATTACH_TIMEOUT_SECONDS = 300 REATTACH_POLL_INTERVAL_SECONDS = 10 GRAPH_API_PORT = 8001 @@ -321,27 +324,41 @@ async def _drain_instance(self, private_ip: str) -> None: """Confirm the old instance has stopped serving before the detach. Two outcomes count as drained: the instance reports zero active - connections, or its graph API is not listening at all — the - maintenance-window procedure stops the container before a tier change - precisely so nothing can be writing when the volume detaches, and every - write reaches the volume through that API. A graph API that is up but - has no drain endpoint is neither, and so is a drain that times out with - connections still open: raise, because detaching under live writes is the - one thing this step exists to prevent. + connections, or its graph API is not listening at all, confirmed across + several attempts so a network blip does not pass for a stopped container + — the maintenance-window procedure stops the container before a tier + change precisely so nothing can be writing when the volume detaches, and + every write reaches the volume through that API. Everything else refuses: + no private IP to ask, a graph API that is up but has no drain endpoint, an + error response, or a drain that times out with connections still open. + Detaching under live writes is the one thing this step exists to prevent. """ if not private_ip: - logger.warning("No private IP for drain, skipping") - return + raise DrainRefusedError( + "No private IP recorded for the source instance; cannot confirm it " + "is drained, so the volume stays attached" + ) base_url = f"http://{private_ip}:{GRAPH_API_PORT}" async with httpx.AsyncClient(timeout=10) as client: - try: - response = await client.post(f"{base_url}/admin/drain") - except httpx.HTTPError as e: + response = None + for attempt in range(1, DRAIN_UNREACHABLE_CONFIRMATIONS + 1): + try: + response = await client.post(f"{base_url}/admin/drain") + break + except httpx.HTTPError as e: + logger.warning( + f"Graph API on {private_ip} not reachable " + f"(attempt {attempt}/{DRAIN_UNREACHABLE_CONFIRMATIONS}): {e}" + ) + if attempt < DRAIN_UNREACHABLE_CONFIRMATIONS: + await asyncio.sleep(DRAIN_POLL_INTERVAL_SECONDS) + if response is None: logger.warning( - f"Graph API on {private_ip} is not reachable ({e}); treating the " - "instance as drained — nothing can be writing through it" + f"Graph API on {private_ip} stayed unreachable across " + f"{DRAIN_UNREACHABLE_CONFIRMATIONS} attempts; treating the instance " + "as drained — nothing can be writing through it" ) return diff --git a/tests/operations/graph/tasks/test_graph_tier_upgrade.py b/tests/operations/graph/tasks/test_graph_tier_upgrade.py index da92d67ba..4ab740d16 100644 --- a/tests/operations/graph/tasks/test_graph_tier_upgrade.py +++ b/tests/operations/graph/tasks/test_graph_tier_upgrade.py @@ -589,9 +589,14 @@ async def test_progress_reporting(self, mock_dynamodb): class TestDrainInstance: """Tests for the _drain_instance helper method.""" - async def test_drain_with_no_private_ip_skips(self): + async def test_drain_with_no_private_ip_refuses(self): + from robosystems.operations.graph.tasks.graph_tier_upgrade import ( + DrainRefusedError, + ) + task = _make_task() - await task._drain_instance("") + with pytest.raises(DrainRefusedError, match="No private IP"): + await task._drain_instance("") async def test_drain_succeeds_when_connections_reach_zero(self): task = _make_task() @@ -626,6 +631,31 @@ async def test_drain_treats_unreachable_api_as_drained(self): await task._drain_instance("10.0.1.5") + # Unreachable is confirmed across several attempts, not taken on faith. + assert client.post.await_count == 3 + client.get.assert_not_called() + + async def test_drain_does_not_mistake_a_blip_for_a_stopped_container(self): + from robosystems.operations.graph.tasks.graph_tier_upgrade import ( + DrainRefusedError, + ) + + task = _make_task() + post_response = MagicMock() + post_response.status_code = 404 + + with patch("httpx.AsyncClient") as MockClient: + client = AsyncMock() + # First attempt fails, second reaches an API with no drain endpoint. + client.post = AsyncMock(side_effect=[httpx.ConnectError("blip"), post_response]) + client.get = AsyncMock() + MockClient.return_value.__aenter__ = AsyncMock(return_value=client) + MockClient.return_value.__aexit__ = AsyncMock(return_value=False) + + with pytest.raises(DrainRefusedError, match="no drain endpoint"): + await task._drain_instance("10.0.1.5") + + assert client.post.await_count == 2 client.get.assert_not_called() async def test_drain_refuses_when_api_has_no_drain_endpoint(self):