1use anyhow::{Context as _, Result, anyhow};
2use assistant_slash_command::{
3 AfterCompletion, ArgumentCompletion, SlashCommand, SlashCommandOutput,
4 SlashCommandOutputSection, SlashCommandResult,
5};
6use collections::HashMap;
7use context_server::{ContextServerId, types::Prompt};
8use gpui::{App, Entity, Task, WeakEntity, Window};
9use language::{BufferSnapshot, CodeLabel, LspAdapterDelegate};
10use project::context_server_store::ContextServerStore;
11use std::sync::Arc;
12use std::sync::atomic::AtomicBool;
13use text::LineEnding;
14use ui::{IconName, SharedString};
15use workspace::Workspace;
16
17use crate::create_label_for_command;
18
19pub struct ContextServerSlashCommand {
20 store: Entity<ContextServerStore>,
21 server_id: ContextServerId,
22 prompt: Prompt,
23}
24
25impl ContextServerSlashCommand {
26 pub fn new(store: Entity<ContextServerStore>, id: ContextServerId, prompt: Prompt) -> Self {
27 Self {
28 server_id: id,
29 prompt,
30 store,
31 }
32 }
33}
34
35impl SlashCommand for ContextServerSlashCommand {
36 fn name(&self) -> String {
37 self.prompt.name.clone()
38 }
39
40 fn label(&self, cx: &App) -> language::CodeLabel {
41 let mut parts = vec![self.prompt.name.as_str()];
42 if let Some(args) = &self.prompt.arguments {
43 if let Some(arg) = args.first() {
44 parts.push(arg.name.as_str());
45 }
46 }
47 create_label_for_command(parts[0], &parts[1..], cx)
48 }
49
50 fn description(&self) -> String {
51 match &self.prompt.description {
52 Some(desc) => desc.clone(),
53 None => format!("Run '{}' from {}", self.prompt.name, self.server_id),
54 }
55 }
56
57 fn menu_text(&self) -> String {
58 match &self.prompt.description {
59 Some(desc) => desc.clone(),
60 None => format!("Run '{}' from {}", self.prompt.name, self.server_id),
61 }
62 }
63
64 fn requires_argument(&self) -> bool {
65 self.prompt.arguments.as_ref().map_or(false, |args| {
66 args.iter().any(|arg| arg.required == Some(true))
67 })
68 }
69
70 fn complete_argument(
71 self: Arc<Self>,
72 arguments: &[String],
73 _cancel: Arc<AtomicBool>,
74 _workspace: Option<WeakEntity<Workspace>>,
75 _window: &mut Window,
76 cx: &mut App,
77 ) -> Task<Result<Vec<ArgumentCompletion>>> {
78 let Ok((arg_name, arg_value)) = completion_argument(&self.prompt, arguments) else {
79 return Task::ready(Err(anyhow!("Failed to complete argument")));
80 };
81
82 let server_id = self.server_id.clone();
83 let prompt_name = self.prompt.name.clone();
84
85 if let Some(server) = self.store.read(cx).get_running_server(&server_id) {
86 cx.foreground_executor().spawn(async move {
87 let protocol = server.client().context("Context server not initialized")?;
88
89 let response = protocol
90 .request::<context_server::types::requests::CompletionComplete>(
91 context_server::types::CompletionCompleteParams {
92 reference: context_server::types::CompletionReference::Prompt(
93 context_server::types::PromptReference {
94 ty: context_server::types::PromptReferenceType::Prompt,
95 name: prompt_name,
96 },
97 ),
98 argument: context_server::types::CompletionArgument {
99 name: arg_name,
100 value: arg_value,
101 },
102 meta: None,
103 },
104 )
105 .await?;
106
107 let completions = response
108 .completion
109 .values
110 .into_iter()
111 .map(|value| ArgumentCompletion {
112 label: CodeLabel::plain(value.clone(), None),
113 new_text: value,
114 after_completion: AfterCompletion::Continue,
115 replace_previous_arguments: false,
116 })
117 .collect();
118 Ok(completions)
119 })
120 } else {
121 Task::ready(Err(anyhow!("Context server not found")))
122 }
123 }
124
125 fn run(
126 self: Arc<Self>,
127 arguments: &[String],
128 _context_slash_command_output_sections: &[SlashCommandOutputSection<language::Anchor>],
129 _context_buffer: BufferSnapshot,
130 _workspace: WeakEntity<Workspace>,
131 _delegate: Option<Arc<dyn LspAdapterDelegate>>,
132 _window: &mut Window,
133 cx: &mut App,
134 ) -> Task<SlashCommandResult> {
135 let server_id = self.server_id.clone();
136 let prompt_name = self.prompt.name.clone();
137
138 let prompt_args = match prompt_arguments(&self.prompt, arguments) {
139 Ok(args) => args,
140 Err(e) => return Task::ready(Err(e)),
141 };
142
143 let store = self.store.read(cx);
144 if let Some(server) = store.get_running_server(&server_id) {
145 cx.foreground_executor().spawn(async move {
146 let protocol = server.client().context("Context server not initialized")?;
147 let response = protocol
148 .request::<context_server::types::requests::PromptsGet>(
149 context_server::types::PromptsGetParams {
150 name: prompt_name.clone(),
151 arguments: Some(prompt_args),
152 meta: None,
153 },
154 )
155 .await?;
156
157 anyhow::ensure!(
158 response
159 .messages
160 .iter()
161 .all(|msg| matches!(msg.role, context_server::types::Role::User)),
162 "Prompt contains non-user roles, which is not supported"
163 );
164
165 // Extract text from user messages into a single prompt string
166 let mut prompt = response
167 .messages
168 .into_iter()
169 .filter_map(|msg| match msg.content {
170 context_server::types::MessageContent::Text { text, .. } => Some(text),
171 _ => None,
172 })
173 .collect::<Vec<String>>()
174 .join("\n\n");
175
176 // We must normalize the line endings here, since servers might return CR characters.
177 LineEnding::normalize(&mut prompt);
178
179 Ok(SlashCommandOutput {
180 sections: vec![SlashCommandOutputSection {
181 range: 0..(prompt.len()),
182 icon: IconName::ZedAssistant,
183 label: SharedString::from(
184 response
185 .description
186 .unwrap_or(format!("Result from {}", prompt_name)),
187 ),
188 metadata: None,
189 }],
190 text: prompt,
191 run_commands_in_text: false,
192 }
193 .to_event_stream())
194 })
195 } else {
196 Task::ready(Err(anyhow!("Context server not found")))
197 }
198 }
199}
200
201fn completion_argument(prompt: &Prompt, arguments: &[String]) -> Result<(String, String)> {
202 anyhow::ensure!(!arguments.is_empty(), "No arguments given");
203
204 match &prompt.arguments {
205 Some(args) if args.len() == 1 => {
206 let arg_name = args[0].name.clone();
207 let arg_value = arguments.join(" ");
208 Ok((arg_name, arg_value))
209 }
210 Some(_) => anyhow::bail!("Prompt must have exactly one argument"),
211 None => anyhow::bail!("Prompt has no arguments"),
212 }
213}
214
215fn prompt_arguments(prompt: &Prompt, arguments: &[String]) -> Result<HashMap<String, String>> {
216 match &prompt.arguments {
217 Some(args) if args.len() > 1 => {
218 anyhow::bail!("Prompt has more than one argument, which is not supported");
219 }
220 Some(args) if args.len() == 1 => {
221 if !arguments.is_empty() {
222 let mut map = HashMap::default();
223 map.insert(args[0].name.clone(), arguments.join(" "));
224 Ok(map)
225 } else if arguments.is_empty() && args[0].required == Some(false) {
226 Ok(HashMap::default())
227 } else {
228 anyhow::bail!("Prompt expects argument but none given");
229 }
230 }
231 Some(_) | None => {
232 anyhow::ensure!(
233 arguments.is_empty(),
234 "Prompt expects no arguments but some were given"
235 );
236 Ok(HashMap::default())
237 }
238 }
239}
240
241/// MCP servers can return prompts with multiple arguments. Since we only
242/// support one argument, we ignore all others. This is the necessary predicate
243/// for this.
244pub fn acceptable_prompt(prompt: &Prompt) -> bool {
245 match &prompt.arguments {
246 None => true,
247 Some(args) if args.len() <= 1 => true,
248 _ => false,
249 }
250}