sysadmin
780ef0b1be
Make active leases queryable workflow state with canonical adopt, release, expire/reclaim, and abandon operations on the #613 control-plane DB substrate. - lease_lifecycle.py: policy, proof gates, safe_next_action vocabulary - control_plane_db: schema v3 provenance columns + list/inspect/adopt/abandon - MCP tools: gitea_*_workflow_lease(s) for list/inspect/adopt/release/expire/abandon/reclaim - issue_lock_store: sanctioned reclaim of expired author locks (dead pid / missing wt) - Tests cover lifecycle, foreign steal refusal, allocator compatibility, merger handoff
1753 lines
66 KiB
Python
1753 lines
66 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), #601
|
||
(first-class lease lifecycle), 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 = 3
|
||
|
||
# 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)
|
||
self._migrate_lease_lifecycle_columns(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()
|
||
|
||
# Observation fields that must agree before collapsing legacy duplicates.
|
||
# Deleting a row would drop any of these; silent loss is forbidden (#619 RC3).
|
||
_INCIDENT_LINK_OBS_COMPARE_FIELDS: tuple[str, ...] = (
|
||
"provider_short_id",
|
||
"provider_permalink",
|
||
"fingerprint",
|
||
"gitea_org",
|
||
"gitea_repo",
|
||
"gitea_issue_number",
|
||
"linked_pr_numbers",
|
||
"first_seen",
|
||
"last_seen",
|
||
"event_count",
|
||
"status",
|
||
"release_resolved_at",
|
||
"last_sync_at",
|
||
)
|
||
|
||
@staticmethod
|
||
def _norm_incident_obs_value(field: str, value: Any) -> Any:
|
||
"""Normalize optional observation values for equality during migration.
|
||
|
||
NULL and empty string are treated as equivalent for optional text
|
||
fields (legacy rows often omit them). Integer fields keep NULL
|
||
distinct from zero so event_count=0 vs NULL is not collapsed away.
|
||
"""
|
||
if field in ("gitea_issue_number", "event_count"):
|
||
if value is None or value == "":
|
||
return None
|
||
return int(value)
|
||
if value is None:
|
||
return ""
|
||
return str(value)
|
||
|
||
def _incident_link_observation_identity(
|
||
self, row: sqlite3.Row, cols: set[str]
|
||
) -> tuple[Any, ...]:
|
||
"""Stable comparable identity of all meaningful observation metadata."""
|
||
parts: list[Any] = []
|
||
for field in self._INCIDENT_LINK_OBS_COMPARE_FIELDS:
|
||
if field not in cols:
|
||
continue
|
||
# Row keys match column names; missing keys treated as absent.
|
||
try:
|
||
raw = row[field]
|
||
except (IndexError, KeyError):
|
||
raw = None
|
||
parts.append((field, self._norm_incident_obs_value(field, raw)))
|
||
return tuple(parts)
|
||
|
||
def _migrate_incident_links_null_scope(self, conn: sqlite3.Connection) -> None:
|
||
"""Collapse legacy NULL-scope duplicates, then normalize to '' (#619 RC).
|
||
|
||
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 **or** any other meaningful observation metadata conflict
|
||
within a group (no silent data loss), then coerce NULLs to ''.
|
||
|
||
Policy: **fail closed** — never invent a merge of fingerprint/status/
|
||
event_count/etc. Only identical (after optional-text NULL≈'') rows may
|
||
collapse to the lowest ``link_id``.
|
||
"""
|
||
cols = {
|
||
row[1]
|
||
for row in conn.execute("PRAGMA table_info(incident_links)").fetchall()
|
||
}
|
||
if not cols:
|
||
return
|
||
|
||
# Build SELECT from present columns so partial legacy schemas still migrate.
|
||
# Scope fields are normalized in the SELECT aliases used for grouping.
|
||
required = ("link_id", "provider", "provider_issue_id", "gitea_org", "gitea_repo", "gitea_issue_number")
|
||
if not all(c in cols for c in required):
|
||
return
|
||
|
||
select_parts = [
|
||
"link_id",
|
||
"provider",
|
||
"IFNULL(provider_base_url, '') AS base_url"
|
||
if "provider_base_url" in cols
|
||
else "'' AS base_url",
|
||
"IFNULL(provider_org, '') AS p_org"
|
||
if "provider_org" in cols
|
||
else "'' AS p_org",
|
||
"IFNULL(provider_project, '') AS p_project"
|
||
if "provider_project" in cols
|
||
else "'' AS p_project",
|
||
"provider_issue_id",
|
||
"gitea_org",
|
||
"gitea_repo",
|
||
"gitea_issue_number",
|
||
]
|
||
for field in self._INCIDENT_LINK_OBS_COMPARE_FIELDS:
|
||
if field in (
|
||
"gitea_org",
|
||
"gitea_repo",
|
||
"gitea_issue_number",
|
||
):
|
||
continue # already selected
|
||
if field in cols:
|
||
select_parts.append(field)
|
||
|
||
rows = conn.execute(
|
||
f"SELECT {', '.join(select_parts)} 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)}"
|
||
)
|
||
# Same Gitea target is not enough: fingerprint/status/event_count/
|
||
# permalinks/timestamps/etc. must also agree or we lose data.
|
||
obs_identities = {
|
||
self._incident_link_observation_identity(m, cols) for m in members
|
||
}
|
||
if len(obs_identities) > 1:
|
||
raise ControlPlaneError(
|
||
"incident_links migration fail closed: conflicting "
|
||
"observation metadata for provider key "
|
||
f"{key!r} (same Gitea target but differing fingerprint/"
|
||
"status/event_count/permalink/timestamps/or related fields); "
|
||
"refusing silent discard of duplicate rows"
|
||
)
|
||
# lowest link_id kept (ORDER BY link_id ASC)
|
||
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,
|
||
worktree_path: str | None = None,
|
||
owner_pid: int | 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),
|
||
)
|
||
# #601 optional lifecycle columns (present after schema v3 migration)
|
||
try:
|
||
lcols = {
|
||
r[1]
|
||
for r in conn.execute("PRAGMA table_info(leases)").fetchall()
|
||
}
|
||
if "worktree_path" in lcols and worktree_path:
|
||
conn.execute(
|
||
"UPDATE leases SET worktree_path = ? WHERE lease_id = ?",
|
||
(worktree_path, lease_id),
|
||
)
|
||
if "owner_pid" in lcols:
|
||
conn.execute(
|
||
"UPDATE leases SET owner_pid = ? WHERE lease_id = ?",
|
||
(
|
||
owner_pid if owner_pid is not None else os.getpid(),
|
||
lease_id,
|
||
),
|
||
)
|
||
if "expected_head_sha" in lcols and expected_head_sha:
|
||
conn.execute(
|
||
"UPDATE leases SET expected_head_sha = ? WHERE lease_id = ?",
|
||
(expected_head_sha, lease_id),
|
||
)
|
||
except sqlite3.Error:
|
||
pass
|
||
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:
|
||
"""Release a lease; unknown ids are no-ops for backward compatibility."""
|
||
with self._tx(immediate=False) as conn:
|
||
row = conn.execute(
|
||
"SELECT lease_id FROM leases WHERE lease_id = ?",
|
||
(lease_id,),
|
||
).fetchone()
|
||
if not row:
|
||
return
|
||
self.release_lease_recorded(lease_id, session_id=session_id)
|
||
|
||
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
|
||
|
||
def get_incident_link_by_provider(
|
||
self,
|
||
*,
|
||
provider: str,
|
||
provider_issue_id: str,
|
||
provider_base_url: str | None = None,
|
||
provider_org: str | None = None,
|
||
provider_project: str | None = None,
|
||
) -> dict[str, Any] | None:
|
||
"""Lookup canonical incident_links row by provider key (#612 / #613)."""
|
||
base_url = _norm_scope(provider_base_url)
|
||
p_org = _norm_scope(provider_org)
|
||
p_project = _norm_scope(provider_project)
|
||
with self._tx(immediate=False) as conn:
|
||
row = conn.execute(
|
||
"""
|
||
SELECT * FROM incident_links
|
||
WHERE provider = ?
|
||
AND provider_base_url = ?
|
||
AND provider_org = ?
|
||
AND provider_project = ?
|
||
AND provider_issue_id = ?
|
||
LIMIT 1
|
||
""",
|
||
(provider, base_url, p_org, p_project, str(provider_issue_id)),
|
||
).fetchone()
|
||
return dict(row) if row else None
|
||
|
||
|
||
# ── lease lifecycle (#601) ────────────────────────────────────────────
|
||
|
||
_LEASE_LIFECYCLE_COLUMNS: tuple[tuple[str, str], ...] = (
|
||
("worktree_path", "TEXT"),
|
||
("owner_pid", "INTEGER"),
|
||
("expected_head_sha", "TEXT"),
|
||
("adopted_from_session_id", "TEXT"),
|
||
("adopted_by_session_id", "TEXT"),
|
||
("provenance_json", "TEXT"),
|
||
("abandon_proof_json", "TEXT"),
|
||
)
|
||
|
||
def _migrate_lease_lifecycle_columns(self, conn: sqlite3.Connection) -> None:
|
||
"""Add provenance/worktree columns to leases for first-class lifecycle (#601)."""
|
||
cols = {
|
||
row[1]
|
||
for row in conn.execute("PRAGMA table_info(leases)").fetchall()
|
||
}
|
||
if not cols:
|
||
return
|
||
for name, decl in self._LEASE_LIFECYCLE_COLUMNS:
|
||
if name not in cols:
|
||
conn.execute(f"ALTER TABLE leases ADD COLUMN {name} {decl}")
|
||
|
||
def _lease_columns(self, conn: sqlite3.Connection) -> set[str]:
|
||
return {
|
||
row[1]
|
||
for row in conn.execute("PRAGMA table_info(leases)").fetchall()
|
||
}
|
||
|
||
def list_leases(
|
||
self,
|
||
*,
|
||
remote: str | None = None,
|
||
org: str | None = None,
|
||
repo: str | None = None,
|
||
role: str | None = None,
|
||
statuses: Sequence[str] | None = None,
|
||
limit: int = 100,
|
||
) -> list[dict[str, Any]]:
|
||
"""List leases joined with work_items as first-class workflow state."""
|
||
clauses: list[str] = []
|
||
params: list[Any] = []
|
||
if remote:
|
||
clauses.append("w.remote = ?")
|
||
params.append(remote)
|
||
if org:
|
||
clauses.append("w.org = ?")
|
||
params.append(org)
|
||
if repo:
|
||
clauses.append("w.repo = ?")
|
||
params.append(repo)
|
||
if role:
|
||
clauses.append("l.role = ?")
|
||
params.append(role)
|
||
if statuses:
|
||
placeholders = ", ".join("?" for _ in statuses)
|
||
clauses.append(f"l.status IN ({placeholders})")
|
||
params.extend(statuses)
|
||
where = ("WHERE " + " AND ".join(clauses)) if clauses else ""
|
||
sql = f"""
|
||
SELECT l.*, w.remote, w.org, w.repo, w.kind AS work_kind,
|
||
w.number AS work_number, w.state AS work_state,
|
||
w.current_head_sha AS work_head_sha,
|
||
s.pid AS session_pid, s.profile AS session_profile,
|
||
s.status AS session_status
|
||
FROM leases l
|
||
JOIN work_items w ON w.work_item_id = l.work_item_id
|
||
LEFT JOIN sessions s ON s.session_id = l.session_id
|
||
{where}
|
||
ORDER BY l.expires_at DESC
|
||
LIMIT ?
|
||
"""
|
||
params.append(max(1, int(limit)))
|
||
with self._tx(immediate=False) as conn:
|
||
rows = conn.execute(sql, params).fetchall()
|
||
return [dict(r) for r in rows]
|
||
|
||
def get_lease_workflow_state(self, lease_id: str) -> dict[str, Any] | None:
|
||
"""Return lease + assignment + work_item + session for one lease id."""
|
||
with self._tx(immediate=False) as conn:
|
||
lease = conn.execute(
|
||
"SELECT * FROM leases WHERE lease_id = ?",
|
||
(lease_id,),
|
||
).fetchone()
|
||
if not lease:
|
||
return None
|
||
work = conn.execute(
|
||
"SELECT * FROM work_items WHERE work_item_id = ?",
|
||
(lease["work_item_id"],),
|
||
).fetchone()
|
||
asn = conn.execute(
|
||
"""
|
||
SELECT * FROM assignments
|
||
WHERE lease_id = ?
|
||
ORDER BY created_at DESC LIMIT 1
|
||
""",
|
||
(lease_id,),
|
||
).fetchone()
|
||
sess = conn.execute(
|
||
"SELECT * FROM sessions WHERE session_id = ?",
|
||
(lease["session_id"],),
|
||
).fetchone()
|
||
provenance = None
|
||
if "provenance_json" in lease.keys() and lease["provenance_json"]:
|
||
try:
|
||
provenance = json.loads(lease["provenance_json"])
|
||
except (TypeError, json.JSONDecodeError):
|
||
provenance = {"raw": lease["provenance_json"]}
|
||
return {
|
||
"lease": dict(lease),
|
||
"work_item": dict(work) if work else None,
|
||
"assignment": dict(asn) if asn else None,
|
||
"session": dict(sess) if sess else None,
|
||
"provenance": provenance,
|
||
}
|
||
|
||
def attach_lease_provenance(
|
||
self, lease_id: str, provenance: dict[str, Any]
|
||
) -> None:
|
||
payload = json.dumps(provenance)
|
||
with self._tx() as conn:
|
||
cols = self._lease_columns(conn)
|
||
if "provenance_json" not in cols:
|
||
return
|
||
conn.execute(
|
||
"UPDATE leases SET provenance_json = ? WHERE lease_id = ?",
|
||
(payload, lease_id),
|
||
)
|
||
if provenance.get("adopted_from_session_id") and "adopted_from_session_id" in cols:
|
||
conn.execute(
|
||
"""
|
||
UPDATE leases
|
||
SET adopted_from_session_id = ?, adopted_by_session_id = ?
|
||
WHERE lease_id = ?
|
||
""",
|
||
(
|
||
provenance.get("adopted_from_session_id"),
|
||
provenance.get("adopted_by_session_id"),
|
||
lease_id,
|
||
),
|
||
)
|
||
if provenance.get("worktree_path") and "worktree_path" in cols:
|
||
conn.execute(
|
||
"UPDATE leases SET worktree_path = ? WHERE lease_id = ?",
|
||
(provenance.get("worktree_path"), lease_id),
|
||
)
|
||
if provenance.get("expected_head_sha") and "expected_head_sha" in cols:
|
||
conn.execute(
|
||
"UPDATE leases SET expected_head_sha = ? WHERE lease_id = ?",
|
||
(provenance.get("expected_head_sha"), lease_id),
|
||
)
|
||
|
||
def release_lease_recorded(
|
||
self, lease_id: str, *, session_id: str
|
||
) -> dict[str, Any]:
|
||
"""Explicit release with audit event + provenance proof (#601)."""
|
||
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"cannot release lease {lease_id} owned by {row['session_id']}"
|
||
)
|
||
if row["status"] not in ("active", "expired"):
|
||
# idempotent-ish for already released
|
||
if row["status"] == "released":
|
||
return {
|
||
"lease_id": lease_id,
|
||
"status": "released",
|
||
"session_id": session_id,
|
||
"released_at": now_s,
|
||
"idempotent": True,
|
||
}
|
||
conn.execute(
|
||
"UPDATE leases SET status = 'released' WHERE lease_id = ?",
|
||
(lease_id,),
|
||
)
|
||
conn.execute(
|
||
"UPDATE assignments SET status = 'released' WHERE lease_id = ?",
|
||
(lease_id,),
|
||
)
|
||
proof = {
|
||
"lease_id": lease_id,
|
||
"status": "released",
|
||
"session_id": session_id,
|
||
"released_at": now_s,
|
||
"prior_status": row["status"],
|
||
"work_item_id": row["work_item_id"],
|
||
}
|
||
cols = self._lease_columns(conn)
|
||
if "provenance_json" in cols:
|
||
prior = {}
|
||
if row["provenance_json"]:
|
||
try:
|
||
prior = json.loads(row["provenance_json"])
|
||
except (TypeError, json.JSONDecodeError):
|
||
prior = {}
|
||
prior["last_release"] = proof
|
||
conn.execute(
|
||
"UPDATE leases SET provenance_json = ? WHERE lease_id = ?",
|
||
(json.dumps(prior), 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}",
|
||
now_s,
|
||
),
|
||
)
|
||
return proof
|
||
|
||
def force_expire_lease(self, lease_id: str, *, reason: str = "") -> None:
|
||
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}")
|
||
conn.execute(
|
||
"UPDATE leases SET status = 'expired' WHERE lease_id = ?",
|
||
(lease_id,),
|
||
)
|
||
conn.execute(
|
||
"UPDATE assignments SET status = 'expired' WHERE lease_id = ?",
|
||
(lease_id,),
|
||
)
|
||
conn.execute(
|
||
"""
|
||
INSERT INTO events(work_item_id, event_type, message, created_at)
|
||
VALUES (?, 'lease_expired', ?, ?)
|
||
""",
|
||
(
|
||
row["work_item_id"],
|
||
f"lease {lease_id} force-expired: {reason or 'unspecified'}",
|
||
now_s,
|
||
),
|
||
)
|
||
|
||
def abandon_lease(
|
||
self,
|
||
*,
|
||
lease_id: str,
|
||
requester_session_id: str,
|
||
proof: dict[str, Any],
|
||
) -> 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}")
|
||
conn.execute(
|
||
"UPDATE leases SET status = 'abandoned' WHERE lease_id = ?",
|
||
(lease_id,),
|
||
)
|
||
conn.execute(
|
||
"UPDATE assignments SET status = 'abandoned' WHERE lease_id = ?",
|
||
(lease_id,),
|
||
)
|
||
cols = self._lease_columns(conn)
|
||
if "abandon_proof_json" in cols:
|
||
conn.execute(
|
||
"UPDATE leases SET abandon_proof_json = ? WHERE lease_id = ?",
|
||
(json.dumps(proof), lease_id),
|
||
)
|
||
msg = (
|
||
f"session {requester_session_id} abandoned lease {lease_id} "
|
||
f"(prior owner {row['session_id']})"
|
||
)
|
||
conn.execute(
|
||
"""
|
||
INSERT INTO events(work_item_id, event_type, message, created_at)
|
||
VALUES (?, 'lease_abandoned', ?, ?)
|
||
""",
|
||
(row["work_item_id"], msg, now_s),
|
||
)
|
||
return {
|
||
"lease_id": lease_id,
|
||
"status": "abandoned",
|
||
"prior_owner_session_id": row["session_id"],
|
||
"requester_session_id": requester_session_id,
|
||
"abandoned_at": now_s,
|
||
"proof": proof,
|
||
}
|
||
|
||
def adopt_lease(
|
||
self,
|
||
*,
|
||
lease_id: str,
|
||
adopter_session_id: str,
|
||
role: str,
|
||
worktree_path: str | None = None,
|
||
expected_head_sha: str | None = None,
|
||
owner_pid: int | None = None,
|
||
provenance: dict[str, Any] | None = None,
|
||
lease_ttl_seconds: int = DEFAULT_LEASE_TTL_SECONDS,
|
||
) -> dict[str, Any]:
|
||
"""Transfer or refresh a lease with provenance (#601).
|
||
|
||
* Same owner + active → refresh (owner-resume).
|
||
* Expired/abandoned/released → create new assignment+lease with provenance.
|
||
* Active foreign → raise ForeignLeaseError (never silent steal).
|
||
"""
|
||
now = _utc_now()
|
||
now_s = _ts(now)
|
||
expires = _ts(now + timedelta(seconds=int(lease_ttl_seconds)))
|
||
provenance = provenance or {}
|
||
|
||
with self._tx(immediate=True) as conn:
|
||
lease = conn.execute(
|
||
"SELECT * FROM leases WHERE lease_id = ?",
|
||
(lease_id,),
|
||
).fetchone()
|
||
if not lease:
|
||
raise ControlPlaneError(f"unknown lease_id {lease_id}")
|
||
work = conn.execute(
|
||
"SELECT * FROM work_items WHERE work_item_id = ?",
|
||
(lease["work_item_id"],),
|
||
).fetchone()
|
||
if not work:
|
||
raise ControlPlaneError("lease has no work_item")
|
||
|
||
# Expire by time if needed
|
||
exp = _parse_ts(lease["expires_at"])
|
||
status = lease["status"]
|
||
if status == "active" and exp and exp <= now:
|
||
conn.execute(
|
||
"UPDATE leases SET status = 'expired' WHERE lease_id = ?",
|
||
(lease_id,),
|
||
)
|
||
conn.execute(
|
||
"UPDATE assignments SET status = 'expired' WHERE lease_id = ?",
|
||
(lease_id,),
|
||
)
|
||
status = "expired"
|
||
|
||
owner = lease["session_id"]
|
||
if status == "active" and owner != adopter_session_id:
|
||
raise ForeignLeaseError(
|
||
f"cannot adopt active foreign lease {lease_id} owned by {owner}"
|
||
)
|
||
|
||
# Ensure adopter session exists
|
||
sess = conn.execute(
|
||
"SELECT session_id FROM sessions WHERE session_id = ?",
|
||
(adopter_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')
|
||
""",
|
||
(adopter_session_id, role, owner_pid or os.getpid(), now_s, now_s),
|
||
)
|
||
|
||
cols = self._lease_columns(conn)
|
||
|
||
if status == "active" and owner == adopter_session_id:
|
||
# Owner-resume refresh
|
||
conn.execute(
|
||
"""
|
||
UPDATE leases
|
||
SET heartbeat_at = ?, expires_at = ?, phase = ?
|
||
WHERE lease_id = ?
|
||
""",
|
||
(now_s, expires, "adopted", lease_id),
|
||
)
|
||
if "worktree_path" in cols and worktree_path:
|
||
conn.execute(
|
||
"UPDATE leases SET worktree_path = ? WHERE lease_id = ?",
|
||
(worktree_path, lease_id),
|
||
)
|
||
if "owner_pid" in cols and owner_pid is not None:
|
||
conn.execute(
|
||
"UPDATE leases SET owner_pid = ? WHERE lease_id = ?",
|
||
(owner_pid, lease_id),
|
||
)
|
||
if "provenance_json" in cols:
|
||
prior = {}
|
||
if lease["provenance_json"]:
|
||
try:
|
||
prior = json.loads(lease["provenance_json"])
|
||
except (TypeError, json.JSONDecodeError):
|
||
prior = {}
|
||
prior["last_adopt"] = provenance
|
||
conn.execute(
|
||
"UPDATE leases SET provenance_json = ? WHERE lease_id = ?",
|
||
(json.dumps(prior), lease_id),
|
||
)
|
||
asn = conn.execute(
|
||
"""
|
||
SELECT * FROM assignments
|
||
WHERE lease_id = ? AND status = 'active'
|
||
ORDER BY created_at DESC LIMIT 1
|
||
""",
|
||
(lease_id,),
|
||
).fetchone()
|
||
lease2 = conn.execute(
|
||
"SELECT * FROM leases WHERE lease_id = ?", (lease_id,)
|
||
).fetchone()
|
||
conn.execute(
|
||
"""
|
||
INSERT INTO events(work_item_id, event_type, message, created_at)
|
||
VALUES (?, 'lease_adopted', ?, ?)
|
||
""",
|
||
(
|
||
lease["work_item_id"],
|
||
f"owner-resume adopt lease {lease_id} by {adopter_session_id}",
|
||
now_s,
|
||
),
|
||
)
|
||
return {
|
||
"outcome": "adopted_owner_resume",
|
||
"lease": dict(lease2) if lease2 else dict(lease),
|
||
"assignment": dict(asn) if asn else None,
|
||
"reasons": ["owner-resume: refreshed lease with provenance"],
|
||
}
|
||
|
||
# Non-active: create new lease + assignment (transfer)
|
||
new_lease_id = f"lease-{uuid.uuid4().hex[:16]}"
|
||
new_asn_id = f"asn-{uuid.uuid4().hex[:16]}"
|
||
# Mark prior non-active if still active somehow
|
||
if status == "active":
|
||
conn.execute(
|
||
"UPDATE leases SET status = 'released' WHERE lease_id = ?",
|
||
(lease_id,),
|
||
)
|
||
conn.execute(
|
||
"UPDATE assignments SET status = 'released' WHERE lease_id = ?",
|
||
(lease_id,),
|
||
)
|
||
|
||
# Prefer prior assignment allowed/forbidden
|
||
prior_asn = conn.execute(
|
||
"""
|
||
SELECT * FROM assignments WHERE lease_id = ?
|
||
ORDER BY created_at DESC LIMIT 1
|
||
""",
|
||
(lease_id,),
|
||
).fetchone()
|
||
allowed = (
|
||
prior_asn["allowed_actions"]
|
||
if prior_asn
|
||
else json.dumps(["implement", "comment", "push", "create_pr"])
|
||
)
|
||
forbidden = (
|
||
prior_asn["forbidden_actions"]
|
||
if prior_asn
|
||
else json.dumps(
|
||
["approve", "merge", "request_changes", "self_select_without_assignment"]
|
||
)
|
||
)
|
||
head = expected_head_sha or (
|
||
prior_asn["expected_head_sha"] if prior_asn else work["current_head_sha"]
|
||
)
|
||
|
||
# Dynamic insert for optional columns
|
||
base_cols = [
|
||
"lease_id",
|
||
"work_item_id",
|
||
"session_id",
|
||
"role",
|
||
"phase",
|
||
"expires_at",
|
||
"heartbeat_at",
|
||
"status",
|
||
]
|
||
base_vals: list[Any] = [
|
||
new_lease_id,
|
||
lease["work_item_id"],
|
||
adopter_session_id,
|
||
role,
|
||
"adopted",
|
||
expires,
|
||
now_s,
|
||
"active",
|
||
]
|
||
optional = {
|
||
"worktree_path": worktree_path,
|
||
"owner_pid": owner_pid,
|
||
"expected_head_sha": head,
|
||
"adopted_from_session_id": owner,
|
||
"adopted_by_session_id": adopter_session_id,
|
||
"provenance_json": json.dumps(provenance),
|
||
}
|
||
for col, val in optional.items():
|
||
if col in cols and val is not None:
|
||
base_cols.append(col)
|
||
base_vals.append(val)
|
||
placeholders = ", ".join("?" for _ in base_cols)
|
||
conn.execute(
|
||
f"INSERT INTO leases({', '.join(base_cols)}) VALUES ({placeholders})",
|
||
base_vals,
|
||
)
|
||
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', ?)
|
||
""",
|
||
(
|
||
new_asn_id,
|
||
lease["work_item_id"],
|
||
adopter_session_id,
|
||
new_lease_id,
|
||
allowed if isinstance(allowed, str) else json.dumps(list(allowed)),
|
||
forbidden if isinstance(forbidden, str) else json.dumps(list(forbidden)),
|
||
head,
|
||
role,
|
||
now_s,
|
||
),
|
||
)
|
||
conn.execute(
|
||
"""
|
||
INSERT INTO events(work_item_id, event_type, message, created_at)
|
||
VALUES (?, 'lease_adopted', ?, ?)
|
||
""",
|
||
(
|
||
lease["work_item_id"],
|
||
f"session {adopter_session_id} adopted from {owner} "
|
||
f"prior={lease_id} new={new_lease_id}",
|
||
now_s,
|
||
),
|
||
)
|
||
lease2 = conn.execute(
|
||
"SELECT * FROM leases WHERE lease_id = ?", (new_lease_id,)
|
||
).fetchone()
|
||
asn2 = conn.execute(
|
||
"SELECT * FROM assignments WHERE assignment_id = ?",
|
||
(new_asn_id,),
|
||
).fetchone()
|
||
return {
|
||
"outcome": "adopted_transfer",
|
||
"lease": dict(lease2) if lease2 else None,
|
||
"assignment": dict(asn2) if asn2 else None,
|
||
"reasons": [
|
||
f"transferred lease ownership from {owner} to {adopter_session_id}"
|
||
],
|
||
}
|
||
|