context_servers: Completion support for context server slash commands (#17085)

This PR adds support for completions via MCP. The protocol now supports
a new request type "completion/complete"
that can either complete a resource URI template (which we currently
don't use in Zed), or a prompt argument.
We use this to add autocompletion to our context server slash commands!


https://github.com/user-attachments/assets/08c9cf04-cbeb-49a7-903f-5049fb3b3d9f



Release Notes:

- context_servers: Added support for argument completions for context
server prompts. These show up as regular completions to slash commands.

crates/assistant/src/slash_command/context_server_command.rs

@@ -1,6 +1,7 @@
 use anyhow::{anyhow, Result};
 use assistant_slash_command::{
-    ArgumentCompletion, SlashCommand, SlashCommandOutput, SlashCommandOutputSection,
+    AfterCompletion, ArgumentCompletion, SlashCommand, SlashCommandOutput,
+    SlashCommandOutputSection,
 };
 use collections::HashMap;
 use context_servers::{
@@ -8,7 +9,7 @@ use context_servers::{
     protocol::PromptInfo,
 };
 use gpui::{Task, WeakView, WindowContext};
-use language::LspAdapterDelegate;
+use language::{CodeLabel, LspAdapterDelegate};
 use std::sync::atomic::AtomicBool;
 use std::sync::Arc;
 use ui::{IconName, SharedString};
@@ -50,12 +51,57 @@ impl SlashCommand for ContextServerSlashCommand {
 
     fn complete_argument(
         self: Arc<Self>,
-        _arguments: &[String],
+        arguments: &[String],
         _cancel: Arc<AtomicBool>,
         _workspace: Option<WeakView<Workspace>>,
-        _cx: &mut WindowContext,
+        cx: &mut WindowContext,
     ) -> Task<Result<Vec<ArgumentCompletion>>> {
-        Task::ready(Ok(Vec::new()))
+        let server_id = self.server_id.clone();
+        let prompt_name = self.prompt.name.clone();
+        let manager = ContextServerManager::global(cx);
+        let manager = manager.read(cx);
+
+        let (arg_name, arg_val) = match completion_argument(&self.prompt, arguments) {
+            Ok(tp) => tp,
+            Err(e) => {
+                return Task::ready(Err(e));
+            }
+        };
+        if let Some(server) = manager.get_server(&server_id) {
+            cx.foreground_executor().spawn(async move {
+                let Some(protocol) = server.client.read().clone() else {
+                    return Err(anyhow!("Context server not initialized"));
+                };
+
+                let completion_result = protocol
+                    .completion(
+                        context_servers::types::CompletionReference::Prompt(
+                            context_servers::types::PromptReference {
+                                r#type: context_servers::types::PromptReferenceType::Prompt,
+                                name: prompt_name,
+                            },
+                        ),
+                        arg_name,
+                        arg_val,
+                    )
+                    .await?;
+
+                let completions = completion_result
+                    .values
+                    .into_iter()
+                    .map(|value| ArgumentCompletion {
+                        label: CodeLabel::plain(value.clone(), None),
+                        new_text: value,
+                        after_completion: AfterCompletion::Continue,
+                        replace_previous_arguments: false,
+                    })
+                    .collect();
+
+                Ok(completions)
+            })
+        } else {
+            Task::ready(Err(anyhow!("Context server not found")))
+        }
     }
 
     fn run(
@@ -102,6 +148,22 @@ impl SlashCommand for ContextServerSlashCommand {
     }
 }
 
+fn completion_argument(prompt: &PromptInfo, arguments: &[String]) -> Result<(String, String)> {
+    if arguments.is_empty() {
+        return Err(anyhow!("No arguments given"));
+    }
+
+    match &prompt.arguments {
+        Some(args) if args.len() == 1 => {
+            let arg_name = args[0].name.clone();
+            let arg_value = arguments.join(" ");
+            Ok((arg_name, arg_value))
+        }
+        Some(_) => Err(anyhow!("Prompt must have exactly one argument")),
+        None => Err(anyhow!("Prompt has no arguments")),
+    }
+}
+
 fn prompt_arguments(prompt: &PromptInfo, arguments: &[String]) -> Result<HashMap<String, String>> {
     match &prompt.arguments {
         Some(args) if args.len() > 1 => Err(anyhow!(