fix(worker): survive a database restart; clean exit when the database is down at start

.gitignore

@@ -17,3 +17,7 @@ id_rsa*
 # visual brainstorming scratch
 .superpowers/
 testing/
+
+# The worker VM's secrets and the contract files copied onto it by hand
+deploy/worker-vm/worker.env
+deploy/worker-vm/deployments/

control-plane/mergit_cp/app.py

@@ -14,6 +14,7 @@
 from . import (
     auth,
     chain_events,
+    health,
     oauth,
     otp,
     redeem,
@@ -110,6 +111,7 @@ async def lifespan(app: FastAPI) -> AsyncGenerator[None]:
         allow_headers=["Content-Type", "Idempotency-Key"],
         expose_headers=["Retry-After"],  # the app says how long to wait after a 429 (redeem codes)
     )
+    app.include_router(health.router)
     app.include_router(auth.router)
     app.include_router(auth.account_router)
     app.include_router(routes_account.router)

control-plane/mergit_cp/health.py

@@ -0,0 +1,37 @@
+"""``GET /healthz``: is this API able to serve, i.e. is its database reachable?
+
+Used by the compose healthcheck and by deploy/server/update.sh. It is deliberately public and says nothing
+beyond ok or not ok: no version, no settings, no error text.
+"""
+
+import asyncio
+import logging
+from typing import Annotated
+
+from fastapi import APIRouter, Depends
+from fastapi.responses import JSONResponse
+
+from .db import Pool
+from .deps import get_pool
+
+log = logging.getLogger(__name__)
+
+router = APIRouter(tags=["health"])
+
+# Under a proxy's own health timeout, so a stuck database answers 503 instead of hanging the check.
+DB_CHECK_TIMEOUT_S = 3.0
+
+
+async def _database_answers(pool: Pool) -> None:
+    async with pool.connection() as conn:
+        await conn.execute("SELECT 1")
+
+
+@router.get("/healthz", include_in_schema=False)
+async def healthz(pool: Annotated[Pool, Depends(get_pool)]) -> JSONResponse:
+    try:
+        await asyncio.wait_for(_database_answers(pool), DB_CHECK_TIMEOUT_S)
+    except Exception:  # any failure, including a timeout or an exhausted pool, means "not healthy"
+        log.warning("healthz: the database did not answer", exc_info=True)
+        return JSONResponse({"ok": False}, status_code=503)
+    return JSONResponse({"ok": True})

control-plane/mergit_cp/pr_store.py

@@ -125,9 +125,15 @@ def connect(cls, database_url: str, worker_id: str = "api") -> PgPrStore:
 
     @staticmethod
     def _open(database_url: str) -> psycopg.Connection[dict[str, Any]]:
+        # connect_timeout: a database that is down or unreachable must not hang this reconnect forever.
         return cast(
             "psycopg.Connection[dict[str, Any]]",
-            psycopg.connect(database_url, autocommit=True, row_factory=dict_row),  # pyright: ignore[reportArgumentType]
+            psycopg.connect(
+                database_url,
+                autocommit=True,
+                row_factory=dict_row,  # pyright: ignore[reportArgumentType]
+                connect_timeout=10,
+            ),
         )
 
     def _retrying(self, fn: Callable[[], _T]) -> _T:

control-plane/mergit_worker/__main__.py

@@ -8,6 +8,8 @@
 import threading
 from dataclasses import dataclass
 
+import psycopg
+
 from mergit_cp.pr_store import PgPrStore
 
 from .agents import validate_agents_are_priced
@@ -23,6 +25,7 @@
 from .providers import llm_provider
 from .publish import ReceiptPublisher
 from .sandbox import Sandbox, SandboxConfig, reap_stale
+from .startup import LOCK_CHECK_INTERVAL_S, WorkerIdLock, agents_without_key, sweep_workspaces, watch_lock
 from .tools import ToolLimits
 from .worker import Worker, WorkerDeps
 from .workspace import GitLimits
@@ -141,8 +144,46 @@ def main() -> int:
         return 2
 
     worker_id = cfg.worker_id or socket.gethostname()
+    if not cfg.worker_id:
+        log.warning(
+            "WORKER_ID is not set, so this worker is named after its hostname (%s). A container's hostname "
+            "changes whenever it is recreated, and a restarted worker would then no longer find the jobs and "
+            "sandbox containers its previous run left. Set WORKER_ID to a fixed name, one per process.",
+            worker_id,
+        )
+    try:
+        no_key = agents_without_key(cfg.agents_file)
+    except ConfigError as exc:
+        log.error("configuration error: %s", exc)
+        return 2
+    if no_key:
+        names = ", ".join(f"MERGIT_AGENT_KEY_{i}" for i in no_key)
+        if cfg.require_all_agent_keys:
+            log.error("REQUIRE_ALL_AGENT_KEYS is set and these are missing or malformed: %s", names)
+            return 2
+        log.warning(
+            "no usable key for agent(s) %s: a job for them fails, and the requester loses an attempt", names
+        )
+    try:
+        lock = WorkerIdLock.acquire(cfg.database_url.get_secret_value(), worker_id)
+    except psycopg.Error as exc:
+        # Exit non-zero without a traceback; Docker's restart policy is the retry loop.
+        log.error("cannot reach the database at start-up: %s", " ".join(str(exc).split()))
+        return 1
+    if lock is None:
+        log.error(
+            "another worker process is already running as %r. Starting here would fail its running jobs and "
+            "remove its sandbox containers. Give each process its own WORKER_ID.",
+            worker_id,
+        )
+        return 2
     log.info("model provider: %s", llm_provider())
-    built = build_worker(cfg, worker_id)
+    try:
+        built = build_worker(cfg, worker_id)
+    except psycopg.Error as exc:
+        log.error("cannot reach the database at start-up: %s", " ".join(str(exc).split()))
+        lock.close()
+        return 1
     log.info("failed %d jobs this worker left running before a restart", built.jobs.fail_own_running())
     log.info(
         "reclaimed %d jobs left running past their lease by any worker",
@@ -160,6 +201,10 @@ def main() -> int:
     else:
         log.info("no GITHUB_TOKEN: pull-request requests will not be handled by this worker")
     log.info("removed %d stale sandbox containers", reap_stale(worker_id))
+    log.info(
+        "removed %d abandoned workspace directories",
+        sweep_workspaces(cfg.sandbox_workspace_root, cfg.job_lease_s),
+    )
 
     stop = threading.Event()
 
@@ -168,11 +213,30 @@ def request_stop(_signum: int, _frame: object) -> None:
 
     for sig in (signal.SIGINT, signal.SIGTERM):
         signal.signal(sig, request_stop)
+    lost = threading.Event()
+
+    def on_lost() -> None:
+        log.error(
+            "another process took the worker id %r while the database was unreachable, and has failed this "
+            "worker's running jobs. Stopping after the current job.",
+            worker_id,
+        )
+        lost.set()
+        stop.set()
+
+    threading.Thread(
+        target=watch_lock,
+        args=(lock, stop, LOCK_CHECK_INTERVAL_S, on_lost),
+        name="worker-id-lock",
+        daemon=True,
+    ).start()
     try:
         built.worker.run_forever(stop)
     finally:
+        stop.set()
         built.close()
-    return 0
+        lock.close()
+    return 2 if lost.is_set() else 0
 
 
 if __name__ == "__main__":

control-plane/mergit_worker/config.py

@@ -63,6 +63,9 @@ class WorkerConfig(BaseSettings):
     # Under the 1800 s dispatch deadline, leaving time to run Φ, seal and publish.
     job_deadline_s: float = Field(default=1500.0, gt=0)
     worker_id: str | None = None  # defaults to the hostname; one process per id
+    # Refuse to start unless every agent in the agents file has a MERGIT_AGENT_KEY_<id>. Off by default: a
+    # worker may deliberately serve only some of the agents (startup.agents_without_key logs the rest).
+    require_all_agent_keys: bool = False
     # Opus 5.5's thinking tokens count against this; 8192 made a truncated turn too likely.
     model_max_tokens: int = Field(default=16_000, ge=256)
     poll_interval_s: float = Field(default=2.0, gt=0)

control-plane/mergit_worker/jobs.py

@@ -97,9 +97,15 @@ def connect(cls, database_url: str, worker_id: str) -> PgJobStore:
 
     @staticmethod
     def _open(database_url: str) -> psycopg.Connection[dict[str, Any]]:
+        # connect_timeout: a database that is down or unreachable must not hang this reconnect forever.
         return cast(
             "psycopg.Connection[dict[str, Any]]",
-            psycopg.connect(database_url, autocommit=True, row_factory=dict_row),  # pyright: ignore[reportArgumentType]
+            psycopg.connect(
+                database_url,
+                autocommit=True,
+                row_factory=dict_row,  # pyright: ignore[reportArgumentType]
+                connect_timeout=10,
+            ),
         )
 
     def _retrying(self, fn: Callable[[], _T]) -> _T:

control-plane/mergit_worker/startup.py

@@ -0,0 +1,165 @@
+"""Start-up checks and housekeeping for one worker process (plan task O10).
+
+The worker runs on its own VM, restarted by Docker and redeployed by hand, so what it leaves behind when it
+stops has to be cleaned up by the next start: jobs, sandbox containers (sandbox.reap_stale) and workspace
+directories (here). Everything below is a plain function over a path, an environment or a database URL, so it
+is tested without a Docker daemon.
+"""
+
+from __future__ import annotations
+
+import contextlib
+import json
+import logging
+import os
+import re
+import shutil
+import threading
+import time
+from collections.abc import Callable, Mapping
+from pathlib import Path
+
+import psycopg
+
+from .config import ConfigError
+
+log = logging.getLogger(__name__)
+
+_AGENT_KEY = re.compile(r"^0x[0-9a-fA-F]{64}$")
+WORKSPACE_PREFIX = "ws-"  # workspace.Workspace.clone names its directories ws-<random>
+
+
+def agents_without_key(agents_file: Path, environ: Mapping[str, str] | None = None) -> list[str]:
+    """Agent ids in the agents file that have no usable ``MERGIT_AGENT_KEY_<id>`` here.
+
+    The key is read when a job for that agent starts, so a missing one is otherwise found only after a
+    requester has picked the agent, funded the plan and burned an attempt on it.
+    """
+    env = os.environ if environ is None else environ
+    try:
+        registry: list[dict[str, object]] = json.loads(agents_file.read_text())["agents"]
+    except (OSError, KeyError, ValueError) as exc:
+        raise ConfigError(f"cannot read agent configuration: {exc}") from exc
+    ids = [str(a["agent_id"]) for a in registry if "agent_id" in a]
+    return [i for i in ids if not _AGENT_KEY.match(env.get(f"MERGIT_AGENT_KEY_{i}", ""))]
+
+
+def _remove_workspace(path: Path) -> None:
+    """Delete one workspace tree, including directories a sandbox left unreadable.
+
+    Never follows a symlink: a sandboxed agent can plant one, and chmod or rmtree through it would reach
+    outside the workspace root.
+    """
+    for dirpath, dirnames, _files in os.walk(path, followlinks=False):
+        os.chmod(dirpath, 0o700)
+        for name in dirnames:
+            child = os.path.join(dirpath, name)
+            if not os.path.islink(child):
+                os.chmod(child, 0o700)
+    shutil.rmtree(path)
+
+
+def sweep_workspaces(root: Path, older_than_s: float, *, now: float | None = None) -> int:
+    """Remove ``ws-*`` directories under ``root`` that nothing can still be using. Returns how many.
+
+    A crash or a SIGKILL skips the normal clean-up, and ``reap_stale`` removes containers only, so without
+    this the VM's disk fills one abandoned checkout at a time. ``older_than_s`` must exceed the longest a
+    job can run (the lease, which is also when its row is reclaimed), so a job still in flight in another
+    process on this host is never touched.
+    """
+    if not root.is_dir():
+        return 0
+    cutoff = (time.time() if now is None else now) - older_than_s
+    removed = 0
+    for entry in root.iterdir():
+        if not entry.name.startswith(WORKSPACE_PREFIX) or entry.is_symlink() or not entry.is_dir():
+            continue
+        try:
+            if entry.stat().st_mtime >= cutoff:
+                continue
+            _remove_workspace(entry)
+            removed += 1
+        except OSError as exc:
+            log.warning("could not remove stale workspace %s: %s", entry, exc)
+    return removed
+
+
+class WorkerIdLock:
+    """A session-level Postgres advisory lock named for the worker id, held while the process lives.
+
+    A worker id must belong to one process. At start-up a worker fails every job still ``running`` under its
+    id and removes every sandbox container labelled with it, so a second process sharing the id would
+    destroy the first one's live work. Postgres releases the lock when the connection ends, whether the
+    process exits cleanly or is killed, so a crashed worker never blocks its own restart.
+
+    The same rule means a database restart drops the lock while this process lives on. ``still_held``
+    notices that, reconnects and takes the lock again, and reports False only when another process has
+    taken the id in the meantime: that process has already failed this one's jobs, so this one must stop.
+    """
+
+    def __init__(self, conn: psycopg.Connection[tuple[object, ...]], database_url: str, key: str) -> None:
+        self._conn: psycopg.Connection[tuple[object, ...]] | None = conn
+        self._url = database_url
+        self._key = key
+
+    @classmethod
+    def acquire(cls, database_url: str, worker_id: str) -> WorkerIdLock | None:
+        key = f"mergit_worker:{worker_id}"
+        conn = _take(database_url, key)
+        return None if conn is None else cls(conn, database_url, key)
+
+    def still_held(self) -> bool | None:
+        """True while the lock is held, False once another process owns the id, None while the database
+        cannot be reached (try again later: nothing can have taken the id from a database that is down)."""
+        try:
+            if self._conn is not None and not self._conn.closed:
+                self._conn.execute("SELECT 1")
+                return True
+        except psycopg.Error:
+            self._drop()
+        try:
+            self._conn = _take(self._url, self._key)
+        except psycopg.Error:
+            return None
+        if self._conn is None:
+            return False
+        log.warning("the worker-id lock was lost with the database connection and has been taken again")
+        return True
+
+    def _drop(self) -> None:
+        if self._conn is not None:
+            with contextlib.suppress(psycopg.Error):
+                self._conn.close()
+        self._conn = None
+
+    def close(self) -> None:
+        self._drop()
+
+
+def _take(database_url: str, key: str) -> psycopg.Connection[tuple[object, ...]] | None:
+    # connect_timeout: a database that is down or unreachable must not hang this call (the initial acquire,
+    # or a retake from still_held) forever instead of letting the caller fail and retry.
+    conn = psycopg.connect(database_url, autocommit=True, connect_timeout=10)  # pyright: ignore[reportArgumentType]
+    try:
+        row = conn.execute("SELECT pg_try_advisory_lock(hashtextextended(%s, 0))", (key,)).fetchone()
+    except BaseException:
+        conn.close()
+        raise
+    if row is not None and row[0] is True:
+        return conn
+    conn.close()
+    return None
+
+
+LOCK_CHECK_INTERVAL_S = 10.0
+
+
+def watch_lock(
+    lock: WorkerIdLock, stop: threading.Event, interval_s: float, on_lost: Callable[[], None]
+) -> None:
+    """Check the lock every ``interval_s`` until ``stop`` is set, calling ``on_lost`` once if another
+    process takes the worker id."""
+    while not stop.wait(interval_s):
+        if lock.still_held() is False:
+            on_lost()
+            return