LLeMbas/src/lembas/services/agent/patch.py

"""Applying a unified diff, and rendering one.

`difflib` produces a unified diff and cannot apply one, so `render` uses it and
`apply` is written here. No new dependency: hard rule 1 is about the browser,
but a patch applier is fifty lines and pulling a package in for it would be
worse than the fifty lines.

Four behaviours carry the whole module, and each of them exists because of how
models actually write patches rather than how the format is specified.

**Fuzzy offset, exact content.** A hunk's `@@ -41,7 +41,8 @@` is a hint and
nothing more. Models get line numbers wrong constantly -- they count from a
truncated read, or from the file as it was three edits ago -- and get the
context lines right. So the hinted position is tried first and then the file is
scanned outward for an exact match of the context block. One match wins; more
than one refuses, because guessing which of two identical blocks was meant is
the one failure that silently corrupts a file.

**Line endings are normalised in and restored out.** A CRLF file otherwise
fails on every single hunk, on context that looks identical in the error
message, which is unfixable from the model's side.

**A blank context line may have lost its leading space.** Trailing whitespace
is stripped by half the things a model's output passes through, so `""` is read
as a blank context line rather than as a malformed one.

**Nothing is written unless every hunk applies.** The new text is built whole in
memory and handed back; a half-applied file is worse than a refused one, and the
model cannot tell the difference without reading it again.
"""

from __future__ import annotations

import difflib
import re
from dataclasses import dataclass

# A patch bigger than this is a rewrite wearing a diff's clothes, and
# `file_write` is the tool for that.
MAX_HUNKS = 60

# How far either side of the hinted line to look for the context block. Wide
# enough for a file that has grown a few hundred lines since the model read it,
# narrow enough that an accidental match is unlikely.
MAX_DRIFT = 200

_HEADER = re.compile(r"^@@\s*-(\d+)(?:,(\d+))?\s+\+(\d+)(?:,(\d+))?\s*@@")
_NO_NEWLINE = "\\ No newline at end of file"


class PatchError(Exception):
    """A patch that did not apply, said precisely enough to retry from."""

    def __init__(self, message: str, *, hunk: int = 0) -> None:
        super().__init__(message)
        self.message = message
        self.hunk = hunk


@dataclass(frozen=True)
class Hunk:
    old_start: int
    old_count: int
    new_start: int
    new_count: int
    # Each line still carrying its ' ', '+' or '-'.
    lines: tuple[str, ...]
    # A `\ No newline at end of file` marker followed a line this hunk *adds*,
    # so the result is meant to end without one. Honoured only when the hunk
    # actually reaches the end of the file -- git emits the marker for the old
    # side too, and reading that as an instruction would strip a newline the
    # patch never touched.
    ends_without_newline: bool = False

    @property
    def before(self) -> tuple[str, ...]:
        """The lines this hunk expects to find, without their markers."""
        return tuple(line[1:] for line in self.lines if line[:1] in (" ", "-"))

    @property
    def after(self) -> tuple[str, ...]:
        return tuple(line[1:] for line in self.lines if line[:1] in (" ", "+"))


def parse(patch: str) -> list[Hunk]:
    """Read a unified diff into hunks.

    File headers are tolerated and ignored -- `diff --git`, `index`, `---`,
    `+++` -- because models emit them by habit and refusing would cost a round
    trip to say so. The `@@` header is required: without one there is nothing to
    anchor against, and the resulting error is at least mechanical to fix.
    """
    hunks: list[Hunk] = []
    state: dict = {"header": None, "body": [], "bare": False}

    def flush() -> None:
        if state["header"] is None:
            return
        hunks.append(
            Hunk(
                *state["header"],
                lines=tuple(state["body"]),
                ends_without_newline=state["bare"],
            )
        )
        state["header"] = None
        state["body"] = []
        state["bare"] = False

    body = (patch or "").replace("\r\n", "\n").replace("\r", "\n").split("\n")
    # The patch's own final newline, not a blank context line. Without this every
    # well-formed patch acquires one phantom line of context at the end and
    # matches nothing -- which looks exactly like the model getting it wrong.
    if body and body[-1] == "":
        body.pop()

    for raw in body:
        matched = _HEADER.match(raw)
        if matched:
            flush()
            state["header"] = (
                int(matched.group(1)),
                int(matched.group(2) or 1),
                int(matched.group(3)),
                int(matched.group(4) or 1),
            )
            continue

        if state["header"] is None:
            # Preamble. Anything before the first @@ is a file header we do not
            # need: the path is a parameter, not something read out of the diff.
            continue

        if raw.startswith(_NO_NEWLINE):
            # It describes whichever side the line above belonged to. Only the
            # new side is an instruction; the old side is a description of the
            # file we are about to read for ourselves.
            if state["body"] and state["body"][-1][:1] in ("+", " "):
                state["bare"] = True
            continue
        if raw[:1] in ("+", "-", " "):
            state["body"].append(raw)
        elif raw == "":
            # A blank line that lost its leading space. Common enough to be the
            # normal case rather than an exceptional one.
            state["body"].append(" ")
        else:
            # A stray line inside a hunk -- a second `diff --git`, a signature.
            # Ends the hunk rather than corrupting it.
            flush()

    flush()

    if not hunks:
        raise PatchError(
            "That patch has no hunks. A patch needs at least one "
            "`@@ -old,count +new,count @@` header, followed by the lines to "
            "change: ' ' for context, '-' to remove, '+' to add."
        )
    if len(hunks) > MAX_HUNKS:
        raise PatchError(
            f"That patch has {len(hunks)} hunks, and {MAX_HUNKS} is the most "
            f"that will be applied at once. Rewrite the file with file_write "
            f"instead, or send the change in pieces."
        )
    return hunks


def _find(lines: list[str], wanted: tuple[str, ...], hint: int, floor: int) -> int:
    """Where `wanted` sits in `lines`, at or after `floor`. Raises if unclear."""
    if not wanted:
        # A pure insertion has no context to match. The hint is all there is.
        return max(floor, min(hint, len(lines)))

    span = len(wanted)
    if hint >= floor and lines[hint : hint + span] == list(wanted):
        return hint

    matches = [
        at
        for at in range(max(floor, hint - MAX_DRIFT), min(len(lines) - span, hint + MAX_DRIFT) + 1)
        if lines[at : at + span] == list(wanted)
    ]
    if len(matches) == 1:
        return matches[0]
    if len(matches) > 1:
        raise PatchError(
            f"Those context lines appear {len(matches)} times in the file, and "
            f"the line numbers in the hunk header do not point at any of them, "
            f"so there is no way to tell which was meant. Include more "
            f"unchanged lines around the change."
        )
    raise PatchError("")  # Filled in by the caller, which knows the hunk number.


def apply(text: str, hunks: list[Hunk]) -> str:
    """The file with every hunk applied, or a PatchError naming the first that
    would not.

    Hunks are applied in order against a cursor, so one cannot match inside
    territory an earlier one already consumed -- which is what a duplicated or
    overlapping hunk would otherwise do, applying the same change twice.
    """
    crlf = "\r\n" in text
    lines = text.replace("\r\n", "\n").replace("\r", "\n").split("\n")
    trailing = lines and lines[-1] == ""
    if trailing:
        lines.pop()

    out: list[str] = []
    cursor = 0
    reached_end = False

    for number, hunk in enumerate(hunks, start=1):
        wanted = hunk.before
        # A pure insertion names the line it goes *after*, not the line it
        # replaces, so it is not off by one the way every other hunk is.
        hint = hunk.old_start if hunk.old_count == 0 else max(hunk.old_start - 1, 0)
        try:
            at = _find(lines, wanted, hint, cursor)
        except PatchError as exc:
            raise _mismatch(number, hunk, lines, hint, exc.message) from None

        out.extend(lines[cursor:at])
        out.extend(hunk.after)
        cursor = at + len(wanted)
        reached_end = hunk.ends_without_newline and cursor >= len(lines)

    out.extend(lines[cursor:])

    result = "\n".join(out)
    if trailing and not reached_end:
        result += "\n"
    return result.replace("\n", "\r\n") if crlf else result


def _mismatch(number: int, hunk: Hunk, lines: list[str], hint: int, why: str) -> PatchError:
    """The message the model retries from, so it has to say what is actually
    there rather than only that something is wrong."""
    if why:
        return PatchError(
            f"Hunk {number} did not apply. {why} Nothing was written.", hunk=number
        )

    expected = next((line[1:] for line in hunk.lines if line[:1] in (" ", "-")), "")
    return PatchError(
        f"Hunk {number} did not apply. It expects line {hint + 1} to be\n"
        f"      {expected}\n"
        f"but the file has\n"
        f"{_around(lines, hint)}\n"
        f"and those lines are nowhere else nearby either. Nothing was written. "
        f"Send a patch whose context matches what is printed above.",
        hunk=number,
    )


# How many lines either side of the hinted position to print back. Three, which
# is what a patch carries as context, so a model can read its next attempt
# straight off the message.
MISMATCH_WINDOW = 3


def _around(lines: list[str], hint: int) -> str:
    """The file as it actually is, around where the hunk expected to land.

    One line was not enough. A model whose line numbers are two out reads "the
    file has X", cannot see where X sits relative to what it wanted, and sends
    the identical patch again -- which is most of the retry loop this tool
    produces in practice. Numbered, because the numbers are what was wrong.
    """
    if not lines:
        return "      (the file is empty)"
    if hint >= len(lines):
        start = max(0, len(lines) - MISMATCH_WINDOW)
        shown = [f"  {n + 1:>5}  {lines[n]}" for n in range(start, len(lines))]
        return "\n".join([*shown, f"      (the file ends at line {len(lines)})"])

    start = max(0, hint - MISMATCH_WINDOW)
    end = min(len(lines), hint + MISMATCH_WINDOW + 1)
    return "\n".join(
        f"{'->' if n == hint else '  '} {n + 1:>5}  {lines[n]}" for n in range(start, end)
    )


def render(before: str, after: str, path: str, *, max_lines: int = 200) -> str:
    """A unified diff of one change, for the transcript.

    Bounded here rather than at render time: this ends up in
    `Message.tool_calls_json`, which is on the row forever and re-parsed on
    every page load, and a generated file's diff can be larger than the file.
    """
    # splitlines, not split("\n"): a file's own final newline would otherwise be
    # an empty last element, which difflib renders as a stray context line at
    # the bottom of every diff -- and as a spurious change whenever one side has
    # it and the other does not. The trailing-newline difference is invisible
    # here as a result, which is right for a display and irrelevant to the write.
    lines = list(
        difflib.unified_diff(
            before.replace("\r\n", "\n").splitlines(),
            after.replace("\r\n", "\n").splitlines(),
            fromfile=f"a/{path}",
            tofile=f"b/{path}",
            lineterm="",
            n=3,
        )
    )
    if len(lines) > max_lines:
        dropped = len(lines) - max_lines
        lines = lines[:max_lines] + [f"… ({dropped} more lines)"]
    return "\n".join(lines)


__all__ = ["MAX_DRIFT", "MAX_HUNKS", "Hunk", "PatchError", "apply", "parse", "render"]