diff options
Diffstat (limited to 'wflib/cloud.py')
| -rw-r--r-- | wflib/cloud.py | 442 |
1 files changed, 442 insertions, 0 deletions
diff --git a/wflib/cloud.py b/wflib/cloud.py new file mode 100644 index 0000000..5878b87 --- /dev/null +++ b/wflib/cloud.py @@ -0,0 +1,442 @@ +"""Cloud-lane ledger: pure logic (spec docs/specs cloud-lane-design §4.4). IO: wf_cloud.py.""" +from __future__ import annotations + +import base64 +import binascii +import dataclasses +import datetime as dt +import gzip +import hashlib +import json +import re + +from . import tasks as TK, usage as U + +OVERHEAD = 1.15 # title generation, setup: what the session can't see +LOST_AFTER = dt.timedelta(hours=24) +ENDED = ("done", "awaiting", "handback", "lost") + + +class CloudError(Exception): + pass + + +def new() -> dict: + return {"budget": 240.0, "spent": 0.0, "reserve_per_task": 4.0, "max_parallel": 3, "entries": []} + + +def loads(text: str) -> dict: + if not text.strip(): + return new() + try: + led = json.loads(text) + led["budget"], led["spent"], led["entries"] + except (ValueError, KeyError, TypeError) as e: + raise CloudError(f"cloud ledger unreadable: {e}") + return {**new(), **led} + + +def dumps(led: dict) -> str: + return json.dumps(led, indent=1) + "\n" + + +def running(led: dict) -> int: + return sum(e["state"] == "running" for e in led["entries"]) + + +def balance(led: dict) -> float: + return led["budget"] - led["spent"] - led["reserve_per_task"] * running(led) + + +def refusal(led: dict) -> str | None: + """One line why `send` must refuse (exit 3), or None.""" + bal, res = balance(led), led["reserve_per_task"] + if bal < res: + return f"cloud busy: balance ${bal:.2f} < reserve ${res:.2f}" + if running(led) >= led["max_parallel"]: + return f"cloud busy: {running(led)} running >= max_parallel {led['max_parallel']}" + return None + + +def add(led: dict, id: str, project: str, sid: str, model: str, now: dt.datetime) -> dict: + e = {"id": id, "project": project, "sid": sid, "model": model, "sent": now.isoformat(), + "state": "running", "usd": None, "usd_source": None} + led["entries"].append(e) + return e + + +def find(led: dict, sid: str) -> dict: + for e in led["entries"]: + if e["sid"] == sid: + return e + raise CloudError(f"no cloud entry with sid {sid}") + + +def charge_usd(model: str, u: U.Usage | None, reserve: float) -> tuple[float, str]: + """(usd, source): WF-USAGE x PRICES + 15%; missing/unknown model -> reserve, 'est'.""" + c = U.cost(model, u) if u is not None else None + return (c * OVERHEAD, "self") if c is not None else (reserve, "est") + + +def end(led: dict, sid: str, state: str, u: U.Usage | None = None) -> dict: + """Running entry ends: state, charge added to spent.""" + if state not in ENDED: + raise CloudError(f"bad end state {state}") + e = find(led, sid) + if e["state"] != "running": + raise CloudError(f"entry {sid} is {e['state']}, not running") + if state == "lost": + usd, src = led["reserve_per_task"], "est" + else: + usd, src = charge_usd(e["model"], u, led["reserve_per_task"]) + e.update(state=state, usd=usd, usd_source=src) + led["spent"] += usd + return e + + +def expire(led: dict, now: dt.datetime) -> list[dict]: + """Running > 24h (no WF-RESULT seen) -> lost, charged reserve. Returns the lost entries.""" + out = [] + for e in led["entries"]: + if e["state"] == "running" and now - dt.datetime.fromisoformat(e["sent"]) > LOST_AFTER: + out.append(end(led, e["sid"], "lost")) + return out + + +def set_balance(led: dict, bal: float, now: dt.datetime) -> dict: + """Owner reconcile (balance read from claude.ai): spent = budget - balance.""" + old = led["spent"] + led["spent"] = led["budget"] - bal + row = {"id": "-", "project": "-", "sid": "-", "model": "-", "sent": now.isoformat(), "state": "done", + "usd": led["spent"] - old, "usd_source": "owner"} + led["entries"].append(row) + return row + + +def summary(led: dict) -> str: + return (f"budget ${led['budget']:.2f} spent ${led['spent']:.2f} balance ${balance(led):.2f} " + f"running {running(led)}/{led['max_parallel']} (reserve ${led['reserve_per_task']:.2f}/task)") + + + +# --- archive ended sessions (spec §4.3 Archive): the CLI's internal archiveRemoteSession, CLI 2.1.291 --- +API = "https://api.anthropic.com" +_ARCH_SID = re.compile(r"session_[A-Za-z0-9]+") + + +def archive_request(sid: str, creds: str, version: str) -> tuple[str, dict]: + """(url, headers) for POST .../archive; creds = the CLI's .credentials.json text. Never logged.""" + if not _ARCH_SID.fullmatch(sid): + raise CloudError(f"{sid!r} is not a cloud session id") + try: + d = json.loads(creds) + tok = d["claudeAiOauth"]["accessToken"] + except (ValueError, KeyError, TypeError): + tok = None + if not tok: + raise CloudError("no claude.ai login token (run claude once)") + h = {"Authorization": f"Bearer {tok}", "Content-Type": "application/json", + "anthropic-version": "2023-06-01", "User-Agent": f"claude-code/{version}"} + if d.get("trustedDeviceToken"): + h["X-Trusted-Device-Token"] = d["trustedDeviceToken"] + return f"{API}/v1/code/sessions/{sid}/archive", h + + +def archive_problem(status: int, body: str) -> str | None: + """None = archived (200, or 409 = already); else one short why.""" + if status in (200, 409): + return None + if status == 401: + return "HTTP 401 (login expired? run claude once)" + return f"HTTP {status}: {' '.join(body.split())}"[:120] + + +def unarchived(led: dict) -> list[str]: + """sids of ended sessions not archived yet (owner reconcile rows and pending sids skipped).""" + return [e["sid"] for e in led["entries"] + if e["state"] in ENDED and _ARCH_SID.fullmatch(e["sid"]) and not e.get("archived")] + +# --- pty driver helpers (spec §3 F1-F3, F7, F8) --------------------------------------------------- +import re # noqa: E402 + +_ANSI = re.compile(r"\x1b(\[[0-?]*[ -/]*[@-~]|\][^\x07\x1b]*(\x07|\x1b\\)|[PX^_][^\x1b]*\x1b\\|[@-Z\\-_])") +_SID = re.compile(r"(session_[A-Za-z0-9]{8,})") + + +def strip_ansi(text: str) -> str: + """Terminal output -> plain text: cursor-forward (CSI n C) -> n spaces, cursor-to-column (CSI n G) + -> one space (the TUI draws word gaps that way), other escape sequences dropped, CR -> LF.""" + text = re.sub(r"\x1b\[(\d*)C", lambda m: " " * int(m.group(1) or 1), text) + text = re.sub(r"\x1b\[\d*G", " ", text) + return _ANSI.sub("", text).replace("\r\n", "\n").replace("\r", "\n") + + +def parse_sid(text: str) -> str | None: + """Session id from `claude --cloud` output: 'Resume with: claude --teleport <sid>', else the View URL.""" + plain = strip_ansi(text) + for pat in (r"Resume with:\s*claude --teleport\s+" + _SID.pattern, r"View:\s*\S*?/code/" + _SID.pattern): + m = re.findall(pat, plain) + if m: + return m[-1] + return None + + +def last_lines(text: str, n: int = 5) -> list[str]: + return [ln.strip() for ln in strip_ansi(text).splitlines() if ln.strip()][-n:] + + +def trust(claude_json: str, path: str) -> str: + """~/.claude.json text with projects[path].hasTrustDialogAccepted = true (F3). Other keys kept.""" + try: + data = json.loads(claude_json) if claude_json.strip() else {} + projects = data.setdefault("projects", {}) + if not isinstance(projects, dict): + raise TypeError("projects is not an object") + except (ValueError, TypeError, AttributeError) as e: + raise CloudError(f"~/.claude.json unreadable: {e}") + projects.setdefault(path, {})["hasTrustDialogAccepted"] = True + return json.dumps(data, indent=2) + "\n" + + +TRUST_PROMPT = re.compile(r"trust\s*(the\s*files\s*in\s*)?this\s*folder|Do\s*you\s*trust", re.I) +RESUMED = re.compile(r"Session\s*resumed", re.I) # teleport reached the prompt (F7) + + +# --- send: snapshot + prompt (spec §4.1, §4.2) ----------------------------------------------------- + +MODEL = "claude-opus-5-5" # what cloud sessions run (F4): ledger rows price with it +CAP = 90_000_000 # bytes: packed snapshot limit (upload ≤ 100 MB, F1) +NEVER = (".wf", ".worktrees", "out") # + TASKS / archive: never in a snapshot (cloud_include re-adds) + + +def include_problem(path: str) -> str | None: + """Why a cloud_include entry is unusable (absolute, leaves the project, empty), None = fine.""" + parts = path.replace("\\", "/").strip("/").split("/") + if not path.strip() or path.startswith("/") or ".." in parts or parts in (["."], [".git"]) or parts[0] == ".git": + return f"cloud_include '{path}': must be a path inside the project" + return None + + +def size_refusal(size: int, cap: int = CAP) -> str | None: + if size > cap: + return f"snapshot {size / 1e6:.0f} MB > {cap / 1e6:.0f} MB; trim cloud_include" + return None + + +def fill(template: str, id: str, base: str, task: str, recipe: str, note: str | None) -> str: + """templates/cloud-prompt.md with its {{…}} fields filled.""" + values = {"id": id, "base": base, "task": task.strip(), "recipe": recipe.strip() or "(none: see CLAUDE.md)", + "note": f"\n{note.strip()}\n" if note and note.strip() else ""} + unknown = sorted(set(re.findall(r"\{\{(\w+)\}\}", template)) - set(values)) + if unknown: + raise CloudError(f"cloud prompt template: unknown field {{{{{unknown[0]}}}}}") + out = re.sub(r"\{\{(\w+)\}\}", lambda m: values[m.group(1)], template) # one pass: task text kept as is + return out + + +# The export renders the prompt as '❯ first line' + ' ' continuations and every assistant message as +# '● first line' + ' ' continuations (F12): only an assistant message can start with '● WF-RESULT'. +RESULT = re.compile(r"^● WF-RESULT\b") + + +def final_message(export: str) -> list[str] | None: + """Lines of the last assistant message that starts with WF-RESULT (bullet / indent and trailing blanks + removed), None = no result yet. The prompt never matches, whatever it contains.""" + lines = export.replace("\r\n", "\n").split("\n") + starts = [i for i, ln in enumerate(lines) if RESULT.match(ln)] + if not starts: + return None + out = [] + for ln in lines[starts[-1]:]: + if out and ln[:2] not in (" ", ""): + break # next message / prompt + out.append(ln[2:].rstrip()) + while out and not out[-1]: + out.pop() + return out + + +def keyless_message(export: str) -> list[str] | None: + """No WF-RESULT anywhere, but the last assistant message ends in a bare WF-PATCH-END line (the session + dropped the key words): its lines (as final_message), else None. A non-empty prompt after it (a redo + sent) -> None: still running.""" + lines = export.replace("\r\n", "\n").split("\n") + blocks: list[tuple[str, list[str]]] = [] + for ln in lines: + if ln[:2] in (" ", "") and blocks: + blocks[-1][1].append(ln[2:].rstrip()) + elif ln.strip(): + blocks.append((ln[:1], [ln[2:].rstrip()])) + found = None + for kind, body in blocks: + if kind == "❯" and any(b.strip() for b in body): + found = None # a later prompt: the message before it is answered + elif kind == "●": + text = [b for b in body if b.strip()] + if text and text[-1].strip() == "WF-PATCH-END": + found = body + if found is None: + return None + out = list(found) + while out and not out[-1]: + out.pop() + return out + + +def keyless_result(lines: list[str]) -> Result: + """keyless_message() lines -> a handback Result: usage + patch (sha256=HEX bytes=N header anywhere before + WF-PATCH-END, base64 after it) when found, so the patch can be kept.""" + text = " ".join(ln.strip() for ln in lines[:-1]) + r = Result("handback", "(no WF-RESULT key)") + if m := _USAGE.search(text): + r.usage = U.Usage(turns=1, inp=int(m[1]), cw5=int(m[2]), cr=int(m[3]), out=int(m[4])) + r.model = m[5] + blob = "".join(text.split()) + if m := _HEADER.search(blob): + r.sha, r.nbytes, r.b64, r.has_patch = m[1].lower(), int(m[2]), blob[m.end():], True + return r + + +# --- pull: result + patch (spec §4.3) --------------------------------------------------------------- + +KEYS = ("WF-RESULT", "WF-REPORT", "WF-USAGE", "WF-PATCH-BEGIN", "WF-PATCH-END") +STATES = ("done", "awaiting", "handback") +_USAGE = re.compile(r"in=(\d+)\s+cw=(\d+)\s+cr=(\d+)\s+out=(\d+)(?:\s+model=(\S+))?") +# whitespace removed: the export may wrap the header anywhere; gzip base64 starts 'H4sI', never a digit +_HEADER = re.compile(r"sha256=([0-9a-fA-F]{64})bytes=(\d+)") + + +@dataclasses.dataclass +class Result: + state: str + report: str + usage: U.Usage | None = None + model: str | None = None + sha: str | None = None + nbytes: int | None = None + b64: str = "" + has_patch: bool = False + + +def parse_result(lines: list[str]) -> Result: + """final_message() lines -> Result. A line starting with a key opens its section, other lines continue + the open one (the export wraps long lines). CloudError: bad state / missing patch markers.""" + sec: dict[str, list[str]] = {} + cur = None + for ln in lines: + s = ln.strip() + key = next((k for k in KEYS if s == k or s.startswith(k + " ")), None) + if key: + cur = key + sec.setdefault(key, []).append(s[len(key):].strip()) + elif cur: + sec[cur].append(s) + state = " ".join(sec.get("WF-RESULT", [""])).strip().lower() + if state not in STATES: + raise CloudError(f"bad WF-RESULT '{state}' (want {'/'.join(STATES)})") + r = Result(state, " ".join(" ".join(sec.get("WF-REPORT", [])).split()) or "(no report)") + if m := _USAGE.search(" ".join(sec.get("WF-USAGE", []))): + r.usage = U.Usage(turns=1, inp=int(m[1]), cw5=int(m[2]), cr=int(m[3]), out=int(m[4])) + r.model = m[5] + if "WF-PATCH-BEGIN" in sec: + if "WF-PATCH-END" not in sec: + raise CloudError("WF-PATCH-BEGIN without WF-PATCH-END (truncated message)") + blob = "".join("".join(sec["WF-PATCH-BEGIN"]).split()) + m = _HEADER.match(blob) + if not m: + raise CloudError("WF-PATCH-BEGIN needs sha256=HEX bytes=N") + r.sha, r.nbytes, r.b64, r.has_patch = m[1].lower(), int(m[2]), blob[m.end():], True + return r + + +def decode_patch(r: Result) -> bytes: + """base64 -> check bytes + sha256 -> gunzip. CloudError on any mismatch.""" + if not r.has_patch: + raise CloudError("no WF-PATCH-BEGIN/END in the result") + try: + gz = base64.b64decode(r.b64, validate=True) + except (binascii.Error, ValueError) as e: + raise CloudError(f"patch base64 broken: {e}") + if len(gz) != r.nbytes: + raise CloudError(f"patch bytes {len(gz)} != {r.nbytes} announced") + if hashlib.sha256(gz).hexdigest() != r.sha: + raise CloudError("patch sha256 mismatch") + try: + return gzip.decompress(gz) + except (OSError, EOFError) as e: + raise CloudError(f"patch gzip broken: {e}") + + +_DIFF = re.compile(r'^diff --git "?a/(.+?)"? "?b/(.+?)"?$') +_MOVE = re.compile(r"^(?:rename|copy) (?:from|to) (.+)$") + + +def patch_files(patch: str) -> list[str]: + """Repo-relative paths a format-patch touches (both sides of renames/copies), in order, unique.""" + out = [] + for ln in patch.split("\n"): + if m := _DIFF.match(ln): + out += [m[1], m[2]] + elif m := _MOVE.match(ln): + out.append(m[1].strip('"')) + return list(dict.fromkeys(out)) + + +def patch_problem(files: list[str], project: str, never: list[str]) -> str | None: + """Why a patch may not be applied, None = fine. files: repo-relative; project: the project folder + relative to the repo ('' = repo root); never: project-relative paths the cloud may not touch + (TASKS, archive, .wf, .worktrees, out, cloud_include).""" + pre = project.strip("/") + for f in files: + parts = f.split("/") + if not f or f.startswith("/") or ".." in parts or parts[0] == ".git": + return f"patch touches {f}: outside the tree" + if pre and not (f + "/").startswith(pre + "/"): + return f"patch touches {f}: outside the project folder {pre}" + rel = f[len(pre) + 1:] if pre else f + for n in never: + n = n.strip("/") + if rel == n or rel.startswith(n + "/"): + return f"patch touches {f}: {n} is never sent (wf bookkeeping / out / cloud_include)" + return None + + +# ------------------------------------------------------------------ fit (spec §4.5) + +# (reason, regex) over title + body; `Cloud: yes` skips them (Model opus still required). wf res before the generic wf-step rule. +UNFIT = ( + ("names `wf res` (local resource ledger)", re.compile(r"\bwf res\b")), + ("names `wf` commands as steps", re.compile(r"\bwf (?!res\b)[a-z][a-z-]*")), + ("needs a GUI", re.compile(r"\bGUI\b|\bscreenshot", re.I)), + ("needs a LAN host", re.compile(r"\bLAN\b|\b192\.168\.\d+\.\d+|\b10\.\d+\.\d+\.\d+|\b[\w-]+\.local\b|\bssh [\w@.-]+")), + ("needs a live server", re.compile(r"live server|running server|\blocalhost\b|127\.0\.0\.1", re.I)), +) + + +def fit(item: TK.Item, cfg) -> str | None: + """None when the task fits the cloud lane, else the reason it does not (spec §4.5).""" + if not cfg.cloud: + return "project not opted in (cloud = true)" + if item.cloud == "no": + return "Cloud: no" + if not item.runner_ready: + return "not runner-ready (needs Done, not owner-bound)" + if item.sessions in ("owner", "solo"): + return f"Sessions: {item.sessions}" + if item.effort not in TK.EFFORTS or TK.EFFORTS.index(item.effort) > TK.EFFORTS.index(cfg.slice_above): + return "slice job (effort above slice_above)" + if item.model != "opus": + return f"Model {item.model} stays local" + if item.cloud == "yes": + return None + text = "\n".join([item.text, *item.body]) + for reason, rx in UNFIT: + if rx.search(text): + return reason + return None + + +def pick_key(item: TK.Item) -> tuple: + """wf orch pick cloud order (spec §4.5): Cloud: yes first, then prio, then larger effort (all fits are opus).""" + eff = TK.EFFORTS.index(item.effort) if item.effort in TK.EFFORTS else -1 + return (item.cloud != "yes", item.prio if item.prio is not None else 9, -eff) |