sysadmin
16cd871fbd
Address remaining PR #619 blockers on head 783ea88:
- Deduplicate legacy NULL-scope incident_links before normalizing to ''
- Fail closed when duplicate rows disagree on Gitea target
- Require non-empty expected_head_sha for PR assign and mutation
1060 lines
40 KiB
Python
1060 lines
40 KiB
Python
"""Control-plane DB substrate for multi-session coordination (#613).
|
||
|
||
Implements the durable coordination store described by
|
||
``docs/architecture/mcp-allocator-control-plane-observability-adr.md``:
|
||
|
||
* DB coordinates live concurrency (sessions, atomic assignment+lease,
|
||
heartbeats, terminal-lock index, events, ``incident_links``).
|
||
* Gitea remains the durable work record and the only assignable work unit
|
||
(``issue`` / ``pr`` — never raw Sentry/GlitchTip incidents).
|
||
* SQLite is the single-writer MVP backend; the API is backend-shaped so a
|
||
shared Postgres or single allocator daemon can replace the connection
|
||
layer later without changing call sites.
|
||
|
||
This module is the **substrate** for #600 (allocator policy/tool) and #612
|
||
(incident bridge). It does **not** implement ``gitea_allocate_next_work``
|
||
routing policy or provider adapters.
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import json
|
||
import os
|
||
import sqlite3
|
||
import threading
|
||
import time
|
||
import uuid
|
||
from contextlib import contextmanager
|
||
from dataclasses import dataclass
|
||
from datetime import datetime, timedelta, timezone
|
||
from typing import Any, Iterator, Sequence
|
||
|
||
SCHEMA_VERSION = 2
|
||
|
||
# Assignable work kinds only — raw monitoring incidents are never work items.
|
||
WORK_KINDS = frozenset({"issue", "pr"})
|
||
|
||
# Work item states that permanently revoke assignment/lease authority.
|
||
TERMINAL_WORK_STATES = frozenset({"merged", "closed"})
|
||
|
||
DEFAULT_LEASE_TTL_SECONDS = 4 * 3600
|
||
|
||
# Environment: path for SQLite MVP (single-writer).
|
||
DB_PATH_ENV = "GITEA_CONTROL_PLANE_DB"
|
||
DEFAULT_DB_PATH = os.path.expanduser("~/.cache/gitea-tools/control-plane/control_plane.sqlite3")
|
||
|
||
_SCHEMA_SQL = """
|
||
PRAGMA foreign_keys = ON;
|
||
|
||
CREATE TABLE IF NOT EXISTS schema_meta (
|
||
key TEXT PRIMARY KEY,
|
||
value TEXT NOT NULL
|
||
);
|
||
|
||
CREATE TABLE IF NOT EXISTS sessions (
|
||
session_id TEXT PRIMARY KEY,
|
||
role TEXT NOT NULL,
|
||
profile TEXT,
|
||
namespace TEXT,
|
||
pid INTEGER,
|
||
started_at TEXT NOT NULL,
|
||
last_heartbeat_at TEXT NOT NULL,
|
||
status TEXT NOT NULL DEFAULT 'active'
|
||
);
|
||
|
||
CREATE TABLE IF NOT EXISTS work_items (
|
||
work_item_id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||
remote TEXT NOT NULL,
|
||
org TEXT NOT NULL,
|
||
repo TEXT NOT NULL,
|
||
kind TEXT NOT NULL CHECK (kind IN ('issue', 'pr')),
|
||
number INTEGER NOT NULL,
|
||
state TEXT NOT NULL DEFAULT 'open',
|
||
priority INTEGER NOT NULL DEFAULT 0,
|
||
current_head_sha TEXT,
|
||
updated_at TEXT NOT NULL,
|
||
UNIQUE (remote, org, repo, kind, number)
|
||
);
|
||
|
||
CREATE TABLE IF NOT EXISTS leases (
|
||
lease_id TEXT PRIMARY KEY,
|
||
work_item_id INTEGER NOT NULL REFERENCES work_items(work_item_id),
|
||
session_id TEXT NOT NULL REFERENCES sessions(session_id),
|
||
role TEXT NOT NULL,
|
||
phase TEXT NOT NULL DEFAULT 'claimed',
|
||
expires_at TEXT NOT NULL,
|
||
heartbeat_at TEXT NOT NULL,
|
||
status TEXT NOT NULL DEFAULT 'active'
|
||
);
|
||
|
||
CREATE TABLE IF NOT EXISTS assignments (
|
||
assignment_id TEXT PRIMARY KEY,
|
||
work_item_id INTEGER NOT NULL REFERENCES work_items(work_item_id),
|
||
session_id TEXT NOT NULL REFERENCES sessions(session_id),
|
||
lease_id TEXT NOT NULL REFERENCES leases(lease_id),
|
||
allowed_actions TEXT NOT NULL,
|
||
forbidden_actions TEXT NOT NULL,
|
||
expected_head_sha TEXT,
|
||
role TEXT NOT NULL,
|
||
status TEXT NOT NULL DEFAULT 'active',
|
||
created_at TEXT NOT NULL
|
||
);
|
||
|
||
CREATE TABLE IF NOT EXISTS terminal_locks (
|
||
terminal_lock_id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||
remote TEXT NOT NULL,
|
||
org TEXT NOT NULL,
|
||
repo TEXT NOT NULL,
|
||
terminal_pr INTEGER NOT NULL,
|
||
review_id TEXT,
|
||
decision TEXT,
|
||
status TEXT NOT NULL DEFAULT 'active',
|
||
cleanup_state TEXT,
|
||
created_at TEXT NOT NULL,
|
||
UNIQUE (remote, org, repo, terminal_pr)
|
||
);
|
||
|
||
CREATE TABLE IF NOT EXISTS events (
|
||
event_id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||
work_item_id INTEGER REFERENCES work_items(work_item_id),
|
||
event_type TEXT NOT NULL,
|
||
message TEXT NOT NULL,
|
||
created_at TEXT NOT NULL
|
||
);
|
||
|
||
-- Provider-neutral incident ↔ Gitea link model (ADR §8–9). Not assignable work.
|
||
-- Optional scope fields are stored as '' (never NULL) so UNIQUE is NULL-safe.
|
||
CREATE TABLE IF NOT EXISTS incident_links (
|
||
link_id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||
provider TEXT NOT NULL,
|
||
provider_base_url TEXT NOT NULL DEFAULT '',
|
||
provider_org TEXT NOT NULL DEFAULT '',
|
||
provider_project TEXT NOT NULL DEFAULT '',
|
||
provider_issue_id TEXT NOT NULL,
|
||
provider_short_id TEXT,
|
||
provider_permalink TEXT,
|
||
fingerprint TEXT,
|
||
gitea_org TEXT NOT NULL,
|
||
gitea_repo TEXT NOT NULL,
|
||
gitea_issue_number INTEGER NOT NULL,
|
||
linked_pr_numbers TEXT,
|
||
first_seen TEXT,
|
||
last_seen TEXT,
|
||
event_count INTEGER,
|
||
status TEXT NOT NULL DEFAULT 'open',
|
||
release_resolved_at TEXT,
|
||
last_sync_at TEXT,
|
||
UNIQUE (provider, provider_base_url, provider_org, provider_project, provider_issue_id)
|
||
);
|
||
|
||
CREATE INDEX IF NOT EXISTS idx_leases_work_status ON leases(work_item_id, status);
|
||
CREATE INDEX IF NOT EXISTS idx_assignments_session ON assignments(session_id, status);
|
||
CREATE INDEX IF NOT EXISTS idx_incident_gitea ON incident_links(gitea_org, gitea_repo, gitea_issue_number);
|
||
"""
|
||
|
||
|
||
def _utc_now() -> datetime:
|
||
return datetime.now(timezone.utc)
|
||
|
||
|
||
def _ts(dt: datetime | None = None) -> str:
|
||
value = dt or _utc_now()
|
||
if value.tzinfo is None:
|
||
value = value.replace(tzinfo=timezone.utc)
|
||
return value.astimezone(timezone.utc).replace(microsecond=0).isoformat().replace("+00:00", "Z")
|
||
|
||
|
||
def _parse_ts(value: str | None) -> datetime | None:
|
||
if not value:
|
||
return None
|
||
text = value.strip()
|
||
if text.endswith("Z"):
|
||
text = text[:-1] + "+00:00"
|
||
try:
|
||
return datetime.fromisoformat(text)
|
||
except ValueError:
|
||
return None
|
||
|
||
|
||
def _norm_scope(value: str | None) -> str:
|
||
"""Normalize optional incident-link scope keys for NULL-safe uniqueness.
|
||
|
||
Empty / whitespace / None all become '' so SQLite UNIQUE treats them as one
|
||
canonical key component (SQLite treats multiple NULLs as distinct).
|
||
"""
|
||
if value is None:
|
||
return ""
|
||
return str(value).strip()
|
||
|
||
|
||
def default_db_path() -> str:
|
||
raw = (os.environ.get(DB_PATH_ENV) or DEFAULT_DB_PATH).strip()
|
||
return raw or DEFAULT_DB_PATH
|
||
|
||
|
||
class ControlPlaneError(RuntimeError):
|
||
"""Base error for control-plane substrate failures."""
|
||
|
||
|
||
class InvalidWorkKindError(ControlPlaneError):
|
||
"""Raised when a non-assignable work kind is requested."""
|
||
|
||
|
||
class LeaseRequiredError(ControlPlaneError):
|
||
"""Raised when a mutation is attempted without a valid assignment/lease."""
|
||
|
||
|
||
class ForeignLeaseError(ControlPlaneError):
|
||
"""Raised when another session holds the active lease."""
|
||
|
||
|
||
@dataclass(frozen=True)
|
||
class AssignmentResult:
|
||
"""Result of an atomic assign+lease transaction."""
|
||
|
||
outcome: str # assigned | wait | no_safe_work
|
||
assignment_id: str | None = None
|
||
lease_id: str | None = None
|
||
session_id: str | None = None
|
||
role: str | None = None
|
||
work_kind: str | None = None
|
||
work_number: int | None = None
|
||
remote: str | None = None
|
||
org: str | None = None
|
||
repo: str | None = None
|
||
expected_head_sha: str | None = None
|
||
allowed_actions: tuple[str, ...] = ()
|
||
forbidden_actions: tuple[str, ...] = ()
|
||
expires_at: str | None = None
|
||
owner_session_id: str | None = None
|
||
reason: str = ""
|
||
|
||
def as_dict(self) -> dict[str, Any]:
|
||
return {
|
||
"outcome": self.outcome,
|
||
"assignment_id": self.assignment_id,
|
||
"lease_id": self.lease_id,
|
||
"session_id": self.session_id,
|
||
"role": self.role,
|
||
"work_kind": self.work_kind,
|
||
"work_number": self.work_number,
|
||
"remote": self.remote,
|
||
"org": self.org,
|
||
"repo": self.repo,
|
||
"expected_head_sha": self.expected_head_sha,
|
||
"allowed_actions": list(self.allowed_actions),
|
||
"forbidden_actions": list(self.forbidden_actions),
|
||
"expires_at": self.expires_at,
|
||
"owner_session_id": self.owner_session_id,
|
||
"reason": self.reason,
|
||
}
|
||
|
||
|
||
class ControlPlaneDB:
|
||
"""SQLite single-writer MVP control-plane store.
|
||
|
||
Use one process as writer for multi-session safety claims, or migrate to
|
||
Postgres / a single allocator daemon for multi-host production (ADR §6).
|
||
"""
|
||
|
||
def __init__(self, db_path: str | None = None) -> None:
|
||
self.db_path = (db_path or default_db_path()).strip()
|
||
parent = os.path.dirname(self.db_path)
|
||
if parent:
|
||
os.makedirs(parent, mode=0o700, exist_ok=True)
|
||
self._lock = threading.RLock()
|
||
self._init_schema()
|
||
|
||
def _connect(self) -> sqlite3.Connection:
|
||
# Default isolation (DEFERRED) so explicit BEGIN IMMEDIATE works.
|
||
conn = sqlite3.connect(self.db_path, timeout=30)
|
||
conn.row_factory = sqlite3.Row
|
||
conn.execute("PRAGMA foreign_keys = ON")
|
||
# Serialize writers even across threads in one process.
|
||
conn.execute("PRAGMA journal_mode = WAL")
|
||
return conn
|
||
|
||
@contextmanager
|
||
def _tx(self, immediate: bool = True) -> Iterator[sqlite3.Connection]:
|
||
with self._lock:
|
||
conn = self._connect()
|
||
try:
|
||
if immediate:
|
||
conn.execute("BEGIN IMMEDIATE")
|
||
else:
|
||
conn.execute("BEGIN")
|
||
yield conn
|
||
conn.commit()
|
||
except Exception:
|
||
try:
|
||
conn.rollback()
|
||
except sqlite3.Error:
|
||
pass
|
||
raise
|
||
finally:
|
||
conn.close()
|
||
|
||
def _init_schema(self) -> None:
|
||
# executescript auto-commits; run schema outside an open txn, then meta.
|
||
with self._lock:
|
||
conn = self._connect()
|
||
try:
|
||
conn.executescript(_SCHEMA_SQL)
|
||
self._migrate_incident_links_null_scope(conn)
|
||
conn.execute(
|
||
"INSERT OR REPLACE INTO schema_meta(key, value) VALUES (?, ?)",
|
||
("schema_version", str(SCHEMA_VERSION)),
|
||
)
|
||
conn.execute(
|
||
"INSERT OR REPLACE INTO schema_meta(key, value) VALUES (?, ?)",
|
||
(
|
||
"architecture",
|
||
"DB coordinates; Gitea records; Sentry/GlitchTip observe; "
|
||
"bridge is only path from observations to Gitea work",
|
||
),
|
||
)
|
||
conn.commit()
|
||
finally:
|
||
conn.close()
|
||
|
||
def _migrate_incident_links_null_scope(self, conn: sqlite3.Connection) -> None:
|
||
"""Collapse legacy NULL-scope duplicates, then normalize to '' (#619 RC2).
|
||
|
||
Order is mandatory: SQLite UNIQUE treats multiple NULLs as distinct, so
|
||
normalizing NULL→'' first can hit the UNIQUE constraint and abort
|
||
migration. Deduplicate under the *normalized* key first, fail closed if
|
||
Gitea targets conflict within a group, then coerce NULLs to ''.
|
||
"""
|
||
cols = {
|
||
row[1]
|
||
for row in conn.execute("PRAGMA table_info(incident_links)").fetchall()
|
||
}
|
||
if not cols:
|
||
return
|
||
|
||
rows = conn.execute(
|
||
"""
|
||
SELECT link_id, provider,
|
||
IFNULL(provider_base_url, '') AS base_url,
|
||
IFNULL(provider_org, '') AS p_org,
|
||
IFNULL(provider_project, '') AS p_project,
|
||
provider_issue_id,
|
||
gitea_org, gitea_repo, gitea_issue_number
|
||
FROM incident_links
|
||
ORDER BY link_id ASC
|
||
"""
|
||
).fetchall()
|
||
|
||
groups: dict[tuple[str, str, str, str, str], list[sqlite3.Row]] = {}
|
||
for row in rows:
|
||
key = (
|
||
str(row["provider"]),
|
||
str(row["base_url"]),
|
||
str(row["p_org"]),
|
||
str(row["p_project"]),
|
||
str(row["provider_issue_id"]),
|
||
)
|
||
groups.setdefault(key, []).append(row)
|
||
|
||
to_delete: list[int] = []
|
||
for key, members in groups.items():
|
||
if len(members) < 2:
|
||
continue
|
||
# Canonical identity for observation links: Gitea issue target.
|
||
targets = {
|
||
(
|
||
str(m["gitea_org"] or ""),
|
||
str(m["gitea_repo"] or ""),
|
||
int(m["gitea_issue_number"]),
|
||
)
|
||
for m in members
|
||
}
|
||
if len(targets) > 1:
|
||
raise ControlPlaneError(
|
||
"incident_links migration fail closed: conflicting Gitea "
|
||
f"targets for provider key {key!r}: {sorted(targets)}"
|
||
)
|
||
keep_id = int(members[0]["link_id"]) # lowest link_id (ORDER BY)
|
||
for m in members[1:]:
|
||
to_delete.append(int(m["link_id"]))
|
||
|
||
for link_id in to_delete:
|
||
conn.execute("DELETE FROM incident_links WHERE link_id = ?", (link_id,))
|
||
|
||
# Safe only after dedupe: normalize NULL scope fields to empty string.
|
||
for col in ("provider_base_url", "provider_org", "provider_project"):
|
||
if col in cols:
|
||
conn.execute(
|
||
f"UPDATE incident_links SET {col} = '' WHERE {col} IS NULL"
|
||
)
|
||
|
||
# ── sessions ──────────────────────────────────────────────────────────
|
||
|
||
def upsert_session(
|
||
self,
|
||
*,
|
||
session_id: str,
|
||
role: str,
|
||
profile: str | None = None,
|
||
namespace: str | None = None,
|
||
pid: int | None = None,
|
||
status: str = "active",
|
||
) -> dict[str, Any]:
|
||
now = _ts()
|
||
with self._tx() as conn:
|
||
existing = conn.execute(
|
||
"SELECT session_id FROM sessions WHERE session_id = ?",
|
||
(session_id,),
|
||
).fetchone()
|
||
if existing:
|
||
conn.execute(
|
||
"""
|
||
UPDATE sessions
|
||
SET role = ?, profile = ?, namespace = ?, pid = ?,
|
||
last_heartbeat_at = ?, status = ?
|
||
WHERE session_id = ?
|
||
""",
|
||
(role, profile, namespace, pid, now, status, session_id),
|
||
)
|
||
else:
|
||
conn.execute(
|
||
"""
|
||
INSERT INTO sessions(
|
||
session_id, role, profile, namespace, pid,
|
||
started_at, last_heartbeat_at, status
|
||
) VALUES (?, ?, ?, ?, ?, ?, ?, ?)
|
||
""",
|
||
(session_id, role, profile, namespace, pid, now, now, status),
|
||
)
|
||
row = conn.execute(
|
||
"SELECT * FROM sessions WHERE session_id = ?",
|
||
(session_id,),
|
||
).fetchone()
|
||
return dict(row)
|
||
|
||
def heartbeat_session(self, session_id: str) -> None:
|
||
with self._tx() as conn:
|
||
conn.execute(
|
||
"UPDATE sessions SET last_heartbeat_at = ? WHERE session_id = ?",
|
||
(_ts(), session_id),
|
||
)
|
||
|
||
# ── work items ────────────────────────────────────────────────────────
|
||
|
||
def upsert_work_item(
|
||
self,
|
||
*,
|
||
remote: str,
|
||
org: str,
|
||
repo: str,
|
||
kind: str,
|
||
number: int,
|
||
state: str = "open",
|
||
priority: int = 0,
|
||
current_head_sha: str | None = None,
|
||
) -> int:
|
||
kind_norm = (kind or "").strip().lower()
|
||
if kind_norm not in WORK_KINDS:
|
||
raise InvalidWorkKindError(
|
||
f"work kind '{kind}' is not assignable; only {sorted(WORK_KINDS)} "
|
||
f"are allowed (raw monitoring incidents are never work items)"
|
||
)
|
||
now = _ts()
|
||
with self._tx() as conn:
|
||
conn.execute(
|
||
"""
|
||
INSERT INTO work_items(
|
||
remote, org, repo, kind, number, state, priority,
|
||
current_head_sha, updated_at
|
||
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
|
||
ON CONFLICT(remote, org, repo, kind, number) DO UPDATE SET
|
||
state = excluded.state,
|
||
priority = excluded.priority,
|
||
current_head_sha = excluded.current_head_sha,
|
||
updated_at = excluded.updated_at
|
||
""",
|
||
(
|
||
remote,
|
||
org,
|
||
repo,
|
||
kind_norm,
|
||
int(number),
|
||
state,
|
||
int(priority),
|
||
current_head_sha,
|
||
now,
|
||
),
|
||
)
|
||
row = conn.execute(
|
||
"""
|
||
SELECT work_item_id FROM work_items
|
||
WHERE remote = ? AND org = ? AND repo = ? AND kind = ? AND number = ?
|
||
""",
|
||
(remote, org, repo, kind_norm, int(number)),
|
||
).fetchone()
|
||
return int(row["work_item_id"])
|
||
|
||
# ── atomic assign + lease ─────────────────────────────────────────────
|
||
|
||
def assign_and_lease(
|
||
self,
|
||
*,
|
||
session_id: str,
|
||
role: str,
|
||
remote: str,
|
||
org: str,
|
||
repo: str,
|
||
kind: str,
|
||
number: int,
|
||
expected_head_sha: str | None = None,
|
||
allowed_actions: Sequence[str] | None = None,
|
||
forbidden_actions: Sequence[str] | None = None,
|
||
lease_ttl_seconds: int = DEFAULT_LEASE_TTL_SECONDS,
|
||
phase: str = "claimed",
|
||
now: datetime | None = None,
|
||
) -> AssignmentResult:
|
||
"""Atomically create assignment + lease for one Gitea work item.
|
||
|
||
Concurrent sessions cannot both receive the same open work item.
|
||
Returns ``wait`` with owner_session_id when a live foreign lease exists.
|
||
"""
|
||
kind_norm = (kind or "").strip().lower()
|
||
if kind_norm not in WORK_KINDS:
|
||
raise InvalidWorkKindError(
|
||
f"cannot assign kind '{kind}'; only {sorted(WORK_KINDS)}"
|
||
)
|
||
head_pin = (expected_head_sha or "").strip()
|
||
if kind_norm == "pr" and not head_pin:
|
||
raise LeaseRequiredError(
|
||
f"pr#{int(number)} assignment requires non-empty expected_head_sha pin"
|
||
)
|
||
|
||
allowed = tuple(allowed_actions or ("implement", "comment"))
|
||
forbidden = tuple(
|
||
forbidden_actions
|
||
or ("approve", "merge", "self_select_without_assignment")
|
||
)
|
||
moment = now or _utc_now()
|
||
now_s = _ts(moment)
|
||
expires = _ts(moment + timedelta(seconds=int(lease_ttl_seconds)))
|
||
|
||
with self._tx(immediate=True) as conn:
|
||
# Ensure session exists
|
||
sess = conn.execute(
|
||
"SELECT session_id FROM sessions WHERE session_id = ?",
|
||
(session_id,),
|
||
).fetchone()
|
||
if not sess:
|
||
conn.execute(
|
||
"""
|
||
INSERT INTO sessions(
|
||
session_id, role, profile, namespace, pid,
|
||
started_at, last_heartbeat_at, status
|
||
) VALUES (?, ?, NULL, NULL, ?, ?, ?, 'active')
|
||
""",
|
||
(session_id, role, os.getpid(), now_s, now_s),
|
||
)
|
||
|
||
# Upsert work item
|
||
conn.execute(
|
||
"""
|
||
INSERT INTO work_items(
|
||
remote, org, repo, kind, number, state, priority,
|
||
current_head_sha, updated_at
|
||
) VALUES (?, ?, ?, ?, ?, 'open', 0, ?, ?)
|
||
ON CONFLICT(remote, org, repo, kind, number) DO UPDATE SET
|
||
current_head_sha = COALESCE(excluded.current_head_sha, work_items.current_head_sha),
|
||
updated_at = excluded.updated_at
|
||
""",
|
||
(remote, org, repo, kind_norm, int(number), expected_head_sha, now_s),
|
||
)
|
||
work = conn.execute(
|
||
"""
|
||
SELECT * FROM work_items
|
||
WHERE remote = ? AND org = ? AND repo = ? AND kind = ? AND number = ?
|
||
""",
|
||
(remote, org, repo, kind_norm, int(number)),
|
||
).fetchone()
|
||
work_item_id = int(work["work_item_id"])
|
||
state = (work["state"] or "open").lower()
|
||
if state in ("merged", "closed"):
|
||
return AssignmentResult(
|
||
outcome="no_safe_work",
|
||
reason=f"work item {kind_norm}#{number} is {state}; never assign",
|
||
)
|
||
|
||
# Expire stale leases on this work item first
|
||
self._expire_stale_leases_conn(conn, work_item_id=work_item_id, now_s=now_s)
|
||
|
||
active = conn.execute(
|
||
"""
|
||
SELECT * FROM leases
|
||
WHERE work_item_id = ? AND status = 'active'
|
||
ORDER BY expires_at DESC
|
||
LIMIT 1
|
||
""",
|
||
(work_item_id,),
|
||
).fetchone()
|
||
if active:
|
||
owner = active["session_id"]
|
||
if owner == session_id:
|
||
# Owner-resume: refresh heartbeat and return existing assignment
|
||
conn.execute(
|
||
"""
|
||
UPDATE leases
|
||
SET heartbeat_at = ?, expires_at = ?, phase = ?
|
||
WHERE lease_id = ?
|
||
""",
|
||
(now_s, expires, phase, active["lease_id"]),
|
||
)
|
||
asn = conn.execute(
|
||
"""
|
||
SELECT * FROM assignments
|
||
WHERE lease_id = ? AND status = 'active'
|
||
ORDER BY created_at DESC LIMIT 1
|
||
""",
|
||
(active["lease_id"],),
|
||
).fetchone()
|
||
if asn:
|
||
return AssignmentResult(
|
||
outcome="assigned",
|
||
assignment_id=asn["assignment_id"],
|
||
lease_id=active["lease_id"],
|
||
session_id=session_id,
|
||
role=asn["role"],
|
||
work_kind=kind_norm,
|
||
work_number=int(number),
|
||
remote=remote,
|
||
org=org,
|
||
repo=repo,
|
||
expected_head_sha=asn["expected_head_sha"],
|
||
allowed_actions=tuple(json.loads(asn["allowed_actions"])),
|
||
forbidden_actions=tuple(json.loads(asn["forbidden_actions"])),
|
||
expires_at=expires,
|
||
owner_session_id=session_id,
|
||
reason="owner-resume: refreshed existing lease",
|
||
)
|
||
return AssignmentResult(
|
||
outcome="wait",
|
||
owner_session_id=owner,
|
||
reason=(
|
||
f"foreign active lease held by session {owner} on "
|
||
f"{kind_norm}#{number}"
|
||
),
|
||
)
|
||
|
||
lease_id = f"lease-{uuid.uuid4().hex[:16]}"
|
||
assignment_id = f"asn-{uuid.uuid4().hex[:16]}"
|
||
conn.execute(
|
||
"""
|
||
INSERT INTO leases(
|
||
lease_id, work_item_id, session_id, role, phase,
|
||
expires_at, heartbeat_at, status
|
||
) VALUES (?, ?, ?, ?, ?, ?, ?, 'active')
|
||
""",
|
||
(lease_id, work_item_id, session_id, role, phase, expires, now_s),
|
||
)
|
||
conn.execute(
|
||
"""
|
||
INSERT INTO assignments(
|
||
assignment_id, work_item_id, session_id, lease_id,
|
||
allowed_actions, forbidden_actions, expected_head_sha,
|
||
role, status, created_at
|
||
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, 'active', ?)
|
||
""",
|
||
(
|
||
assignment_id,
|
||
work_item_id,
|
||
session_id,
|
||
lease_id,
|
||
json.dumps(list(allowed)),
|
||
json.dumps(list(forbidden)),
|
||
expected_head_sha or work["current_head_sha"],
|
||
role,
|
||
now_s,
|
||
),
|
||
)
|
||
conn.execute(
|
||
"""
|
||
INSERT INTO events(work_item_id, event_type, message, created_at)
|
||
VALUES (?, 'assigned', ?, ?)
|
||
""",
|
||
(
|
||
work_item_id,
|
||
f"session {session_id} assigned {kind_norm}#{number} lease={lease_id}",
|
||
now_s,
|
||
),
|
||
)
|
||
return AssignmentResult(
|
||
outcome="assigned",
|
||
assignment_id=assignment_id,
|
||
lease_id=lease_id,
|
||
session_id=session_id,
|
||
role=role,
|
||
work_kind=kind_norm,
|
||
work_number=int(number),
|
||
remote=remote,
|
||
org=org,
|
||
repo=repo,
|
||
expected_head_sha=expected_head_sha or work["current_head_sha"],
|
||
allowed_actions=allowed,
|
||
forbidden_actions=forbidden,
|
||
expires_at=expires,
|
||
owner_session_id=session_id,
|
||
reason="atomic assign+lease created",
|
||
)
|
||
|
||
def _expire_stale_leases_conn(
|
||
self,
|
||
conn: sqlite3.Connection,
|
||
*,
|
||
work_item_id: int | None = None,
|
||
now_s: str | None = None,
|
||
) -> int:
|
||
now_s = now_s or _ts()
|
||
if work_item_id is None:
|
||
rows = conn.execute(
|
||
"SELECT lease_id FROM leases WHERE status = 'active' AND expires_at <= ?",
|
||
(now_s,),
|
||
).fetchall()
|
||
else:
|
||
rows = conn.execute(
|
||
"""
|
||
SELECT lease_id FROM leases
|
||
WHERE status = 'active' AND expires_at <= ? AND work_item_id = ?
|
||
""",
|
||
(now_s, work_item_id),
|
||
).fetchall()
|
||
for row in rows:
|
||
lid = row["lease_id"]
|
||
conn.execute(
|
||
"UPDATE leases SET status = 'expired' WHERE lease_id = ?",
|
||
(lid,),
|
||
)
|
||
conn.execute(
|
||
"UPDATE assignments SET status = 'expired' WHERE lease_id = ?",
|
||
(lid,),
|
||
)
|
||
return len(rows)
|
||
|
||
def expire_stale_leases(self) -> int:
|
||
with self._tx() as conn:
|
||
return self._expire_stale_leases_conn(conn)
|
||
|
||
def heartbeat_lease(self, lease_id: str, *, session_id: str) -> dict[str, Any]:
|
||
now_s = _ts()
|
||
with self._tx() as conn:
|
||
row = conn.execute(
|
||
"SELECT * FROM leases WHERE lease_id = ?",
|
||
(lease_id,),
|
||
).fetchone()
|
||
if not row:
|
||
raise ControlPlaneError(f"unknown lease_id {lease_id}")
|
||
if row["session_id"] != session_id:
|
||
raise ForeignLeaseError(
|
||
f"lease {lease_id} owned by {row['session_id']}, not {session_id}"
|
||
)
|
||
if row["status"] != "active":
|
||
raise ControlPlaneError(f"lease {lease_id} status is {row['status']}")
|
||
exp = _parse_ts(row["expires_at"])
|
||
if exp and exp <= _utc_now():
|
||
conn.execute(
|
||
"UPDATE leases SET status = 'expired' WHERE lease_id = ?",
|
||
(lease_id,),
|
||
)
|
||
raise ControlPlaneError(f"lease {lease_id} already expired")
|
||
# Extend TTL on heartbeat
|
||
new_exp = _ts(_utc_now() + timedelta(seconds=DEFAULT_LEASE_TTL_SECONDS))
|
||
conn.execute(
|
||
"""
|
||
UPDATE leases SET heartbeat_at = ?, expires_at = ?
|
||
WHERE lease_id = ?
|
||
""",
|
||
(now_s, new_exp, lease_id),
|
||
)
|
||
conn.execute(
|
||
"UPDATE sessions SET last_heartbeat_at = ? WHERE session_id = ?",
|
||
(now_s, session_id),
|
||
)
|
||
return {
|
||
"lease_id": lease_id,
|
||
"heartbeat_at": now_s,
|
||
"expires_at": new_exp,
|
||
"session_id": session_id,
|
||
}
|
||
|
||
def release_lease(self, lease_id: str, *, session_id: str) -> None:
|
||
with self._tx() as conn:
|
||
row = conn.execute(
|
||
"SELECT * FROM leases WHERE lease_id = ?",
|
||
(lease_id,),
|
||
).fetchone()
|
||
if not row:
|
||
return
|
||
if row["session_id"] != session_id:
|
||
raise ForeignLeaseError(
|
||
f"cannot release lease {lease_id} owned by {row['session_id']}"
|
||
)
|
||
conn.execute(
|
||
"UPDATE leases SET status = 'released' WHERE lease_id = ?",
|
||
(lease_id,),
|
||
)
|
||
conn.execute(
|
||
"UPDATE assignments SET status = 'released' WHERE lease_id = ?",
|
||
(lease_id,),
|
||
)
|
||
conn.execute(
|
||
"""
|
||
INSERT INTO events(work_item_id, event_type, message, created_at)
|
||
VALUES (?, 'lease_released', ?, ?)
|
||
""",
|
||
(row["work_item_id"], f"session {session_id} released {lease_id}", _ts()),
|
||
)
|
||
|
||
def require_valid_assignment(
|
||
self,
|
||
*,
|
||
session_id: str,
|
||
remote: str,
|
||
org: str,
|
||
repo: str,
|
||
kind: str,
|
||
number: int,
|
||
action: str,
|
||
) -> dict[str, Any]:
|
||
"""Gate a mutation: require active assignment+lease for this session/work.
|
||
|
||
Fail-closed when the work item is terminal (merged/closed) or when the
|
||
assignment's expected_head_sha no longer matches the work item head
|
||
(stale-head after assignment time).
|
||
"""
|
||
kind_norm = (kind or "").strip().lower()
|
||
if kind_norm not in WORK_KINDS:
|
||
raise InvalidWorkKindError(f"invalid kind '{kind}'")
|
||
with self._tx(immediate=False) as conn:
|
||
work = conn.execute(
|
||
"""
|
||
SELECT * FROM work_items
|
||
WHERE remote = ? AND org = ? AND repo = ? AND kind = ? AND number = ?
|
||
""",
|
||
(remote, org, repo, kind_norm, int(number)),
|
||
).fetchone()
|
||
if not work:
|
||
raise LeaseRequiredError(
|
||
f"no work_item for {kind_norm}#{number}; assign first"
|
||
)
|
||
work_state = (work["state"] or "").strip().lower()
|
||
if work_state in TERMINAL_WORK_STATES:
|
||
raise LeaseRequiredError(
|
||
f"work item {kind_norm}#{number} is terminal state "
|
||
f"'{work_state}'; assignment no longer authorizes mutations"
|
||
)
|
||
self._expire_stale_leases_conn(conn, work_item_id=int(work["work_item_id"]))
|
||
asn = conn.execute(
|
||
"""
|
||
SELECT a.*, l.status AS lease_status, l.expires_at, l.lease_id
|
||
FROM assignments a
|
||
JOIN leases l ON l.lease_id = a.lease_id
|
||
WHERE a.work_item_id = ?
|
||
AND a.session_id = ?
|
||
AND a.status = 'active'
|
||
AND l.status = 'active'
|
||
ORDER BY a.created_at DESC
|
||
LIMIT 1
|
||
""",
|
||
(int(work["work_item_id"]), session_id),
|
||
).fetchone()
|
||
if not asn:
|
||
raise LeaseRequiredError(
|
||
f"session {session_id} has no active assignment/lease for "
|
||
f"{kind_norm}#{number}"
|
||
)
|
||
exp = _parse_ts(asn["expires_at"])
|
||
if exp and exp <= _utc_now():
|
||
raise LeaseRequiredError(f"lease {asn['lease_id']} expired")
|
||
# Head pin: PRs require a non-empty expected_head_sha (fail closed).
|
||
assigned_head = (asn["expected_head_sha"] or "").strip()
|
||
current_head = (work["current_head_sha"] or "").strip()
|
||
if kind_norm == "pr" and not assigned_head:
|
||
raise LeaseRequiredError(
|
||
f"assignment {asn['assignment_id']} for pr#{number} has no "
|
||
f"expected_head_sha pin; PR mutations require a head pin"
|
||
)
|
||
# Stale-head: pinned expected_head_sha must still match live work item.
|
||
if assigned_head and assigned_head != current_head:
|
||
raise LeaseRequiredError(
|
||
f"assignment {asn['assignment_id']} stale head: expected "
|
||
f"{assigned_head!r} but work item head is {current_head!r}"
|
||
)
|
||
allowed = json.loads(asn["allowed_actions"])
|
||
forbidden = json.loads(asn["forbidden_actions"])
|
||
if action in forbidden:
|
||
raise LeaseRequiredError(
|
||
f"action '{action}' is forbidden on assignment {asn['assignment_id']}"
|
||
)
|
||
if allowed and action not in allowed and action != "heartbeat":
|
||
raise LeaseRequiredError(
|
||
f"action '{action}' not in allowed_actions {allowed}"
|
||
)
|
||
return dict(asn)
|
||
|
||
# ── terminal locks ────────────────────────────────────────────────────
|
||
|
||
def set_terminal_lock(
|
||
self,
|
||
*,
|
||
remote: str,
|
||
org: str,
|
||
repo: str,
|
||
terminal_pr: int,
|
||
review_id: str | None = None,
|
||
decision: str | None = None,
|
||
status: str = "active",
|
||
cleanup_state: str | None = None,
|
||
) -> dict[str, Any]:
|
||
now_s = _ts()
|
||
with self._tx() as conn:
|
||
conn.execute(
|
||
"""
|
||
INSERT INTO terminal_locks(
|
||
remote, org, repo, terminal_pr, review_id, decision,
|
||
status, cleanup_state, created_at
|
||
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
|
||
ON CONFLICT(remote, org, repo, terminal_pr) DO UPDATE SET
|
||
review_id = excluded.review_id,
|
||
decision = excluded.decision,
|
||
status = excluded.status,
|
||
cleanup_state = excluded.cleanup_state
|
||
""",
|
||
(
|
||
remote,
|
||
org,
|
||
repo,
|
||
int(terminal_pr),
|
||
review_id,
|
||
decision,
|
||
status,
|
||
cleanup_state,
|
||
now_s,
|
||
),
|
||
)
|
||
row = conn.execute(
|
||
"""
|
||
SELECT * FROM terminal_locks
|
||
WHERE remote = ? AND org = ? AND repo = ? AND terminal_pr = ?
|
||
""",
|
||
(remote, org, repo, int(terminal_pr)),
|
||
).fetchone()
|
||
return dict(row)
|
||
|
||
def get_active_terminal_lock(
|
||
self, *, remote: str, org: str, repo: str
|
||
) -> dict[str, Any] | None:
|
||
with self._tx(immediate=False) as conn:
|
||
row = conn.execute(
|
||
"""
|
||
SELECT * FROM terminal_locks
|
||
WHERE remote = ? AND org = ? AND repo = ? AND status = 'active'
|
||
ORDER BY created_at DESC LIMIT 1
|
||
""",
|
||
(remote, org, repo),
|
||
).fetchone()
|
||
return dict(row) if row else None
|
||
|
||
# ── incident_links (provider-neutral; not assignable) ─────────────────
|
||
|
||
def upsert_incident_link(
|
||
self,
|
||
*,
|
||
provider: str,
|
||
provider_issue_id: str,
|
||
gitea_org: str,
|
||
gitea_repo: str,
|
||
gitea_issue_number: int,
|
||
provider_base_url: str | None = None,
|
||
provider_org: str | None = None,
|
||
provider_project: str | None = None,
|
||
provider_short_id: str | None = None,
|
||
provider_permalink: str | None = None,
|
||
fingerprint: str | None = None,
|
||
linked_pr_numbers: Sequence[int] | None = None,
|
||
first_seen: str | None = None,
|
||
last_seen: str | None = None,
|
||
event_count: int | None = None,
|
||
status: str = "open",
|
||
) -> dict[str, Any]:
|
||
now_s = _ts()
|
||
# Canonical key: never store NULL in UNIQUE scope columns.
|
||
base_url = _norm_scope(provider_base_url)
|
||
p_org = _norm_scope(provider_org)
|
||
p_project = _norm_scope(provider_project)
|
||
with self._tx() as conn:
|
||
conn.execute(
|
||
"""
|
||
INSERT INTO incident_links(
|
||
provider, provider_base_url, provider_org, provider_project,
|
||
provider_issue_id, provider_short_id, provider_permalink,
|
||
fingerprint, gitea_org, gitea_repo, gitea_issue_number,
|
||
linked_pr_numbers, first_seen, last_seen, event_count,
|
||
status, last_sync_at
|
||
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
||
ON CONFLICT(provider, provider_base_url, provider_org, provider_project, provider_issue_id)
|
||
DO UPDATE SET
|
||
gitea_issue_number = excluded.gitea_issue_number,
|
||
gitea_org = excluded.gitea_org,
|
||
gitea_repo = excluded.gitea_repo,
|
||
provider_permalink = COALESCE(excluded.provider_permalink, incident_links.provider_permalink),
|
||
fingerprint = COALESCE(excluded.fingerprint, incident_links.fingerprint),
|
||
linked_pr_numbers = COALESCE(excluded.linked_pr_numbers, incident_links.linked_pr_numbers),
|
||
last_seen = COALESCE(excluded.last_seen, incident_links.last_seen),
|
||
event_count = COALESCE(excluded.event_count, incident_links.event_count),
|
||
status = excluded.status,
|
||
last_sync_at = excluded.last_sync_at
|
||
""",
|
||
(
|
||
provider,
|
||
base_url,
|
||
p_org,
|
||
p_project,
|
||
provider_issue_id,
|
||
provider_short_id,
|
||
provider_permalink,
|
||
fingerprint,
|
||
gitea_org,
|
||
gitea_repo,
|
||
int(gitea_issue_number),
|
||
json.dumps(list(linked_pr_numbers or [])),
|
||
first_seen,
|
||
last_seen,
|
||
event_count,
|
||
status,
|
||
now_s,
|
||
),
|
||
)
|
||
row = conn.execute(
|
||
"""
|
||
SELECT * FROM incident_links
|
||
WHERE provider = ?
|
||
AND provider_base_url = ?
|
||
AND provider_org = ?
|
||
AND provider_project = ?
|
||
AND provider_issue_id = ?
|
||
""",
|
||
(provider, base_url, p_org, p_project, provider_issue_id),
|
||
).fetchone()
|
||
return dict(row)
|
||
|
||
def get_incident_link_for_gitea_issue(
|
||
self, *, gitea_org: str, gitea_repo: str, gitea_issue_number: int
|
||
) -> dict[str, Any] | None:
|
||
with self._tx(immediate=False) as conn:
|
||
row = conn.execute(
|
||
"""
|
||
SELECT * FROM incident_links
|
||
WHERE gitea_org = ? AND gitea_repo = ? AND gitea_issue_number = ?
|
||
ORDER BY link_id DESC LIMIT 1
|
||
""",
|
||
(gitea_org, gitea_repo, int(gitea_issue_number)),
|
||
).fetchone()
|
||
return dict(row) if row else None
|