fix: exclude foreign-claimed work from allocation instead of blockading the queue (Closes #765)

Recovered preserved candidate d06198b onto current master 0c2f45a via clean cherry-pick. Active foreign leases are excluded before allocator ranking; controller_instance_id ownership is persisted; dashboard matches allocator exclusion. Stable patch-id bacafc5f… preserved. Co-Authored-By: Claude Opus 4.8 (1M context) <[email protected]>
2026-07-20 13:52:39 -05:00
parent 0c2f45abb7
commit ad13d872df
5 changed files with 735 additions and 22 deletions
@@ -302,6 +302,7 @@ class ControlPlaneDB:
                conn.executescript(_SCHEMA_SQL)
                self._migrate_incident_links_null_scope(conn)
                self._migrate_lease_lifecycle_columns(conn)
+                self._migrate_session_ownership_columns(conn)
                conn.execute(
                    "INSERT OR REPLACE INTO schema_meta(key, value) VALUES (?, ?)",
                    ("schema_version", str(SCHEMA_VERSION)),
@@ -492,32 +493,62 @@ class ControlPlaneDB:
        namespace: str | None = None,
        pid: int | None = None,
        status: str = "active",
+        controller_instance_id: str | None = None,
    ) -> dict[str, Any]:
+        """Register/refresh a session row.
+
+        *controller_instance_id* (#765) is the stable identity of the
+        controller that owns this session. Session ids are regenerated per
+        invocation, so they cannot express "my own in-progress task"; the
+        controller instance can. It is never overwritten with ``None``, so a
+        heartbeat from a caller that does not supply one cannot erase
+        ownership.
+        """
        now = _ts()
+        instance = (controller_instance_id or "").strip() or None
        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),
-                )
+                if instance is None:
+                    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(
+                        """
+                        UPDATE sessions
+                        SET role = ?, profile = ?, namespace = ?, pid = ?,
+                            last_heartbeat_at = ?, status = ?,
+                            controller_instance_id = ?
+                        WHERE session_id = ?
+                        """,
+                        (
+                            role, profile, namespace, pid, now, status,
+                            instance, session_id,
+                        ),
+                    )
            else:
                conn.execute(
                    """
                    INSERT INTO sessions(
                        session_id, role, profile, namespace, pid,
-                        started_at, last_heartbeat_at, status
-                    ) VALUES (?, ?, ?, ?, ?, ?, ?, ?)
+                        started_at, last_heartbeat_at, status,
+                        controller_instance_id
+                    ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
                    """,
-                    (session_id, role, profile, namespace, pid, now, now, status),
+                    (
+                        session_id, role, profile, namespace, pid, now, now,
+                        status, instance,
+                    ),
                )
            row = conn.execute(
                "SELECT * FROM sessions WHERE session_id = ?",
@@ -1215,6 +1246,73 @@ class ControlPlaneDB:
            if name not in cols:
                conn.execute(f"ALTER TABLE leases ADD COLUMN {name} {decl}")

+    _SESSION_OWNERSHIP_COLUMNS: tuple[tuple[str, str], ...] = (
+        ("controller_instance_id", "TEXT"),
+    )
+
+    def _migrate_session_ownership_columns(self, conn: sqlite3.Connection) -> None:
+        """Add the stable controller identity to sessions (#765).
+
+        Pre-existing rows migrate with ``NULL``. A NULL instance is treated as
+        *unknown ownership* by the allocator and is never silently adopted.
+        """
+        cols = {
+            row[1]
+            for row in conn.execute("PRAGMA table_info(sessions)").fetchall()
+        }
+        if not cols:
+            return
+        for name, decl in self._SESSION_OWNERSHIP_COLUMNS:
+            if name not in cols:
+                conn.execute(f"ALTER TABLE sessions ADD COLUMN {name} {decl}")
+
+    def list_active_claims(
+        self,
+        *,
+        remote: str | None = None,
+        org: str | None = None,
+        repo: str | None = None,
+        role: str | None = None,
+        limit: int = 500,
+    ) -> dict[tuple[str, int], dict[str, Any]]:
+        """Map ``(work_kind, work_number)`` to its live claim (#765).
+
+        Only ``active`` leases count as claims; released/expired rows never
+        withhold work. Callers compare the returned ``controller_instance_id``
+        against their own to decide own-task vs foreign-task.
+        """
+        claims: dict[tuple[str, int], dict[str, Any]] = {}
+        for row in self.list_leases(
+            remote=remote,
+            org=org,
+            repo=repo,
+            role=role,
+            statuses=("active",),
+            limit=limit,
+        ):
+            kind = str(row.get("work_kind") or "").strip().lower()
+            number = row.get("work_number")
+            if not kind or number is None:
+                continue
+            key = (kind, int(number))
+            claim = {
+                "lease_id": row.get("lease_id"),
+                "session_id": row.get("session_id"),
+                "controller_instance_id": row.get("session_controller_instance_id"),
+                "role": row.get("role"),
+                "profile": row.get("session_profile"),
+                "expires_at": row.get("expires_at"),
+                "work_kind": kind,
+                "work_number": int(number),
+            }
+            # Keep the longest-lived claim when duplicates exist.
+            previous = claims.get(key)
+            if previous is None or str(claim["expires_at"] or "") > str(
+                previous["expires_at"] or ""
+            ):
+                claims[key] = claim
+        return claims
+
    def _lease_columns(self, conn: sqlite3.Connection) -> set[str]:
        return {
            row[1]
@@ -1256,7 +1354,8 @@ class ControlPlaneDB:
                   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
+                   s.status AS session_status,
+                   s.controller_instance_id AS session_controller_instance_id
            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