ep: Fix moving cursor to a predicted position

crates/cloud_llm_client/src/predict_edits_v3.rs

@@ -46,6 +46,9 @@ pub struct PredictEditsV3Response {
     pub editable_range: Range<usize>,
     #[serde(default, skip_serializing_if = "Option::is_none")]
     pub model_version: Option<String>,
+    /// Predicted cursor offset within `output`.
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub cursor_offset: Option<usize>,
 }
 
 #[derive(Debug, Deserialize, Serialize)]

crates/edit_prediction/src/edit_prediction_tests.rs

@@ -2356,6 +2356,7 @@ fn model_response(request: &PredictEditsV3Request, diff_to_apply: &str) -> Predi
         request_id: Uuid::new_v4().to_string(),
         editable_range,
         output: new_excerpt,
+        cursor_offset: None,
         model_version: None,
     }
 }
@@ -2365,6 +2366,7 @@ fn empty_response() -> PredictEditsV3Response {
         request_id: Uuid::new_v4().to_string(),
         editable_range: 0..0,
         output: String::new(),
+        cursor_offset: None,
         model_version: None,
     }
 }
@@ -2713,6 +2715,7 @@ async fn test_edit_prediction_no_spurious_trailing_newline(cx: &mut TestAppConte
         output: "hello world\n".to_string(),
         editable_range: 0..excerpt_length,
         model_version: None,
+        cursor_offset: None,
     };
     respond_tx.send(response).unwrap();
 
@@ -2771,9 +2774,10 @@ async fn test_v3_prediction_strips_cursor_marker_from_edit_text(cx: &mut TestApp
     respond_tx
         .send(PredictEditsV3Response {
             request_id: Uuid::new_v4().to_string(),
-            output: "hello<|user_cursor|> world".to_string(),
+            output: "hello world".to_string(),
             editable_range: 0..excerpt_length,
             model_version: None,
+            cursor_offset: Some(5),
         })
         .unwrap();
 
@@ -2878,6 +2882,7 @@ async fn make_test_ep_store(
                                     editable_range: 0..req.input.cursor_excerpt.len(),
                                     output: completion_response.lock().clone(),
                                     model_version: None,
+                                    cursor_offset: None,
                                 })
                                 .unwrap()
                                 .into(),

crates/edit_prediction/src/zeta.rs

@@ -25,8 +25,7 @@ use zeta_prompt::{ParsedOutput, ZetaPromptInput};
 use std::{env, ops::Range, path::Path, sync::Arc};
 use zeta_prompt::{
     ZetaFormat, format_zeta_prompt, get_prefill, parse_zeta2_model_output,
-    parsed_output_from_editable_region, prompt_input_contains_special_tokens,
-    stop_tokens_for_format,
+    prompt_input_contains_special_tokens, stop_tokens_for_format,
     zeta1::{self, EDITABLE_REGION_END_MARKER},
 };
 
@@ -280,12 +279,12 @@ pub fn request_prediction_with_zeta(
                     .await?;
 
                     let request_id = EditPredictionId(response.request_id.into());
-                    let output_text = Some(response.output).filter(|s| !s.is_empty());
                     let model_version = response.model_version;
-                    let parsed_output = parsed_output_from_editable_region(
-                        response.editable_range,
-                        output_text.unwrap_or_default(),
-                    );
+                    let parsed_output = ParsedOutput {
+                        new_editable_region: response.output,
+                        range_in_excerpt: response.editable_range,
+                        cursor_offset_in_new_editable_region: response.cursor_offset,
+                    };
 
                     Some((request_id, Some(parsed_output), model_version, usage))
                 })