subtask-mcp 0.2.3__py3-none-any.whl

subtask/install/installer.py

@@ -0,0 +1,204 @@
+"""Plan and apply install / uninstall for one client, with rollback and ownership."""
+
+from __future__ import annotations
+
+import os
+from dataclasses import dataclass, field
+from importlib import resources
+from pathlib import Path
+from typing import Any
+
+from . import ownership
+from .clients import Client, InstallError, Step
+from .fileops import Snapshot, atomic_write, read_bytes, sha256, snapshot
+
+
+def skill_source() -> bytes:
+    return resources.files("subtask").joinpath("skill/SKILL.md").read_bytes()
+
+
+@dataclass
+class Outcome:
+    client: str
+    scope: str
+    status: str = "planned"  # planned | done | unchanged | collision | skipped | error
+    steps: list[str] = field(default_factory=list)
+    notes: list[str] = field(default_factory=list)
+
+    def as_dict(self) -> dict[str, Any]:
+        return {"client": self.client, "scope": self.scope, "status": self.status,
+                "steps": self.steps, "notes": self.notes}  # fmt: skip
+
+
+@dataclass
+class _Plan:
+    mcp_steps: list[Step] = field(default_factory=list)
+    skill_write: bytes | None = None
+    skill_remove: bool = False
+    relink: str | None = None  # a symlinked skill folder replaced by a real one (target kept for rollback)
+    blocked: bool = False
+
+
+def _mcp_plan(c: Client, scope: str, argv: list[str], force: bool, out: Outcome, plan: _Plan) -> None:
+    existing = c.read_entry(scope)
+    rec = ownership.get(ownership.key(c.name, scope, "mcp", c.config_path(scope)))
+    if existing is None:
+        plan.mcp_steps = c.add_steps(scope, argv)
+        return
+    if Client.target(existing) == argv:
+        out.notes.append("MCP entry already present with this command")
+        return
+    ours = rec is not None and rec.get("entry") == existing
+    if not ours and not force:
+        out.notes.append(f"collision: a '{c.name}' entry named subtask runs {Client.target(existing)}; "
+                         "use --force to replace it")  # fmt: skip
+        plan.blocked = True
+        return
+    plan.mcp_steps = c.remove_steps(scope) + c.add_steps(scope, argv)
+
+
+def _skill_plan(c: Client, scope: str, force: bool, out: Outcome, plan: _Plan) -> None:
+    dst = c.skill_path(scope)
+    src = skill_source()
+    if dst.parent.is_symlink():
+        broken = not dst.parent.exists()
+        if not broken and not force:
+            out.notes.append(
+                f"collision: {dst.parent} is a link to {os.readlink(dst.parent)}; use --force to replace it"
+            )
+            plan.blocked = True
+            return
+        plan.relink = str(os.readlink(dst.parent))
+        out.notes.append(
+            f"{'broken ' if broken else ''}link {dst.parent} -> {plan.relink} replaced by a folder"
+        )
+        plan.skill_write = src
+        return
+    cur = read_bytes(dst)
+    if cur == src:
+        return
+    rec = ownership.get(ownership.key(c.name, scope, "skill", dst))
+    if cur is not None and not force and not (rec and rec.get("sha") == sha256(cur)):
+        out.notes.append(f"collision: {dst} exists and is not ours; use --force to replace it")
+        plan.blocked = True
+        return
+    plan.skill_write = src
+
+
+def _apply(c: Client, scope: str, argv: list[str], plan: _Plan, out: Outcome) -> None:
+    dst = c.skill_path(scope)
+    snaps: list[Snapshot] = []
+    try:
+        if plan.mcp_steps:
+            snaps.append(snapshot(c.config_path(scope)))
+        if (plan.skill_write is not None or plan.skill_remove) and plan.relink is None:
+            snaps.append(snapshot(dst, keep_copy=False))  # skill dirs are scanned: no stray copies
+        for step in plan.mcp_steps:
+            c.apply(step, scope, argv)
+        if plan.relink is not None:
+            dst.parent.unlink()  # only the link: its target is never touched
+        if plan.skill_write is not None:
+            atomic_write(dst, plan.skill_write)
+        elif plan.skill_remove:
+            dst.unlink(missing_ok=True)
+            _rmdir_empty(dst.parent)
+    except Exception as exc:  # noqa: BLE001 - any failure must roll back, then is reported
+        for snap in reversed(snaps):
+            snap.restore()
+        if plan.relink is not None and not dst.parent.is_symlink():
+            dst.unlink(missing_ok=True)
+            _rmdir_empty(dst.parent)
+            if not dst.parent.exists():
+                os.symlink(plan.relink, dst.parent)
+        out.status = "error"
+        out.notes.append(f"{exc}; rolled back")
+        return
+    out.notes += [f"backup: {s.backup}" for s in snaps if s.backup]
+    out.status = "done"
+
+
+def _rmdir_empty(path: Path) -> None:
+    try:
+        path.rmdir()
+    except OSError:
+        pass
+
+
+def install(c: Client, scope: str, argv: list[str], yes: bool, force: bool) -> Outcome:
+    out = Outcome(c.name, scope)
+    try:
+        c.check_scope(scope)
+        plan = _Plan()
+        _mcp_plan(c, scope, argv, force, out, plan)
+        _skill_plan(c, scope, force, out, plan)
+    except InstallError as exc:
+        out.status = "error"
+        out.notes.append(str(exc))
+        return out
+    out.steps = [s.describe() for s in plan.mcp_steps]
+    if plan.skill_write is not None:
+        out.steps.append(f"copy: SKILL.md -> {c.skill_path(scope)}")
+    if plan.blocked:
+        out.status = "collision"
+        return out
+    if not out.steps:
+        out.status = "unchanged"
+    elif not yes:
+        return out
+    else:
+        _apply(c, scope, argv, plan, out)
+    if out.status in ("done", "unchanged"):
+        _record(c, scope)
+    return out
+
+
+def _record(c: Client, scope: str) -> None:
+    entry = c.read_entry(scope)
+    if entry is not None:
+        ownership.put(ownership.key(c.name, scope, "mcp", c.config_path(scope)), {"entry": entry})
+    cur = read_bytes(c.skill_path(scope))
+    if cur is not None and cur == skill_source():
+        ownership.put(ownership.key(c.name, scope, "skill", c.skill_path(scope)), {"sha": sha256(cur)})
+
+
+def uninstall(c: Client, scope: str, yes: bool) -> Outcome:
+    out = Outcome(c.name, scope)
+    plan = _Plan()
+    mkey = ownership.key(c.name, scope, "mcp", c.config_path(scope))
+    skey = ownership.key(c.name, scope, "skill", c.skill_path(scope))
+    try:
+        c.check_scope(scope)
+        existing = c.read_entry(scope)
+        rec = ownership.get(mkey)
+        if existing is not None:
+            if rec is None:
+                out.notes.append("MCP entry 'subtask' was not installed by subtask: left alone")
+            elif rec.get("entry") != existing:
+                out.notes.append("MCP entry 'subtask' was modified since install: left alone")
+            else:
+                plan.mcp_steps = c.remove_steps(scope)
+        cur = read_bytes(c.skill_path(scope))
+        srec = ownership.get(skey)
+        if cur is not None:
+            if srec and srec.get("sha") == sha256(cur):
+                plan.skill_remove = True
+            else:
+                out.notes.append(f"{c.skill_path(scope)} is not ours or was modified: left alone")
+    except InstallError as exc:
+        out.status = "error"
+        out.notes.append(str(exc))
+        return out
+    out.steps = [s.describe() for s in plan.mcp_steps]
+    if plan.skill_remove:
+        out.steps.append(f"delete: {c.skill_path(scope)}")
+    if not out.steps:
+        out.status = "unchanged"
+        return out
+    if yes:
+        _apply(c, scope, [], plan, out)
+        if out.status == "done":
+            if plan.mcp_steps:
+                ownership.drop(mkey)
+            if plan.skill_remove:
+                ownership.drop(skey)
+    return out

subtask/install/ownership.py

@@ -0,0 +1,44 @@
+"""Record of what `subtask install` wrote, so uninstall removes only that."""
+
+from __future__ import annotations
+
+import json
+import os
+from pathlib import Path
+from typing import Any
+
+from .fileops import atomic_write
+
+
+def store_path() -> Path:
+    base = os.environ.get("XDG_DATA_HOME") or str(Path.home() / ".local" / "share")
+    return Path(base) / "subtask" / "installs.json"
+
+
+def load() -> dict[str, Any]:
+    path = store_path()
+    try:
+        data = json.loads(path.read_text())
+    except FileNotFoundError:
+        return {}
+    except (OSError, ValueError) as exc:
+        raise RuntimeError(f"ownership file {path} is unreadable: {exc}") from exc
+    return data if isinstance(data, dict) else {}
+
+
+def key(client: str, scope: str, kind: str, where: Path) -> str:
+    return f"{client}:{scope}:{kind}:{where}"
+
+
+def get(k: str) -> dict[str, Any] | None:
+    return load().get(k)
+
+
+def put(k: str, record: dict[str, Any]) -> None:
+    data = {**load(), k: record}
+    atomic_write(store_path(), (json.dumps(data, indent=2, sort_keys=True) + "\n").encode())
+
+
+def drop(k: str) -> None:
+    data = {name: rec for name, rec in load().items() if name != k}
+    atomic_write(store_path(), (json.dumps(data, indent=2, sort_keys=True) + "\n").encode())

subtask/install/tomledit.py

@@ -0,0 +1,44 @@
+"""Minimal text-level edits of a Codex config.toml: touch one key, keep every other byte."""
+
+from __future__ import annotations
+
+import re
+
+_HEADER = re.compile(r"^\s*\[\[?\s*([^\]]+?)\s*\]\]?\s*(#.*)?$")
+
+
+def _table_span(lines: list[str], table: str) -> tuple[int, int] | None:
+    """Index of the `[table]` header line and the index where its body ends."""
+    start = None
+    for i, line in enumerate(lines):
+        m = _HEADER.match(line)
+        if not m:
+            continue
+        name = m.group(1).replace('"', "").replace(" ", "")
+        if start is not None:
+            return start, i
+        if name == table:
+            start = i
+    return (start, len(lines)) if start is not None else None
+
+
+def set_key(text: str, table: str, key: str, value: str) -> str:
+    """Set `key = value` (value already TOML-encoded) inside `[table]`; other lines untouched."""
+    lines = text.splitlines(keepends=True)
+    span = _table_span(lines, table)
+    if span is None:
+        raise ValueError(f"table [{table}] not found")
+    start, end = span
+    key_re = re.compile(rf"^\s*{re.escape(key)}\s*=")
+    for i in range(start + 1, end):
+        if key_re.match(lines[i]):
+            eol = "\n" if lines[i].endswith("\n") else ""
+            lines[i] = f"{key} = {value}{eol}"
+            return "".join(lines)
+    insert_at = end
+    while insert_at > start + 1 and not lines[insert_at - 1].strip():
+        insert_at -= 1
+    if insert_at > 0 and not lines[insert_at - 1].endswith("\n"):
+        lines[insert_at - 1] += "\n"
+    lines.insert(insert_at, f"{key} = {value}\n")
+    return "".join(lines)

subtask/jobs.py

@@ -0,0 +1,256 @@
+"""Job store on disk: one folder per job, JSON files under a per-job lock.
+
+status.json v1: state, attempt, attempts[], proc (pid + start time + boot id),
+liveness, cost_usd, session_id... A worker is declared dead only with proof
+(procid.alive == "dead"); "unknown" stays visible and keeps its slot and money.
+"""
+
+from __future__ import annotations
+
+import fcntl
+import json
+import os
+import secrets
+import signal
+import time
+from collections.abc import Iterator
+from contextlib import contextmanager
+from pathlib import Path
+
+from . import config, procid
+
+FINAL = {"done", "failed", "stopped"}
+ACTIVE = {"queued", "running"}
+VERSION = 1
+
+
+def new_id() -> str:
+    return "st-" + secrets.token_hex(5)
+
+
+def job_dir(job_id: str) -> Path:
+    if not isinstance(job_id, str) or not job_id.startswith("st-") or "/" in job_id or ".." in job_id:
+        raise ValueError(f"invalid id: {job_id}")
+    return config.JOBS / job_id
+
+
+def read_json(path: Path, default: dict | None = None) -> dict:
+    try:
+        data = json.loads(path.read_text())
+    except (FileNotFoundError, json.JSONDecodeError, NotADirectoryError):
+        return dict(default or {})
+    return data if isinstance(data, dict) else dict(default or {})
+
+
+def write_json(path: Path, data: dict) -> None:
+    tmp = path.with_suffix(f"{path.suffix}.{os.getpid()}.tmp")
+    tmp.write_text(json.dumps(data, ensure_ascii=False, indent=1))
+    os.replace(tmp, path)
+
+
+def spec(job_id: str) -> dict:
+    return read_json(job_dir(job_id) / "spec.json")
+
+
+@contextmanager
+def job_lock(job_id: str) -> Iterator[None]:
+    with (job_dir(job_id) / ".lock").open("a") as lock:
+        fcntl.flock(lock, fcntl.LOCK_EX)
+        yield
+
+
+def update_status(job_id: str, locked: bool = False, **fields: object) -> dict:
+    """Read-modify-write under the job lock (pass locked=True if already held)."""
+    path = job_dir(job_id) / "status.json"
+    if locked:
+        data = {"v": VERSION, **read_json(path), **fields, "updated": time.time()}
+        write_json(path, data)
+        return data
+    with job_lock(job_id):
+        data = {"v": VERSION, **read_json(path), **fields, "updated": time.time()}
+        write_json(path, data)
+    return data
+
+
+def _proc(data: dict) -> dict | None:
+    if data.get("proc"):
+        return data["proc"]
+    if data.get("pid"):  # v0 status: pid only, guarded by the cmdline check below
+        return {"pid": data["pid"], "start": None, "boot": None, "legacy": True}
+    return None
+
+
+def liveness(job_id: str, data: dict) -> str:
+    proc = _proc(data)
+    if not proc:
+        return procid.UNKNOWN
+    verdict = procid.alive(proc)
+    if verdict == procid.UNKNOWN and proc.get("legacy"):
+        try:
+            cmdline = Path(f"/proc/{proc['pid']}/cmdline").read_bytes()
+        except FileNotFoundError:
+            return procid.DEAD
+        except OSError:
+            return verdict
+        return procid.ALIVE if b"subtask.worker" in cmdline and job_id.encode() in cmdline else procid.DEAD
+    return verdict
+
+
+def status(job_id: str) -> dict:
+    """Current state; books a provably dead worker as failed and settles its reservation."""
+    data = read_json(job_dir(job_id) / "status.json", {"state": "unknown"})
+    if data.get("state") not in ACTIVE:
+        return data
+    if _proc(data):
+        verdict = liveness(job_id, data)
+        if verdict == procid.DEAD:
+            return _close_dead(job_id, data)
+        if verdict != data.get("liveness"):
+            return update_status(job_id, liveness=verdict)
+        return data
+    if data.get("state") == "queued" and time.time() - float(data.get("updated") or 0) > 60:
+        return _close_dead(job_id, data, error="worker never started")
+    return data
+
+
+def _close_dead(job_id: str, seen: dict, error: str = "worker died before finishing") -> dict:
+    from . import admission, events, ledger
+
+    with job_lock(job_id):
+        data = read_json(job_dir(job_id) / "status.json")
+        if data.get("state") not in ACTIVE or _proc(data) != _proc(seen):
+            return data  # it finished or was restarted meanwhile
+        charged = 0.0
+        for row in ledger.open_for(job_id):
+            # Spend possible but unknown: charge the whole reservation (prudential), else release.
+            # v0 statuses (no "v") never recorded spend_possible: stay prudential for them.
+            amount = float(row.get("usd") or 0) if data.get("spend_possible", "v" not in data) else 0.0
+            booked = ledger.close(row["rid"], amount, "settle" if amount else "release")
+            charged += booked or 0.0
+        data = update_status(
+            job_id,
+            locked=True,
+            state="failed",
+            pid=None,
+            proc=None,
+            liveness=procid.DEAD,
+            cost_usd=round(float(data.get("cost_usd") or 0) + charged, 4),
+            error=error,
+        )
+    admission.release(job_id)
+    events.emit(job_id, "end", error, state="failed")
+    return data
+
+
+def _group_alive(pgid: int) -> bool:
+    try:
+        os.killpg(pgid, 0)
+    except ProcessLookupError:
+        return False
+    except PermissionError:
+        return True
+    return True
+
+
+def _kill_group(pgid: int, grace_s: float) -> bool:
+    """SIGTERM the worker's process group, then SIGKILL whatever is left. True when nothing is left."""
+    for sig, wait in ((signal.SIGTERM, grace_s), (signal.SIGKILL, 3.0)):
+        try:
+            os.killpg(pgid, sig)
+        except ProcessLookupError:
+            return True
+        except PermissionError:
+            return False
+        deadline = time.time() + wait
+        while time.time() < deadline:
+            try:
+                os.waitpid(pgid, os.WNOHANG)  # reap our own child, if it is one
+            except (ChildProcessError, OSError):
+                pass
+            if not _group_alive(pgid):
+                return True
+            time.sleep(0.1)
+    return not _group_alive(pgid)
+
+
+def stop(job_id: str, grace_s: float = 8.0) -> bool:
+    """Stop the work. True only when no process of the job is left."""
+    current = status(job_id)
+    if current.get("state") not in ACTIVE:
+        return False
+    proc = _proc(current)
+    if not proc:
+        # Spawned but not recorded yet: ask the worker to stop itself, then wait for its identity.
+        update_status(job_id, stop_requested=True)
+        deadline = time.time() + grace_s
+        while time.time() < deadline:
+            current = read_json(job_dir(job_id) / "status.json")
+            proc = _proc(current)
+            if proc or current.get("state") not in ACTIVE:
+                break
+            time.sleep(0.1)
+        if current.get("state") not in ACTIVE:
+            return current.get("state") == "stopped"
+        if not proc:
+            return False  # no worker to prove stopped; the request stays and the worker honours it
+    if liveness(job_id, current) == procid.ALIVE and not _kill_group(int(proc["pid"]), grace_s):
+        return False  # something survived SIGKILL: do not pretend
+    data = status(job_id)  # books a killed worker as failed and settles its reservation
+    if data.get("state") in ACTIVE and _proc(data) and liveness(job_id, data) != procid.DEAD:
+        return False
+    if read_json(job_dir(job_id) / "status.json").get("state") != "stopped":
+        from . import admission, events
+
+        update_status(job_id, state="stopped", pid=None, proc=None)
+        admission.release(job_id)
+        events.emit(job_id, "end", "stopped by request", state="stopped")
+    return True
+
+
+def all_jobs() -> list[str]:
+    if not config.JOBS.is_dir():
+        return []
+    dirs = []
+    for d in config.JOBS.iterdir():
+        if d.name.startswith("st-"):
+            try:
+                dirs.append((d.stat().st_mtime, d.name))
+            except FileNotFoundError:
+                continue
+    return [name for _, name in sorted(dirs, reverse=True)]
+
+
+def running_count() -> int:
+    from . import admission
+
+    return admission.count()
+
+
+# ---- compatibility wrappers over ledger.py (v0 API)
+
+
+def reserve(job_id: str, amount: float) -> float | None:
+    from . import ledger
+
+    return ledger.reserve(f"{job_id}#1", job_id, amount)
+
+
+def unsettled(job_id: str) -> bool:
+    from . import ledger
+
+    return bool(ledger.open_for(job_id))
+
+
+def settle_open(job_id: str, cost: float) -> float:
+    from . import ledger
+
+    booked = 0.0
+    for i, row in enumerate(ledger.open_for(job_id)):
+        booked += ledger.close(row["rid"], cost if i == 0 else 0.0) or 0.0
+    return booked
+
+
+def month_spent() -> float:
+    from . import ledger
+
+    return ledger.committed()