1#!/usr/bin/env python3
2"""Shared helpers for the continuous research loop scripts.
3 
4All queue mutations go through atomic writes (`write_json`, `write_jsonl`) and
5optional file locks (`queue_lock`). The lock uses fcntl exclusive flock so
6concurrent loop_step + loop_discover invocations cannot race on the inbox or
7parked queue.
8"""
9 
10from __future__ import annotations
11 
12import errno
13import fcntl
14import hashlib
15import json
16import os
17import tempfile
18from contextlib import contextmanager
19from datetime import datetime, timezone
20from pathlib import Path
21from typing import Any, Iterable, Iterator
22 
23 
24ROOT = Path(__file__).resolve().parents[2]
25RESEARCHER = ROOT / "researcher"
26QUEUE_DIR = RESEARCHER / "queue"
27ORCH_DIR = RESEARCHER / "orchestration"
28REPORTS_DIR = RESEARCHER / "reports"
29RUNS_DIR = RESEARCHER / "runs"
30SNAPSHOTS_DIR = REPORTS_DIR / "snapshots"
31LOCK_DIR = QUEUE_DIR / ".locks"
32JSONL_QUARANTINE_DIR = REPORTS_DIR / "jsonl-quarantine"
33 
34 
35def utc_now() -> str:
36    return datetime.now(timezone.utc).replace(microsecond=0).isoformat()
37 
38 
39def today_utc() -> str:
40    return datetime.now(timezone.utc).strftime("%Y-%m-%d")
41 
42 
43def load_json(path: Path) -> Any:
44    return json.loads(path.read_text(encoding="utf-8"))
45 
46 
47def _atomic_write(path: Path, payload: str) -> None:
48    path.parent.mkdir(parents=True, exist_ok=True)
49    fd, tmp_name = tempfile.mkstemp(
50        prefix=f".{path.name}.",
51        suffix=".tmp",
52        dir=str(path.parent),
53    )
54    try:
55        with os.fdopen(fd, "w", encoding="utf-8") as handle:
56            handle.write(payload)
57            handle.flush()
58            os.fsync(handle.fileno())
59        os.replace(tmp_name, path)
60    except Exception:
61        try:
62            os.unlink(tmp_name)
63        except FileNotFoundError:
64            pass
65        raise
66 
67 
68def write_json(path: Path, data: Any) -> None:
69    _atomic_write(path, json.dumps(data, indent=2) + "\n")
70 
71 
72def _quarantine_bad_line(path: Path, line_number: int, line: str, exc: Exception) -> None:
73    JSONL_QUARANTINE_DIR.mkdir(parents=True, exist_ok=True)
74    timestamp = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%S")
75    target = JSONL_QUARANTINE_DIR / f"{path.name}.{timestamp}.{line_number}.txt"
76    target.write_text(f"# {exc}\n{line}\n", encoding="utf-8")
77 
78 
79def read_jsonl(path: Path, *, tolerant: bool = True) -> list[dict[str, Any]]:
80    if not path.exists():
81        return []
82    records: list[dict[str, Any]] = []
83    for line_number, raw in enumerate(path.read_text(encoding="utf-8").splitlines(), start=1):
84        line = raw.strip()
85        if not line:
86            continue
87        try:
88            value = json.loads(line)
89        except json.JSONDecodeError as exc:
90            if not tolerant:
91                raise
92            _quarantine_bad_line(path, line_number, raw, exc)
93            continue
94        if not isinstance(value, dict):
95            if not tolerant:
96                raise ValueError(f"{path}:{line_number} expected object, got {type(value).__name__}")
97            _quarantine_bad_line(path, line_number, raw, ValueError("not an object"))
98            continue
99        records.append(value)
100    return records
101 
102 
103def write_jsonl(path: Path, records: Iterable[dict[str, Any]]) -> None:
104    lines = [json.dumps(record, sort_keys=True) for record in records]
105    _atomic_write(path, ("\n".join(lines) + "\n") if lines else "")
106 
107 
108def append_jsonl(path: Path, record: dict[str, Any]) -> None:
109    path.parent.mkdir(parents=True, exist_ok=True)
110    with path.open("a", encoding="utf-8") as handle:
111        try:
112            fcntl.flock(handle.fileno(), fcntl.LOCK_EX)
113        except OSError:
114            pass
115        handle.write(json.dumps(record, sort_keys=True) + "\n")
116        handle.flush()
117        try:
118            fcntl.flock(handle.fileno(), fcntl.LOCK_UN)
119        except OSError:
120            pass
121 
122 
123@contextmanager
124def queue_lock(name: str) -> Iterator[None]:
125    """Exclusive lock for queue mutations. Use one lock per queue file family."""
126 
127    LOCK_DIR.mkdir(parents=True, exist_ok=True)
128    lock_path = LOCK_DIR / f"{name}.lock"
129    handle = open(lock_path, "a+")
130    try:
131        try:
132            fcntl.flock(handle.fileno(), fcntl.LOCK_EX)
133        except OSError as exc:
134            if exc.errno not in {errno.EACCES, errno.EAGAIN}:
135                raise
136        yield
137    finally:
138        try:
139            fcntl.flock(handle.fileno(), fcntl.LOCK_UN)
140        except OSError:
141            pass
142        handle.close()
143 
144 
145def load_config() -> dict[str, Any]:
146    return load_json(ORCH_DIR / "config.json")
147 
148 
149def source_id_for(url: str) -> str:
150    digest = hashlib.sha256(url.strip().lower().encode("utf-8")).hexdigest()
151    return digest[:16]
152 
153 
154def list_run_dirs() -> list[Path]:
155    if not RUNS_DIR.exists():
156        return []
157    return sorted(path for path in RUNS_DIR.iterdir() if path.is_dir())
158 
159 
160def load_run_state(run_dir: Path) -> dict[str, Any] | None:
161    state_file = run_dir / "run-state.json"
162    if not state_file.exists():
163        return None
164    try:
165        return load_json(state_file)
166    except json.JSONDecodeError:
167        return None
168 
169 
170def active_run_urls() -> set[str]:
171    urls: set[str] = set()
172    for run_dir in list_run_dirs():
173        state = load_run_state(run_dir)
174        if not state:
175            continue
176        if state.get("current_state") == "closed":
177            continue
178        url = state.get("source_url")
179        if isinstance(url, str) and url:
180            urls.add(url)
181    return urls
182 
183 
184def closed_run_urls() -> set[str]:
185    urls: set[str] = set()
186    for run_dir in list_run_dirs():
187        state = load_run_state(run_dir)
188        if not state:
189            continue
190        if state.get("current_state") != "closed":
191            continue
192        url = state.get("source_url")
193        if isinstance(url, str) and url:
194            urls.add(url)
195    return urls
196 
197 
198def categorize_runs() -> dict[str, list[Path]]:
199    buckets: dict[str, list[Path]] = {
200        "active": [],
201        "parked": [],
202        "closed": [],
203        "unknown": [],
204    }
205    parked_ids = {record.get("run_id") for record in read_jsonl(QUEUE_DIR / "parked.jsonl")}
206    for run_dir in list_run_dirs():
207        state = load_run_state(run_dir)
208        if not state:
209            buckets["unknown"].append(run_dir)
210            continue
211        current = state.get("current_state")
212        if current == "closed":
213            buckets["closed"].append(run_dir)
214        elif run_dir.name in parked_ids:
215            buckets["parked"].append(run_dir)
216        else:
217            buckets["active"].append(run_dir)
218    return buckets
219 
220 
221def runs_created_today() -> int:
222    today = today_utc()
223    count = 0
224    for run_dir in list_run_dirs():
225        state = load_run_state(run_dir)
226        if state and state.get("created_at", "").startswith(today):
227            count += 1
228    return count
229