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 completion_result = protocol
90 .completion(
91 context_server::types::CompletionReference::Prompt(
92 context_server::types::PromptReference {
93 r#type: context_server::types::PromptReferenceType::Prompt,
94 name: prompt_name,
95 },
96 ),
97 arg_name,
98 arg_value,
99 )
100 .await?;
101
102 let completions = completion_result
103 .values
104 .into_iter()
105 .map(|value| ArgumentCompletion {
106 label: CodeLabel::plain(value.clone(), None),
107 new_text: value,
108 after_completion: AfterCompletion::Continue,
109 replace_previous_arguments: false,
110 })
111 .collect();
112 Ok(completions)
113 })
114 } else {
115 Task::ready(Err(anyhow!("Context server not found")))
116 }
117 }
118
119 fn run(
120 self: Arc<Self>,
121 arguments: &[String],
122 _context_slash_command_output_sections: &[SlashCommandOutputSection<language::Anchor>],
123 _context_buffer: BufferSnapshot,
124 _workspace: WeakEntity<Workspace>,
125 _delegate: Option<Arc<dyn LspAdapterDelegate>>,
126 _window: &mut Window,
127 cx: &mut App,
128 ) -> Task<SlashCommandResult> {
129 let server_id = self.server_id.clone();
130 let prompt_name = self.prompt.name.clone();
131
132 let prompt_args = match prompt_arguments(&self.prompt, arguments) {
133 Ok(args) => args,
134 Err(e) => return Task::ready(Err(e)),
135 };
136
137 let store = self.store.read(cx);
138 if let Some(server) = store.get_running_server(&server_id) {
139 cx.foreground_executor().spawn(async move {
140 let protocol = server.client().context("Context server not initialized")?;
141 let result = protocol.run_prompt(&prompt_name, prompt_args).await?;
142
143 anyhow::ensure!(
144 result
145 .messages
146 .iter()
147 .all(|msg| matches!(msg.role, context_server::types::Role::User)),
148 "Prompt contains non-user roles, which is not supported"
149 );
150
151 // Extract text from user messages into a single prompt string
152 let mut prompt = result
153 .messages
154 .into_iter()
155 .filter_map(|msg| match msg.content {
156 context_server::types::MessageContent::Text { text, .. } => Some(text),
157 _ => None,
158 })
159 .collect::<Vec<String>>()
160 .join("\n\n");
161
162 // We must normalize the line endings here, since servers might return CR characters.
163 LineEnding::normalize(&mut prompt);
164
165 Ok(SlashCommandOutput {
166 sections: vec![SlashCommandOutputSection {
167 range: 0..(prompt.len()),
168 icon: IconName::ZedAssistant,
169 label: SharedString::from(
170 result
171 .description
172 .unwrap_or(format!("Result from {}", prompt_name)),
173 ),
174 metadata: None,
175 }],
176 text: prompt,
177 run_commands_in_text: false,
178 }
179 .to_event_stream())
180 })
181 } else {
182 Task::ready(Err(anyhow!("Context server not found")))
183 }
184 }
185}
186
187fn completion_argument(prompt: &Prompt, arguments: &[String]) -> Result<(String, String)> {
188 anyhow::ensure!(!arguments.is_empty(), "No arguments given");
189
190 match &prompt.arguments {
191 Some(args) if args.len() == 1 => {
192 let arg_name = args[0].name.clone();
193 let arg_value = arguments.join(" ");
194 Ok((arg_name, arg_value))
195 }
196 Some(_) => anyhow::bail!("Prompt must have exactly one argument"),
197 None => anyhow::bail!("Prompt has no arguments"),
198 }
199}
200
201fn prompt_arguments(prompt: &Prompt, arguments: &[String]) -> Result<HashMap<String, String>> {
202 match &prompt.arguments {
203 Some(args) if args.len() > 1 => {
204 anyhow::bail!("Prompt has more than one argument, which is not supported");
205 }
206 Some(args) if args.len() == 1 => {
207 if !arguments.is_empty() {
208 let mut map = HashMap::default();
209 map.insert(args[0].name.clone(), arguments.join(" "));
210 Ok(map)
211 } else if arguments.is_empty() && args[0].required == Some(false) {
212 Ok(HashMap::default())
213 } else {
214 anyhow::bail!("Prompt expects argument but none given");
215 }
216 }
217 Some(_) | None => {
218 anyhow::ensure!(
219 arguments.is_empty(),
220 "Prompt expects no arguments but some were given"
221 );
222 Ok(HashMap::default())
223 }
224 }
225}
226
227/// MCP servers can return prompts with multiple arguments. Since we only
228/// support one argument, we ignore all others. This is the necessary predicate
229/// for this.
230pub fn acceptable_prompt(prompt: &Prompt) -> bool {
231 match &prompt.arguments {
232 None => true,
233 Some(args) if args.len() <= 1 => true,
234 _ => false,
235 }
236}