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

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"

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

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
@@ -36,11 +37,19 @@
 
 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
 
 
+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,24 +321,58 @@ 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, 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:
-        await client.post(f"{base_url}/admin/drain")
-      except httpx.HTTPError as e:
-        logger.warning(f"Drain request failed (may be expected): {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} stayed unreachable across "
+          f"{DRAIN_UNREACHABLE_CONFIRMATIONS} attempts; 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 +389,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."""

robosystems/routers/billing/checkout.py

@@ -18,29 +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__)
 
-
-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"
-
-
 router = APIRouter(prefix="/billing", tags=["Billing"])
 
 

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"
@@ -588,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()
@@ -601,7 +607,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 +617,92 @@ 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")
+
+    # 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):
+    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

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"