DRU-856 -- Refuse a create at a provider's host limit

docs/api.md

@@ -40,6 +40,10 @@ or claimed the host, `admin` for an admin key, or `null` for an unclaimed
 warm host. Callers cannot set it. An `Idempotency-Key` belongs to the
 service account that first used it. Another one reusing it gets `409`.
 
+`POST /hosts` returns `503` with `CAPACITY` and `Retry-After` when the
+provider has `HOST_LIMITS` live hosts. A live host is not in `error` and
+has not expired.
+
 ## Refresh a host secret
 
 `POST /hosts/{host_id}/secrets/{service}/refresh` makes the exchange drop

docs/deploy.md

@@ -586,6 +586,7 @@ Core, optional:
 | `LEASE_DEFAULT_TTL` | `86400` | Lease TTL in seconds for hosts created without an explicit `expires_at`, and the extension applied by an empty `POST /hosts/{id}/renew`. An explicit `expires_at: null` at create time opts out of expiry entirely. |
 | `IDEMPOTENCY_KEY_TTL_HOURS` | `24` | Retention period for successful `Idempotency-Key` mappings. |
 | `POOL_SIZES` | `{}` | Warm hosts to keep ready per provider, as JSON (e.g. `{"exe": 2, "hetzner": 1}`). Overrides `POOL_SIZE` for the providers it names. |
+| `HOST_LIMITS` | `{}` | Most live hosts per provider, as JSON (e.g. `{"exe": 28, "docker": 6}`). A create at the limit returns `503` with `CAPACITY`. Warm pool hosts count. A provider it does not name has no limit. |
 | `POOL_SIZE` | `0` | Warm hosts to keep ready for the default provider. `0` disables its pool. |
 | `POOL_HOST_MAX_AGE_HOURS` | `4` | Max age before the janitor reaps an unclaimed pool host. |
 | `POOL_MAX_CREATES_PER_TICK` | `2` | Upper bound on pool provisions per tick, across all providers; caps over-provision blast radius when ticks overlap. |

src/core/settings.py

@@ -165,6 +165,14 @@ class Settings(BaseSettings):
             "Overrides POOL_SIZE for the providers it names."
         ),
     )
+    host_limits: dict[str, Annotated[int, Field(ge=0)]] = Field(
+        default_factory=dict,
+        validation_alias="HOST_LIMITS",
+        description=(
+            'Most live hosts per provider, as JSON (e.g. {"exe": 28, "docker": 6}). '
+            "A provider it does not name has no limit."
+        ),
+    )
     pool_size: int = Field(
         default=0,
         ge=0,

src/hosts/exceptions.py

@@ -1,3 +1,5 @@
+from typing import ClassVar
+
 from core.exceptions import AppException
 
 
@@ -16,6 +18,12 @@ class HostTeardownError(AppException):
     error_code = "HOST_TEARDOWN"
 
 
+class HostCapacityError(AppException):
+    status_code = 503
+    error_code = "CAPACITY"
+    headers: ClassVar[dict[str, str]] = {"Retry-After": "30"}
+
+
 class ProvisioningFailedError(AppException):
     status_code = 502
     error_code = "PROVISIONING_FAILED"

src/hosts/service.py

@@ -8,7 +8,7 @@
 from typing import Any
 
 import httpx
-from sqlalchemy import delete, or_, select, update
+from sqlalchemy import delete, func, or_, select, update
 from sqlalchemy.exc import IntegrityError
 from sqlalchemy.ext.asyncio import AsyncSession
 from uuid6 import uuid7
@@ -20,7 +20,12 @@
 from host_secrets import catalog
 from host_secrets.exceptions import SecretsProxyNotConfiguredError, SecretStaticError
 from host_secrets.placeholder import Placeholder
-from hosts.exceptions import HostStateError, IdempotencyKeyConflictError, ProvisioningFailedError
+from hosts.exceptions import (
+    HostCapacityError,
+    HostStateError,
+    IdempotencyKeyConflictError,
+    ProvisioningFailedError,
+)
 from hosts.models import Host, HostStatus, IdempotencyKey
 from networking.tailscale import (
     DeviceDiscoveryTimeoutError,
@@ -286,6 +291,25 @@ async def create_host(
             expires_at=initial_expires_at,
             pool_member=pool_member,
         )
+        if vm.name in self.settings.host_limits:
+            # The lock serializes this provider's creates until the insert
+            # commits, so two requests cannot both pass the count. SQLite has
+            # no advisory locks; it serves single-process quickstarts.
+            if self.session.get_bind().dialect.name == "postgresql":
+                await self.session.execute(
+                    select(func.pg_advisory_xact_lock(func.hashtext(vm.name)))
+                )
+            live = (
+                await self.session.execute(
+                    select(func.count())
+                    .select_from(Host)
+                    .where(Host.provider == vm.name)
+                    .where(Host.status != HostStatus.ERROR.value)
+                    .where(or_(Host.expires_at.is_(None), Host.expires_at > now))
+                )
+            ).scalar_one()
+            if live >= self.settings.host_limits[vm.name]:
+                raise HostCapacityError(f"provider {vm.name!r} is at its host limit")
         self.session.add(host)
         await self.session.commit()
         await self.session.refresh(host)

src/hosts/tests/test_pool.py

@@ -7,7 +7,8 @@
 from uuid6 import uuid7
 
 from core.database import async_session_factory
-from core.settings import get_settings
+from core.settings import Settings, get_settings
+from hosts.exceptions import HostCapacityError
 from hosts.models import Host, HostStatus
 from hosts.pool import _MAX_CONSECUTIVE_CREATE_FAILURES, maintain_pool
 from hosts.service import HostService, utc_now
@@ -208,6 +209,27 @@ async def test_create_host_never_claims_another_providers_pool_member(
     mocked_provision.assert_awaited_once()
 
 
+async def test_create_host_refuses_a_provider_at_its_host_limit(monkeypatch, stub_provider):
+    # The warm pool member counts toward exe's limit; stub has none.
+    monkeypatch.setenv("HOST_LIMITS", '{"exe": 1}')
+    settings = Settings()  # pyright: ignore[reportCallIssue]
+    await _seed_pool_host(name="lb-pool-exe")
+    monkeypatch.setattr("hosts.service.HostService.provision", AsyncMock())
+
+    async with async_session_factory() as session:
+        with pytest.raises(HostCapacityError):
+            await HostService(session, settings=settings).create_host(
+                env={}, image=None, provider="exe"
+            )
+    async with async_session_factory() as session:
+        await HostService(session, settings=settings).create_host(
+            env={}, image=None, provider="stub"
+        )
+        counted = await session.execute(select(Host.provider, func.count()).group_by(Host.provider))
+
+    assert {provider: count for provider, count in counted.all()} == {"exe": 1, "stub": 1}
+
+
 async def test_create_host_skips_pool_when_sized(multi_pool_settings, monkeypatch, stub_provider):
     # Pool members are warmed at the provider's default size, so a sized
     # request must provision fresh even when its provider has a warm host.