Harden ops run identity and leases

Assistant: codex
Assistant-Model: gpt-5.6-sol
Assistant-Session: 01a028de-e2c8-7732-8521-46a7fc5db82f

src/activity_core/ops_run_queue.py

@@ -36,6 +36,20 @@ def default_lease_seconds() -> int:
         return 900
 
 
+def _utcnow() -> datetime:
+    return datetime.now(timezone.utc)
+
+
+def _has_active_lease(row: OpsRun, *, now: datetime) -> bool:
+    """Return whether ``row`` has a lease strictly after ``now``."""
+    lease_until = row.lease_until
+    if not isinstance(lease_until, datetime):
+        return False
+    if lease_until.tzinfo is None:
+        lease_until = lease_until.replace(tzinfo=timezone.utc)
+    return lease_until > now
+
+
 def max_attempts() -> int:
     try:
         return max(1, int(os.environ.get("OPS_RUN_MAX_ATTEMPTS", "3")))
@@ -163,14 +177,14 @@ async def create_ops_run_from_spec(
 
 async def reopen_stale_claims(session: AsyncSession) -> int:
     """Return claimed rows with expired leases to open."""
-    now = datetime.now(timezone.utc)
+    now = _utcnow()
     stmt = (
         update(OpsRun)
         .where(
             and_(
                 OpsRun.state == "claimed",
                 OpsRun.lease_until.is_not(None),
-                OpsRun.lease_until < now,
+                OpsRun.lease_until <= now,
             )
         )
         .values(
@@ -252,13 +266,15 @@ async def heartbeat_ops_run(
     worker_id: str,
     lease_seconds: int | None = None,
 ) -> OpsRun | None:
-    row = await session.get(OpsRun, run_id)
+    row = await session.get(OpsRun, run_id, with_for_update=True)
     if row is None:
         return None
     if row.state != "claimed" or row.claim_owner != worker_id:
         return None
+    now = _utcnow()
+    if not _has_active_lease(row, now=now):
+        return None
     lease_seconds = lease_seconds or default_lease_seconds()
-    now = datetime.now(timezone.utc)
     row.lease_until = now + timedelta(seconds=lease_seconds)
     row.updated_at = now
     return row
@@ -271,12 +287,14 @@ async def complete_ops_run(
     worker_id: str,
     result: dict[str, Any] | None = None,
 ) -> OpsRun | None:
-    row = await session.get(OpsRun, run_id)
+    row = await session.get(OpsRun, run_id, with_for_update=True)
     if row is None:
         return None
     if row.state != "claimed" or row.claim_owner != worker_id:
         return None
-    now = datetime.now(timezone.utc)
+    now = _utcnow()
+    if not _has_active_lease(row, now=now):
+        return None
     row.state = "succeeded"
     row.lease_until = None
     row.result = normalise_ops_result(result)
@@ -293,12 +311,14 @@ async def fail_ops_run(
     reopen: bool = False,
     result: dict[str, Any] | None = None,
 ) -> OpsRun | None:
-    row = await session.get(OpsRun, run_id)
+    row = await session.get(OpsRun, run_id, with_for_update=True)
     if row is None:
         return None
     if row.state != "claimed" or row.claim_owner != worker_id:
         return None
-    now = datetime.now(timezone.utc)
+    now = _utcnow()
+    if not _has_active_lease(row, now=now):
+        return None
     payload = normalise_ops_result(result)
     if error:
         payload["error"] = error[:2000]

src/activity_core/ops_runs_api.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+import hmac
 import os
 import uuid
 from datetime import datetime
@@ -16,6 +17,7 @@ from activity_core.ops_auth import (
     extract_operator_token,
     extract_sso_principal,
     operator_token_configured,
+    require_operator,
 )
 from activity_core.ops_run_queue import (
     claim_ops_runs,
@@ -34,6 +36,7 @@ router = APIRouter(prefix="/ops-runs", tags=["ops-runs"])
 _get_db: Callable[[], async_sessionmaker[AsyncSession]] | None = None
 
 WORKER_TOKEN_ENV = "ACTIVITY_CORE_WORKER_TOKEN"
+WORKER_ID_ENV = "ACTIVITY_CORE_WORKER_ID"
 
 
 def bind_ops_runs_deps(
@@ -53,6 +56,11 @@ def _worker_token_configured() -> bool:
     return bool((os.environ.get(WORKER_TOKEN_ENV) or "").strip())
 
 
+def configured_worker_id() -> str | None:
+    value = (os.environ.get(WORKER_ID_ENV) or "").strip()
+    return value or None
+
+
 def _extract_worker_token(
     *,
     x_worker_token: str | None,
@@ -77,8 +85,18 @@ def require_worker_or_operator(
         x_worker_token=x_worker_token, authorization=authorization
     )
     expected_worker = (os.environ.get(WORKER_TOKEN_ENV) or "").strip()
-    if expected_worker and worker_tok and worker_tok == expected_worker:
-        return f"worker:{worker_tok[:8]}"
+    if (
+        expected_worker
+        and worker_tok
+        and hmac.compare_digest(worker_tok, expected_worker)
+    ):
+        worker_id = configured_worker_id()
+        if not worker_id:
+            raise HTTPException(
+                status_code=503,
+                detail=f"worker identity not configured ({WORKER_ID_ENV})",
+            )
+        return f"worker:{worker_id}"
 
     sso = extract_sso_principal(request)
     if sso:
@@ -88,18 +106,61 @@ def require_worker_or_operator(
         x_operator_token=x_operator_token, authorization=authorization
     )
     expected_op = (os.environ.get("ACTIVITY_CORE_OPERATOR_TOKEN") or "").strip()
-    if expected_op and op_tok and op_tok == expected_op:
+    if expected_op and op_tok and hmac.compare_digest(op_tok, expected_op):
         return "operator:token"
 
     if not _worker_token_configured() and not operator_token_configured():
         if allow_unauth_mutations():
             return "dev:unauth"
-        # Dev convenience when no tokens configured at all
-        return "dev:open"
+
+        raise HTTPException(status_code=503, detail="ops run auth not configured")
 
     raise HTTPException(status_code=401, detail="worker or operator auth required")
 
 
+def require_worker(
+    request: Request,
+    *,
+    x_worker_token: str | None = None,
+    authorization: str | None = None,
+) -> str:
+    """Return the configured worker identity for a valid worker credential."""
+    del request  # reserved for future mTLS/proxy-bound worker principals
+    expected_worker = (os.environ.get(WORKER_TOKEN_ENV) or "").strip()
+    worker_tok = _extract_worker_token(
+        x_worker_token=x_worker_token, authorization=authorization
+    )
+
+    if expected_worker:
+        if not worker_tok or not hmac.compare_digest(worker_tok, expected_worker):
+            raise HTTPException(status_code=401, detail="worker auth required")
+        worker_id = configured_worker_id()
+        if not worker_id:
+            raise HTTPException(
+                status_code=503,
+                detail=f"worker identity not configured ({WORKER_ID_ENV})",
+            )
+        return worker_id
+
+    if allow_unauth_mutations() and not operator_token_configured():
+        return "dev:unauth"
+
+    raise HTTPException(status_code=503, detail="worker auth not configured")
+
+
+def bind_worker_id(requested_worker_id: str, authenticated_worker_id: str) -> str:
+    """Bind a body compatibility field to its authenticated queue principal."""
+    requested = requested_worker_id.strip()
+    if authenticated_worker_id == "dev:unauth":
+        return requested
+    if requested != authenticated_worker_id:
+        raise HTTPException(
+            status_code=403,
+            detail="worker_id does not match authenticated worker",
+        )
+    return authenticated_worker_id
+
+
 class ClaimBody(BaseModel):
     worker_id: str = Field(..., min_length=1, max_length=256)
     labels: list[str] | None = None
@@ -195,21 +256,20 @@ async def post_claim(
     body: ClaimBody,
     request: Request,
     x_worker_token: str | None = Header(default=None, alias="X-Worker-Token"),
-    x_operator_token: str | None = Header(default=None, alias="X-Operator-Token"),
     authorization: str | None = Header(default=None),
 ) -> dict[str, Any]:
-    require_worker_or_operator(
+    authenticated_worker_id = require_worker(
         request,
         x_worker_token=x_worker_token,
-        x_operator_token=x_operator_token,
         authorization=authorization,
     )
+    worker_id = bind_worker_id(body.worker_id, authenticated_worker_id)
     Session = _db()
     async with Session() as session:
         async with session.begin():
             claimed = await claim_ops_runs(
                 session,
-                worker_id=body.worker_id,
+                worker_id=worker_id,
                 labels=body.labels,
                 labels_mode=body.labels_mode,
                 limit=body.limit,
@@ -227,28 +287,27 @@ async def post_heartbeat(
     body: HeartbeatBody,
     request: Request,
     x_worker_token: str | None = Header(default=None, alias="X-Worker-Token"),
-    x_operator_token: str | None = Header(default=None, alias="X-Operator-Token"),
     authorization: str | None = Header(default=None),
 ) -> dict[str, Any]:
-    require_worker_or_operator(
+    authenticated_worker_id = require_worker(
         request,
         x_worker_token=x_worker_token,
-        x_operator_token=x_operator_token,
         authorization=authorization,
     )
+    worker_id = bind_worker_id(body.worker_id, authenticated_worker_id)
     Session = _db()
     async with Session() as session:
         async with session.begin():
             row = await heartbeat_ops_run(
                 session,
                 run_id,
-                worker_id=body.worker_id,
+                worker_id=worker_id,
                 lease_seconds=body.lease_seconds,
             )
             if row is None:
                 raise HTTPException(
                     status_code=409,
-                    detail="not claimed by this worker or not found",
+                    detail="not actively leased by this worker or not found",
                 )
             return ops_run_to_dict(row)
 
@@ -259,28 +318,27 @@ async def post_complete(
     body: CompleteBody,
     request: Request,
     x_worker_token: str | None = Header(default=None, alias="X-Worker-Token"),
-    x_operator_token: str | None = Header(default=None, alias="X-Operator-Token"),
     authorization: str | None = Header(default=None),
 ) -> dict[str, Any]:
-    require_worker_or_operator(
+    authenticated_worker_id = require_worker(
         request,
         x_worker_token=x_worker_token,
-        x_operator_token=x_operator_token,
         authorization=authorization,
     )
+    worker_id = bind_worker_id(body.worker_id, authenticated_worker_id)
     Session = _db()
     async with Session() as session:
         async with session.begin():
             row = await complete_ops_run(
                 session,
                 run_id,
-                worker_id=body.worker_id,
+                worker_id=worker_id,
                 result=body.result,
             )
             if row is None:
                 raise HTTPException(
                     status_code=409,
-                    detail="not claimed by this worker or not found",
+                    detail="not actively leased by this worker or not found",
                 )
             return ops_run_to_dict(row)
 
@@ -291,22 +349,21 @@ async def post_fail(
     body: FailBody,
     request: Request,
     x_worker_token: str | None = Header(default=None, alias="X-Worker-Token"),
-    x_operator_token: str | None = Header(default=None, alias="X-Operator-Token"),
     authorization: str | None = Header(default=None),
 ) -> dict[str, Any]:
-    require_worker_or_operator(
+    authenticated_worker_id = require_worker(
         request,
         x_worker_token=x_worker_token,
-        x_operator_token=x_operator_token,
         authorization=authorization,
     )
+    worker_id = bind_worker_id(body.worker_id, authenticated_worker_id)
     Session = _db()
     async with Session() as session:
         async with session.begin():
             row = await fail_ops_run(
                 session,
                 run_id,
-                worker_id=body.worker_id,
+                worker_id=worker_id,
                 error=body.error,
                 reopen=body.reopen,
                 result=body.result,
@@ -314,7 +371,7 @@ async def post_fail(
             if row is None:
                 raise HTTPException(
                     status_code=409,
-                    detail="not claimed by this worker or not found",
+                    detail="not actively leased by this worker or not found",
                 )
             return ops_run_to_dict(row)
 
@@ -322,13 +379,11 @@ async def post_fail(
 @router.post("/expire-leases")
 async def post_expire_leases(
     request: Request,
-    x_worker_token: str | None = Header(default=None, alias="X-Worker-Token"),
     x_operator_token: str | None = Header(default=None, alias="X-Operator-Token"),
     authorization: str | None = Header(default=None),
 ) -> dict[str, Any]:
-    require_worker_or_operator(
+    await require_operator(
         request,
-        x_worker_token=x_worker_token,
         x_operator_token=x_operator_token,
         authorization=authorization,
     )