patch-cc 0.1.0__py3-none-any.whl

patch_cc/patcher.py

@@ -0,0 +1,244 @@
+"""Orchestration: read a binary, run selected patches, write it back safely."""
+
+from __future__ import annotations
+
+import hashlib
+import json
+import os
+import shutil
+from dataclasses import dataclass, field
+from pathlib import Path
+
+from . import __version__, locate
+from .bun import Bundle, container
+from .patches import ALL_PATCHES, DEFAULT_SUFFIX, SENTINEL, Options, Outcome, Patch
+
+#: Every patched bundle ends with one comment line recording exactly what was
+#: applied. Comments cannot collide with code, survive re-extraction, and make
+#: ``status`` a parse instead of a guess -- value-flip patches leave no other
+#: fingerprint.
+MANIFEST_PREFIX = "//patch-cc "
+
+#: Fingerprints of binaries patched by versions before the manifest existed.
+_LEGACY_MARKER = "(Claude Code)\\n" + DEFAULT_SUFFIX
+
+
+class AlreadyPatchedError(RuntimeError):
+    """The only source available is already patched; patching it would stack.
+
+    Our edits change lengths, so a second pass over a patched bundle corrupts
+    rather than updates. There is deliberately no force-override: when no
+    pristine backup exists the only honest fixes are ``restore`` or a
+    reinstall.
+    """
+
+
+@dataclass(slots=True)
+class PatchReport:
+    version: str | None
+    kind: str
+    original_size: int
+    patched_size: int = 0
+    results: list[tuple[Patch, Outcome]] = field(default_factory=list)
+    backup: Path | None = None
+    output: Path | None = None
+
+    @property
+    def landed_ids(self) -> list[str]:
+        return [p.id for p, o in self.results if o.landed]
+
+    @property
+    def regressions(self) -> list[Patch]:
+        """Selected patches that were expected to change something but did not."""
+        return [p for p, o in self.results if not o.landed]
+
+    @property
+    def partial(self) -> list[tuple[Patch, list[str]]]:
+        """Patches that landed but had some sub-step miss."""
+        out = []
+        for patch, outcome in self.results:
+            missed = outcome.missed_steps()
+            if outcome.landed and missed:
+                out.append((patch, missed))
+        return out
+
+
+def build_manifest(landed: list[str], options: Options) -> str:
+    payload: dict = {"v": 1, "tool": __version__, "patches": landed}
+    if options.rebrands:
+        payload["brand"] = options.brand
+    if options.version_suffix != DEFAULT_SUFFIX:
+        payload["suffix"] = options.version_suffix
+    if options.subagent_models:
+        payload["models"] = options.subagent_models
+    return "\n" + MANIFEST_PREFIX + json.dumps(payload, separators=(",", ":")) + "\n"
+
+
+def read_manifest(source: str) -> dict | None:
+    """The applied-patch record, or ``None`` for pristine/legacy binaries."""
+    start = source.rfind("\n" + MANIFEST_PREFIX)
+    if start == -1:
+        return None
+    start += 1 + len(MANIFEST_PREFIX)
+    end = source.find("\n", start)
+    line = source[start:] if end == -1 else source[start:end]
+    try:
+        data = json.loads(line)
+    except json.JSONDecodeError:
+        return None
+    return data if isinstance(data, dict) else None
+
+
+def is_patched(source: str) -> bool:
+    return (
+        ("\n" + MANIFEST_PREFIX) in source
+        or SENTINEL in source
+        or _LEGACY_MARKER in source
+    )
+
+
+def selected_patches(ids: list[str]) -> list[Patch]:
+    """Resolve ids to patches, preserving registry (run) order."""
+    wanted = set(ids)
+    return [patch for patch in ALL_PATCHES if patch.id in wanted]
+
+
+def run_patches(
+    source: str, patches: list[Patch], options: Options
+) -> tuple[str, list[tuple[Patch, Outcome]]]:
+    results: list[tuple[Patch, Outcome]] = []
+    current = source
+    for patch in patches:
+        current, outcome = patch.run(current, options)
+        results.append((patch, outcome))
+    return current, results
+
+
+def _backup_dir() -> Path:
+    base = os.environ.get("XDG_DATA_HOME")
+    root = Path(base) if base else Path.home() / ".local" / "share"
+    return root / "patch-cc" / "backups"
+
+
+def backup_path_for(install: locate.Installation) -> Path:
+    """The single source of truth for where a binary's backup lives.
+
+    Canonical native installs are version-named, giving a clean
+    ``<name>.<version>.orig``. When the name is not a version we cannot tell two
+    unrelated ``claude`` binaries apart by name alone, so a short hash of the
+    absolute path is mixed in to keep their backups distinct.
+    """
+    root = _backup_dir()
+    if install.version:
+        stem = f"{install.binary.name}.{install.version}"
+    else:
+        digest = hashlib.sha256(str(install.binary.resolve()).encode()).hexdigest()[:8]
+        stem = f"{install.binary.name}.unknown-{digest}"
+    return root / f"{stem}.orig"
+
+
+def read_pristine(install: locate.Installation) -> Bundle:
+    """The bundle patching starts from: the backup when one exists.
+
+    Patching never stacks edits on edits -- each apply begins at this pristine
+    source, so the selected set is always exactly what ends up in the binary.
+    """
+    backup = backup_path_for(install)
+    return container.read(str(backup if backup.exists() else install.binary))
+
+
+def _backup(install: locate.Installation, *, pristine: bool) -> Path | None:
+    """Record the pristine original once, so ``restore`` is a plain copy back.
+
+    Only ever captures a binary that is actually unpatched: backing up an
+    already-marked binary would enshrine a poisoned "original" that ``restore``
+    would later hand back as clean.
+    """
+    dest = backup_path_for(install)
+    if dest.exists():
+        return dest
+    if not pristine:
+        return None
+    dest.parent.mkdir(parents=True, exist_ok=True)
+    shutil.copy2(install.binary, dest)
+    return dest
+
+
+def patch_installation(
+    install: locate.Installation,
+    selected: list[str],
+    options: Options,
+    *,
+    bundle: Bundle | None = None,
+    out_path: Path | None = None,
+    make_backup: bool = True,
+) -> PatchReport:
+    """Patch ``install`` (or write to ``out_path``) with the ``selected`` patches.
+
+    Patching always starts from a pristine source (:func:`read_pristine`), so
+    re-applying replaces the previous patch set instead of stacking on it.
+    ``bundle`` may carry that already-read source (the CLI reads it for
+    validation first).
+    """
+    source = bundle if bundle is not None else read_pristine(install)
+    if is_patched(source.source):
+        raise AlreadyPatchedError(
+            f"{install.binary} is already patched and no pristine backup exists, "
+            "so there is nothing clean to patch from. Run `patch-cc restore`, "
+            "or reinstall Claude to get a clean binary."
+        )
+
+    patches = selected_patches(selected)
+    patched_source, results = run_patches(source.source, patches, options)
+
+    report = PatchReport(
+        version=install.version,
+        kind=source.kind,
+        original_size=source.binary_size,
+        results=results,
+    )
+
+    landed = report.landed_ids
+    if not landed:
+        # Nothing changed; writing would only strip bytecode for no benefit.
+        return report
+
+    patched_source += build_manifest(landed, options)
+
+    target = out_path or install.binary
+    if make_backup and out_path is None:
+        report.backup = _backup(install, pristine=not is_patched(source.source))
+
+    container.write(source, patched_source, str(target))
+    report.output = Path(target)
+    report.patched_size = Path(target).stat().st_size
+    return report
+
+
+def clean_source_path(install: locate.Installation) -> Path | None:
+    """A binary whose bundle is guaranteed unpatched, for matcher-health tests.
+
+    If the installed binary is already patched, our own edits have removed the
+    anchors the matchers look for, so a dry-run against it conflates
+    "already applied" with "anchor gone". The pristine backup is the honest
+    thing to test against.
+    """
+    backup = backup_path_for(install)
+    return backup if backup.exists() else None
+
+
+def restore(install: locate.Installation) -> Path:
+    """Copy the pristine backup back over the installed binary."""
+    backup = backup_path_for(install)
+    if not backup.exists():
+        raise FileNotFoundError(
+            f"No backup found for {install.binary.name} "
+            f"{install.version or '(unknown version)'} at {backup}. "
+            "If Claude auto-updated, the original for this version was never saved -- "
+            "reinstall to get a clean binary."
+        )
+    # A full-file copy-back, so it works for both ELF and Mach-O.
+    from .bun.elf import atomic_write  # noqa: PLC0415
+
+    atomic_write(str(install.binary), backup.read_bytes(), mode_from=str(backup))
+    return install.binary

patch_cc/patches/__init__.py

@@ -0,0 +1,75 @@
+"""The patch registry.
+
+Patches run in registration order and each sees the previous one's output.
+The order below is the upstream order; do not reorder casually -- some patches
+depend on regions an earlier one leaves untouched.
+"""
+
+from __future__ import annotations
+
+from . import agents, chrome, output, streaming, thinking
+from .base import (
+    DEFAULT_BRAND,
+    DEFAULT_SUFFIX,
+    GROUP_AGENTS,
+    GROUP_CHROME,
+    GROUP_LIVE,
+    GROUP_OUTPUT,
+    GROUP_THINKING,
+    SENTINEL,
+    Options,
+    Outcome,
+    Patch,
+    derived_brand,
+)
+
+ALL_PATCHES: list[Patch] = [
+    *output.PATCHES,
+    *thinking.PATCHES,
+    *streaming.PATCHES,
+    *agents.PATCHES,
+    *chrome.PATCHES,
+]
+
+GROUP_ORDER = [GROUP_OUTPUT, GROUP_THINKING, GROUP_LIVE, GROUP_AGENTS, GROUP_CHROME]
+
+_BY_ID = {patch.id: patch for patch in ALL_PATCHES}
+
+
+def get(patch_id: str) -> Patch:
+    try:
+        return _BY_ID[patch_id]
+    except KeyError:
+        raise KeyError(f"unknown patch id: {patch_id}") from None
+
+
+def ids() -> list[str]:
+    return [patch.id for patch in ALL_PATCHES]
+
+
+def default_ids() -> list[str]:
+    return [patch.id for patch in ALL_PATCHES if patch.default]
+
+
+def by_group() -> dict[str, list[Patch]]:
+    grouped: dict[str, list[Patch]] = {group: [] for group in GROUP_ORDER}
+    for patch in ALL_PATCHES:
+        grouped.setdefault(patch.group, []).append(patch)
+    return grouped
+
+
+__all__ = [
+    "ALL_PATCHES",
+    "DEFAULT_BRAND",
+    "DEFAULT_SUFFIX",
+    "GROUP_ORDER",
+    "Options",
+    "Outcome",
+    "Patch",
+    "SENTINEL",
+    "derived_brand",
+    "get",
+    "ids",
+    "default_ids",
+    "by_group",
+]

patch_cc/patches/agents.py

@@ -0,0 +1,289 @@
+"""Subagent patches: prompt visibility, and overriding built-in models.
+
+Everything the model override offers is discovered from the bundle itself:
+
+* **Agents** come from the built-in definition shape
+  ``agentType:"<name>",whenToUse:...`` carrying ``source:"built-in"``.
+* **Models** come from the Task tool's own input schema -- the
+  ``model:enum([...])`` whose describe-string starts "Optional model override".
+
+So a new upstream agent or model shows up here without a code change, and we
+can never offer a name the binary would reject.
+"""
+
+from __future__ import annotations
+
+import re
+from dataclasses import dataclass
+
+from .base import GROUP_AGENTS, IDENT, Options, Outcome, Patch, compile_js, splice
+
+# ------------------------------------------------------- prompt visibility
+
+_BACKGROUNDED = '"Backgrounded agent"'
+_LIVE_PROMPT_MOUNT = compile_js(
+    rf"({IDENT})&&({IDENT})&&({IDENT})\.createElement\(m,\{{marginBottom:1\}},"
+    rf"\3\.createElement\(({IDENT}),\{{prompt:\2\}}\)\)"
+)
+_EMPTY_STATE = compile_js(
+    rf"if\(({IDENT})\.length===0&&!?\(({IDENT})&&({IDENT})\)\)return"
+)
+_TRANSCRIPT_MODE = compile_js(rf"isTranscriptMode:({IDENT})=!1")
+
+
+def _subagent_prompt(content: str, _options: Options, outcome: Outcome) -> str:
+    """Show the subagent ``Prompt`` block outside transcript mode."""
+    gate = outcome.step("gate")
+    output = content
+    index = 0
+
+    while True:
+        anchor = output.find(_BACKGROUNDED, index)
+        if anchor == -1:
+            break
+        fn_start = output.rfind("function ", 0, anchor)
+        fn_end_candidate = output.find("function ", anchor + len(_BACKGROUNDED))
+        fn_end = len(output) if fn_end_candidate == -1 else fn_end_candidate
+        index = anchor + len(_BACKGROUNDED)
+        if fn_start == -1 or fn_end <= fn_start:
+            continue
+
+        segment = output[fn_start:fn_end]
+        relevant = (
+            'action:"app:toggleTranscript"' in segment
+            and 'fallback:"ctrl+o"' in segment
+            and "isTranscriptMode:" in segment
+            and "{prompt:" in segment
+            and ",theme:" in segment
+        )
+        if not relevant:
+            continue
+
+        transcript = _TRANSCRIPT_MODE.search(segment)
+        if not transcript:
+            continue
+        transcript_var = transcript.group(1)
+        gate_pattern = compile_js(rf"{re.escape(transcript_var)}&&({IDENT})&&")
+
+        def drop_gate(match: re.Match[str]) -> str:
+            prompt_var = match.group(1)
+            nearby = segment[match.end() : match.end() + 260]
+            if f"{{prompt:{prompt_var},theme:" not in nearby:
+                return match.group(0)
+            gate.candidates += 1
+            gate.applied += 1
+            return f"{prompt_var}&&"
+
+        next_segment = gate_pattern.sub(drop_gate, segment)
+        if next_segment != segment:
+            output = splice(output, fn_start, fn_end, next_segment)
+            index = fn_start + len(next_segment)
+
+    mount = outcome.step("mount")
+
+    def rewrite_mount(match: re.Match[str]) -> str:
+        _transcript, prompt_var, ns, component = match.groups()
+        mount.candidates += 1
+        replacement = (
+            f"{prompt_var}&&{ns}.createElement(m,{{marginBottom:1}},"
+            f"{ns}.createElement({component},{{prompt:{prompt_var}}}))"
+        )
+        if replacement != match.group(0):
+            mount.applied += 1
+        return replacement
+
+    output = _LIVE_PROMPT_MOUNT.sub(rewrite_mount, output)
+
+    empty = outcome.step("empty-state")
+
+    def rewrite_empty(match: re.Match[str]) -> str:
+        rows, _transcript, prompt_var = match.groups()
+        empty.candidates += 1
+        replacement = f"if({rows}.length===0&&!{prompt_var})return"
+        if replacement != match.group(0):
+            empty.applied += 1
+        return replacement
+
+    return _EMPTY_STATE.sub(rewrite_empty, output)
+
+
+# ----------------------------------------------------------- discovery
+
+#: Always offered besides the discovered aliases: keep the agent on whatever
+#: the main loop runs.
+INHERIT = "inherit"
+
+_AGENT_DEF = compile_js(r'agentType:"([\w-]+)",whenToUse:')
+_MODEL_FIELD = compile_js(r'model:"([\w\[\].-]+)"')
+#: A definition object is scanned at most this far; every known definition fits
+#: well within it, and the cap keeps a moved anchor from swallowing a neighbour.
+_DEF_WINDOW = 3000
+
+_MODEL_ENUM = compile_js(
+    rf'model:{IDENT}\.enum\(\[((?:"[\w\[\]]+",?)+)\]\)\.optional\(\)'
+    rf'\.describe\([`"]Optional model override'
+)
+#: Used only if the Task-tool schema anchor ever disappears.
+_FALLBACK_MODELS = ("haiku", "sonnet", "opus")
+
+
+@dataclass(slots=True, frozen=True)
+class BuiltinAgent:
+    """One built-in agent definition as found in a bundle."""
+
+    name: str
+    #: Current ``model:"..."`` literal, or ``None`` when the definition has no
+    #: model field (which the runtime treats as inherit).
+    model: str | None
+    #: Offset of the definition anchor in the scanned source.
+    start: int
+    #: Offset of the model *value* inside the source, ``-1`` when absent.
+    model_start: int
+    #: Where a ``model:"...",`` property would be inserted.
+    insert_at: int
+
+    @property
+    def effective_model(self) -> str:
+        return self.model or INHERIT
+
+
+def discover_agents(source: str) -> list[BuiltinAgent]:
+    """Built-in agent definitions as they exist in *this* bundle.
+
+    Definitions marked internal (their ``whenToUse`` says so) are not offered:
+    they are orchestration plumbing, not agents a user chooses.
+    """
+    agents: list[BuiltinAgent] = []
+    seen: set[str] = set()
+    for match in _AGENT_DEF.finditer(source):
+        name = match.group(1)
+        window = source[match.start() : match.start() + _DEF_WINDOW]
+        stop = window.find("getSystemPrompt:")
+        span = window if stop == -1 else window[:stop]
+        if 'source:"built-in"' not in span or 'whenToUse:"Internal' in span:
+            continue
+        if name in seen:
+            continue
+        seen.add(name)
+        field = _MODEL_FIELD.search(span)
+        agents.append(
+            BuiltinAgent(
+                name=name,
+                model=field.group(1) if field else None,
+                start=match.start(),
+                model_start=match.start() + field.start(1) if field else -1,
+                insert_at=match.end() - len("whenToUse:"),
+            )
+        )
+    return agents
+
+
+def discover_models(source: str) -> list[str]:
+    """Model aliases the binary's own Task tool accepts for subagents."""
+    match = _MODEL_ENUM.search(source)
+    if not match:
+        return list(_FALLBACK_MODELS)
+    return re.findall(r'"([\w\[\]]+)"', match.group(1))
+
+
+# --------------------------------------------------------- model overrides
+
+# One helper resolves a built-in agent's default model and, for exactly one
+# agent (Explore today), ignores the definition's model field in favour of its
+# own pin. Overriding that agent means neutralising this bypass so the
+# definition -- which we just rewrote -- is authoritative again.
+_MODEL_BYPASS = compile_js(
+    rf"function ({IDENT})\(({IDENT}),({IDENT})\)\{{"
+    rf'if\(\2\.agentType!==({IDENT})\.agentType\|\|\2\.source!=="built-in"\)return \2\.model;'
+    rf'return {IDENT}\(\3\)\?{IDENT}:"inherit"\}}'
+)
+
+
+def bypassed_agent(source: str) -> str | None:
+    """Name of the agent whose definition the bypass helper overrides, if any."""
+    match = _MODEL_BYPASS.search(source)
+    if not match:
+        return None
+    def_var = re.escape(match.group(4))
+    assign = compile_js(rf'(?<![\w$]){def_var}=\{{agentType:"([\w-]+)"').search(source)
+    return assign.group(1) if assign else None
+
+
+def _neutralize_bypass(content: str, step: Outcome) -> str:
+    def rewrite(match: re.Match[str]) -> str:
+        step.candidates += 1
+        step.applied += 1
+        name, obj, model = match.group(1, 2, 3)
+        return f"function {name}({obj},{model}){{return {obj}.model}}"
+
+    output = _MODEL_BYPASS.sub(rewrite, content, count=1)
+    if step.candidates == 0:
+        step.note(
+            "model-bypass helper not found; the pinned agent may keep its own default"
+        )
+    return output
+
+
+def _subagent_models(content: str, options: Options, outcome: Outcome) -> str:
+    """Write the chosen model into each overridden built-in definition.
+
+    Definitions with a ``model:"..."`` literal get it rewritten; definitions
+    without one get it inserted. Both target offsets from a fresh discovery
+    pass, so this never guesses about the bytes between anchor and value.
+    """
+    if not options.subagent_models:
+        outcome.note("no subagent model overrides configured")
+        return content
+
+    output = content
+    offered = {INHERIT, *discover_models(output)}
+
+    for agent, target in sorted(options.subagent_models.items()):
+        step = outcome.step(agent)
+        if target not in offered:
+            step.note(f"model {target!r} not offered by this bundle; skipped")
+            continue
+        info = next((a for a in discover_agents(output) if a.name == agent), None)
+        if info is None:
+            step.note(f"no built-in agent {agent!r} in this bundle; skipped")
+            continue
+
+        step.candidates += 1
+        if info.effective_model == target:
+            continue  # already the desired model
+        if info.model is None:
+            output = splice(
+                output, info.insert_at, info.insert_at, f'model:"{target}",'
+            )
+        else:
+            output = splice(
+                output, info.model_start, info.model_start + len(info.model), target
+            )
+        step.applied += 1
+
+    pinned = bypassed_agent(output)
+    if pinned is not None and pinned in options.subagent_models:
+        output = _neutralize_bypass(output, outcome.step("model-bypass"))
+
+    return output
+
+
+PATCHES = [
+    Patch(
+        id="subagent-prompt",
+        title="Show subagent prompts",
+        summary="Show a subagent's Prompt block during normal use, not only in transcript mode.",
+        group=GROUP_AGENTS,
+        fn=_subagent_prompt,
+        anchors=('"Backgrounded agent"', 'action:"app:toggleTranscript"'),
+    ),
+    Patch(
+        id="subagent-models",
+        title="Override subagent models",
+        summary="Choose the default model for the built-in agents found in your binary.",
+        group=GROUP_AGENTS,
+        fn=_subagent_models,
+        default=False,
+        anchors=('agentType:"', "Optional model override"),
+    ),
+]