From c3301077af6893095093bddb41e37fa46038dff1 Mon Sep 17 00:00:00 2001
From: Antonio Scandurra <me@as-cii.com>
Date: Fri, 10 Jan 2025 15:07:17 +0100
Subject: [PATCH] Log errors when a prediction fails (#22961)

Release Notes:

- N/A

---------

Co-authored-by: Thorsten <thorsten@zed.dev>
---
 Cargo.lock               |  1 +
 crates/collab/src/llm.rs |  2 +-
 crates/zeta/Cargo.toml   |  1 +
 crates/zeta/src/zeta.rs  | 26 +++++++++++++++-----------
 4 files changed, 18 insertions(+), 12 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index df669c8d8a..8029cf7538 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -16517,6 +16517,7 @@ dependencies = [
  "tree-sitter-go",
  "tree-sitter-rust",
  "ui",
+ "util",
  "uuid",
  "workspace",
  "worktree",
diff --git a/crates/collab/src/llm.rs b/crates/collab/src/llm.rs
index 9cdbad6634..2ece048064 100644
--- a/crates/collab/src/llm.rs
+++ b/crates/collab/src/llm.rs
@@ -477,7 +477,7 @@ async fn predict_edits(
         open_ai::CompletionRequest {
             model: model.to_string(),
             prompt: prompt.clone(),
-            max_tokens: 1024,
+            max_tokens: 2048,
             temperature: 0.,
             prediction: Some(open_ai::Prediction::Content {
                 content: params.input_excerpt,
diff --git a/crates/zeta/Cargo.toml b/crates/zeta/Cargo.toml
index 5db6fddbbb..d7ea7db078 100644
--- a/crates/zeta/Cargo.toml
+++ b/crates/zeta/Cargo.toml
@@ -39,6 +39,7 @@ telemetry.workspace = true
 telemetry_events.workspace = true
 theme.workspace = true
 ui.workspace = true
+util.workspace = true
 uuid.workspace = true
 workspace.workspace = true
 
diff --git a/crates/zeta/src/zeta.rs b/crates/zeta/src/zeta.rs
index bcd638aa12..546f6a2549 100644
--- a/crates/zeta/src/zeta.rs
+++ b/crates/zeta/src/zeta.rs
@@ -30,6 +30,7 @@ use std::{
     time::{Duration, Instant},
 };
 use telemetry_events::InlineCompletionRating;
+use util::ResultExt;
 use uuid::Uuid;
 
 const CURSOR_MARKER: &'static str = "<|user_cursor_is_here|>";
@@ -975,7 +976,7 @@ impl CurrentInlineCompletion {
 
 struct PendingCompletion {
     id: usize,
-    _task: Task<Result<()>>,
+    _task: Task<()>,
 }
 
 pub struct ZetaInlineCompletionProvider {
@@ -1053,13 +1054,16 @@ impl inline_completion::InlineCompletionProvider for ZetaInlineCompletionProvide
                 })
             });
 
-            let mut completion = None;
-            if let Ok(completion_request) = completion_request {
-                completion = Some(CurrentInlineCompletion {
-                    buffer_id: buffer.entity_id(),
-                    completion: completion_request.await?,
-                });
-            }
+            let completion = match completion_request {
+                Ok(completion_request) => {
+                    let completion_request = completion_request.await;
+                    completion_request.map(|completion| CurrentInlineCompletion {
+                        buffer_id: buffer.entity_id(),
+                        completion,
+                    })
+                }
+                Err(error) => Err(error),
+            };
 
             this.update(&mut cx, |this, cx| {
                 if this.pending_completions[0].id == pending_completion_id {
@@ -1068,7 +1072,8 @@ impl inline_completion::InlineCompletionProvider for ZetaInlineCompletionProvide
                     this.pending_completions.clear();
                 }
 
-                if let Some(new_completion) = completion {
+                if let Some(new_completion) = completion.context("zeta prediction failed").log_err()
+                {
                     if let Some(old_completion) = this.current_completion.as_ref() {
                         let snapshot = buffer.read(cx).snapshot();
                         if new_completion.should_replace_completion(&old_completion, &snapshot) {
@@ -1083,12 +1088,11 @@ impl inline_completion::InlineCompletionProvider for ZetaInlineCompletionProvide
                         });
                         this.current_completion = Some(new_completion);
                     }
-                } else {
-                    this.current_completion = None;
                 }
 
                 cx.notify();
             })
+            .ok();
         });
 
         // We always maintain at most two pending completions. When we already