Remember max mode setting per-thread and add a user setting (#30042)

Supersedes: https://github.com/zed-industries/zed/pull/29936

Thanks for your contribution @imumesh18, but we had a slightly different
take on it :)

Release Notes:

- N/A

Co-authored-by: Ben Brandt <benjamin.j.brandt@gmail.com>

Cargo.lock

@@ -612,6 +612,7 @@ dependencies = [
  "serde_json_lenient",
  "settings",
  "workspace-hack",
+ "zed_llm_client",
 ]
 
 [[package]]

assets/settings/default.json

@@ -637,6 +637,8 @@
     "version": "2",
     // Whether the agent is enabled.
     "enabled": true,
+    /// What completion mode to start new threads in, if available. Can be 'normal' or 'max'.
+    "preferred_completion_mode": "normal",
     // Whether to show the agent panel button in the status bar.
     "button": true,
     // Where to dock the agent panel. Can be 'left', 'right' or 'bottom'.

crates/agent/src/message_editor.rs

@@ -8,6 +8,7 @@ use crate::ui::{
     AnimatedLabel, MaxModeTooltip,
     preview::{AgentPreview, UsageCallout},
 };
+use assistant_settings::CompletionMode;
 use buffer_diff::BufferDiff;
 use client::UserStore;
 use collections::{HashMap, HashSet};
@@ -42,7 +43,6 @@ use ui::{Disclosure, DocumentationSide, KeyBinding, PopoverMenuHandle, Tooltip,
 use util::{ResultExt as _, maybe};
 use workspace::dock::DockPosition;
 use workspace::{CollaboratorId, Workspace};
-use zed_llm_client::CompletionMode;
 
 use crate::context_picker::{ContextPicker, ContextPickerCompletionProvider, crease_for_mention};
 use crate::context_store::ContextStore;

crates/agent/src/thread.rs

@@ -5,7 +5,7 @@ use std::sync::Arc;
 use std::time::Instant;
 
 use anyhow::{Result, anyhow};
-use assistant_settings::AssistantSettings;
+use assistant_settings::{AssistantSettings, CompletionMode};
 use assistant_tool::{ActionLog, AnyToolCard, Tool, ToolWorkingSet};
 use chrono::{DateTime, Utc};
 use collections::HashMap;
@@ -37,7 +37,7 @@ use settings::Settings;
 use thiserror::Error;
 use util::{ResultExt as _, TryFutureExt as _, post_inc};
 use uuid::Uuid;
-use zed_llm_client::{CompletionMode, CompletionRequestStatus};
+use zed_llm_client::CompletionRequestStatus;
 
 use crate::ThreadStore;
 use crate::context::{AgentContext, AgentContextHandle, ContextLoadResult, LoadedContext};
@@ -312,14 +312,6 @@ pub enum TokenUsageRatio {
     Exceeded,
 }
 
-fn default_completion_mode(cx: &App) -> CompletionMode {
-    if cx.is_staff() {
-        CompletionMode::Max
-    } else {
-        CompletionMode::Normal
-    }
-}
-
 #[derive(Debug, Clone, Copy)]
 pub enum QueueState {
     Sending,
@@ -336,7 +328,7 @@ pub struct Thread {
     detailed_summary_task: Task<Option<()>>,
     detailed_summary_tx: postage::watch::Sender<DetailedSummaryState>,
     detailed_summary_rx: postage::watch::Receiver<DetailedSummaryState>,
-    completion_mode: CompletionMode,
+    completion_mode: assistant_settings::CompletionMode,
     messages: Vec<Message>,
     next_message_id: MessageId,
     last_prompt_id: PromptId,
@@ -395,7 +387,7 @@ impl Thread {
             detailed_summary_task: Task::ready(None),
             detailed_summary_tx,
             detailed_summary_rx,
-            completion_mode: default_completion_mode(cx),
+            completion_mode: AssistantSettings::get_global(cx).preferred_completion_mode,
             messages: Vec::new(),
             next_message_id: MessageId(0),
             last_prompt_id: PromptId::new(),
@@ -464,6 +456,10 @@ impl Thread {
                 .or_else(|| registry.default_model())
         });
 
+        let completion_mode = serialized
+            .completion_mode
+            .unwrap_or_else(|| AssistantSettings::get_global(cx).preferred_completion_mode);
+
         Self {
             id,
             updated_at: serialized.updated_at,
@@ -472,7 +468,7 @@ impl Thread {
             detailed_summary_task: Task::ready(None),
             detailed_summary_tx,
             detailed_summary_rx,
-            completion_mode: default_completion_mode(cx),
+            completion_mode,
             messages: serialized
                 .messages
                 .into_iter()
@@ -1095,6 +1091,7 @@ impl Thread {
                         provider: model.provider.id().0.to_string(),
                         model: model.model.id().0.to_string(),
                     }),
+                completion_mode: Some(this.completion_mode),
             })
         })
     }
@@ -1246,9 +1243,9 @@ impl Thread {
 
         request.tools = available_tools;
         request.mode = if model.supports_max_mode() {
-            Some(self.completion_mode)
+            Some(self.completion_mode.into())
         } else {
-            Some(CompletionMode::Normal)
+            Some(CompletionMode::Normal.into())
         };
 
         request

crates/agent/src/thread_store.rs

@@ -5,7 +5,7 @@ use std::rc::Rc;
 use std::sync::Arc;
 
 use anyhow::{Context as _, Result, anyhow};
-use assistant_settings::{AgentProfile, AgentProfileId, AssistantSettings};
+use assistant_settings::{AgentProfile, AgentProfileId, AssistantSettings, CompletionMode};
 use assistant_tool::{ToolId, ToolSource, ToolWorkingSet};
 use chrono::{DateTime, Utc};
 use collections::HashMap;
@@ -651,6 +651,8 @@ pub struct SerializedThread {
     pub exceeded_window_error: Option<ExceededWindowError>,
     #[serde(default)]
     pub model: Option<SerializedLanguageModel>,
+    #[serde(default)]
+    pub completion_mode: Option<CompletionMode>,
 }
 
 #[derive(Serialize, Deserialize, Debug)]
@@ -794,6 +796,7 @@ impl LegacySerializedThread {
             detailed_summary_state: DetailedSummaryState::default(),
             exceeded_window_error: None,
             model: None,
+            completion_mode: None,
         }
     }
 }

crates/assistant_settings/Cargo.toml

@@ -27,6 +27,7 @@ schemars.workspace = true
 serde.workspace = true
 settings.workspace = true
 workspace-hack.workspace = true
+zed_llm_client.workspace = true
 
 [dev-dependencies]
 fs.workspace = true

crates/assistant_settings/src/assistant_settings.rs

@@ -89,6 +89,7 @@ pub struct AssistantSettings {
     pub notify_when_agent_waiting: NotifyWhenAgentWaiting,
     pub stream_edits: bool,
     pub single_file_review: bool,
+    pub preferred_completion_mode: CompletionMode,
 }
 
 impl AssistantSettings {
@@ -226,6 +227,7 @@ impl AssistantSettingsContent {
                     notify_when_agent_waiting: None,
                     stream_edits: None,
                     single_file_review: None,
+                    preferred_completion_mode: None,
                 },
                 VersionedAssistantSettingsContent::V2(ref settings) => settings.clone(),
             },
@@ -255,6 +257,7 @@ impl AssistantSettingsContent {
                 notify_when_agent_waiting: None,
                 stream_edits: None,
                 single_file_review: None,
+                preferred_completion_mode: None,
             },
             None => AssistantSettingsContentV2::default(),
         }
@@ -520,6 +523,7 @@ impl Default for VersionedAssistantSettingsContent {
             notify_when_agent_waiting: None,
             stream_edits: None,
             single_file_review: None,
+            preferred_completion_mode: None,
         })
     }
 }
@@ -583,6 +587,28 @@ pub struct AssistantSettingsContentV2 {
     ///
     /// Default: true
     single_file_review: Option<bool>,
+
+    /// What completion mode to enable for new threads
+    ///
+    /// Default: normal
+    preferred_completion_mode: Option<CompletionMode>,
+}
+
+#[derive(Clone, Copy, Debug, Serialize, Deserialize, JsonSchema, PartialEq, Default)]
+#[serde(rename_all = "snake_case")]
+pub enum CompletionMode {
+    #[default]
+    Normal,
+    Max,
+}
+
+impl From<CompletionMode> for zed_llm_client::CompletionMode {
+    fn from(value: CompletionMode) -> Self {
+        match value {
+            CompletionMode::Normal => zed_llm_client::CompletionMode::Normal,
+            CompletionMode::Max => zed_llm_client::CompletionMode::Max,
+        }
+    }
 }
 
 #[derive(Clone, Debug, Serialize, Deserialize, JsonSchema, PartialEq)]
@@ -750,6 +776,10 @@ impl Settings for AssistantSettings {
             merge(&mut settings.stream_edits, value.stream_edits);
             merge(&mut settings.single_file_review, value.single_file_review);
             merge(&mut settings.default_profile, value.default_profile);
+            merge(
+                &mut settings.preferred_completion_mode,
+                value.preferred_completion_mode,
+            );
 
             if let Some(profiles) = value.profiles {
                 settings
@@ -883,6 +913,7 @@ mod tests {
                                 notify_when_agent_waiting: None,
                                 stream_edits: None,
                                 single_file_review: None,
+                                preferred_completion_mode: None,
                             },
                         )),
                     }