Record the token name that creates each host

alembic/versions/0007_host_service_account.py

@@ -0,0 +1,36 @@
+"""Record the service account that holds each host.
+
+Revision ID: 0007_host_service_account
+Revises: 0006_service_accounts
+"""
+
+from collections.abc import Sequence
+
+import sqlalchemy as sa
+from alembic import op
+
+revision: str = "0007_host_service_account"
+down_revision: str | None = "0006_service_accounts"
+branch_labels: str | Sequence[str] | None = None
+depends_on: str | Sequence[str] | None = None
+
+
+def upgrade() -> None:
+    with op.batch_alter_table("service_accounts") as service_accounts:
+        service_accounts.alter_column("fingerprint", nullable=True)
+    op.execute(sa.text("INSERT INTO service_accounts (name) VALUES ('admin')"))
+    op.add_column("hosts", sa.Column("service_account", sa.String(64)))
+    # Before this revision only admin keys could create or claim a host.
+    op.execute(
+        sa.text(
+            "UPDATE hosts SET service_account = 'admin' "
+            "WHERE NOT pool_member OR claimed_at IS NOT NULL"
+        )
+    )
+
+
+def downgrade() -> None:
+    op.drop_column("hosts", "service_account")
+    op.execute(sa.text("DELETE FROM service_accounts WHERE name = 'admin'"))
+    with op.batch_alter_table("service_accounts") as service_accounts:
+        service_accounts.alter_column("fingerprint", nullable=False)

api-tests/tests/full-api.spec.js

@@ -28,6 +28,7 @@ const EXPECTED_OPENAPI_OPERATIONS = [
 const HOST_KEYS = [
   "id",
   "name",
+  "service_account",
   "status",
   "provider",
   "image",

docs/api.md

@@ -18,7 +18,8 @@ service accounts. A missing token returns `401`, a rejected one `403`.
 and the SHA-256 fingerprint of the token. Names are 1–64 lowercase
 letters, digits, or hyphens. A duplicate name returns `409`.
 `DELETE /service-accounts/ci` revokes the token on the next request, or
-returns `404` for an unknown name.
+returns `404` for an unknown name. The `admin` account exists from the
+migration, has no token, and cannot be removed. Admin keys act as it.
 
 ## Endpoints
 
@@ -33,6 +34,11 @@ returns `404` for an unknown name.
 - `GET /doctor` — read-only dependency diagnostics
 - `GET /healthz` — unauthenticated liveness probe
 
+Host responses carry `service_account`: the service account that created
+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`.
+
 ## The secrets exchange
 
 The exchange is a second process, `python -m secrets_exchange`, on a private

src/conftest.py

@@ -39,14 +39,17 @@ def load_test_env() -> None:
 
 @pytest.fixture(autouse=True)
 async def reset_database() -> AsyncGenerator[None]:
+    from sqlalchemy import insert
+
     from core.database import Base, engine
     from hosts import models  # noqa: F401
-    from service_accounts import models as service_account_models  # noqa: F401
+    from service_accounts.models import ServiceAccount
     from templates import models as template_models  # noqa: F401
 
     async with engine.begin() as connection:
         await connection.run_sync(Base.metadata.drop_all)
         await connection.run_sync(Base.metadata.create_all)
+        await connection.execute(insert(ServiceAccount).values(name=ServiceAccount.ADMIN))
 
     yield
 

src/hosts/api.py

@@ -2,14 +2,14 @@
 import uuid
 from typing import Annotated
 
-from fastapi import APIRouter, Depends, Header, HTTPException, Response, status
+from fastapi import APIRouter, Body, Depends, Header, HTTPException, Response, status
 from sqlalchemy.exc import SQLAlchemyError
 
 from hosts.auth import require_auth
 from hosts.deps import get_host_service
 from hosts.exceptions import HostTeardownError
 from hosts.models import Host
-from hosts.schemas import HostCreate, HostOut, HostRenew
+from hosts.schemas import ExpiresAt, HostCreate, HostOut
 from hosts.service import HostService
 from networking.tailscale import NetworkError
 from providers.exceptions import ProviderError, UnknownProviderError, UnsupportedSizingError
@@ -24,6 +24,7 @@
 @router.post("", response_model=HostOut, status_code=status.HTTP_201_CREATED)
 async def create_host(
     service: HostServiceDep,
+    service_account: Annotated[str | None, Depends(require_auth)],
     payload: HostCreate | None = None,
     idempotency_key: Annotated[
         str | None,
@@ -47,6 +48,7 @@ async def create_host(
 
     try:
         return await service.get_or_create_host(
+            service_account=service_account,
             env=host_create.env,
             secrets={name: entry.to_storage() for name, entry in host_create.secrets.items()},
             image=host_create.image,
@@ -83,10 +85,18 @@ async def get_host(host_id: uuid.UUID, service: HostServiceDep) -> Host:
 async def renew_host(
     host_id: uuid.UUID,
     service: HostServiceDep,
-    payload: HostRenew | None = None,
+    expires_at: Annotated[
+        ExpiresAt,
+        Body(
+            embed=True,
+            description=(
+                "Omit or null: extend by LEASE_DEFAULT_TTL from now. "
+                "Renewal never makes a host permanent."
+            ),
+        ),
+    ] = None,
 ) -> Host:
-    host_renew = payload or HostRenew()
-    return await service.renew_host(host_id, expires_at=host_renew.expires_at)
+    return await service.renew_host(host_id, expires_at=expires_at)
 
 
 @router.delete("/{host_id}", status_code=status.HTTP_204_NO_CONTENT)

src/hosts/auth.py

@@ -28,19 +28,17 @@ def is_admin_key(token: str) -> bool:
 async def require_auth(
     credentials: Annotated[HTTPAuthorizationCredentials | None, Depends(bearer_scheme)],
     session: Annotated[AsyncSession, Depends(get_session)],
-) -> None:
+) -> str | None:
     if not credentials:
         raise HTTPException(status_code=401, detail="admin key or service account token required")
     if is_admin_key(credentials.credentials):
-        return
+        return ServiceAccount.ADMIN
 
     try:
-        if await ServiceAccount.authenticate(session, credentials.credentials):
-            return
+        return await ServiceAccount.authenticate(session, credentials.credentials)
     except SQLAlchemyError:
         logger.exception("service account token lookup failed")
         raise HTTPException(status_code=503, detail="service account tokens unavailable") from None
-    raise HTTPException(status_code=403, detail="admin key or service account token rejected")
 
 
 def require_admin_auth(

src/hosts/exceptions.py

@@ -6,6 +6,11 @@ class HostStateError(AppException):
     error_code = "HOST_STATE"
 
 
+class IdempotencyKeyConflictError(AppException):
+    status_code = 409
+    error_code = "IDEMPOTENCY_KEY_CONFLICT"
+
+
 class HostTeardownError(AppException):
     status_code = 503
     error_code = "HOST_TEARDOWN"

src/hosts/models.py

@@ -74,6 +74,7 @@ class Host(Base):
     env: Mapped[dict[str, str]] = mapped_column(_JSONType, default=dict)
     secrets: Mapped[SecretsMapping] = EncryptedJsonField()
     name: Mapped[str] = mapped_column(String(100), unique=True, index=True)
+    service_account: Mapped[str | None] = mapped_column(String(64))
     status: Mapped[str] = mapped_column(String(32), default=HostStatus.PROVISIONING.value)
     provider: Mapped[str] = mapped_column(String(20), default="exe")
     image: Mapped[str] = mapped_column(Text)

src/hosts/schemas.py

@@ -1,8 +1,9 @@
 import re
 import uuid
 from datetime import UTC, datetime
+from typing import Annotated
 
-from pydantic import BaseModel, ConfigDict, Field, ValidationInfo, field_validator
+from pydantic import AfterValidator, BaseModel, ConfigDict, Field, ValidationInfo, field_validator
 
 from host_secrets import catalog
 from host_secrets.schemas import SECRET_NAME_PATTERN, SecretEntry
@@ -13,15 +14,18 @@
 
 
 def _expires_at_must_be_future_and_tz_aware(expires_at: datetime | None) -> datetime | None:
-    if expires_at is None:
-        return None
-    if expires_at.tzinfo is None:
+    if not expires_at:
+        return
+    if not expires_at.tzinfo:
         raise ValueError("expires_at must include a timezone offset")
     if expires_at <= datetime.now(UTC):
         raise ValueError("expires_at must be in the future")
     return expires_at
 
 
+ExpiresAt = Annotated[datetime | None, AfterValidator(_expires_at_must_be_future_and_tz_aware)]
+
+
 class HostCreate(BaseModel):
     image: str | None = None
     template: uuid.UUID | None = Field(
@@ -36,7 +40,7 @@ class HostCreate(BaseModel):
             "service handle. The sandbox receives a placeholder per entry, never the value."
         ),
     )
-    expires_at: datetime | None = None
+    expires_at: ExpiresAt = None
     provider: str | None = Field(
         default=None,
         description="VM provider to provision on. Omit to use the service default.",
@@ -92,22 +96,13 @@ def reject_reserved_env_keys(cls, env: dict[str, str]) -> dict[str, str]:
         environment.get_persist(env)
         return env
 
-    _validate_expires_at = field_validator("expires_at")(_expires_at_must_be_future_and_tz_aware)
-
-
-class HostRenew(BaseModel):
-    # Omitted (or null) means "extend by LEASE_DEFAULT_TTL from now"; renewal
-    # never makes a host permanent — that is a create-time choice.
-    expires_at: datetime | None = None
-
-    _validate_expires_at = field_validator("expires_at")(_expires_at_must_be_future_and_tz_aware)
-
 
 class HostOut(BaseModel):
     model_config = ConfigDict(from_attributes=True)
 
     id: uuid.UUID
     name: str
+    service_account: str | None
     status: str
     provider: str
     image: str

src/hosts/service.py

@@ -20,7 +20,7 @@
 from host_secrets import catalog
 from host_secrets.exceptions import SecretsProxyNotConfiguredError
 from host_secrets.placeholder import Placeholder
-from hosts.exceptions import HostStateError, ProvisioningFailedError
+from hosts.exceptions import HostStateError, IdempotencyKeyConflictError, ProvisioningFailedError
 from hosts.models import Host, HostStatus, IdempotencyKey
 from networking.tailscale import (
     DeviceDiscoveryTimeoutError,
@@ -101,6 +101,7 @@ def _default_lease_expires_at(self) -> datetime:
     async def get_or_create_host(
         self,
         *,
+        service_account: str | None = None,
         env: dict[str, str],
         secrets: dict[str, dict[str, Any]] | None = None,
         image: str | None,
@@ -123,7 +124,7 @@ async def get_or_create_host(
                 raise UnknownProviderError(f"unknown provider {provider!r}; available: {available}")
 
         if idempotency_key:
-            existing = await self._lookup_idempotency_key(idempotency_key)
+            existing = await self._lookup_idempotency_key(idempotency_key, service_account)
             if existing:
                 return existing
 
@@ -136,10 +137,11 @@ async def get_or_create_host(
         customized = env or secrets or image or template or instance_type or disk_gb
         if not customized and self.settings.get_pool_targets().get(requested_provider):
             host = await self._try_claim_pool_host(
-                provider=requested_provider, expires_at=expires_at
+                service_account=service_account, provider=requested_provider, expires_at=expires_at
             )
         if not host:
             host = await self.create_host(
+                service_account=service_account,
                 env=env,
                 secrets=secrets,
                 image=image,
@@ -158,14 +160,18 @@ async def get_or_create_host(
                 host.claimed_at,
             )
             await self._release_idempotency_loser(host)
-            winner = await self._lookup_idempotency_key(idempotency_key)
+            winner = await self._lookup_idempotency_key(idempotency_key, service_account)
             if not winner:
                 raise HostStateError("idempotency race could not be resolved") from None
             return winner
         return host
 
     async def _try_claim_pool_host(
-        self, *, provider: str, expires_at: datetime | None | EllipsisType
+        self,
+        *,
+        service_account: str | None = None,
+        provider: str,
+        expires_at: datetime | None | EllipsisType,
     ) -> Host | None:
         # Pick a candidate, then atomically claim it with UPDATE ... WHERE
         # claimed_at IS NULL ... RETURNING. The WHERE predicate is the actual
@@ -198,7 +204,12 @@ async def _try_claim_pool_host(
             update(Host)
             .where(Host.id == candidate_id)
             .where(Host.claimed_at.is_(None))
-            .values(claimed_at=now, updated_at=now, expires_at=expires_at)
+            .values(
+                service_account=service_account,
+                claimed_at=now,
+                updated_at=now,
+                expires_at=expires_at,
+            )
             .returning(Host)
         )
         host = result.scalar_one_or_none()
@@ -212,6 +223,7 @@ async def _try_claim_pool_host(
     async def create_host(
         self,
         *,
+        service_account: str | None = None,
         env: dict[str, str],
         secrets: dict[str, dict[str, Any]] | None = None,
         image: str | None,
@@ -259,6 +271,7 @@ async def create_host(
         )
         host = Host(
             id=uid,
+            service_account=service_account,
             env=env,
             secrets=secrets or {},
             name=name,
@@ -321,18 +334,23 @@ async def _resolve_template_image(self, *, template_id: uuid.UUID, provider: str
         template.last_used_at = utc_now()
         return template.image
 
-    async def _lookup_idempotency_key(self, key: str) -> Host | None:
+    async def _lookup_idempotency_key(self, key: str, service_account: str | None) -> Host | None:
         record = (
             await self.session.execute(select(IdempotencyKey).where(IdempotencyKey.key == key))
         ).scalar_one_or_none()
 
         if not record:
             return
 
-        if record.expires_at > utc_now():
-            host = await self.session.get(Host, record.host_id)
-            if host:
-                return host
+        host = (
+            await self.session.get(Host, record.host_id) if record.expires_at > utc_now() else None
+        )
+        if host:
+            if host.service_account != service_account:
+                raise IdempotencyKeyConflictError(
+                    f"idempotency key {key} belongs to another service account"
+                )
+            return host
         # Stale: expired, or the host vanished without the FK cascade firing.
         # GC in a dedicated session so we don't autoflush the caller's pending
         # state on `self.session`.
@@ -374,6 +392,7 @@ async def _release_idempotency_loser(self, host: Host) -> None:
             now = utc_now()
             if fresh.claimed_at:
                 fresh.claimed_at = None
+                fresh.service_account = None
                 fresh.expires_at = now + timedelta(hours=self.settings.pool_host_max_age_hours)
                 fresh.updated_at = now
                 logger.info(