diff --git a/crates/edit_prediction/src/edit_prediction.rs b/crates/edit_prediction/src/edit_prediction.rs index 37964e484a..8b57c442c2 100644 --- a/crates/edit_prediction/src/edit_prediction.rs +++ b/crates/edit_prediction/src/edit_prediction.rs @@ -1698,16 +1698,6 @@ impl EditPredictionStore { .. } = pending_prediction; let settled_editable_region_for_metrics = settled_editable_region.clone(); - let kept_rate_result = cx - .background_spawn(async move { - compute_kept_rate( - &editable_region_before_prediction, - &predicted_editable_region, - &settled_editable_region_for_metrics, - ) - }) - .await; - #[cfg(test)] { let request_id = request_id.clone(); @@ -1718,12 +1708,17 @@ impl EditPredictionStore { } }); } - cx.background_spawn({ let client = client.clone(); let llm_token = llm_token.clone(); let app_version = app_version.clone(); async move { + let kept_rate_result = compute_kept_rate( + &editable_region_before_prediction, + &predicted_editable_region, + &settled_editable_region_for_metrics, + ); + let result: anyhow::Result<()> = async { let settled_editable_region = can_collect_data.then_some(settled_editable_region); @@ -1756,6 +1751,10 @@ impl EditPredictionStore { model_version, e2e_latency_ms: e2e_latency.as_millis(), }; + + let json_bytes = serde_json::to_vec(&body)?; + let compressed = zstd::encode_all(&json_bytes[..], 3)?; + let url = client .http_client() .build_zed_llm_url("/predict_edits/settled", &[])?; @@ -1763,7 +1762,8 @@ impl EditPredictionStore { |builder| { Ok(builder .uri(url.as_ref()) - .body(serde_json::to_string(&body)?.into())?) + .header("Content-Encoding", "zstd") + .body(compressed.clone().into())?) }, client, llm_token, diff --git a/crates/edit_prediction/src/edit_prediction_tests.rs b/crates/edit_prediction/src/edit_prediction_tests.rs index 21af57374e..86f218b401 100644 --- a/crates/edit_prediction/src/edit_prediction_tests.rs +++ b/crates/edit_prediction/src/edit_prediction_tests.rs @@ -2541,6 +2541,11 @@ fn init_test_with_fake_client_and_legacy_data_collection( let http_client = FakeHttpClient::create({ move |req| { let uri = req.uri().path().to_string(); + let content_encoding = req + .headers() + .get("Content-Encoding") + .and_then(|value| value.to_str().ok()) + .map(str::to_owned); let mut body = req.into_body(); let predict_req_tx = predict_req_tx.clone(); let reject_req_tx = reject_req_tx.clone(); @@ -2573,7 +2578,12 @@ fn init_test_with_fake_client_and_legacy_data_collection( "/predict_edits/settled" => { let mut buf = Vec::new(); body.read_to_end(&mut buf).await.ok(); - let req = serde_json::from_slice(&buf).unwrap(); + let body = if content_encoding.as_deref() == Some("zstd") { + zstd::decode_all(&buf[..]).unwrap() + } else { + buf + }; + let req = serde_json::from_slice(&body).unwrap(); settled_req_tx.unbounded_send(req).unwrap(); serde_json::to_string(&SubmitEditPredictionSettledResponse {}).unwrap() }