462 lines
22 KiB
Python
462 lines
22 KiB
Python
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
import hashlib
|
|
import hmac
|
|
import json
|
|
from pathlib import Path
|
|
import secrets
|
|
import sqlite3
|
|
import threading
|
|
import time
|
|
from typing import Any
|
|
|
|
from .errors import AuthenticationError, ConfigurationError, ConflictError, ValidationError
|
|
from .models import Decision, Event
|
|
|
|
|
|
REMOTE_SCOPES = frozenset({"claim", "heartbeat", "decision", "fail"})
|
|
|
|
|
|
def require_idempotency_key(value: str | None) -> str:
|
|
if not isinstance(value, str) or not 16 <= len(value) <= 128:
|
|
raise ValidationError("invalid_idempotency_key")
|
|
if any(ch not in "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789._:-" for ch in value):
|
|
raise ValidationError("invalid_idempotency_key")
|
|
return value
|
|
|
|
|
|
def authenticate_runner(config: Any, runner_id: str | None, scope: str, header: str) -> bytes:
|
|
if scope not in REMOTE_SCOPES:
|
|
raise AuthenticationError("invalid_runner_scope")
|
|
if not isinstance(runner_id, str) or not runner_id:
|
|
raise AuthenticationError("missing_runner_id")
|
|
if not header.startswith("Bearer "):
|
|
raise AuthenticationError("missing_bearer")
|
|
try:
|
|
expected = config.remote_runner_token(runner_id, scope)
|
|
except ConfigurationError as exc:
|
|
raise AuthenticationError("runner_not_authorized") from exc
|
|
supplied = header[7:].encode("utf-8")
|
|
if not hmac.compare_digest(supplied, expected):
|
|
raise AuthenticationError("invalid_bearer")
|
|
return expected
|
|
|
|
|
|
def validate_remote_decision(value: dict[str, Any], event: Event) -> Decision:
|
|
decision = Decision.from_dict(value)
|
|
if decision.contract_commit != event.contract_commit:
|
|
raise ValidationError("contract_version_mismatch")
|
|
if decision.event_id != event.event_id or decision.project_key != event.project_key:
|
|
raise ValidationError("event_identity_mismatch")
|
|
if decision.expected_state_version != event.state_version:
|
|
raise ValidationError("state_version_mismatch")
|
|
allowed = event.payload.get("allowed_work_unit_ids", [])
|
|
if not isinstance(allowed, list) or any(not isinstance(item, str) for item in allowed):
|
|
raise ValidationError("invalid_allowed_work_units")
|
|
if any(unit.id not in allowed for unit in decision.next_work_units):
|
|
raise ValidationError("work_unit_not_allowed")
|
|
return decision
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class RemoteLease:
|
|
job_id: str
|
|
runner_id: str
|
|
lease_version: int
|
|
lease_until: int
|
|
lease_token: str
|
|
event: dict[str, Any]
|
|
|
|
def to_dict(self) -> dict[str, Any]:
|
|
return {
|
|
"job_id": self.job_id,
|
|
"runner_id": self.runner_id,
|
|
"lease_version": self.lease_version,
|
|
"lease_until": self.lease_until,
|
|
"lease_token": self.lease_token,
|
|
"event": self.event,
|
|
}
|
|
|
|
|
|
class RemoteCodexBroker:
|
|
"""Durable, exact-once remote decision job broker.
|
|
|
|
Lease credentials are derived with HMAC from the scoped runner token and
|
|
are never logged or persisted in plaintext.
|
|
"""
|
|
|
|
def __init__(self, database_path: Path, *, contract_commit: str,
|
|
lease_seconds: int, max_retries: int):
|
|
self.database_path = Path(database_path)
|
|
self.contract_commit = contract_commit
|
|
self.lease_seconds = lease_seconds
|
|
self.max_retries = max_retries
|
|
self._local = threading.local()
|
|
self._initialize()
|
|
|
|
def _connect(self) -> sqlite3.Connection:
|
|
connection = sqlite3.connect(self.database_path, timeout=30, isolation_level=None)
|
|
connection.row_factory = sqlite3.Row
|
|
connection.execute("PRAGMA journal_mode=WAL")
|
|
connection.execute("PRAGMA foreign_keys=ON")
|
|
connection.execute("PRAGMA busy_timeout=30000")
|
|
return connection
|
|
|
|
def _conn(self) -> sqlite3.Connection:
|
|
connection = getattr(self._local, "connection", None)
|
|
if connection is None:
|
|
connection = self._connect()
|
|
self._local.connection = connection
|
|
return connection
|
|
|
|
def close(self) -> None:
|
|
connection = getattr(self._local, "connection", None)
|
|
if connection is not None:
|
|
connection.close()
|
|
self._local.connection = None
|
|
|
|
def _initialize(self) -> None:
|
|
connection = self._connect()
|
|
try:
|
|
connection.executescript(
|
|
"""
|
|
CREATE TABLE IF NOT EXISTS remote_jobs (
|
|
job_id TEXT PRIMARY KEY,
|
|
event_id TEXT NOT NULL UNIQUE,
|
|
event_digest TEXT NOT NULL,
|
|
event_json TEXT NOT NULL,
|
|
contract_commit TEXT NOT NULL,
|
|
expected_state_version INTEGER NOT NULL,
|
|
status TEXT NOT NULL,
|
|
runner_id TEXT,
|
|
lease_version INTEGER NOT NULL DEFAULT 0,
|
|
lease_until INTEGER,
|
|
attempts INTEGER NOT NULL DEFAULT 0,
|
|
next_attempt_at INTEGER NOT NULL DEFAULT 0,
|
|
failure_category TEXT,
|
|
decision_digest TEXT,
|
|
created_at INTEGER NOT NULL,
|
|
updated_at INTEGER NOT NULL
|
|
);
|
|
CREATE INDEX IF NOT EXISTS idx_remote_jobs_claim
|
|
ON remote_jobs(status, next_attempt_at, lease_until, created_at);
|
|
CREATE TABLE IF NOT EXISTS remote_requests (
|
|
runner_id TEXT NOT NULL,
|
|
scope TEXT NOT NULL,
|
|
idempotency_key TEXT NOT NULL,
|
|
request_digest TEXT NOT NULL,
|
|
job_id TEXT,
|
|
lease_version INTEGER,
|
|
response_json TEXT NOT NULL,
|
|
created_at INTEGER NOT NULL,
|
|
PRIMARY KEY(runner_id, scope, idempotency_key)
|
|
);
|
|
CREATE TABLE IF NOT EXISTS remote_runners (
|
|
runner_id TEXT PRIMARY KEY,
|
|
last_seen_at INTEGER NOT NULL,
|
|
last_scope TEXT NOT NULL
|
|
);
|
|
"""
|
|
)
|
|
finally:
|
|
connection.close()
|
|
|
|
@staticmethod
|
|
def _digest(value: dict[str, Any]) -> str:
|
|
canonical = json.dumps(value, separators=(",", ":"), sort_keys=True).encode()
|
|
return hashlib.sha256(canonical).hexdigest()
|
|
|
|
@staticmethod
|
|
def _lease_token(token: bytes, job_id: str, runner_id: str,
|
|
lease_version: int, claim_key: str) -> str:
|
|
message = f"{job_id}\0{runner_id}\0{lease_version}\0{claim_key}".encode()
|
|
return hmac.new(token, message, hashlib.sha256).hexdigest()
|
|
|
|
def enqueue(self, event: Event) -> tuple[bool, str]:
|
|
if event.contract_commit != self.contract_commit:
|
|
raise ConflictError("contract_version_mismatch")
|
|
now = int(time.time())
|
|
digest = event.digest()
|
|
connection = self._conn()
|
|
connection.execute("BEGIN IMMEDIATE")
|
|
try:
|
|
existing = connection.execute(
|
|
"SELECT event_digest FROM remote_jobs WHERE event_id=?", (event.event_id,)
|
|
).fetchone()
|
|
if existing is not None:
|
|
if existing["event_digest"] != digest:
|
|
raise ConflictError("remote_job_idempotency_conflict")
|
|
connection.execute("COMMIT")
|
|
return False, event.event_id
|
|
connection.execute(
|
|
"""
|
|
INSERT INTO remote_jobs(
|
|
job_id,event_id,event_digest,event_json,contract_commit,
|
|
expected_state_version,status,next_attempt_at,created_at,updated_at
|
|
) VALUES(?,?,?,?,?,?,'PENDING',?,?,?)
|
|
""",
|
|
(event.event_id, event.event_id, digest, event.canonical_json(),
|
|
event.contract_commit, event.state_version, now, now, now),
|
|
)
|
|
connection.execute("COMMIT")
|
|
return True, event.event_id
|
|
except Exception:
|
|
connection.execute("ROLLBACK")
|
|
raise
|
|
|
|
def _touch_runner(self, connection: sqlite3.Connection, runner_id: str, scope: str, now: int) -> None:
|
|
connection.execute(
|
|
"""
|
|
INSERT INTO remote_runners(runner_id,last_seen_at,last_scope) VALUES(?,?,?)
|
|
ON CONFLICT(runner_id) DO UPDATE SET
|
|
last_seen_at=excluded.last_seen_at,last_scope=excluded.last_scope
|
|
""", (runner_id, now, scope),
|
|
)
|
|
|
|
def _replay(self, connection: sqlite3.Connection, runner_id: str, scope: str,
|
|
key: str, digest: str) -> dict[str, Any] | None:
|
|
row = connection.execute(
|
|
"SELECT request_digest,response_json FROM remote_requests "
|
|
"WHERE runner_id=? AND scope=? AND idempotency_key=?",
|
|
(runner_id, scope, key),
|
|
).fetchone()
|
|
if row is None:
|
|
return None
|
|
if row["request_digest"] != digest:
|
|
raise ConflictError("idempotency_payload_conflict")
|
|
return json.loads(row["response_json"])
|
|
|
|
def _record(self, connection: sqlite3.Connection, *, runner_id: str, scope: str,
|
|
key: str, digest: str, response: dict[str, Any], job_id: str | None,
|
|
lease_version: int | None, now: int) -> None:
|
|
connection.execute(
|
|
"""
|
|
INSERT INTO remote_requests(
|
|
runner_id,scope,idempotency_key,request_digest,job_id,
|
|
lease_version,response_json,created_at
|
|
) VALUES(?,?,?,?,?,?,?,?)
|
|
""",
|
|
(runner_id, scope, key, digest, job_id, lease_version,
|
|
json.dumps(response, separators=(",", ":"), sort_keys=True), now),
|
|
)
|
|
|
|
def claim(self, *, runner_id: str, idempotency_key: str,
|
|
runner_token: bytes) -> dict[str, Any]:
|
|
key = require_idempotency_key(idempotency_key)
|
|
request = {"runner_id": runner_id, "contract_commit": self.contract_commit}
|
|
digest = self._digest(request)
|
|
now = int(time.time())
|
|
connection = self._conn()
|
|
connection.execute("BEGIN IMMEDIATE")
|
|
try:
|
|
replay = self._replay(connection, runner_id, "claim", key, digest)
|
|
if replay is not None:
|
|
if replay == {"claimed": False}:
|
|
connection.execute("COMMIT")
|
|
return replay
|
|
row = connection.execute(
|
|
"SELECT status,runner_id,lease_version,lease_until FROM remote_jobs WHERE job_id=?",
|
|
(replay.get("job_id"),),
|
|
).fetchone()
|
|
if row is not None and row["runner_id"] == runner_id and row["status"] == "LEASED":
|
|
replay["lease_token"] = self._lease_token(
|
|
runner_token, str(replay["job_id"]), runner_id,
|
|
int(row["lease_version"]), key,
|
|
)
|
|
connection.execute("COMMIT")
|
|
return replay
|
|
raise ConflictError("claim_replay_no_longer_active")
|
|
|
|
connection.execute(
|
|
"""
|
|
UPDATE remote_jobs SET
|
|
status=CASE WHEN attempts + 1 > ? THEN 'DEAD_LETTER' ELSE 'RETRY' END,
|
|
attempts=attempts+1,runner_id=NULL,lease_until=NULL,
|
|
failure_category='STALE_LEASE',next_attempt_at=?,updated_at=?
|
|
WHERE status='LEASED' AND lease_until < ?
|
|
""", (self.max_retries, now, now, now),
|
|
)
|
|
row = connection.execute(
|
|
"""
|
|
SELECT * FROM remote_jobs
|
|
WHERE status IN ('PENDING','RETRY') AND next_attempt_at <= ?
|
|
AND contract_commit=?
|
|
ORDER BY created_at LIMIT 1
|
|
""", (now, self.contract_commit),
|
|
).fetchone()
|
|
self._touch_runner(connection, runner_id, "claim", now)
|
|
if row is None:
|
|
response = {"claimed": False}
|
|
self._record(connection, runner_id=runner_id, scope="claim", key=key,
|
|
digest=digest, response=response, job_id=None,
|
|
lease_version=None, now=now)
|
|
connection.execute("COMMIT")
|
|
return response
|
|
lease_version = int(row["lease_version"]) + 1
|
|
lease_until = now + self.lease_seconds
|
|
connection.execute(
|
|
"""
|
|
UPDATE remote_jobs SET status='LEASED',runner_id=?,lease_version=?,
|
|
lease_until=?,updated_at=? WHERE job_id=? AND status IN ('PENDING','RETRY')
|
|
""", (runner_id, lease_version, lease_until, now, row["job_id"]),
|
|
)
|
|
lease = RemoteLease(
|
|
job_id=str(row["job_id"]), runner_id=runner_id,
|
|
lease_version=lease_version, lease_until=lease_until,
|
|
lease_token=self._lease_token(runner_token, str(row["job_id"]), runner_id,
|
|
lease_version, key),
|
|
event=json.loads(row["event_json"]),
|
|
)
|
|
response = {"claimed": True, **lease.to_dict(), "claim_idempotency_key": key}
|
|
persisted_response = dict(response)
|
|
persisted_response.pop("lease_token")
|
|
self._record(connection, runner_id=runner_id, scope="claim", key=key,
|
|
digest=digest, response=persisted_response, job_id=lease.job_id,
|
|
lease_version=lease_version, now=now)
|
|
connection.execute("COMMIT")
|
|
return response
|
|
except Exception:
|
|
connection.execute("ROLLBACK")
|
|
raise
|
|
|
|
def _validate_lease(self, connection: sqlite3.Connection, *, job_id: str,
|
|
runner_id: str, lease_token: str, claim_key: str,
|
|
runner_token: bytes, allow_completed: bool = False) -> sqlite3.Row:
|
|
row = connection.execute("SELECT * FROM remote_jobs WHERE job_id=?", (job_id,)).fetchone()
|
|
if row is None:
|
|
raise ConflictError("remote_job_not_found")
|
|
if row["runner_id"] != runner_id:
|
|
raise ConflictError("remote_job_not_leased")
|
|
expected = self._lease_token(runner_token, job_id, runner_id,
|
|
int(row["lease_version"]), claim_key)
|
|
if not hmac.compare_digest(lease_token, expected):
|
|
raise AuthenticationError("invalid_lease_token")
|
|
if allow_completed and row["status"] == "COMPLETED":
|
|
return row
|
|
if row["status"] != "LEASED":
|
|
raise ConflictError("remote_job_not_leased")
|
|
if row["lease_until"] is None or int(row["lease_until"]) < int(time.time()):
|
|
raise ConflictError("remote_lease_expired")
|
|
return row
|
|
|
|
def heartbeat(self, *, job_id: str, runner_id: str, lease_token: str,
|
|
claim_key: str, idempotency_key: str, runner_token: bytes) -> dict[str, Any]:
|
|
key = require_idempotency_key(idempotency_key)
|
|
request = {"job_id": job_id, "lease_token": lease_token, "claim_key": claim_key}
|
|
digest = self._digest(request); now = int(time.time())
|
|
connection = self._conn(); connection.execute("BEGIN IMMEDIATE")
|
|
try:
|
|
replay = self._replay(connection, runner_id, "heartbeat", key, digest)
|
|
if replay is not None:
|
|
connection.execute("COMMIT"); return replay
|
|
row = self._validate_lease(connection, job_id=job_id, runner_id=runner_id,
|
|
lease_token=lease_token, claim_key=claim_key,
|
|
runner_token=runner_token)
|
|
lease_until = now + self.lease_seconds
|
|
connection.execute("UPDATE remote_jobs SET lease_until=?,updated_at=? WHERE job_id=?",
|
|
(lease_until, now, job_id))
|
|
self._touch_runner(connection, runner_id, "heartbeat", now)
|
|
response = {"accepted": True, "job_id": job_id,
|
|
"lease_version": int(row["lease_version"]), "lease_until": lease_until}
|
|
self._record(connection, runner_id=runner_id, scope="heartbeat", key=key,
|
|
digest=digest, response=response, job_id=job_id,
|
|
lease_version=int(row["lease_version"]), now=now)
|
|
connection.execute("COMMIT"); return response
|
|
except Exception:
|
|
connection.execute("ROLLBACK"); raise
|
|
|
|
def event_for_submission(self, *, job_id: str, runner_id: str, lease_token: str,
|
|
claim_key: str, runner_token: bytes) -> Event:
|
|
row = self._validate_lease(self._conn(), job_id=job_id, runner_id=runner_id,
|
|
lease_token=lease_token, claim_key=claim_key,
|
|
runner_token=runner_token, allow_completed=True)
|
|
return Event.from_dict(json.loads(row["event_json"]))
|
|
|
|
def complete(self, *, job_id: str, runner_id: str, idempotency_key: str,
|
|
request_digest: str, decision_digest: str,
|
|
state_version_after: int) -> dict[str, Any]:
|
|
key = require_idempotency_key(idempotency_key); now = int(time.time())
|
|
connection = self._conn(); connection.execute("BEGIN IMMEDIATE")
|
|
try:
|
|
replay = self._replay(connection, runner_id, "decision", key, request_digest)
|
|
if replay is not None:
|
|
connection.execute("COMMIT"); return replay
|
|
row = connection.execute("SELECT * FROM remote_jobs WHERE job_id=?", (job_id,)).fetchone()
|
|
if row is None:
|
|
raise ConflictError("remote_job_not_found")
|
|
if row["status"] == "COMPLETED":
|
|
if row["decision_digest"] != decision_digest:
|
|
raise ConflictError("remote_decision_idempotency_conflict")
|
|
elif row["status"] == "LEASED" and row["runner_id"] == runner_id:
|
|
connection.execute(
|
|
"UPDATE remote_jobs SET status='COMPLETED',decision_digest=?,lease_until=NULL,updated_at=? WHERE job_id=?",
|
|
(decision_digest, now, job_id),
|
|
)
|
|
else:
|
|
raise ConflictError("remote_job_not_leased")
|
|
response = {"accepted": True, "job_id": job_id,
|
|
"state_version": state_version_after, "completed": True}
|
|
self._record(connection, runner_id=runner_id, scope="decision", key=key,
|
|
digest=request_digest, response=response, job_id=job_id,
|
|
lease_version=int(row["lease_version"]), now=now)
|
|
self._touch_runner(connection, runner_id, "decision", now)
|
|
connection.execute("COMMIT"); return response
|
|
except Exception:
|
|
connection.execute("ROLLBACK"); raise
|
|
|
|
def fail(self, *, job_id: str, runner_id: str, lease_token: str, claim_key: str,
|
|
category: str, idempotency_key: str, runner_token: bytes,
|
|
retry_after: int) -> dict[str, Any]:
|
|
if not isinstance(category, str) or not category or len(category) > 64 or not category.replace("_", "").isalnum():
|
|
raise ValidationError("invalid_failure_category")
|
|
key = require_idempotency_key(idempotency_key)
|
|
request = {"job_id": job_id, "lease_token": lease_token,
|
|
"claim_key": claim_key, "category": category}
|
|
digest = self._digest(request); now = int(time.time())
|
|
connection = self._conn(); connection.execute("BEGIN IMMEDIATE")
|
|
try:
|
|
replay = self._replay(connection, runner_id, "fail", key, digest)
|
|
if replay is not None:
|
|
connection.execute("COMMIT"); return replay
|
|
row = self._validate_lease(connection, job_id=job_id, runner_id=runner_id,
|
|
lease_token=lease_token, claim_key=claim_key,
|
|
runner_token=runner_token)
|
|
attempts = int(row["attempts"]) + 1
|
|
status = "DEAD_LETTER" if attempts > self.max_retries else "RETRY"
|
|
connection.execute(
|
|
"""
|
|
UPDATE remote_jobs SET status=?,attempts=?,failure_category=?,
|
|
next_attempt_at=?,runner_id=NULL,lease_until=NULL,updated_at=? WHERE job_id=?
|
|
""", (status, attempts, category, now + max(0, retry_after), now, job_id),
|
|
)
|
|
response = {"accepted": True, "job_id": job_id, "status": status,
|
|
"attempts": attempts}
|
|
self._record(connection, runner_id=runner_id, scope="fail", key=key,
|
|
digest=digest, response=response, job_id=job_id,
|
|
lease_version=int(row["lease_version"]), now=now)
|
|
self._touch_runner(connection, runner_id, "fail", now)
|
|
connection.execute("COMMIT"); return response
|
|
except Exception:
|
|
connection.execute("ROLLBACK"); raise
|
|
|
|
def health_snapshot(self, *, healthy_within_seconds: int = 300) -> dict[str, Any]:
|
|
now = int(time.time())
|
|
connection = self._conn()
|
|
rows = connection.execute(
|
|
"SELECT runner_id,last_seen_at,last_scope FROM remote_runners ORDER BY runner_id"
|
|
).fetchall()
|
|
counts = {
|
|
str(row["status"]): int(row["count"])
|
|
for row in connection.execute(
|
|
"SELECT status,COUNT(*) AS count FROM remote_jobs GROUP BY status"
|
|
).fetchall()
|
|
}
|
|
return {
|
|
"channel": "remote_codex",
|
|
"contract_commit": self.contract_commit,
|
|
"healthy_runners": sum(1 for row in rows if now - int(row["last_seen_at"]) <= healthy_within_seconds),
|
|
"runner_count": len(rows),
|
|
"job_counts": counts,
|
|
}
|