logiguard fork: GPUI xdg-activation keyboard-focus serial fix

Single-commit orphan branch: full zed-industries/zed @ 8c74db0 source tree
with a 3-file patch applied (no upstream history).

Patch (crates/gpui_linux/src/linux/wayland/):
  - serial.rs: add SerialKind::KeyboardEnter
  - client.rs: store wl_keyboard.enter serial; add latest_serial_of()
  - window.rs: activate() uses keyboard-enter serial (Mutter focus gate)

Mutter honors window activation only when the token carries the keyboard-
focus serial from wl_keyboard.enter; GPUI used a stale mouse-press serial.
See docs/tray-window-focus-wayland.md in logiguard.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

crates/edit_prediction_cli/Cargo.toml

@@ -0,0 +1,93 @@
+[package]
+name = "edit_prediction_cli"
+version = "0.1.0"
+edition.workspace = true
+publish.workspace = true
+license = "GPL-3.0-or-later"
+
+[lints]
+workspace = true
+
+[lib]
+path = "src/lib.rs"
+
+[[bin]]
+name = "ep"
+path = "src/main.rs"
+
+[dependencies]
+anyhow.workspace = true
+anthropic.workspace = true
+http_client.workspace = true
+chrono.workspace = true
+clap = "4"
+client.workspace = true
+cloud_llm_client = { workspace = true, features = ["predict-edits"] }
+collections.workspace = true
+db.workspace = true
+debug_adapter_extension.workspace = true
+dirs.workspace = true
+extension.workspace = true
+fs.workspace = true
+futures.workspace = true
+gpui.workspace = true
+gpui_platform.workspace = true
+gpui_tokio.workspace = true
+indoc.workspace = true
+language.workspace = true
+language_extension.workspace = true
+language_model.workspace = true
+language_models.workspace = true
+languages = { workspace = true, features = ["load-grammars"] }
+libc.workspace = true
+log.workspace = true
+node_runtime.workspace = true
+open_ai.workspace = true
+
+paths.workspace = true
+project.workspace = true
+prompt_store.workspace = true
+release_channel.workspace = true
+reqwest_client.workspace = true
+serde.workspace = true
+serde_json.workspace = true
+settings.workspace = true
+shellexpand.workspace = true
+smol.workspace = true
+sqlez.workspace = true
+strum.workspace = true
+sqlez_macros.workspace = true
+terminal_view.workspace = true
+util.workspace = true
+watch.workspace = true
+edit_prediction = { workspace = true, features = ["cli-support"] }
+edit_prediction_metrics.workspace = true
+telemetry_events.workspace = true
+wasmtime.workspace = true
+zeta_prompt.workspace = true
+rand.workspace = true
+similar = "2.7.0"
+flate2 = "1.1.8"
+toml.workspace = true
+rust-embed.workspace = true
+gaoya = "0.2.0"
+
+# Wasmtime is included as a dependency in order to enable the same
+# features that are enabled in Zed.
+#
+# If we don't enable these features we get crashes when creating
+# a Tree-sitter WasmStore.
+[features]
+dynamic_prompts = []
+
+[package.metadata.cargo-machete]
+ignored = ["wasmtime"]
+
+[dev-dependencies]
+gpui = { workspace = true, features = ["test-support"] }
+indoc.workspace = true
+pretty_assertions.workspace = true
+project = { workspace = true, features = ["test-support"] }
+tempfile.workspace = true
+workspace = { workspace = true, features = ["test-support"] }
+

crates/edit_prediction_cli/LICENSE-GPL

@@ -0,0 +1 @@
+../../LICENSE-GPL

crates/edit_prediction_cli/build.rs

@@ -0,0 +1,14 @@
+fn main() {
+    let cargo_toml =
+        std::fs::read_to_string("../zed/Cargo.toml").expect("Failed to read crates/zed/Cargo.toml");
+    let version = cargo_toml
+        .lines()
+        .find(|line| line.starts_with("version = "))
+        .expect("Version not found in crates/zed/Cargo.toml")
+        .split('=')
+        .nth(1)
+        .expect("Invalid version format")
+        .trim()
+        .trim_matches('"');
+    println!("cargo:rustc-env=ZED_PKG_VERSION={}", version);
+}

crates/edit_prediction_cli/evals/.zed/settings.json

@@ -0,0 +1,4 @@
+{
+  "remove_trailing_whitespace_on_save": false,
+  "soft_wrap": "none",
+}

crates/edit_prediction_cli/evals/codex-acp--add-derive.md

@@ -0,0 +1,69 @@
++++
+repository_url = "https://github.com/zed-industries/codex-acp"
+revision = "c3d24ee70928fc9da08c131fc632d624413ccc43"
++++
+
+## Edit History
+
+```diff
+--- a/src/prompt_args.rs
++++ b/src/prompt_args.rs
+@@ -28,7 +28,7 @@ impl PromptArgsError {
+     }
+ }
+
+-#[derive(Debug)]
++#[derive(Debug, Serialize)]
+ pub enum PromptExpansionError {
+     Args {
+         command: String,
+```
+
+## Cursor Position
+
+```src/prompt_args.rs
+#[derive(Debug)]
+pub enum PromptArgsError {
+//                      ^[CURSOR_POSITION]
+    MissingAssignment { token: String },
+    MissingKey { token: String },
+}
+```
+
+## Expected Patch
+
+```diff
+--- a/src/prompt_args.rs
++++ b/src/prompt_args.rs
+@@ -9,7 +9,7 @@ use std::sync::LazyLock;
+ static PROMPT_ARG_REGEX: LazyLock<Regex> =
+     LazyLock::new(|| Regex::new(r"\$[A-Z][A-Z0-9_]*").unwrap_or_else(|_| std::process::abort()));
+
+-#[derive(Debug)]
++#[derive(Debug, Serialize)]
+ pub enum PromptArgsError {
+     MissingAssignment { token: String },
+     MissingKey { token: String },
+```
+
+```diff
+--- a/src/prompt_args.rs
++++ b/src/prompt_args.rs
+@@ -3,19 +3,20 @@
+ use regex_lite::Regex;
++use serde::Serialize;
+ use shlex::Shlex;
+ use std::collections::HashMap;
+ use std::collections::HashSet;
+ use std::sync::LazyLock;
+
+ static PROMPT_ARG_REGEX: LazyLock<Regex> =
+     LazyLock::new(|| Regex::new(r"\$[A-Z][A-Z0-9_]*").unwrap_or_else(|_| std::process::abort()));
+
+-#[derive(Debug)]
++#[derive(Debug, Serialize)]
+ pub enum PromptArgsError {
+     MissingAssignment { token: String },
+     MissingKey { token: String },
+ }
+```

crates/edit_prediction_cli/evals/flask--add-and-rename-test-function.md

@@ -0,0 +1,98 @@
++++
+repository_url = "https://github.com/pallets/flask"
+revision = "2fec0b206c6e83ea813ab26597e15c96fab08be7"
++++
+
+## Edit History
+
+```diff
+--- a/tests/test_basic.py
++++ b/tests/test_basic.py
+@@ -1376,5 +1376,8 @@
+ def test_static_files(app, client):
+     rv = client.get("/static/index.html")
+     assert rv.status_code == 200
+     assert rv.data.strip() == b"<h1>Hello World!</h1>"
+     with app.test_request_context():
+         assert flask.url_for("static", filename="index.html") == "/static/index.html"
+     rv.close()
+
+
++de
++
++
+ def test_static_url_path():
+--- a/tests/test_basic.py
++++ b/tests/test_basic.py
+@@ -1372,15 +1372,15 @@
+-de
++def test_():
++    pass
+
+
+ def test_static_url_path():
+--- a/tests/test_basic.py
++++ b/tests/test_basic.py
+@@ -1372,15 +1372,15 @@
+-def test_():
++def test_static_file_not_found():
+     pass
+
+
+ def test_static_url_path():
+```
+
+## Cursor Position
+
+```tests/test_basic.py
+def test_static_file_not_found():
+#                             ^[CURSOR_POSITION]
+    pass
+
+
+def test_static_url_path():
+```
+
+## Expected Patch
+
+```diff
+--- a/tests/test_basic.py
++++ b/tests/test_basic.py
+@@ -1372,15 +1372,15 @@
+-def test_static_file_not_found():
+-    pass
++def test_static_file_not_found(app, client):
++    rv = client.get("/static/non_existent.html")
++    assert rv.status_code == 404
++    rv.close()
+```
+
+```diff
+--- a/tests/test_basic.py
++++ b/tests/test_basic.py
+@@ -1372,15 +1372,15 @@
+-def test_static_file_not_found():
+-    pass
++def test_static_file_not_found(app, client):
++    rv = client.get("/static/not_found.html")
++    assert rv.status_code == 404
++    rv.close()
+```
+
+```diff
+--- a/tests/test_basic.py
++++ b/tests/test_basic.py
+@@ -1376,8 +1376,13 @@
+-def test_static_file_not_found():
+-    pass
++def test_static_file_not_found(app, client):
++    rv = client.get("/static/nonexistent.html")
++    assert rv.status_code == 404
++    assert rv.data.strip() == b"<h1>Not Found</h1>"
++    with app.test_request_context():
++        pytest.raises(BuildError, flask.url_for, "static", filename="nonexistent.html")
++    rv.close()
+
+
+ def test_static_url_path():
+```

crates/edit_prediction_cli/evals/flask--add-import-statement.md

@@ -0,0 +1,76 @@
++++
+repository_url = "https://github.com/pallets/flask"
+revision = "2fec0b206c6e83ea813ab26597e15c96fab08be7"
++++
+
+## Edit History
+
+```diff
+--- a/src/flask/logging.py
++++ b/src/flask/logging.py
+@@ -4,7 +4,7 @@
+ import sys
+ import typing as t
+
+-from werkzeug.local import LocalProxy
++imfrom werkzeug.local import LocalProxy
+
+ from .globals import request
+
+```
+
+## Cursor Position
+
+```src/flask/logging.py
+from __future__ import annotations
+
+import logging
+import sys
+import typing as t
+
+imfrom werkzeug.local import LocalProxy
+# ^[CURSOR_POSITION]
+
+from .globals import request
+
+if t.TYPE_CHECKING:  # pragma: no cover
+    from .sansio.app import App
+```
+
+## Expected Patch
+
+```diff
+--- a/src/flask/logging.py
++++ b/src/flask/logging.py
+@@ -1,21 +1,21 @@
+ from __future__ import annotations
+
+ import logging
+ import sys
+ import typing as t
+
+-imfrom werkzeug.local import LocalProxy
++import
+#       ^[CURSOR_POSITION]
++from werkzeug.local import LocalProxy
+
+ from .globals import request
+```
+
+```diff
+--- a/src/flask/logging.py
++++ b/src/flask/logging.py
+@@ -1,21 +1,21 @@
+ from __future__ import annotations
+
+ import logging
+ import sys
+ import typing as t
+-
+-imfrom werkzeug.local import LocalProxy
++import werkzeug
+#               ^[CURSOR_POSITION]
++from werkzeug.local import LocalProxy
+
+ from .globals import request
+```

crates/edit_prediction_cli/evals/flask--add-test-function.md

@@ -0,0 +1,177 @@
++++
+repository_url = "https://github.com/pallets/flask"
+revision = "2fec0b206c6e83ea813ab26597e15c96fab08be7"
++++
+
+## Edit History
+
+```diff
+--- a/tests/test_basic.py
++++ b/tests/test_basic.py
+@@ -1376,5 +1376,8 @@
+ def test_static_files(app, client):
+     rv = client.get("/static/index.html")
+     assert rv.status_code == 200
+     assert rv.data.strip() == b"<h1>Hello World!</h1>"
+     with app.test_request_context():
+         assert flask.url_for("static", filename="index.html") == "/static/index.html"
+     rv.close()
+
+
++de
++
++
+ def test_static_url_path():
+```
+
+## Cursor Position
+
+```tests/test_basic.py
+def test_static_files(app, client):
+    rv = client.get("/static/index.html")
+    assert rv.status_code == 200
+    assert rv.data.strip() == b"<h1>Hello World!</h1>"
+    with app.test_request_context():
+        assert flask.url_for("static", filename="index.html") == "/static/index.html"
+    rv.close()
+
+
+de
+# ^[CURSOR_POSITION]
+
+
+def test_static_url_path():
+```
+
+## Expected Patch
+
+```diff
+--- a/tests/test_basic.py
++++ b/tests/test_basic.py
+@@ -1372,15 +1372,15 @@
+-de
++def test_():
+#         ^[CURSOR_POSITION]
++    
+
+
+def test_static_url_path():
+```
+
+```diff
+--- a/tests/test_basic.py
++++ b/tests/test_basic.py
+@@ -1372,15 +1372,15 @@
+-de
++def test_():
+#         ^[CURSOR_POSITION]
++    pass
+
+
+def test_static_url_path():
+```
+
+
+```diff
+--- a/tests/test_basic.py
++++ b/tests/test_basic.py
+@@ -1372,15 +1372,15 @@
+-de
++def test_(app, client):
+#         ^[CURSOR_POSITION]
++    
+
+
+def test_static_url_path():
+```
+
+```diff
+--- a/tests/test_basic.py
++++ b/tests/test_basic.py
+@@ -1372,15 +1372,15 @@
+-de
++def test_(app, client):
+#         ^[CURSOR_POSITION]
++    pass
+
+
+def test_static_url_path():
+```
+
+```diff
+--- a/tests/test_basic.py
++++ b/tests/test_basic.py
+@@ -1372,15 +1372,15 @@
+-de
++def test_static_():
+#                ^[CURSOR_POSITION]
++    
+
+
+def test_static_url_path():
+```
+
+```diff
+--- a/tests/test_basic.py
++++ b/tests/test_basic.py
+@@ -1372,15 +1372,15 @@
+-de
++def test_static_():
+#                ^[CURSOR_POSITION]
++    pass
+
+
+def test_static_url_path():
+```
+
+```diff
+--- a/tests/test_basic.py
++++ b/tests/test_basic.py
+@@ -1372,15 +1372,15 @@
+-de
++def test_static_(app, client):
+#                ^[CURSOR_POSITION]
++    
+
+
+def test_static_url_path():
+```
+
+```diff
+--- a/tests/test_basic.py
++++ b/tests/test_basic.py
+@@ -1372,15 +1372,15 @@
+-de
++def test_static_(app, client):
+#                ^[CURSOR_POSITION]
++    pass
+
+
+def test_static_url_path():
+```
+
+```diff
+--- a/tests/test_basic.py
++++ b/tests/test_basic.py
+@@ -1372,15 +1372,15 @@
+-de
++def test_static_folder():
+#                       ^[CURSOR_POSITION]
++    pass
+
+
+def test_static_url_path():
+```
+
+```diff
+--- a/tests/test_basic.py
++++ b/tests/test_basic.py
+@@ -1372,15 +1372,15 @@
+-de
++def test_static_route_with_host_matching(app, client):
++    
+#    ^[CURSOR_POSITION]
+
+
+def test_static_url_path():
+```

crates/edit_prediction_cli/evals/flask--rename-accepted-prediction.md

@@ -0,0 +1,98 @@
++++
+repository_url = "https://github.com/pallets/flask"
+revision = "2fec0b206c6e83ea813ab26597e15c96fab08be7"
++++
+
+## Edit History
+
+```diff
+--- a/tests/test_basic.py
++++ b/tests/test_basic.py
+@@ -356,3 +356,6 @@
+     cookie = rv.headers["set-cookie"].lower()
+     assert "samesite=lax" in cookie
+
+
++de
++
++
+ def test_missing_session(app):
+```
+
+// User accepted prediction:
+```diff
+--- a/tests/test_basic.py
++++ b/tests/test_basic.py
+@@ -358,6 +358,14 @@
+
+
+-de
++def test_session_cookie_httponly(app, client):
++    app.config["SESSION_COOKIE_HTTPONLY"] = True
++
++    @app.route("/")
++    def index():
++        flask.session["testing"] = 42
++        return "Hello World"
++
++    rv = client.get("/")
++    assert "httponly" in rv.headers["set-cookie"].lower()
+
+
+ def test_missing_session(app):
+```
+
+```diff
+--- a/tests/test_basic.py
++++ b/tests/test_basic.py
+@@ -358,14 +358,14 @@
+
+
+-def test_session_cookie_httponly(app, client):
++def test_session_cookie_secur(app, client):
+     app.config["SESSION_COOKIE_HTTPONLY"] = True
+```
+
+## Cursor Position
+
+```tests/test_basic.py
+    cookie = rv.headers["set-cookie"].lower()
+    assert "samesite=lax" in cookie
+
+
+def test_session_cookie_secur(app, client):
+#                            ^[CURSOR_POSITION]
+    app.config["SESSION_COOKIE_HTTPONLY"] = True
+
+    @app.route("/")
+    def index():
+        flask.session["testing"] = 42
+        return "Hello World"
+
+    rv = client.get("/")
+    assert "httponly" in rv.headers["set-cookie"].lower()
+
+
+def test_missing_session(app):
+```
+
+## Expected Patch
+
+```diff
+--- a/tests/test_basic.py
++++ b/tests/test_basic.py
+@@ -358,14 +358,14 @@
+-def test_session_cookie_secur(app, client):
+-    app.config["SESSION_COOKIE_HTTPONLY"] = True
++def test_session_cookie_secure(app, client):
++    app.config["SESSION_COOKIE_SECURE"] = True
+
+     @app.route("/")
+     def index():
+         flask.session["testing"] = 42
+         return "Hello World"
+
+     rv = client.get("/")
+-    assert "httponly" in rv.headers["set-cookie"].lower()
++    assert "secure" in rv.headers["set-cookie"].lower()
+```

crates/edit_prediction_cli/evals/hello-world--rename-accepted-group-by.md

@@ -0,0 +1,81 @@
++++
+repository_url = "https://github.com/octocat/hello-world"
+revision = "7fd1a60b01f91b314f59955a4e4d4e80d8edf11d"
++++
+
+## Edit History
+
+```diff
+--- a/README
++++ b/README
+@@ -1,1 +1,6 @@
+-Hello World!
++function filterByStatus(items, status) {
++    return items.filter(item => item.status === status);
++}
++
++function groupBy
++
+```
+
+// User accepted prediction:
+```diff
+--- a/README
++++ b/README
+@@ -4,3 +4,9 @@
+ 
+-function groupBy
++function groupByStatus(items) {
++    return items.reduce((groups, item) => {
++        const key = item.status;
++        (groups[key] = groups[key] || []).push(item);
++        return groups;
++    }, {});
++}
+ 
+```
+
+```diff
+--- a/README
++++ b/README
+@@ -4,4 +4,4 @@
+ 
+-function groupByStatus(items) {
++function groupByCat(items) {
+     return items.reduce((groups, item) => {
+```
+
+## Cursor Position
+
+```README
+function filterByStatus(items, status) {
+    return items.filter(item => item.status === status);
+}
+
+function groupByCat(items) {
+#                  ^[CURSOR_POSITION]
+    return items.reduce((groups, item) => {
+        const key = item.status;
+        (groups[key] = groups[key] || []).push(item);
+        return groups;
+    }, {});
+}
+
+```
+
+## Expected Patch
+
+```diff
+--- a/README
++++ b/README
+@@ -5,7 +5,7 @@
+-function groupByCat(items) {
++function groupByCategory(items) {
+#                        ^[CURSOR_POSITION]
+     return items.reduce((groups, item) => {
+-        const key = item.status;
++        const key = item.category;
+         (groups[key] = groups[key] || []).push(item);
+         return groups;
+     }, {});
+```

crates/edit_prediction_cli/evals/terraform--add-comment.md

@@ -0,0 +1,140 @@
++++
+repository_url = "https://github.com/hashicorp/terraform"
+revision = "a3dc571150a7651a1a4a8b302342d26089c97795"
++++
+
+## Edit History
+
+```diff
+--- a/internal/actions/actions.go
++++ b/internal/actions/actions.go
+@@ -63,6 +63,7 @@
+ 	a.mu.Lock()
+ 	defer a.mu.Unlock()
+
++	/
+ 	result := []addrs.AbsActionInstance{}
+ 	for _, data := range a.actionInstances.Elements() {
+ 		if data.Key.ContainingAction().Equal(addr) {
+```
+
+## Cursor Position
+
+```internal/actions/actions.go
+	defer a.mu.Unlock()
+
+	data, ok := a.actionInstances.GetOk(addr)
+
+	if !ok {
+		return nil, false
+	}
+
+	return &data, true
+}
+
+func (a *Actions) GetActionInstanceKeys(addr addrs.AbsAction) []addrs.AbsActionInstance {
+	a.mu.Lock()
+	defer a.mu.Unlock()
+
+	/
+   // <[CURSOR_POSITION]
+	result := []addrs.AbsActionInstance{}
+	for _, data := range a.actionInstances.Elements() {
+		if data.Key.ContainingAction().Equal(addr) {
+			result = append(result, data.Key)
+		}
+	}
+
+	return result
+}
+```
+
+## Expected Patch
+
+```diff
+--- a/internal/actions/actions.go
++++ b/internal/actions/actions.go
+@@ -51,26 +51,26 @@
+ func (a *Actions) GetActionInstanceKeys(addr addrs.AbsAction) []addrs.AbsActionInstance {
+ 	a.mu.Lock()
+ 	defer a.mu.Unlock()
+
+-	/
++	// Filter action instances by the given action.
+ 	result := []addrs.AbsActionInstance{}
+ 	for _, data := range a.actionInstances.Elements() {
+ 		if data.Key.ContainingAction().Equal(addr) {
+ 			result = append(result, data.Key)
+ 		}
+ 	}
+```
+
+```diff
+--- a/internal/actions/actions.go
++++ b/internal/actions/actions.go
+@@ -54,25 +54,25 @@
+ func (a *Actions) GetActionInstanceKeys(addr addrs.AbsAction) []addrs.AbsActionInstance {
+ 	a.mu.Lock()
+ 	defer a.mu.Unlock()
+
+-	/
++	// Filter action instances that belong to the given action
+ 	result := []addrs.AbsActionInstance{}
+ 	for _, data := range a.actionInstances.Elements() {
+ 		if data.Key.ContainingAction().Equal(addr) {
+ 			result = append(result, data.Key)
+ 		}
+ 	}
+```
+
+```diff
+--- a/internal/actions/actions.go
++++ b/internal/actions/actions.go
+@@ -54,25 +54,25 @@
+ func (a *Actions) GetActionInstanceKeys(addr addrs.AbsAction) []addrs.AbsActionInstance {
+ 	a.mu.Lock()
+ 	defer a.mu.Unlock()
+
+-	/
++	// Iterate through all action instances and filter by the containing action
+ 	result := []addrs.AbsActionInstance{}
+ 	for _, data := range a.actionInstances.Elements() {
+ 		if data.Key.ContainingAction().Equal(addr) {
+ 			result = append(result, data.Key)
+ 		}
+ 	}
+```
+
+```diff
+--- a/internal/actions/actions.go
++++ b/internal/actions/actions.go
+ func (a *Actions) GetActionInstanceKeys(addr addrs.AbsAction) []addrs.AbsActionInstance {
+ 	a.mu.Lock()
+ 	defer a.mu.Unlock()
+
+-	/
++	// Iterate through all action instances and return those that belong to the given action
+ 	result := []addrs.AbsActionInstance{}
+ 	for _, data := range a.actionInstances.Elements() {
+ 		if data.Key.ContainingAction().Equal(addr) {
+ 			result = append(result, data.Key)
+ 		}
+ 	}
+```
+
+```diff
+--- a/internal/actions/actions.go
++++ b/internal/actions/actions.go
+ func (a *Actions) GetActionInstanceKeys(addr addrs.AbsAction) []addrs.AbsActionInstance {
+ 	a.mu.Lock()
+ 	defer a.mu.Unlock()
+
+-	/
++	// Collect all action instances that belong to the given action
+ 	result := []addrs.AbsActionInstance{}
+ 	for _, data := range a.actionInstances.Elements() {
+ 		if data.Key.ContainingAction().Equal(addr) {
+ 			result = append(result, data.Key)
+ 		}
+ 	}
+```

crates/edit_prediction_cli/evals/tree-sitter--if-let-to-match.md

@@ -0,0 +1,136 @@
++++
+repository_url = "git@github.com:tree-sitter/tree-sitter"
+revision = "17e3c7a5c56527a179fa6e37ce7ee934493e5047"
++++
+
+## Edit History
+
+```diff
+--- a/crates/loader/src/loader.rs
++++ b/crates/loader/src/loader.rs
+@@ -729,15 +729,16 @@
+             ));
+         }
+         for parser_container_dir in &config.parser_directories {
+-            if let Ok(entries) = fs::read_dir(parser_container_dir) {
+-                for entry in entries {
+-                    let entry = entry.map_err(|e| LoaderError::IO(IoError::new(e, None)))?;
+-                    if let Some(parser_dir_name) = entry.file_name().to_str() {
+-                        if parser_dir_name.starts_with("tree-sitter-") {
+-                            self.find_language_configurations_at_path(
+-                                &parser_container_dir.join(parser_dir_name),
+-                                false,
+-                            )
++            match fs::read_dir(parser_container_dir) {
++                Ok(entries) => {
++                    for entry in entries {
++                        let entry = entry.map_err(|e| LoaderError::IO(IoError::new(e, None)))?;
++                        if let Some(parser_dir_name) = entry.file_name().to_str() {
++                            if parser_dir_name.starts_with("tree-sitter-") {
++                                self.find_language_configurations_at_path(
++                                    &parser_container_dir.join(parser_dir_name),
++                                    false,
++                                )
+                             .ok();
+                         }
+                     }
+--- a/crates/loader/src/loader.rs
++++ b/crates/loader/src/loader.rs
+@@ -739,7 +739,8 @@
+                                     &parser_container_dir.join(parser_dir_name),
+                                     false,
+                                 )
+-                            .ok();
++                                .ok();
++                            }
+                         }
+                     }
+                 }
+```
+
+## Cursor Position
+
+```crates/loader/src/loader.rs
+                        if let Some(parser_dir_name) = entry.file_name().to_str() {
+                            if parser_dir_name.starts_with("tree-sitter-") {
+                                self.find_language_configurations_at_path(
+                                    &parser_container_dir.join(parser_dir_name),
+                                    false,
+                                )
+                                .ok();
+                            }
+//                           ^[CURSOR_POSITION]
+                        }
+                    }
+                }
+            }
+        }
+```
+
+## Expected Patch
+
+```diff
+--- a/crates/loader/src/loader.rs
++++ b/crates/loader/src/loader.rs
+@@ -736,13 +736,13 @@
+                         if let Some(parser_dir_name) = entry.file_name().to_str() {
+                             if parser_dir_name.starts_with("tree-sitter-") {
+                                 self.find_language_configurations_at_path(
+                                     &parser_container_dir.join(parser_dir_name),
+                                     false,
+                                 )
+                                 .ok();
+                             }
+                         }
+                     }
+                 }
++                Err(error) => {}
+#                               ^[CURSOR_POSITION]
+             }
+         }
+```
+
+```diff
+--- a/crates/loader/src/loader.rs
++++ b/crates/loader/src/loader.rs
+@@ -736,13 +736,13 @@
+                         if let Some(parser_dir_name) = entry.file_name().to_str() {
+                             if parser_dir_name.starts_with("tree-sitter-") {
+                                 self.find_language_configurations_at_path(
+                                     &parser_container_dir.join(parser_dir_name),
+                                     false,
+                                 )
+                                 .ok();
+                             }
+                         }
+                     }
+                 }
++                Err(_) => {}
+#                           ^[CURSOR_POSITION]
+             }
+         }
+```
+
+
+```diff
+--- a/crates/loader/src/loader.rs
++++ b/crates/loader/src/loader.rs
+@@ -736,13 +736,13 @@
+                         if let Some(parser_dir_name) = entry.file_name().to_str() {
+                             if parser_dir_name.starts_with("tree-sitter-") {
+                                 self.find_language_configurations_at_path(
+                                     &parser_container_dir.join(parser_dir_name),
+                                     false,
+                                 )
+                                 .ok();
+                             }
+                         }
+                     }
+                 }
++                Err(e) => {
++                    
+#                    ^[CURSOR_POSITION]
++                }
+             }
+         }
+```

crates/edit_prediction_cli/evals/tree-sitter--tuple-to-struct-definition.md

@@ -0,0 +1,125 @@
++++
+repository_url = "git@github.com:tree-sitter/tree-sitter"
+revision = "24007727d42b4caceda3095ac685c463fae1ba1a"
++++
+
+## Edit History
+
+```diff
+--- a/tree-sitter/crates/loader/src/loader.rs
++++ b/tree-sitter/crates/loader/src/loader.rs
+@@ -604,7 +604,7 @@
+
+ pub struct Loader {
+     pub parser_lib_path: PathBuf,
+-    languages_by_id: Vec<(PathBuf, OnceCell<Language>, Option<Vec<PathBuf>>)>,
++    languages_by_id: Vec<LanguageEntry>,
+     language_configurations: Vec<LanguageConfiguration<'static>>,
+     language_configuration_ids_by_file_type: HashMap<String, Vec<usize>>,
+     language_configuration_in_current_path: Option<usize>,
+--- a/tree-sitter/crates/loader/src/loader.rs
++++ b/tree-sitter/crates/loader/src/loader.rs
+@@ -621,6 +621,8 @@
+     wasm_store: Mutex<Option<tree_sitter::WasmStore>>,
+ }
+
++str
+ pub struct CompileConfig<'a> {
+     pub src_path: &'a Path,
+     pub header_paths: Vec<&'a Path>,
+```
+
+## Cursor Position
+
+```tree-sitter/crates/loader/src/loader.rs
+    sanitize_build: bool,
+    force_rebuild: bool,
+
+    #[cfg(feature = "wasm")]
+    wasm_store: Mutex<Option<tree_sitter::WasmStore>>,
+}
+
+str
+// ^[CURSOR_POSITION]
+pub struct CompileConfig<'a> {
+    pub src_path: &'a Path,
+    pub header_paths: Vec<&'a Path>,
+    pub parser_path: PathBuf,
+    pub scanner_path: Option<PathBuf>,
+    pub external_files: Option<&'a [PathBuf]>,
+```
+
+## Expected Patch
+
+```diff
+--- a/tree-sitter/crates/loader/src/loader.rs
++++ b/tree-sitter/crates/loader/src/loader.rs
+@@ -621,6 +621,8 @@
+     wasm_store: Mutex<Option<tree_sitter::WasmStore>>,
+ }
+
+-str
++struct LanguageEntry {
++    path: PathBuf,
++    language: OnceCell<Language>,
++    external_files: Option<Vec<PathBuf>>,
++}
++
+ pub struct CompileConfig<'a> {
+     pub src_path: &'a Path,
+     pub header_paths: Vec<&'a Path>,
+```
+
+```diff
+--- a/tree-sitter/crates/loader/src/loader.rs
++++ b/tree-sitter/crates/loader/src/loader.rs
+@@ -621,6 +621,8 @@
+     wasm_store: Mutex<Option<tree_sitter::WasmStore>>,
+ }
+
+-str
++struct LanguageEntry {
++    path: PathBuf,
++    language: OnceCell<Language>,
++    dependencies: Option<Vec<PathBuf>>,
++}
++
+ pub struct CompileConfig<'a> {
+     pub src_path: &'a Path,
+     pub header_paths: Vec<&'a Path>,
+```
+
+
+```diff
+--- a/tree-sitter/crates/loader/src/loader.rs
++++ b/tree-sitter/crates/loader/src/loader.rs
+@@ -621,6 +621,8 @@
+     wasm_store: Mutex<Option<tree_sitter::WasmStore>>,
+ }
+
+-str
++struct LanguageEntry {
++    path: PathBuf,
++    language: OnceCell<Language>,
++    extra_files: Option<Vec<PathBuf>>,
++}
++
+ pub struct CompileConfig<'a> {
+     pub src_path: &'a Path,
+     pub header_paths: Vec<&'a Path>,
+```
+
+```diff
+--- a/tree-sitter/crates/loader/src/loader.rs
++++ b/tree-sitter/crates/loader/src/loader.rs
+@@ -621,6 +621,8 @@
+     wasm_store: Mutex<Option<tree_sitter::WasmStore>>,
+ }
+
+-str
++struct LanguageEntry(PathBuf, OnceCell<Language>, Option<Vec<PathBuf>>);
++
+ pub struct CompileConfig<'a> {
+     pub src_path: &'a Path,
+     pub header_paths: Vec<&'a Path>,
+```

crates/edit_prediction_cli/evals/tree-sitter--tuple-to-struct-destructuring.md

@@ -0,0 +1,153 @@
++++
+repository_url = "git@github.com:tree-sitter/tree-sitter"
+revision = "24007727d42b4caceda3095ac685c463fae1ba1a"
++++
+
+## Edit History
+
+```diff
+--- a/tree-sitter/crates/loader/src/loader.rs
++++ b/tree-sitter/crates/loader/src/loader.rs
+@@ -604,7 +604,7 @@
+
+ pub struct Loader {
+     pub parser_lib_path: PathBuf,
+-    languages_by_id: Vec<(PathBuf, OnceCell<Language>, Option<Vec<PathBuf>>)>,
++    languages_by_id: Vec<LanguageEntry>,
+     language_configurations: Vec<LanguageConfiguration<'static>>,
+     language_configuration_ids_by_file_type: HashMap<String, Vec<usize>>,
+     language_configuration_in_current_path: Option<usize>,
+@@ -619,6 +619,12 @@
+
+     #[cfg(feature = "wasm")]
+     wasm_store: Mutex<Option<tree_sitter::WasmStore>>,
+ }
++
++struct LanguageEntry {
++    path: PathBuf,
++    language: OnceCell<Language>,
++    external_files: Option<Vec<PathBuf>>,
++}
+
+ pub struct CompileConfig<'a> {
+@@ -767,7 +773,7 @@
+     pub fn get_all_language_configurations(&self) -> Vec<(&LanguageConfiguration, &Path)> {
+         self.language_configurations
+             .iter()
+-            .map(|c| (c, self.languages_by_id[c.language_id].0.as_ref()))
++            .map(|c| (c, self.languages_by_id[c.language_id].path.as_ref()))
+             .collect()
+     }
+
+```
+
+## Cursor Position
+
+```tree-sitter/crates/loader/src/loader.rs
+    fn language_for_id(&self, id: usize) -> LoaderResult<Language> {
+        let (path, language, externals) = &self.languages_by_id[id];
+        //  ^[CURSOR_POSITION]
+        language
+            .get_or_try_init(|| {
+                let src_path = path.join("src");
+                self.load_language_at_path(CompileConfig::new(
+                    &src_path,
+                    externals.as_deref(),
+                    None,
+                ))
+            })
+            .cloned()
+    }
+```
+
+## Expected Patch
+
+```diff
+--- a/tree-sitter/crates/loader/src/loader.rs
++++ b/tree-sitter/crates/loader/src/loader.rs
+@@ -926,7 +926,11 @@
+     }
+
+     fn language_for_id(&self, id: usize) -> LoaderResult<Language> {
+-        let (path, language, externals) = &self.languages_by_id[id];
++        let LanguageEntry {
++            path,
++            language,
++            external_files,
++        } = &self.languages_by_id[id];
+         language
+             .get_or_try_init(|| {
+                 let src_path = path.join("src");
+                 self.load_language_at_path(CompileConfig::new(
+                     &src_path,
+-                    externals.as_deref(),
++                    external_files.as_deref(),
+                     None,
+                 ))
+             })
+             .cloned()
+```
+
+```diff
+--- a/tree-sitter/crates/loader/src/loader.rs
++++ b/tree-sitter/crates/loader/src/loader.rs
+@@ -926,7 +926,11 @@
+     }
+
+     fn language_for_id(&self, id: usize) -> LoaderResult<Language> {
+-        let (path, language, externals) = &self.languages_by_id[id];
++        let LanguageEntry {
++            path,
++            language,
++            external_files: externals,
++        } = &self.languages_by_id[id];
+         language
+             .get_or_try_init(|| {
+                 let src_path = path.join("src");
+```
+
+```diff
+--- a/tree-sitter/crates/loader/src/loader.rs
++++ b/tree-sitter/crates/loader/src/loader.rs
+@@ -926,13 +926,14 @@
+     }
+
+     fn language_for_id(&self, id: usize) -> LoaderResult<Language> {
+-        let (path, language, externals) = &self.languages_by_id[id];
+-        language
++        let entry = &self.languages_by_id[id];
++        entry
++            .language
+             .get_or_try_init(|| {
+-                let src_path = path.join("src");
++                let src_path = entry.path.join("src");
+                 self.load_language_at_path(CompileConfig::new(
+                     &src_path,
+-                    externals.as_deref(),
++                    entry.external_files.as_deref(),
+                     None,
+                 ))
+             })
+```
+
+```diff
+--- a/tree-sitter/crates/loader/src/loader.rs
++++ b/tree-sitter/crates/loader/src/loader.rs
+@@ -926,0 +926,0 @@
+     }
+
+     fn language_for_id(&self, id: usize) -> LoaderResult<Language> {
+-        let (path, language, externals) = &self.languages_by_id[id];
++        let LanguageEntry { path, language, external_files  } = &self.languages_by_id[id];
+         language
+             .get_or_try_init(|| {
+                 let src_path = path.join("src");
+                 self.load_language_at_path(CompileConfig::new(
+                     &src_path,
+-                    externals.as_deref(),
++                    external_files.as_deref(),
+                     None,
+                 ))
+             })
+             .cloned()
+```

crates/edit_prediction_cli/evals/tree-sitter--tuple-to-struct-field-access.md

@@ -0,0 +1,62 @@
++++
+repository_url = "git@github.com:tree-sitter/tree-sitter"
+revision = "24007727d42b4caceda3095ac685c463fae1ba1a"
++++
+
+## Edit History
+
+```diff
+--- a/tree-sitter/crates/loader/src/loader.rs
++++ b/tree-sitter/crates/loader/src/loader.rs
+@@ -604,7 +604,7 @@
+
+ pub struct Loader {
+     pub parser_lib_path: PathBuf,
+-    languages_by_id: Vec<(PathBuf, OnceCell<Language>, Option<Vec<PathBuf>>)>,
++    languages_by_id: Vec<LanguageEntry>,
+     language_configurations: Vec<LanguageConfiguration<'static>>,
+     language_configuration_ids_by_file_type: HashMap<String, Vec<usize>>,
+     language_configuration_in_current_path: Option<usize>,
+@@ -619,6 +619,12 @@
+
+     #[cfg(feature = "wasm")]
+     wasm_store: Mutex<Option<tree_sitter::WasmStore>>,
+ }
++
++struct LanguageEntry {
++    path: PathBuf,
++    language: OnceCell<Language>,
++    external_files: Option<Vec<PathBuf>>,
++}
+
+ pub struct CompileConfig<'a> {
+```
+
+## Cursor Position
+
+```tree-sitter/crates/loader/src/loader.rs
+    #[must_use]
+    pub fn get_all_language_configurations(&self) -> Vec<(&LanguageConfiguration, &Path)> {
+        self.language_configurations
+            .iter()
+            .map(|c| (c, self.languages_by_id[c.language_id].0.as_ref()))
+            //                                               ^[CURSOR_POSITION]
+            .collect()
+    }
+```
+
+## Expected Patch
+
+```diff
+--- a/tree-sitter/crates/loader/src/loader.rs
++++ b/tree-sitter/crates/loader/src/loader.rs
+@@ -773,7 +773,7 @@
+     pub fn get_all_language_configurations(&self) -> Vec<(&LanguageConfiguration, &Path)> {
+         self.language_configurations
+             .iter()
+-            .map(|c| (c, self.languages_by_id[c.language_id].0.as_ref()))
++            .map(|c| (c, self.languages_by_id[c.language_id].path.as_ref()))
+             .collect()
+     }
+
+```

crates/edit_prediction_cli/evals/tree-sitter--tuple-to-struct-for-loop.md

@@ -0,0 +1,106 @@
++++
+repository_url = "git@github.com:tree-sitter/tree-sitter"
+revision = "24007727d42b4caceda3095ac685c463fae1ba1a"
++++
+
+## Edit History
+
+```diff
+--- a/tree-sitter/crates/loader/src/loader.rs
++++ b/tree-sitter/crates/loader/src/loader.rs
+@@ -604,7 +604,7 @@
+
+ pub struct Loader {
+     pub parser_lib_path: PathBuf,
+-    languages_by_id: Vec<(PathBuf, OnceCell<Language>, Option<Vec<PathBuf>>)>,
++    languages_by_id: Vec<LanguageEntry>,
+     language_configurations: Vec<LanguageConfiguration<'static>>,
+     language_configuration_ids_by_file_type: HashMap<String, Vec<usize>>,
+     language_configuration_in_current_path: Option<usize>,
+@@ -619,6 +619,12 @@
+
+     #[cfg(feature = "wasm")]
+     wasm_store: Mutex<Option<tree_sitter::WasmStore>>,
+ }
++
++struct LanguageEntry {
++    path: PathBuf,
++    language: OnceCell<Language>,
++    external_files: Option<Vec<PathBuf>>,
++}
+
+ pub struct CompileConfig<'a> {
+@@ -767,7 +773,7 @@
+     pub fn get_all_language_configurations(&self) -> Vec<(&LanguageConfiguration, &Path)> {
+         self.language_configurations
+             .iter()
+-            .map(|c| (c, self.languages_by_id[c.language_id].0.as_ref()))
++            .map(|c| (c, self.languages_by_id[c.language_id].path.as_ref()))
+             .collect()
+     }
+
+@@ -920,7 +926,11 @@
+     }
+
+     fn language_for_id(&self, id: usize) -> LoaderResult<Language> {
+-        let (path, language, externals) = &self.languages_by_id[id];
++        let LanguageEntry {
++            path,
++            language,
++            external_files,
++        } = &self.languages_by_id[id];
+         language
+             .get_or_try_init(|| {
+                 let src_path = path.join("src");
+```
+
+## Cursor Position
+
+```tree-sitter/crates/loader/src/loader.rs
+                    let mut language_id = None;
+                    for (id, (path, _, _)) in
+                    //       ^[CURSOR_POSITION]
+                        self.languages_by_id.iter().enumerate().skip(language_count)
+                    {
+                        if language_path == *path {
+                            language_id = Some(id);
+                        }
+                    }
+```
+
+## Expected Patch
+
+```diff
+--- a/tree-sitter/crates/loader/src/loader.rs
++++ b/tree-sitter/crates/loader/src/loader.rs
+@@ -1542,8 +1542,7 @@
+                     // Determine if a previous language configuration in this package.json file
+                     // already uses the same language.
+                     let mut language_id = None;
+-                    for (id, (path, _, _)) in
+-                        self.languages_by_id.iter().enumerate().skip(language_count)
++                    for (id, entry) in self.languages_by_id.iter().enumerate().skip(language_count)
+                     {
+-                        if language_path == *path {
++                        if language_path == entry.path {
+                             language_id = Some(id);
+                         }
+                     }
+```
+
+```diff
+--- a/tree-sitter/crates/loader/src/loader.rs
++++ b/tree-sitter/crates/loader/src/loader.rs
+@@ -1542,8 +1542,7 @@
+                     // Determine if a previous language configuration in this package.json file
+                     // already uses the same language.
+                     let mut language_id = None;
+-                    for (id, (path, _, _)) in
++                    for (id, LanguageEntry { path, .. }) in
+                         self.languages_by_id.iter().enumerate().skip(language_count)
+                     {
+                         if language_path == *path {
+                             language_id = Some(id);
+                         }
+                     }
+```

crates/edit_prediction_cli/evals/tree-sitter--tuple-to-struct-literal.md

@@ -0,0 +1,144 @@
++++
+repository_url = "git@github.com:tree-sitter/tree-sitter"
+revision = "24007727d42b4caceda3095ac685c463fae1ba1a"
++++
+
+## Edit History
+
+```diff
+--- a/tree-sitter/crates/loader/src/loader.rs
++++ b/tree-sitter/crates/loader/src/loader.rs
+@@ -604,7 +604,7 @@
+
+ pub struct Loader {
+     pub parser_lib_path: PathBuf,
+-    languages_by_id: Vec<(PathBuf, OnceCell<Language>, Option<Vec<PathBuf>>)>,
++    languages_by_id: Vec<LanguageEntry>,
+     language_configurations: Vec<LanguageConfiguration<'static>>,
+     language_configuration_ids_by_file_type: HashMap<String, Vec<usize>>,
+     language_configuration_in_current_path: Option<usize>,
+@@ -619,6 +619,12 @@
+
+     #[cfg(feature = "wasm")]
+     wasm_store: Mutex<Option<tree_sitter::WasmStore>>,
++}
++
++struct LanguageEntry {
++    path: PathBuf,
++    language: OnceCell<Language>,
++    external_files: Option<Vec<PathBuf>>,
+ }
+
+ pub struct CompileConfig<'a> {
+@@ -767,7 +773,7 @@
+     pub fn get_all_language_configurations(&self) -> Vec<(&LanguageConfiguration, &Path)> {
+         self.language_configurations
+             .iter()
+-            .map(|c| (c, self.languages_by_id[c.language_id].0.as_ref()))
++            .map(|c| (c, self.languages_by_id[c.language_id].path.as_ref()))
+             .collect()
+     }
+
+@@ -920,13 +926,17 @@
+     }
+
+     fn language_for_id(&self, id: usize) -> LoaderResult<Language> {
+-        let (path, language, externals) = &self.languages_by_id[id];
++        let LanguageEntry {
++            path,
++            language,
++            external_files,
++        } = &self.languages_by_id[id];
+         language
+             .get_or_try_init(|| {
+                 let src_path = path.join("src");
+                 self.load_language_at_path(CompileConfig::new(
+                     &src_path,
+-                    externals.as_deref(),
++                    external_files.as_deref(),
+                     None,
+                 ))
+             })
+@@ -1532,10 +1542,9 @@
+                     // Determine if a previous language configuration in this package.json file
+                     // already uses the same language.
+                     let mut language_id = None;
+-                    for (id, (path, _, _)) in
+-                        self.languages_by_id.iter().enumerate().skip(language_count)
++                    for (id, entry) in self.languages_by_id.iter().enumerate().skip(language_count)
+                     {
+-                        if language_path == *path {
++                        if language_path == entry.path {
+                             language_id = Some(id);
+                         }
+                     }
+--- a/tree-sitter/crates/loader/src/loader.rs
++++ b/tree-sitter/crates/loader/src/loader.rs
+@@ -1553,10 +1553,10 @@
+                     let language_id = if let Some(language_id) = language_id {
+                         language_id
+                     } else {
+-                        self.languages_by_id.push((
+-                            language_path,
+-                            OnceCell::new(),
+-                            grammar
++                        self.languages_by_id.push(LanguageEntry {
++                            path: language_path,
++                            language: OnceCell::new(),
++                            external_files: grammar
+                                 .external_files
+                                 .clone()
+                                 .into_vec()
+```
+
+## Cursor Position
+
+```tree-sitter/crates/loader/src/loader.rs
+                    let language_id = if let Some(language_id) = language_id {
+                        language_id
+                    } else {
+                        self.languages_by_id.push(LanguageEntry {
+                            path: language_path,
+                            language: OnceCell::new(),
+                            external_files: grammar
+                                .external_files
+                                .clone()
+                                .into_vec()
+                                .map(|files| {
+                                    files
+                                        .into_iter()
+                                        .map(|path| {
+                                            let path = parser_path.join(path);
+                                            // prevent p being above/outside of parser_path
+                                            if path.starts_with(parser_path) {
+                                                Ok(path)
+                                            } else {
+                                                Err(LoaderError::ExternalFile(
+                                                    path.to_string_lossy().to_string(),
+                                                    parser_path.to_string_lossy().to_string(),
+                                                ))
+                                            }
+                                        })
+                                        .collect::<LoaderResult<Vec<_>>>()
+                                })
+                                .transpose()?,
+                                //            ^[CURSOR_POSITION]
+                        ));
+                        self.languages_by_id.len() - 1
+                    };
+```
+
+## Expected Patch
+
+```diff
+--- a/tree-sitter/crates/loader/src/loader.rs
++++ b/tree-sitter/crates/loader/src/loader.rs
+@@ -1578,7 +1578,7 @@
+                                         .collect::<LoaderResult<Vec<_>>>()
+                                 })
+                                 .transpose()?,
+-                        ));
++                        });
+                         self.languages_by_id.len() - 1
+                     };
+```

crates/edit_prediction_cli/evals/vscode--add-async-and-await.md

@@ -0,0 +1,88 @@
++++
+repository_url = "https://github.com/microsoft/vscode"
+revision = "29e6da6efa2287aaa981635a475d425ff4fd5d5c"
++++
+
+## Edit History
+
+```diff
+--- a/src/vs/workbench/contrib/debug/browser/debugCommands.ts
++++ b/src/vs/workbench/contrib/debug/browser/debugCommands.ts
+@@ -304,8 +304,8 @@ CommandsRegistry.registerCommand({
+ 
+ CommandsRegistry.registerCommand({
+ 	id: REVERSE_CONTINUE_ID,
+-	handler: (accessor: ServicesAccessor, _: string, context: CallStackContext | unknown) => {
+-		getThreadAndRun(accessor, context, thread => thread.reverseContinue());
++	handler: async (accessor: ServicesAccessor, _: string, context: CallStackContext | unknown) => {
++		await getThreadAndRun(accessor, context, thread => thread.reverseContinue());
+ 	}
+ });
+--- a/src/vs/workbench/contrib/debug/browser/debugCommands.ts
++++ b/src/vs/workbench/contrib/debug/browser/debugCommands.ts
+@@ -311,11 +311,11 @@ CommandsRegistry.registerCommand({
+ 
+ CommandsRegistry.registerCommand({
+ 	id: STEP_BACK_ID,
+-	handler: (accessor: ServicesAccessor, _: string, context: CallStackContext | unknown) => {
++	handler: async (accessor: ServicesAccessor, _: string, context: CallStackContext | unknown) => {
+ 		const contextKeyService = accessor.get(IContextKeyService);
+ 		if (CONTEXT_DISASSEMBLY_VIEW_FOCUS.getValue(contextKeyService)) {
+-			getThreadAndRun(accessor, context, (thread: IThread) => thread.stepBack('instruction'));
++			await getThreadAndRun(accessor, context, (thread: IThread) => thread.stepBack('instruction'));
+ 		} else {
+-			getThreadAndRun(accessor, context, (thread: IThread) => thread.stepBack());
++			await getThreadAndRun(accessor, context, (thread: IThread) => thread.stepBack());
+ 		}
+ 	}
+ });
+--- a/src/vs/workbench/contrib/debug/browser/debugCommands.ts
++++ b/src/vs/workbench/contrib/debug/browser/debugCommands.ts
+@@ -323,8 +323,8 @@ CommandsRegistry.registerCommand({
+ 
+ CommandsRegistry.registerCommand({
+ 	id: TERMINATE_THREAD_ID,
+-	handler: (accessor: ServicesAccessor, _: string, context: CallStackContext | unknown) => {
+-		getThreadAndRun(accessor, context, thread => thread.terminate());
++	handler: async (accessor: ServicesAccessor, _: string, context: CallStackContext | unknown) => {
++		await getThreadAndRun(accessor, context, thread => thread.terminate());
+ 	}
+ });
+```
+
+## Cursor Position
+
+```src/vs/workbench/contrib/debug/browser/debugCommands.ts
+	weight: KeybindingWeight.WorkbenchContrib,
+	primary: isWeb ? (KeyMod.Alt | KeyCode.F10) : KeyCode.F10, // Browsers do not allow F10 to be binded so we have to bind an alternative
+	when: CONTEXT_DEBUG_STATE.isEqualTo('stopped'),
+	handler: (accessor: ServicesAccessor, _: string, context: CallStackContext | unknown) => {
+	//       ^[CURSOR_POSITION]
+		const contextKeyService = accessor.get(IContextKeyService);
+		if (CONTEXT_DISASSEMBLY_VIEW_FOCUS.getValue(contextKeyService)) {
+			getThreadAndRun(accessor, context, (thread: IThread) => thread.next('instruction'));
+		} else {
+```
+
+## Expected Patch
+
+```diff
+--- a/src/vs/workbench/contrib/debug/browser/debugCommands.ts
++++ b/src/vs/workbench/contrib/debug/browser/debugCommands.ts
+@@ -467,10 +467,10 @@ KeybindingsRegistry.registerCommandAndKeybindingRule({
+ 	weight: KeybindingWeight.WorkbenchContrib,
+ 	primary: isWeb ? (KeyMod.Alt | KeyCode.F10) : KeyCode.F10, // Browsers do not allow F10 to be binded so we have to bind an alternative
+ 	when: CONTEXT_DEBUG_STATE.isEqualTo('stopped'),
+-	handler: (accessor: ServicesAccessor, _: string, context: CallStackContext | unknown) => {
++	handler: async (accessor: ServicesAccessor, _: string, context: CallStackContext | unknown) => {
+ 		const contextKeyService = accessor.get(IContextKeyService);
+ 		if (CONTEXT_DISASSEMBLY_VIEW_FOCUS.getValue(contextKeyService)) {
+-			getThreadAndRun(accessor, context, (thread: IThread) => thread.next('instruction'));
++			await getThreadAndRun(accessor, context, (thread: IThread) => thread.next('instruction'));
+ 		} else {
+-			getThreadAndRun(accessor, context, (thread: IThread) => thread.next());
++			await getThreadAndRun(accessor, context, (thread: IThread) => thread.next());
+ 		}
+ 	}
+ });
+```

crates/edit_prediction_cli/evals/vscode--add-class-decorator.md

@@ -0,0 +1,74 @@
++++
+repository_url = "https://github.com/microsoft/vscode"
+revision = "6f6e26fcdf0a7ca5084e0da284cd7a5b2d41ae4d"
++++
+
+## Edit History
+
+```diff
+--- a/src/vs/workbench/api/common/extHostTypes.ts
++++ b/src/vs/workbench/api/common/extHostTypes.ts
+@@ -18,6 +18,14 @@ import { FileSystemProviderErrorCode, markAsFileSystemProviderError } from 'vs/
+ import type * as vscode from 'vscode';
+
++function es5ClassCompat(target: Function): any {
++	///@ts-expect-error
++	function _() { return Reflect.construct(target, arguments, this.constructor); }
++	Object.defineProperty(_, 'name', Object.getOwnPropertyDescriptor(target, 'name')!);
++	Object.setPrototypeOf(_, target);
++	Object.setPrototypeOf(_.prototype, target.prototype);
++	return _;
++}
++
++@es5ClassCompat
+ export class Disposable {
+--- a/src/vs/workbench/api/common/extHostTypes.ts
++++ b/src/vs/workbench/api/common/extHostTypes.ts
+@@ -50,6 +58,7 @@ export class Disposable {
+ 	}
+ }
+
++@es5ClassCompat
+ export class Position {
+
+ 	static Min(...positions: Position[]): Position {
+--- a/src/vs/workbench/api/common/extHostTypes.ts
++++ b/src/vs/workbench/api/common/extHostTypes.ts
+@@ -220,6 +229,7 @@ export class Position {
+ 	}
+ }
+
++@es5ClassCompat
+ export class Range {
+
+ 	static isRange(thing: any): thing is vscode.Range {
+```
+
+## Cursor Position
+
+```src/vs/workbench/api/common/extHostTypes.ts
+	Prepend = 3
+}
+
+export class TextEdit {
+// <[CURSOR_POSITION]
+
+	static isTextEdit(thing: any): thing is TextEdit {
+		if (thing instanceof TextEdit) {
+			return true;
+```
+
+## Expected Patch
+
+```diff
+--- a/src/vs/workbench/api/common/extHostTypes.ts
++++ b/src/vs/workbench/api/common/extHostTypes.ts
+@@ -475,6 +485,7 @@ export enum EnvironmentVariableMutatorType {
+ 	Prepend = 3
+ }
+
++@es5ClassCompat
+ export class TextEdit {
+
+ 	static isTextEdit(thing: any): thing is TextEdit {
+```

crates/edit_prediction_cli/evals/vscode--add-interface-method.md

@@ -0,0 +1,113 @@
++++
+repository_url = "https://github.com/microsoft/vscode"
+revision = "b64eaf598008e2d600a81d846108f72cb37b48e2"
++++
+
+## Edit History
+
+```diff
+--- a/src/vs/platform/window/electron-main/window.ts
++++ b/src/vs/platform/window/electron-main/window.ts
+@@ -1,49 +1,50 @@
+ export interface ICodeWindow extends IDisposable {
+ 
+ 	readonly onWillLoad: Event<ILoadEvent>;
+ 	readonly onDidSignalReady: Event<void>;
++	readonly onDidTriggerSystemContextMenu: Event<{ x: number; y: number }>;
+ 	readonly onDidClose: Event<void>;
+ 	readonly onDidDestroy: Event<void>;
+ 
+ 	readonly whenClosedOrLoaded: Promise<void>;
+--- a/src/vs/platform/windows/electron-main/window.ts
++++ b/src/vs/platform/windows/electron-main/window.ts
+@@ -63,60 +63,63 @@ const enum ReadyState {
+ export class CodeWindow extends Disposable implements ICodeWindow {
+ 
+ 	//#region Events
+ 
+ 	private readonly _onWillLoad = this._register(new Emitter<ILoadEvent>());
+ 	readonly onWillLoad = this._onWillLoad.event;
+ 
+ 	private readonly _onDidSignalReady = this._register(new Emitter<void>());
+ 	readonly onDidSignalReady = this._onDidSignalReady.event;
+ 
++	private readonly _onDidTriggerSystemContextMenu = this._register(new Emitter<{ x: number; y: number }>());
++	readonly onDidTriggerSystemContextMenu = this._onDidTriggerSystemContextMenu.event;
++
+ 	private readonly _onDidClose = this._register(new Emitter<void>());
+ 	readonly onDidClose = this._onDidClose.event;
+ 
+ 	private readonly _onDidDestroy = this._register(new Emitter<void>());
+ 	readonly onDidDestroy = this._onDidDestroy.event;
+ 
+ 	//#endregion
+--- a/src/vs/platform/windows/electron-main/windows.ts
++++ b/src/vs/platform/windows/electron-main/windows.ts
+@@ -1,54 +1,55 @@
+ export interface IWindowsMainService {
+ 
+ 	readonly _serviceBrand: undefined;
+ 
+ 	readonly onDidChangeWindowsCount: Event<IWindowsCountChangedEvent>;
+ 
+ 	readonly onDidOpenWindow: Event<ICodeWindow>;
+ 	readonly onDidSignalReadyWindow: Event<ICodeWindow>;
++	readonly onDidTriggerSystemContextMenu: Event<{ window: ICodeWindow; x: number; y: number }>;
+ 	readonly onDidDestroyWindow: Event<ICodeWindow>;
+--- a/src/vs/platform/windows/electron-main/windowsMainService.ts
++++ b/src/vs/platform/windows/electron-main/windowsMainService.ts
+@@ -160,60 +160,63 @@ interface ISingleFolderWorkspacePathToOpen extends IPathToOpen {
+ export class WindowsMainService extends Disposable implements IWindowsMainService {
+ 
+ 	declare readonly _serviceBrand: undefined;
+ 
+ 	private static readonly WINDOWS: ICodeWindow[] = [];
+ 
+ 	private readonly _onDidOpenWindow = this._register(new Emitter<ICodeWindow>());
+ 	readonly onDidOpenWindow = this._onDidOpenWindow.event;
+ 
+ 	private readonly _onDidSignalReadyWindow = this._register(new Emitter<ICodeWindow>());
+ 	readonly onDidSignalReadyWindow = this._onDidSignalReadyWindow.event;
+ 
+ 	private readonly _onDidDestroyWindow = this._register(new Emitter<ICodeWindow>());
+ 	readonly onDidDestroyWindow = this._onDidDestroyWindow.event;
+ 
+ 	private readonly _onDidChangeWindowsCount = this._register(new Emitter<IWindowsCountChangedEvent>());
+ 	readonly onDidChangeWindowsCount = this._onDidChangeWindowsCount.event;
+ 
++	private readonly _onDidTriggerSystemContextMenu = this._register(new Emitter<{ window: ICodeWindow; x: number; y: number }>());
++	readonly onDidTriggerSystemContextMenu = this._onDidTriggerSystemContextMenu.event;
++
+ 	private readonly windowsStateHandler = this._register(new WindowsStateHandler(this, this.stateMainService, this.lifecycleMainService, this.logService, this.configurationService));
+```
+
+## Cursor Position
+
+```src/vs/platform/windows/test/electron-main/windowsFinder.test.ts
+	function createTestCodeWindow(options: { lastFocusTime: number; openedFolderUri?: URI; openedWorkspace?: IWorkspaceIdentifier }): ICodeWindow {
+		return new class implements ICodeWindow {
+			onWillLoad: Event<ILoadEvent> = Event.None;
+			onDidSignalReady: Event<void> = Event.None;
+			// <[CURSOR_POSITION]
+			onDidClose: Event<void> = Event.None;
+			onDidDestroy: Event<void> = Event.None;
+			whenClosedOrLoaded: Promise<void> = Promise.resolve();
+			id: number = -1;
+```
+
+## Expected Patch
+
+```diff
+--- a/src/vs/platform/windows/test/electron-main/windowsFinder.test.ts
++++ b/src/vs/platform/windows/test/electron-main/windowsFinder.test.ts
+@@ -7,60 +7,61 @@ import * as assert from 'assert';
+ 	function createTestCodeWindow(options: { lastFocusTime: number; openedFolderUri?: URI; openedWorkspace?: IWorkspaceIdentifier }): ICodeWindow {
+ 		return new class implements ICodeWindow {
+ 			onWillLoad: Event<ILoadEvent> = Event.None;
++			onDidTriggerSystemContextMenu: Event<{ x: number; y: number }> = Event.None;
+ 			onDidSignalReady: Event<void> = Event.None;
+ 			onDidClose: Event<void> = Event.None;
+ 			onDidDestroy: Event<void> = Event.None;
+ 			whenClosedOrLoaded: Promise<void> = Promise.resolve();
+ 			id: number = -1;
+```

crates/edit_prediction_cli/evals/vscode--log-object-property.md

@@ -0,0 +1,56 @@
++++
+repository_url = "https://github.com/microsoft/vscode"
+revision = "e28a92fc1fbe9de11eca2f8ad19899334bff8525"
++++
+
+This prediction requires the model to see the `IDiffComputationResult` type definition.
+
+## Edit History
+
+```diff
+--- a/src/vs/editor/browser/widget/diffEditorWidget.ts
++++ b/src/vs/editor/browser/widget/diffEditorWidget.ts
+@@ -1117,6 +1117,7 @@
+ 				&& currentModifiedModel === this._modifiedEditor.getModel()
+ 			) {
+ 				this._setState(editorBrowser.DiffEditorState.DiffComputed);
++				console.log("did quit:")
+ 				this._diffComputationResult = result;
+ 				this._updateDecorationsRunner.schedule();
+ 				this._onDidUpdateDiff.fire();
+```
+
+## Cursor Position
+
+```src/vs/editor/browser/widget/diffEditorWidget.ts
+			if (currentToken === this._diffComputationToken
+				&& currentOriginalModel === this._originalEditor.getModel()
+				&& currentModifiedModel === this._modifiedEditor.getModel()
+			) {
+				this._setState(editorBrowser.DiffEditorState.DiffComputed);
+				console.log("did quit:")
+				//                    ^[CURSOR_POSITION]
+				this._diffComputationResult = result;
+				this._updateDecorationsRunner.schedule();
+				this._onDidUpdateDiff.fire();
+			}
+```
+
+## Expected Patch
+
+```diff
+--- a/src/vs/editor/browser/widget/diffEditorWidget.ts
++++ b/src/vs/editor/browser/widget/diffEditorWidget.ts
+@@ -1115,10 +1115,10 @@
+ 			if (currentToken === this._diffComputationToken
+ 				&& currentOriginalModel === this._originalEditor.getModel()
+ 				&& currentModifiedModel === this._modifiedEditor.getModel()
+ 			) {
+ 				this._setState(editorBrowser.DiffEditorState.DiffComputed);
+-				console.log("did quit:")
++				console.log("did quit:", result.quitEarly)
+ 				this._diffComputationResult = result;
+ 				this._updateDecorationsRunner.schedule();
+ 				this._onDidUpdateDiff.fire();
+ 			}
+```

crates/edit_prediction_cli/evals/zed--add-eprintln.md

@@ -0,0 +1,54 @@
++++
+repository_url = "git@github.com:zed-industries/zed"
+revision = "b7090c9fae7390a82021b994994c0f587744d96c"
++++
+
+This example shows the model's preference for making conservative predictions, and ability to place
+the cursor within the predicted output.
+
+## Edit History
+
+```diff
+--- a/crates/edit_prediction_ui/src/rate_prediction_modal.rs
++++ b/crates/edit_prediction_ui/src/rate_prediction_modal.rs
+@@ -144,7 +144,7 @@
+     fn select_next_edit(&mut self, _: &NextEdit, _: &mut Window, cx: &mut Context<Self>) {
++        epr
+         let next_index = self
+             .ep_store
+             .read(cx)
+```
+
+## Cursor Position
+
+```crates/edit_prediction_ui/src/rate_prediction_modal.rs
+    fn select_next_edit(&mut self, _: &NextEdit, _: &mut Window, cx: &mut Context<Self>) {
+        epr
+        // ^[CURSOR_POSITION]
+        let next_index = self
+            .ep_store
+            .read(cx)
+            .shown_predictions()
+            .skip(self.selected_index)
+            .enumerate()
+            .skip(1) // Skip straight to the next item
+```
+
+## Expected Patch
+
+```diff
+--- a/crates/edit_prediction_ui/src/rate_prediction_modal.rs
++++ b/crates/edit_prediction_ui/src/rate_prediction_modal.rs
+@@ -144,14 +144,14 @@
+     fn select_next_edit(&mut self, _: &NextEdit, _: &mut Window, cx: &mut Context<Self>) {
+-        epr
++        eprintln!("");
+#                   ^[CURSOR_POSITION]
+         let next_index = self
+             .ep_store
+             .read(cx)
+             .shown_predictions()
+             .skip(self.selected_index)
+             .enumerate()
+             .skip(1) // Skip straight to the next item
+```

crates/edit_prediction_cli/evals/zed--change-match-arm.md

@@ -0,0 +1,68 @@
++++
+repository_url = "git@github.com:zed-industries/zed"
+revision = "be5763632dccb33470ca233c36ccd9e5e790e3b2"
++++
+
+This prediction requires the model to see the `project::Event` enum.
+
+## Edit History
+
+```diff
+--- a/crates/edit_prediction/src/edit_prediction.rs
++++ b/crates/edit_prediction/src/edit_prediction.rs
+@@ -1035,7 +1035,7 @@
+                     project_state.recent_paths.push_front(path);
+                 }
+             }
+-            project::Event::DiagnosticsUpdated { .. } => {
++            project::Event::Disk { .. } => {
+                 if cx.has_flag::<EditPredictionJumpsFeatureFlag>() {
+                     self.refresh_prediction_from_diagnostics(
+                         project,
+```
+
+## Cursor Position
+
+```crates/edit_prediction/src/edit_prediction.rs
+                    {
+                        project_state.recent_paths.remove(ix);
+                    }
+                    project_state.recent_paths.push_front(path);
+                }
+            }
+            project::Event::Disk { .. } => {
+                //              ^[CURSOR_POSITION]
+                if cx.has_flag::<EditPredictionJumpsFeatureFlag>() {
+                    self.refresh_prediction_from_diagnostics(
+                        project,
+```
+
+## Expected Patch
+
+```diff
+--- a/crates/edit_prediction/src/edit_prediction.rs
++++ b/crates/edit_prediction/src/edit_prediction.rs
+@@ -1032,10 +1032,10 @@
+                     project_state.recent_paths.push_front(path);
+                 }
+             }
+-            project::Event::Disk { .. } => {
++            project::Event::DiskBasedDiagnosticsFinished { .. } => {
+                 if cx.has_flag::<EditPredictionJumpsFeatureFlag>() {
+                     self.refresh_prediction_from_diagnostics(
+                         project,
+```
+
+```diff
+--- a/crates/edit_prediction/src/edit_prediction.rs
++++ b/crates/edit_prediction/src/edit_prediction.rs
+@@ -1032,10 +1032,10 @@
+                     project_state.recent_paths.push_front(path);
+                 }
+             }
+-            project::Event::Disk { .. } => {
++            project::Event::DiskBasedDiagnosticsStarted { .. } => {
+                 if cx.has_flag::<EditPredictionJumpsFeatureFlag>() {
+                     self.refresh_prediction_from_diagnostics(
+                         project,
+```

crates/edit_prediction_cli/src/anthropic_client.rs

@@ -0,0 +1,866 @@
+use anthropic::{
+    ANTHROPIC_API_URL, Event, Message, Request as AnthropicRequest, RequestContent,
+    Response as AnthropicResponse, ResponseContent, Role, non_streaming_completion,
+    stream_completion,
+};
+use anyhow::Result;
+use futures::StreamExt as _;
+use http_client::HttpClient;
+use indoc::indoc;
+use reqwest_client::ReqwestClient;
+use sqlez::bindable::Bind;
+use sqlez::bindable::StaticColumnCount;
+use sqlez_macros::sql;
+use std::collections::HashSet;
+use std::hash::Hash;
+use std::hash::Hasher;
+use std::path::Path;
+use std::sync::{Arc, Mutex};
+
+pub struct PlainLlmClient {
+    pub http_client: Arc<dyn HttpClient>,
+    pub api_key: String,
+}
+
+impl PlainLlmClient {
+    pub fn new() -> Result<Self> {
+        let http_client: Arc<dyn http_client::HttpClient> = Arc::new(ReqwestClient::new());
+        let api_key = std::env::var("ANTHROPIC_API_KEY")
+            .map_err(|_| anyhow::anyhow!("ANTHROPIC_API_KEY environment variable not set"))?;
+        Ok(Self {
+            http_client,
+            api_key,
+        })
+    }
+
+    pub async fn generate(
+        &self,
+        model: &str,
+        max_tokens: u64,
+        messages: Vec<Message>,
+    ) -> Result<AnthropicResponse> {
+        let request = AnthropicRequest {
+            model: model.to_string(),
+            max_tokens,
+            messages,
+            tools: Vec::new(),
+            thinking: None,
+            tool_choice: None,
+            system: None,
+            cache_control: None,
+            metadata: None,
+            output_config: None,
+            stop_sequences: Vec::new(),
+            speed: None,
+            temperature: None,
+            top_k: None,
+            top_p: None,
+        };
+
+        let response = non_streaming_completion(
+            self.http_client.as_ref(),
+            ANTHROPIC_API_URL,
+            &self.api_key,
+            request,
+            None,
+        )
+        .await
+        .map_err(|e| anyhow::anyhow!("{:?}", e))?;
+
+        Ok(response)
+    }
+
+    pub async fn generate_streaming<F>(
+        &self,
+        model: &str,
+        max_tokens: u64,
+        messages: Vec<Message>,
+        mut on_progress: F,
+    ) -> Result<AnthropicResponse>
+    where
+        F: FnMut(usize, &str),
+    {
+        let request = AnthropicRequest {
+            model: model.to_string(),
+            max_tokens,
+            messages,
+            tools: Vec::new(),
+            thinking: None,
+            tool_choice: None,
+            system: None,
+            cache_control: None,
+            metadata: None,
+            output_config: None,
+            stop_sequences: Vec::new(),
+            speed: None,
+            temperature: None,
+            top_k: None,
+            top_p: None,
+        };
+
+        let mut stream = stream_completion(
+            self.http_client.as_ref(),
+            ANTHROPIC_API_URL,
+            &self.api_key,
+            request,
+            None,
+        )
+        .await
+        .map_err(|e| anyhow::anyhow!("{:?}", e))?;
+
+        let mut response: Option<AnthropicResponse> = None;
+        let mut text_content = String::new();
+
+        while let Some(event_result) = stream.next().await {
+            let event = event_result.map_err(|e| anyhow::anyhow!("{:?}", e))?;
+
+            match event {
+                Event::MessageStart { message } => {
+                    response = Some(message);
+                }
+                Event::ContentBlockDelta { delta, .. } => {
+                    if let anthropic::ContentDelta::TextDelta { text } = delta {
+                        text_content.push_str(&text);
+                        on_progress(text_content.len(), &text_content);
+                    }
+                }
+                _ => {}
+            }
+        }
+
+        let mut response = response.ok_or_else(|| anyhow::anyhow!("No response received"))?;
+
+        if response.content.is_empty() && !text_content.is_empty() {
+            response
+                .content
+                .push(ResponseContent::Text { text: text_content });
+        }
+
+        Ok(response)
+    }
+}
+
+pub struct BatchingLlmClient {
+    connection: Mutex<sqlez::connection::Connection>,
+    http_client: Arc<dyn HttpClient>,
+    api_key: String,
+}
+
+struct CacheRow {
+    request_hash: String,
+    request: Option<String>,
+    response: Option<String>,
+    batch_id: Option<String>,
+}
+
+impl StaticColumnCount for CacheRow {
+    fn column_count() -> usize {
+        4
+    }
+}
+
+impl Bind for CacheRow {
+    fn bind(&self, statement: &sqlez::statement::Statement, start_index: i32) -> Result<i32> {
+        let next_index = statement.bind(&self.request_hash, start_index)?;
+        let next_index = statement.bind(&self.request, next_index)?;
+        let next_index = statement.bind(&self.response, next_index)?;
+        let next_index = statement.bind(&self.batch_id, next_index)?;
+        Ok(next_index)
+    }
+}
+
+#[derive(serde::Serialize, serde::Deserialize)]
+struct SerializableRequest {
+    model: String,
+    max_tokens: u64,
+    messages: Vec<SerializableMessage>,
+}
+
+#[derive(serde::Serialize, serde::Deserialize)]
+struct SerializableMessage {
+    role: String,
+    content: String,
+}
+
+impl BatchingLlmClient {
+    fn new(cache_path: &Path) -> Result<Self> {
+        let http_client: Arc<dyn http_client::HttpClient> = Arc::new(ReqwestClient::new());
+        let api_key = std::env::var("ANTHROPIC_API_KEY")
+            .map_err(|_| anyhow::anyhow!("ANTHROPIC_API_KEY environment variable not set"))?;
+
+        let connection = sqlez::connection::Connection::open_file(&cache_path.to_str().unwrap());
+        let mut statement = sqlez::statement::Statement::prepare(
+            &connection,
+            indoc! {"
+                CREATE TABLE IF NOT EXISTS cache (
+                    request_hash TEXT PRIMARY KEY,
+                    request TEXT,
+                    response TEXT,
+                    batch_id TEXT
+                );
+                "},
+        )?;
+        statement.exec()?;
+        drop(statement);
+
+        Ok(Self {
+            connection: Mutex::new(connection),
+            http_client,
+            api_key,
+        })
+    }
+
+    pub fn lookup(
+        &self,
+        model: &str,
+        max_tokens: u64,
+        messages: &[Message],
+        seed: Option<usize>,
+    ) -> Result<Option<AnthropicResponse>> {
+        let request_hash_str = Self::request_hash(model, max_tokens, messages, seed);
+        let connection = self.connection.lock().unwrap();
+        let response: Vec<String> = connection.select_bound(
+            &sql!(SELECT response FROM cache WHERE request_hash = ?1 AND response IS NOT NULL;),
+        )?(request_hash_str.as_str())?;
+        Ok(response
+            .into_iter()
+            .next()
+            .and_then(|text| serde_json::from_str(&text).ok()))
+    }
+
+    pub fn mark_for_batch(
+        &self,
+        model: &str,
+        max_tokens: u64,
+        messages: &[Message],
+        seed: Option<usize>,
+    ) -> Result<()> {
+        let request_hash = Self::request_hash(model, max_tokens, messages, seed);
+
+        let serializable_messages: Vec<SerializableMessage> = messages
+            .iter()
+            .map(|msg| SerializableMessage {
+                role: match msg.role {
+                    Role::User => "user".to_string(),
+                    Role::Assistant => "assistant".to_string(),
+                },
+                content: message_content_to_string(&msg.content),
+            })
+            .collect();
+
+        let serializable_request = SerializableRequest {
+            model: model.to_string(),
+            max_tokens,
+            messages: serializable_messages,
+        };
+
+        let request = Some(serde_json::to_string(&serializable_request)?);
+        let cache_row = CacheRow {
+            request_hash,
+            request,
+            response: None,
+            batch_id: None,
+        };
+        let connection = self.connection.lock().unwrap();
+        connection.exec_bound::<CacheRow>(sql!(
+            INSERT OR IGNORE INTO cache(request_hash, request, response, batch_id) VALUES (?, ?, ?, ?)))?(
+            cache_row,
+        )
+    }
+
+    async fn generate(
+        &self,
+        model: &str,
+        max_tokens: u64,
+        messages: Vec<Message>,
+        seed: Option<usize>,
+        cache_only: bool,
+    ) -> Result<Option<AnthropicResponse>> {
+        let response = self.lookup(model, max_tokens, &messages, seed)?;
+        if let Some(response) = response {
+            return Ok(Some(response));
+        }
+
+        if !cache_only {
+            self.mark_for_batch(model, max_tokens, &messages, seed)?;
+        }
+
+        Ok(None)
+    }
+
+    /// Uploads pending requests as batches (chunked to 16k each); downloads finished batches if any.
+    async fn sync_batches(&self) -> Result<()> {
+        let _batch_ids = self.upload_pending_requests().await?;
+        self.download_finished_batches().await
+    }
+
+    pub fn pending_batch_count(&self) -> Result<usize> {
+        let connection = self.connection.lock().unwrap();
+        let counts: Vec<i32> = connection.select(
+            sql!(SELECT COUNT(*) FROM cache WHERE batch_id IS NOT NULL AND response IS NULL),
+        )?()?;
+        Ok(counts.into_iter().next().unwrap_or(0) as usize)
+    }
+
+    /// Import batch results from external batch IDs (useful for recovering after database loss)
+    pub async fn import_batches(&self, batch_ids: &[String]) -> Result<()> {
+        for batch_id in batch_ids {
+            log::info!("Importing batch {}", batch_id);
+
+            let batch_status = anthropic::batches::retrieve_batch(
+                self.http_client.as_ref(),
+                ANTHROPIC_API_URL,
+                &self.api_key,
+                batch_id,
+            )
+            .await
+            .map_err(|e| anyhow::anyhow!("Failed to retrieve batch {}: {:?}", batch_id, e))?;
+
+            log::info!(
+                "Batch {} status: {}",
+                batch_id,
+                batch_status.processing_status
+            );
+
+            if batch_status.processing_status != "ended" {
+                log::warn!(
+                    "Batch {} is not finished (status: {}), skipping",
+                    batch_id,
+                    batch_status.processing_status
+                );
+                continue;
+            }
+
+            let results = anthropic::batches::retrieve_batch_results(
+                self.http_client.as_ref(),
+                ANTHROPIC_API_URL,
+                &self.api_key,
+                batch_id,
+            )
+            .await
+            .map_err(|e| {
+                anyhow::anyhow!("Failed to retrieve batch results for {}: {:?}", batch_id, e)
+            })?;
+
+            let mut updates: Vec<(String, String, String)> = Vec::new();
+            let mut success_count = 0;
+            let mut error_count = 0;
+
+            for result in results {
+                let request_hash = result
+                    .custom_id
+                    .strip_prefix("req_hash_")
+                    .unwrap_or(&result.custom_id)
+                    .to_string();
+
+                match result.result {
+                    anthropic::batches::BatchResult::Succeeded { message } => {
+                        let response_json = serde_json::to_string(&message)?;
+                        updates.push((request_hash, response_json, batch_id.clone()));
+                        success_count += 1;
+                    }
+                    anthropic::batches::BatchResult::Errored { error } => {
+                        log::error!(
+                            "Batch request {} failed: {}: {}",
+                            request_hash,
+                            error.error.error_type,
+                            error.error.message
+                        );
+                        let error_json = serde_json::json!({
+                            "error": {
+                                "type": error.error.error_type,
+                                "message": error.error.message
+                            }
+                        })
+                        .to_string();
+                        updates.push((request_hash, error_json, batch_id.clone()));
+                        error_count += 1;
+                    }
+                    anthropic::batches::BatchResult::Canceled => {
+                        log::warn!("Batch request {} was canceled", request_hash);
+                        error_count += 1;
+                    }
+                    anthropic::batches::BatchResult::Expired => {
+                        log::warn!("Batch request {} expired", request_hash);
+                        error_count += 1;
+                    }
+                }
+            }
+
+            let connection = self.connection.lock().unwrap();
+            connection.with_savepoint("batch_import", || {
+                // Use INSERT OR REPLACE to handle both new entries and updating existing ones
+                let q = sql!(
+                    INSERT OR REPLACE INTO cache(request_hash, request, response, batch_id)
+                    VALUES (?, (SELECT request FROM cache WHERE request_hash = ?), ?, ?)
+                );
+                let mut exec = connection.exec_bound::<(&str, &str, &str, &str)>(q)?;
+                for (request_hash, response_json, batch_id) in &updates {
+                    exec((
+                        request_hash.as_str(),
+                        request_hash.as_str(),
+                        response_json.as_str(),
+                        batch_id.as_str(),
+                    ))?;
+                }
+                Ok(())
+            })?;
+
+            log::info!(
+                "Imported batch {}: {} successful, {} errors",
+                batch_id,
+                success_count,
+                error_count
+            );
+        }
+
+        Ok(())
+    }
+
+    async fn download_finished_batches(&self) -> Result<()> {
+        let batch_ids: Vec<String> = {
+            let connection = self.connection.lock().unwrap();
+            let q = sql!(SELECT DISTINCT batch_id FROM cache WHERE batch_id IS NOT NULL AND response IS NULL);
+            connection.select(q)?()?
+        };
+
+        for batch_id in &batch_ids {
+            let batch_status = anthropic::batches::retrieve_batch(
+                self.http_client.as_ref(),
+                ANTHROPIC_API_URL,
+                &self.api_key,
+                &batch_id,
+            )
+            .await
+            .map_err(|e| anyhow::anyhow!("{:?}", e))?;
+
+            log::info!(
+                "Batch {} status: {}",
+                batch_id,
+                batch_status.processing_status
+            );
+
+            if batch_status.processing_status == "ended" {
+                let results = anthropic::batches::retrieve_batch_results(
+                    self.http_client.as_ref(),
+                    ANTHROPIC_API_URL,
+                    &self.api_key,
+                    &batch_id,
+                )
+                .await
+                .map_err(|e| anyhow::anyhow!("{:?}", e))?;
+
+                let mut updates: Vec<(String, String)> = Vec::new();
+                let mut success_count = 0;
+                for result in results {
+                    let request_hash = result
+                        .custom_id
+                        .strip_prefix("req_hash_")
+                        .unwrap_or(&result.custom_id)
+                        .to_string();
+
+                    match result.result {
+                        anthropic::batches::BatchResult::Succeeded { message } => {
+                            let response_json = serde_json::to_string(&message)?;
+                            updates.push((response_json, request_hash));
+                            success_count += 1;
+                        }
+                        anthropic::batches::BatchResult::Errored { error } => {
+                            log::error!(
+                                "Batch request {} failed: {}: {}",
+                                request_hash,
+                                error.error.error_type,
+                                error.error.message
+                            );
+                            let error_json = serde_json::json!({
+                                "error": {
+                                    "type": error.error.error_type,
+                                    "message": error.error.message
+                                }
+                            })
+                            .to_string();
+                            updates.push((error_json, request_hash));
+                        }
+                        anthropic::batches::BatchResult::Canceled => {
+                            log::warn!("Batch request {} was canceled", request_hash);
+                            let error_json = serde_json::json!({
+                                "error": {
+                                    "type": "canceled",
+                                    "message": "Batch request was canceled"
+                                }
+                            })
+                            .to_string();
+                            updates.push((error_json, request_hash));
+                        }
+                        anthropic::batches::BatchResult::Expired => {
+                            log::warn!("Batch request {} expired", request_hash);
+                            let error_json = serde_json::json!({
+                                "error": {
+                                    "type": "expired",
+                                    "message": "Batch request expired"
+                                }
+                            })
+                            .to_string();
+                            updates.push((error_json, request_hash));
+                        }
+                    }
+                }
+
+                let connection = self.connection.lock().unwrap();
+                connection.with_savepoint("batch_download", || {
+                    let q = sql!(UPDATE cache SET response = ? WHERE request_hash = ?);
+                    let mut exec = connection.exec_bound::<(&str, &str)>(q)?;
+                    for (response_json, request_hash) in &updates {
+                        exec((response_json.as_str(), request_hash.as_str()))?;
+                    }
+                    Ok(())
+                })?;
+                log::info!("Downloaded {} successful requests", success_count);
+            }
+        }
+
+        Ok(())
+    }
+
+    async fn upload_pending_requests(&self) -> Result<Vec<String>> {
+        const BATCH_CHUNK_SIZE: i32 = 16_000;
+        const MAX_BATCH_SIZE_BYTES: usize = 100 * 1024 * 1024;
+
+        let mut all_batch_ids = Vec::new();
+        let mut total_uploaded = 0;
+
+        let mut current_batch_rows = Vec::new();
+        let mut current_batch_size = 0usize;
+        let mut pending_hashes: HashSet<String> = HashSet::new();
+        loop {
+            let rows: Vec<(String, String)> = {
+                let connection = self.connection.lock().unwrap();
+                let q = sql!(
+                    SELECT request_hash, request FROM cache
+                    WHERE batch_id IS NULL AND response IS NULL
+                    LIMIT ?
+                );
+                connection.select_bound(q)?(BATCH_CHUNK_SIZE)?
+            };
+
+            if rows.is_empty() {
+                break;
+            }
+
+            // Split rows into sub-batches based on size
+            let mut batches_to_upload = Vec::new();
+            let mut new_rows_added = 0;
+
+            for row in rows {
+                let (hash, request_str) = row;
+
+                // Skip rows already added to current_batch_rows but not yet uploaded
+                if pending_hashes.contains(&hash) {
+                    continue;
+                }
+                let serializable_request: SerializableRequest = serde_json::from_str(&request_str)?;
+
+                let messages: Vec<Message> = serializable_request
+                    .messages
+                    .into_iter()
+                    .map(|msg| Message {
+                        role: match msg.role.as_str() {
+                            "user" => Role::User,
+                            "assistant" => Role::Assistant,
+                            _ => Role::User,
+                        },
+                        content: vec![RequestContent::Text {
+                            text: msg.content,
+                            cache_control: None,
+                        }],
+                    })
+                    .collect();
+
+                let params = AnthropicRequest {
+                    model: serializable_request.model,
+                    max_tokens: serializable_request.max_tokens,
+                    messages,
+                    tools: Vec::new(),
+                    thinking: None,
+                    tool_choice: None,
+                    system: None,
+                    cache_control: None,
+                    metadata: None,
+                    output_config: None,
+                    stop_sequences: Vec::new(),
+                    temperature: None,
+                    top_k: None,
+                    top_p: None,
+                    speed: None,
+                };
+
+                let custom_id = format!("req_hash_{}", hash);
+                let batch_request = anthropic::batches::BatchRequest { custom_id, params };
+
+                // Estimate the serialized size of this request
+                let estimated_size = serde_json::to_string(&batch_request)?.len();
+
+                // If adding this request would exceed the limit, start a new batch
+                if !current_batch_rows.is_empty()
+                    && current_batch_size + estimated_size > MAX_BATCH_SIZE_BYTES
+                {
+                    batches_to_upload.push((current_batch_rows, current_batch_size));
+                    current_batch_rows = Vec::new();
+                    current_batch_size = 0;
+                }
+
+                pending_hashes.insert(hash.clone());
+                current_batch_rows.push((hash, batch_request));
+                current_batch_size += estimated_size;
+                new_rows_added += 1;
+            }
+
+            // If no new rows were added this iteration, all pending requests are already
+            // in current_batch_rows, so we should break to avoid an infinite loop
+            if new_rows_added == 0 {
+                break;
+            }
+
+            // Only upload full batches, keep the partial batch for the next iteration
+            // Upload each sub-batch
+            for (batch_rows, batch_size) in batches_to_upload {
+                let request_hashes: Vec<String> =
+                    batch_rows.iter().map(|(hash, _)| hash.clone()).collect();
+
+                // Remove uploaded hashes from pending set
+                for hash in &request_hashes {
+                    pending_hashes.remove(hash);
+                }
+                let batch_requests: Vec<anthropic::batches::BatchRequest> =
+                    batch_rows.into_iter().map(|(_, req)| req).collect();
+
+                let batch_len = batch_requests.len();
+                log::info!(
+                    "Uploading batch with {} requests (~{:.2} MB)",
+                    batch_len,
+                    batch_size as f64 / (1024.0 * 1024.0)
+                );
+
+                let batch = anthropic::batches::create_batch(
+                    self.http_client.as_ref(),
+                    ANTHROPIC_API_URL,
+                    &self.api_key,
+                    anthropic::batches::CreateBatchRequest {
+                        requests: batch_requests,
+                    },
+                )
+                .await
+                .map_err(|e| anyhow::anyhow!("{:?}", e))?;
+
+                {
+                    let connection = self.connection.lock().unwrap();
+                    connection.with_savepoint("batch_upload", || {
+                        let q = sql!(UPDATE cache SET batch_id = ? WHERE request_hash = ?);
+                        let mut exec = connection.exec_bound::<(&str, &str)>(q)?;
+                        for hash in &request_hashes {
+                            exec((batch.id.as_str(), hash.as_str()))?;
+                        }
+                        Ok(())
+                    })?;
+                }
+
+                total_uploaded += batch_len;
+                log::info!(
+                    "Uploaded batch {} with {} requests ({} total)",
+                    batch.id,
+                    batch_len,
+                    total_uploaded
+                );
+
+                all_batch_ids.push(batch.id);
+            }
+        }
+
+        // Upload any remaining partial batch at the end
+        if !current_batch_rows.is_empty() {
+            let request_hashes: Vec<String> = current_batch_rows
+                .iter()
+                .map(|(hash, _)| hash.clone())
+                .collect();
+            let batch_requests: Vec<anthropic::batches::BatchRequest> =
+                current_batch_rows.into_iter().map(|(_, req)| req).collect();
+
+            let batch_len = batch_requests.len();
+            log::info!(
+                "Uploading final batch with {} requests (~{:.2} MB)",
+                batch_len,
+                current_batch_size as f64 / (1024.0 * 1024.0)
+            );
+
+            let batch = anthropic::batches::create_batch(
+                self.http_client.as_ref(),
+                ANTHROPIC_API_URL,
+                &self.api_key,
+                anthropic::batches::CreateBatchRequest {
+                    requests: batch_requests,
+                },
+            )
+            .await
+            .map_err(|e| anyhow::anyhow!("{:?}", e))?;
+
+            {
+                let connection = self.connection.lock().unwrap();
+                connection.with_savepoint("batch_upload", || {
+                    let q = sql!(UPDATE cache SET batch_id = ? WHERE request_hash = ?);
+                    let mut exec = connection.exec_bound::<(&str, &str)>(q)?;
+                    for hash in &request_hashes {
+                        exec((batch.id.as_str(), hash.as_str()))?;
+                    }
+                    Ok(())
+                })?;
+            }
+
+            total_uploaded += batch_len;
+            log::info!(
+                "Uploaded batch {} with {} requests ({} total)",
+                batch.id,
+                batch_len,
+                total_uploaded
+            );
+
+            all_batch_ids.push(batch.id);
+        }
+
+        if !all_batch_ids.is_empty() {
+            log::info!(
+                "Finished uploading {} batches with {} total requests",
+                all_batch_ids.len(),
+                total_uploaded
+            );
+        }
+
+        Ok(all_batch_ids)
+    }
+
+    fn request_hash(
+        model: &str,
+        max_tokens: u64,
+        messages: &[Message],
+        seed: Option<usize>,
+    ) -> String {
+        let mut hasher = std::hash::DefaultHasher::new();
+        model.hash(&mut hasher);
+        max_tokens.hash(&mut hasher);
+        for msg in messages {
+            message_content_to_string(&msg.content).hash(&mut hasher);
+        }
+        if let Some(seed) = seed {
+            seed.hash(&mut hasher);
+        }
+        let request_hash = hasher.finish();
+        format!("{request_hash:016x}")
+    }
+}
+
+fn message_content_to_string(content: &[RequestContent]) -> String {
+    content
+        .iter()
+        .filter_map(|c| match c {
+            RequestContent::Text { text, .. } => Some(text.clone()),
+            _ => None,
+        })
+        .collect::<Vec<String>>()
+        .join("\n")
+}
+
+pub enum AnthropicClient {
+    // No batching
+    Plain(PlainLlmClient),
+    Batch(BatchingLlmClient),
+    Dummy,
+}
+
+impl AnthropicClient {
+    pub fn plain() -> Result<Self> {
+        Ok(Self::Plain(PlainLlmClient::new()?))
+    }
+
+    pub fn batch(cache_path: &Path) -> Result<Self> {
+        Ok(Self::Batch(BatchingLlmClient::new(cache_path)?))
+    }
+
+    #[allow(dead_code)]
+    pub fn dummy() -> Self {
+        Self::Dummy
+    }
+
+    pub async fn generate(
+        &self,
+        model: &str,
+        max_tokens: u64,
+        messages: Vec<Message>,
+        seed: Option<usize>,
+        cache_only: bool,
+    ) -> Result<Option<AnthropicResponse>> {
+        match self {
+            AnthropicClient::Plain(plain_llm_client) => plain_llm_client
+                .generate(model, max_tokens, messages)
+                .await
+                .map(Some),
+            AnthropicClient::Batch(batching_llm_client) => {
+                batching_llm_client
+                    .generate(model, max_tokens, messages, seed, cache_only)
+                    .await
+            }
+            AnthropicClient::Dummy => panic!("Dummy LLM client is not expected to be used"),
+        }
+    }
+
+    #[allow(dead_code)]
+    pub async fn generate_streaming<F>(
+        &self,
+        model: &str,
+        max_tokens: u64,
+        messages: Vec<Message>,
+        on_progress: F,
+    ) -> Result<Option<AnthropicResponse>>
+    where
+        F: FnMut(usize, &str),
+    {
+        match self {
+            AnthropicClient::Plain(plain_llm_client) => plain_llm_client
+                .generate_streaming(model, max_tokens, messages, on_progress)
+                .await
+                .map(Some),
+            AnthropicClient::Batch(_) => {
+                anyhow::bail!("Streaming not supported with batching client")
+            }
+            AnthropicClient::Dummy => panic!("Dummy LLM client is not expected to be used"),
+        }
+    }
+
+    pub async fn sync_batches(&self) -> Result<()> {
+        match self {
+            AnthropicClient::Plain(_) => Ok(()),
+            AnthropicClient::Batch(batching_llm_client) => batching_llm_client.sync_batches().await,
+            AnthropicClient::Dummy => panic!("Dummy LLM client is not expected to be used"),
+        }
+    }
+
+    pub fn pending_batch_count(&self) -> Result<usize> {
+        match self {
+            AnthropicClient::Plain(_) => Ok(0),
+            AnthropicClient::Batch(batching_llm_client) => {
+                batching_llm_client.pending_batch_count()
+            }
+            AnthropicClient::Dummy => panic!("Dummy LLM client is not expected to be used"),
+        }
+    }
+
+    pub async fn import_batches(&self, batch_ids: &[String]) -> Result<()> {
+        match self {
+            AnthropicClient::Plain(_) => {
+                anyhow::bail!("Import batches is only supported with batching client")
+            }
+            AnthropicClient::Batch(batching_llm_client) => {
+                batching_llm_client.import_batches(batch_ids).await
+            }
+            AnthropicClient::Dummy => panic!("Dummy LLM client is not expected to be used"),
+        }
+    }
+}

crates/edit_prediction_cli/src/distill.rs

@@ -0,0 +1,36 @@
+use anyhow::Result;
+
+use crate::{PredictionProvider, example::Example};
+
+pub async fn run_distill(example: &mut Example) -> Result<()> {
+    let has_repair = example
+        .predictions
+        .iter()
+        .find(|p| p.provider == PredictionProvider::Repair);
+    let predictions = if let Some(has_repair) = has_repair {
+        vec![has_repair]
+    } else {
+        example.predictions.iter().collect()
+    };
+
+    let expected_patches = predictions
+        .into_iter()
+        .filter_map(|p| {
+            Some((
+                p.actual_patch.clone()?,
+                p.actual_cursor
+                    .as_ref()
+                    .and_then(|c| c.editable_region_offset),
+            ))
+        })
+        .collect();
+
+    example
+        .spec
+        .set_expected_patches_with_cursor_positions(expected_patches);
+    example.prompt = None;
+    example.predictions = Vec::new();
+    example.score = Vec::new();
+    example.qa = Vec::new();
+    Ok(())
+}

crates/edit_prediction_cli/src/example.rs

@@ -0,0 +1,317 @@
+use crate::PredictionProvider;
+use crate::paths::WORKTREES_DIR;
+use crate::qa::QaResult;
+use anyhow::{Context as _, Result};
+use collections::HashMap;
+use edit_prediction::example_spec::ExampleSpec;
+use edit_prediction::udiff::OpenedBuffers;
+use gpui::Entity;
+use http_client::Url;
+use language::{Anchor, Buffer};
+use project::Project;
+use serde::{Deserialize, Serialize};
+use std::{
+    borrow::Cow,
+    collections::VecDeque,
+    io::Read,
+    path::{Path, PathBuf},
+};
+use zeta_prompt::ZetaPromptInput;
+
+#[derive(Clone, Debug, Serialize, Deserialize)]
+pub struct Example {
+    #[serde(flatten)]
+    pub spec: ExampleSpec,
+
+    /// The full content of the file where an edit is being predicted, and the
+    /// actual cursor offset.
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub prompt_inputs: Option<ZetaPromptInput>,
+
+    /// The input and expected output from the edit prediction model.
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub prompt: Option<ExamplePrompt>,
+
+    /// The actual predictions from the model.
+    #[serde(default, skip_serializing_if = "Vec::is_empty")]
+    pub predictions: Vec<ExamplePrediction>,
+
+    /// The scores, for how well the actual predictions match the expected
+    /// predictions.
+    #[serde(default, skip_serializing_if = "Vec::is_empty")]
+    pub score: Vec<ExampleScore>,
+
+    /// QA evaluation results for each prediction (indexed parallel to `predictions`).
+    #[serde(default, skip_serializing_if = "Vec::is_empty")]
+    pub qa: Vec<Option<QaResult>>,
+
+    /// The Zed version used to generate this example.
+    pub zed_version: Option<String>,
+
+    /// The application state used to process this example.
+    #[serde(skip)]
+    pub state: Option<ExampleState>,
+}
+
+#[derive(Clone, Debug)]
+pub struct ExampleState {
+    pub project: Entity<Project>,
+    pub buffer: Entity<Buffer>,
+    pub cursor_position: Anchor,
+    pub _open_buffers: OpenedBuffers,
+}
+
+#[derive(Clone, Debug, Serialize, Deserialize)]
+pub struct ExamplePrompt {
+    pub input: String,
+    #[serde(default)]
+    pub expected_output: Option<String>,
+    pub rejected_output: Option<String>, // For DPO
+    #[serde(default)]
+    pub prefill: Option<String>,
+    pub provider: PredictionProvider,
+}
+
+#[derive(Clone, Debug, Serialize, Deserialize)]
+pub struct ExamplePrediction {
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub actual_patch: Option<String>,
+    #[serde(deserialize_with = "deserialize_null_as_empty_string")]
+    pub actual_output: String,
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub actual_cursor: Option<ActualCursor>,
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub error: Option<String>,
+    pub provider: PredictionProvider,
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub cumulative_logprob: Option<f64>,
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub avg_logprob: Option<f64>,
+}
+
+#[derive(Clone, Debug, Serialize, Deserialize)]
+pub struct ActualCursor {
+    pub path: String,
+    pub row: u32,
+    pub column: u32,
+    pub offset: usize,
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub editable_region_offset: Option<usize>,
+}
+
+impl ActualCursor {
+    /// Construct an `ActualCursor` from a cursor offset within the new editable region.
+    ///
+    /// - `path`: file path the cursor is in
+    /// - `editable_region_cursor_offset`: byte offset of the cursor within the new editable region text
+    /// - `new_editable_region`: the full new editable region text (after marker removal)
+    /// - `content`: the full file content (before the edit)
+    /// - `editable_region_byte_offset`: byte offset where the editable region starts in `content`
+    /// - `editable_region_start_line`: 0-based line number where the editable region starts in `content`
+    pub fn from_editable_region(
+        path: &std::path::Path,
+        editable_region_cursor_offset: usize,
+        new_editable_region: &str,
+        content: &str,
+        editable_region_byte_offset: usize,
+        editable_region_start_line: usize,
+    ) -> Self {
+        let global_offset = editable_region_byte_offset + editable_region_cursor_offset;
+        let new_region_prefix = &new_editable_region[..editable_region_cursor_offset];
+        let row = (editable_region_start_line + new_region_prefix.matches('\n').count()) as u32;
+        let column = match new_region_prefix.rfind('\n') {
+            Some(pos) => (editable_region_cursor_offset - pos - 1) as u32,
+            None => {
+                let content_prefix = &content[..editable_region_byte_offset];
+                let content_column = match content_prefix.rfind('\n') {
+                    Some(pos) => editable_region_byte_offset - pos - 1,
+                    None => editable_region_byte_offset,
+                };
+                (content_column + editable_region_cursor_offset) as u32
+            }
+        };
+        ActualCursor {
+            path: path.to_string_lossy().to_string(),
+            row,
+            column,
+            offset: global_offset,
+            editable_region_offset: Some(editable_region_cursor_offset),
+        }
+    }
+}
+
+fn deserialize_null_as_empty_string<'de, D>(deserializer: D) -> Result<String, D::Error>
+where
+    D: serde::Deserializer<'de>,
+{
+    let opt = Option::<String>::deserialize(deserializer)?;
+    Ok(opt.unwrap_or_default())
+}
+
+pub type ExampleScore = edit_prediction_metrics::PredictionScore;
+
+impl Example {
+    pub fn repo_name(&self) -> Result<RepoName<'_>> {
+        // git@github.com:owner/repo.git
+        if self.spec.repository_url.contains('@') {
+            let (owner, repo) = self
+                .spec
+                .repository_url
+                .split_once(':')
+                .context("expected : in git url")?
+                .1
+                .split_once('/')
+                .context("expected / in git url")?;
+            Ok(RepoName {
+                owner: Cow::Borrowed(owner),
+                name: Cow::Borrowed(repo.trim_end_matches(".git")),
+            })
+        // http://github.com/owner/repo.git
+        } else {
+            let url = Url::parse(&self.spec.repository_url)?;
+            let mut segments = url.path_segments().context("empty http url")?;
+            let owner = segments
+                .next()
+                .context("expected owner path segment")?
+                .to_string();
+            let repo = segments
+                .next()
+                .context("expected repo path segment")?
+                .trim_end_matches(".git")
+                .to_string();
+            assert!(segments.next().is_none());
+
+            Ok(RepoName {
+                owner: Cow::Owned(owner),
+                name: Cow::Owned(repo),
+            })
+        }
+    }
+}
+
+pub struct RepoName<'a> {
+    pub owner: Cow<'a, str>,
+    pub name: Cow<'a, str>,
+}
+
+impl RepoName<'_> {
+    pub fn worktree_path(&self) -> PathBuf {
+        WORKTREES_DIR
+            .join(self.owner.as_ref())
+            .join(self.name.as_ref())
+    }
+}
+
+pub fn read_example_files(inputs: &[PathBuf]) -> Vec<Example> {
+    let mut examples = Vec::new();
+
+    for path in inputs {
+        let is_stdin = path.as_path() == Path::new("-");
+        let content = if is_stdin {
+            let mut buffer = String::new();
+            std::io::stdin()
+                .read_to_string(&mut buffer)
+                .expect("Failed to read from stdin");
+            buffer
+        } else {
+            std::fs::read_to_string(path)
+                .unwrap_or_else(|_| panic!("Failed to read path: {:?}", &path))
+        };
+        let filename = path.file_stem().unwrap().to_string_lossy().to_string();
+        let ext = if !is_stdin {
+            path.extension()
+                .map(|ext| ext.to_string_lossy().to_string())
+                .unwrap_or_else(|| panic!("{} should have an extension", path.display()))
+        } else {
+            "jsonl".to_string()
+        };
+
+        match ext.as_ref() {
+            "json" => {
+                let mut example =
+                    serde_json::from_str::<Example>(&content).unwrap_or_else(|error| {
+                        panic!("Failed to parse example file: {}\n{error}", path.display())
+                    });
+                if example.spec.name.is_empty() {
+                    example.spec.name = filename;
+                }
+                examples.push(example);
+            }
+            "jsonl" => examples.extend(
+                content
+                    .lines()
+                    .enumerate()
+                    .map(|(line_ix, line)| {
+                        let mut example =
+                            serde_json::from_str::<Example>(line).unwrap_or_else(|error| {
+                                panic!(
+                                    "Failed to parse example on {}:{}\n{error}",
+                                    path.display(),
+                                    line_ix + 1
+                                )
+                            });
+                        if example.spec.name.is_empty() {
+                            example.spec.name = format!("{filename}-{line_ix}")
+                        }
+                        example
+                    })
+                    .collect::<Vec<Example>>(),
+            ),
+            "md" => {
+                let mut example = parse_markdown_example(&content).unwrap();
+                if example.spec.name.is_empty() {
+                    example.spec.name = filename;
+                }
+                examples.push(example);
+            }
+            ext => {
+                panic!("{} has invalid example extension `{ext}`", path.display())
+            }
+        }
+    }
+
+    examples
+}
+
+pub fn sort_examples_by_repo_and_rev(examples: &mut [Example]) {
+    examples.sort_by(|a, b| {
+        a.spec
+            .repository_url
+            .cmp(&b.spec.repository_url)
+            .then(b.spec.revision.cmp(&a.spec.revision))
+    });
+}
+
+pub fn group_examples_by_repo(examples: Vec<Example>) -> VecDeque<Vec<Example>> {
+    let mut examples_by_repo: HashMap<String, Vec<Example>> = HashMap::default();
+    let mut ungrouped = Vec::new();
+    for example in examples {
+        if example.spec.repository_url.is_empty() {
+            ungrouped.push(example);
+        } else {
+            examples_by_repo
+                .entry(example.spec.repository_url.clone())
+                .or_insert_with(Vec::new)
+                .push(example);
+        }
+    }
+    let mut result: VecDeque<Vec<Example>> = examples_by_repo.into_values().collect();
+    for example in ungrouped {
+        result.push_back(vec![example]);
+    }
+    result
+}
+
+fn parse_markdown_example(input: &str) -> Result<Example> {
+    let spec = ExampleSpec::from_markdown(input)?;
+    Ok(Example {
+        spec,
+        prompt_inputs: None,
+        prompt: None,
+        predictions: Vec::new(),
+        score: Vec::new(),
+        qa: Vec::new(),
+        state: None,
+        zed_version: None,
+    })
+}

crates/edit_prediction_cli/src/filter_languages.rs

@@ -0,0 +1,551 @@
+//! `ep filter-languages` implementation.
+//!
+//! This command filters a JSONL dataset to include only examples where the
+//! cursor is in a file of a specified programming language.
+//!
+//! # Usage
+//!
+//! ```text
+//! ep filter-languages [input.jsonl] --languages rust,python,go
+//! ```
+//!
+//! # Language Detection
+//!
+//! Language is detected based on file extension of the `cursor_path` field.
+//! The extension-to-language mapping is built from the embedded language
+//! config files in the `grammars` crate.
+
+use anyhow::{Context as _, Result, bail};
+use clap::Args;
+use collections::HashMap;
+use serde::Deserialize;
+use std::ffi::OsStr;
+use std::fs::File;
+use std::io::{self, BufRead, BufReader, BufWriter, Write};
+use std::path::{Path, PathBuf};
+
+#[cfg(not(feature = "dynamic_prompts"))]
+mod language_configs_embedded {
+    use rust_embed::RustEmbed;
+
+    #[derive(RustEmbed)]
+    #[folder = "../grammars/src/"]
+    #[include = "*/config.toml"]
+    pub struct LanguageConfigs;
+}
+
+#[cfg(not(feature = "dynamic_prompts"))]
+use language_configs_embedded::LanguageConfigs;
+
+#[derive(Debug, Deserialize)]
+struct LanguageConfig {
+    name: String,
+    #[serde(default)]
+    path_suffixes: Vec<String>,
+}
+
+/// `ep filter-languages` CLI args.
+#[derive(Debug, Args, Clone)]
+#[command(
+    about = "Filter a JSONL dataset by programming language (based on cursor_path extension)",
+    after_help = r#"EXAMPLES:
+  # Filter to only Rust, Python, and Go examples
+  ep filter-languages input.jsonl --languages rust,python,go -o filtered.jsonl
+
+  # Filter by language and also include specific extensions
+  ep filter-languages input.jsonl --languages rust,python --extensions txt,md -o filtered.jsonl
+
+  # Filter by extensions only (no language filter)
+  ep filter-languages input.jsonl --extensions cs,java,swift -o filtered.jsonl
+
+  # List available languages
+  ep filter-languages --list
+
+  # Show statistics about languages in the input
+  ep filter-languages input.jsonl --stats
+
+NOTES:
+  Language names are case-insensitive.
+  Extensions should be specified without the leading dot (e.g., 'txt' not '.txt').
+  Use --list to see all available language names.
+"#
+)]
+pub struct FilterLanguagesArgs {
+    /// Comma-separated list of languages to include
+    #[arg(long, short = 'l', value_delimiter = ',')]
+    pub languages: Option<Vec<String>>,
+
+    /// Comma-separated list of file extensions to include (without leading dot)
+    #[arg(long, short = 'e', value_delimiter = ',')]
+    pub extensions: Option<Vec<String>>,
+
+    /// List all available languages and their extensions
+    #[arg(long)]
+    pub list: bool,
+
+    /// Show statistics about language distribution in the input
+    #[arg(long)]
+    pub stats: bool,
+
+    /// Include examples where language could not be detected
+    #[arg(long)]
+    pub include_unknown: bool,
+
+    /// Show top N excluded file extensions after filtering
+    #[arg(long, value_name = "N")]
+    pub show_top_excluded: Option<usize>,
+}
+
+#[cfg(not(feature = "dynamic_prompts"))]
+fn build_extension_to_language_map() -> HashMap<String, String> {
+    let mut map = HashMap::default();
+
+    for file_path in LanguageConfigs::iter() {
+        if let Some(content) = LanguageConfigs::get(&file_path) {
+            let content_str = match std::str::from_utf8(&content.data) {
+                Ok(s) => s,
+                Err(_) => continue,
+            };
+
+            let config: LanguageConfig = match toml::from_str(content_str) {
+                Ok(c) => c,
+                Err(_) => continue,
+            };
+
+            for suffix in &config.path_suffixes {
+                map.insert(suffix.to_lowercase(), config.name.clone());
+            }
+        }
+    }
+
+    map
+}
+
+#[cfg(feature = "dynamic_prompts")]
+fn build_extension_to_language_map() -> HashMap<String, String> {
+    const LANGUAGES_SRC_DIR: &str = concat!(env!("CARGO_MANIFEST_DIR"), "/../grammars/src");
+
+    let mut map = HashMap::default();
+
+    let languages_dir = Path::new(LANGUAGES_SRC_DIR);
+    let entries = match std::fs::read_dir(languages_dir) {
+        Ok(e) => e,
+        Err(_) => return map,
+    };
+
+    for entry in entries.flatten() {
+        let config_path = entry.path().join("config.toml");
+        if !config_path.exists() {
+            continue;
+        }
+
+        let content_str = match std::fs::read_to_string(&config_path) {
+            Ok(s) => s,
+            Err(_) => continue,
+        };
+
+        let config: LanguageConfig = match toml::from_str(&content_str) {
+            Ok(c) => c,
+            Err(_) => continue,
+        };
+
+        for suffix in &config.path_suffixes {
+            map.insert(suffix.to_lowercase(), config.name.clone());
+        }
+    }
+
+    map
+}
+
+fn get_all_languages(extension_map: &HashMap<String, String>) -> Vec<(String, Vec<String>)> {
+    let mut language_to_extensions: HashMap<String, Vec<String>> = HashMap::default();
+
+    for (ext, lang) in extension_map {
+        language_to_extensions
+            .entry(lang.clone())
+            .or_default()
+            .push(ext.clone());
+    }
+
+    let mut result: Vec<_> = language_to_extensions.into_iter().collect();
+    result.sort_by_key(|res| res.0.to_lowercase());
+    for (_, extensions) in &mut result {
+        extensions.sort();
+    }
+    result
+}
+
+fn detect_language(cursor_path: &str, extension_map: &HashMap<String, String>) -> Option<String> {
+    let path = Path::new(cursor_path);
+
+    if let Some(ext) = path.extension().and_then(OsStr::to_str) {
+        if let Some(lang) = extension_map.get(&ext.to_lowercase()) {
+            return Some(lang.clone());
+        }
+    }
+
+    if let Some(file_name) = path.file_name().and_then(OsStr::to_str) {
+        if let Some(lang) = extension_map.get(&file_name.to_lowercase()) {
+            return Some(lang.clone());
+        }
+    }
+
+    None
+}
+
+fn get_extension(cursor_path: &str) -> Option<String> {
+    let path = Path::new(cursor_path);
+
+    if let Some(ext) = path.extension().and_then(OsStr::to_str) {
+        return Some(ext.to_lowercase());
+    }
+
+    if let Some(file_name) = path.file_name().and_then(OsStr::to_str) {
+        return Some(file_name.to_lowercase());
+    }
+
+    None
+}
+
+fn read_lines_streaming(
+    input: Option<&Path>,
+) -> Result<Box<dyn Iterator<Item = io::Result<String>>>> {
+    let reader: Box<dyn BufRead> = match input {
+        Some(path) => {
+            let file =
+                File::open(path).with_context(|| format!("failed to open '{}'", path.display()))?;
+            Box::new(BufReader::new(file))
+        }
+        None => Box::new(BufReader::new(io::stdin())),
+    };
+    Ok(Box::new(reader.lines()))
+}
+
+fn get_cursor_path(line: &str) -> Option<String> {
+    let value: serde_json::Value = serde_json::from_str(line).ok()?;
+    value
+        .get("cursor_path")
+        .and_then(|v| v.as_str())
+        .map(|s| s.to_string())
+}
+
+pub fn run_filter_languages(
+    args: &FilterLanguagesArgs,
+    inputs: &[PathBuf],
+    output: Option<&PathBuf>,
+) -> Result<()> {
+    let extension_map = build_extension_to_language_map();
+
+    if args.list {
+        let languages = get_all_languages(&extension_map);
+        println!("Available languages ({}):", languages.len());
+        println!();
+        for (lang, extensions) in languages {
+            println!("  {}: {}", lang, extensions.join(", "));
+        }
+        return Ok(());
+    }
+
+    let input_path: Option<&Path> = match inputs.first().map(|p| p.as_path()) {
+        Some(p) if p.as_os_str() == "-" => None,
+        Some(p) => Some(p),
+        None => None,
+    };
+
+    if args.stats {
+        let stats_input =
+            input_path.with_context(|| "input file is required for --stats (cannot use stdin)")?;
+        return run_stats(stats_input, &extension_map);
+    }
+
+    if args.languages.is_none() && args.extensions.is_none() {
+        bail!(
+            "--languages and/or --extensions is required (use --list to see available languages, or --stats to see input distribution)"
+        );
+    }
+
+    let allowed_languages: std::collections::HashSet<String> = args
+        .languages
+        .as_ref()
+        .map(|langs| langs.iter().map(|l| l.to_lowercase()).collect())
+        .unwrap_or_default();
+
+    let allowed_extensions: std::collections::HashSet<String> = args
+        .extensions
+        .as_ref()
+        .map(|exts| {
+            exts.iter()
+                .map(|e| e.trim_start_matches('.').to_lowercase())
+                .collect()
+        })
+        .unwrap_or_default();
+
+    let language_name_lower_map: HashMap<String, String> = get_all_languages(&extension_map)
+        .into_iter()
+        .map(|(lang, _)| (lang.to_lowercase(), lang))
+        .collect();
+
+    if !allowed_languages.is_empty() {
+        for lang in &allowed_languages {
+            if !language_name_lower_map.contains_key(lang) {
+                eprintln!(
+                    "Warning: '{}' is not a recognized language name. Use --list to see available languages.",
+                    lang
+                );
+            }
+        }
+    }
+
+    let lines = read_lines_streaming(input_path)?;
+
+    let mut writer: Box<dyn Write> = match output {
+        Some(path) => {
+            if let Some(parent) = path.parent() {
+                if !parent.as_os_str().is_empty() {
+                    std::fs::create_dir_all(parent).with_context(|| {
+                        format!("failed to create directory '{}'", parent.display())
+                    })?;
+                }
+            }
+            let file = File::create(path)
+                .with_context(|| format!("failed to create '{}'", path.display()))?;
+            Box::new(BufWriter::new(file))
+        }
+        None => Box::new(BufWriter::new(io::stdout())),
+    };
+
+    let mut total_count = 0usize;
+    let mut included_count = 0usize;
+    let mut unknown_count = 0usize;
+    let mut excluded_extensions: HashMap<String, usize> = HashMap::default();
+
+    for line_result in lines {
+        let line = line_result.context("failed to read line")?;
+        if line.trim().is_empty() {
+            continue;
+        }
+
+        total_count += 1;
+
+        let cursor_path = match get_cursor_path(&line) {
+            Some(p) => p,
+            None => {
+                if args.include_unknown {
+                    unknown_count += 1;
+                    included_count += 1;
+                    writeln!(writer, "{}", line)?;
+                }
+                continue;
+            }
+        };
+
+        let language = detect_language(&cursor_path, &extension_map);
+        let extension = get_extension(&cursor_path);
+
+        let matches_language = match &language {
+            Some(lang) => allowed_languages.contains(&lang.to_lowercase()),
+            None => false,
+        };
+
+        let matches_extension = match &extension {
+            Some(ext) => allowed_extensions.contains(ext),
+            None => false,
+        };
+
+        let should_include = if matches_language || matches_extension {
+            true
+        } else if language.is_none() && args.include_unknown {
+            unknown_count += 1;
+            true
+        } else {
+            if let Some(ext) = &extension {
+                *excluded_extensions.entry(ext.clone()).or_default() += 1;
+            }
+            false
+        };
+
+        if should_include {
+            included_count += 1;
+            writeln!(writer, "{}", line)?;
+        }
+    }
+
+    writer.flush()?;
+
+    eprintln!(
+        "Filtered {} examples to {} ({} unknown language)",
+        total_count, included_count, unknown_count
+    );
+
+    if let Some(top_n) = args.show_top_excluded {
+        if !excluded_extensions.is_empty() {
+            let mut sorted: Vec<_> = excluded_extensions.into_iter().collect();
+            sorted.sort_by_key(|res| std::cmp::Reverse(res.1));
+            eprintln!("\nTop {} excluded extensions:", top_n.min(sorted.len()));
+            for (ext, count) in sorted.into_iter().take(top_n) {
+                eprintln!("  {:>6}  .{}", count, ext);
+            }
+        }
+    }
+
+    Ok(())
+}
+
+fn run_stats(input: &Path, extension_map: &HashMap<String, String>) -> Result<()> {
+    let lines = read_lines_streaming(Some(input))?;
+
+    let mut language_counts: HashMap<String, usize> = HashMap::default();
+    let mut unknown_extensions: HashMap<String, usize> = HashMap::default();
+    let mut total_count = 0usize;
+
+    for line_result in lines {
+        let line = line_result.context("failed to read line")?;
+        if line.trim().is_empty() {
+            continue;
+        }
+
+        total_count += 1;
+
+        let cursor_path = match get_cursor_path(&line) {
+            Some(p) => p,
+            None => {
+                *language_counts
+                    .entry("<no cursor_path>".to_string())
+                    .or_default() += 1;
+                continue;
+            }
+        };
+
+        match detect_language(&cursor_path, extension_map) {
+            Some(lang) => {
+                *language_counts.entry(lang).or_default() += 1;
+            }
+            None => {
+                let ext = Path::new(&cursor_path)
+                    .extension()
+                    .and_then(OsStr::to_str)
+                    .map(|s| s.to_string())
+                    .unwrap_or_else(|| {
+                        Path::new(&cursor_path)
+                            .file_name()
+                            .and_then(OsStr::to_str)
+                            .map(|s| s.to_string())
+                            .unwrap_or_else(|| "<no extension>".to_string())
+                    });
+                *unknown_extensions.entry(ext).or_default() += 1;
+                *language_counts.entry("<unknown>".to_string()).or_default() += 1;
+            }
+        }
+    }
+
+    let mut sorted_counts: Vec<_> = language_counts.into_iter().collect();
+    sorted_counts.sort_by_key(|res| std::cmp::Reverse(res.1));
+
+    println!("Language distribution ({} total examples):", total_count);
+    println!();
+    for (lang, count) in &sorted_counts {
+        let pct = (*count as f64 / total_count as f64) * 100.0;
+        println!("  {:>6} ({:>5.1}%)  {}", count, pct, lang);
+    }
+
+    if !unknown_extensions.is_empty() {
+        println!();
+        println!("Unknown extensions:");
+        let mut sorted_unknown: Vec<_> = unknown_extensions.into_iter().collect();
+        sorted_unknown.sort_by_key(|res| std::cmp::Reverse(res.1));
+        for (ext, count) in sorted_unknown.iter().take(30) {
+            println!("  {:>6}  .{}", count, ext);
+        }
+        if sorted_unknown.len() > 30 {
+            println!("  ... and {} more", sorted_unknown.len() - 30);
+        }
+    }
+
+    Ok(())
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_build_extension_map() {
+        let map = build_extension_to_language_map();
+        assert!(!map.is_empty());
+        assert_eq!(map.get("rs"), Some(&"Rust".to_string()));
+        assert_eq!(map.get("py"), Some(&"Python".to_string()));
+        assert_eq!(map.get("go"), Some(&"Go".to_string()));
+    }
+
+    #[test]
+    fn test_detect_language_by_extension() {
+        let map = build_extension_to_language_map();
+
+        assert_eq!(
+            detect_language("src/main.rs", &map),
+            Some("Rust".to_string())
+        );
+        assert_eq!(
+            detect_language("lib/foo.py", &map),
+            Some("Python".to_string())
+        );
+        assert_eq!(
+            detect_language("cmd/server.go", &map),
+            Some("Go".to_string())
+        );
+        assert_eq!(detect_language("index.tsx", &map), Some("TSX".to_string()));
+    }
+
+    #[test]
+    fn test_detect_language_by_filename() {
+        let map = build_extension_to_language_map();
+
+        // PKGBUILD is a filename-based match for Shell Script
+        assert_eq!(
+            detect_language("PKGBUILD", &map),
+            Some("Shell Script".to_string())
+        );
+        assert_eq!(
+            detect_language("project/PKGBUILD", &map),
+            Some("Shell Script".to_string())
+        );
+        // .env files are also Shell Script
+        assert_eq!(
+            detect_language(".env", &map),
+            Some("Shell Script".to_string())
+        );
+    }
+
+    #[test]
+    fn test_detect_language_unknown() {
+        let map = build_extension_to_language_map();
+
+        assert_eq!(detect_language("file.xyz123", &map), None);
+        assert_eq!(detect_language("random_file", &map), None);
+    }
+
+    #[test]
+    fn test_get_cursor_path() {
+        let line = r#"{"cursor_path": "src/main.rs", "other": "data"}"#;
+        assert_eq!(get_cursor_path(line), Some("src/main.rs".to_string()));
+
+        let line_no_cursor = r#"{"other": "data"}"#;
+        assert_eq!(get_cursor_path(line_no_cursor), None);
+
+        let invalid_json = "not json";
+        assert_eq!(get_cursor_path(invalid_json), None);
+    }
+
+    #[test]
+    fn test_get_all_languages() {
+        let map = build_extension_to_language_map();
+        let languages = get_all_languages(&map);
+
+        assert!(!languages.is_empty());
+
+        let rust_entry = languages.iter().find(|(name, _)| name == "Rust");
+        assert!(rust_entry.is_some());
+        let (_, rust_extensions) = rust_entry.unwrap();
+        assert!(rust_extensions.contains(&"rs".to_string()));
+    }
+}

crates/edit_prediction_cli/src/format_prompt.rs

@@ -0,0 +1,962 @@
+use crate::{
+    FormatPromptArgs, PredictionProvider,
+    example::{ActualCursor, Example, ExamplePrompt},
+    headless::EpAppState,
+    progress::{ExampleProgress, Step},
+    retrieve_context::run_context_retrieval,
+};
+use anyhow::{Context as _, Result, anyhow};
+use gpui::AsyncApp;
+use std::ops::Range;
+use std::sync::Arc;
+use zeta_prompt::{
+    ZetaFormat, format_expected_output, format_zeta_prompt, multi_region, resolve_cursor_region,
+};
+
+fn resolved_excerpt_ranges_for_format(
+    input: &zeta_prompt::ZetaPromptInput,
+    format: ZetaFormat,
+) -> (Range<usize>, Range<usize>) {
+    let (_, editable_range_in_context, context_range, _) = resolve_cursor_region(input, format);
+    let editable_range = (context_range.start + editable_range_in_context.start)
+        ..(context_range.start + editable_range_in_context.end);
+    (editable_range, context_range)
+}
+
+pub async fn run_format_prompt(
+    example: &mut Example,
+    args: &FormatPromptArgs,
+    app_state: Arc<EpAppState>,
+    example_progress: &ExampleProgress,
+    cx: AsyncApp,
+) -> Result<()> {
+    run_context_retrieval(example, app_state.clone(), example_progress, cx.clone()).await?;
+
+    let step_progress = example_progress.start(Step::FormatPrompt);
+
+    let prompt_inputs = example
+        .prompt_inputs
+        .as_ref()
+        .context("prompt_inputs must be set after context retrieval")?;
+
+    match args.provider {
+        PredictionProvider::Teacher(_, zeta_format)
+        | PredictionProvider::TeacherNonBatching(_, zeta_format) => {
+            step_progress.set_substatus("formatting teacher prompt");
+
+            let (editable_range, context_range) =
+                resolved_excerpt_ranges_for_format(prompt_inputs, zeta_format);
+
+            let include_diagnostics = matches!(zeta_format, ZetaFormat::V0420Diagnostics);
+
+            let prompt = TeacherPrompt::format_prompt(
+                example,
+                editable_range,
+                context_range,
+                include_diagnostics,
+            );
+            example.prompt = Some(ExamplePrompt {
+                input: prompt,
+                expected_output: None,
+                rejected_output: None,
+                prefill: None,
+                provider: args.provider,
+            });
+        }
+        PredictionProvider::TeacherMultiRegion(_)
+        | PredictionProvider::TeacherMultiRegionNonBatching(_) => {
+            step_progress.set_substatus("formatting teacher multi-region prompt");
+
+            let zeta_format = ZetaFormat::default();
+            let (editable_range, context_range) =
+                resolved_excerpt_ranges_for_format(prompt_inputs, zeta_format);
+
+            let include_diagnostics = matches!(zeta_format, ZetaFormat::V0420Diagnostics);
+
+            let prompt = TeacherMultiRegionPrompt::format_prompt(
+                example,
+                editable_range,
+                context_range,
+                include_diagnostics,
+            );
+            example.prompt = Some(ExamplePrompt {
+                input: prompt,
+                expected_output: None,
+                rejected_output: None,
+                prefill: None,
+                provider: args.provider,
+            });
+        }
+        PredictionProvider::Zeta2(zeta_format) => {
+            step_progress.set_substatus("formatting zeta2 prompt");
+
+            let prompt = format_zeta_prompt(prompt_inputs, zeta_format);
+            let prefill = zeta_prompt::get_prefill(prompt_inputs, zeta_format);
+            let expected_output = example
+                .spec
+                .expected_patches_with_cursor_positions()
+                .into_iter()
+                .next()
+                .and_then(|(expected_patch, expected_cursor_offset)| {
+                    format_expected_output(
+                        prompt_inputs,
+                        zeta_format,
+                        &expected_patch,
+                        expected_cursor_offset,
+                    )
+                    .ok()
+                });
+
+            let rejected_output = example.spec.rejected_patch.as_ref().and_then(|patch| {
+                format_expected_output(prompt_inputs, zeta_format, patch, None).ok()
+            });
+
+            example.prompt = prompt.map(|prompt| ExamplePrompt {
+                input: prompt,
+                expected_output,
+                rejected_output,
+                provider: args.provider,
+                prefill: Some(prefill),
+            });
+        }
+        _ => {
+            panic!("Cannot format prompt for {:?}", args.provider);
+        }
+    };
+    Ok(())
+}
+
+pub struct TeacherPrompt;
+
+impl TeacherPrompt {
+    pub(crate) const EDITABLE_REGION_START: &str = "<|editable_region_start|>\n";
+    pub(crate) const EDITABLE_REGION_END: &str = "\n<|editable_region_end|>";
+    pub(crate) const USER_CURSOR_MARKER: &str = "<|user_cursor|>";
+    pub(crate) const NO_EDITS: &str = "NO_EDITS";
+
+    /// Truncate edit history to this number of last lines
+    const MAX_HISTORY_LINES: usize = 128;
+
+    pub fn format_prompt(
+        example: &Example,
+        editable_range: Range<usize>,
+        context_range: Range<usize>,
+        include_diagnostics: bool,
+    ) -> String {
+        let edit_history = Self::format_edit_history(&example.spec.edit_history);
+        let context = Self::format_context(example);
+        let cursor_excerpt = Self::format_cursor_excerpt(example, editable_range, context_range);
+        let diagnostics = include_diagnostics
+            .then(|| Self::format_diagnostics(example))
+            .map(|diagnostics| format!("# 4. Diagnostics\n\n{diagnostics}"));
+
+        let prompt_template = crate::prompt_assets::get_prompt("teacher.md");
+        let prompt = prompt_template
+            .replace("{{context}}", &context)
+            .replace("{{edit_history}}", &edit_history)
+            .replace("{{diagnostics}}", diagnostics.as_deref().unwrap_or(""))
+            .replace("{{cursor_excerpt}}", &cursor_excerpt);
+
+        prompt
+    }
+
+    pub fn parse(example: &Example, response: &str) -> Result<(String, Option<ActualCursor>)> {
+        // Check if the model indicated no edits are needed
+        let no_edits = (String::new(), None);
+        if let Some(last_codeblock) = extract_last_codeblock(&response) {
+            if last_codeblock.trim() == Self::NO_EDITS {
+                return Ok(no_edits);
+            }
+        }
+
+        if response
+            .trim_end_matches(&[' ', '\n', '`'])
+            .ends_with(Self::NO_EDITS)
+        {
+            return Ok(no_edits);
+        }
+
+        // Extract updated (new) editable region from the model response.
+        let new_editable_region = Self::extract_editable_region(&response)?;
+        let cursor_offset = new_editable_region.find(Self::USER_CURSOR_MARKER);
+        let mut new_editable_region = new_editable_region.replace(Self::USER_CURSOR_MARKER, "");
+        let old_editable_region = Self::extract_editable_region(
+            &example
+                .prompt
+                .as_ref()
+                .context("example prompt missing")?
+                .input,
+        )?
+        .replace(Self::USER_CURSOR_MARKER, "");
+
+        let prompt_inputs = example
+            .prompt_inputs
+            .as_ref()
+            .context("example is missing prompt inputs")?;
+
+        // Normalize leading newlines: if old starts with newline but new doesn't,
+        // prepend newline to new to preserve whitespace structure.
+        // This handles the case where the model drops the leading blank line.
+        if old_editable_region.starts_with('\n') && !new_editable_region.starts_with('\n') {
+            new_editable_region.insert(0, '\n');
+        }
+
+        let excerpt = prompt_inputs.cursor_excerpt.as_ref();
+        let (editable_region_offset, _) = excerpt
+            .match_indices(&old_editable_region)
+            .min_by_key(|(index, _)| index.abs_diff(prompt_inputs.cursor_offset_in_excerpt))
+            .context("editable region not found in prompt content")?;
+        let editable_region_start_line = excerpt[..editable_region_offset].matches('\n').count();
+
+        let editable_region_lines = old_editable_region.lines().count() as u32;
+        let diff = language::unified_diff_with_context(
+            &old_editable_region,
+            &new_editable_region,
+            editable_region_start_line as u32,
+            editable_region_start_line as u32,
+            editable_region_lines,
+        );
+
+        let diff = indoc::formatdoc! {"
+            --- a/{path}
+            +++ b/{path}
+            {diff}",
+            path = example.spec.cursor_path.to_string_lossy(),
+            diff = diff,
+        };
+
+        let actual_cursor = cursor_offset.map(|editable_region_cursor_offset| {
+            ActualCursor::from_editable_region(
+                &example.spec.cursor_path,
+                editable_region_cursor_offset,
+                &new_editable_region,
+                excerpt,
+                editable_region_offset,
+                editable_region_start_line,
+            )
+        });
+
+        Ok((diff, actual_cursor))
+    }
+
+    fn format_edit_history(edit_history: &str) -> String {
+        let lines: Vec<&str> = edit_history.lines().collect();
+
+        if lines.is_empty() {
+            return "(No edit history)".to_string();
+        }
+
+        if lines.len() > Self::MAX_HISTORY_LINES {
+            let truncated = lines[lines.len() - Self::MAX_HISTORY_LINES..].join("\n");
+            format!("{truncated}\n[...truncated...]")
+        } else {
+            lines.join("\n")
+        }
+    }
+
+    pub fn format_context(example: &Example) -> String {
+        let related_files = example
+            .prompt_inputs
+            .as_ref()
+            .and_then(|pi| pi.related_files.as_deref());
+
+        let Some(related_files) = related_files else {
+            return "(No context)".to_string();
+        };
+
+        if related_files.is_empty() {
+            return "(No context)".to_string();
+        }
+
+        let prefix = "`````";
+        let suffix = "`````\n\n";
+        let max_tokens = 1024;
+        zeta_prompt::format_related_files_within_budget(related_files, &prefix, &suffix, max_tokens)
+    }
+
+    fn format_cursor_excerpt(
+        example: &Example,
+        editable_range: Range<usize>,
+        context_range: Range<usize>,
+    ) -> String {
+        let mut result = String::new();
+
+        let prompt_inputs = example.prompt_inputs.as_ref().unwrap();
+        let excerpt = prompt_inputs.cursor_excerpt.as_ref();
+        let cursor_offset = prompt_inputs.cursor_offset_in_excerpt;
+
+        let path_str = example.spec.cursor_path.to_string_lossy();
+        result.push_str(&format!("`````{path_str}\n"));
+        result.push_str(&excerpt[context_range.start..editable_range.start]);
+        result.push_str(Self::EDITABLE_REGION_START);
+        result.push_str(&excerpt[editable_range.start..cursor_offset]);
+        result.push_str(Self::USER_CURSOR_MARKER);
+        result.push_str(&excerpt[cursor_offset..editable_range.end]);
+        result.push_str(Self::EDITABLE_REGION_END);
+        result.push_str(&excerpt[editable_range.end..context_range.end]);
+        result.push_str("\n`````");
+
+        result
+    }
+
+    pub fn extract_editable_region(text: &str) -> Result<String> {
+        let start = text
+            .rfind(Self::EDITABLE_REGION_START)
+            .map_or(0, |pos| pos + Self::EDITABLE_REGION_START.len());
+        let end = text.rfind(Self::EDITABLE_REGION_END).unwrap_or(text.len());
+
+        if start >= end {
+            return Err(anyhow!("Invalid editable region markers"));
+        }
+
+        let region = &text[start..end];
+        Ok(region.strip_suffix('\n').unwrap_or(region).to_string())
+    }
+
+    fn format_diagnostics(example: &Example) -> String {
+        let Some(prompt_inputs) = example.prompt_inputs.as_ref() else {
+            return "No Diagnostics".to_string();
+        };
+
+        let cursor_buffer_row = prompt_inputs.excerpt_start_row.map(|excerpt_start_row| {
+            excerpt_start_row
+                + prompt_inputs.cursor_excerpt[..prompt_inputs.cursor_offset_in_excerpt]
+                    .bytes()
+                    .filter(|byte| *byte == b'\n')
+                    .count() as u32
+        });
+        let diagnostics = zeta_prompt::format_active_buffer_diagnostics_with_budget(
+            &prompt_inputs.active_buffer_diagnostics,
+            cursor_buffer_row,
+            2_000,
+        );
+
+        let diagnostics = diagnostics
+            .strip_prefix("<filename>diagnostics\n")
+            .unwrap_or(&diagnostics);
+
+        if diagnostics.is_empty() {
+            "No Diagnostics".to_string()
+        } else {
+            diagnostics.to_string()
+        }
+    }
+}
+
+pub struct TeacherMultiRegionPrompt;
+
+impl TeacherMultiRegionPrompt {
+    pub(crate) const USER_CURSOR_MARKER: &str = "<|user_cursor|>";
+    pub(crate) const NO_EDITS: &str = "NO_EDITS";
+
+    /// Truncate edit history to this number of last lines
+    const MAX_HISTORY_LINES: usize = 128;
+
+    pub fn format_prompt(
+        example: &Example,
+        editable_range: Range<usize>,
+        context_range: Range<usize>,
+        include_diagnostics: bool,
+    ) -> String {
+        let edit_history = Self::format_edit_history(&example.spec.edit_history);
+        let context = Self::format_context(example);
+        let cursor_excerpt = Self::format_cursor_excerpt(example, editable_range, context_range);
+        let diagnostics = include_diagnostics
+            .then(|| TeacherPrompt::format_diagnostics(example))
+            .map(|diagnostics| format!("# 4. Diagnostics\n\n{diagnostics}"));
+
+        let prompt_template = crate::prompt_assets::get_prompt("teacher_multi_region.md");
+        let prompt = prompt_template
+            .replace("{{context}}", &context)
+            .replace("{{edit_history}}", &edit_history)
+            .replace("{{diagnostics}}", diagnostics.as_deref().unwrap_or(""))
+            .replace("{{cursor_excerpt}}", &cursor_excerpt);
+
+        prompt
+    }
+
+    pub fn parse(example: &Example, response: &str) -> Result<(String, Option<ActualCursor>)> {
+        let no_edits = (String::new(), None);
+        if let Some(last_codeblock) = extract_last_codeblock(&response) {
+            if last_codeblock.trim() == Self::NO_EDITS {
+                return Ok(no_edits);
+            }
+        }
+
+        if response.trim().ends_with(Self::NO_EDITS) {
+            return Ok(no_edits);
+        }
+
+        let prompt_inputs = example
+            .prompt_inputs
+            .as_ref()
+            .context("example is missing prompt inputs")?;
+
+        let zeta_format = ZetaFormat::default();
+        let (editable_range, _) = resolved_excerpt_ranges_for_format(prompt_inputs, zeta_format);
+        let excerpt = prompt_inputs.cursor_excerpt.as_ref();
+        let old_editable_region = &excerpt[editable_range.clone()];
+        let marker_offsets = multi_region::compute_marker_offsets(old_editable_region);
+
+        let codeblock =
+            extract_last_codeblock(&response).context("no codeblock found in model response")?;
+        let (start_num, end_num, raw_new_span) = multi_region::extract_marker_span(&codeblock)?;
+
+        let start_idx = start_num
+            .checked_sub(1)
+            .context("marker numbers are 1-indexed")?;
+        let end_idx = end_num
+            .checked_sub(1)
+            .context("marker numbers are 1-indexed")?;
+        let start_byte = *marker_offsets
+            .get(start_idx)
+            .context("start marker number out of range")?;
+        let end_byte = *marker_offsets
+            .get(end_idx)
+            .context("end marker number out of range")?;
+
+        if start_byte > end_byte {
+            return Err(anyhow!("start marker must come before end marker"));
+        }
+
+        let cursor_in_span = raw_new_span.find(Self::USER_CURSOR_MARKER);
+        let new_span = raw_new_span.replace(Self::USER_CURSOR_MARKER, "");
+
+        let old_span = &old_editable_region[start_byte..end_byte];
+        let mut new_span = new_span;
+        if old_span.ends_with('\n') && !new_span.ends_with('\n') && !new_span.is_empty() {
+            new_span.push('\n');
+        }
+        if !old_span.ends_with('\n') && new_span.ends_with('\n') {
+            new_span.pop();
+        }
+
+        let mut new_editable_region = String::new();
+        new_editable_region.push_str(&old_editable_region[..start_byte]);
+        new_editable_region.push_str(&new_span);
+        new_editable_region.push_str(&old_editable_region[end_byte..]);
+
+        let cursor_offset = cursor_in_span.map(|pos| start_byte + pos);
+
+        if old_editable_region.starts_with('\n') && !new_editable_region.starts_with('\n') {
+            new_editable_region.insert(0, '\n');
+        }
+
+        let editable_region_offset = editable_range.start;
+        let editable_region_start_line = excerpt[..editable_region_offset].matches('\n').count();
+
+        let editable_region_lines = old_editable_region.lines().count() as u32;
+        let diff = language::unified_diff_with_context(
+            old_editable_region,
+            &new_editable_region,
+            editable_region_start_line as u32,
+            editable_region_start_line as u32,
+            editable_region_lines,
+        );
+
+        let diff = indoc::formatdoc! {"
+            --- a/{path}
+            +++ b/{path}
+            {diff}",
+            path = example.spec.cursor_path.to_string_lossy(),
+            diff = diff,
+        };
+
+        let actual_cursor = cursor_offset.map(|editable_region_cursor_offset| {
+            ActualCursor::from_editable_region(
+                &example.spec.cursor_path,
+                editable_region_cursor_offset,
+                &new_editable_region,
+                excerpt,
+                editable_region_offset,
+                editable_region_start_line,
+            )
+        });
+
+        Ok((diff, actual_cursor))
+    }
+
+    fn format_edit_history(edit_history: &str) -> String {
+        let lines: Vec<&str> = edit_history.lines().collect();
+
+        if lines.is_empty() {
+            return "(No edit history)".to_string();
+        }
+
+        if lines.len() > Self::MAX_HISTORY_LINES {
+            let truncated = lines[lines.len() - Self::MAX_HISTORY_LINES..].join("\n");
+            format!("{truncated}\n[...truncated...]")
+        } else {
+            lines.join("\n")
+        }
+    }
+
+    pub fn format_context(example: &Example) -> String {
+        let related_files = example
+            .prompt_inputs
+            .as_ref()
+            .and_then(|pi| pi.related_files.as_deref());
+        let Some(related_files) = related_files else {
+            return "(No context)".to_string();
+        };
+
+        if related_files.is_empty() {
+            return "(No context)".to_string();
+        }
+
+        let prefix = "`````";
+        let suffix = "`````\n\n";
+        let max_tokens = 1024;
+        zeta_prompt::format_related_files_within_budget(related_files, &prefix, &suffix, max_tokens)
+    }
+
+    fn format_cursor_excerpt(
+        example: &Example,
+        editable_range: Range<usize>,
+        context_range: Range<usize>,
+    ) -> String {
+        let mut result = String::new();
+
+        let prompt_inputs = example.prompt_inputs.as_ref().unwrap();
+        let excerpt = prompt_inputs.cursor_excerpt.as_ref();
+        let cursor_offset = prompt_inputs.cursor_offset_in_excerpt;
+
+        let editable_text = &excerpt[editable_range.clone()];
+        let cursor_in_editable = cursor_offset - editable_range.start;
+
+        let path_str = example.spec.cursor_path.to_string_lossy();
+        result.push_str(&format!("`````{path_str}\n"));
+
+        result.push_str(&excerpt[context_range.start..editable_range.start]);
+
+        multi_region::write_editable_with_markers(
+            &mut result,
+            editable_text,
+            cursor_in_editable,
+            Self::USER_CURSOR_MARKER,
+        );
+
+        result.push_str(&excerpt[editable_range.end..context_range.end]);
+        result.push_str("\n`````");
+
+        result
+    }
+}
+
+/// Extract the cursor excerpt from an example.
+/// First tries to extract from an existing prompt, then falls back to constructing from prompt_inputs.
+pub fn extract_cursor_excerpt_from_example(example: &Example) -> Option<String> {
+    // If we have the original prompt, extract the cursor excerpt from it
+    if let Some(prompt) = &example.prompt {
+        // Find "# 3. Current File" section and extract the content
+        if let Some(start) = prompt.input.find("# 3. Current File") {
+            let content_start = prompt.input[start..].find('`').map(|i| start + i)?;
+            let backtick_count = prompt.input[content_start..]
+                .chars()
+                .take_while(|&c| c == '`')
+                .count();
+            let content_start = content_start + backtick_count;
+
+            // Find the path line and skip it
+            let newline_pos = prompt.input[content_start..].find('\n')?;
+            let text_start = content_start + newline_pos + 1;
+
+            // Find the closing backticks
+            let closing_pattern = "`".repeat(backtick_count);
+            let text_end = prompt.input[text_start..].find(&closing_pattern)?;
+            let cursor_excerpt = &prompt.input[text_start..text_start + text_end];
+
+            let path_str = example.spec.cursor_path.to_string_lossy();
+            return Some(format!("`````{path_str}\n{cursor_excerpt}`````"));
+        }
+    }
+
+    // Fallback: construct from prompt_inputs if available
+    let prompt_inputs = example.prompt_inputs.as_ref()?;
+    let excerpt = prompt_inputs.cursor_excerpt.as_ref();
+    let cursor_offset = prompt_inputs.cursor_offset_in_excerpt;
+
+    // Simple fallback: just show content around cursor with markers
+    let path_str = example.spec.cursor_path.to_string_lossy();
+    let mut result = format!("`````{path_str}\n");
+    result.push_str(TeacherPrompt::EDITABLE_REGION_START);
+    result.push_str(&excerpt[..cursor_offset]);
+    result.push_str(TeacherPrompt::USER_CURSOR_MARKER);
+    result.push_str(&excerpt[cursor_offset..]);
+    result.push_str(TeacherPrompt::EDITABLE_REGION_END);
+    result.push_str("\n`````");
+
+    Some(result)
+}
+
+pub(crate) fn extract_last_codeblock(text: &str) -> Option<String> {
+    let lines: Vec<&str> = text.lines().collect();
+
+    // Search from the end for a closing fence (line containing only backticks, 3+)
+    let mut closing_line_idx = None;
+    let mut backtick_count = 0;
+
+    for i in (0..lines.len()).rev() {
+        let line = lines[i].trim();
+        if line.len() >= 3 && line.chars().all(|c| c == '`') {
+            closing_line_idx = Some(i);
+            backtick_count = line.len();
+            break;
+        }
+    }
+
+    let closing_idx = closing_line_idx?;
+
+    // Search backwards for matching opening fence
+    // Opening fence starts with same backtick count, possibly followed by language/metadata
+    let opening_pattern = "`".repeat(backtick_count);
+
+    for i in (0..closing_idx).rev() {
+        let line = lines[i];
+        if line.starts_with(&opening_pattern) {
+            // Ensure it's exactly the right number of backticks (not more)
+            let rest = &line[backtick_count..];
+            if rest.is_empty() || !rest.starts_with('`') {
+                // Found matching opening fence
+                // Extract content between opening and closing (exclusive)
+                if closing_idx > i + 1 {
+                    let content = lines[i + 1..closing_idx].join("\n");
+                    // Preserve trailing newline to match previous behavior
+                    return Some(format!("{}\n", content));
+                } else {
+                    // Empty block
+                    return Some(String::new());
+                }
+            }
+        }
+    }
+
+    None
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_extract_last_code_block() {
+        let text = indoc::indoc! {"
+            Some thinking
+
+            ```
+            first block
+            ```
+
+            `````path='something' lines=1:2
+            last block
+            `````
+            "};
+        let last_block = extract_last_codeblock(text).unwrap();
+        assert_eq!(last_block, "last block\n");
+    }
+
+    #[test]
+    fn test_extract_codeblock_with_nested_fences() {
+        let text = indoc::indoc! {"
+            `````
+            content with ``` inline
+            and ```python nested
+            more content
+            `````
+            "};
+        let last_block = extract_last_codeblock(text).unwrap();
+        assert_eq!(
+            last_block,
+            "content with ``` inline\nand ```python nested\nmore content\n"
+        );
+    }
+
+    #[test]
+    fn test_extract_codeblock_ignores_inline_backticks() {
+        let text = indoc::indoc! {"
+            `````
+            here is some `code` with inline backticks
+            and here```more```stuff
+            `````
+            "};
+        let last_block = extract_last_codeblock(text).unwrap();
+        assert_eq!(
+            last_block,
+            "here is some `code` with inline backticks\nand here```more```stuff\n"
+        );
+    }
+
+    #[test]
+    fn test_extract_editable_region_old_format() {
+        let text = indoc::indoc! {"
+            some lines
+            are
+            here
+            <|editable_region_start|>
+            one
+            two three
+
+            <|editable_region_end|>
+            more
+            lines here
+            "};
+        let parsed = TeacherPrompt::extract_editable_region(text).unwrap();
+        assert_eq!(
+            parsed,
+            indoc::indoc! {"
+            one
+            two three"}
+        );
+    }
+
+    #[test]
+    fn test_extract_editable_region_marker_format() {
+        let text = indoc::indoc! {"
+            some context
+            <|marker_1|>
+            one
+            two three
+            <|marker_2|>
+            more context
+            "};
+        let parsed = multi_region::extract_editable_region_from_markers(text).unwrap();
+        assert_eq!(parsed, "one\ntwo three");
+    }
+
+    #[test]
+    fn test_extract_editable_region_multi_markers() {
+        let text = indoc::indoc! {"
+            prefix
+            <|marker_1|>
+            aaa
+            bbb
+            <|marker_2|>
+            ccc
+            ddd
+            <|marker_3|>
+            suffix
+            "};
+        let parsed = multi_region::extract_editable_region_from_markers(text).unwrap();
+        // Intermediate marker and its trailing \n are stripped
+        assert_eq!(parsed, "aaa\nbbb\nccc\nddd");
+    }
+
+    #[test]
+    fn test_extract_last_codeblock_nested_bibtex() {
+        let text = indoc::indoc! {r#"
+            Looking at the edit history, I can see that a Citation section was just added.
+
+            `````
+            ## Collaborations
+            Our mission is to create a 4D generative model.
+
+            ## Citation
+
+            If you found Unique3D helpful, please cite our report:
+            ```bibtex
+            @misc{wu2024unique3d,
+                  title={Unique3D},
+            }
+            ```
+            `````
+            "#};
+        let last_block = extract_last_codeblock(text).unwrap();
+        assert_eq!(
+            last_block,
+            indoc::indoc! {r#"
+            ## Collaborations
+            Our mission is to create a 4D generative model.
+
+            ## Citation
+
+            If you found Unique3D helpful, please cite our report:
+            ```bibtex
+            @misc{wu2024unique3d,
+                  title={Unique3D},
+            }
+            ```
+            "#}
+        );
+    }
+
+    #[test]
+    fn test_extract_editable_region_no_markers() {
+        let text = indoc::indoc! {"
+            one
+            two three"};
+        let parsed = TeacherPrompt::extract_editable_region(text).unwrap();
+        assert_eq!(
+            parsed,
+            indoc::indoc! {"
+            one
+            two three"}
+        );
+    }
+
+    #[test]
+    fn test_parse_no_edits_response() {
+        let response = indoc::indoc! {"
+            The code is already complete. There is no clear next edit to make.
+
+            `````
+            NO_EDITS
+            `````
+        "};
+        let codeblock = extract_last_codeblock(response).unwrap();
+        assert_eq!(codeblock.trim(), TeacherPrompt::NO_EDITS);
+    }
+
+    #[test]
+    fn test_extract_codeblock_no_valid_block() {
+        // Text with no code blocks should return None
+        let text = "Just some plain text without any code blocks";
+        assert!(extract_last_codeblock(text).is_none());
+
+        // Unclosed code block should return None
+        let text = indoc::indoc! {"
+            ```
+            unclosed block
+        "};
+        assert!(extract_last_codeblock(text).is_none());
+
+        // Analysis text with nested markdown but no proper outer block
+        let text = indoc::indoc! {"
+            # Analysis
+            Looking at this:
+            ```
+            some code
+            ```
+            But then more analysis without wrapping block
+        "};
+        // This should find the inner block
+        let result = extract_last_codeblock(text).unwrap();
+        assert_eq!(result, "some code\n");
+    }
+
+    #[test]
+    fn test_extract_codeblock_no_trailing_newline() {
+        // Text ending without trailing newline after closing fence
+        let text = "`````\ncontent here\n`````";
+        let result = extract_last_codeblock(text).unwrap();
+        assert_eq!(result, "content here\n");
+    }
+
+    #[test]
+    fn test_parse_no_edits_response_with_trailing_backticks() {
+        let response = "NO_EDITS```";
+
+        let parsed = TeacherPrompt::parse(
+            &Example {
+                spec: edit_prediction::example_spec::ExampleSpec {
+                    name: "test".to_string(),
+                    repository_url: "https://github.com/zed-industries/zed.git".to_string(),
+                    revision: "HEAD".to_string(),
+                    tags: Vec::new(),
+                    reasoning: None,
+                    uncommitted_diff: String::new(),
+                    cursor_path: std::sync::Arc::from(std::path::Path::new("src/main.rs")),
+                    cursor_position: "0:0".to_string(),
+                    edit_history: String::new(),
+                    expected_patches: Vec::new(),
+                    rejected_patch: None,
+                    telemetry: None,
+                    human_feedback: Vec::new(),
+                    rating: None,
+                },
+                prompt_inputs: None,
+                prompt: None,
+                predictions: Vec::new(),
+                score: Vec::new(),
+                qa: Vec::new(),
+                zed_version: None,
+                state: None,
+            },
+            response,
+        )
+        .unwrap();
+
+        assert!(parsed.0.is_empty());
+        assert!(parsed.1.is_none());
+    }
+
+    #[test]
+    fn test_v0327_teacher_prompt_uses_resolved_ranges() {
+        let excerpt = (0..80)
+            .map(|index| format!("line{index:02}\n"))
+            .collect::<String>();
+        let cursor_offset = excerpt.find("line40").expect("cursor line exists");
+        let prompt_inputs = zeta_prompt::ZetaPromptInput {
+            cursor_path: std::path::Path::new("src/main.rs").into(),
+            cursor_excerpt: excerpt.clone().into(),
+            cursor_offset_in_excerpt: cursor_offset,
+            excerpt_start_row: None,
+            events: Vec::new(),
+            related_files: Some(Vec::new()),
+            active_buffer_diagnostics: Vec::new(),
+            excerpt_ranges: zeta_prompt::ExcerptRanges {
+                editable_150: 0..32,
+                editable_180: 0..32,
+                editable_350: 0..32,
+                editable_512: None,
+                editable_150_context_350: 0..48,
+                editable_180_context_350: 0..48,
+                editable_350_context_150: 20..50,
+                editable_350_context_512: None,
+                editable_350_context_1024: None,
+                context_4096: None,
+                context_8192: Some(30..excerpt.len()),
+            },
+            syntax_ranges: None,
+            in_open_source_repo: false,
+            can_collect_data: false,
+            repo_url: None,
+        };
+
+        let (stored_editable_range, stored_context_range) = zeta_prompt::excerpt_range_for_format(
+            ZetaFormat::V0327SingleFile,
+            &prompt_inputs.excerpt_ranges,
+        );
+        assert!(stored_context_range.start > stored_editable_range.start);
+
+        let (editable_range, context_range) =
+            resolved_excerpt_ranges_for_format(&prompt_inputs, ZetaFormat::V0327SingleFile);
+        assert_eq!(context_range, 0..excerpt.len());
+        assert!(editable_range.start < cursor_offset);
+        assert!(editable_range.end > cursor_offset);
+
+        let prompt = TeacherPrompt::format_prompt(
+            &Example {
+                spec: edit_prediction::example_spec::ExampleSpec {
+                    name: "test".to_string(),
+                    repository_url: "https://github.com/zed-industries/zed.git".to_string(),
+                    revision: "HEAD".to_string(),
+                    tags: Vec::new(),
+                    reasoning: None,
+                    uncommitted_diff: String::new(),
+                    cursor_path: std::sync::Arc::from(std::path::Path::new("src/main.rs")),
+                    cursor_position: "0:0".to_string(),
+                    edit_history: String::new(),
+                    expected_patches: Vec::new(),
+                    rejected_patch: None,
+                    telemetry: None,
+                    human_feedback: Vec::new(),
+                    rating: None,
+                },
+                prompt_inputs: Some(prompt_inputs),
+                prompt: None,
+                predictions: Vec::new(),
+                score: Vec::new(),
+                qa: Vec::new(),
+                zed_version: None,
+                state: None,
+            },
+            editable_range,
+            context_range,
+            false,
+        );
+
+        assert!(prompt.contains(TeacherPrompt::EDITABLE_REGION_START));
+        assert!(prompt.contains(TeacherPrompt::USER_CURSOR_MARKER));
+        assert!(prompt.contains("line40"));
+    }
+}

crates/edit_prediction_cli/src/git.rs

@@ -0,0 +1,129 @@
+use anyhow::{Context as _, Result};
+use collections::HashMap;
+use futures::lock::{Mutex, OwnedMutexGuard};
+use std::{
+    cell::RefCell,
+    path::{Path, PathBuf},
+    sync::Arc,
+};
+
+use crate::paths::REPOS_DIR;
+
+thread_local! {
+    static REPO_LOCKS: RefCell<HashMap<PathBuf, Arc<Mutex<()>>>> = RefCell::new(HashMap::default());
+}
+
+#[must_use]
+pub async fn lock_repo(path: impl AsRef<Path>) -> OwnedMutexGuard<()> {
+    REPO_LOCKS
+        .with(|cell| {
+            cell.borrow_mut()
+                .entry(path.as_ref().to_path_buf())
+                .or_default()
+                .clone()
+        })
+        .lock_owned()
+        .await
+}
+
+pub async fn run_git(repo_path: &Path, args: &[&str]) -> Result<String> {
+    let output = smol::process::Command::new("git")
+        .current_dir(repo_path)
+        .args(args)
+        .output()
+        .await?;
+
+    anyhow::ensure!(
+        output.status.success(),
+        "`git {}` within `{}` failed with status: {}\nstderr:\n{}\nstdout:\n{}",
+        args.join(" "),
+        repo_path.display(),
+        output.status,
+        String::from_utf8_lossy(&output.stderr),
+        String::from_utf8_lossy(&output.stdout),
+    );
+    Ok(String::from_utf8(output.stdout)?.trim().to_string())
+}
+
+pub fn parse_repo_url(url: &str) -> Result<(String, String)> {
+    if url.contains('@') {
+        let (_, path) = url.split_once(':').context("expected : in git url")?;
+        let (owner, repo) = path.split_once('/').context("expected / in git url")?;
+        Ok((owner.to_string(), repo.trim_end_matches(".git").to_string()))
+    } else {
+        let parsed = http_client::Url::parse(url)?;
+        let mut segments = parsed.path_segments().context("empty http url")?;
+        let owner = segments.next().context("expected owner")?;
+        let repo = segments.next().context("expected repo")?;
+        Ok((owner.to_string(), repo.trim_end_matches(".git").to_string()))
+    }
+}
+
+pub fn repo_path_for_url(url: &str) -> Result<PathBuf> {
+    let (owner, name) = parse_repo_url(url)?;
+    Ok(REPOS_DIR.join(&owner).join(&name))
+}
+
+pub async fn ensure_repo_cloned(repo_url: &str) -> Result<PathBuf> {
+    let repo_path = repo_path_for_url(repo_url)?;
+    let _lock = lock_repo(&repo_path).await;
+
+    // Validate existing repo has correct origin, otherwise remove and re-init.
+    let mut git_repo_exists = false;
+    if repo_path.is_dir() {
+        if run_git(&repo_path, &["remote", "get-url", "origin"])
+            .await
+            .map_or(false, |origin| origin.trim() == repo_url)
+        {
+            git_repo_exists = true;
+        } else {
+            std::fs::remove_dir_all(&repo_path).ok();
+        }
+    }
+
+    if !git_repo_exists {
+        log::info!("Cloning {} into {:?}", repo_url, repo_path);
+        std::fs::create_dir_all(&repo_path)?;
+        run_git(&repo_path, &["init"]).await?;
+        run_git(&repo_path, &["remote", "add", "origin", repo_url])
+            .await
+            .ok();
+    }
+
+    // Always fetch to get latest commits
+    run_git(&repo_path, &["fetch", "--depth", "1000", "origin"]).await?;
+
+    // Check if we have a valid HEAD, if not checkout FETCH_HEAD
+    let has_head = run_git(&repo_path, &["rev-parse", "HEAD"]).await.is_ok();
+    if !has_head {
+        // Use reset to set HEAD without needing a branch
+        run_git(&repo_path, &["reset", "--hard", "FETCH_HEAD"]).await?;
+    }
+
+    Ok(repo_path)
+}
+
+pub async fn fetch_if_needed(repo_path: &Path, revision: &str) -> Result<String> {
+    let resolved = run_git(
+        repo_path,
+        &["rev-parse", &format!("{}^{{commit}}", revision)],
+    )
+    .await;
+
+    if let Ok(sha) = resolved {
+        return Ok(sha);
+    }
+
+    if run_git(repo_path, &["fetch", "--depth", "1", "origin", revision])
+        .await
+        .is_err()
+    {
+        run_git(repo_path, &["fetch", "origin"]).await?;
+    }
+
+    run_git(
+        repo_path,
+        &["rev-parse", &format!("{}^{{commit}}", revision)],
+    )
+    .await
+}

crates/edit_prediction_cli/src/headless.rs

@@ -0,0 +1,126 @@
+use client::{Client, ProxySettings, RefreshLlmTokenListener, UserStore};
+use db::AppDatabase;
+use extension::ExtensionHostProxy;
+use fs::RealFs;
+use gpui::http_client::read_proxy_from_env;
+use gpui::{App, AppContext, Entity};
+use gpui_tokio::Tokio;
+use language::LanguageRegistry;
+use language_extension::LspAccess;
+use node_runtime::{NodeBinaryOptions, NodeRuntime};
+use project::project_settings::ProjectSettings;
+use release_channel::{AppCommitSha, AppVersion};
+use reqwest_client::ReqwestClient;
+use settings::{Settings, SettingsStore};
+use std::path::PathBuf;
+use std::sync::Arc;
+use util::ResultExt as _;
+
+/// Headless subset of `workspace::AppState`.
+pub struct EpAppState {
+    pub languages: Arc<LanguageRegistry>,
+    pub client: Arc<Client>,
+    pub user_store: Entity<UserStore>,
+    pub fs: Arc<dyn fs::Fs>,
+    pub node_runtime: NodeRuntime,
+}
+
+pub fn init(cx: &mut App) -> EpAppState {
+    let app_commit_sha = option_env!("ZED_COMMIT_SHA").map(|s| AppCommitSha::new(s.to_owned()));
+
+    let app_version = AppVersion::load(
+        env!("ZED_PKG_VERSION"),
+        option_env!("ZED_BUILD_ID"),
+        app_commit_sha,
+    );
+    release_channel::init(app_version.clone(), cx);
+    gpui_tokio::init(cx);
+
+    let settings_store = SettingsStore::new(cx, &settings::default_settings());
+    cx.set_global(settings_store);
+
+    // Set User-Agent so we can download language servers from GitHub
+    let user_agent = format!(
+        "Zeta CLI/{} ({}; {})",
+        app_version,
+        std::env::consts::OS,
+        std::env::consts::ARCH
+    );
+    let proxy_str = ProxySettings::get_global(cx).proxy.to_owned();
+    let proxy_url = proxy_str
+        .as_ref()
+        .and_then(|input| input.parse().ok())
+        .or_else(read_proxy_from_env);
+    let http = {
+        let _guard = Tokio::handle(cx).enter();
+
+        ReqwestClient::proxy_and_user_agent(proxy_url, &user_agent)
+            .expect("could not start HTTP client")
+    };
+    cx.set_http_client(Arc::new(http));
+
+    let client = Client::production(cx);
+    cx.set_http_client(client.http_client());
+
+    let app_db = AppDatabase::new();
+    cx.set_global(app_db);
+
+    let git_binary_path = None;
+    let fs = Arc::new(RealFs::new(
+        git_binary_path,
+        cx.background_executor().clone(),
+    ));
+
+    let mut languages = LanguageRegistry::new(cx.background_executor().clone());
+    languages.set_language_server_download_dir(paths::languages_dir().clone());
+    let languages = Arc::new(languages);
+
+    let user_store = cx.new(|cx| UserStore::new(client.clone(), cx));
+
+    extension::init(cx);
+
+    let (mut tx, rx) = watch::channel(None);
+    cx.observe_global::<SettingsStore>(move |cx| {
+        let settings = &ProjectSettings::get_global(cx).node;
+        let options = NodeBinaryOptions {
+            allow_path_lookup: !settings.ignore_system_version,
+            allow_binary_download: true,
+            use_paths: settings.path.as_ref().map(|node_path| {
+                let node_path = PathBuf::from(shellexpand::tilde(node_path).as_ref());
+                let npm_path = settings
+                    .npm_path
+                    .as_ref()
+                    .map(|path| PathBuf::from(shellexpand::tilde(&path).as_ref()));
+                (
+                    node_path.clone(),
+                    npm_path.unwrap_or_else(|| {
+                        let base_path = PathBuf::new();
+                        node_path.parent().unwrap_or(&base_path).join("npm")
+                    }),
+                )
+            }),
+        };
+        tx.send(Some(options)).log_err();
+    })
+    .detach();
+    let node_runtime = NodeRuntime::new(client.http_client(), None, rx);
+
+    let extension_host_proxy = ExtensionHostProxy::global(cx);
+
+    debug_adapter_extension::init(extension_host_proxy.clone(), cx);
+    language_extension::init(LspAccess::Noop, extension_host_proxy, languages.clone());
+    language_model::init(cx);
+    RefreshLlmTokenListener::register(client.clone(), user_store.clone(), cx);
+    language_models::init(user_store.clone(), client.clone(), cx);
+    languages::init(languages.clone(), fs.clone(), node_runtime.clone(), cx);
+    prompt_store::init(cx);
+    terminal_view::init(cx);
+
+    EpAppState {
+        languages,
+        client,
+        user_store,
+        fs,
+        node_runtime,
+    }
+}

crates/edit_prediction_cli/src/lib.rs

@@ -0,0 +1,2 @@
+#[allow(dead_code)]
+mod word_diff;

crates/edit_prediction_cli/src/load_project.rs

@@ -0,0 +1,393 @@
+use crate::{
+    example::{Example, ExampleState},
+    git,
+    headless::EpAppState,
+    progress::{ExampleProgress, InfoStyle, Step, StepProgress},
+};
+use anyhow::{Context as _, Result};
+use edit_prediction::{
+    EditPredictionStore,
+    cursor_excerpt::{compute_cursor_excerpt, compute_syntax_ranges},
+    udiff::{OpenedBuffers, refresh_worktree_entries, strip_diff_path_prefix},
+};
+use futures::AsyncWriteExt as _;
+use gpui::{AsyncApp, Entity};
+use language::{Anchor, Buffer, LanguageNotFound, ToOffset};
+use project::{Project, ProjectPath, buffer_store::BufferStoreEvent};
+use std::{fs, path::PathBuf, sync::Arc};
+use zeta_prompt::ZetaPromptInput;
+
+pub async fn run_load_project(
+    example: &mut Example,
+    app_state: Arc<EpAppState>,
+    example_progress: &ExampleProgress,
+    mut cx: AsyncApp,
+) -> Result<()> {
+    if example.state.is_some() {
+        return Ok(());
+    }
+
+    let progress = example_progress.start(Step::LoadProject);
+
+    let project = setup_project(example, &app_state, &progress, &mut cx).await?;
+
+    progress.set_substatus("applying edit history");
+    let open_buffers = apply_edit_history(example, &project, &mut cx).await?;
+
+    let ep_store = cx
+        .update(|cx| EditPredictionStore::try_global(cx))
+        .context("EditPredictionStore not initialized")?;
+
+    let recent_paths: Vec<ProjectPath> = open_buffers
+        .buffers()
+        .filter_map(|buffer| {
+            buffer.read_with(&cx, |buffer, cx| {
+                buffer
+                    .file()
+                    .map(|file| ProjectPath::from_file(file.as_ref(), cx))
+            })
+        })
+        .collect();
+
+    ep_store.update(&mut cx, |store, cx| {
+        store.set_recent_paths_for_project(&project, recent_paths, cx);
+    });
+
+    progress.set_substatus("resolving cursor");
+    let (buffer, cursor_position) =
+        cursor_position(example, &project, &open_buffers, &mut cx).await?;
+    buffer
+        .read_with(&cx, |buffer, _| buffer.parsing_idle())
+        .await;
+
+    let events: Vec<Arc<zeta_prompt::Event>> = ep_store.update(&mut cx, |store, cx| {
+        store
+            .edit_history_for_project(&project, cx)
+            .into_iter()
+            .map(|e| e.event)
+            .collect()
+    });
+
+    let existing_related_files = example
+        .prompt_inputs
+        .take()
+        .and_then(|inputs| inputs.related_files);
+
+    let (prompt_inputs, language_name) = buffer.read_with(&cx, |buffer, _cx| {
+        let snapshot = buffer.snapshot();
+        let cursor_offset = cursor_position.to_offset(&snapshot);
+        let language_name = buffer
+            .language()
+            .map(|l| l.name().to_string())
+            .unwrap_or_else(|| "Unknown".to_string());
+
+        let (excerpt_point_range, excerpt_offset_range, cursor_offset_in_excerpt) =
+            compute_cursor_excerpt(&snapshot, cursor_offset);
+
+        let cursor_excerpt: Arc<str> = buffer
+            .text_for_range(excerpt_offset_range.clone())
+            .collect::<String>()
+            .into();
+        let syntax_ranges = compute_syntax_ranges(&snapshot, cursor_offset, &excerpt_offset_range);
+        let excerpt_ranges = zeta_prompt::compute_legacy_excerpt_ranges(
+            &cursor_excerpt,
+            cursor_offset_in_excerpt,
+            &syntax_ranges,
+        );
+
+        (
+            ZetaPromptInput {
+                cursor_path: example.spec.cursor_path.clone(),
+                cursor_excerpt,
+                cursor_offset_in_excerpt,
+                excerpt_start_row: Some(excerpt_point_range.start.row),
+                events,
+                related_files: existing_related_files,
+                active_buffer_diagnostics: vec![],
+                excerpt_ranges,
+                syntax_ranges: Some(syntax_ranges),
+                in_open_source_repo: false,
+                can_collect_data: false,
+                repo_url: None,
+            },
+            language_name,
+        )
+    });
+
+    progress.set_info(language_name, InfoStyle::Normal);
+
+    example.prompt_inputs = Some(prompt_inputs);
+    example.state = Some(ExampleState {
+        buffer,
+        project,
+        cursor_position,
+        _open_buffers: open_buffers,
+    });
+    Ok(())
+}
+
+async fn cursor_position(
+    example: &Example,
+    project: &Entity<Project>,
+    open_buffers: &OpenedBuffers,
+    cx: &mut AsyncApp,
+) -> Result<(Entity<Buffer>, Anchor)> {
+    let language_registry = project.read_with(cx, |project, _| project.languages().clone());
+    let result = language_registry
+        .load_language_for_file_path(&example.spec.cursor_path)
+        .await;
+
+    if let Err(error) = result
+        && !error.is::<LanguageNotFound>()
+    {
+        return Err(error);
+    }
+
+    let cursor_path_str = example.spec.cursor_path.to_string_lossy();
+    // Also try cursor path with first component stripped - old examples may have
+    // paths like "zed/crates/foo.rs" instead of "crates/foo.rs".
+    let cursor_path_without_prefix: PathBuf =
+        example.spec.cursor_path.components().skip(1).collect();
+    let cursor_path_without_prefix_str = cursor_path_without_prefix.to_string_lossy();
+
+    // We try open_buffers first because the file might be new and not saved to disk
+    let cursor_buffer = if let Some(buffer) = open_buffers.get(cursor_path_str.as_ref()) {
+        buffer.clone()
+    } else if let Some(buffer) = open_buffers.get(cursor_path_without_prefix_str.as_ref()) {
+        buffer.clone()
+    } else {
+        // Since the worktree scanner is disabled, manually refresh entries for the cursor path.
+        if let Some(worktree) = project.read_with(cx, |project, cx| project.worktrees(cx).next()) {
+            refresh_worktree_entries(&worktree, [&*example.spec.cursor_path], cx).await?;
+        }
+
+        let cursor_path = project
+            .read_with(cx, |project, cx| {
+                project
+                    .find_project_path(&example.spec.cursor_path, cx)
+                    .or_else(|| project.find_project_path(&cursor_path_without_prefix, cx))
+            })
+            .with_context(|| {
+                format!(
+                    "failed to find cursor path {}",
+                    example.spec.cursor_path.display()
+                )
+            })?;
+
+        project
+            .update(cx, |project, cx| project.open_buffer(cursor_path, cx))
+            .await?
+    };
+
+    let (cursor_excerpt, cursor_offset_within_excerpt) = example.spec.cursor_excerpt()?;
+
+    let excerpt_offset = cursor_buffer.read_with(&*cx, |buffer, _cx| {
+        let text = buffer.text();
+
+        let mut matches = text.match_indices(&cursor_excerpt);
+        let (excerpt_offset, _) = matches.next().with_context(|| {
+            format!("Cursor excerpt did not exist in buffer:\n\n{cursor_excerpt}\n",)
+        })?;
+        anyhow::ensure!(
+            matches.next().is_none(),
+            "More than one cursor position match found",
+        );
+        Ok(excerpt_offset)
+    })?;
+
+    let cursor_offset = excerpt_offset + cursor_offset_within_excerpt;
+    let cursor_anchor =
+        cursor_buffer.read_with(&*cx, |buffer, _| buffer.anchor_after(cursor_offset));
+
+    Ok((cursor_buffer, cursor_anchor))
+}
+
+async fn setup_project(
+    example: &mut Example,
+    app_state: &Arc<EpAppState>,
+    step_progress: &StepProgress,
+    cx: &mut AsyncApp,
+) -> Result<Entity<Project>> {
+    let ep_store = cx
+        .update(|cx| EditPredictionStore::try_global(cx))
+        .context("Store should be initialized at init")?;
+
+    let worktree_path = setup_worktree(example, step_progress).await?;
+
+    let project = cx.update(|cx| {
+        Project::local(
+            app_state.client.clone(),
+            app_state.node_runtime.clone(),
+            app_state.user_store.clone(),
+            app_state.languages.clone(),
+            app_state.fs.clone(),
+            None,
+            project::LocalProjectFlags {
+                init_worktree_trust: false,
+                watch_global_configs: false,
+            },
+            cx,
+        )
+    });
+
+    project
+        .update(cx, |project, cx| {
+            project.disable_worktree_scanner(cx);
+            project.create_worktree(&worktree_path, true, cx)
+        })
+        .await?;
+
+    let buffer_store = project.read_with(cx, |project, _| project.buffer_store().clone());
+    cx.subscribe(&buffer_store, {
+        let project = project.downgrade();
+        let ep_store = ep_store.downgrade();
+        move |_, event, cx| match event {
+            BufferStoreEvent::BufferAdded(buffer) => {
+                let Some(project) = project.upgrade() else {
+                    return;
+                };
+                ep_store
+                    .update(cx, |store, cx| store.register_buffer(&buffer, &project, cx))
+                    .ok();
+            }
+            _ => {}
+        }
+    })
+    .detach();
+
+    Ok(project)
+}
+
+async fn setup_worktree(example: &Example, step_progress: &StepProgress) -> Result<PathBuf> {
+    let repo_name = example.repo_name().context("failed to get repo name")?;
+    let repo_dir = git::repo_path_for_url(&example.spec.repository_url)?;
+    let worktree_path = repo_name.worktree_path();
+    let repo_lock = git::lock_repo(&repo_dir).await;
+
+    // Clean up any stale git lock files from previous crashed runs.
+    // Safe-ish since we have our own lock.
+    // WARNING: Can corrupt worktrees if multiple processes of the CLI are running.
+    let worktree_git_dir = repo_dir
+        .join(".git/worktrees")
+        .join(repo_name.name.as_ref());
+    for lock_file in &["index.lock", "HEAD.lock", "config.lock"] {
+        let worktree_lock_path = worktree_git_dir.join(lock_file);
+        let repo_lock_path = repo_dir.join(".git").join(lock_file);
+        if worktree_lock_path.exists() {
+            fs::remove_file(&worktree_lock_path).ok();
+        }
+        if repo_lock_path.exists() {
+            fs::remove_file(&repo_lock_path).ok();
+        }
+    }
+
+    let mut git_repo_exists = false;
+    if repo_dir.is_dir() {
+        if git::run_git(&repo_dir, &["remote", "get-url", "origin"])
+            .await
+            .map_or(false, |origin| origin.trim() == example.spec.repository_url)
+        {
+            git_repo_exists = true;
+        } else {
+            fs::remove_dir_all(&repo_dir).ok();
+        }
+    }
+
+    if !git_repo_exists {
+        step_progress.set_substatus(format!("cloning {}", repo_name.name));
+        fs::create_dir_all(&repo_dir)?;
+        git::run_git(&repo_dir, &["init"]).await?;
+        git::run_git(
+            &repo_dir,
+            &["remote", "add", "origin", &example.spec.repository_url],
+        )
+        .await?;
+    }
+
+    // Resolve the example to a revision, fetching it if needed.
+    step_progress.set_substatus("fetching");
+    let revision = git::fetch_if_needed(&repo_dir, &example.spec.revision).await?;
+
+    // Clean up any stale worktree registrations from previous crashed runs.
+    git::run_git(&repo_dir, &["worktree", "prune"]).await.ok();
+
+    // Create the worktree for this example if needed.
+    step_progress.set_substatus("preparing worktree");
+
+    // Check if worktree exists and is valid (not just a directory from a crashed run).
+    let worktree_valid = worktree_path.is_dir()
+        && git::run_git(&worktree_path, &["rev-parse", "--git-dir"])
+            .await
+            .is_ok();
+
+    if worktree_valid {
+        git::run_git(&worktree_path, &["clean", "--force", "-d"]).await?;
+        git::run_git(&worktree_path, &["reset", "--hard", "HEAD"]).await?;
+        git::run_git(&worktree_path, &["checkout", revision.as_str()]).await?;
+    } else {
+        let worktree_path_string = worktree_path.to_string_lossy();
+
+        // Clean up invalid worktree directory and registration if they exist.
+        if worktree_path.exists() {
+            fs::remove_dir_all(&worktree_path).ok();
+        }
+        git::run_git(
+            &repo_dir,
+            &["worktree", "remove", "--force", &worktree_path_string],
+        )
+        .await
+        .ok();
+
+        let branch_name = example.spec.filename();
+        git::run_git(
+            &repo_dir,
+            &["branch", "-f", &branch_name, revision.as_str()],
+        )
+        .await?;
+        git::run_git(
+            &repo_dir,
+            &["worktree", "add", "-f", &worktree_path_string, &branch_name],
+        )
+        .await?;
+    }
+    drop(repo_lock);
+
+    if !example.spec.uncommitted_diff.is_empty() {
+        step_progress.set_substatus("applying diff");
+
+        // old examples had full paths in the uncommitted diff.
+        let uncommitted_diff =
+            strip_diff_path_prefix(&example.spec.uncommitted_diff, &repo_name.name);
+
+        let mut apply_process = smol::process::Command::new("git")
+            .current_dir(&worktree_path)
+            .args(&["apply", "-"])
+            .stdin(std::process::Stdio::piped())
+            .spawn()?;
+
+        let mut stdin = apply_process.stdin.take().context("Failed to get stdin")?;
+        stdin.write_all(uncommitted_diff.as_bytes()).await?;
+        stdin.close().await?;
+        drop(stdin);
+
+        let apply_result = apply_process.output().await?;
+        anyhow::ensure!(
+            apply_result.status.success(),
+            "Failed to apply uncommitted diff patch with status: {}\nstderr:\n{}\nstdout:\n{}",
+            apply_result.status,
+            String::from_utf8_lossy(&apply_result.stderr),
+            String::from_utf8_lossy(&apply_result.stdout),
+        );
+    }
+
+    step_progress.clear_substatus();
+    Ok(worktree_path)
+}
+
+async fn apply_edit_history(
+    example: &Example,
+    project: &Entity<Project>,
+    cx: &mut AsyncApp,
+) -> Result<OpenedBuffers> {
+    edit_prediction::udiff::apply_diff(&example.spec.edit_history, project, cx).await
+}

crates/edit_prediction_cli/src/metrics.rs

@@ -0,0 +1,16 @@
+#![allow(unused_imports)]
+
+pub use edit_prediction_metrics::ClassificationMetrics;
+pub use edit_prediction_metrics::Counts;
+pub use edit_prediction_metrics::DeltaChrFMetrics;
+pub use edit_prediction_metrics::KeptRateResult;
+pub use edit_prediction_metrics::TokenChangeCounts;
+pub use edit_prediction_metrics::braces_disbalance;
+pub use edit_prediction_metrics::compute_kept_rate;
+pub use edit_prediction_metrics::count_patch_token_changes;
+pub use edit_prediction_metrics::delta_chr_f;
+pub use edit_prediction_metrics::delta_chr_f_beta;
+pub use edit_prediction_metrics::exact_lines_match;
+pub use edit_prediction_metrics::extract_changed_lines_from_diff;
+pub use edit_prediction_metrics::has_isolated_whitespace_changes;
+pub use edit_prediction_metrics::is_editable_region_correct;

crates/edit_prediction_cli/src/openai_client.rs

@@ -0,0 +1,701 @@
+use anyhow::Result;
+use http_client::HttpClient;
+use indoc::indoc;
+use open_ai::{
+    MessageContent, OPEN_AI_API_URL, Request as OpenAiRequest, RequestMessage,
+    Response as OpenAiResponse, batches, non_streaming_completion,
+};
+use reqwest_client::ReqwestClient;
+use sqlez::bindable::Bind;
+use sqlez::bindable::StaticColumnCount;
+use sqlez_macros::sql;
+use std::hash::Hash;
+use std::hash::Hasher;
+use std::path::Path;
+use std::sync::{Arc, Mutex};
+
+pub struct PlainOpenAiClient {
+    pub http_client: Arc<dyn HttpClient>,
+    pub api_key: String,
+}
+
+impl PlainOpenAiClient {
+    pub fn new() -> Result<Self> {
+        let http_client: Arc<dyn http_client::HttpClient> = Arc::new(ReqwestClient::new());
+        let api_key = std::env::var("OPENAI_API_KEY")
+            .map_err(|_| anyhow::anyhow!("OPENAI_API_KEY environment variable not set"))?;
+        Ok(Self {
+            http_client,
+            api_key,
+        })
+    }
+
+    pub async fn generate(
+        &self,
+        model: &str,
+        max_tokens: u64,
+        messages: Vec<RequestMessage>,
+    ) -> Result<OpenAiResponse> {
+        let request = OpenAiRequest {
+            model: model.to_string(),
+            messages,
+            stream: false,
+            stream_options: None,
+            max_completion_tokens: Some(max_tokens),
+            stop: Vec::new(),
+            temperature: None,
+            tool_choice: None,
+            parallel_tool_calls: None,
+            tools: Vec::new(),
+            prompt_cache_key: None,
+            reasoning_effort: None,
+        };
+
+        let response = non_streaming_completion(
+            self.http_client.as_ref(),
+            OPEN_AI_API_URL,
+            &self.api_key,
+            request,
+        )
+        .await
+        .map_err(|e| anyhow::anyhow!("{:?}", e))?;
+
+        Ok(response)
+    }
+}
+
+pub struct BatchingOpenAiClient {
+    connection: Mutex<sqlez::connection::Connection>,
+    http_client: Arc<dyn HttpClient>,
+    api_key: String,
+}
+
+struct CacheRow {
+    request_hash: String,
+    request: Option<String>,
+    response: Option<String>,
+    batch_id: Option<String>,
+}
+
+impl StaticColumnCount for CacheRow {
+    fn column_count() -> usize {
+        4
+    }
+}
+
+impl Bind for CacheRow {
+    fn bind(&self, statement: &sqlez::statement::Statement, start_index: i32) -> Result<i32> {
+        let next_index = statement.bind(&self.request_hash, start_index)?;
+        let next_index = statement.bind(&self.request, next_index)?;
+        let next_index = statement.bind(&self.response, next_index)?;
+        let next_index = statement.bind(&self.batch_id, next_index)?;
+        Ok(next_index)
+    }
+}
+
+#[derive(serde::Serialize, serde::Deserialize)]
+struct SerializableRequest {
+    model: String,
+    max_tokens: u64,
+    messages: Vec<SerializableMessage>,
+}
+
+#[derive(serde::Serialize, serde::Deserialize)]
+struct SerializableMessage {
+    role: String,
+    content: String,
+}
+
+impl BatchingOpenAiClient {
+    fn new(cache_path: &Path) -> Result<Self> {
+        let http_client: Arc<dyn http_client::HttpClient> = Arc::new(ReqwestClient::new());
+        let api_key = std::env::var("OPENAI_API_KEY")
+            .map_err(|_| anyhow::anyhow!("OPENAI_API_KEY environment variable not set"))?;
+
+        let connection = sqlez::connection::Connection::open_file(cache_path.to_str().unwrap());
+        let mut statement = sqlez::statement::Statement::prepare(
+            &connection,
+            indoc! {"
+                CREATE TABLE IF NOT EXISTS openai_cache (
+                    request_hash TEXT PRIMARY KEY,
+                    request TEXT,
+                    response TEXT,
+                    batch_id TEXT
+                );
+                "},
+        )?;
+        statement.exec()?;
+        drop(statement);
+
+        Ok(Self {
+            connection: Mutex::new(connection),
+            http_client,
+            api_key,
+        })
+    }
+
+    pub fn lookup(
+        &self,
+        model: &str,
+        max_tokens: u64,
+        messages: &[RequestMessage],
+        seed: Option<usize>,
+    ) -> Result<Option<OpenAiResponse>> {
+        let request_hash_str = Self::request_hash(model, max_tokens, messages, seed);
+        let connection = self.connection.lock().unwrap();
+        let response: Vec<String> = connection.select_bound(
+            &sql!(SELECT response FROM openai_cache WHERE request_hash = ?1 AND response IS NOT NULL;),
+        )?(request_hash_str.as_str())?;
+        Ok(response
+            .into_iter()
+            .next()
+            .and_then(|text| serde_json::from_str(&text).ok()))
+    }
+
+    pub fn mark_for_batch(
+        &self,
+        model: &str,
+        max_tokens: u64,
+        messages: &[RequestMessage],
+        seed: Option<usize>,
+    ) -> Result<()> {
+        let request_hash = Self::request_hash(model, max_tokens, messages, seed);
+
+        let serializable_messages: Vec<SerializableMessage> = messages
+            .iter()
+            .map(|msg| SerializableMessage {
+                role: message_role_to_string(msg),
+                content: message_content_to_string(msg),
+            })
+            .collect();
+
+        let serializable_request = SerializableRequest {
+            model: model.to_string(),
+            max_tokens,
+            messages: serializable_messages,
+        };
+
+        let request = Some(serde_json::to_string(&serializable_request)?);
+        let cache_row = CacheRow {
+            request_hash,
+            request,
+            response: None,
+            batch_id: None,
+        };
+        let connection = self.connection.lock().unwrap();
+        connection.exec_bound::<CacheRow>(sql!(
+            INSERT OR IGNORE INTO openai_cache(request_hash, request, response, batch_id) VALUES (?, ?, ?, ?)))?(
+            cache_row,
+        )
+    }
+
+    async fn generate(
+        &self,
+        model: &str,
+        max_tokens: u64,
+        messages: Vec<RequestMessage>,
+        seed: Option<usize>,
+        cache_only: bool,
+    ) -> Result<Option<OpenAiResponse>> {
+        let response = self.lookup(model, max_tokens, &messages, seed)?;
+        if let Some(response) = response {
+            return Ok(Some(response));
+        }
+
+        if !cache_only {
+            self.mark_for_batch(model, max_tokens, &messages, seed)?;
+        }
+
+        Ok(None)
+    }
+
+    async fn sync_batches(&self) -> Result<()> {
+        let _batch_ids = self.upload_pending_requests().await?;
+        self.download_finished_batches().await
+    }
+
+    pub fn pending_batch_count(&self) -> Result<usize> {
+        let connection = self.connection.lock().unwrap();
+        let counts: Vec<i32> = connection.select(
+            sql!(SELECT COUNT(*) FROM openai_cache WHERE batch_id IS NOT NULL AND response IS NULL),
+        )?()?;
+        Ok(counts.into_iter().next().unwrap_or(0) as usize)
+    }
+
+    pub async fn import_batches(&self, batch_ids: &[String]) -> Result<()> {
+        for batch_id in batch_ids {
+            log::info!("Importing OpenAI batch {}", batch_id);
+
+            let batch_status = batches::retrieve_batch(
+                self.http_client.as_ref(),
+                OPEN_AI_API_URL,
+                &self.api_key,
+                batch_id,
+            )
+            .await
+            .map_err(|e| anyhow::anyhow!("Failed to retrieve batch {}: {:?}", batch_id, e))?;
+
+            log::info!("Batch {} status: {}", batch_id, batch_status.status);
+
+            if batch_status.status != "completed" {
+                log::warn!(
+                    "Batch {} is not completed (status: {}), skipping",
+                    batch_id,
+                    batch_status.status
+                );
+                continue;
+            }
+
+            let output_file_id = batch_status.output_file_id.ok_or_else(|| {
+                anyhow::anyhow!("Batch {} completed but has no output file", batch_id)
+            })?;
+
+            let results_content = batches::download_file(
+                self.http_client.as_ref(),
+                OPEN_AI_API_URL,
+                &self.api_key,
+                &output_file_id,
+            )
+            .await
+            .map_err(|e| {
+                anyhow::anyhow!("Failed to download batch results for {}: {:?}", batch_id, e)
+            })?;
+
+            let results = batches::parse_batch_output(&results_content)
+                .map_err(|e| anyhow::anyhow!("Failed to parse batch output: {:?}", e))?;
+
+            let mut updates: Vec<(String, String, String)> = Vec::new();
+            let mut success_count = 0;
+            let mut error_count = 0;
+
+            for result in results {
+                let request_hash = result
+                    .custom_id
+                    .strip_prefix("req_hash_")
+                    .unwrap_or(&result.custom_id)
+                    .to_string();
+
+                if let Some(response_body) = result.response {
+                    if response_body.status_code == 200 {
+                        let response_json = serde_json::to_string(&response_body.body)?;
+                        updates.push((request_hash, response_json, batch_id.clone()));
+                        success_count += 1;
+                    } else {
+                        log::error!(
+                            "Batch request {} failed with status {}",
+                            request_hash,
+                            response_body.status_code
+                        );
+                        let error_json = serde_json::json!({
+                            "error": {
+                                "type": "http_error",
+                                "status_code": response_body.status_code
+                            }
+                        })
+                        .to_string();
+                        updates.push((request_hash, error_json, batch_id.clone()));
+                        error_count += 1;
+                    }
+                } else if let Some(error) = result.error {
+                    log::error!(
+                        "Batch request {} failed: {}: {}",
+                        request_hash,
+                        error.code,
+                        error.message
+                    );
+                    let error_json = serde_json::json!({
+                        "error": {
+                            "type": error.code,
+                            "message": error.message
+                        }
+                    })
+                    .to_string();
+                    updates.push((request_hash, error_json, batch_id.clone()));
+                    error_count += 1;
+                }
+            }
+
+            let connection = self.connection.lock().unwrap();
+            connection.with_savepoint("batch_import", || {
+                let q = sql!(
+                    INSERT OR REPLACE INTO openai_cache(request_hash, request, response, batch_id)
+                    VALUES (?, (SELECT request FROM openai_cache WHERE request_hash = ?), ?, ?)
+                );
+                let mut exec = connection.exec_bound::<(&str, &str, &str, &str)>(q)?;
+                for (request_hash, response_json, batch_id) in &updates {
+                    exec((
+                        request_hash.as_str(),
+                        request_hash.as_str(),
+                        response_json.as_str(),
+                        batch_id.as_str(),
+                    ))?;
+                }
+                Ok(())
+            })?;
+
+            log::info!(
+                "Imported batch {}: {} successful, {} errors",
+                batch_id,
+                success_count,
+                error_count
+            );
+        }
+
+        Ok(())
+    }
+
+    async fn download_finished_batches(&self) -> Result<()> {
+        let batch_ids: Vec<String> = {
+            let connection = self.connection.lock().unwrap();
+            let q = sql!(SELECT DISTINCT batch_id FROM openai_cache WHERE batch_id IS NOT NULL AND response IS NULL);
+            connection.select(q)?()?
+        };
+
+        for batch_id in &batch_ids {
+            let batch_status = batches::retrieve_batch(
+                self.http_client.as_ref(),
+                OPEN_AI_API_URL,
+                &self.api_key,
+                batch_id,
+            )
+            .await
+            .map_err(|e| anyhow::anyhow!("{:?}", e))?;
+
+            log::info!("Batch {} status: {}", batch_id, batch_status.status);
+
+            if batch_status.status == "completed" {
+                let output_file_id = match batch_status.output_file_id {
+                    Some(id) => id,
+                    None => {
+                        log::warn!("Batch {} completed but has no output file", batch_id);
+                        continue;
+                    }
+                };
+
+                let results_content = batches::download_file(
+                    self.http_client.as_ref(),
+                    OPEN_AI_API_URL,
+                    &self.api_key,
+                    &output_file_id,
+                )
+                .await
+                .map_err(|e| anyhow::anyhow!("{:?}", e))?;
+
+                let results = batches::parse_batch_output(&results_content)
+                    .map_err(|e| anyhow::anyhow!("Failed to parse batch output: {:?}", e))?;
+
+                let mut updates: Vec<(String, String)> = Vec::new();
+                let mut success_count = 0;
+
+                for result in results {
+                    let request_hash = result
+                        .custom_id
+                        .strip_prefix("req_hash_")
+                        .unwrap_or(&result.custom_id)
+                        .to_string();
+
+                    if let Some(response_body) = result.response {
+                        if response_body.status_code == 200 {
+                            let response_json = serde_json::to_string(&response_body.body)?;
+                            updates.push((response_json, request_hash));
+                            success_count += 1;
+                        } else {
+                            log::error!(
+                                "Batch request {} failed with status {}",
+                                request_hash,
+                                response_body.status_code
+                            );
+                            let error_json = serde_json::json!({
+                                "error": {
+                                    "type": "http_error",
+                                    "status_code": response_body.status_code
+                                }
+                            })
+                            .to_string();
+                            updates.push((error_json, request_hash));
+                        }
+                    } else if let Some(error) = result.error {
+                        log::error!(
+                            "Batch request {} failed: {}: {}",
+                            request_hash,
+                            error.code,
+                            error.message
+                        );
+                        let error_json = serde_json::json!({
+                            "error": {
+                                "type": error.code,
+                                "message": error.message
+                            }
+                        })
+                        .to_string();
+                        updates.push((error_json, request_hash));
+                    }
+                }
+
+                let connection = self.connection.lock().unwrap();
+                connection.with_savepoint("batch_download", || {
+                    let q = sql!(UPDATE openai_cache SET response = ? WHERE request_hash = ?);
+                    let mut exec = connection.exec_bound::<(&str, &str)>(q)?;
+                    for (response_json, request_hash) in &updates {
+                        exec((response_json.as_str(), request_hash.as_str()))?;
+                    }
+                    Ok(())
+                })?;
+                log::info!("Downloaded {} successful requests", success_count);
+            }
+        }
+
+        Ok(())
+    }
+
+    async fn upload_pending_requests(&self) -> Result<Vec<String>> {
+        const BATCH_CHUNK_SIZE: i32 = 16_000;
+        let mut all_batch_ids = Vec::new();
+        let mut total_uploaded = 0;
+
+        loop {
+            let rows: Vec<(String, String)> = {
+                let connection = self.connection.lock().unwrap();
+                let q = sql!(
+                    SELECT request_hash, request FROM openai_cache
+                    WHERE batch_id IS NULL AND response IS NULL
+                    LIMIT ?
+                );
+                connection.select_bound(q)?(BATCH_CHUNK_SIZE)?
+            };
+
+            if rows.is_empty() {
+                break;
+            }
+
+            let request_hashes: Vec<String> = rows.iter().map(|(hash, _)| hash.clone()).collect();
+
+            let mut jsonl_content = String::new();
+            for (hash, request_str) in &rows {
+                let serializable_request: SerializableRequest =
+                    serde_json::from_str(request_str).unwrap();
+
+                let messages: Vec<RequestMessage> = serializable_request
+                    .messages
+                    .into_iter()
+                    .map(|msg| match msg.role.as_str() {
+                        "user" => RequestMessage::User {
+                            content: MessageContent::Plain(msg.content),
+                        },
+                        "assistant" => RequestMessage::Assistant {
+                            content: Some(MessageContent::Plain(msg.content)),
+                            tool_calls: Vec::new(),
+                            reasoning_content: None,
+                        },
+                        "system" => RequestMessage::System {
+                            content: MessageContent::Plain(msg.content),
+                        },
+                        _ => RequestMessage::User {
+                            content: MessageContent::Plain(msg.content),
+                        },
+                    })
+                    .collect();
+
+                let request = OpenAiRequest {
+                    model: serializable_request.model,
+                    messages,
+                    stream: false,
+                    stream_options: None,
+                    max_completion_tokens: Some(serializable_request.max_tokens),
+                    stop: Vec::new(),
+                    temperature: None,
+                    tool_choice: None,
+                    parallel_tool_calls: None,
+                    tools: Vec::new(),
+                    prompt_cache_key: None,
+                    reasoning_effort: None,
+                };
+
+                let custom_id = format!("req_hash_{}", hash);
+                let batch_item = batches::BatchRequestItem::new(custom_id, request);
+                let line = batch_item
+                    .to_jsonl_line()
+                    .map_err(|e| anyhow::anyhow!("Failed to serialize batch item: {:?}", e))?;
+                jsonl_content.push_str(&line);
+                jsonl_content.push('\n');
+            }
+
+            let filename = format!("batch_{}.jsonl", chrono::Utc::now().timestamp());
+            let file_obj = batches::upload_batch_file(
+                self.http_client.as_ref(),
+                OPEN_AI_API_URL,
+                &self.api_key,
+                &filename,
+                jsonl_content.into_bytes(),
+            )
+            .await
+            .map_err(|e| anyhow::anyhow!("Failed to upload batch file: {:?}", e))?;
+
+            let batch = batches::create_batch(
+                self.http_client.as_ref(),
+                OPEN_AI_API_URL,
+                &self.api_key,
+                batches::CreateBatchRequest::new(file_obj.id),
+            )
+            .await
+            .map_err(|e| anyhow::anyhow!("Failed to create batch: {:?}", e))?;
+
+            {
+                let connection = self.connection.lock().unwrap();
+                connection.with_savepoint("batch_upload", || {
+                    let q = sql!(UPDATE openai_cache SET batch_id = ? WHERE request_hash = ?);
+                    let mut exec = connection.exec_bound::<(&str, &str)>(q)?;
+                    for hash in &request_hashes {
+                        exec((batch.id.as_str(), hash.as_str()))?;
+                    }
+                    Ok(())
+                })?;
+            }
+
+            let batch_len = rows.len();
+            total_uploaded += batch_len;
+            log::info!(
+                "Uploaded batch {} with {} requests ({} total)",
+                batch.id,
+                batch_len,
+                total_uploaded
+            );
+
+            all_batch_ids.push(batch.id);
+        }
+
+        if !all_batch_ids.is_empty() {
+            log::info!(
+                "Finished uploading {} batches with {} total requests",
+                all_batch_ids.len(),
+                total_uploaded
+            );
+        }
+
+        Ok(all_batch_ids)
+    }
+
+    fn request_hash(
+        model: &str,
+        max_tokens: u64,
+        messages: &[RequestMessage],
+        seed: Option<usize>,
+    ) -> String {
+        let mut hasher = std::hash::DefaultHasher::new();
+        "openai".hash(&mut hasher);
+        model.hash(&mut hasher);
+        max_tokens.hash(&mut hasher);
+        for msg in messages {
+            message_content_to_string(msg).hash(&mut hasher);
+        }
+        if let Some(seed) = seed {
+            seed.hash(&mut hasher);
+        }
+        let request_hash = hasher.finish();
+        format!("{request_hash:016x}")
+    }
+}
+
+fn message_role_to_string(msg: &RequestMessage) -> String {
+    match msg {
+        RequestMessage::User { .. } => "user".to_string(),
+        RequestMessage::Assistant { .. } => "assistant".to_string(),
+        RequestMessage::System { .. } => "system".to_string(),
+        RequestMessage::Tool { .. } => "tool".to_string(),
+    }
+}
+
+fn message_content_to_string(msg: &RequestMessage) -> String {
+    match msg {
+        RequestMessage::User { content } => content_to_string(content),
+        RequestMessage::Assistant { content, .. } => {
+            content.as_ref().map(content_to_string).unwrap_or_default()
+        }
+        RequestMessage::System { content } => content_to_string(content),
+        RequestMessage::Tool { content, .. } => content_to_string(content),
+    }
+}
+
+fn content_to_string(content: &MessageContent) -> String {
+    match content {
+        MessageContent::Plain(text) => text.clone(),
+        MessageContent::Multipart(parts) => parts
+            .iter()
+            .filter_map(|part| match part {
+                open_ai::MessagePart::Text { text } => Some(text.clone()),
+                _ => None,
+            })
+            .collect::<Vec<String>>()
+            .join("\n"),
+    }
+}
+
+pub enum OpenAiClient {
+    Plain(PlainOpenAiClient),
+    Batch(BatchingOpenAiClient),
+    #[allow(dead_code)]
+    Dummy,
+}
+
+impl OpenAiClient {
+    pub fn plain() -> Result<Self> {
+        Ok(Self::Plain(PlainOpenAiClient::new()?))
+    }
+
+    pub fn batch(cache_path: &Path) -> Result<Self> {
+        Ok(Self::Batch(BatchingOpenAiClient::new(cache_path)?))
+    }
+
+    #[allow(dead_code)]
+    pub fn dummy() -> Self {
+        Self::Dummy
+    }
+
+    pub async fn generate(
+        &self,
+        model: &str,
+        max_tokens: u64,
+        messages: Vec<RequestMessage>,
+        seed: Option<usize>,
+        cache_only: bool,
+    ) -> Result<Option<OpenAiResponse>> {
+        match self {
+            OpenAiClient::Plain(plain_client) => plain_client
+                .generate(model, max_tokens, messages)
+                .await
+                .map(Some),
+            OpenAiClient::Batch(batching_client) => {
+                batching_client
+                    .generate(model, max_tokens, messages, seed, cache_only)
+                    .await
+            }
+            OpenAiClient::Dummy => panic!("Dummy OpenAI client is not expected to be used"),
+        }
+    }
+
+    pub async fn sync_batches(&self) -> Result<()> {
+        match self {
+            OpenAiClient::Plain(_) => Ok(()),
+            OpenAiClient::Batch(batching_client) => batching_client.sync_batches().await,
+            OpenAiClient::Dummy => panic!("Dummy OpenAI client is not expected to be used"),
+        }
+    }
+
+    pub fn pending_batch_count(&self) -> Result<usize> {
+        match self {
+            OpenAiClient::Plain(_) => Ok(0),
+            OpenAiClient::Batch(batching_client) => batching_client.pending_batch_count(),
+            OpenAiClient::Dummy => panic!("Dummy OpenAI client is not expected to be used"),
+        }
+    }
+
+    pub async fn import_batches(&self, batch_ids: &[String]) -> Result<()> {
+        match self {
+            OpenAiClient::Plain(_) => {
+                anyhow::bail!("Import batches is only supported with batching client")
+            }
+            OpenAiClient::Batch(batching_client) => batching_client.import_batches(batch_ids).await,
+            OpenAiClient::Dummy => panic!("Dummy OpenAI client is not expected to be used"),
+        }
+    }
+}

crates/edit_prediction_cli/src/parse_output.rs

@@ -0,0 +1,92 @@
+use crate::{
+    PredictionProvider,
+    example::{ActualCursor, Example},
+    format_prompt::{TeacherMultiRegionPrompt, TeacherPrompt},
+    repair,
+};
+use anyhow::{Context as _, Result};
+use zeta_prompt::{ZetaFormat, parse_zeta2_model_output, parsed_output_to_patch};
+
+pub fn run_parse_output(example: &mut Example) -> Result<()> {
+    example
+        .prompt_inputs
+        .as_ref()
+        .context("prompt_inputs required")?;
+
+    let to_parse: Vec<_> = example
+        .predictions
+        .iter()
+        .enumerate()
+        .filter(|(_, p)| !p.actual_output.is_empty())
+        .map(|(ix, p)| (ix, p.actual_output.clone(), p.provider))
+        .collect();
+
+    for (ix, actual_output, provider) in to_parse {
+        let (actual_patch, actual_cursor) =
+            parse_prediction_output(example, &actual_output, provider)?;
+        example.predictions[ix].actual_patch = Some(actual_patch);
+        example.predictions[ix].actual_cursor = actual_cursor;
+    }
+
+    Ok(())
+}
+
+pub fn parse_prediction_output(
+    example: &Example,
+    actual_output: &str,
+    provider: PredictionProvider,
+) -> Result<(String, Option<ActualCursor>)> {
+    match provider {
+        PredictionProvider::Teacher(_, _) | PredictionProvider::TeacherNonBatching(_, _) => {
+            TeacherPrompt::parse(example, actual_output)
+        }
+        PredictionProvider::TeacherMultiRegion(_)
+        | PredictionProvider::TeacherMultiRegionNonBatching(_) => {
+            TeacherMultiRegionPrompt::parse(example, actual_output)
+        }
+        PredictionProvider::Zeta2(version) => parse_zeta2_output(example, actual_output, version),
+        PredictionProvider::Repair => repair::parse(example, actual_output),
+        _ => anyhow::bail!(
+            "parse-output only supports Teacher and Zeta2 providers, got {:?}",
+            provider
+        ),
+    }
+}
+
+fn parse_zeta2_output(
+    example: &Example,
+    actual_output: &str,
+    format: ZetaFormat,
+) -> Result<(String, Option<ActualCursor>)> {
+    let prompt_inputs = example
+        .prompt_inputs
+        .as_ref()
+        .context("prompt_inputs required")?;
+
+    let parsed = parse_zeta2_model_output(actual_output, format, prompt_inputs)?;
+    let range_in_excerpt = parsed.range_in_excerpt.clone();
+    let excerpt = prompt_inputs.cursor_excerpt.as_ref();
+    let editable_region_offset = range_in_excerpt.start;
+    let editable_region_start_line = excerpt[..editable_region_offset].matches('\n').count();
+
+    let mut new_text = parsed.new_editable_region.clone();
+    if !new_text.is_empty() && !new_text.ends_with('\n') {
+        new_text.push('\n');
+    }
+
+    let cursor_offset = parsed.cursor_offset_in_new_editable_region;
+    let formatted_diff = parsed_output_to_patch(prompt_inputs, parsed)?;
+
+    let actual_cursor = cursor_offset.map(|editable_region_cursor_offset| {
+        ActualCursor::from_editable_region(
+            &example.spec.cursor_path,
+            editable_region_cursor_offset,
+            &new_text,
+            excerpt,
+            editable_region_offset,
+            editable_region_start_line,
+        )
+    });
+
+    Ok((formatted_diff, actual_cursor))
+}

crates/edit_prediction_cli/src/paths.rs

@@ -0,0 +1,31 @@
+use std::{
+    path::{Path, PathBuf},
+    sync::LazyLock,
+};
+
+pub static DATA_DIR: LazyLock<PathBuf> = LazyLock::new(|| {
+    let dir = dirs::home_dir().unwrap().join(".zed_ep");
+    ensure_dir(&dir)
+});
+pub static CACHE_DIR: LazyLock<PathBuf> = LazyLock::new(|| ensure_dir(&DATA_DIR.join("cache")));
+pub static REPOS_DIR: LazyLock<PathBuf> = LazyLock::new(|| ensure_dir(&DATA_DIR.join("repos")));
+pub static WORKTREES_DIR: LazyLock<PathBuf> =
+    LazyLock::new(|| ensure_dir(&DATA_DIR.join("worktrees")));
+pub static RUN_DIR: LazyLock<PathBuf> = LazyLock::new(|| {
+    DATA_DIR
+        .join("runs")
+        .join(chrono::Local::now().format("%d-%m-%y-%H_%M_%S").to_string())
+});
+pub static LATEST_EXAMPLE_RUN_DIR: LazyLock<PathBuf> = LazyLock::new(|| DATA_DIR.join("latest"));
+pub static LATEST_FAILED_EXAMPLES_DIR: LazyLock<PathBuf> =
+    LazyLock::new(|| DATA_DIR.join("latest_failed"));
+pub static LLM_CACHE_DB: LazyLock<PathBuf> = LazyLock::new(|| CACHE_DIR.join("llm_cache.sqlite"));
+pub static SYNTHESIZE_STATE_FILE: LazyLock<PathBuf> =
+    LazyLock::new(|| DATA_DIR.join("synthesize_state.json"));
+pub static FAILED_EXAMPLES_DIR: LazyLock<PathBuf> =
+    LazyLock::new(|| ensure_dir(&RUN_DIR.join("failed")));
+
+fn ensure_dir(path: &Path) -> PathBuf {
+    std::fs::create_dir_all(path).expect("Failed to create directory");
+    path.to_path_buf()
+}

crates/edit_prediction_cli/src/predict.rs

@@ -0,0 +1,789 @@
+use crate::{
+    FormatPromptArgs, PredictArgs, PredictionProvider, TeacherBackend,
+    anthropic_client::AnthropicClient,
+    example::{Example, ExamplePrediction, ExamplePrompt},
+    format_prompt::{TeacherMultiRegionPrompt, TeacherPrompt, run_format_prompt},
+    headless::EpAppState,
+    load_project::run_load_project,
+    openai_client::OpenAiClient,
+    parse_output::parse_prediction_output,
+    paths::{LATEST_EXAMPLE_RUN_DIR, RUN_DIR},
+    progress::{ExampleProgress, InfoStyle, Progress, Step, StepProgress},
+    retrieve_context::run_context_retrieval,
+};
+use anyhow::Context as _;
+use cloud_llm_client::predict_edits_v3::{RawCompletionRequest, RawCompletionResponse};
+use edit_prediction::{DebugEvent, EditPredictionStore, Zeta2RawConfig};
+use futures::{AsyncReadExt as _, FutureExt as _, StreamExt as _, future::Shared};
+use gpui::{AppContext as _, AsyncApp, Task};
+use http_client::{AsyncBody, HttpClient, Method};
+use reqwest_client::ReqwestClient;
+use std::{
+    fs,
+    sync::{
+        Arc, Mutex, OnceLock,
+        atomic::{AtomicUsize, Ordering::SeqCst},
+    },
+};
+use zeta_prompt::ZetaFormat;
+
+static ANTHROPIC_CLIENT: OnceLock<AnthropicClient> = OnceLock::new();
+static OPENAI_CLIENT: OnceLock<OpenAiClient> = OnceLock::new();
+
+pub async fn run_prediction(
+    example: &mut Example,
+    args: &PredictArgs,
+    app_state: Arc<EpAppState>,
+    example_progress: &ExampleProgress,
+    mut cx: AsyncApp,
+) -> anyhow::Result<()> {
+    let repetition_count = args.repetitions;
+
+    if let Some(existing_prediction) = example.predictions.first() {
+        let has_prediction = existing_prediction.actual_patch.is_some()
+            || !existing_prediction.actual_output.is_empty();
+        if has_prediction {
+            match args.provider {
+                None => return Ok(()),
+                Some(provider) if existing_prediction.provider == provider => return Ok(()),
+                Some(_) => example.predictions.clear(),
+            }
+        }
+    }
+
+    let Some(provider) = args.provider else {
+        anyhow::bail!(
+            "No existing predictions found. Use --provider to specify which model to use for prediction."
+        );
+    };
+
+    if matches!(
+        provider,
+        PredictionProvider::TeacherMultiRegion(..)
+            | PredictionProvider::TeacherMultiRegionNonBatching(..)
+    ) {
+        anyhow::bail!("Teacher multi-region providers are not supported for prediction.");
+    }
+
+    if let PredictionProvider::Teacher(backend, _)
+    | PredictionProvider::TeacherNonBatching(backend, _) = provider
+    {
+        run_context_retrieval(example, app_state.clone(), example_progress, cx.clone()).await?;
+        run_format_prompt(
+            example,
+            &FormatPromptArgs { provider },
+            app_state.clone(),
+            example_progress,
+            cx,
+        )
+        .await?;
+
+        let step_progress = example_progress.start(Step::Predict);
+        let batched = matches!(
+            provider,
+            PredictionProvider::Teacher(..) | PredictionProvider::TeacherMultiRegion(..)
+        );
+        return predict_teacher(
+            example,
+            backend,
+            batched,
+            repetition_count,
+            args.cache_only,
+            &step_progress,
+        )
+        .await;
+    }
+
+    if let PredictionProvider::Baseten(format) = provider {
+        run_format_prompt(
+            example,
+            &FormatPromptArgs {
+                provider: PredictionProvider::Zeta2(format),
+            },
+            app_state.clone(),
+            example_progress,
+            cx,
+        )
+        .await?;
+
+        let step_progress = example_progress.start(Step::Predict);
+        return predict_baseten(example, format, &step_progress).await;
+    }
+
+    run_load_project(example, app_state.clone(), example_progress, cx.clone()).await?;
+    run_context_retrieval(example, app_state.clone(), example_progress, cx.clone()).await?;
+
+    let step_progress = example_progress.start(Step::Predict);
+
+    if matches!(
+        provider,
+        PredictionProvider::Zeta1 | PredictionProvider::Zeta2(_)
+    ) {
+        step_progress.set_substatus("authenticating");
+        static AUTHENTICATED: OnceLock<Shared<Task<()>>> = OnceLock::new();
+        AUTHENTICATED
+            .get_or_init(|| {
+                let client = app_state.client.clone();
+                cx.spawn(async move |cx| {
+                    if let Err(e) = client.sign_in_with_optional_connect(true, cx).await {
+                        eprintln!("Authentication failed: {}", e);
+                    }
+                })
+                .shared()
+            })
+            .clone()
+            .await;
+    }
+
+    let ep_store = cx
+        .update(|cx| EditPredictionStore::try_global(cx))
+        .context("EditPredictionStore not initialized")?;
+
+    ep_store.update(&mut cx, |store, _cx| {
+        let model = match provider {
+            PredictionProvider::Zeta1 => edit_prediction::EditPredictionModel::Zeta,
+            PredictionProvider::Zeta2(_) => edit_prediction::EditPredictionModel::Zeta,
+            PredictionProvider::Mercury => edit_prediction::EditPredictionModel::Mercury,
+            PredictionProvider::Teacher(..)
+            | PredictionProvider::TeacherMultiRegion(..)
+            | PredictionProvider::TeacherNonBatching(..)
+            | PredictionProvider::TeacherMultiRegionNonBatching(..)
+            | PredictionProvider::Repair
+            | PredictionProvider::Baseten(_) => {
+                unreachable!()
+            }
+        };
+        store.set_edit_prediction_model(model);
+
+        // If user specified a non-default Zeta2 version, configure raw endpoint.
+        // ZED_ZETA_MODEL env var is optional.
+        if let PredictionProvider::Zeta2(format) = provider {
+            if format != ZetaFormat::default() {
+                let model_id = std::env::var("ZED_ZETA_MODEL").ok();
+                let environment = std::env::var("ZED_ZETA_ENVIRONMENT").ok();
+                store.set_zeta2_raw_config(Zeta2RawConfig {
+                    model_id,
+                    environment,
+                    format,
+                });
+            }
+        }
+    });
+    step_progress.set_substatus("configuring model");
+    let state = example.state.as_ref().context("state must be set")?;
+    let run_dir = RUN_DIR.join(&example.spec.name);
+
+    let updated_example = Arc::new(Mutex::new(example.clone()));
+    let current_run_ix = Arc::new(AtomicUsize::new(0));
+
+    let mut debug_rx = ep_store.update(&mut cx, |store, cx| store.debug_info(&state.project, cx));
+    let debug_task = cx.background_spawn({
+        let updated_example = updated_example.clone();
+        let current_run_ix = current_run_ix.clone();
+        let run_dir = run_dir.clone();
+        async move {
+            while let Some(event) = debug_rx.next().await {
+                let run_ix = current_run_ix.load(SeqCst);
+                let mut updated_example = updated_example.lock().unwrap();
+
+                let run_dir = if repetition_count > 1 {
+                    run_dir.join(format!("{:03}", run_ix))
+                } else {
+                    run_dir.clone()
+                };
+
+                match event {
+                    DebugEvent::EditPredictionStarted(request) => {
+                        assert_eq!(updated_example.predictions.len(), run_ix + 1);
+
+                        if let Some(prompt) = request.prompt {
+                            fs::write(run_dir.join("prediction_prompt.md"), &prompt)?;
+                            if matches!(provider, PredictionProvider::Zeta2(_)) {
+                                updated_example.prompt.get_or_insert(ExamplePrompt {
+                                    input: prompt,
+                                    expected_output: None,
+                                    rejected_output: None,
+                                    provider,
+                                    prefill: None,
+                                });
+                            }
+                        }
+                    }
+                    DebugEvent::EditPredictionFinished(request) => {
+                        assert_eq!(updated_example.predictions.len(), run_ix + 1);
+
+                        if let Some(output) = request.model_output {
+                            fs::write(run_dir.join("prediction_response.md"), &output)?;
+                            updated_example
+                                .predictions
+                                .last_mut()
+                                .unwrap()
+                                .actual_output = output;
+                        }
+                        if run_ix >= repetition_count {
+                            break;
+                        }
+                    }
+                    _ => {}
+                }
+            }
+            anyhow::Ok(())
+        }
+    });
+
+    for ix in 0..repetition_count {
+        current_run_ix.store(ix, SeqCst);
+        let run_dir = if repetition_count > 1 {
+            run_dir.join(format!("{:03}", ix))
+        } else {
+            run_dir.clone()
+        };
+
+        if repetition_count > 1 {
+            step_progress.set_substatus(format!(
+                "running prediction {}/{}",
+                ix + 1,
+                repetition_count
+            ));
+        } else {
+            step_progress.set_substatus("running prediction");
+        }
+
+        fs::create_dir_all(&run_dir)?;
+        if LATEST_EXAMPLE_RUN_DIR.is_symlink() {
+            fs::remove_file(&*LATEST_EXAMPLE_RUN_DIR)?;
+        }
+        #[cfg(unix)]
+        std::os::unix::fs::symlink(&run_dir, &*LATEST_EXAMPLE_RUN_DIR)?;
+        #[cfg(windows)]
+        std::os::windows::fs::symlink_dir(&run_dir, &*LATEST_EXAMPLE_RUN_DIR)?;
+
+        updated_example
+            .lock()
+            .unwrap()
+            .predictions
+            .push(ExamplePrediction {
+                actual_patch: None,
+                actual_output: String::new(),
+                actual_cursor: None,
+                error: None,
+                provider,
+                cumulative_logprob: None,
+                avg_logprob: None,
+            });
+
+        step_progress.set_substatus("requesting prediction");
+        let prediction = ep_store
+            .update(&mut cx, |store, cx| {
+                store.request_prediction(
+                    &state.project,
+                    &state.buffer,
+                    state.cursor_position,
+                    cloud_llm_client::PredictEditsRequestTrigger::Cli,
+                    cx,
+                )
+            })
+            .await?;
+
+        let actual_patch = prediction.and_then(|prediction| {
+            let prediction = prediction.prediction.ok()?;
+            prediction
+                .edit_preview
+                .as_unified_diff(prediction.snapshot.file(), &prediction.edits)
+        });
+
+        let has_prediction = actual_patch.as_ref().is_some_and(|p| !p.is_empty());
+
+        updated_example
+            .lock()
+            .unwrap()
+            .predictions
+            .last_mut()
+            .unwrap()
+            .actual_patch = actual_patch;
+
+        if ix == repetition_count - 1 {
+            let (info, style) = if has_prediction {
+                ("predicted", InfoStyle::Normal)
+            } else {
+                ("no prediction", InfoStyle::Warning)
+            };
+            step_progress.set_info(info, style);
+        }
+    }
+
+    ep_store.update(&mut cx, |store, _| {
+        store.remove_project(&state.project);
+    });
+    debug_task.await?;
+
+    *example = Arc::into_inner(updated_example)
+        .ok_or_else(|| anyhow::anyhow!("Failed to unwrap Arc"))?
+        .into_inner()
+        .map_err(|_| anyhow::anyhow!("Failed to unwrap Mutex"))?;
+    Ok(())
+}
+
+async fn predict_teacher(
+    example: &mut Example,
+    backend: TeacherBackend,
+    batched: bool,
+    repetition_count: usize,
+    cache_only: bool,
+    step_progress: &crate::progress::StepProgress,
+) -> anyhow::Result<()> {
+    match backend {
+        TeacherBackend::Sonnet45 | TeacherBackend::Sonnet46 => {
+            predict_anthropic(
+                example,
+                backend,
+                batched,
+                repetition_count,
+                cache_only,
+                step_progress,
+            )
+            .await
+        }
+        TeacherBackend::Gpt52 => {
+            predict_openai(
+                example,
+                backend,
+                batched,
+                repetition_count,
+                cache_only,
+                step_progress,
+            )
+            .await
+        }
+    }
+}
+
+async fn predict_anthropic(
+    example: &mut Example,
+    backend: TeacherBackend,
+    batched: bool,
+    repetition_count: usize,
+    cache_only: bool,
+    step_progress: &crate::progress::StepProgress,
+) -> anyhow::Result<()> {
+    let llm_model_name = backend.model_name();
+    let max_tokens = 16384;
+    let llm_client = ANTHROPIC_CLIENT.get_or_init(|| {
+        let client = if batched {
+            AnthropicClient::batch(&crate::paths::LLM_CACHE_DB)
+        } else {
+            AnthropicClient::plain()
+        };
+        client.expect("Failed to create Anthropic client")
+    });
+
+    let prompt = example.prompt.as_ref().context("Prompt is required")?;
+
+    for ix in 0..repetition_count {
+        if repetition_count > 1 {
+            step_progress.set_substatus(format!(
+                "running prediction {}/{}",
+                ix + 1,
+                repetition_count
+            ));
+        } else {
+            step_progress.set_substatus("running prediction");
+        }
+
+        let messages = vec![anthropic::Message {
+            role: anthropic::Role::User,
+            content: vec![anthropic::RequestContent::Text {
+                text: prompt.input.clone(),
+                cache_control: None,
+            }],
+        }];
+
+        let seed = if repetition_count > 1 { Some(ix) } else { None };
+        let Some(response) = llm_client
+            .generate(llm_model_name, max_tokens, messages, seed, cache_only)
+            .await?
+        else {
+            // Request stashed for batched processing
+            continue;
+        };
+
+        let actual_output = response
+            .content
+            .into_iter()
+            .filter_map(|content| match content {
+                anthropic::ResponseContent::Text { text } => Some(text),
+                _ => None,
+            })
+            .collect::<Vec<String>>()
+            .join("\n");
+
+        let parser_provider = if batched {
+            example
+                .prompt
+                .as_ref()
+                .map(|prompt| prompt.provider)
+                .unwrap_or(PredictionProvider::Teacher(backend, ZetaFormat::default()))
+        } else {
+            match example.prompt.as_ref().map(|prompt| prompt.provider) {
+                Some(PredictionProvider::TeacherMultiRegion(_))
+                | Some(PredictionProvider::TeacherMultiRegionNonBatching(_)) => {
+                    PredictionProvider::TeacherMultiRegionNonBatching(backend)
+                }
+                _ => PredictionProvider::TeacherNonBatching(backend, ZetaFormat::default()),
+            }
+        };
+
+        let (actual_patch, actual_cursor) = match parser_provider {
+            PredictionProvider::TeacherMultiRegion(_)
+            | PredictionProvider::TeacherMultiRegionNonBatching(_) => {
+                TeacherMultiRegionPrompt::parse(example, &actual_output)?
+            }
+            _ => TeacherPrompt::parse(example, &actual_output)?,
+        };
+
+        let prediction = ExamplePrediction {
+            actual_patch: Some(actual_patch),
+            actual_output,
+            actual_cursor,
+            error: None,
+            provider: if batched {
+                match example.prompt.as_ref().map(|prompt| prompt.provider) {
+                    Some(PredictionProvider::TeacherMultiRegion(_)) => {
+                        PredictionProvider::TeacherMultiRegion(backend)
+                    }
+                    _ => PredictionProvider::Teacher(backend, ZetaFormat::default()),
+                }
+            } else {
+                match example.prompt.as_ref().map(|prompt| prompt.provider) {
+                    Some(PredictionProvider::TeacherMultiRegion(_))
+                    | Some(PredictionProvider::TeacherMultiRegionNonBatching(_)) => {
+                        PredictionProvider::TeacherMultiRegionNonBatching(backend)
+                    }
+                    _ => PredictionProvider::TeacherNonBatching(backend, ZetaFormat::default()),
+                }
+            },
+            cumulative_logprob: None,
+            avg_logprob: None,
+        };
+
+        example.predictions.push(prediction);
+    }
+    Ok(())
+}
+
+async fn predict_openai(
+    example: &mut Example,
+    backend: TeacherBackend,
+    batched: bool,
+    repetition_count: usize,
+    cache_only: bool,
+    step_progress: &crate::progress::StepProgress,
+) -> anyhow::Result<()> {
+    let llm_model_name = backend.model_name();
+    let max_tokens = 16384;
+    let llm_client = OPENAI_CLIENT.get_or_init(|| {
+        let client = if batched {
+            OpenAiClient::batch(&crate::paths::LLM_CACHE_DB)
+        } else {
+            OpenAiClient::plain()
+        };
+        client.expect("Failed to create OpenAI client")
+    });
+
+    let prompt = example.prompt.as_ref().context("Prompt is required")?;
+
+    for ix in 0..repetition_count {
+        if repetition_count > 1 {
+            step_progress.set_substatus(format!(
+                "running prediction {}/{}",
+                ix + 1,
+                repetition_count
+            ));
+        } else {
+            step_progress.set_substatus("running prediction");
+        }
+
+        let messages = vec![open_ai::RequestMessage::User {
+            content: open_ai::MessageContent::Plain(prompt.input.clone()),
+        }];
+
+        let seed = if repetition_count > 1 { Some(ix) } else { None };
+        let Some(response) = llm_client
+            .generate(llm_model_name, max_tokens, messages, seed, cache_only)
+            .await?
+        else {
+            // Request stashed for batched processing
+            continue;
+        };
+
+        let actual_output = response
+            .choices
+            .into_iter()
+            .filter_map(|choice| match choice.message {
+                open_ai::RequestMessage::Assistant { content, .. } => content.map(|c| match c {
+                    open_ai::MessageContent::Plain(text) => text,
+                    open_ai::MessageContent::Multipart(parts) => parts
+                        .into_iter()
+                        .filter_map(|p| match p {
+                            open_ai::MessagePart::Text { text } => Some(text),
+                            _ => None,
+                        })
+                        .collect::<Vec<_>>()
+                        .join(""),
+                }),
+                _ => None,
+            })
+            .collect::<Vec<String>>()
+            .join("\n");
+
+        let parser_provider = if batched {
+            example
+                .prompt
+                .as_ref()
+                .map(|prompt| prompt.provider)
+                .unwrap_or(PredictionProvider::Teacher(backend, ZetaFormat::default()))
+        } else {
+            match example.prompt.as_ref().map(|prompt| prompt.provider) {
+                Some(PredictionProvider::TeacherMultiRegion(_))
+                | Some(PredictionProvider::TeacherMultiRegionNonBatching(_)) => {
+                    PredictionProvider::TeacherMultiRegionNonBatching(backend)
+                }
+                _ => PredictionProvider::TeacherNonBatching(backend, ZetaFormat::default()),
+            }
+        };
+
+        let (actual_patch, actual_cursor) = match parser_provider {
+            PredictionProvider::TeacherMultiRegion(_)
+            | PredictionProvider::TeacherMultiRegionNonBatching(_) => {
+                TeacherMultiRegionPrompt::parse(example, &actual_output)?
+            }
+            _ => TeacherPrompt::parse(example, &actual_output)?,
+        };
+
+        let prediction = ExamplePrediction {
+            actual_patch: Some(actual_patch),
+            actual_output,
+            actual_cursor,
+            error: None,
+            provider: if batched {
+                match example.prompt.as_ref().map(|prompt| prompt.provider) {
+                    Some(PredictionProvider::TeacherMultiRegion(_)) => {
+                        PredictionProvider::TeacherMultiRegion(backend)
+                    }
+                    _ => PredictionProvider::Teacher(backend, ZetaFormat::default()),
+                }
+            } else {
+                match example.prompt.as_ref().map(|prompt| prompt.provider) {
+                    Some(PredictionProvider::TeacherMultiRegion(_))
+                    | Some(PredictionProvider::TeacherMultiRegionNonBatching(_)) => {
+                        PredictionProvider::TeacherMultiRegionNonBatching(backend)
+                    }
+                    _ => PredictionProvider::TeacherNonBatching(backend, ZetaFormat::default()),
+                }
+            },
+            cumulative_logprob: None,
+            avg_logprob: None,
+        };
+
+        example.predictions.push(prediction);
+    }
+    Ok(())
+}
+
+pub async fn predict_baseten(
+    example: &mut Example,
+    format: ZetaFormat,
+    step_progress: &StepProgress,
+) -> anyhow::Result<()> {
+    let model_id =
+        std::env::var("ZED_ZETA_MODEL").context("ZED_ZETA_MODEL environment variable required")?;
+
+    let api_key =
+        std::env::var("BASETEN_API_KEY").context("BASETEN_API_KEY environment variable not set")?;
+
+    let prompt = example.prompt.as_ref().context("Prompt is required")?;
+    let prompt_text = prompt.input.clone();
+    let prefill = prompt.prefill.clone().unwrap_or_default();
+
+    step_progress.set_substatus("running prediction via baseten");
+
+    let environment: String = <&'static str>::from(&format).to_lowercase();
+    let url = format!(
+        "https://model-{model_id}.api.baseten.co/environments/{environment}/sync/v1/completions"
+    );
+
+    let request_body = RawCompletionRequest {
+        model: model_id,
+        prompt: prompt_text.clone(),
+        max_tokens: Some(2048),
+        temperature: Some(0.),
+        stop: vec![],
+        environment: None,
+    };
+
+    let body_bytes =
+        serde_json::to_vec(&request_body).context("Failed to serialize request body")?;
+
+    let http_client: Arc<dyn HttpClient> = Arc::new(ReqwestClient::new());
+    let request = http_client::Request::builder()
+        .method(Method::POST)
+        .uri(&url)
+        .header("Content-Type", "application/json")
+        .header("Authorization", format!("Api-Key {api_key}"))
+        .body(AsyncBody::from(body_bytes))?;
+
+    let mut response = http_client.send(request).await?;
+    let status = response.status();
+
+    let mut body = String::new();
+    response
+        .body_mut()
+        .read_to_string(&mut body)
+        .await
+        .context("Failed to read Baseten response body")?;
+
+    if !status.is_success() {
+        anyhow::bail!("Baseten API returned {status}: {body}");
+    }
+
+    let completion: RawCompletionResponse =
+        serde_json::from_str(&body).context("Failed to parse Baseten response")?;
+
+    let actual_output = completion
+        .choices
+        .into_iter()
+        .next()
+        .map(|choice| choice.text)
+        .unwrap_or_default();
+
+    let actual_output = format!("{prefill}{actual_output}");
+
+    let (actual_patch, actual_cursor) =
+        parse_prediction_output(example, &actual_output, PredictionProvider::Zeta2(format))?;
+
+    let prediction = ExamplePrediction {
+        actual_patch: Some(actual_patch),
+        actual_output,
+        actual_cursor,
+        error: None,
+        provider: PredictionProvider::Baseten(format),
+        cumulative_logprob: None,
+        avg_logprob: None,
+    };
+
+    example.predictions.push(prediction);
+    Ok(())
+}
+
+pub async fn sync_batches(provider: Option<&PredictionProvider>) -> anyhow::Result<()> {
+    match provider {
+        Some(PredictionProvider::Teacher(backend, _))
+        | Some(PredictionProvider::TeacherMultiRegion(backend)) => match backend {
+            TeacherBackend::Sonnet45 | TeacherBackend::Sonnet46 => {
+                let llm_client = ANTHROPIC_CLIENT.get_or_init(|| {
+                    AnthropicClient::batch(&crate::paths::LLM_CACHE_DB)
+                        .expect("Failed to create Anthropic client")
+                });
+                llm_client
+                    .sync_batches()
+                    .await
+                    .context("Failed to sync Anthropic batches")?;
+            }
+            TeacherBackend::Gpt52 => {
+                let llm_client = OPENAI_CLIENT.get_or_init(|| {
+                    OpenAiClient::batch(&crate::paths::LLM_CACHE_DB)
+                        .expect("Failed to create OpenAI client")
+                });
+                llm_client
+                    .sync_batches()
+                    .await
+                    .context("Failed to sync OpenAI batches")?;
+            }
+        },
+        _ => (),
+    };
+    Ok(())
+}
+
+pub async fn reprocess_after_batch_wait(
+    examples: &mut [Example],
+    args: &PredictArgs,
+) -> anyhow::Result<()> {
+    let Some(PredictionProvider::Teacher(backend, _)) = args.provider else {
+        return Ok(());
+    };
+
+    let mut reprocessed = 0;
+    for example in examples.iter_mut() {
+        let has_prediction = example
+            .predictions
+            .iter()
+            .any(|p| p.actual_patch.is_some() || !p.actual_output.is_empty());
+        if has_prediction || example.prompt.is_none() {
+            continue;
+        }
+
+        let example_progress = Progress::global().start_group(&example.spec.name);
+        let step_progress = example_progress.start(Step::Predict);
+        predict_teacher(
+            example,
+            backend,
+            true,
+            args.repetitions,
+            false,
+            &step_progress,
+        )
+        .await?;
+        reprocessed += 1;
+    }
+
+    if reprocessed > 0 {
+        eprintln!("Reprocessed {} example(s) with batch results", reprocessed);
+    }
+
+    Ok(())
+}
+
+pub async fn wait_for_batches(provider: Option<&PredictionProvider>) -> anyhow::Result<()> {
+    let poll_interval = std::time::Duration::from_secs(30);
+
+    loop {
+        let pending = pending_batch_count(provider)?;
+        if pending == 0 {
+            break;
+        }
+
+        eprintln!(
+            "Waiting for {} pending batch request(s) to complete... (polling every {}s)",
+            pending,
+            poll_interval.as_secs()
+        );
+        std::thread::sleep(poll_interval);
+
+        sync_batches(provider).await?;
+    }
+
+    Ok(())
+}
+
+fn pending_batch_count(provider: Option<&PredictionProvider>) -> anyhow::Result<usize> {
+    match provider {
+        Some(PredictionProvider::Teacher(backend, _)) => match backend {
+            TeacherBackend::Sonnet45 | TeacherBackend::Sonnet46 => {
+                let llm_client = ANTHROPIC_CLIENT.get_or_init(|| {
+                    AnthropicClient::batch(&crate::paths::LLM_CACHE_DB)
+                        .expect("Failed to create Anthropic client")
+                });
+                llm_client.pending_batch_count()
+            }
+            TeacherBackend::Gpt52 => {
+                let llm_client = OPENAI_CLIENT.get_or_init(|| {
+                    OpenAiClient::batch(&crate::paths::LLM_CACHE_DB)
+                        .expect("Failed to create OpenAI client")
+                });
+                llm_client.pending_batch_count()
+            }
+        },
+        _ => Ok(0),
+    }
+}

crates/edit_prediction_cli/src/progress.rs

@@ -0,0 +1,651 @@
+use std::{
+    borrow::Cow,
+    collections::HashMap,
+    io::{IsTerminal, Write},
+    sync::{Arc, Mutex, OnceLock},
+    time::{Duration, Instant},
+};
+
+use crate::paths::RUN_DIR;
+
+use log::{Level, Log, Metadata, Record};
+
+pub struct Progress {
+    inner: Mutex<ProgressInner>,
+}
+
+struct ProgressInner {
+    completed: Vec<CompletedTask>,
+    in_progress: HashMap<String, InProgressTask>,
+    is_tty: bool,
+    terminal_width: usize,
+    max_example_name_len: usize,
+    status_lines_displayed: usize,
+    total_examples: usize,
+    completed_examples: usize,
+    failed_examples: usize,
+    last_line_is_logging: bool,
+    ticker: Option<std::thread::JoinHandle<()>>,
+}
+
+#[derive(Clone)]
+struct InProgressTask {
+    step: Step,
+    started_at: Instant,
+    substatus: Option<String>,
+    info: Option<(String, InfoStyle)>,
+}
+
+struct CompletedTask {
+    step: Step,
+    example_name: String,
+    duration: Duration,
+    info: Option<(String, InfoStyle)>,
+}
+
+#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
+pub enum Step {
+    LoadProject,
+    Context,
+    FormatPrompt,
+    Predict,
+    Score,
+    Qa,
+    Repair,
+    Synthesize,
+    PullExamples,
+}
+
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+pub enum InfoStyle {
+    Normal,
+    Warning,
+}
+
+impl Step {
+    pub fn label(&self) -> &'static str {
+        match self {
+            Step::LoadProject => "Load",
+            Step::Context => "Context",
+            Step::FormatPrompt => "Format",
+            Step::Predict => "Predict",
+            Step::Score => "Score",
+            Step::Qa => "QA",
+            Step::Repair => "Repair",
+            Step::Synthesize => "Synthesize",
+            Step::PullExamples => "Pull",
+        }
+    }
+
+    fn color_code(&self) -> &'static str {
+        match self {
+            Step::LoadProject => "\x1b[33m",
+            Step::Context => "\x1b[35m",
+            Step::FormatPrompt => "\x1b[34m",
+            Step::Predict => "\x1b[32m",
+            Step::Score => "\x1b[31m",
+            Step::Qa => "\x1b[36m",
+            Step::Repair => "\x1b[95m",
+            Step::Synthesize => "\x1b[36m",
+            Step::PullExamples => "\x1b[36m",
+        }
+    }
+}
+
+static GLOBAL: OnceLock<Arc<Progress>> = OnceLock::new();
+static LOGGER: ProgressLogger = ProgressLogger;
+
+const MARGIN: usize = 4;
+const MAX_STATUS_LINES: usize = 10;
+const STATUS_TICK_INTERVAL: Duration = Duration::from_millis(300);
+
+impl Progress {
+    /// Returns the global Progress instance, initializing it if necessary.
+    pub fn global() -> Arc<Progress> {
+        GLOBAL
+            .get_or_init(|| {
+                let progress = Arc::new(Self {
+                    inner: Mutex::new(ProgressInner {
+                        completed: Vec::new(),
+                        in_progress: HashMap::new(),
+                        is_tty: std::env::var("COLOR").is_ok()
+                            || (std::env::var("NO_COLOR").is_err()
+                                && std::io::stderr().is_terminal()),
+                        terminal_width: get_terminal_width(),
+                        max_example_name_len: 0,
+                        status_lines_displayed: 0,
+                        total_examples: 0,
+                        completed_examples: 0,
+                        failed_examples: 0,
+                        last_line_is_logging: false,
+                        ticker: None,
+                    }),
+                });
+                let _ = log::set_logger(&LOGGER);
+                log::set_max_level(log::LevelFilter::Info);
+                progress
+            })
+            .clone()
+    }
+
+    pub fn start_group(self: &Arc<Self>, example_name: &str) -> ExampleProgress {
+        ExampleProgress {
+            progress: self.clone(),
+            example_name: example_name.to_string(),
+        }
+    }
+
+    fn increment_completed(&self) {
+        let mut inner = self.inner.lock().unwrap();
+        inner.completed_examples += 1;
+    }
+
+    pub fn set_total_examples(&self, total: usize) {
+        let mut inner = self.inner.lock().unwrap();
+        inner.total_examples = total;
+    }
+
+    pub fn set_max_example_name_len(&self, example_names: impl Iterator<Item = impl AsRef<str>>) {
+        let mut inner = self.inner.lock().unwrap();
+        let max_name_width = inner
+            .terminal_width
+            .saturating_sub(MARGIN * 2)
+            .saturating_div(3)
+            .max(1);
+        inner.max_example_name_len = example_names
+            .map(|name| name.as_ref().len().min(max_name_width))
+            .max()
+            .unwrap_or(0);
+    }
+
+    pub fn increment_failed(&self) {
+        let mut inner = self.inner.lock().unwrap();
+        inner.failed_examples += 1;
+    }
+
+    /// Prints a message to stderr, clearing and redrawing status lines to avoid corruption.
+    /// This should be used for any output that needs to appear above the status lines.
+    fn log(&self, message: &str) {
+        let mut inner = self.inner.lock().unwrap();
+        Self::clear_status_lines(&mut inner);
+
+        if !inner.last_line_is_logging {
+            let reset = "\x1b[0m";
+            let dim = "\x1b[2m";
+            let divider = "─".repeat(inner.terminal_width.saturating_sub(MARGIN));
+            eprintln!("{dim}{divider}{reset}");
+            inner.last_line_is_logging = true;
+        }
+
+        let max_width = inner.terminal_width.saturating_sub(MARGIN);
+        for line in message.lines() {
+            let truncated = truncate_to_visible_width(line, max_width);
+            if truncated.len() < line.len() {
+                eprintln!("{}…", truncated);
+            } else {
+                eprintln!("{}", truncated);
+            }
+        }
+
+        Self::print_status_lines(&mut inner);
+    }
+
+    pub fn start(self: &Arc<Self>, step: Step, example_name: &str) -> StepProgress {
+        let mut inner = self.inner.lock().unwrap();
+
+        Self::clear_status_lines(&mut inner);
+
+        // Update max_example_name_len if not already set via set_max_example_name_len
+        if inner.max_example_name_len == 0 {
+            let max_name_width = inner
+                .terminal_width
+                .saturating_sub(MARGIN * 2)
+                .saturating_div(3)
+                .max(1);
+            inner.max_example_name_len = example_name.len().min(max_name_width);
+        }
+        inner.in_progress.insert(
+            example_name.to_string(),
+            InProgressTask {
+                step,
+                started_at: Instant::now(),
+                substatus: None,
+                info: None,
+            },
+        );
+
+        if inner.is_tty && inner.ticker.is_none() {
+            let progress = self.clone();
+            inner.ticker = Some(std::thread::spawn(move || {
+                loop {
+                    std::thread::sleep(STATUS_TICK_INTERVAL);
+
+                    let mut inner = progress.inner.lock().unwrap();
+                    if inner.in_progress.is_empty() {
+                        break;
+                    }
+
+                    Progress::clear_status_lines(&mut inner);
+                    Progress::print_status_lines(&mut inner);
+                }
+            }));
+        }
+
+        Self::print_status_lines(&mut inner);
+
+        StepProgress {
+            progress: self.clone(),
+            step,
+            example_name: example_name.to_string(),
+        }
+    }
+
+    fn finish(&self, step: Step, example_name: &str) {
+        let mut inner = self.inner.lock().unwrap();
+
+        let Some(task) = inner.in_progress.remove(example_name) else {
+            return;
+        };
+
+        if task.step == step {
+            let duration = task.started_at.elapsed();
+
+            // Skip logging for tasks that complete quickly (under 500ms)
+            let should_print = duration >= Duration::from_millis(500);
+
+            inner.completed.push(CompletedTask {
+                step: task.step,
+                example_name: example_name.to_string(),
+                duration,
+                info: task.info,
+            });
+
+            Self::clear_status_lines(&mut inner);
+            if should_print {
+                Self::print_logging_closing_divider(&mut inner);
+                if let Some(last_completed) = inner.completed.last() {
+                    Self::print_completed(&inner, last_completed);
+                }
+            }
+            Self::print_status_lines(&mut inner);
+        } else {
+            inner.in_progress.insert(example_name.to_string(), task);
+        }
+    }
+
+    fn print_logging_closing_divider(inner: &mut ProgressInner) {
+        if inner.last_line_is_logging {
+            let reset = "\x1b[0m";
+            let dim = "\x1b[2m";
+            let divider = "─".repeat(inner.terminal_width.saturating_sub(MARGIN));
+            eprintln!("{dim}{divider}{reset}");
+            inner.last_line_is_logging = false;
+        }
+    }
+
+    fn clear_status_lines(inner: &mut ProgressInner) {
+        if inner.is_tty && inner.status_lines_displayed > 0 {
+            // Move up and clear each line we previously displayed
+            for _ in 0..inner.status_lines_displayed {
+                eprint!("\x1b[A\x1b[K");
+            }
+            inner.status_lines_displayed = 0;
+        }
+    }
+
+    fn print_completed(inner: &ProgressInner, task: &CompletedTask) {
+        let duration = format_duration(task.duration);
+        let name_width = inner.max_example_name_len;
+        let truncated_name = truncate_with_ellipsis(&task.example_name, name_width);
+
+        if inner.is_tty {
+            let reset = "\x1b[0m";
+            let bold = "\x1b[1m";
+            let dim = "\x1b[2m";
+
+            let yellow = "\x1b[33m";
+            let info_part = task
+                .info
+                .as_ref()
+                .map(|(s, style)| {
+                    if *style == InfoStyle::Warning {
+                        format!("{yellow}{s}{reset}")
+                    } else {
+                        s.to_string()
+                    }
+                })
+                .unwrap_or_default();
+
+            let prefix = format!(
+                "{bold}{color}{label:>12}{reset} {name:<name_width$} {dim}│{reset} {info_part}",
+                color = task.step.color_code(),
+                label = task.step.label(),
+                name = truncated_name,
+            );
+
+            let duration_with_margin = format!("{duration} ");
+            let padding_needed = inner
+                .terminal_width
+                .saturating_sub(MARGIN)
+                .saturating_sub(duration_with_margin.len())
+                .saturating_sub(strip_ansi_len(&prefix));
+            let padding = " ".repeat(padding_needed);
+
+            eprintln!("{prefix}{padding}{dim}{duration_with_margin}{reset}");
+        } else {
+            let info_part = task
+                .info
+                .as_ref()
+                .map(|(s, _)| format!(" | {}", s))
+                .unwrap_or_default();
+
+            eprintln!(
+                "{label:>12} {name:<name_width$}{info_part} {duration}",
+                label = task.step.label(),
+                name = truncate_with_ellipsis(&task.example_name, name_width),
+            );
+        }
+    }
+
+    fn print_status_lines(inner: &mut ProgressInner) {
+        if !inner.is_tty || inner.in_progress.is_empty() {
+            inner.status_lines_displayed = 0;
+            return;
+        }
+
+        let reset = "\x1b[0m";
+        let bold = "\x1b[1m";
+        let dim = "\x1b[2m";
+
+        // Build the done/in-progress/total label
+        let done_count = inner.completed_examples;
+        let in_progress_count = inner.in_progress.len();
+        let failed_count = inner.failed_examples;
+
+        let failed_label = if failed_count > 0 {
+            format!(" {} failed ", failed_count)
+        } else {
+            String::new()
+        };
+
+        let range_label = format!(
+            " {}/{}/{} ",
+            done_count, in_progress_count, inner.total_examples
+        );
+
+        // Print a divider line with failed count on left, range label on right
+        let failed_visible_len = strip_ansi_len(&failed_label);
+        let range_visible_len = range_label.len();
+        let middle_divider_len = inner
+            .terminal_width
+            .saturating_sub(MARGIN * 2)
+            .saturating_sub(failed_visible_len)
+            .saturating_sub(range_visible_len);
+        let left_divider = "─".repeat(MARGIN);
+        let middle_divider = "─".repeat(middle_divider_len);
+        let right_divider = "─".repeat(MARGIN);
+        eprintln!(
+            "{dim}{left_divider}{reset}{failed_label}{dim}{middle_divider}{reset}{range_label}{dim}{right_divider}{reset}"
+        );
+
+        let mut tasks: Vec<_> = inner.in_progress.iter().collect();
+        tasks.sort_by_key(|(name, _)| *name);
+
+        let total_tasks = tasks.len();
+        let mut lines_printed = 0;
+
+        for (name, task) in tasks.iter().take(MAX_STATUS_LINES) {
+            let elapsed = format_duration(task.started_at.elapsed());
+            let substatus_part = task
+                .substatus
+                .as_ref()
+                .map(|s| truncate_with_ellipsis(s, 30))
+                .unwrap_or_default();
+
+            let step_label = task.step.label();
+            let step_color = task.step.color_code();
+            let name_width = inner.max_example_name_len;
+            let truncated_name = truncate_with_ellipsis(name, name_width);
+
+            let prefix = format!(
+                "{bold}{step_color}{step_label:>12}{reset} {name:<name_width$} {dim}│{reset} {substatus_part}",
+                name = truncated_name,
+            );
+
+            let duration_with_margin = format!("{elapsed} ");
+            let padding_needed = inner
+                .terminal_width
+                .saturating_sub(MARGIN)
+                .saturating_sub(duration_with_margin.len())
+                .saturating_sub(strip_ansi_len(&prefix));
+            let padding = " ".repeat(padding_needed);
+
+            eprintln!("{prefix}{padding}{dim}{duration_with_margin}{reset}");
+            lines_printed += 1;
+        }
+
+        // Show "+N more" on its own line if there are more tasks
+        if total_tasks > MAX_STATUS_LINES {
+            let remaining = total_tasks - MAX_STATUS_LINES;
+            eprintln!("{:>12} +{remaining} more", "");
+            lines_printed += 1;
+        }
+
+        inner.status_lines_displayed = lines_printed + 1; // +1 for the divider line
+        let _ = std::io::stderr().flush();
+    }
+
+    pub fn finalize(&self) {
+        let ticker = {
+            let mut inner = self.inner.lock().unwrap();
+            inner.ticker.take()
+        };
+
+        if let Some(ticker) = ticker {
+            let _ = ticker.join();
+        }
+
+        let mut inner = self.inner.lock().unwrap();
+        Self::clear_status_lines(&mut inner);
+
+        // Print summary if there were failures
+        if inner.failed_examples > 0 {
+            let total_examples = inner.total_examples;
+            let percentage = if total_examples > 0 {
+                inner.failed_examples as f64 / total_examples as f64 * 100.0
+            } else {
+                0.0
+            };
+            let failed_jsonl_path = RUN_DIR.join("failed.jsonl");
+            eprintln!(
+                "\n{} of {} examples failed ({:.1}%)\nFailed examples: {}",
+                inner.failed_examples,
+                total_examples,
+                percentage,
+                failed_jsonl_path.display()
+            );
+        }
+    }
+}
+
+pub struct ExampleProgress {
+    progress: Arc<Progress>,
+    example_name: String,
+}
+
+impl ExampleProgress {
+    pub fn start(&self, step: Step) -> StepProgress {
+        self.progress.start(step, &self.example_name)
+    }
+}
+
+impl Drop for ExampleProgress {
+    fn drop(&mut self) {
+        self.progress.increment_completed();
+    }
+}
+
+pub struct StepProgress {
+    progress: Arc<Progress>,
+    step: Step,
+    example_name: String,
+}
+
+impl StepProgress {
+    pub fn set_substatus(&self, substatus: impl Into<Cow<'static, str>>) {
+        let mut inner = self.progress.inner.lock().unwrap();
+        if let Some(task) = inner.in_progress.get_mut(&self.example_name) {
+            task.substatus = Some(substatus.into().into_owned());
+            Progress::clear_status_lines(&mut inner);
+            Progress::print_status_lines(&mut inner);
+        }
+    }
+
+    pub fn clear_substatus(&self) {
+        let mut inner = self.progress.inner.lock().unwrap();
+        if let Some(task) = inner.in_progress.get_mut(&self.example_name) {
+            task.substatus = None;
+            Progress::clear_status_lines(&mut inner);
+            Progress::print_status_lines(&mut inner);
+        }
+    }
+
+    pub fn set_info(&self, info: impl Into<String>, style: InfoStyle) {
+        let mut inner = self.progress.inner.lock().unwrap();
+        if let Some(task) = inner.in_progress.get_mut(&self.example_name) {
+            task.info = Some((info.into(), style));
+        }
+    }
+}
+
+impl Drop for StepProgress {
+    fn drop(&mut self) {
+        self.progress.finish(self.step, &self.example_name);
+    }
+}
+
+struct ProgressLogger;
+
+impl Log for ProgressLogger {
+    fn enabled(&self, metadata: &Metadata) -> bool {
+        metadata.level() <= Level::Info
+    }
+
+    fn log(&self, record: &Record) {
+        if !self.enabled(record.metadata()) {
+            return;
+        }
+
+        let level_color = match record.level() {
+            Level::Error => "\x1b[31m",
+            Level::Warn => "\x1b[33m",
+            Level::Info => "\x1b[32m",
+            Level::Debug => "\x1b[34m",
+            Level::Trace => "\x1b[35m",
+        };
+        let reset = "\x1b[0m";
+        let bold = "\x1b[1m";
+
+        let level_label = match record.level() {
+            Level::Error => "Error",
+            Level::Warn => "Warn",
+            Level::Info => "Info",
+            Level::Debug => "Debug",
+            Level::Trace => "Trace",
+        };
+
+        let message = format!(
+            "{bold}{level_color}{level_label:>12}{reset} {}",
+            record.args()
+        );
+
+        if let Some(progress) = GLOBAL.get() {
+            progress.log(&message);
+        } else {
+            eprintln!("{}", message);
+        }
+    }
+
+    fn flush(&self) {
+        let _ = std::io::stderr().flush();
+    }
+}
+
+#[cfg(unix)]
+fn get_terminal_width() -> usize {
+    unsafe {
+        let mut winsize: libc::winsize = std::mem::zeroed();
+        if libc::ioctl(libc::STDERR_FILENO, libc::TIOCGWINSZ, &mut winsize) == 0
+            && winsize.ws_col > 0
+        {
+            winsize.ws_col as usize
+        } else {
+            80
+        }
+    }
+}
+
+#[cfg(not(unix))]
+fn get_terminal_width() -> usize {
+    80
+}
+
+fn strip_ansi_len(s: &str) -> usize {
+    let mut len = 0;
+    let mut in_escape = false;
+    for c in s.chars() {
+        if c == '\x1b' {
+            in_escape = true;
+        } else if in_escape {
+            if c == 'm' {
+                in_escape = false;
+            }
+        } else {
+            len += 1;
+        }
+    }
+    len
+}
+
+fn truncate_with_ellipsis(s: &str, max_len: usize) -> Cow<'_, str> {
+    if s.len() <= max_len {
+        Cow::Borrowed(s)
+    } else {
+        Cow::Owned(format!("{}…", &s[..max_len.saturating_sub(1)]))
+    }
+}
+
+fn truncate_to_visible_width(s: &str, max_visible_len: usize) -> &str {
+    let mut visible_len = 0;
+    let mut in_escape = false;
+    let mut last_byte_index = 0;
+    for (byte_index, c) in s.char_indices() {
+        if c == '\x1b' {
+            in_escape = true;
+        } else if in_escape {
+            if c == 'm' {
+                in_escape = false;
+            }
+        } else {
+            if visible_len >= max_visible_len {
+                return &s[..last_byte_index];
+            }
+            visible_len += 1;
+        }
+        last_byte_index = byte_index + c.len_utf8();
+    }
+    s
+}
+
+fn format_duration(duration: Duration) -> String {
+    const MINUTE_IN_MILLIS: f32 = 60. * 1000.;
+
+    let millis = duration.as_millis() as f32;
+    if millis < 1000.0 {
+        format!("{}ms", millis)
+    } else if millis < MINUTE_IN_MILLIS {
+        format!("{:.1}s", millis / 1_000.0)
+    } else {
+        format!("{:.1}m", millis / MINUTE_IN_MILLIS)
+    }
+}

crates/edit_prediction_cli/src/prompt_assets.rs

@@ -0,0 +1,46 @@
+use std::borrow::Cow;
+
+#[cfg(feature = "dynamic_prompts")]
+pub fn get_prompt(name: &'static str) -> Cow<'static, str> {
+    use anyhow::Context;
+    use std::collections::HashMap;
+    use std::path::Path;
+    use std::sync::{LazyLock, RwLock};
+
+    const PROMPTS_DIR: &str = concat!(env!("CARGO_MANIFEST_DIR"), "/src/prompts");
+
+    static PROMPT_CACHE: LazyLock<RwLock<HashMap<&'static str, &'static str>>> =
+        LazyLock::new(|| RwLock::new(HashMap::default()));
+
+    let filesystem_path = Path::new(PROMPTS_DIR).join(name);
+    if let Some(cached_contents) = PROMPT_CACHE.read().unwrap().get(name) {
+        return Cow::Borrowed(cached_contents);
+    }
+    let contents = std::fs::read_to_string(&filesystem_path)
+        .context(name)
+        .expect("Failed to read prompt");
+    let leaked = contents.leak();
+    PROMPT_CACHE.write().unwrap().insert(name, leaked);
+    return Cow::Borrowed(leaked);
+}
+
+#[cfg(not(feature = "dynamic_prompts"))]
+pub fn get_prompt(name: &'static str) -> Cow<'static, str> {
+    use rust_embed::RustEmbed;
+
+    #[derive(RustEmbed)]
+    #[folder = "src/prompts"]
+    struct EmbeddedPrompts;
+
+    match EmbeddedPrompts::get(name) {
+        Some(file) => match file.data {
+            Cow::Borrowed(bytes) => {
+                Cow::Borrowed(std::str::from_utf8(bytes).expect("prompt file is not valid UTF-8"))
+            }
+            Cow::Owned(bytes) => {
+                Cow::Owned(String::from_utf8(bytes).expect("prompt file is not valid UTF-8"))
+            }
+        },
+        None => panic!("prompt file not found: {name}"),
+    }
+}

crates/edit_prediction_cli/src/prompts/qa.md

@@ -0,0 +1,43 @@
+You are evaluating an edit prediction model for a code editor. The model observes a programmer's recent edit history and predicts what edit they will make next.
+
+All diffs are in the word-diff format.
+
+The model is instructed to:
+- Complete partially-applied refactoring or changes
+- Maintain consistency with established patterns and style
+- NOT delete or revert text that was just added (unless the user explicitly undid it themselves)
+
+## Edit History (chronological)
+```````
+{edit_history}
+```````
+
+## Current File
+The file where the prediction will be applied, with editable region markers showing where edits can occur:
+{cursor_excerpt}
+
+## Predicted Next Edit
+```````
+{actual_patch_word_diff}
+```````
+
+## Evaluate
+
+1. **reverts_edits**: Does the prediction undo, or revert changes the user intentionally made in the **edit history**?
+
+2. **confidence**: How likely is the user to accept this suggestion?
+   - 1 = Definitely reject (wrong, nonsensical, or harmful)
+   - 2 = Probably reject (doesn't fit intent or pattern)
+   - 3 = Uncertain (plausible but not clearly correct)
+   - 4 = Probably accept (reasonable next step)
+   - 5 = Definitely accept (obvious continuation)
+
+Output JSON in this format:
+
+```
+{
+    "reasoning": "your reasoning here",
+    "reverts_edits": true/false,
+    "confidence": 1-5
+}
+```

crates/edit_prediction_cli/src/prompts/repair.md

@@ -0,0 +1,44 @@
+# Repair Request
+
+Your previous prediction has quality issues that need to be addressed. Please generate an improved prediction.
+
+## Quality Feedback
+
+{quality_feedback}
+
+{token_change_info}
+
+## Your Previous Prediction (word-diff format)
+
+`````
+{actual_patch_word_diff}
+`````
+
+## Instructions
+
+Generate an improved prediction following the same rules and output format from the original instructions. The key rules remain:
+
+- **NEVER undo or revert the user's recent edits** — if a line was removed in the edit history, do NOT restore it
+- If your prediction would make the code more similar to what it was BEFORE the user's edit, output `NO_EDITS` instead
+- When uncertain, predict only the minimal, high-confidence portion of the edit
+
+## Output Format
+
+Follow the same output format as before, with one addition:
+
+- If the code is complete as-is and no edits should be made, output `NO_EDITS`
+- **NEW: If your previous prediction was actually correct** (the quality feedback was overly cautious), output `KEEP_PREVIOUS`:
+
+  `````
+  KEEP_PREVIOUS
+  `````
+
+  Use `KEEP_PREVIOUS` when you determine the original prediction correctly addresses the user's intent despite the feedback.
+
+**Important:** `NO_EDITS` and `KEEP_PREVIOUS` are NOT interchangeable:
+- `NO_EDITS` = make zero changes to the code (discard the previous prediction)
+- `KEEP_PREVIOUS` = the previous prediction is correct, use it as-is
+
+## Your Improved Prediction
+
+Briefly explain what was wrong with your previous prediction (or why it was actually correct), then provide the improved output.

crates/edit_prediction_cli/src/prompts/teacher.md

@@ -0,0 +1,358 @@
+# Instructions
+
+You are an edit prediction assistant in a code editor. Your task is to predict the next edit to a given region of code surrounding the user's cursor.
+
+1. Analyze the edit history to understand what the programmer is trying to achieve
+2. Identify any incomplete refactoring or changes that need to be finished
+3. Make the remaining edits that a human programmer would logically make next (by rewriting the code around their cursor)
+
+## Focus on
+
+- Completing any partially-applied changes made
+- Ensuring consistency with the programming style and patterns already established
+- Making edits that maintain or improve code quality
+
+## Rules
+
+- **NEVER undo or revert the user's recent edits.** Examine the diff in the edit history carefully:
+  - If a line was removed (starts with `-`), do NOT restore that content—even if the code now appears incomplete or broken without it
+  - If a line was added (starts with `+`), do NOT delete or significantly modify it
+  - If code appears broken or incomplete after the user's edit, output `NO_EDITS` rather than "fixing" it by reverting
+  - Only add NEW content that extends the user's work forward; never restore what they removed
+  - **Key test**: if your prediction would make the code more similar to what it was BEFORE the user's edit, output `NO_EDITS` instead
+  - **Never assume a deletion was accidental.** Even if removing content breaks the code, breaks a pattern, or leaves text looking "incomplete", respect it. The user may be mid-rewrite. Do NOT "complete" partial text by restoring what was deleted.
+- Auto-generated code can be modified: Hunks marked with `// User accepted prediction:` contain code from a previous prediction the user accepted. Unlike user-typed content, these hunks CAN be edited, corrected, or replaced if it improves the code. The "never undo/revert" rule protects the user's *current typing intent*—auto-generated code doesn't carry this protection
+- Do not just mechanically apply patterns - reason about what changes make sense given the context and the programmer's apparent goals.
+- Do not just fix syntax errors - look for the broader refactoring pattern and apply it systematically throughout the code.
+- Keep existing formatting unless it's absolutely necessary
+- When edit history and surrounding code suggest different edits, prioritize the most recent edits in the history as they best reflect current intent.
+- Treat partial text at or near the cursor as the beginning of something the user is actively typing. Complete the code the user appears to be creating based on context.
+- When completing partial code, prefer predictions that save meaningful keystrokes, even if this requires making educated guesses about the user's intent.
+- For code, it's better to make a substantive prediction that might be rejected than to make a minimal prediction that saves only a few keystrokes.
+- When the user is editing prose or documentation (e.g. Markdown, comments, plain text), predict conservatively. Complete the current fragment or sentence, but do not generate additional lines of free-form content since prose is less constrained than code and more prone to incorrect continuations.
+
+# Input Format
+
+You will be provided with:
+1. The user's *edit history*, in chronological order. Use this to infer the user's trajectory and predict the next most logical edit.
+  - Hunks preceded by `// User accepted prediction:` indicate code that was auto-generated by a previous prediction and accepted by the user. These are treated differently than user-typed edits (see Rules).
+2. A set of *related excerpts* from the user's codebase. Some of these may be needed for correctly predicting the next edit.
+  - `…` may appear within a related file to indicate that some code has been skipped.
+3. An excerpt from the user's *current file*.
+    - Within the user's current file, there is an *editable region* delimited by the `<|editable_region_start|>` and `<|editable_region_end|>` tags. You can only predict edits in this region.
+    - The `<|user_cursor|>` tag marks the user's current cursor position, as it stands after the last edit in the history.
+
+# Output Format
+
+- Briefly explain the user's current intent based on the edit history and their current cursor location.
+- Output a markdown codeblock containing **only** the editable region with your predicted edits applied. The codeblock must start with `<|editable_region_start|>` and end with `<|editable_region_end|>`. Do not include any content before or after these tags.
+- If no edit is needed (the code is already complete and correct, or there is no clear next edit to make), output a codeblock containing only `NO_EDITS`:
+  `````
+  NO_EDITS
+  `````
+- If there is a specific place in the predicted output where the user is likely to edit next, indicate it using the `<|user_cursor|>` tag.
+
+## Example 1
+
+There is code missing at the cursor location. The related excerpts includes the definition of a relevant type. You should fill in the missing code.
+
+### Related Excerpts
+
+`````
+struct Product {
+    name: String,
+    price: u32,
+}
+`````
+
+### User Edit History
+
+`````
+--- a/src/calculate.rs
++++ b/src/calculate.rs
+@@ -100,6 +100,7 @@
+ fn calculate_total(products: &[Product]) -> u32 {
+     let mut total = 0;
+     for product in products {
++        total += ;
+     }
+     total
+ }
+`````
+
+### Current File
+
+`````src/calculate.rs
+fn calculate_total(products: &[Product]) -> u32 {
+<|editable_region_start|>
+    let mut total = 0;
+    for product in products {
+        total += <|user_cursor|>;
+    }
+    total
+<|editable_region_end|>
+}
+`````
+
+### Output
+
+The user is computing a sum based on a list of products. The only numeric field on `Product` is `price`, so they must intend to sum the prices.
+
+`````
+<|editable_region_start|>
+    let mut total = 0;
+    for product in products {
+        total += product.price;
+    }
+    total
+<|editable_region_end|>
+`````
+
+## Example 2
+
+The user appears to be in the process of typing an eprintln call. Rather than fixing the spelling issue by deleting the newly-inserted content, you must continue the user's trajectory. It's not clear what data they intend to print. You should fill in as much code as is obviously intended, and position the cursor so that the user can fill in the rest.
+
+### User Edit History
+
+`````
+--- a/src/modal.rs
++++ b/src/modal.rs
+@@ -100,4 +100,4 @@
+ fn handle_close_button_click(modal_state: &mut ModalState, evt: &Event) {
+     modal_state.close();
+-     modal_state.dismiss();
++     eprmodal_state.dismiss();
+ }
+`````
+
+### Current File
+
+`````src/modal.rs
+// handle the close button click
+<|editable_region_start|>
+fn handle_close_button_click(modal_state: &mut ModalState, evt: &Event) {
+    modal_state.close();
+    epr<|user_cursor|>modal_state.dismiss();
+<|editable_region_end|>
+}
+`````
+
+### Output
+
+The user is clearly starting to type `eprintln!()`, however, what they intend to print is not obvious. I should fill in the print call and string literal, with the cursor positioned inside the string literal so the user can print whatever they want.
+
+`````
+<|editable_region_start|>
+fn handle_close_button_click(modal_state: &mut ModalState, evt: &Event) {
+    modal_state.close();
+    eprintln!("<|user_cursor|>");
+    modal_state.dismiss();
+<|editable_region_end|>
+`````
+
+## Example 3
+
+Here, the user is adding a function. There's no way to tell for sure what the function's name will be. In this situation, you should make a reasonable guess at the function's name and signature, and place the user's cursor in the function body. This way, if you guess correctly, it will save the user a meaningful number of keystrokes, and the file will be left in a coherent state.
+
+### User Edit History
+
+`````
+--- a/src/modal.rs
++++ b/src/modal.rs
+@@ -100,4 +100,4 @@
+ fn handle_close_button_click(modal_state: &mut ModalState, evt: &Event) {
+     modal_state.close();
+     modal_state.dismiss();
+ }
++
++fn
+
+ fn handle_keystroke(modal_state: &mut ModalState, evt: &Event) {
+`````
+
+### Current File
+
+`````src/modal.rs
+// handle the close button click
+fn handle_close_button_click(modal_state: &mut ModalState, evt: &Event) {
+    modal_state.close();
+<|editable_region_start|>
+    modal_state.dismiss();
+}
+
+fn<|user_cursor|>
+
+fn handle_keystroke(modal_state: &mut ModalState, evt: &Event) {
+<|editable_region_end|>
+    modal_state.begin_edit();
+`````
+
+### Output
+
+The user is adding a new function. The existing functions I see are `handle_close_button_click` and `handle_keystroke`, which have similar signatures. One possible function they might be adding is `handle_submit`.
+
+`````
+<|editable_region_start|>
+    modal_state.dismiss();
+}
+
+fn handle_submit(modal_state: &mut ModalState, evt: &Event) {
+    <|user_cursor|>
+}
+
+fn handle_keystroke(modal_state: &mut ModalState, evt: &Event) {
+<|editable_region_end|>
+`````
+
+## Example 4
+
+The code is already complete and there is no clear next edit to make. You should output NO_EDITS.
+
+### User Edit History
+
+`````
+--- a/src/utils.rs
++++ b/src/utils.rs
+@@ -10,7 +10,7 @@
+ fn add(a: i32, b: i32) -> i32 {
+-    a - b
++    a + b
+ }
+`````
+
+### Current File
+
+`````src/utils.rs
+<|editable_region_start|>
+fn add(a: i32, b: i32) -> i32 {
+    a + b<|user_cursor|>
+}
+<|editable_region_end|>
+`````
+
+### Output
+
+The user just fixed a bug in the `add` function, changing subtraction to addition. The code is now correct and complete. There is no clear next edit to make.
+
+`````
+NO_EDITS
+`````
+
+## Example 5
+
+The user just deleted code, leaving behind what looks incomplete. You must NOT "complete" it by restoring deleted content—that would undo their edit. Output NO_EDITS. **This is the correct response even though the code appears broken.**
+
+### User Edit History
+
+`````
+--- a/config.nix
++++ b/config.nix
+@@ -10,7 +10,7 @@
+     # /etc/modular/crashdb needs to be mutable
+-    ln -s /tmp/crashdb $out/etc/modular/crashdb
++    ln -s /tmp/cr $out/etc/modular/crashdb
+   '';
+`````
+
+### Current File
+
+`````config.nix
+<|editable_region_start|>
+    # /etc/modular/crashdb needs to be mutable
+    ln -s /tmp/cr<|user_cursor|> $out/etc/modular/crashdb
+  '';
+<|editable_region_end|>
+`````
+
+### Output
+
+The user deleted `ashdb` from `/tmp/crashdb`, leaving `/tmp/cr`. Although this looks like incomplete text that I could "complete", doing so would restore deleted content. The user intentionally removed that text—I must not undo their deletion.
+
+`````
+NO_EDITS
+`````
+
+## Example 6
+
+The user accepted a prediction for a function, then started renaming it. The original arguments were auto-generated (marked with `// User accepted prediction:`), so they CAN be updated to match the new function name. This is NOT reverting user input—it's improving auto-generated scaffolding.
+
+### User Edit History
+
+`````
+--- a/math_utils.py
++++ b/math_utils.py
+@@ -3,3 +3,5 @@
+ def calculate_rectangle_area(width, height):
+     return width * height
+
++de
+
+// User accepted prediction:
+--- a/math_utils.py
++++ b/math_utils.py
+@@ -3,5 +3,7 @@
+ def calculate_rectangle_area(width, height):
+     return width * height
+
+-de
++def calculate_rectangle_perimeter(width, height):
++
+
+--- a/math_utils.py
++++ b/math_utils.py
+@@ -5,5 +5,5 @@
+     return width * height
+
+-def calculate_rectangle_perimeter(width, height):
++def calculate_sq_perimeter(width, height):
+
+`````
+
+### Current File
+
+`````math_utils.py
+def calculate_rectangle_area(width, height):
+    return width * height
+
+<|editable_region_start|>
+def calculate_sq<|user_cursor|>_perimeter(width, height):
+
+<|editable_region_end|>
+`````
+
+### Output
+
+The user accepted a prediction for `calculate_rectangle_perimeter(width, height)`, then started renaming `rectangle` to `square`. Since squares have equal sides, the arguments should change from `(width, height)` to `(side)`. The arguments were auto-generated (from an accepted prediction), so modifying them is appropriate.
+
+`````
+<|editable_region_start|>
+def calculate_square_perimeter(side):
+    <|user_cursor|>
+<|editable_region_end|>
+`````
+
+
+
+# Your task:
+
+# 1. User Edit History
+
+`````
+{{edit_history}}
+`````
+
+# 2. Related excerpts
+
+{{context}}
+
+# 3. Current File
+
+{{cursor_excerpt}}
+
+
+{{diagnostics}}
+
+
+-----
+
+Based on the edit history and context above, predict the user's next edit within the editable region.

crates/edit_prediction_cli/src/prompts/teacher_multi_region.md

@@ -0,0 +1,366 @@
+# Instructions
+
+You are an edit prediction assistant in a code editor. Your task is to predict the next edit to a given region of code surrounding the user's cursor.
+
+1. Analyze the edit history to understand what the programmer is trying to achieve
+2. Identify any incomplete refactoring or changes that need to be finished
+3. Make the remaining edits that a human programmer would logically make next (by rewriting a region of code near their cursor)
+
+## Focus on
+
+- Completing any partially-applied changes made
+- Ensuring consistency with the programming style and patterns already established
+- Making edits that maintain or improve code quality
+
+## Rules
+
+- **NEVER undo or revert the user's recent edits.** Examine the diff in the edit history carefully:
+  - If a line was removed (starts with `-`), do NOT restore that content—even if the code now appears incomplete or broken without it
+  - If a line was added (starts with `+`), do NOT delete or significantly modify it
+  - If code appears broken or incomplete after the user's edit, output `NO_EDITS` rather than "fixing" it by reverting
+  - Only add NEW content that extends the user's work forward; never restore what they removed
+  - **Key test**: if your prediction would make the code more similar to what it was BEFORE the user's edit, output `NO_EDITS` instead
+  - **Never assume a deletion was accidental.** Even if removing content breaks the code, breaks a pattern, or leaves text looking "incomplete", respect it. The user may be mid-rewrite. Do NOT "complete" partial text by restoring what was deleted.
+- Auto-generated code can be modified: Hunks marked with `// User accepted prediction:` contain code from a previous prediction the user accepted. Unlike user-typed content, these hunks CAN be edited, corrected, or replaced if it improves the code. The "never undo/revert" rule protects the user's *current typing intent*—auto-generated code doesn't carry this protection
+- Do not just mechanically apply patterns - reason about what changes make sense given the context and the programmer's apparent goals.
+- Do not just fix syntax errors - look for the broader refactoring pattern and apply it systematically throughout the code.
+- Keep existing formatting unless it's absolutely necessary
+- When edit history and surrounding code suggest different edits, prioritize the most recent edits in the history as they best reflect current intent.
+- Treat partial text at or near the cursor as the beginning of something the user is actively typing. Complete the code the user appears to be creating based on context.
+- When completing partial code, prefer predictions that save meaningful keystrokes, even if this requires making educated guesses about the user's intent.
+- For code, it's better to make a substantive prediction that might be rejected than to make a minimal prediction that saves only a few keystrokes.
+- When the user is editing prose or documentation (e.g. Markdown, comments, plain text), predict conservatively. Complete the current fragment or sentence, but do not generate additional lines of free-form content since prose is less constrained than code and more prone to incorrect continuations.
+
+# Input Format
+
+You will be provided with:
+1. The user's *edit history*, in chronological order. Use this to infer the user's trajectory and predict the next most logical edit.
+  - Hunks preceded by `// User accepted prediction:` indicate code that was auto-generated by a previous prediction and accepted by the user. These are treated differently than user-typed edits (see Rules).
+2. A set of *related excerpts* from the user's codebase. Some of these may be needed for correctly predicting the next edit.
+  - `…` may appear within a related file to indicate that some code has been skipped.
+3. An excerpt from the user's *current file*.
+    - The excerpt contains numbered *marker* tags (`<|marker_1|>`, `<|marker_2|>`, etc.) placed at block boundaries throughout the code. These markers divide the excerpt into spans that you can target for editing.
+    - Code that appears before the first marker or after the last marker is read-only context and cannot be edited.
+    - The `<|user_cursor|>` tag marks the user's current cursor position, as it stands after the last edit in the history.
+
+# Output Format
+
+- Briefly explain the user's current intent based on the edit history and their current cursor location.
+- Output a markdown codeblock containing your predicted edit as a **marker-bounded span**:
+  - The codeblock must **start** with a marker tag (e.g. `<|marker_2|>`) and **end** with a marker tag (e.g. `<|marker_4|>`).
+  - The content between these two markers is the full replacement for that span in the original file.
+  - Choose the **narrowest** pair of markers that fully contains your predicted edits, to minimize unnecessary output.
+  - Reproduce any unchanged lines within the chosen span faithfully — do not omit or alter them.
+  - Do not include any intermediate marker tags in your output — only the start and end markers.
+- If no edit is needed (the code is already complete and correct, or there is no clear next edit to make), output a codeblock containing only `NO_EDITS`:
+  `````
+  NO_EDITS
+  `````
+- If there is a specific place in the predicted output where the user is likely to edit next, indicate it using the `<|user_cursor|>` tag.
+
+## Example 1
+
+There is code missing at the cursor location. The related excerpts includes the definition of a relevant type. You should fill in the missing code.
+
+### Related Excerpts
+
+`````
+struct Product {
+    name: String,
+    price: u32,
+}
+`````
+
+### User Edit History
+
+`````
+--- a/src/calculate.rs
++++ b/src/calculate.rs
+@@ -100,6 +100,7 @@
+ fn calculate_total(products: &[Product]) -> u32 {
+     let mut total = 0;
+     for product in products {
++        total += ;
+     }
+     total
+ }
+`````
+
+### Current File
+
+`````src/calculate.rs
+fn calculate_total(products: &[Product]) -> u32 {
+<|marker_1|>
+    let mut total = 0;
+    for product in products {
+        total += <|user_cursor|>;
+    }
+    total
+<|marker_2|>
+}
+`````
+
+### Output
+
+The user is computing a sum based on a list of products. The only numeric field on `Product` is `price`, so they must intend to sum the prices.
+
+`````
+<|marker_1|>
+    let mut total = 0;
+    for product in products {
+        total += product.price;
+    }
+    total
+<|marker_2|>
+`````
+
+## Example 2
+
+The user appears to be in the process of typing an eprintln call. Rather than fixing the spelling issue by deleting the newly-inserted content, you must continue the user's trajectory. It's not clear what data they intend to print. You should fill in as much code as is obviously intended, and position the cursor so that the user can fill in the rest.
+
+### User Edit History
+
+`````
+--- a/src/modal.rs
++++ b/src/modal.rs
+@@ -100,4 +100,4 @@
+ fn handle_close_button_click(modal_state: &mut ModalState, evt: &Event) {
+     modal_state.close();
+-     modal_state.dismiss();
++     eprmodal_state.dismiss();
+ }
+`````
+
+### Current File
+
+`````src/modal.rs
+<|marker_1|>
+// handle the close button click
+fn handle_close_button_click(modal_state: &mut ModalState, evt: &Event) {
+<|marker_2|>
+    modal_state.close();
+    epr<|user_cursor|>modal_state.dismiss();
+}
+<|marker_3|>
+`````
+
+### Output
+
+The user is clearly starting to type `eprintln!()`, however, what they intend to print is not obvious. I should fill in the print call and string literal, with the cursor positioned inside the string literal so the user can print whatever they want.
+
+`````
+<|marker_2|>
+    modal_state.close();
+    eprintln!("<|user_cursor|>");
+    modal_state.dismiss();
+}
+<|marker_3|>
+`````
+
+## Example 3
+
+Here, the user is adding a function. There's no way to tell for sure what the function's name will be. In this situation, you should make a reasonable guess at the function's name and signature, and place the user's cursor in the function body. This way, if you guess correctly, it will save the user a meaningful number of keystrokes, and the file will be left in a coherent state.
+
+### User Edit History
+
+`````
+--- a/src/modal.rs
++++ b/src/modal.rs
+@@ -100,4 +100,4 @@
+ fn handle_close_button_click(modal_state: &mut ModalState, evt: &Event) {
+     modal_state.close();
+     modal_state.dismiss();
+ }
++
++fn
+
+ fn handle_keystroke(modal_state: &mut ModalState, evt: &Event) {
+`````
+
+### Current File
+
+`````src/modal.rs
+// handle the close button click
+fn handle_close_button_click(modal_state: &mut ModalState, evt: &Event) {
+    modal_state.close();
+<|marker_1|>
+    modal_state.dismiss();
+}
+
+fn<|user_cursor|>
+
+<|marker_2|>
+fn handle_keystroke(modal_state: &mut ModalState, evt: &Event) {
+    modal_state.begin_edit();
+<|marker_3|>
+`````
+
+### Output
+
+The user is adding a new function. The existing functions I see are `handle_close_button_click` and `handle_keystroke`, which have similar signatures. One possible function they might be adding is `handle_submit`.
+
+`````
+<|marker_1|>
+    modal_state.dismiss();
+}
+
+fn handle_submit(modal_state: &mut ModalState, evt: &Event) {
+    <|user_cursor|>
+}
+
+<|marker_2|>
+`````
+
+## Example 4
+
+The code is already complete and there is no clear next edit to make. You should output NO_EDITS.
+
+### User Edit History
+
+`````
+--- a/src/utils.rs
++++ b/src/utils.rs
+@@ -10,7 +10,7 @@
+ fn add(a: i32, b: i32) -> i32 {
+-    a - b
++    a + b
+ }
+`````
+
+### Current File
+
+`````src/utils.rs
+<|marker_1|>
+fn add(a: i32, b: i32) -> i32 {
+    a + b<|user_cursor|>
+}
+<|marker_2|>
+`````
+
+### Output
+
+The user just fixed a bug in the `add` function, changing subtraction to addition. The code is now correct and complete. There is no clear next edit to make.
+
+`````
+NO_EDITS
+`````
+
+## Example 5
+
+The user just deleted code, leaving behind what looks incomplete. You must NOT "complete" it by restoring deleted content—that would undo their edit. Output NO_EDITS. **This is the correct response even though the code appears broken.**
+
+### User Edit History
+
+`````
+--- a/config.nix
++++ b/config.nix
+@@ -10,7 +10,7 @@
+     # /etc/modular/crashdb needs to be mutable
+-    ln -s /tmp/crashdb $out/etc/modular/crashdb
++    ln -s /tmp/cr $out/etc/modular/crashdb
+   '';
+`````
+
+### Current File
+
+`````config.nix
+<|marker_1|>
+    # /etc/modular/crashdb needs to be mutable
+    ln -s /tmp/cr<|user_cursor|> $out/etc/modular/crashdb
+  '';
+<|marker_2|>
+`````
+
+### Output
+
+The user deleted `ashdb` from `/tmp/crashdb`, leaving `/tmp/cr`. Although this looks like incomplete text that I could "complete", doing so would restore deleted content. The user intentionally removed that text—I must not undo their deletion.
+
+`````
+NO_EDITS
+`````
+
+## Example 6
+
+The user accepted a prediction for a function, then started renaming it. The original arguments were auto-generated (marked with `// User accepted prediction:`), so they CAN be updated to match the new function name. This is NOT reverting user input—it's improving auto-generated scaffolding.
+
+### User Edit History
+
+`````
+--- a/math_utils.py
++++ b/math_utils.py
+@@ -3,3 +3,5 @@
+ def calculate_rectangle_area(width, height):
+     return width * height
+
+
++de
+
+// User accepted prediction:
+--- a/math_utils.py
++++ b/math_utils.py
+@@ -3,5 +3,7 @@
+ def calculate_rectangle_area(width, height):
+     return width * height
+
+-de
++def calculate_rectangle_perimeter(width, height):
++
+
+--- a/math_utils.py
++++ b/math_utils.py
+@@ -5,5 +5,5 @@
+     return width * height
+
+-def calculate_rectangle_perimeter(width, height):
++def calculate_sq_perimeter(width, height):
+
+`````
+
+### Current File
+
+`````math_utils.py
+<|marker_1|>
+def calculate_rectangle_area(width, height):
+    return width * height
+
+<|marker_2|>
+def calculate_sq<|user_cursor|>_perimeter(width, height):
+
+<|marker_3|>
+`````
+
+### Output
+
+The user accepted a prediction for `calculate_rectangle_perimeter(width, height)`, then started renaming `rectangle` to `square`. Since squares have equal sides, the arguments should change from `(width, height)` to `(side)`. The arguments were auto-generated (from an accepted prediction), so modifying them is appropriate.
+
+`````
+<|marker_2|>
+def calculate_square_perimeter(side):
+    <|user_cursor|>
+<|marker_3|>
+`````
+
+
+
+# Your task:
+
+# 1. User Edit History
+
+`````
+{{edit_history}}
+`````
+
+# 2. Related excerpts
+
+{{context}}
+
+# 3. Current File
+
+{{cursor_excerpt}}
+
+
+
+
+-----
+
+Based on the edit history and context above, predict the user's next edit within the marker-bounded spans.

crates/edit_prediction_cli/src/qa.rs

@@ -0,0 +1,301 @@
+//! Quality assessment of predictions using LLM-as-a-judge.
+//!
+//! This module uses LLM Batch APIs to evaluate prediction quality.
+//! Caching is handled by the underlying client.
+
+use crate::{
+    BatchProvider,
+    anthropic_client::AnthropicClient,
+    example::Example,
+    format_prompt::extract_cursor_excerpt_from_example,
+    openai_client::OpenAiClient,
+    parse_output::run_parse_output,
+    paths::LLM_CACHE_DB,
+    progress::{ExampleProgress, Step},
+    word_diff::unified_to_word_diff,
+};
+use anyhow::{Context as _, Result};
+use serde::{Deserialize, Serialize};
+use std::sync::OnceLock;
+
+/// Arguments for the QA command.
+#[derive(Debug, Clone, clap::Args)]
+pub struct QaArgs {
+    /// Use synchronous API instead of batch
+    #[clap(long)]
+    pub no_batch: bool,
+
+    /// Which LLM provider to use (anthropic or openai)
+    #[clap(long, default_value = "openai")]
+    pub backend: BatchProvider,
+}
+
+fn model_for_backend(backend: BatchProvider) -> &'static str {
+    match backend {
+        BatchProvider::Anthropic => "claude-sonnet-4-5",
+        BatchProvider::Openai => "gpt-5.2",
+    }
+}
+
+/// Result of QA evaluation for a single prediction.
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct QaResult {
+    /// Free-form reasoning from the judge.
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub reasoning: Option<String>,
+
+    /// Does the prediction undo/revert changes the user intentionally made?
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub reverts_edits: Option<bool>,
+
+    /// Confidence score (1-5) for user acceptance likelihood.
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub confidence: Option<u8>,
+
+    /// The raw response from the model.
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub response: Option<String>,
+
+    /// Error message if parsing or request failed.
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub error: Option<String>,
+}
+
+/// Build the assessment prompt for an example.
+pub fn build_prompt(example: &Example) -> Result<String> {
+    let prediction = example
+        .predictions
+        .first()
+        .context("no predictions available")?;
+    let actual_patch = prediction
+        .actual_patch
+        .as_ref()
+        .context("no actual_patch available (run predict first)")?;
+    let prompt_inputs = example
+        .prompt_inputs
+        .as_ref()
+        .context("prompt_inputs missing (run context retrieval first)")?;
+
+    let actual_patch_word_diff = unified_to_word_diff(actual_patch);
+
+    let cursor_excerpt =
+        extract_cursor_excerpt_from_example(example).context("failed to extract cursor excerpt")?;
+
+    let mut edit_history = String::new();
+    for event in &prompt_inputs.events {
+        let zeta_prompt::Event::BufferChange {
+            path,
+            old_path,
+            diff,
+            predicted: _,
+            in_open_source_repo: _,
+        } = event.as_ref();
+        edit_history.push_str(&format!("--- a{}\n", old_path.display()));
+        edit_history.push_str(&format!("+++ b{}\n", path.display()));
+        let diff_word_diff = unified_to_word_diff(&diff);
+        edit_history.push_str(&diff_word_diff);
+        edit_history.push_str("\n\n");
+    }
+
+    let prompt_template = crate::prompt_assets::get_prompt("qa.md");
+    Ok(prompt_template
+        .replace("{edit_history}", &edit_history)
+        .replace("{cursor_excerpt}", &cursor_excerpt)
+        .replace("{actual_patch_word_diff}", &actual_patch_word_diff))
+}
+
+fn extract_codeblock(response: &str) -> Option<String> {
+    let lines: Vec<&str> = response.lines().collect();
+    for (i, line) in lines.iter().enumerate() {
+        if line.starts_with("```") {
+            let start = i + 1;
+            for (j, end_line) in lines[start..].iter().enumerate() {
+                if end_line.starts_with("```") {
+                    return Some(lines[start..start + j].join("\n"));
+                }
+            }
+            return Some(lines[start..].join("\n"));
+        }
+    }
+    None
+}
+
+fn parse_response(response_text: &str) -> QaResult {
+    let codeblock = extract_codeblock(response_text);
+
+    for text_to_parse in [codeblock.as_deref(), Some(response_text.trim())] {
+        let Some(text) = text_to_parse else {
+            continue;
+        };
+
+        if let Ok(parsed) = serde_json::from_str::<serde_json::Value>(text) {
+            return QaResult {
+                reasoning: parsed
+                    .get("reasoning")
+                    .and_then(|v| v.as_str())
+                    .map(|s| s.to_string()),
+                reverts_edits: parsed.get("reverts_edits").and_then(|v| v.as_bool()),
+                confidence: parsed
+                    .get("confidence")
+                    .and_then(|v| v.as_u64())
+                    .map(|v| v as u8),
+                response: Some(response_text.to_string()),
+                error: None,
+            };
+        }
+    }
+
+    QaResult {
+        reasoning: Some(response_text.to_string()),
+        reverts_edits: None,
+        confidence: None,
+        response: Some(response_text.to_string()),
+        error: Some("Could not parse JSON from response".to_string()),
+    }
+}
+
+static ANTHROPIC_CLIENT_BATCH: OnceLock<AnthropicClient> = OnceLock::new();
+static ANTHROPIC_CLIENT_PLAIN: OnceLock<AnthropicClient> = OnceLock::new();
+static OPENAI_CLIENT_BATCH: OnceLock<OpenAiClient> = OnceLock::new();
+static OPENAI_CLIENT_PLAIN: OnceLock<OpenAiClient> = OnceLock::new();
+
+/// Run QA evaluation for a single example.
+pub async fn run_qa(
+    example: &mut Example,
+    args: &QaArgs,
+    example_progress: &ExampleProgress,
+) -> Result<()> {
+    if example
+        .qa
+        .first()
+        .and_then(|q| q.as_ref())
+        .and_then(|q| q.confidence)
+        .is_some()
+    {
+        return Ok(());
+    }
+
+    run_parse_output(example).context("Failed to execute run_parse_output")?;
+
+    if example.prompt_inputs.is_none() {
+        anyhow::bail!("prompt_inputs missing (run context retrieval first)");
+    }
+
+    let step_progress = example_progress.start(Step::Qa);
+
+    let model = model_for_backend(args.backend);
+    let prompt = build_prompt(example).context("Failed to build QA prompt")?;
+
+    step_progress.set_substatus("generating");
+
+    let response = match args.backend {
+        BatchProvider::Anthropic => {
+            let client = if args.no_batch {
+                ANTHROPIC_CLIENT_PLAIN.get_or_init(|| {
+                    AnthropicClient::plain().expect("Failed to create Anthropic client")
+                })
+            } else {
+                ANTHROPIC_CLIENT_BATCH.get_or_init(|| {
+                    AnthropicClient::batch(&LLM_CACHE_DB)
+                        .expect("Failed to create Anthropic client")
+                })
+            };
+
+            let messages = vec![anthropic::Message {
+                role: anthropic::Role::User,
+                content: vec![anthropic::RequestContent::Text {
+                    text: prompt,
+                    cache_control: None,
+                }],
+            }];
+
+            let Some(response) = client.generate(model, 1024, messages, None, false).await? else {
+                return Ok(());
+            };
+
+            response
+                .content
+                .iter()
+                .filter_map(|c| match c {
+                    anthropic::ResponseContent::Text { text } => Some(text.as_str()),
+                    _ => None,
+                })
+                .collect::<Vec<_>>()
+                .join("")
+        }
+        BatchProvider::Openai => {
+            let client = if args.no_batch {
+                OPENAI_CLIENT_PLAIN
+                    .get_or_init(|| OpenAiClient::plain().expect("Failed to create OpenAI client"))
+            } else {
+                OPENAI_CLIENT_BATCH.get_or_init(|| {
+                    OpenAiClient::batch(&LLM_CACHE_DB).expect("Failed to create OpenAI client")
+                })
+            };
+
+            let messages = vec![open_ai::RequestMessage::User {
+                content: open_ai::MessageContent::Plain(prompt),
+            }];
+
+            let Some(response) = client.generate(model, 1024, messages, None, false).await? else {
+                return Ok(());
+            };
+
+            response
+                .choices
+                .into_iter()
+                .filter_map(|choice| match choice.message {
+                    open_ai::RequestMessage::Assistant { content, .. } => {
+                        content.map(|c| match c {
+                            open_ai::MessageContent::Plain(text) => text,
+                            open_ai::MessageContent::Multipart(parts) => parts
+                                .into_iter()
+                                .filter_map(|p| match p {
+                                    open_ai::MessagePart::Text { text } => Some(text),
+                                    _ => None,
+                                })
+                                .collect::<Vec<_>>()
+                                .join(""),
+                        })
+                    }
+                    _ => None,
+                })
+                .collect::<Vec<_>>()
+                .join("")
+        }
+    };
+
+    let result = parse_response(&response);
+
+    example.qa = example
+        .predictions
+        .iter()
+        .enumerate()
+        .map(|(i, _)| if i == 0 { Some(result.clone()) } else { None })
+        .collect();
+
+    Ok(())
+}
+
+/// Sync batches for QA (upload pending requests, download finished results).
+pub async fn sync_batches(args: &QaArgs) -> Result<()> {
+    if args.no_batch {
+        return Ok(());
+    }
+
+    match args.backend {
+        BatchProvider::Anthropic => {
+            let client = ANTHROPIC_CLIENT_BATCH.get_or_init(|| {
+                AnthropicClient::batch(&LLM_CACHE_DB).expect("Failed to create Anthropic client")
+            });
+            client.sync_batches().await?;
+        }
+        BatchProvider::Openai => {
+            let client = OPENAI_CLIENT_BATCH.get_or_init(|| {
+                OpenAiClient::batch(&LLM_CACHE_DB).expect("Failed to create OpenAI client")
+            });
+            client.sync_batches().await?;
+        }
+    }
+    Ok(())
+}

crates/edit_prediction_cli/src/repair.rs

@@ -0,0 +1,605 @@
+//! Repair predictions that received poor quality signals.
+//!
+//! This module takes examples with predictions, identifies predictions that need
+//! improvement, and uses an LLM to generate improved predictions. It supports
+//! two sources of quality signals:
+//! - QA feedback (reverts_edits or low confidence)
+//! - Computed scores when QA is unavailable (high reversal_ratio or wrong_editable_region)
+
+use crate::{
+    BatchProvider, PredictionProvider,
+    anthropic_client::AnthropicClient,
+    example::{ActualCursor, Example, ExamplePrediction},
+    format_prompt::TeacherPrompt,
+    metrics::count_patch_token_changes,
+    openai_client::OpenAiClient,
+    parse_output::run_parse_output,
+    paths::LLM_CACHE_DB,
+    progress::{ExampleProgress, Progress, Step},
+    word_diff::unified_to_word_diff,
+};
+use anyhow::{Context as _, Result};
+use std::sync::OnceLock;
+
+const KEEP_PREVIOUS: &str = "KEEP_PREVIOUS";
+
+/// Print a summary report of repair results across all examples.
+pub fn print_report(examples: &[Example], confidence_threshold: u8) {
+    let total = examples.len();
+    let mut no_repair_needed = 0;
+    let mut repaired = 0;
+    let mut repair_failed = 0;
+
+    for example in examples {
+        if !needs_repair(example, confidence_threshold) {
+            no_repair_needed += 1;
+            continue;
+        }
+
+        if has_successful_repair(example) {
+            repaired += 1;
+        } else {
+            repair_failed += 1;
+        }
+    }
+
+    let needed_repair = total - no_repair_needed;
+
+    eprintln!();
+    eprintln!("Repair summary ({total} examples):");
+    eprintln!(
+        "  {no_repair_needed}/{total} didn't need repair (confidence > {confidence_threshold})"
+    );
+    if needed_repair > 0 {
+        eprintln!("  {needed_repair}/{total} needed repair:");
+        if repaired > 0 {
+            eprintln!("    {repaired} repaired successfully");
+        }
+        if repair_failed > 0 {
+            eprintln!("    {repair_failed} failed to repair");
+        }
+    }
+}
+
+/// Arguments for the repair command.
+#[derive(Debug, Clone, clap::Args)]
+pub struct RepairArgs {
+    /// Use synchronous API instead of batch
+    #[clap(long)]
+    pub no_batch: bool,
+
+    /// Confidence threshold: repair predictions with confidence <= this value (1-5)
+    #[clap(long, default_value = "2")]
+    pub confidence_threshold: u8,
+
+    /// Which LLM provider to use (anthropic or openai)
+    #[clap(long, default_value = "anthropic")]
+    pub backend: BatchProvider,
+    /// Wait for all batches to complete before exiting
+    #[clap(long)]
+    pub wait: bool,
+}
+
+fn model_for_backend(backend: BatchProvider) -> &'static str {
+    match backend {
+        BatchProvider::Anthropic => "claude-sonnet-4-6",
+        BatchProvider::Openai => "gpt-5.2",
+    }
+}
+
+/// Build the quality feedback string from QA results.
+fn build_qa_feedback(example: &Example) -> Option<String> {
+    let qa = example.qa.first()?.as_ref()?;
+
+    let qa_reasoning = qa.reasoning.as_deref().unwrap_or("No reasoning provided");
+    let reverts_edits = qa
+        .reverts_edits
+        .map_or("unknown", |v| if v { "yes" } else { "no" });
+    let confidence = qa
+        .confidence
+        .map_or("unknown".to_string(), |v| v.to_string());
+
+    Some(format!(
+        "- **Reverts user edits**: {reverts_edits}\n\
+         - **Confidence score**: {confidence}/5\n\
+         - **Reasoning**: {qa_reasoning}"
+    ))
+}
+
+/// Build the quality feedback string from computed scores when QA is unavailable.
+fn build_score_feedback(example: &Example) -> Option<String> {
+    let score = example.score.first()?;
+
+    let mut issues = Vec::new();
+
+    if score.reversal_ratio > 0.9 {
+        issues.push(format!(
+            "Automated analysis detected a high reversal ratio ({:.2}), which suggests this \
+             prediction may be reverting changes the user intentionally made. Double-check that \
+             the prediction doesn't undo the user's recent edits. If the prediction is actually \
+             fine and the edits are intentional completions rather than reversals, keep it as-is. \
+             If it truly reverts the user's changes, generate an improved prediction that \
+             continues the user's intent instead.",
+            score.reversal_ratio
+        ));
+    }
+
+    if score.wrong_editable_region == Some(true) {
+        issues.push(
+            "Automated analysis detected that the prediction may be modifying code outside \
+             the expected editable region, or producing changes misaligned with the editable \
+             region boundaries. Make sure the prediction only modifies code within the editable \
+             region and is properly aligned."
+                .to_string(),
+        );
+    }
+
+    if score.discarded_chars.unwrap_or(0) > 80 && score.exact_lines_fp > 5 {
+        issues.push(
+            "Automated analysis detected that this prediction might be too large or speculative. \
+            Please review it and think if we should keep it or generate a more focused prediction. \
+            Examples of more focused predictions: \
+            - Predicting a function outline but not its body. \
+            - Predicting only the first logical step and not speculating about further steps.
+            In general, the smaller the prediction you make, the higher the chance it will be correct."
+                .to_string(),
+        );
+    }
+
+    if issues.is_empty() {
+        return None;
+    }
+
+    let mut feedback = String::from(
+        "No human quality assessment is available, but automated scoring flagged potential issues:\n\n",
+    );
+    for issue in &issues {
+        feedback.push_str(&format!("- {issue}\n"));
+    }
+    feedback.push_str(
+        "\nRemember: if the previous prediction was actually correct, output `KEEP_PREVIOUS`. \
+         If no edits should be made at all and you are unsure how to improve it, output `NO_EDITS`.",
+    );
+
+    Some(feedback)
+}
+
+/// Build the repair message (Turn 3) for a multi-turn conversation.
+///
+/// This message is sent after the original teacher prompt (Turn 1) and
+/// teacher response (Turn 2) to request an improved prediction.
+pub fn build_repair_message(example: &Example) -> Result<String> {
+    let prediction = example
+        .predictions
+        .first()
+        .context("no predictions available")?;
+    let actual_patch = prediction
+        .actual_patch
+        .as_ref()
+        .context("no actual_patch available (run predict first)")?;
+
+    let quality_feedback = build_qa_feedback(example)
+        .or_else(|| build_score_feedback(example))
+        .context("no quality feedback available (need either QA results or computed scores)")?;
+
+    let actual_patch_word_diff = unified_to_word_diff(actual_patch);
+
+    let token_counts = count_patch_token_changes(actual_patch);
+    let mut token_change_info = format!(
+        "\n## Token Change Statistics\n\n\
+         - **Deleted tokens**: {}\n\
+         - **Inserted tokens**: {}",
+        token_counts.deleted_tokens, token_counts.inserted_tokens,
+    );
+    if token_counts.deleted_tokens > 100 || token_counts.inserted_tokens > 100 {
+        token_change_info.push_str(
+            "\n\n> **Note:** The token change count is high. \
+             Consider producing a more scoped edit that targets only the lines \
+             that truly need to change, rather than rewriting large sections.",
+        );
+    }
+
+    let prompt_template = crate::prompt_assets::get_prompt("repair.md");
+    Ok(prompt_template
+        .replace("{actual_patch_word_diff}", &actual_patch_word_diff)
+        .replace("{quality_feedback}", &quality_feedback)
+        .replace("{token_change_info}", &token_change_info))
+}
+
+/// Check if an example needs repair based on QA feedback or computed scores.
+pub fn needs_repair(example: &Example, confidence_threshold: u8) -> bool {
+    // Check QA-based signals first.
+    if let Some(qa) = example.qa.first().and_then(|q| q.as_ref()) {
+        if qa.reverts_edits == Some(true) {
+            return true;
+        }
+
+        if let Some(confidence) = qa.confidence {
+            if confidence <= confidence_threshold {
+                return true;
+            }
+        }
+
+        return false;
+    }
+
+    // When QA is unavailable, fall back to computed score signals.
+    if let Some(score) = example.score.first() {
+        if score.reversal_ratio > 0.9 {
+            return true;
+        }
+
+        if score.wrong_editable_region == Some(true) {
+            return true;
+        }
+    }
+
+    false
+}
+
+/// Parse repair model output into a patch and optional cursor.
+///
+/// Handles the `KEEP_PREVIOUS` sentinel by copying the teacher's prediction,
+/// and delegates normal output to `TeacherPrompt::parse`.
+pub fn parse(example: &Example, actual_output: &str) -> Result<(String, Option<ActualCursor>)> {
+    if actual_output.contains(KEEP_PREVIOUS) {
+        let original = example
+            .predictions
+            .first()
+            .context("no original prediction to keep")?;
+        let patch = original.actual_patch.clone().unwrap_or_default();
+        let cursor = original.actual_cursor.clone();
+        return Ok((patch, cursor));
+    }
+
+    TeacherPrompt::parse(example, actual_output)
+}
+
+/// Check if an example already has a successful repair prediction.
+fn has_successful_repair(example: &Example) -> bool {
+    example
+        .predictions
+        .iter()
+        .any(|p| p.provider == PredictionProvider::Repair && p.actual_patch.is_some())
+}
+
+static ANTHROPIC_CLIENT_BATCH: OnceLock<AnthropicClient> = OnceLock::new();
+static ANTHROPIC_CLIENT_PLAIN: OnceLock<AnthropicClient> = OnceLock::new();
+static OPENAI_CLIENT_BATCH: OnceLock<OpenAiClient> = OnceLock::new();
+static OPENAI_CLIENT_PLAIN: OnceLock<OpenAiClient> = OnceLock::new();
+
+/// Run repair for a single example.
+///
+/// This sends a multi-turn conversation to the LLM:
+/// - Turn 1 (User): Original teacher prompt
+/// - Turn 2 (Assistant): Original teacher response
+/// - Turn 3 (User): Repair critique and instructions
+/// - Turn 4 (Assistant): Improved prediction (the response we parse)
+pub async fn run_repair(
+    example: &mut Example,
+    args: &RepairArgs,
+    example_progress: &ExampleProgress,
+) -> Result<()> {
+    if has_successful_repair(example) {
+        return Ok(());
+    }
+
+    if !needs_repair(example, args.confidence_threshold) {
+        return Ok(());
+    }
+
+    run_parse_output(example).context("Failed to execute run_parse_output")?;
+
+    if example.prompt_inputs.is_none() {
+        anyhow::bail!("prompt_inputs missing (run context retrieval first)");
+    }
+
+    if example.predictions.is_empty() {
+        anyhow::bail!("no predictions available (run predict first)");
+    }
+
+    let teacher_prompt = example
+        .prompt
+        .as_ref()
+        .context("prompt missing (run format_prompt first)")?;
+
+    let teacher_response = &example.predictions[0].actual_output;
+    if teacher_response.is_empty() {
+        anyhow::bail!("teacher response is empty (run predict first)");
+    }
+
+    let step_progress = example_progress.start(Step::Repair);
+
+    let model = model_for_backend(args.backend);
+    let repair_message = build_repair_message(example).context("Failed to build repair message")?;
+
+    step_progress.set_substatus("generating");
+
+    let response = match args.backend {
+        BatchProvider::Anthropic => {
+            let client = if args.no_batch {
+                ANTHROPIC_CLIENT_PLAIN.get_or_init(|| {
+                    AnthropicClient::plain().expect("Failed to create Anthropic client")
+                })
+            } else {
+                ANTHROPIC_CLIENT_BATCH.get_or_init(|| {
+                    AnthropicClient::batch(&LLM_CACHE_DB)
+                        .expect("Failed to create Anthropic client")
+                })
+            };
+
+            let messages = vec![
+                // Turn 1: Original teacher prompt
+                anthropic::Message {
+                    role: anthropic::Role::User,
+                    content: vec![anthropic::RequestContent::Text {
+                        text: teacher_prompt.input.clone(),
+                        cache_control: None,
+                    }],
+                },
+                // Turn 2: Original teacher response
+                anthropic::Message {
+                    role: anthropic::Role::Assistant,
+                    content: vec![anthropic::RequestContent::Text {
+                        text: teacher_response.clone(),
+                        cache_control: None,
+                    }],
+                },
+                // Turn 3: Repair critique and instructions
+                anthropic::Message {
+                    role: anthropic::Role::User,
+                    content: vec![anthropic::RequestContent::Text {
+                        text: repair_message,
+                        cache_control: None,
+                    }],
+                },
+            ];
+
+            let Some(response) = client.generate(model, 16384, messages, None, false).await? else {
+                return Ok(());
+            };
+
+            response
+                .content
+                .iter()
+                .filter_map(|c| match c {
+                    anthropic::ResponseContent::Text { text } => Some(text.as_str()),
+                    _ => None,
+                })
+                .collect::<Vec<_>>()
+                .join("")
+        }
+        BatchProvider::Openai => {
+            let client = if args.no_batch {
+                OPENAI_CLIENT_PLAIN
+                    .get_or_init(|| OpenAiClient::plain().expect("Failed to create OpenAI client"))
+            } else {
+                OPENAI_CLIENT_BATCH.get_or_init(|| {
+                    OpenAiClient::batch(&LLM_CACHE_DB).expect("Failed to create OpenAI client")
+                })
+            };
+
+            let messages = vec![
+                // Turn 1: Original teacher prompt
+                open_ai::RequestMessage::User {
+                    content: open_ai::MessageContent::Plain(teacher_prompt.input.clone()),
+                },
+                // Turn 2: Original teacher response
+                open_ai::RequestMessage::Assistant {
+                    content: Some(open_ai::MessageContent::Plain(teacher_response.clone())),
+                    tool_calls: vec![],
+                    reasoning_content: None,
+                },
+                // Turn 3: Repair critique and instructions
+                open_ai::RequestMessage::User {
+                    content: open_ai::MessageContent::Plain(repair_message),
+                },
+            ];
+
+            let Some(response) = client.generate(model, 16384, messages, None, false).await? else {
+                return Ok(());
+            };
+
+            response
+                .choices
+                .into_iter()
+                .filter_map(|choice| match choice.message {
+                    open_ai::RequestMessage::Assistant { content, .. } => {
+                        content.map(|c| match c {
+                            open_ai::MessageContent::Plain(text) => text,
+                            open_ai::MessageContent::Multipart(parts) => parts
+                                .into_iter()
+                                .filter_map(|p| match p {
+                                    open_ai::MessagePart::Text { text } => Some(text),
+                                    _ => None,
+                                })
+                                .collect::<Vec<_>>()
+                                .join(""),
+                        })
+                    }
+                    _ => None,
+                })
+                .collect::<Vec<_>>()
+                .join("")
+        }
+    };
+
+    let parse_result = parse(example, &response);
+    let err = parse_result
+        .as_ref()
+        .err()
+        .map(|e| format!("Failed to parse repair response: {}", e));
+
+    let (actual_patch, actual_cursor) = parse_result.ok().unzip();
+    let actual_cursor = actual_cursor.flatten();
+
+    example.predictions.push(ExamplePrediction {
+        actual_patch,
+        actual_output: response,
+        actual_cursor,
+        error: err,
+        provider: PredictionProvider::Repair,
+        cumulative_logprob: None,
+        avg_logprob: None,
+    });
+
+    Ok(())
+}
+
+/// Sync batches for repair (upload pending requests, download finished results).
+pub async fn sync_batches(args: &RepairArgs) -> Result<()> {
+    if args.no_batch {
+        return Ok(());
+    }
+
+    match args.backend {
+        BatchProvider::Anthropic => {
+            let client = ANTHROPIC_CLIENT_BATCH.get_or_init(|| {
+                AnthropicClient::batch(&LLM_CACHE_DB).expect("Failed to create Anthropic client")
+            });
+            client.sync_batches().await?;
+        }
+        BatchProvider::Openai => {
+            let client = OPENAI_CLIENT_BATCH.get_or_init(|| {
+                OpenAiClient::batch(&LLM_CACHE_DB).expect("Failed to create OpenAI client")
+            });
+            client.sync_batches().await?;
+        }
+    }
+
+    Ok(())
+}
+
+pub async fn reprocess_after_batch_wait(examples: &mut [Example], args: &RepairArgs) -> Result<()> {
+    let mut reprocessed = 0;
+    for example in examples.iter_mut() {
+        if has_successful_repair(example) || !needs_repair(example, args.confidence_threshold) {
+            continue;
+        }
+
+        let example_progress = Progress::global().start_group(&example.spec.name);
+        run_repair(example, args, &example_progress).await?;
+        reprocessed += 1;
+    }
+
+    if reprocessed > 0 {
+        eprintln!("Reprocessed {} example(s) with batch results", reprocessed);
+    }
+
+    Ok(())
+}
+
+pub async fn wait_for_batches(args: &RepairArgs) -> Result<()> {
+    if args.no_batch {
+        return Ok(());
+    }
+
+    let poll_interval = std::time::Duration::from_secs(30);
+
+    loop {
+        let pending = pending_batch_count(args)?;
+        if pending == 0 {
+            break;
+        }
+
+        eprintln!(
+            "Waiting for {} pending repair batch request(s) to complete... (polling every {}s)",
+            pending,
+            poll_interval.as_secs()
+        );
+        std::thread::sleep(poll_interval);
+
+        sync_batches(args).await?;
+    }
+
+    Ok(())
+}
+
+fn pending_batch_count(args: &RepairArgs) -> Result<usize> {
+    match args.backend {
+        BatchProvider::Anthropic => {
+            let client = ANTHROPIC_CLIENT_BATCH.get_or_init(|| {
+                AnthropicClient::batch(&LLM_CACHE_DB).expect("Failed to create Anthropic client")
+            });
+            client.pending_batch_count()
+        }
+        BatchProvider::Openai => {
+            let client = OPENAI_CLIENT_BATCH.get_or_init(|| {
+                OpenAiClient::batch(&LLM_CACHE_DB).expect("Failed to create OpenAI client")
+            });
+            client.pending_batch_count()
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::{PredictionProvider, TeacherBackend};
+    use edit_prediction::example_spec::ExampleSpec;
+    use std::{path::Path, sync::Arc};
+    use zeta_prompt::ZetaFormat;
+
+    fn example_with_previous_prediction() -> Example {
+        Example {
+            spec: ExampleSpec {
+                name: "example".to_string(),
+                repository_url: "https://github.com/zed-industries/zed.git".to_string(),
+                revision: "HEAD".to_string(),
+                tags: Vec::new(),
+                reasoning: None,
+                uncommitted_diff: String::new(),
+                cursor_path: Arc::from(Path::new("src/main.rs")),
+                cursor_position: "0:0".to_string(),
+                edit_history: String::new(),
+                expected_patches: Vec::new(),
+                rejected_patch: None,
+                telemetry: None,
+                human_feedback: Vec::new(),
+                rating: None,
+            },
+            prompt_inputs: None,
+            prompt: None,
+            predictions: vec![ExamplePrediction {
+                actual_patch: Some("previous patch".to_string()),
+                actual_output: String::new(),
+                actual_cursor: Some(ActualCursor {
+                    path: "src/main.rs".to_string(),
+                    row: 1,
+                    column: 2,
+                    offset: 3,
+                    editable_region_offset: Some(4),
+                }),
+                error: None,
+                provider: PredictionProvider::Teacher(
+                    TeacherBackend::Sonnet45,
+                    ZetaFormat::default(),
+                ),
+                cumulative_logprob: None,
+                avg_logprob: None,
+            }],
+            score: Vec::new(),
+            qa: Vec::new(),
+            zed_version: None,
+            state: None,
+        }
+    }
+
+    #[test]
+    fn test_parse_keeps_previous_when_sentinel_appears_outside_last_codeblock() {
+        let example = example_with_previous_prediction();
+        let actual_output = indoc::indoc! {"
+            After reviewing the feedback, the previous prediction is still correct.
+            Use `KEEP_PREVIOUS`.
+
+            ```
+            unrelated trailing code block
+            ```
+        "};
+
+        let (patch, cursor) = parse(&example, actual_output).unwrap();
+
+        assert_eq!(patch, "previous patch");
+        assert_eq!(cursor.unwrap().offset, 3);
+    }
+}

crates/edit_prediction_cli/src/retrieve_context.rs

@@ -0,0 +1,186 @@
+use crate::{
+    example::Example,
+    headless::EpAppState,
+    load_project::run_load_project,
+    progress::{ExampleProgress, InfoStyle, Step, StepProgress},
+};
+use anyhow::Context as _;
+use collections::HashSet;
+use edit_prediction::{DebugEvent, EditPredictionStore};
+use futures::{FutureExt as _, StreamExt as _, channel::mpsc};
+use gpui::{AsyncApp, Entity};
+use language::Buffer;
+use project::Project;
+use std::sync::Arc;
+use std::time::Duration;
+
+pub async fn run_context_retrieval(
+    example: &mut Example,
+    app_state: Arc<EpAppState>,
+    example_progress: &ExampleProgress,
+    mut cx: AsyncApp,
+) -> anyhow::Result<()> {
+    if example
+        .prompt_inputs
+        .as_ref()
+        .is_some_and(|inputs| inputs.related_files.is_some())
+        || example.spec.repository_url.is_empty()
+    {
+        return Ok(());
+    }
+
+    run_load_project(example, app_state.clone(), example_progress, cx.clone()).await?;
+
+    let step_progress: Arc<StepProgress> = example_progress.start(Step::Context).into();
+
+    let state = example.state.as_ref().unwrap();
+    let project = state.project.clone();
+
+    let _lsp_handle = project.update(&mut cx, |project, cx| {
+        project.register_buffer_with_language_servers(&state.buffer, cx)
+    });
+    wait_for_language_servers_to_start(&project, &state.buffer, &step_progress, &mut cx).await?;
+
+    let ep_store = cx
+        .update(|cx| EditPredictionStore::try_global(cx))
+        .context("EditPredictionStore not initialized")?;
+
+    let mut events = ep_store.update(&mut cx, |store, cx| {
+        store.register_buffer(&state.buffer, &project, cx);
+        store.refresh_context(&project, &state.buffer, state.cursor_position, cx);
+        store.debug_info(&project, cx)
+    });
+
+    while let Some(event) = events.next().await {
+        match event {
+            DebugEvent::ContextRetrievalFinished(_) => {
+                break;
+            }
+            _ => {}
+        }
+    }
+
+    let context_files =
+        ep_store.update(&mut cx, |store, cx| store.context_for_project(&project, cx));
+
+    let excerpt_count: usize = context_files.iter().map(|f| f.excerpts.len()).sum();
+    step_progress.set_info(format!("{} excerpts", excerpt_count), InfoStyle::Normal);
+
+    if let Some(prompt_inputs) = example.prompt_inputs.as_mut() {
+        prompt_inputs.related_files = Some(context_files);
+    }
+    Ok(())
+}
+
+async fn wait_for_language_servers_to_start(
+    project: &Entity<Project>,
+    buffer: &Entity<Buffer>,
+    step_progress: &Arc<StepProgress>,
+    cx: &mut AsyncApp,
+) -> anyhow::Result<()> {
+    let lsp_store = project.read_with(cx, |project, _| project.lsp_store());
+
+    // Determine which servers exist for this buffer, and which are still starting.
+    let mut servers_pending_start = HashSet::default();
+    let mut servers_pending_diagnostics = HashSet::default();
+    buffer.update(cx, |buffer, cx| {
+        lsp_store.update(cx, |lsp_store, cx| {
+            let ids = lsp_store.language_servers_for_local_buffer(buffer, cx);
+            for &id in &ids {
+                match lsp_store.language_server_statuses.get(&id) {
+                    None => {
+                        servers_pending_start.insert(id);
+                        servers_pending_diagnostics.insert(id);
+                    }
+                    Some(status) if status.has_pending_diagnostic_updates => {
+                        servers_pending_diagnostics.insert(id);
+                    }
+                    Some(_) => {}
+                }
+            }
+        });
+    });
+
+    step_progress.set_substatus(format!(
+        "waiting for {} LSPs",
+        servers_pending_diagnostics.len()
+    ));
+
+    let timeout_duration = if servers_pending_start.is_empty() {
+        Duration::from_secs(30)
+    } else {
+        Duration::from_secs(60 * 5)
+    };
+    let timeout = cx.background_executor().timer(timeout_duration).shared();
+
+    let (mut started_tx, mut started_rx) = mpsc::channel(servers_pending_start.len().max(1));
+    let (mut diag_tx, mut diag_rx) = mpsc::channel(servers_pending_diagnostics.len().max(1));
+    let subscriptions = [cx.subscribe(&lsp_store, {
+        let step_progress = step_progress.clone();
+        move |lsp_store, event, cx| match event {
+            project::LspStoreEvent::LanguageServerAdded(id, name, _) => {
+                step_progress.set_substatus(format!("LSP started: {}", name));
+                started_tx.try_send(*id).ok();
+            }
+            project::LspStoreEvent::DiskBasedDiagnosticsFinished { language_server_id } => {
+                let name = lsp_store
+                    .read(cx)
+                    .language_server_adapter_for_id(*language_server_id)
+                    .unwrap()
+                    .name();
+                step_progress.set_substatus(format!("LSP idle: {}", name));
+                diag_tx.try_send(*language_server_id).ok();
+            }
+            project::LspStoreEvent::LanguageServerUpdate {
+                message:
+                    client::proto::update_language_server::Variant::WorkProgress(
+                        client::proto::LspWorkProgress {
+                            message: Some(message),
+                            ..
+                        },
+                    ),
+                ..
+            } => {
+                step_progress.set_substatus(message.clone());
+            }
+            _ => {}
+        }
+    })];
+
+    // Phase 1: wait for all servers to start.
+    while !servers_pending_start.is_empty() {
+        futures::select! {
+            id = started_rx.next() => {
+                if let Some(id) = id {
+                    servers_pending_start.remove(&id);
+                }
+            },
+            _ = timeout.clone().fuse() => {
+                return Err(anyhow::anyhow!("LSP wait timed out after {} minutes", timeout_duration.as_secs() / 60));
+            }
+        }
+    }
+
+    // Save the buffer so the server sees the current content and kicks off diagnostics.
+    project
+        .update(cx, |project, cx| project.save_buffer(buffer.clone(), cx))
+        .await?;
+
+    // Phase 2: wait for all servers to finish their diagnostic pass.
+    while !servers_pending_diagnostics.is_empty() {
+        futures::select! {
+            id = diag_rx.next() => {
+                if let Some(id) = id {
+                    servers_pending_diagnostics.remove(&id);
+                }
+            },
+            _ = timeout.clone().fuse() => {
+                return Err(anyhow::anyhow!("LSP wait timed out after {} minutes", timeout_duration.as_secs() / 60));
+            }
+        }
+    }
+
+    drop(subscriptions);
+    step_progress.clear_substatus();
+    Ok(())
+}

crates/edit_prediction_cli/src/score.rs

@@ -0,0 +1,521 @@
+use crate::{
+    PredictArgs, PredictionProvider,
+    example::Example,
+    format_prompt::TeacherPrompt,
+    headless::EpAppState,
+    parse_output::parse_prediction_output,
+    predict::run_prediction,
+    progress::{ExampleProgress, Step},
+};
+use anyhow::Context as _;
+use edit_prediction_metrics::{
+    ActualPredictionCursor, PredictionReversalContext, PredictionScoringInput,
+};
+use gpui::{AppContext as _, AsyncApp};
+use std::fs::File;
+use std::io::BufWriter;
+use std::path::Path;
+use std::sync::Arc;
+
+pub async fn run_scoring(
+    example: &mut Example,
+    args: &PredictArgs,
+    app_state: Arc<EpAppState>,
+    example_progress: &ExampleProgress,
+    cx: AsyncApp,
+) -> anyhow::Result<()> {
+    run_prediction(example, args, app_state, example_progress, cx.clone()).await?;
+
+    let progress = example_progress.start(Step::Score);
+
+    progress.set_substatus("computing metrics");
+    let example_for_scoring = example.clone();
+    example.score = cx
+        .background_spawn(async move {
+            let prompt_inputs = example_for_scoring
+                .prompt_inputs
+                .as_ref()
+                .context("prompt_inputs is required for scoring - run prediction first or ensure JSON includes prompt_inputs")?;
+            let original_text: &str = prompt_inputs.cursor_excerpt.as_ref();
+            let expected_patches_with_cursors = example_for_scoring
+                .spec
+                .expected_patches_with_cursor_positions();
+
+            let old_editable_region = if let Some(p) = example_for_scoring.prompt.as_ref() {
+                if matches!(
+                    p.provider,
+                    PredictionProvider::Teacher(_, _) | PredictionProvider::TeacherNonBatching(_, _)
+                ) {
+                    Some(
+                        TeacherPrompt::extract_editable_region(&p.input)?
+                            .replace(TeacherPrompt::USER_CURSOR_MARKER, ""),
+                    )
+                } else {
+                    None
+                }
+            } else {
+                None
+            };
+
+            let prepared_expected_patches = edit_prediction_metrics::prepare_expected_patches(
+                &expected_patches_with_cursors,
+                original_text,
+                old_editable_region.as_deref(),
+            )
+            .with_context(|| {
+                format!(
+                    "Expected patch did not apply for {}",
+                    example_for_scoring.spec.name
+                )
+            })?;
+
+            let cursor_path = example_for_scoring.spec.cursor_path.as_ref();
+
+            let mut scores = vec![];
+            for prediction in &example_for_scoring.predictions {
+                let actual_patch = prediction.actual_patch.clone().or_else(|| {
+                    parse_prediction_output(
+                        &example_for_scoring,
+                        &prediction.actual_output,
+                        prediction.provider,
+                    )
+                    .ok()
+                    .map(|(patch, _)| patch)
+                });
+
+                let actual_cursor = prediction.actual_cursor.as_ref().map(|cursor| {
+                    ActualPredictionCursor {
+                        row: cursor.row,
+                        editable_region_offset: cursor.editable_region_offset,
+                    }
+                });
+
+                scores.push(edit_prediction_metrics::score_prediction(
+                    PredictionScoringInput {
+                        original_text,
+                        expected_patches: &prepared_expected_patches,
+                        actual_patch: actual_patch.as_deref(),
+                        actual_cursor,
+                        reversal_context: Some(PredictionReversalContext {
+                            edit_history: &prompt_inputs.events,
+                            excerpt_start_row: prompt_inputs.excerpt_start_row,
+                            cursor_path,
+                        }),
+                        cumulative_logprob: prediction.cumulative_logprob,
+                        avg_logprob: prediction.avg_logprob,
+                    },
+                ));
+            }
+
+            anyhow::Ok(scores)
+        })
+        .await?;
+    Ok(())
+}
+
+pub fn print_report(examples: &[Example], verbose: bool) {
+    const MAX_EXAMPLES_DEFAULT: usize = 20;
+    use crate::metrics::ClassificationMetrics;
+
+    const LINE_WIDTH: usize = 101;
+    let separator = "─".repeat(LINE_WIDTH);
+
+    println!("{}", separator);
+    println!(
+        "{:<40} {:>8} {:>5} {:>7} {:>7} {:>7} {:>7} {:>6} {:>5}",
+        "Example", "DeltaChrF", "Brace", "F1", "Revert", "QaRev", "QaConf", "Cursor", "WrgER"
+    );
+    println!("{}", separator);
+
+    let mut all_delta_chr_f_scores = Vec::new();
+    let mut all_reversal_ratios = Vec::new();
+    let mut braces_disbalance_sum: usize = 0;
+    let mut total_delta_chr_f = ClassificationMetrics::default();
+    let mut total_delta_chr_f_precision = 0.0;
+    let mut total_delta_chr_f_recall = 0.0;
+    let mut delta_chr_f_beta = 0.0;
+    let mut total_exact_lines = ClassificationMetrics::default();
+    let mut total_scores: usize = 0;
+    let mut qa_reverts_count: usize = 0;
+    let mut qa_reverts_total: usize = 0;
+    let mut qa_confidence_sum: u64 = 0;
+    let mut qa_confidence_count: usize = 0;
+    let mut cursor_exact_matches: usize = 0;
+    let mut cursor_total: usize = 0;
+    let mut cursor_distance_sum: usize = 0;
+    let mut cursor_distance_count: usize = 0;
+    let mut wrong_editable_region_count: usize = 0;
+    let mut wrong_editable_region_total: usize = 0;
+    let mut isolated_whitespace_count: usize = 0;
+    let mut kept_rate_sum: f64 = 0.0;
+    let mut kept_rate_count: usize = 0;
+    let mut kept_chars_total: usize = 0;
+    let mut correctly_deleted_chars_total: usize = 0;
+    let mut discarded_chars_total: usize = 0;
+    let mut recall_rate_sum: f64 = 0.0;
+    let mut recall_rate_count: usize = 0;
+    let mut patch_inserted_tokens: Vec<usize> = Vec::new();
+    let mut patch_deleted_tokens: Vec<usize> = Vec::new();
+    let mut predictions_with_patch: usize = 0;
+
+    let mut printed_lines: usize = 0;
+    let mut skipped_lines: usize = 0;
+
+    for example in examples {
+        for (score_idx, score) in example.score.iter().enumerate() {
+            let exact_lines = score.exact_lines_counts();
+
+            // Get QA results for this prediction if available
+            let qa_result = example.qa.get(score_idx).and_then(|q| q.as_ref());
+            let qa_reverts_str = qa_result
+                .and_then(|q| q.reverts_edits)
+                .map(|v| if v { "yes" } else { "no" })
+                .unwrap_or("-");
+            let qa_conf_str = qa_result
+                .and_then(|q| q.confidence)
+                .map(|v| format!("{}", v))
+                .unwrap_or("-".to_string());
+
+            // Format wrong editable region metric
+            let wrong_er_str = match score.wrong_editable_region {
+                Some(true) => "✗",
+                Some(false) => "",
+                None => "",
+            };
+
+            // Format cursor metric
+            let cursor_str = match (score.cursor_exact_match, score.cursor_distance) {
+                (Some(true), _) => "✓".to_string(),
+                (Some(false), Some(dist)) => format!("±{}", dist),
+                (Some(false), None) => "✗".to_string(),
+                (None, _) => "-".to_string(),
+            };
+
+            if verbose || printed_lines < MAX_EXAMPLES_DEFAULT {
+                println!(
+                    "{:<40} {:>8.2} {:>5} {:>6.1}% {:>6.1}% {:>7} {:>7} {:>6} {:>5}",
+                    truncate_name(&example.spec.name, 40),
+                    score.delta_chr_f,
+                    score.braces_disbalance,
+                    exact_lines.f1() * 100.0,
+                    score.reversal_ratio * 100.0,
+                    qa_reverts_str,
+                    qa_conf_str,
+                    cursor_str,
+                    wrong_er_str
+                );
+                printed_lines += 1;
+            } else {
+                skipped_lines += 1;
+            }
+
+            all_delta_chr_f_scores.push(score.delta_chr_f);
+            all_reversal_ratios.push(score.reversal_ratio);
+            total_scores += 1;
+            braces_disbalance_sum += score.braces_disbalance;
+            total_delta_chr_f.accumulate(&score.delta_chr_f_counts());
+            total_delta_chr_f_precision += score.delta_chr_f_precision;
+            total_delta_chr_f_recall += score.delta_chr_f_recall;
+            delta_chr_f_beta = score.delta_chr_f_beta;
+            total_exact_lines.accumulate(&score.exact_lines_counts());
+
+            // Accumulate QA metrics
+            if let Some(qa) = qa_result {
+                if let Some(reverts) = qa.reverts_edits {
+                    qa_reverts_total += 1;
+                    if reverts {
+                        qa_reverts_count += 1;
+                    }
+                }
+                if let Some(conf) = qa.confidence {
+                    qa_confidence_sum += conf as u64;
+                    qa_confidence_count += 1;
+                }
+            }
+
+            // Accumulate wrong editable region metrics
+            if let Some(wrong) = score.wrong_editable_region {
+                wrong_editable_region_total += 1;
+                if wrong {
+                    wrong_editable_region_count += 1;
+                }
+            }
+
+            // Accumulate isolated whitespace metrics
+            if score.has_isolated_whitespace_changes {
+                isolated_whitespace_count += 1;
+            }
+
+            // Accumulate kept and recall rate metrics
+            if let Some(kr) = score.kept_rate {
+                kept_rate_sum += kr;
+                kept_rate_count += 1;
+            }
+            if let Some(kept_chars) = score.kept_chars {
+                kept_chars_total += kept_chars;
+            }
+            if let Some(correctly_deleted_chars) = score.correctly_deleted_chars {
+                correctly_deleted_chars_total += correctly_deleted_chars;
+            }
+            if let Some(discarded_chars) = score.discarded_chars {
+                discarded_chars_total += discarded_chars;
+            }
+            if let Some(rr) = score.recall_rate {
+                recall_rate_sum += rr;
+                recall_rate_count += 1;
+            }
+
+            // Accumulate token change metrics (only for predictions that produced a patch)
+            let has_patch = example
+                .predictions
+                .get(score_idx)
+                .and_then(|p| p.actual_patch.as_ref())
+                .is_some_and(|p| !p.is_empty());
+            if has_patch {
+                predictions_with_patch += 1;
+                patch_inserted_tokens.push(score.inserted_tokens);
+                patch_deleted_tokens.push(score.deleted_tokens);
+            }
+
+            // Accumulate cursor metrics
+            if let Some(exact_match) = score.cursor_exact_match {
+                cursor_total += 1;
+                if exact_match {
+                    cursor_exact_matches += 1;
+                }
+            }
+            if let Some(dist) = score.cursor_distance {
+                cursor_distance_sum += dist;
+                cursor_distance_count += 1;
+            }
+        }
+    }
+
+    if skipped_lines > 0 {
+        println!(
+            "{:<40} (use --verbose to see all {} examples)",
+            format!("... and {} more", skipped_lines),
+            printed_lines + skipped_lines
+        );
+    }
+    println!("{}", separator);
+
+    if !all_delta_chr_f_scores.is_empty() {
+        let avg_delta_chr_f: f32 =
+            all_delta_chr_f_scores.iter().sum::<f32>() / all_delta_chr_f_scores.len() as f32;
+        let avg_reversal_ratio: f32 =
+            all_reversal_ratios.iter().sum::<f32>() / all_reversal_ratios.len() as f32;
+        let braces_disbalance_avg: f32 = braces_disbalance_sum as f32 / total_scores as f32;
+
+        let qa_reverts_str = if qa_reverts_total > 0 {
+            format!(
+                "{:.1}%",
+                qa_reverts_count as f32 / qa_reverts_total as f32 * 100.0
+            )
+        } else {
+            "-".to_string()
+        };
+        let qa_conf_str = if qa_confidence_count > 0 {
+            format!(
+                "{:.1}",
+                qa_confidence_sum as f32 / qa_confidence_count as f32
+            )
+        } else {
+            "-".to_string()
+        };
+        let cursor_str = if cursor_total > 0 {
+            format!(
+                "{:.0}%",
+                cursor_exact_matches as f32 / cursor_total as f32 * 100.0
+            )
+        } else {
+            "-".to_string()
+        };
+        let wrong_er_str = if wrong_editable_region_total > 0 {
+            format!(
+                "{:.2}%",
+                wrong_editable_region_count as f32 / wrong_editable_region_total as f32 * 100.0
+            )
+        } else {
+            "-".to_string()
+        };
+        let isolated_ws_str = if total_scores > 0 {
+            format!(
+                "{}/{} ({:.1}%)",
+                isolated_whitespace_count,
+                total_scores,
+                isolated_whitespace_count as f32 / total_scores as f32 * 100.0
+            )
+        } else {
+            "-".to_string()
+        };
+        let avg_cursor_distance = if cursor_distance_count > 0 {
+            Some(cursor_distance_sum as f32 / cursor_distance_count as f32)
+        } else {
+            None
+        };
+
+        println!(
+            "{:<40} {:>8.2} {:>5.1} {:>6.1}% {:>6.1}% {:>7} {:>7} {:>6} {:>5}",
+            "TOTAL / AVERAGE",
+            avg_delta_chr_f,
+            braces_disbalance_avg,
+            total_exact_lines.f1() * 100.0,
+            avg_reversal_ratio * 100.0,
+            qa_reverts_str,
+            qa_conf_str,
+            cursor_str,
+            wrong_er_str
+        );
+        println!("{}", separator);
+        println!(
+            "Delta chrF (β={:.1}): TP={}, FP={}, FN={}, P={:.1}%, R={:.1}%",
+            delta_chr_f_beta,
+            total_delta_chr_f.true_positives,
+            total_delta_chr_f.false_positives,
+            total_delta_chr_f.false_negatives,
+            total_delta_chr_f_precision / total_scores as f64 * 100.0,
+            total_delta_chr_f_recall / total_scores as f64 * 100.0
+        );
+
+        // Print additional cursor metrics if available
+        if let Some(avg_dist) = avg_cursor_distance {
+            println!(
+                "Cursor: {}/{} exact matches ({:.0}%), avg distance: {:.1} bytes",
+                cursor_exact_matches,
+                cursor_total,
+                cursor_exact_matches as f32 / cursor_total as f32 * 100.0,
+                avg_dist
+            );
+        }
+
+        // Print isolated whitespace metrics
+        if total_scores > 0 {
+            println!("Isolated whitespace changes: {}", isolated_ws_str);
+        }
+
+        // Print kept and recall rate metrics
+        if kept_rate_count > 0 {
+            let avg_kept_rate = kept_rate_sum / kept_rate_count as f64;
+            println!(
+                "Kept rate: {:.1}% avg ({} evaluated, kept chars: {}, correctly deleted chars: {}, discarded chars: {})",
+                avg_kept_rate * 100.0,
+                kept_rate_count,
+                kept_chars_total,
+                correctly_deleted_chars_total,
+                discarded_chars_total
+            );
+        }
+        if recall_rate_count > 0 {
+            let avg_recall_rate = recall_rate_sum / recall_rate_count as f64;
+            println!(
+                "Recall rate: {:.1}% avg ({} evaluated)",
+                avg_recall_rate * 100.0,
+                recall_rate_count
+            );
+        }
+
+        // Print token change percentile summary (only for predictions with a patch)
+        if !patch_inserted_tokens.is_empty() {
+            patch_inserted_tokens.sort_unstable();
+            patch_deleted_tokens.sort_unstable();
+            let mut patch_total_tokens: Vec<usize> = patch_inserted_tokens
+                .iter()
+                .zip(patch_deleted_tokens.iter())
+                .map(|(i, d)| i + d)
+                .collect();
+            patch_total_tokens.sort_unstable();
+
+            let patch_rate = predictions_with_patch as f32 / total_scores as f32 * 100.0;
+            println!();
+            println!(
+                "Token changes ({}/{} predictions produced a patch, {:.1}% — table includes only those)",
+                predictions_with_patch, total_scores, patch_rate
+            );
+            println!(
+                "{:<20} {:>8} {:>8} {:>8} {:>8} {:>8}",
+                "", "p25", "p50", "p75", "p90", "p99"
+            );
+            println!("{}", "─".repeat(LINE_WIDTH));
+            println!(
+                "{:<20} {:>8} {:>8} {:>8} {:>8} {:>8}",
+                "Inserted tokens",
+                percentile(&patch_inserted_tokens, 25),
+                percentile(&patch_inserted_tokens, 50),
+                percentile(&patch_inserted_tokens, 75),
+                percentile(&patch_inserted_tokens, 90),
+                percentile(&patch_inserted_tokens, 99),
+            );
+            println!(
+                "{:<20} {:>8} {:>8} {:>8} {:>8} {:>8}",
+                "Deleted tokens",
+                percentile(&patch_deleted_tokens, 25),
+                percentile(&patch_deleted_tokens, 50),
+                percentile(&patch_deleted_tokens, 75),
+                percentile(&patch_deleted_tokens, 90),
+                percentile(&patch_deleted_tokens, 99),
+            );
+            println!(
+                "{:<20} {:>8} {:>8} {:>8} {:>8} {:>8}",
+                "Total tokens",
+                percentile(&patch_total_tokens, 25),
+                percentile(&patch_total_tokens, 50),
+                percentile(&patch_total_tokens, 75),
+                percentile(&patch_total_tokens, 90),
+                percentile(&patch_total_tokens, 99),
+            );
+        }
+    }
+
+    println!("\n");
+}
+
+fn percentile(sorted_values: &[usize], p: usize) -> usize {
+    if sorted_values.is_empty() {
+        return 0;
+    }
+    let idx = (p as f64 / 100.0 * (sorted_values.len() as f64 - 1.0)).round() as usize;
+    sorted_values[idx.min(sorted_values.len() - 1)]
+}
+
+fn truncate_name(name: &str, max_len: usize) -> String {
+    if name.len() <= max_len {
+        name.to_string()
+    } else {
+        format!("{}...", &name[..max_len - 3])
+    }
+}
+
+pub type SummaryJson = edit_prediction_metrics::SummaryJson;
+
+pub fn compute_summary(examples: &[Example]) -> SummaryJson {
+    edit_prediction_metrics::compute_summary(examples.iter().flat_map(|example| {
+        example
+            .score
+            .iter()
+            .enumerate()
+            .map(move |(score_idx, score)| {
+                let qa = example
+                    .qa
+                    .get(score_idx)
+                    .and_then(|qa| qa.as_ref())
+                    .map(|qa| edit_prediction_metrics::QaSummaryData {
+                        reverts_edits: qa.reverts_edits,
+                        confidence: qa.confidence,
+                    });
+
+                edit_prediction_metrics::PredictionSummaryInput { score, qa }
+            })
+    }))
+}
+
+pub fn write_summary_json(examples: &[Example], path: &Path) -> anyhow::Result<()> {
+    let summary = compute_summary(examples);
+    let file = File::create(path)
+        .with_context(|| format!("Failed to create summary JSON file: {}", path.display()))?;
+    let writer = BufWriter::new(file);
+    serde_json::to_writer_pretty(writer, &summary)
+        .with_context(|| format!("Failed to write summary JSON to: {}", path.display()))?;
+    eprintln!("Wrote summary JSON to: {}", path.display());
+    Ok(())
+}

crates/edit_prediction_cli/src/split_dataset.rs

@@ -0,0 +1,679 @@
+//! `ep split` implementation.
+//!
+//! This command splits a JSONL dataset into multiple files based on size specifications,
+//! with optional stratification by a JSON field.
+//!
+//! # Usage
+//!
+//! ```text
+//! ep split [--stratify=<field>] [input.jsonl] <out1>=<size1> <out2>=<size2> ...
+//! ```
+//!
+//! If `input.jsonl` is not provided or is `-`, reads from stdin.
+//!
+//! # Size specifications
+//!
+//! - `80%` - percentage of total examples (lines)
+//! - `100` - approximate absolute count of examples (lines)
+//! - `rest` - all remaining items (only one split can use this)
+//!
+//! # Stratification
+//!
+//! The `--stratify` flag controls how examples are grouped before splitting:
+//!
+//! - `cursor-path` (default): group by the `cursor_path` JSON field
+//! - `project`: group by the first component of the `cursor_path` JSON field
+//! - `repo`: group by the `repository_url` JSON field
+//! - `none`: no grouping, split individual examples
+//!
+//! When stratifying, the split ensures each output file contains examples from
+//! non-overlapping groups. Size specifications always apply to the number of
+//! examples (lines), with whole groups assigned greedily to meet the target.
+//! Examples missing the stratification field are treated as individual groups.
+
+use anyhow::{Context as _, Result, bail};
+use clap::Args;
+use rand::SeedableRng;
+use rand::seq::SliceRandom;
+use serde_json::Value;
+use std::collections::BTreeMap;
+use std::fs::File;
+use std::io::{self, BufRead, BufReader, BufWriter, Write};
+use std::path::{Path, PathBuf};
+
+/// `ep split` CLI args.
+#[derive(Debug, Args, Clone)]
+#[command(
+    about = "Split a JSONL dataset into multiple files with optional stratification",
+    after_help = r#"SIZE SPECIFICATIONS:
+  <percentage>%    Percentage of total (e.g., 80%)
+  <count>          Absolute number (e.g., 100)
+  rest             All remaining items (only one output can use this)
+
+  Sizes always apply to examples (lines). When stratifying, whole groups
+  are assigned greedily to approximate the target count.
+
+EXAMPLES:
+  # Split 80% train, 20% validation (default: stratify by cursor_path)
+  ep split input.jsonl train.jsonl=80% valid.jsonl=rest
+
+  # Split into train/valid/test
+  ep split input.jsonl train.jsonl=80% valid.jsonl=10% test.jsonl=rest
+
+  # Stratify by repository_url instead of cursor_path
+  ep split --stratify=repo input.jsonl train.jsonl=80% valid.jsonl=rest
+
+  # No stratification (split by individual examples)
+  ep split --stratify=none input.jsonl train.jsonl=80% valid.jsonl=rest
+
+  # Read from stdin
+  cat input.jsonl | ep split train.jsonl=80% valid.jsonl=rest
+
+  # Reproducible split with seed
+  ep split --seed 42 input.jsonl train.jsonl=80% valid.jsonl=rest
+
+STRATIFICATION:
+  Controls how examples are grouped before splitting:
+    cursor-path  Group by "cursor_path" field (default)
+    project      Group by the first component of the "cursor_path" field
+    repo         Group by "repository_url" field
+    none         No grouping, split individual examples
+
+  When stratifying, the split ensures each output file contains examples
+  from non-overlapping groups. This prevents data leakage between
+  train/test splits.
+"#
+)]
+pub struct SplitArgs {
+    /// Random seed for reproducibility
+    #[arg(long)]
+    pub seed: Option<u64>,
+
+    /// Stratification field for splitting the dataset
+    #[arg(long, default_value = "cursor-path")]
+    pub stratify: Stratify,
+}
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq, clap::ValueEnum, strum::Display)]
+pub enum Stratify {
+    #[strum(serialize = "cursor_path")]
+    CursorPath,
+    #[strum(serialize = "project")]
+    Project,
+    #[strum(serialize = "repo")]
+    Repo,
+    #[strum(serialize = "none")]
+    None,
+}
+
+#[derive(Debug, Clone)]
+pub enum SplitSize {
+    Percentage(f64),
+    Absolute(usize),
+    Rest,
+}
+
+#[derive(Debug, Clone)]
+pub struct SplitSpec {
+    pub path: PathBuf,
+    pub size: SplitSize,
+}
+
+fn parse_split_spec(spec: &str) -> Result<SplitSpec> {
+    let (path, size_str) = spec
+        .rsplit_once('=')
+        .with_context(|| format!("invalid split spec '{}': expected <path>=<size>", spec))?;
+
+    let size = if size_str == "rest" {
+        SplitSize::Rest
+    } else if size_str.ends_with('%') {
+        let pct_str = size_str.trim_end_matches('%');
+        let pct: f64 = pct_str
+            .parse()
+            .with_context(|| format!("invalid percentage '{}' in '{}'", pct_str, spec))?;
+        if !(0.0..=100.0).contains(&pct) {
+            bail!("percentage must be between 0 and 100, got {}", pct);
+        }
+        SplitSize::Percentage(pct / 100.0)
+    } else {
+        let count: usize = size_str
+            .parse()
+            .with_context(|| format!("invalid count '{}' in '{}'", size_str, spec))?;
+        SplitSize::Absolute(count)
+    };
+
+    Ok(SplitSpec {
+        path: PathBuf::from(path),
+        size,
+    })
+}
+
+fn read_lines_from_input(input: Option<&Path>) -> Result<Vec<String>> {
+    let reader: Box<dyn BufRead> = match input {
+        Some(path) => {
+            let file =
+                File::open(path).with_context(|| format!("failed to open '{}'", path.display()))?;
+            Box::new(BufReader::new(file))
+        }
+        None => Box::new(BufReader::new(io::stdin())),
+    };
+
+    let lines: Vec<String> = reader
+        .lines()
+        .collect::<io::Result<Vec<_>>>()
+        .context("failed to read input lines")?;
+
+    Ok(lines)
+}
+
+fn compute_split_counts(specs: &[SplitSpec], total: usize) -> Result<Vec<usize>> {
+    let mut counts = vec![0usize; specs.len()];
+    let mut remaining = total;
+    let mut rest_index: Option<usize> = None;
+
+    for (i, spec) in specs.iter().enumerate() {
+        match &spec.size {
+            SplitSize::Percentage(pct) => {
+                let count = (total as f64 * pct).round() as usize;
+                counts[i] = count.min(remaining);
+                remaining = remaining.saturating_sub(counts[i]);
+            }
+            SplitSize::Absolute(count) => {
+                counts[i] = (*count).min(remaining);
+                remaining = remaining.saturating_sub(counts[i]);
+            }
+            SplitSize::Rest => {
+                if rest_index.is_some() {
+                    bail!("only one split can use 'rest'");
+                }
+                rest_index = Some(i);
+            }
+        }
+    }
+
+    if let Some(idx) = rest_index {
+        counts[idx] = remaining;
+    }
+
+    Ok(counts)
+}
+
+fn write_lines_to_file(path: &Path, lines: &[String]) -> Result<()> {
+    if let Some(parent) = path.parent() {
+        if !parent.as_os_str().is_empty() {
+            std::fs::create_dir_all(parent)
+                .with_context(|| format!("failed to create directory '{}'", parent.display()))?;
+        }
+    }
+
+    let file =
+        File::create(path).with_context(|| format!("failed to create '{}'", path.display()))?;
+    let mut writer = BufWriter::new(file);
+
+    for line in lines {
+        writeln!(writer, "{}", line)
+            .with_context(|| format!("failed to write to '{}'", path.display()))?;
+    }
+
+    writer
+        .flush()
+        .with_context(|| format!("failed to flush '{}'", path.display()))?;
+
+    Ok(())
+}
+
+pub fn run_split(args: &SplitArgs, inputs: &[PathBuf]) -> Result<()> {
+    if inputs.is_empty() {
+        bail!("usage: ep split [input.jsonl] train.jsonl=80% valid.jsonl=rest");
+    }
+
+    let (input_path, split_specs_raw): (Option<&Path>, &[PathBuf]) =
+        if inputs.first().is_some_and(|p| {
+            let s = p.to_string_lossy();
+            !s.contains('=')
+        }) {
+            let first = inputs.first().map(|p| p.as_path());
+            let first = if first == Some(Path::new("-")) {
+                None
+            } else {
+                first
+            };
+            (first, &inputs[1..])
+        } else {
+            (None, inputs)
+        };
+
+    if split_specs_raw.is_empty() {
+        bail!("no split specifications provided");
+    }
+
+    let specs: Vec<SplitSpec> = split_specs_raw
+        .iter()
+        .map(|p| parse_split_spec(&p.to_string_lossy()))
+        .collect::<Result<Vec<_>>>()?;
+
+    let lines = read_lines_from_input(input_path)?;
+    let total_lines = lines.len();
+
+    if total_lines == 0 {
+        for spec in &specs {
+            write_lines_to_file(&spec.path, &[])?;
+        }
+        return Ok(());
+    }
+
+    let mut grouped_lines = group_lines(&lines, args.stratify);
+
+    if args.stratify != Stratify::None {
+        eprintln!(
+            "Stratifying by {} ({} unique groups, {} examples)",
+            args.stratify,
+            grouped_lines.len(),
+            total_lines
+        );
+    } else {
+        eprintln!(
+            "No stratification, splitting {} examples by line",
+            total_lines
+        );
+    }
+
+    let mut rng = match args.seed {
+        Some(seed) => rand::rngs::StdRng::seed_from_u64(seed),
+        None => rand::rngs::StdRng::from_os_rng(),
+    };
+
+    grouped_lines.shuffle(&mut rng);
+
+    let line_targets = compute_split_counts(&specs, total_lines)?;
+    let rest_index = specs.iter().position(|s| matches!(s.size, SplitSize::Rest));
+    let mut split_outputs: Vec<Vec<String>> = vec![Vec::new(); specs.len()];
+    let mut group_iter = grouped_lines.into_iter();
+
+    for (split_idx, &target) in line_targets.iter().enumerate() {
+        if Some(split_idx) == rest_index {
+            continue;
+        }
+        let mut accumulated = 0;
+        while accumulated < target {
+            if let Some(group) = group_iter.next() {
+                accumulated += group.len();
+                split_outputs[split_idx].extend(group);
+            } else {
+                break;
+            }
+        }
+    }
+
+    if let Some(idx) = rest_index {
+        for group in group_iter {
+            split_outputs[idx].extend(group);
+        }
+    }
+
+    for (spec, output_lines) in specs.iter().zip(split_outputs.iter()) {
+        write_lines_to_file(&spec.path, output_lines)?;
+        eprintln!("{}: {} examples", spec.path.display(), output_lines.len());
+    }
+
+    Ok(())
+}
+
+/// Groups lines by the specified stratification field.
+///
+/// When `stratify` is `None`, each line becomes its own group.
+/// When a line is missing the stratification field, it is also placed in its own group.
+fn group_lines(lines: &[String], stratify: Stratify) -> Vec<Vec<String>> {
+    if stratify == Stratify::None {
+        return lines.iter().map(|line| vec![line.clone()]).collect();
+    }
+
+    let get_key = |line: &str| {
+        let json: Value = serde_json::from_str(line).unwrap_or_default();
+        match stratify {
+            Stratify::Repo => json
+                .get("repository_url")
+                .and_then(|v| v.as_str())
+                .map(|s| s.to_string()),
+            Stratify::CursorPath => json
+                .get("cursor_path")
+                .and_then(|v| v.as_str())
+                .map(|s| s.to_string()),
+            Stratify::Project => json
+                .get("cursor_path")
+                .and_then(|v| v.as_str())
+                .and_then(|s| s.split(['/', '\\']).next())
+                .map(|s| s.to_string()),
+            Stratify::None => unreachable!(),
+        }
+    };
+
+    let mut groups: BTreeMap<String, Vec<String>> = BTreeMap::new();
+    let mut ungrouped: Vec<Vec<String>> = Vec::new();
+
+    for line in lines {
+        let key = get_key(line);
+        match key {
+            Some(key) => groups.entry(key).or_default().push(line.clone()),
+            None => ungrouped.push(vec![line.clone()]),
+        }
+    }
+
+    let mut result: Vec<Vec<String>> = groups.into_values().collect();
+    result.extend(ungrouped);
+    result
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use std::io::Write;
+    use tempfile::NamedTempFile;
+
+    fn create_temp_jsonl(lines: &[&str]) -> NamedTempFile {
+        let mut file = NamedTempFile::new().unwrap();
+        for line in lines {
+            writeln!(file, "{}", line).unwrap();
+        }
+        file.flush().unwrap();
+        file
+    }
+
+    #[test]
+    fn test_parse_split_spec_percentage() {
+        let spec = parse_split_spec("train.jsonl=80%").unwrap();
+        assert_eq!(spec.path, PathBuf::from("train.jsonl"));
+        match spec.size {
+            SplitSize::Percentage(p) => assert!((p - 0.8).abs() < 0.001),
+            _ => panic!("expected percentage"),
+        }
+    }
+
+    #[test]
+    fn test_parse_split_spec_absolute() {
+        let spec = parse_split_spec("test.jsonl=100").unwrap();
+        assert_eq!(spec.path, PathBuf::from("test.jsonl"));
+        match spec.size {
+            SplitSize::Absolute(n) => assert_eq!(n, 100),
+            _ => panic!("expected absolute"),
+        }
+    }
+
+    #[test]
+    fn test_parse_split_spec_rest() {
+        let spec = parse_split_spec("valid.jsonl=rest").unwrap();
+        assert_eq!(spec.path, PathBuf::from("valid.jsonl"));
+        assert!(matches!(spec.size, SplitSize::Rest));
+    }
+
+    #[test]
+    fn test_group_lines_none() {
+        let lines = vec!["a".to_string(), "b".to_string(), "c".to_string()];
+        let groups = group_lines(&lines, Stratify::None);
+        assert_eq!(groups.len(), 3);
+        assert!(groups.iter().all(|g| g.len() == 1));
+    }
+
+    #[test]
+    fn test_compute_split_counts_percentage() {
+        let specs = vec![
+            SplitSpec {
+                path: PathBuf::from("a"),
+                size: SplitSize::Percentage(0.8),
+            },
+            SplitSpec {
+                path: PathBuf::from("b"),
+                size: SplitSize::Percentage(0.2),
+            },
+        ];
+        let counts = compute_split_counts(&specs, 100).unwrap();
+        assert_eq!(counts, vec![80, 20]);
+    }
+
+    #[test]
+    fn test_compute_split_counts_with_rest() {
+        let specs = vec![
+            SplitSpec {
+                path: PathBuf::from("a"),
+                size: SplitSize::Percentage(0.8),
+            },
+            SplitSpec {
+                path: PathBuf::from("b"),
+                size: SplitSize::Rest,
+            },
+        ];
+        let counts = compute_split_counts(&specs, 100).unwrap();
+        assert_eq!(counts, vec![80, 20]);
+    }
+
+    #[test]
+    fn test_compute_split_counts_absolute() {
+        let specs = vec![
+            SplitSpec {
+                path: PathBuf::from("a"),
+                size: SplitSize::Absolute(50),
+            },
+            SplitSpec {
+                path: PathBuf::from("b"),
+                size: SplitSize::Rest,
+            },
+        ];
+        let counts = compute_split_counts(&specs, 100).unwrap();
+        assert_eq!(counts, vec![50, 50]);
+    }
+
+    #[test]
+    fn test_group_lines_by_repo() {
+        let lines = vec![
+            r#"{"repository_url": "repo1", "id": 1}"#.to_string(),
+            r#"{"repository_url": "repo1", "id": 2}"#.to_string(),
+            r#"{"repository_url": "repo2", "id": 3}"#.to_string(),
+            r#"{"id": 4}"#.to_string(),
+        ];
+
+        let groups = group_lines(&lines, Stratify::Repo);
+
+        let grouped_count: usize = groups.iter().filter(|g| g.len() > 1).count();
+        let ungrouped_count: usize = groups.iter().filter(|g| g.len() == 1).count();
+        let total_lines: usize = groups.iter().map(|g| g.len()).sum();
+
+        assert_eq!(grouped_count, 1); // repo1 has 2 lines
+        assert_eq!(ungrouped_count, 2); // repo2 (1 line) + line without repo
+        assert_eq!(total_lines, 4);
+    }
+
+    #[test]
+    fn test_group_lines_by_cursor_path() {
+        let lines = vec![
+            r#"{"cursor_path": "src/main.rs", "id": 1}"#.to_string(),
+            r#"{"cursor_path": "src/main.rs", "id": 2}"#.to_string(),
+            r#"{"cursor_path": "src/lib.rs", "id": 3}"#.to_string(),
+        ];
+
+        let groups = group_lines(&lines, Stratify::CursorPath);
+
+        let total_lines: usize = groups.iter().map(|g| g.len()).sum();
+        assert_eq!(groups.len(), 2);
+        assert_eq!(total_lines, 3);
+    }
+
+    #[test]
+    fn test_run_split_basic() {
+        let input = create_temp_jsonl(&[
+            r#"{"repository_url": "repo1", "id": 1}"#,
+            r#"{"repository_url": "repo1", "id": 2}"#,
+            r#"{"repository_url": "repo2", "id": 3}"#,
+            r#"{"repository_url": "repo2", "id": 4}"#,
+            r#"{"repository_url": "repo3", "id": 5}"#,
+            r#"{"repository_url": "repo3", "id": 6}"#,
+            r#"{"repository_url": "repo4", "id": 7}"#,
+            r#"{"repository_url": "repo4", "id": 8}"#,
+        ]);
+
+        let temp_dir = tempfile::tempdir().unwrap();
+        let train_path = temp_dir.path().join("train.jsonl");
+        let valid_path = temp_dir.path().join("valid.jsonl");
+
+        let args = SplitArgs {
+            seed: Some(42),
+            stratify: Stratify::Repo,
+        };
+        let inputs = vec![
+            input.path().to_path_buf(),
+            PathBuf::from(format!("{}=50%", train_path.display())),
+            PathBuf::from(format!("{}=rest", valid_path.display())),
+        ];
+
+        run_split(&args, &inputs).unwrap();
+
+        let train_content = std::fs::read_to_string(&train_path).unwrap();
+        let valid_content = std::fs::read_to_string(&valid_path).unwrap();
+
+        let train_lines: Vec<&str> = train_content.lines().collect();
+        let valid_lines: Vec<&str> = valid_content.lines().collect();
+
+        assert_eq!(train_lines.len() + valid_lines.len(), 8);
+
+        let get_repo = |line: &str| -> Option<String> {
+            let value: Value = serde_json::from_str(line).ok()?;
+            value
+                .get("repository_url")
+                .and_then(|v| v.as_str())
+                .map(|s| s.to_string())
+        };
+
+        let train_repos: std::collections::HashSet<_> =
+            train_lines.iter().filter_map(|l| get_repo(l)).collect();
+        let valid_repos: std::collections::HashSet<_> =
+            valid_lines.iter().filter_map(|l| get_repo(l)).collect();
+
+        assert!(
+            train_repos.is_disjoint(&valid_repos),
+            "train and valid should have non-overlapping repos"
+        );
+    }
+
+    #[test]
+    fn test_multiple_rest_fails() {
+        let specs = vec![
+            SplitSpec {
+                path: PathBuf::from("a"),
+                size: SplitSize::Rest,
+            },
+            SplitSpec {
+                path: PathBuf::from("b"),
+                size: SplitSize::Rest,
+            },
+        ];
+        assert!(compute_split_counts(&specs, 100).is_err());
+    }
+
+    #[test]
+    fn test_absolute_targets_lines_not_groups() {
+        // 5 repos × 3 lines each = 15 total lines.
+        // `train=6` should target ~6 lines (2 groups), NOT 6 groups (all 15 lines).
+        let input = create_temp_jsonl(&[
+            r#"{"repository_url": "r1", "id": 1}"#,
+            r#"{"repository_url": "r1", "id": 2}"#,
+            r#"{"repository_url": "r1", "id": 3}"#,
+            r#"{"repository_url": "r2", "id": 4}"#,
+            r#"{"repository_url": "r2", "id": 5}"#,
+            r#"{"repository_url": "r2", "id": 6}"#,
+            r#"{"repository_url": "r3", "id": 7}"#,
+            r#"{"repository_url": "r3", "id": 8}"#,
+            r#"{"repository_url": "r3", "id": 9}"#,
+            r#"{"repository_url": "r4", "id": 10}"#,
+            r#"{"repository_url": "r4", "id": 11}"#,
+            r#"{"repository_url": "r4", "id": 12}"#,
+            r#"{"repository_url": "r5", "id": 13}"#,
+            r#"{"repository_url": "r5", "id": 14}"#,
+            r#"{"repository_url": "r5", "id": 15}"#,
+        ]);
+
+        let temp_dir = tempfile::tempdir().unwrap();
+        let train_path = temp_dir.path().join("train.jsonl");
+        let valid_path = temp_dir.path().join("valid.jsonl");
+
+        let args = SplitArgs {
+            seed: Some(42),
+            stratify: Stratify::Repo,
+        };
+        let inputs = vec![
+            input.path().to_path_buf(),
+            PathBuf::from(format!("{}=6", train_path.display())),
+            PathBuf::from(format!("{}=rest", valid_path.display())),
+        ];
+
+        run_split(&args, &inputs).unwrap();
+
+        let train_content = std::fs::read_to_string(&train_path).unwrap();
+        let valid_content = std::fs::read_to_string(&valid_path).unwrap();
+
+        let train_lines: Vec<&str> = train_content.lines().collect();
+        let valid_lines: Vec<&str> = valid_content.lines().collect();
+
+        // With 3-line groups, train should get 2 groups (6 lines) to meet the
+        // target of 6, NOT 6 groups (which don't even exist). Valid gets the rest.
+        assert_eq!(train_lines.len(), 6);
+        assert_eq!(valid_lines.len(), 9);
+    }
+
+    #[test]
+    fn test_stratify_by_project() {
+        // 5 repos × 3 lines each = 15 total lines.
+        // `train=6` should target ~6 lines (2 groups), NOT 6 groups (all 15 lines).
+        let input = create_temp_jsonl(&[
+            r#"{"cursor_path": "project1/some/file.rs", "id": 1}"#,
+            r#"{"cursor_path": "project2/some/file.rs", "id": 2}"#,
+            r#"{"cursor_path": "project3/some/file.rs", "id": 3}"#,
+            r#"{"cursor_path": "project1/other/file.rs", "id": 4}"#,
+            r#"{"cursor_path": "project2/other/file.rs", "id": 5}"#,
+            r#"{"cursor_path": "project3/other/file.rs", "id": 6}"#,
+            r#"{"cursor_path": "project3/another/file.rs", "id": 7}"#,
+            r#"{"cursor_path": "project3/even/more.rs", "id": 8}"#,
+        ]);
+
+        let temp_dir = tempfile::tempdir().unwrap();
+        let train_path = temp_dir.path().join("train.jsonl");
+        let valid_path = temp_dir.path().join("valid.jsonl");
+
+        let args = SplitArgs {
+            seed: Some(1),
+            stratify: Stratify::Project,
+        };
+        let inputs = vec![
+            input.path().to_path_buf(),
+            PathBuf::from(format!("{}=4", train_path.display())),
+            PathBuf::from(format!("{}=rest", valid_path.display())),
+        ];
+
+        run_split(&args, &inputs).unwrap();
+
+        let train_content = std::fs::read_to_string(&train_path).unwrap();
+        let valid_content = std::fs::read_to_string(&valid_path).unwrap();
+
+        // Make sure project 1 and project 2 are in the train set, and project 3 is in the valid set.
+        let mut train_ids: Vec<u64> = train_content
+            .lines()
+            .map(|l| {
+                serde_json::from_str::<serde_json::Value>(l).unwrap()["id"]
+                    .as_u64()
+                    .unwrap()
+            })
+            .collect();
+        let mut valid_ids: Vec<u64> = valid_content
+            .lines()
+            .map(|l| {
+                serde_json::from_str::<serde_json::Value>(l).unwrap()["id"]
+                    .as_u64()
+                    .unwrap()
+            })
+            .collect();
+
+        train_ids.sort();
+        valid_ids.sort();
+
+        assert_eq!(train_ids, vec![1, 2, 4, 5]);
+        assert_eq!(valid_ids, vec![3, 6, 7, 8]);
+    }
+}

crates/edit_prediction_cli/src/synthesize.rs

@@ -0,0 +1,967 @@
+use crate::{
+    anthropic_client::PlainLlmClient,
+    git::{ensure_repo_cloned, run_git},
+    paths::{FAILED_EXAMPLES_DIR, LATEST_FAILED_EXAMPLES_DIR, SYNTHESIZE_STATE_FILE},
+    progress::{InfoStyle, Progress, Step, StepProgress},
+};
+use anthropic::ResponseContent;
+use anyhow::{Context as _, Result};
+use chrono::Local;
+use collections::{HashMap, HashSet};
+use edit_prediction::{
+    example_spec::ExampleSpec,
+    udiff::{apply_diff_to_string, edits_for_diff},
+};
+use futures::stream::{FuturesUnordered, StreamExt};
+use indoc::indoc;
+use serde::{Deserialize, Serialize};
+use std::{
+    path::{Path, PathBuf},
+    sync::Arc,
+};
+
+#[derive(Debug, Clone)]
+pub struct SynthesizeConfig {
+    pub repo_urls: Vec<String>,
+    /// Number of examples to generate per repository
+    pub count: usize,
+    pub max_commits: usize,
+    pub output_dir: PathBuf,
+    pub fresh: bool,
+}
+
+#[derive(Debug, Default, Serialize, Deserialize)]
+struct SynthesizeState {
+    repositories: HashMap<String, RepoState>,
+}
+
+#[derive(Debug, Default, Serialize, Deserialize)]
+struct RepoState {
+    processed_commits: HashSet<String>,
+    examples_generated: usize,
+}
+
+impl SynthesizeState {
+    fn load() -> Self {
+        if SYNTHESIZE_STATE_FILE.exists() {
+            std::fs::read_to_string(&*SYNTHESIZE_STATE_FILE)
+                .ok()
+                .and_then(|s| serde_json::from_str(&s).ok())
+                .unwrap_or_default()
+        } else {
+            Self::default()
+        }
+    }
+
+    fn save(&self) -> Result<()> {
+        let content = serde_json::to_string_pretty(self)?;
+        std::fs::write(&*SYNTHESIZE_STATE_FILE, content)?;
+        Ok(())
+    }
+
+    fn take_repo_state(&mut self, repo_url: &str) -> RepoState {
+        self.repositories.remove(repo_url).unwrap_or_default()
+    }
+
+    fn merge_repo_state(&mut self, repo_url: String, repo_state: RepoState) {
+        self.repositories.insert(repo_url, repo_state);
+    }
+}
+
+impl RepoState {
+    fn is_processed(&self, commit_sha: &str) -> bool {
+        self.processed_commits.contains(commit_sha)
+    }
+
+    fn mark_processed(&mut self, commit_sha: &str, examples_count: usize) {
+        self.processed_commits.insert(commit_sha.to_string());
+        self.examples_generated += examples_count;
+    }
+}
+
+#[derive(Debug)]
+struct CommitInfo {
+    sha: String,
+    parent_sha: String,
+    message: String,
+    diff: String,
+    expanded_diff: String,
+}
+
+/// Claude's response parsed into structured form
+#[derive(Debug)]
+struct ClaudeResponse {
+    name: String,
+    reasoning: String,
+    edit_history_hunks: Vec<String>,
+    expected_patch_hunks: Vec<String>,
+}
+
+pub async fn run_synthesize(config: SynthesizeConfig) -> Result<()> {
+    let mut state = if config.fresh {
+        SynthesizeState::default()
+    } else {
+        SynthesizeState::load()
+    };
+
+    std::fs::create_dir_all(&config.output_dir)?;
+    std::fs::create_dir_all(&*FAILED_EXAMPLES_DIR)?;
+
+    // Create "latest_failed" symlink pointing to this run's failed directory
+    if LATEST_FAILED_EXAMPLES_DIR.is_symlink() {
+        std::fs::remove_file(&*LATEST_FAILED_EXAMPLES_DIR)?;
+    }
+    #[cfg(unix)]
+    std::os::unix::fs::symlink(&*FAILED_EXAMPLES_DIR, &*LATEST_FAILED_EXAMPLES_DIR)?;
+    #[cfg(windows)]
+    std::os::windows::fs::symlink_dir(&*FAILED_EXAMPLES_DIR, &*LATEST_FAILED_EXAMPLES_DIR)?;
+
+    let progress = Progress::global();
+    let total_examples = config.count * config.repo_urls.len();
+    progress.set_total_examples(total_examples);
+
+    let client = Arc::new(PlainLlmClient::new()?);
+    let config = Arc::new(config);
+
+    let mut futures: FuturesUnordered<_> = config
+        .repo_urls
+        .iter()
+        .map(|repo_url| {
+            let client = client.clone();
+            let repo_state = state.take_repo_state(repo_url);
+            let config = config.clone();
+            let repo_url = repo_url.clone();
+            async move {
+                let result = synthesize_repo(&client, repo_state, &config, &repo_url).await;
+                (repo_url, result)
+            }
+        })
+        .collect();
+
+    let mut errors = Vec::new();
+    while let Some((repo_url, result)) = futures.next().await {
+        match result {
+            Ok(repo_state) => {
+                state.merge_repo_state(repo_url, repo_state);
+            }
+            Err(e) => {
+                errors.push(e);
+            }
+        }
+    }
+
+    state.save()?;
+
+    progress.finalize();
+
+    if let Some(first_error) = errors.into_iter().next() {
+        return Err(first_error);
+    }
+
+    Ok(())
+}
+
+async fn synthesize_repo(
+    client: &PlainLlmClient,
+    mut repo_state: RepoState,
+    config: &SynthesizeConfig,
+    repo_url: &str,
+) -> Result<RepoState> {
+    let progress = Progress::global();
+    let batch_size = config.max_commits;
+
+    let clone_progress = progress.start(Step::Synthesize, &format!("clone {}", repo_url));
+    let repo_path = ensure_repo_cloned(repo_url).await?;
+    drop(clone_progress);
+
+    let mut examples_generated = 0;
+    let mut commits_skipped = 0;
+
+    'outer: loop {
+        let list_progress = progress.start(
+            Step::Synthesize,
+            &format!("{}: list-commits", repo_name_from_url(repo_url)),
+        );
+        let commits = list_commits(&repo_path, batch_size, commits_skipped).await?;
+        drop(list_progress);
+
+        if commits.is_empty() {
+            break;
+        }
+
+        commits_skipped += commits.len();
+
+        for commit in commits {
+            if examples_generated >= config.count {
+                break 'outer;
+            }
+
+            if !config.fresh && repo_state.is_processed(&commit.sha) {
+                continue;
+            }
+
+            if should_skip_commit(&commit) {
+                continue;
+            }
+
+            let repo_name = repo_name_from_url(repo_url);
+            let commit_label = format!(
+                "{}: {} {}",
+                repo_name,
+                &commit.sha[..8],
+                truncate_message(&commit.message, 40)
+            );
+            let step_progress = Arc::new(progress.start(Step::Synthesize, &commit_label));
+
+            // Single Claude call to identify and copy hunks
+            step_progress.set_substatus("analyzing...");
+            let claude_response =
+                match analyze_commit(client, repo_url, &commit, step_progress.clone()).await {
+                    Ok(Some(response)) => response,
+                    Ok(None) => {
+                        step_progress.set_info("no pattern", InfoStyle::Normal);
+                        repo_state.mark_processed(&commit.sha, 0);
+                        continue;
+                    }
+                    Err(e) => {
+                        step_progress.set_info(format!("error: {:?}", e), InfoStyle::Warning);
+                        repo_state.mark_processed(&commit.sha, 0);
+                        continue;
+                    }
+                };
+
+            // Validate and build the example
+            step_progress.set_substatus("validating...");
+            match build_example(repo_url, &commit, &repo_path, &claude_response).await {
+                Ok(spec) => {
+                    let timestamp = Local::now().format("%Y-%m-%d--%H-%M-%S");
+                    let filename = format!("{}--{}.md", repo_name, timestamp);
+                    let path = config.output_dir.join(&filename);
+                    std::fs::write(&path, spec.to_markdown())?;
+                    examples_generated += 1;
+                    step_progress.set_info(filename, InfoStyle::Normal);
+                }
+                Err(rejection_reason) => {
+                    log::debug!("Example rejected: {}", rejection_reason);
+                    let timestamp = Local::now().format("%Y-%m-%d--%H-%M-%S%.3f");
+                    let filename = format!("{}--{}.md", repo_name, timestamp);
+                    let path = FAILED_EXAMPLES_DIR.join(&filename);
+                    let content = format_rejected_example(&claude_response, &rejection_reason);
+                    if let Err(e) = std::fs::write(&path, content) {
+                        log::warn!("Failed to write rejected example: {:?}", e);
+                    }
+                    step_progress.set_info(format!("rejected: {}", filename), InfoStyle::Warning);
+                }
+            }
+
+            repo_state.mark_processed(&commit.sha, 1);
+        }
+    }
+
+    Ok(repo_state)
+}
+
+fn repo_name_from_url(url: &str) -> String {
+    url.rsplit('/')
+        .next()
+        .unwrap_or(url)
+        .trim_end_matches(".git")
+        .to_string()
+}
+
+fn truncate_message(msg: &str, max_len: usize) -> String {
+    let first_line = msg.lines().next().unwrap_or("");
+    if first_line.len() <= max_len {
+        first_line.to_string()
+    } else {
+        format!("{}...", &first_line[..max_len - 3])
+    }
+}
+
+fn should_skip_commit(commit: &CommitInfo) -> bool {
+    let lines_changed = commit
+        .diff
+        .lines()
+        .filter(|l| l.starts_with('+') || l.starts_with('-'))
+        .count();
+    lines_changed < 30
+        || lines_changed > 1000
+        || is_non_code_commit(commit)
+        || is_rename_commit(commit)
+}
+
+fn is_non_code_commit(commit: &CommitInfo) -> bool {
+    let non_code_extensions = [
+        ".md", ".txt", ".json", ".yaml", ".yml", ".toml", ".lock", ".svg", ".png", ".jpg", ".gif",
+        ".ico", ".woff", ".ttf", ".eot",
+    ];
+
+    let diff_files: Vec<&str> = commit
+        .diff
+        .lines()
+        .filter(|l| l.starts_with("+++ b/") || l.starts_with("--- a/"))
+        .filter_map(|l| {
+            l.strip_prefix("+++ b/")
+                .or_else(|| l.strip_prefix("--- a/"))
+        })
+        .collect();
+
+    if diff_files.is_empty() {
+        return false;
+    }
+
+    diff_files
+        .iter()
+        .all(|f| non_code_extensions.iter().any(|ext| f.ends_with(ext)))
+}
+
+fn is_rename_commit(commit: &CommitInfo) -> bool {
+    commit.diff.contains("similarity index")
+        || commit.diff.contains("rename from")
+        || commit.diff.contains("rename to")
+}
+
+async fn list_commits(
+    repo_path: &Path,
+    max_commits: usize,
+    skip: usize,
+) -> Result<Vec<CommitInfo>> {
+    let output = run_git(
+        repo_path,
+        &[
+            "log",
+            "--no-merges",
+            &format!("--skip={}", skip),
+            &format!("-{}", max_commits),
+            "--format=%H|%P|%s",
+        ],
+    )
+    .await?;
+
+    let mut commits = Vec::new();
+    for line in output.lines() {
+        let parts: Vec<&str> = line.splitn(3, '|').collect();
+        if parts.len() < 3 {
+            continue;
+        }
+        let sha = parts[0].to_string();
+        let parent_sha = parts[1].split_whitespace().next().unwrap_or("").to_string();
+        if parent_sha.is_empty() {
+            continue;
+        }
+
+        // Get standard diff (for skip checks)
+        let diff = run_git(repo_path, &["show", "--format=", &sha])
+            .await
+            .unwrap_or_default();
+
+        // Get expanded diff with 30 lines of context
+        let expanded_diff = run_git(repo_path, &["show", "-U30", "--format=", &sha])
+            .await
+            .unwrap_or_default();
+
+        commits.push(CommitInfo {
+            sha,
+            parent_sha,
+            message: parts[2].to_string(),
+            diff,
+            expanded_diff,
+        });
+    }
+
+    Ok(commits)
+}
+
+fn build_prompt(repo_url: &str, commit: &CommitInfo) -> String {
+    format!(
+        indoc! {r#"
+            You are analyzing a git commit to construct a realistic edit prediction example.
+
+            Your goal is to tell the story of a programmer's editing session: what sequence
+            of changes did they make, and what change logically comes next? We use these examples
+            to train a model to predict edits, so the quality of the EDIT HISTORY is what matters most.
+
+            An edit prediction example consists of:
+            1. **Edit History**: 2-6 hunks showing what the programmer did BEFORE making the expected patch.
+               This is the most important part - it must tell a coherent story of the changes leading up to the prediction.
+            2. **Expected Patch**: One small hunk that logically follows from the edit history.
+
+            Both single-file and multi-file patterns are acceptable.
+
+            ## What Makes a Good Example
+
+            The edit history should read like a story: "First the programmer changed X, then Y, then Z, and now they need to change W."
+
+            GOOD examples (rich sequences with 3+ steps):
+            - Removing a parameter: docstring update → constructor change → field removal → (predict) usage site update
+            - Adding a feature: type definition → first usage → second usage → (predict) third usage
+            - Bug fix pattern: fix in file A → fix in file B → fix in file C → (predict) fix in file D
+
+            BAD examples (respond NO_PATTERN):
+            - Commits where all changes are independent (no narrative thread)
+            - Simple find-and-replace (renaming, version bumps)
+            - Documentation-only or config-only changes
+            - Changes where you can only find 1-2 hunks for the edit history
+
+            ## Commit Information
+
+            Repository: {repo_url}
+            Commit: {sha}
+            Message: {message}
+
+            ## Diff (30 lines context)
+
+            ```diff
+            {expanded_diff}
+            ```
+
+            ## Your Task
+
+            First, THINK through whether this commit can support a good example:
+
+            1. What is the high-level pattern in this commit?
+            2. Can you identify at least 3 related hunks (2 or more for edit history + 1 for expected patch)?
+            3. What would be the narrative? (First... then... then... finally predict...)
+            4. Which specific hunk should be the expected patch (the "punchline")?
+
+            If you cannot construct a coherent 3+ hunk story, respond with just:
+            NO_PATTERN: <brief reason>
+
+            If you CAN construct a good example, respond in this format:
+
+            ANALYSIS:
+            Pattern: <one sentence describing the pattern>
+            Steps:
+            1. <file:line-range> - <what this hunk does>
+            2. <file:line-range> - <what this hunk does>
+            3. <file:line-range> - <what this hunk does>
+            4. [EXPECTED PATCH] <file:line-range> - <what this hunk does>
+
+            NAME: <short description, like a commit message, under 60 chars>
+
+            EDIT_HISTORY:
+
+            Hunk 1:
+            ```diff
+            --- a/src/models/user.py
+            +++ b/src/models/user.py
+            @@ -15,7 +15,6 @@ class User:
+                 """A user in the system.
+
+                 Attributes:
+            -        email: The user's email address.
+                     name: The user's display name.
+                 """
+            ```
+
+            Hunk 2:
+            ```diff
+            --- a/src/models/user.py
+            +++ b/src/models/user.py
+            @@ -25,10 +24,9 @@ class User:
+                 def __init__(
+                     self,
+                     name: str,
+            -        email: str,
+                     created_at: datetime,
+                 ):
+                     self.name = name
+            -        self.email = email
+                     self.created_at = created_at
+            ```
+
+            Hunk 3:
+            ```diff
+            --- a/src/api/handlers.py
+            +++ b/src/api/handlers.py
+            @@ -42,7 +42,6 @@ def create_user(request):
+                 data = request.json()
+                 user = User(
+                     name=data["name"],
+            -        email=data["email"],
+                     created_at=datetime.now(),
+                 )
+                 return user.save()
+            ```
+
+            EXPECTED_PATCH:
+            ```diff
+            --- a/src/api/handlers.py
+            +++ b/src/api/handlers.py
+            @@ -58,7 +57,6 @@ def update_user(request, user_id):
+                 user = User.get(user_id)
+                 user.name = data.get("name", user.name)
+            -    user.email = data.get("email", user.email)
+                 user.save()
+                 return user
+            ```
+
+            ## Requirements for the diffs
+
+            Edit history:
+            - MUST have 3-6 hunks (if you cannot find 3+, respond NO_PATTERN instead)
+            - Each hunk needs file headers (--- a/path and +++ b/path)
+            - Hunks must be valid unified diffs that apply to the parent commit
+            - Order hunks as a programmer would naturally make the changes
+
+            Expected patch:
+            - Must be a SINGLE hunk from a SINGLE file
+            - Must be SMALL: 1-15 changed lines (not counting context)
+            - Must be clearly predictable from the edit history narrative
+        "#},
+        repo_url = repo_url,
+        sha = commit.sha,
+        message = commit.message,
+        expanded_diff = commit.expanded_diff,
+    )
+}
+
+async fn analyze_commit(
+    client: &PlainLlmClient,
+    repo_url: &str,
+    commit: &CommitInfo,
+    step_progress: Arc<StepProgress>,
+) -> Result<Option<ClaudeResponse>> {
+    use anthropic::{Message, RequestContent, Role};
+
+    let prompt = build_prompt(repo_url, commit);
+    let messages = vec![Message {
+        role: Role::User,
+        content: vec![RequestContent::Text {
+            text: prompt,
+            cache_control: None,
+        }],
+    }];
+
+    let response = client
+        .generate_streaming("claude-sonnet-4-5", 8192, messages, |chars, _text| {
+            step_progress.set_substatus(format!("analyzing: {:.1}K", chars as f64 / 1000.0));
+        })
+        .await?;
+
+    // Extract text content from response
+    let response_text: String = response
+        .content
+        .iter()
+        .filter_map(|block| {
+            if let ResponseContent::Text { text } = block {
+                Some(text.as_str())
+            } else {
+                None
+            }
+        })
+        .collect::<Vec<_>>()
+        .join("\n");
+
+    parse_claude_response(&response_text)
+}
+
+fn parse_claude_response(response: &str) -> Result<Option<ClaudeResponse>> {
+    // Check for NO_PATTERN
+    if response.contains("NO_PATTERN:") {
+        return Ok(None);
+    }
+
+    // Parse NAME
+    let name = response
+        .lines()
+        .find(|l| l.starts_with("NAME:"))
+        .map(|l| l.strip_prefix("NAME:").unwrap_or("").trim().to_string())
+        .unwrap_or_else(|| "unnamed example".to_string());
+
+    // Parse ANALYSIS section (Claude's planning) - this is the primary reasoning
+    let reasoning = extract_section(
+        response,
+        "ANALYSIS:",
+        &["NAME:", "REASONING:", "EDIT_HISTORY:", "EXPECTED_PATCH:"],
+    )
+    .unwrap_or_default();
+
+    // Parse EDIT_HISTORY diff block
+    let edit_history_hunks = extract_diff_block(response, "EDIT_HISTORY:")?;
+
+    // Parse EXPECTED_PATCH diff block
+    let expected_patch_hunks = extract_diff_block(response, "EXPECTED_PATCH:")?;
+
+    if edit_history_hunks.is_empty() {
+        anyhow::bail!("No edit history hunks found in response");
+    }
+    if expected_patch_hunks.is_empty() {
+        anyhow::bail!("No expected patch hunks found in response");
+    }
+
+    Ok(Some(ClaudeResponse {
+        name,
+        reasoning,
+        edit_history_hunks,
+        expected_patch_hunks,
+    }))
+}
+
+fn extract_section(text: &str, start_marker: &str, end_markers: &[&str]) -> Option<String> {
+    let start_idx = text.find(start_marker)?;
+    let content_start = start_idx + start_marker.len();
+
+    let end_idx = end_markers
+        .iter()
+        .filter_map(|marker| text[content_start..].find(marker))
+        .min()
+        .map(|idx| content_start + idx)
+        .unwrap_or(text.len());
+
+    Some(text[content_start..end_idx].trim().to_string())
+}
+
+fn extract_diff_block(text: &str, section_marker: &str) -> Result<Vec<String>> {
+    let section_start = text
+        .find(section_marker)
+        .context(format!("Section {} not found", section_marker))?;
+
+    let after_marker = &text[section_start + section_marker.len()..];
+
+    // Find where the next major section starts (to bound our search)
+    let section_end = ["EXPECTED_PATCH:", "## "]
+        .iter()
+        .filter(|&&m| m != section_marker)
+        .filter_map(|marker| after_marker.find(marker))
+        .min()
+        .unwrap_or(after_marker.len());
+
+    let section_content = &after_marker[..section_end];
+
+    // Collect all ```diff blocks in this section
+    let mut hunks = Vec::new();
+    let mut search_start = 0;
+
+    while let Some(diff_start) = section_content[search_start..].find("```diff") {
+        let abs_diff_start = search_start + diff_start;
+        let block_content_start = section_content[abs_diff_start..]
+            .find('\n')
+            .map(|i| abs_diff_start + i + 1)
+            .unwrap_or(abs_diff_start);
+
+        if let Some(block_end_rel) = section_content[block_content_start..].find("```") {
+            let block_end = block_content_start + block_end_rel;
+            let diff_content = section_content[block_content_start..block_end].trim();
+
+            // Split this block into hunks (in case multiple hunks in one block)
+            hunks.extend(split_into_hunks(diff_content));
+
+            search_start = block_end + 3;
+        } else {
+            break;
+        }
+    }
+
+    if hunks.is_empty() {
+        anyhow::bail!("No diff blocks found in section {}", section_marker);
+    }
+
+    Ok(hunks)
+}
+
+/// Split a diff block into individual hunks, preserving file headers
+fn split_into_hunks(diff: &str) -> Vec<String> {
+    let mut hunks = Vec::new();
+    let mut current_file_header: Option<String> = None;
+    let mut current_hunk: Vec<String> = Vec::new();
+    let mut in_hunk = false;
+
+    for line in diff.lines() {
+        if line.starts_with("--- a/") || line.starts_with("--- /") {
+            // Start of file header - flush previous hunk
+            if in_hunk && !current_hunk.is_empty() {
+                let mut hunk_text = String::new();
+                if let Some(ref header) = current_file_header {
+                    hunk_text.push_str(header);
+                    hunk_text.push('\n');
+                }
+                hunk_text.push_str(&current_hunk.join("\n"));
+                hunks.push(hunk_text);
+                current_hunk.clear();
+            }
+            current_file_header = Some(line.to_string());
+            in_hunk = false;
+        } else if line.starts_with("+++ b/") || line.starts_with("+++ /") {
+            if let Some(ref mut header) = current_file_header {
+                header.push('\n');
+                header.push_str(line);
+            }
+        } else if line.starts_with("@@ ") {
+            // New hunk - flush previous
+            if in_hunk && !current_hunk.is_empty() {
+                let mut hunk_text = String::new();
+                if let Some(ref header) = current_file_header {
+                    hunk_text.push_str(header);
+                    hunk_text.push('\n');
+                }
+                hunk_text.push_str(&current_hunk.join("\n"));
+                hunks.push(hunk_text);
+                current_hunk.clear();
+            }
+            current_hunk.push(line.to_string());
+            in_hunk = true;
+        } else if in_hunk {
+            current_hunk.push(line.to_string());
+        }
+    }
+
+    // Flush final hunk
+    if !current_hunk.is_empty() {
+        let mut hunk_text = String::new();
+        if let Some(ref header) = current_file_header {
+            hunk_text.push_str(header);
+            hunk_text.push('\n');
+        }
+        hunk_text.push_str(&current_hunk.join("\n"));
+        hunks.push(hunk_text);
+    }
+
+    hunks
+}
+
+/// Validate Claude's output by applying diffs and build the ExampleSpec
+async fn build_example(
+    repo_url: &str,
+    commit: &CommitInfo,
+    repo_path: &Path,
+    response: &ClaudeResponse,
+) -> Result<ExampleSpec, String> {
+    // Validate expected patch hunks
+    if response.expected_patch_hunks.len() != 1 {
+        return Err(format!(
+            "Expected exactly 1 expected patch hunk, got {}",
+            response.expected_patch_hunks.len()
+        ));
+    }
+
+    // Parse the expected patch to determine cursor file
+    let expected_patch = &response.expected_patch_hunks[0];
+    let cursor_file = extract_file_from_hunk(expected_patch)
+        .ok_or_else(|| "Could not determine file from expected patch".to_string())?;
+
+    // Get the file content before the commit
+    let before_content = run_git(
+        repo_path,
+        &["show", &format!("{}^:{}", commit.sha, cursor_file)],
+    )
+    .await
+    .map_err(|e| format!("Failed to get file content for {}: {}", cursor_file, e))?;
+
+    // Build edit history diff from Claude's hunks
+    let edit_history = response.edit_history_hunks.join("\n");
+
+    // Apply edit history to get intermediate state (validates edit history)
+    let intermediate_state =
+        apply_edit_history_to_content(&before_content, &edit_history, &cursor_file)?;
+
+    // Validate expected patch applies to intermediate state
+    let expected_patch_with_header = ensure_diff_header(expected_patch, &cursor_file);
+    apply_diff_to_string(&expected_patch_with_header, &intermediate_state)
+        .map_err(|e| format!("Expected patch failed to apply: {}", e))?;
+
+    // Find where the expected patch edits would apply in the intermediate state
+    let edits = edits_for_diff(&intermediate_state, &expected_patch_with_header)
+        .map_err(|e| format!("Failed to parse expected patch: {}", e))?;
+    if edits.is_empty() {
+        return Err(
+            "Could not locate expected patch in file (context not found or ambiguous)".to_string(),
+        );
+    }
+
+    // Use the start of the first edit for cursor positioning
+    let cursor_byte_offset = edits[0].0.start;
+
+    // Extract excerpt around the edit location
+    let (excerpt, cursor_offset) = extract_cursor_excerpt(&intermediate_state, cursor_byte_offset)?;
+
+    // Build the ExampleSpec and use set_cursor_excerpt to format with comment marker
+    let comment_prefix = line_comment_prefix(&cursor_file);
+    let reasoning_with_source = format!(
+        "Source commit: {} ({})\n\n{}",
+        commit.sha,
+        truncate_message(&commit.message, 60),
+        response.reasoning
+    );
+    let mut spec = ExampleSpec {
+        name: response.name.clone(),
+        repository_url: repo_url.to_string(),
+        revision: commit.parent_sha.clone(),
+        tags: Vec::new(),
+        reasoning: Some(reasoning_with_source),
+        uncommitted_diff: String::new(),
+        cursor_path: Arc::from(Path::new(&cursor_file)),
+        cursor_position: String::new(),
+        edit_history,
+        expected_patches: vec![expected_patch_with_header],
+        rejected_patch: None,
+
+        telemetry: None,
+        human_feedback: Vec::new(),
+        rating: None,
+    };
+    spec.set_cursor_excerpt(&excerpt, cursor_offset, comment_prefix);
+
+    Ok(spec)
+}
+
+/// Extract file path from a hunk (looks for --- a/path or +++ b/path)
+fn extract_file_from_hunk(hunk: &str) -> Option<String> {
+    for line in hunk.lines() {
+        if let Some(path) = line.strip_prefix("+++ b/") {
+            return Some(path.to_string());
+        }
+        if let Some(path) = line.strip_prefix("--- a/") {
+            return Some(path.to_string());
+        }
+    }
+    None
+}
+
+/// Ensure a hunk has proper file headers
+fn ensure_diff_header(hunk: &str, file_path: &str) -> String {
+    if hunk.contains("--- a/") || hunk.contains("+++ b/") {
+        return hunk.to_string();
+    }
+    format!("--- a/{}\n+++ b/{}\n{}", file_path, file_path, hunk)
+}
+
+/// Apply edit history to file content, only if hunks affect this file
+fn apply_edit_history_to_content(
+    content: &str,
+    edit_history: &str,
+    cursor_file: &str,
+) -> Result<String, String> {
+    // Extract just the hunks for this file from the edit history
+    let file_diff = extract_file_diff_from_combined(edit_history, cursor_file);
+
+    if file_diff.is_empty() {
+        return Ok(content.to_string());
+    }
+
+    apply_diff_to_string(&file_diff, content)
+        .map_err(|e| format!("Failed to apply edit history: {}", e))
+}
+
+/// Extract hunks for a specific file from a combined diff
+fn extract_file_diff_from_combined(combined_diff: &str, target_file: &str) -> String {
+    let mut result = String::new();
+    let mut in_target_file = false;
+    let mut found_header = false;
+
+    for line in combined_diff.lines() {
+        if line.starts_with("--- a/") {
+            let file = line.strip_prefix("--- a/").unwrap_or("");
+            in_target_file = file == target_file;
+            if in_target_file {
+                result.push_str(line);
+                result.push('\n');
+                found_header = false;
+            }
+        } else if line.starts_with("+++ b/") && in_target_file {
+            result.push_str(line);
+            result.push('\n');
+            found_header = true;
+        } else if in_target_file && found_header {
+            if line.starts_with("--- a/") {
+                break;
+            }
+            result.push_str(line);
+            result.push('\n');
+        }
+    }
+
+    result
+}
+
+/// Extract a cursor position excerpt from content around a byte offset.
+/// Returns the excerpt and the cursor offset within the excerpt.
+fn extract_cursor_excerpt(
+    content: &str,
+    cursor_byte_offset: usize,
+) -> Result<(String, usize), String> {
+    // Find the line containing the cursor
+    let line_start = content[..cursor_byte_offset]
+        .rfind('\n')
+        .map(|pos| pos + 1)
+        .unwrap_or(0);
+    let line_end = content[cursor_byte_offset..]
+        .find('\n')
+        .map(|pos| cursor_byte_offset + pos)
+        .unwrap_or(content.len());
+
+    // Get context lines before
+    let lines_before: Vec<&str> = content[..line_start].lines().collect();
+    let context_before: Vec<&str> = lines_before.iter().rev().take(3).rev().cloned().collect();
+
+    // Get context lines after
+    let after_line_end = if line_end < content.len() {
+        line_end + 1
+    } else {
+        line_end
+    };
+    let context_after: Vec<&str> = content[after_line_end..].lines().take(4).collect();
+
+    // The line containing the cursor
+    let cursor_line = &content[line_start..line_end];
+    let cursor_column = cursor_byte_offset - line_start;
+
+    // Build the excerpt
+    let mut excerpt = String::new();
+    for line in context_before {
+        excerpt.push_str(line);
+        excerpt.push('\n');
+    }
+    // Track where cursor will be in the excerpt
+    let cursor_offset_in_excerpt = excerpt.len() + cursor_column;
+    // Line containing cursor
+    excerpt.push_str(cursor_line);
+    excerpt.push('\n');
+    for line in context_after {
+        excerpt.push_str(line);
+        excerpt.push('\n');
+    }
+
+    // Trim trailing newline
+    if excerpt.ends_with('\n') {
+        excerpt.pop();
+    }
+
+    Ok((excerpt, cursor_offset_in_excerpt))
+}
+
+/// Get the line comment prefix for a file based on its extension
+fn line_comment_prefix(file_path: &str) -> &'static str {
+    let extension = file_path.rsplit('.').next().unwrap_or("");
+    match extension {
+        "rs" | "c" | "cpp" | "cc" | "h" | "hpp" | "js" | "ts" | "tsx" | "jsx" | "go" | "java"
+        | "swift" | "kt" | "kts" | "scala" | "cs" | "m" | "mm" | "zig" | "v" | "d" => "//",
+        "py" | "rb" | "sh" | "bash" | "zsh" | "pl" | "pm" | "r" | "jl" | "yaml" | "yml"
+        | "toml" | "coffee" | "cr" | "ex" | "exs" | "elixir" => "#",
+        "lua" | "hs" | "sql" => "--",
+        "lisp" | "clj" | "cljs" | "scm" | "rkt" | "el" => ";",
+        "erl" | "hrl" => "%",
+        _ => "//",
+    }
+}
+
+fn format_rejected_example(response: &ClaudeResponse, rejection_reason: &str) -> String {
+    let mut content = String::new();
+    content.push_str("# Rejected Example\n\n");
+    content.push_str(&format!("## Name\n\n{}\n\n", response.name));
+    content.push_str(&format!("## Reasoning\n\n{}\n\n", response.reasoning));
+    content.push_str("## Edit History Hunks\n\n```diff\n");
+    for hunk in &response.edit_history_hunks {
+        content.push_str(hunk);
+        content.push_str("\n\n");
+    }
+    content.push_str("```\n\n");
+    content.push_str("## Expected Patch Hunks\n\n```diff\n");
+    for hunk in &response.expected_patch_hunks {
+        content.push_str(hunk);
+        content.push_str("\n\n");
+    }
+    content.push_str("```\n\n");
+    content.push_str(&format!("## Rejection Reason\n\n{}\n", rejection_reason));
+    content
+}

crates/edit_prediction_cli/src/truncate_expected_patch.rs

@@ -0,0 +1,75 @@
+use crate::example::{Example, read_example_files};
+use crate::reorder_patch::{Hunk, Patch, PatchLine};
+use clap::Args;
+use std::path::PathBuf;
+
+#[derive(Args, Debug, Clone)]
+pub struct TruncatePatchArgs {
+    /// Number of logical groups ahead to leave
+    #[arg(long)]
+    pub num_groups: usize,
+
+    /// Leave only edits in the file under the cursor
+    #[arg(long, default_value_t = false)]
+    pub current_file_only: bool,
+}
+
+pub fn run_truncate_expected_patch(
+    args: &TruncatePatchArgs,
+    inputs: &[PathBuf],
+) -> anyhow::Result<()> {
+    let stdin_path = PathBuf::from("-");
+    let inputs = if inputs.is_empty() {
+        std::slice::from_ref(&stdin_path)
+    } else {
+        inputs
+    };
+
+    let mut examples = read_example_files(inputs);
+    for example in &mut examples {
+        run_one_input(example, args)?;
+
+        println!("{}", serde_json::to_string(&example)?);
+    }
+    Ok(())
+}
+
+fn run_one_input(example: &mut Example, args: &TruncatePatchArgs) -> anyhow::Result<()> {
+    let mut patch = Patch::parse_unified_diff(&example.spec.expected_patches[0]);
+    let mut groups_left = args.num_groups;
+
+    patch.hunks.retain(|hunk| {
+        if groups_left == 0 {
+            return false;
+        }
+        if starts_new_group(hunk) {
+            groups_left -= 1;
+        }
+
+        if args.current_file_only {
+            return hunk.filename == example.spec.cursor_path.display().to_string();
+        }
+
+        true
+    });
+
+    // Remove all group headers
+    patch.header = String::new();
+    patch.hunks.iter_mut().for_each(|hunk| {
+        hunk.lines.retain(|line| match line {
+            PatchLine::Garbage(line) => !line.starts_with("//"),
+            _ => true,
+        });
+    });
+
+    example.spec.expected_patches[0] = patch.to_string();
+
+    Ok(())
+}
+
+fn starts_new_group(hunk: &Hunk) -> bool {
+    hunk.lines.iter().any(|line| match line {
+        PatchLine::Garbage(content) => content.starts_with("///"),
+        _ => false,
+    })
+}

crates/edit_prediction_cli/src/word_diff.rs

@@ -0,0 +1,270 @@
+//! Word-diff utilities for converting unified diffs to word-diff format.
+
+use similar::{DiffTag, TextDiff};
+
+/// Convert unified diff to word-diff format.
+///
+/// This transforms line-based diffs into word-level diffs where:
+/// - Deletions are marked with `[-...-]`
+/// - Insertions are marked with `{+...+}`
+pub fn unified_to_word_diff(unified_diff: &str) -> String {
+    let lines: Vec<&str> = unified_diff.lines().collect();
+    let mut result = String::new();
+    let mut old_lines: Vec<&str> = Vec::new();
+    let mut new_lines: Vec<&str> = Vec::new();
+
+    let flush_changes =
+        |old_lines: &mut Vec<&str>, new_lines: &mut Vec<&str>, result: &mut String| {
+            if old_lines.is_empty() && new_lines.is_empty() {
+                return;
+            }
+
+            // Strip the leading '-' or '+' from each line
+            let old_text: String = old_lines
+                .iter()
+                .map(|line| if line.len() > 1 { &line[1..] } else { "" })
+                .collect::<Vec<_>>()
+                .join("\n");
+
+            let new_text: String = new_lines
+                .iter()
+                .map(|line| if line.len() > 1 { &line[1..] } else { "" })
+                .collect::<Vec<_>>()
+                .join("\n");
+
+            if !old_text.is_empty() || !new_text.is_empty() {
+                let word_diff = compute_word_diff(&old_text, &new_text);
+                result.push_str(&word_diff);
+            }
+
+            old_lines.clear();
+            new_lines.clear();
+        };
+
+    for line in lines {
+        if line.starts_with("---") || line.starts_with("+++") {
+            flush_changes(&mut old_lines, &mut new_lines, &mut result);
+            result.push_str(line);
+            result.push('\n');
+        } else if line.starts_with("@@") {
+            flush_changes(&mut old_lines, &mut new_lines, &mut result);
+            result.push_str(line);
+            result.push('\n');
+        } else if line.starts_with('-') {
+            old_lines.push(line);
+        } else if line.starts_with('+') {
+            new_lines.push(line);
+        } else if line.starts_with(' ') || line.is_empty() {
+            flush_changes(&mut old_lines, &mut new_lines, &mut result);
+            result.push_str(line);
+            result.push('\n');
+        } else {
+            // Header lines (diff --git, index, etc.)
+            flush_changes(&mut old_lines, &mut new_lines, &mut result);
+            result.push_str(line);
+            result.push('\n');
+        }
+    }
+
+    flush_changes(&mut old_lines, &mut new_lines, &mut result);
+    result
+}
+
+/// Compute word-level diff between two text blocks.
+///
+/// Words and whitespace are treated as separate tokens. The output uses:
+/// - `[-...-]` for deleted content
+/// - `{+...+}` for inserted content
+fn compute_word_diff(old_text: &str, new_text: &str) -> String {
+    // Split into words while preserving whitespace
+    let old_words = tokenize(old_text);
+    let new_words = tokenize(new_text);
+
+    let ops = diff_tokens(&old_words, &new_words);
+    let mut result = String::new();
+
+    for op in ops {
+        match op {
+            DiffOp::Equal(start, end) => {
+                for token in &old_words[start..end] {
+                    result.push_str(token);
+                }
+            }
+            DiffOp::Delete(start, end) => {
+                result.push_str("[-");
+                for token in &old_words[start..end] {
+                    result.push_str(token);
+                }
+                result.push_str("-]");
+            }
+            DiffOp::Insert(start, end) => {
+                result.push_str("{+");
+                for token in &new_words[start..end] {
+                    result.push_str(token);
+                }
+                result.push_str("+}");
+            }
+            DiffOp::Replace {
+                old_start,
+                old_end,
+                new_start,
+                new_end,
+            } => {
+                result.push_str("[-");
+                for token in &old_words[old_start..old_end] {
+                    result.push_str(token);
+                }
+                result.push_str("-]");
+                result.push_str("{+");
+                for token in &new_words[new_start..new_end] {
+                    result.push_str(token);
+                }
+                result.push_str("+}");
+            }
+        }
+    }
+
+    if !result.is_empty() && !result.ends_with('\n') {
+        result.push('\n');
+    }
+
+    result
+}
+
+/// Classify a character into one of three token classes:
+/// - 0: identifier (alphanumeric or `_`)
+/// - 1: whitespace
+/// - 2: punctuation (everything else, each character becomes its own token)
+fn char_class(ch: char) -> u8 {
+    if ch.is_alphanumeric() || ch == '_' {
+        0
+    } else if ch.is_whitespace() {
+        1
+    } else {
+        2
+    }
+}
+
+/// Tokenize text into identifier words, whitespace runs, and individual punctuation characters.
+///
+/// This splitting aligns with the syntactic atoms of source code so that the
+/// LCS-based diff can produce fine-grained, meaningful change regions.
+pub(crate) fn tokenize(text: &str) -> Vec<&str> {
+    let mut tokens = Vec::new();
+    let mut chars = text.char_indices().peekable();
+
+    while let Some((start, ch)) = chars.next() {
+        let class = char_class(ch);
+        if class == 2 {
+            // Punctuation: each character is a separate token
+            tokens.push(&text[start..start + ch.len_utf8()]);
+        } else {
+            // Identifier or whitespace: collect contiguous run of same class
+            let mut end = start + ch.len_utf8();
+            while let Some(&(_, next_ch)) = chars.peek() {
+                if char_class(next_ch) == class {
+                    end += next_ch.len_utf8();
+                    chars.next();
+                } else {
+                    break;
+                }
+            }
+            tokens.push(&text[start..end]);
+        }
+    }
+
+    tokens
+}
+
+#[derive(Debug)]
+pub(crate) enum DiffOp {
+    Equal(usize, usize),
+    Delete(usize, usize),
+    Insert(usize, usize),
+    Replace {
+        old_start: usize,
+        old_end: usize,
+        new_start: usize,
+        new_end: usize,
+    },
+}
+
+/// Compute diff operations between two token sequences using `similar`'s Myers diff.
+pub(crate) fn diff_tokens<'a>(old: &[&'a str], new: &[&'a str]) -> Vec<DiffOp> {
+    let diff = TextDiff::from_slices(old, new);
+    diff.ops()
+        .iter()
+        .map(|op| {
+            let tag = op.tag();
+            let old_range = op.old_range();
+            let new_range = op.new_range();
+            match tag {
+                DiffTag::Equal => DiffOp::Equal(old_range.start, old_range.end),
+                DiffTag::Delete => DiffOp::Delete(old_range.start, old_range.end),
+                DiffTag::Insert => DiffOp::Insert(new_range.start, new_range.end),
+                DiffTag::Replace => DiffOp::Replace {
+                    old_start: old_range.start,
+                    old_end: old_range.end,
+                    new_start: new_range.start,
+                    new_end: new_range.end,
+                },
+            }
+        })
+        .collect()
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_tokenize() {
+        let tokens = tokenize("hello world");
+        assert_eq!(tokens, vec!["hello", " ", "world"]);
+
+        let tokens = tokenize("  multiple   spaces  ");
+        assert_eq!(tokens, vec!["  ", "multiple", "   ", "spaces", "  "]);
+
+        let tokens = tokenize("self.name");
+        assert_eq!(tokens, vec!["self", ".", "name"]);
+
+        let tokens = tokenize("foo(bar, baz)");
+        assert_eq!(tokens, vec!["foo", "(", "bar", ",", " ", "baz", ")"]);
+
+        let tokens = tokenize("hello_world");
+        assert_eq!(tokens, vec!["hello_world"]);
+
+        let tokens = tokenize("fn();");
+        assert_eq!(tokens, vec!["fn", "(", ")", ";"]);
+
+        let tokens = tokenize("foo_bar123 + baz");
+        assert_eq!(tokens, vec!["foo_bar123", " ", "+", " ", "baz"]);
+
+        let tokens = tokenize("print(\"hello\")");
+        assert_eq!(tokens, vec!["print", "(", "\"", "hello", "\"", ")"]);
+    }
+
+    #[test]
+    fn test_compute_word_diff_simple() {
+        let result = compute_word_diff("hello world", "hello there");
+        assert!(result.contains("[-world-]"));
+        assert!(result.contains("{+there+}"));
+    }
+
+    #[test]
+    fn test_unified_to_word_diff() {
+        let unified = "\
+--- a/file.txt
++++ b/file.txt
+@@ -1,3 +1,3 @@
+ context line
+-old text here
++new text here
+ more context";
+
+        let result = unified_to_word_diff(unified);
+        assert!(result.contains("--- a/file.txt"));
+        assert!(result.contains("+++ b/file.txt"));
+        assert!(result.contains("@@"));
+    }
+}