Apply model edits to a file with an AST and model routing

The setting: an AI code editor. A user asks for a change in chat or in an inline prompt. A model writes the change, usually as a sketch: the edited functions in full, with comments like # ... existing code ... standing in for the parts it left alone. Something must turn that sketch into an exact edit of the file and show it as a diff. A second part of the same question: many models are available (a fast small one, frontier ones), and each request must go to the right one.

Functional requirements

• Apply a model's sketch edit to the current file and show a reviewable diff within a few seconds.
• Never silently corrupt the file: the result parses, and anything the sketch did not mention is unchanged.
• Route each request (tab completion, inline edit, chat, agent step) to a model by latency budget, difficulty and cost, with fallbacks when a provider is slow or down.
• Learn from outcomes: accepted and rejected diffs feed back into routing and the apply model.

Non-functional requirements

• Inline edit: diff visible in under 2 s for a 400-line file. Tab completion: under 300 ms end to end.
• Apply must succeed on more than 99% of sketches (by falling back, not by guessing).
• Cost per request bounded: most requests should not touch the most expensive model.

Back-of-the-envelope estimation

• A 400-line file is about 4,000 tokens. A frontier model streams about 80 tokens/s, so rewriting the whole file costs 50 s. Too slow: the frontier model should only write the sketch (about 300 tokens, under 4 s).
• A small fine-tuned apply model that rewrites the full file at 250 tokens/s takes 16 s. With speculative decoding that drafts from the original file (the output mostly copies it), acceptance is about 95% on unchanged regions; at 16-token drafts that is about 11 tokens per forward pass, so roughly 2,000+ tokens/s: about 2 s.
• A structural (AST) apply takes milliseconds and costs nothing, but only works when the sketch is made of whole definitions.
• Load: 1M daily users × 50 inline or chat edits = 50M applies a day, about 600/s on average and 2,000/s at peak. If 60% apply structurally, the apply model serves 800/s at peak.

• POST /v1/edits {file_uri, file_text, sketch, language} → stream of {status} then {new_text, diff, method: "ast" | "model", checks}.
• POST /v1/complete {kind: "tab" | "inline" | "chat" | "agent", context, user_model_choice} → streamed tokens. The router is internal behind this call.
• POST /v1/feedback {edit_id, accepted, final_text} for learning.

• edit_events(edit_id, user_id, file_hash, sketch_hash, method, latency_ms, checks, accepted, created_at): the outcome log.
• routing_policies(version, rules, model_weights): the router's config, versioned for rollback and A/B tests.
• model_health(model, region, p50_ms, p99_ms, error_rate, queue_depth): refreshed every few seconds from serving metrics.

Build it one requirement at a time. Each drawing adds to the one before; the shaded boxes are what that step adds.

1) Sketch, then a model rewrites the file

The editor sends the request to a strong edit model, which writes a sketch. A fast apply model rewrites the whole file from the original plus the sketch, using speculative decoding that drafts from the original text. The editor shows the diff.

{"type":"excalidraw","elements":[
 {"type":"rectangle","id":"editor","x":0,"y":0,"width":180,"height":70,"label":{"text":"Editor"}},
 {"type":"rectangle","id":"edit","x":680,"y":0,"width":180,"height":70,"label":{"text":"Edit model"}},
 {"type":"rectangle","id":"apply","x":1020,"y":160,"width":180,"height":70,"label":{"text":"Apply model"}},
 {"type":"rectangle","id":"diff","x":1360,"y":160,"width":180,"height":70,"label":{"text":"Diff review"}},
 {"type":"arrow","id":"e1","x":180,"y":35,"width":500,"height":0,"start":{"id":"editor"},"end":{"id":"edit"},"label":{"text":"request"}},
 {"type":"arrow","id":"e2","x":860,"y":60,"width":160,"height":130,"start":{"id":"edit"},"end":{"id":"apply"},"label":{"text":"sketch"}},
 {"type":"arrow","id":"e3","x":1200,"y":195,"width":160,"height":0,"start":{"id":"apply"},"end":{"id":"diff"}}
]}

2) Apply structurally first, check every result

Most sketches replace whole functions or classes. An AST apply step matches definitions by name and splices them in, in milliseconds. Only sketches it cannot place go to the apply model. Every result passes a validator (parse, lint, and "unchanged outside the edited spans") before the diff is shown.

{"type":"excalidraw","elements":[
 {"type":"text","id":"legend","x":0,"y":-80,"text":"Shaded: added in this step. Dashed: fallback.","fontSize":16},
 {"type":"rectangle","id":"editor","x":0,"y":0,"width":180,"height":70,"label":{"text":"Editor"}},
 {"type":"rectangle","id":"edit","x":680,"y":0,"width":180,"height":70,"label":{"text":"Edit model"}},
 {"type":"rectangle","id":"ast","x":1020,"y":0,"width":180,"height":70,"label":{"text":"AST apply"},"backgroundColor":"#d0ebff","fillStyle":"solid"},
 {"type":"rectangle","id":"check","x":1360,"y":0,"width":180,"height":70,"label":{"text":"Validator"},"backgroundColor":"#d0ebff","fillStyle":"solid"},
 {"type":"rectangle","id":"apply","x":1020,"y":160,"width":180,"height":70,"label":{"text":"Apply model"}},
 {"type":"rectangle","id":"diff","x":1360,"y":160,"width":180,"height":70,"label":{"text":"Diff review"}},
 {"type":"arrow","id":"e1","x":180,"y":35,"width":500,"height":0,"start":{"id":"editor"},"end":{"id":"edit"},"label":{"text":"request"}},
 {"type":"arrow","id":"e2","x":860,"y":35,"width":160,"height":0,"start":{"id":"edit"},"end":{"id":"ast"},"label":{"text":"sketch"}},
 {"type":"arrow","id":"e4","x":1110,"y":70,"width":1,"height":90,"start":{"id":"ast"},"end":{"id":"apply"},"strokeStyle":"dashed"},
 {"type":"arrow","id":"e5","x":1200,"y":35,"width":160,"height":0,"start":{"id":"ast"},"end":{"id":"check"}},
 {"type":"arrow","id":"e6","x":1200,"y":170,"width":160,"height":-110,"start":{"id":"apply"},"end":{"id":"check"}},
 {"type":"arrow","id":"e7","x":1450,"y":70,"width":1,"height":90,"start":{"id":"check"},"end":{"id":"diff"}}
]}

3) Route each request to the right model

A router in front picks the model: easy requests (tab completion, small inline edits) go to a small fast model, hard ones (multi-file changes, agent steps) to a frontier model. It also pulls context from the codebase index. Both models produce sketches for the same apply path.

{"type":"excalidraw","elements":[
 {"type":"text","id":"legend","x":0,"y":-80,"text":"Shaded: added in this step. Dashed: fallback.","fontSize":16},
 {"type":"rectangle","id":"editor","x":0,"y":0,"width":180,"height":70,"label":{"text":"Editor"}},
 {"type":"rectangle","id":"router","x":340,"y":0,"width":180,"height":70,"label":{"text":"Router"},"backgroundColor":"#d0ebff","fillStyle":"solid"},
 {"type":"ellipse","id":"index","x":330,"y":170,"width":200,"height":80,"label":{"text":"Code index"},"backgroundColor":"#d0ebff","fillStyle":"solid"},
 {"type":"rectangle","id":"edit","x":680,"y":0,"width":180,"height":70,"label":{"text":"Edit model"}},
 {"type":"rectangle","id":"small","x":680,"y":160,"width":180,"height":70,"label":{"text":"Small model"},"backgroundColor":"#d0ebff","fillStyle":"solid"},
 {"type":"rectangle","id":"ast","x":1020,"y":0,"width":180,"height":70,"label":{"text":"AST apply"}},
 {"type":"rectangle","id":"check","x":1360,"y":0,"width":180,"height":70,"label":{"text":"Validator"}},
 {"type":"rectangle","id":"apply","x":1020,"y":160,"width":180,"height":70,"label":{"text":"Apply model"}},
 {"type":"rectangle","id":"diff","x":1360,"y":160,"width":180,"height":70,"label":{"text":"Diff review"}},
 {"type":"arrow","id":"e1","x":180,"y":35,"width":160,"height":0,"start":{"id":"editor"},"end":{"id":"router"},"label":{"text":"request"}},
 {"type":"arrow","id":"e8","x":430,"y":70,"width":1,"height":100,"start":{"id":"router"},"end":{"id":"index"}},
 {"type":"arrow","id":"e9","x":520,"y":35,"width":160,"height":0,"start":{"id":"router"},"end":{"id":"edit"},"label":{"text":"hard"}},
 {"type":"arrow","id":"e10","x":520,"y":60,"width":160,"height":130,"start":{"id":"router"},"end":{"id":"small"},"label":{"text":"easy"}},
 {"type":"arrow","id":"e2","x":860,"y":35,"width":160,"height":0,"start":{"id":"edit"},"end":{"id":"ast"},"label":{"text":"sketch"}},
 {"type":"arrow","id":"e11","x":860,"y":180,"width":160,"height":-120,"start":{"id":"small"},"end":{"id":"ast"}},
 {"type":"arrow","id":"e4","x":1110,"y":70,"width":1,"height":90,"start":{"id":"ast"},"end":{"id":"apply"},"strokeStyle":"dashed"},
 {"type":"arrow","id":"e5","x":1200,"y":35,"width":160,"height":0,"start":{"id":"ast"},"end":{"id":"check"}},
 {"type":"arrow","id":"e6","x":1200,"y":170,"width":160,"height":-110,"start":{"id":"apply"},"end":{"id":"check"}},
 {"type":"arrow","id":"e7","x":1450,"y":70,"width":1,"height":90,"start":{"id":"check"},"end":{"id":"diff"}}
]}

Final design

{"type":"excalidraw","elements":[
 {"type":"text","id":"legend","x":0,"y":-80,"text":"Dashed: fallback when the AST apply cannot place the sketch.","fontSize":16},
 {"type":"rectangle","id":"editor","x":0,"y":0,"width":180,"height":70,"label":{"text":"Editor"}},
 {"type":"rectangle","id":"router","x":340,"y":0,"width":180,"height":70,"label":{"text":"Router"}},
 {"type":"ellipse","id":"index","x":330,"y":170,"width":200,"height":80,"label":{"text":"Code index"}},
 {"type":"rectangle","id":"edit","x":680,"y":0,"width":180,"height":70,"label":{"text":"Edit model"}},
 {"type":"rectangle","id":"small","x":680,"y":160,"width":180,"height":70,"label":{"text":"Small model"}},
 {"type":"rectangle","id":"ast","x":1020,"y":0,"width":180,"height":70,"label":{"text":"AST apply"}},
 {"type":"rectangle","id":"check","x":1360,"y":0,"width":180,"height":70,"label":{"text":"Validator"}},
 {"type":"rectangle","id":"apply","x":1020,"y":160,"width":180,"height":70,"label":{"text":"Apply model"}},
 {"type":"rectangle","id":"diff","x":1360,"y":160,"width":180,"height":70,"label":{"text":"Diff review"}},
 {"type":"arrow","id":"e1","x":180,"y":35,"width":160,"height":0,"start":{"id":"editor"},"end":{"id":"router"},"label":{"text":"request"}},
 {"type":"arrow","id":"e8","x":430,"y":70,"width":1,"height":100,"start":{"id":"router"},"end":{"id":"index"}},
 {"type":"arrow","id":"e9","x":520,"y":35,"width":160,"height":0,"start":{"id":"router"},"end":{"id":"edit"},"label":{"text":"hard"}},
 {"type":"arrow","id":"e10","x":520,"y":60,"width":160,"height":130,"start":{"id":"router"},"end":{"id":"small"},"label":{"text":"easy"}},
 {"type":"arrow","id":"e2","x":860,"y":35,"width":160,"height":0,"start":{"id":"edit"},"end":{"id":"ast"},"label":{"text":"sketch"}},
 {"type":"arrow","id":"e11","x":860,"y":180,"width":160,"height":-120,"start":{"id":"small"},"end":{"id":"ast"}},
 {"type":"arrow","id":"e4","x":1110,"y":70,"width":1,"height":90,"start":{"id":"ast"},"end":{"id":"apply"},"strokeStyle":"dashed"},
 {"type":"arrow","id":"e5","x":1200,"y":35,"width":160,"height":0,"start":{"id":"ast"},"end":{"id":"check"}},
 {"type":"arrow","id":"e6","x":1200,"y":170,"width":160,"height":-110,"start":{"id":"apply"},"end":{"id":"check"}},
 {"type":"arrow","id":"e7","x":1450,"y":70,"width":1,"height":90,"start":{"id":"check"},"end":{"id":"diff"}}
]}

The structural apply

Parse both the file and the sketch. Walk the sketch's top-level definitions: a function or class whose name exists in the file replaces that definition's line span (decorators included); a class merges member by member, so a sketch that shows one method of a 40-method class touches only that method; a new name is appended. Comments like # ... existing code ... disappear in the syntax tree, which is exactly what we want. Anything the matcher cannot place (a bare statement, a changed import block, a sketch that does not parse) raises, and the request falls back to the apply model. The rule is: apply structurally only when the result is unambiguous; otherwise hand off.

In production this runs on tree-sitter, which parses dozens of languages and tolerates syntax errors in the file being edited. Here is the idea in Python with the standard ast module:

import ast
import textwrap
from typing import List, Tuple

DEFS = (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)


class NoMatch(Exception):
    """The snippet can't be applied structurally; fall back to the apply model."""


def _span(node: ast.AST) -> Tuple[int, int]:
    """1-based inclusive line span, decorators included."""
    start = min([node.lineno] + [d.lineno for d in getattr(node, "decorator_list", [])])
    return start, node.end_lineno


def _text(lines: List[str], node: ast.AST, indent: int) -> List[str]:
    start, end = _span(node)
    block = textwrap.dedent("\n".join(lines[start - 1:end]))
    return textwrap.indent(block, " " * indent).split("\n")


def _is_filler(node: ast.AST) -> bool:
    """`...` or `pass` that a model writes in place of omitted code."""
    return isinstance(node, ast.Pass) or (
        isinstance(node, ast.Expr) and isinstance(node.value, ast.Constant) and node.value.value is Ellipsis)


def apply_ast_edit(source: str, snippet: str) -> str:
    """Merge a partial-file edit into `source` by matching definitions by name.

    Functions and classes in the snippet replace the same-named definitions
    (classes merge member by member); new names are appended. Anything else
    (top-level statements, class attributes) raises NoMatch.
    """
    try:
        src_tree, snip_tree = ast.parse(source), ast.parse(snippet)
    except SyntaxError as e:
        raise NoMatch(f"does not parse: {e}") from None
    src_lines, snip_lines = source.split("\n"), snippet.split("\n")
    edits: List[Tuple[int, int, List[str]]] = []  # (start, end, new lines); end < start = insert

    def merge(src_body, snip_body, indent: int, insert_at: int, gap: int) -> None:
        targets = {n.name: n for n in src_body if isinstance(n, DEFS)}
        added: List[str] = []
        for node in snip_body:
            if _is_filler(node):
                continue
            if not isinstance(node, DEFS):
                raise NoMatch(f"statement at line {node.lineno} is not a definition")
            target = targets.get(node.name)
            if target is None:
                added += [""] * gap + _text(snip_lines, node, indent)
            elif isinstance(node, ast.ClassDef) and isinstance(target, ast.ClassDef):
                body_indent = target.body[0].col_offset
                merge(target.body, node.body, body_indent, target.end_lineno + 1, 1)
            else:
                start, end = _span(target)
                edits.append((start, end, _text(snip_lines, node, target.col_offset)))
        if added:
            edits.append((insert_at, insert_at - 1, added))

    last = len(src_lines)
    while last > 0 and not src_lines[last - 1].strip():
        last -= 1
    merge(src_tree.body, snip_tree.body, 0, last + 1, 2)

    lines = list(src_lines)
    # Bottom-up so earlier line numbers stay valid; at equal positions the
    # later (outer) insert goes first, so inner additions land above it.
    for _, (start, end, new) in sorted(enumerate(edits), key=lambda ie: (-ie[1][0], -ie[0])):
        lines[start - 1:end] = new  # works for inserts too (end = start - 1)
    result = "\n".join(lines).rstrip("\n") + "\n"
    ast.parse(result)  # validate before showing the diff
    return result

_SRC = (
    "import os\n\n\n"
    "class Store:\n"
    "    def __init__(self):\n"
    "        self.items = {}\n\n"
    "    def get(self, key):\n"
    "        return self.items.get(key)\n\n\n"
    "@cache\n"
    "def helper(x):\n"
    "    return x + 1\n"
)


def test_replaces_one_method_and_adds_another():
    sketch = (
        "class Store:\n"
        "    # ... existing code ...\n"
        "    def get(self, key, default=None):\n"
        "        return self.items.get(key, default)\n\n"
        "    def put(self, key, value):\n"
        "        self.items[key] = value\n"
    )
    out = apply_ast_edit(_SRC, sketch)
    assert "def __init__(self):" in out                     # untouched
    assert "return self.items.get(key, default)" in out
    assert out.index("def put") < out.index("@cache")       # added inside the class
    assert "return x + 1" in out


def test_replaces_decorated_function_and_appends_new_one():
    sketch = "@lru_cache(None)\ndef helper(x):\n    return x * 2\n\n\ndef extra():\n    return 3\n"
    out = apply_ast_edit(_SRC, sketch)
    assert "@cache" not in out and "@lru_cache(None)" in out and "return x * 2" in out
    assert out.rstrip().endswith("return 3")


def test_unplaceable_sketch_falls_back():
    for sketch in ("x = 3\n", "def broken(:\n"):
        try:
            apply_ast_edit(_SRC, sketch)
            assert False
        except NoMatch:
            pass

The apply model

When the structural step gives up, a small model (fine-tuned on pairs of (original file, sketch) -> edited file) rewrites the whole file. Two choices make it fast enough:

• Full rewrite, not a diff format. Models are poor at line numbers and at reproducing search/replace anchors exactly. Rewriting the whole file is the easiest format for a model to get right, and speculative decoding makes the copying nearly free.
• Speculative decoding from the original. The draft for the next tokens is "whatever followed this position in the original file". Unchanged regions are accepted 16 tokens at a time; only the edited lines are generated token by token. A 4,000-token file goes from 16 s to about 2 s.

For files over about 1,500 lines, rewrite only the region around the sketch (the enclosing definitions plus a margin) and splice it back, to cap latency and cost.

Validation

Before the user sees the diff:

1. The result parses (tree-sitter, or the language server).
2. Lines outside the edited definitions are byte-identical to the original. If the apply model changed something the sketch did not mention, show a warning or retry.
3. Optional fast checks: linter errors introduced by the edit, type errors from the language server. Run them in parallel with showing the diff and annotate it, rather than blocking.

The router

Inputs per request: the surface (tab, inline, chat, agent), prompt and context size, the user's explicit model choice, plan limits, a difficulty score, and live model health.

• Hard constraints first. Tab completion has a 300 ms budget, so only the small in-house model qualifies. An explicit user choice wins unless that model is down.
• Difficulty classifier. A cheap classifier (or the small model itself) scores the request: rename a variable versus refactor a module across files. Train it on outcomes: requests where the small model's diff was rejected and the user retried on a bigger model are positive examples for "hard".
• Cost arithmetic. If a frontier call costs 10× a small one and 60% of chat requests go to the small model with the same acceptance rate, cost falls to 0.4 + 0.6 × 0.1 = 46% of all-frontier.
• Fallbacks. Health data (p99 latency, error rate, queue depth) feeds a circuit breaker per model and region. When a provider degrades, route to the next model in the same quality tier, then to a lower tier with a notice.
• Evaluation. Change routing policies only through A/B tests on acceptance rate, retry rate, latency and cost per accepted edit.

• AST apply versus apply model. The AST path is instant, free and exact, but only covers whole-definition edits in languages with a parser. The model path covers everything at a cost of about 2 s and a small risk of drift, which validation catches.
• Full rewrite versus diff format. Diffs are cheaper to generate but fail often on anchors and line numbers. Full rewrites with speculative decoding are reliable and fast enough up to a few thousand lines.
• Static rules versus learned routing. Rules are predictable and easy to debug; a learned classifier saves more money but needs outcome data and guardrails against routing hard requests to weak models.

Mid-level (L4)

• Separates "write the change" from "apply the change" and explains why the frontier model should not rewrite whole files.
• Proposes a working apply approach (structural or model) and validates that the result parses.
• Routes by surface and latency budget with a basic fallback.

Senior (L5)

• Gets to the hybrid: structural apply when unambiguous, apply model otherwise, with speculative decoding from the original file and the latency numbers to back it up.
• Defines validation precisely (parse, unchanged outside edited spans, async lint) and what happens when it fails.
• Designs routing with hard constraints, a difficulty signal, health-based circuit breakers and the cost arithmetic.

Staff (L6+)

• Closes the loop: accepted and rejected diffs train the apply model and the routing classifier; policies change through A/B tests on cost per accepted edit.
• Handles scale and edge cases: huge files (region rewrite), many languages (tree-sitter), partial syntax errors in the user's file, provider outages.
• Explains which numbers decide each trade-off and how they would measure them before committing.