|
13 | 13 |
|
14 | 14 | from __future__ import annotations |
15 | 15 |
|
| 16 | +import json |
16 | 17 | import logging |
| 18 | +import os |
17 | 19 | import subprocess |
| 20 | +import threading |
18 | 21 | from dataclasses import dataclass |
| 22 | +from datetime import datetime, timezone |
19 | 23 | from pathlib import Path |
20 | 24 | from typing import Optional |
21 | 25 |
|
22 | 26 | logger = logging.getLogger(__name__) |
23 | 27 |
|
24 | 28 | WORKTREE_DIR = ".codeframe/worktrees" |
| 29 | +_REGISTRY_FILE = ".codeframe/worktrees.json" |
| 30 | +_registry_lock = threading.Lock() |
25 | 31 |
|
26 | 32 |
|
27 | 33 | @dataclass |
@@ -173,3 +179,98 @@ def cleanup( |
173 | 179 | ) |
174 | 180 | except Exception as exc: |
175 | 181 | logger.warning("Failed to delete branch %s: %s", branch_name, exc) |
| 182 | + |
| 183 | + |
| 184 | +def get_base_branch(workspace_path: Path) -> str: |
| 185 | + """Return the current HEAD branch name, defaulting to 'main' on failure. |
| 186 | +
|
| 187 | + Returns 'main' when git is unavailable, the directory is not a repo, |
| 188 | + or HEAD is detached (rev-parse returns 'HEAD' literally). |
| 189 | + """ |
| 190 | + result = subprocess.run( |
| 191 | + ["git", "rev-parse", "--abbrev-ref", "HEAD"], |
| 192 | + cwd=str(workspace_path), |
| 193 | + capture_output=True, |
| 194 | + text=True, |
| 195 | + ) |
| 196 | + branch = result.stdout.strip() |
| 197 | + if result.returncode != 0 or not branch or branch == "HEAD": |
| 198 | + return "main" |
| 199 | + return branch |
| 200 | + |
| 201 | + |
| 202 | +def list_worktrees(workspace_path: Path) -> list[dict]: |
| 203 | + """Return all entries in the worktree registry, or [] if absent/corrupt/malformed.""" |
| 204 | + registry_file = workspace_path / _REGISTRY_FILE |
| 205 | + if not registry_file.exists(): |
| 206 | + return [] |
| 207 | + try: |
| 208 | + data = json.loads(registry_file.read_text()) |
| 209 | + return data if isinstance(data, list) else [] |
| 210 | + except Exception: |
| 211 | + return [] |
| 212 | + |
| 213 | + |
| 214 | +class WorktreeRegistry: |
| 215 | + """Atomic registry of active worktrees for orphan detection. |
| 216 | +
|
| 217 | + Stores entries in ``.codeframe/worktrees.json``. All mutations are |
| 218 | + protected by a module-level threading.Lock and written atomically |
| 219 | + (write to .tmp then os.replace). |
| 220 | + """ |
| 221 | + |
| 222 | + def register(self, workspace_path: Path, task_id: str, batch_id: str) -> None: |
| 223 | + """Add or refresh a worktree entry for task_id.""" |
| 224 | + with _registry_lock: |
| 225 | + entries = list_worktrees(workspace_path) |
| 226 | + # Remove any existing entry for this task_id (idempotent) |
| 227 | + entries = [e for e in entries if e["task_id"] != task_id] |
| 228 | + entries.append({ |
| 229 | + "task_id": task_id, |
| 230 | + "batch_id": batch_id, |
| 231 | + "created_at": datetime.now(timezone.utc).isoformat(), |
| 232 | + "pid": os.getpid(), |
| 233 | + }) |
| 234 | + self._write(workspace_path, entries) |
| 235 | + |
| 236 | + def unregister(self, workspace_path: Path, task_id: str) -> None: |
| 237 | + """Remove the worktree entry for task_id. Safe if absent.""" |
| 238 | + with _registry_lock: |
| 239 | + entries = list_worktrees(workspace_path) |
| 240 | + entries = [e for e in entries if e["task_id"] != task_id] |
| 241 | + self._write(workspace_path, entries) |
| 242 | + |
| 243 | + def list_stale(self, workspace_path: Path) -> list[dict]: |
| 244 | + """Return entries whose registered PID is no longer alive.""" |
| 245 | + stale = [] |
| 246 | + for entry in list_worktrees(workspace_path): |
| 247 | + pid = entry.get("pid") |
| 248 | + if pid is None: |
| 249 | + continue |
| 250 | + try: |
| 251 | + os.kill(pid, 0) |
| 252 | + except ProcessLookupError: |
| 253 | + stale.append(entry) |
| 254 | + except PermissionError: |
| 255 | + pass # Process is alive but owned by another user — not stale |
| 256 | + return stale |
| 257 | + |
| 258 | + def cleanup_stale(self, workspace_path: Path) -> None: |
| 259 | + """Remove stale worktree directories and unregister their entries.""" |
| 260 | + stale = self.list_stale(workspace_path) |
| 261 | + if not stale: |
| 262 | + return |
| 263 | + wt = TaskWorktree() |
| 264 | + for entry in stale: |
| 265 | + task_id = entry["task_id"] |
| 266 | + logger.info("Cleaning up orphaned worktree for task %s (pid %s)", task_id, entry.get("pid")) |
| 267 | + wt.cleanup(workspace_path, task_id) |
| 268 | + self.unregister(workspace_path, task_id) |
| 269 | + |
| 270 | + @staticmethod |
| 271 | + def _write(workspace_path: Path, entries: list[dict]) -> None: |
| 272 | + registry_file = workspace_path / _REGISTRY_FILE |
| 273 | + registry_file.parent.mkdir(parents=True, exist_ok=True) |
| 274 | + tmp_file = registry_file.with_suffix(".json.tmp") |
| 275 | + tmp_file.write_text(json.dumps(entries, indent=2)) |
| 276 | + os.replace(tmp_file, registry_file) |
0 commit comments