1#![cfg_attr(target_os = "windows", allow(unused, dead_code))]
2
3pub mod assistant_panel;
4pub mod assistant_settings;
5mod context;
6pub mod context_store;
7mod inline_assistant;
8mod model_selector;
9mod patch;
10mod prompt_library;
11mod prompts;
12mod slash_command;
13pub(crate) mod slash_command_picker;
14pub mod slash_command_settings;
15mod streaming_diff;
16mod terminal_inline_assistant;
17mod tools;
18
19pub use assistant_panel::{AssistantPanel, AssistantPanelEvent};
20use assistant_settings::AssistantSettings;
21use assistant_slash_command::SlashCommandRegistry;
22use assistant_tool::ToolRegistry;
23use client::{proto, Client};
24use command_palette_hooks::CommandPaletteFilter;
25pub use context::*;
26use context_servers::ContextServerRegistry;
27pub use context_store::*;
28use feature_flags::FeatureFlagAppExt;
29use fs::Fs;
30use gpui::{actions, AppContext, Global, SharedString, UpdateGlobal};
31use gpui::{impl_actions, Context as _};
32use indexed_docs::IndexedDocsRegistry;
33pub(crate) use inline_assistant::*;
34use language_model::{
35 LanguageModelId, LanguageModelProviderId, LanguageModelRegistry, LanguageModelResponseMessage,
36};
37pub(crate) use model_selector::*;
38pub use patch::*;
39pub use prompts::PromptBuilder;
40use prompts::PromptLoadingParams;
41use semantic_index::{CloudEmbeddingProvider, SemanticDb};
42use serde::{Deserialize, Serialize};
43use settings::{update_settings_file, Settings, SettingsStore};
44use slash_command::workflow_command::WorkflowSlashCommand;
45use slash_command::{
46 auto_command, cargo_workspace_command, context_server_command, default_command, delta_command,
47 diagnostics_command, docs_command, fetch_command, file_command, now_command, project_command,
48 prompt_command, search_command, symbols_command, tab_command, terminal_command,
49 workflow_command,
50};
51use std::path::PathBuf;
52use std::sync::Arc;
53pub(crate) use streaming_diff::*;
54use util::ResultExt;
55
56use crate::slash_command_settings::SlashCommandSettings;
57
58actions!(
59 assistant,
60 [
61 Assist,
62 Split,
63 CopyCode,
64 CycleMessageRole,
65 QuoteSelection,
66 InsertIntoEditor,
67 ToggleFocus,
68 InsertActivePrompt,
69 DeployHistory,
70 DeployPromptLibrary,
71 ConfirmCommand,
72 NewContext,
73 ToggleModelSelector,
74 CycleNextInlineAssist,
75 CyclePreviousInlineAssist
76 ]
77);
78
79#[derive(PartialEq, Clone, Deserialize)]
80pub enum InsertDraggedFiles {
81 ProjectPaths(Vec<PathBuf>),
82 ExternalFiles(Vec<PathBuf>),
83}
84
85impl_actions!(assistant, [InsertDraggedFiles]);
86
87const DEFAULT_CONTEXT_LINES: usize = 50;
88
89#[derive(Copy, Clone, Debug, Eq, PartialEq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
90pub struct MessageId(clock::Lamport);
91
92impl MessageId {
93 pub fn as_u64(self) -> u64 {
94 self.0.as_u64()
95 }
96}
97
98#[derive(Deserialize, Debug)]
99pub struct LanguageModelUsage {
100 pub prompt_tokens: u32,
101 pub completion_tokens: u32,
102 pub total_tokens: u32,
103}
104
105#[derive(Deserialize, Debug)]
106pub struct LanguageModelChoiceDelta {
107 pub index: u32,
108 pub delta: LanguageModelResponseMessage,
109 pub finish_reason: Option<String>,
110}
111
112#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
113pub enum MessageStatus {
114 Pending,
115 Done,
116 Error(SharedString),
117 Canceled,
118}
119
120impl MessageStatus {
121 pub fn from_proto(status: proto::ContextMessageStatus) -> MessageStatus {
122 match status.variant {
123 Some(proto::context_message_status::Variant::Pending(_)) => MessageStatus::Pending,
124 Some(proto::context_message_status::Variant::Done(_)) => MessageStatus::Done,
125 Some(proto::context_message_status::Variant::Error(error)) => {
126 MessageStatus::Error(error.message.into())
127 }
128 Some(proto::context_message_status::Variant::Canceled(_)) => MessageStatus::Canceled,
129 None => MessageStatus::Pending,
130 }
131 }
132
133 pub fn to_proto(&self) -> proto::ContextMessageStatus {
134 match self {
135 MessageStatus::Pending => proto::ContextMessageStatus {
136 variant: Some(proto::context_message_status::Variant::Pending(
137 proto::context_message_status::Pending {},
138 )),
139 },
140 MessageStatus::Done => proto::ContextMessageStatus {
141 variant: Some(proto::context_message_status::Variant::Done(
142 proto::context_message_status::Done {},
143 )),
144 },
145 MessageStatus::Error(message) => proto::ContextMessageStatus {
146 variant: Some(proto::context_message_status::Variant::Error(
147 proto::context_message_status::Error {
148 message: message.to_string(),
149 },
150 )),
151 },
152 MessageStatus::Canceled => proto::ContextMessageStatus {
153 variant: Some(proto::context_message_status::Variant::Canceled(
154 proto::context_message_status::Canceled {},
155 )),
156 },
157 }
158 }
159}
160
161/// The state pertaining to the Assistant.
162#[derive(Default)]
163struct Assistant {
164 /// Whether the Assistant is enabled.
165 enabled: bool,
166}
167
168impl Global for Assistant {}
169
170impl Assistant {
171 const NAMESPACE: &'static str = "assistant";
172
173 fn set_enabled(&mut self, enabled: bool, cx: &mut AppContext) {
174 if self.enabled == enabled {
175 return;
176 }
177
178 self.enabled = enabled;
179
180 if !enabled {
181 CommandPaletteFilter::update_global(cx, |filter, _cx| {
182 filter.hide_namespace(Self::NAMESPACE);
183 });
184
185 return;
186 }
187
188 CommandPaletteFilter::update_global(cx, |filter, _cx| {
189 filter.show_namespace(Self::NAMESPACE);
190 });
191 }
192}
193
194pub fn init(
195 fs: Arc<dyn Fs>,
196 client: Arc<Client>,
197 stdout_is_a_pty: bool,
198 cx: &mut AppContext,
199) -> Arc<PromptBuilder> {
200 cx.set_global(Assistant::default());
201 AssistantSettings::register(cx);
202 SlashCommandSettings::register(cx);
203
204 // TODO: remove this when 0.148.0 is released.
205 if AssistantSettings::get_global(cx).using_outdated_settings_version {
206 update_settings_file::<AssistantSettings>(fs.clone(), cx, {
207 let fs = fs.clone();
208 |content, cx| {
209 content.update_file(fs, cx);
210 }
211 });
212 }
213
214 cx.spawn(|mut cx| {
215 let client = client.clone();
216 async move {
217 let embedding_provider = CloudEmbeddingProvider::new(client.clone());
218 let semantic_index = SemanticDb::new(
219 paths::embeddings_dir().join("semantic-index-db.0.mdb"),
220 Arc::new(embedding_provider),
221 &mut cx,
222 )
223 .await?;
224
225 cx.update(|cx| cx.set_global(semantic_index))
226 }
227 })
228 .detach();
229
230 context_store::init(&client.clone().into());
231 prompt_library::init(cx);
232 init_language_model_settings(cx);
233 assistant_slash_command::init(cx);
234 assistant_tool::init(cx);
235 assistant_panel::init(cx);
236 context_servers::init(cx);
237
238 let prompt_builder = prompts::PromptBuilder::new(Some(PromptLoadingParams {
239 fs: fs.clone(),
240 repo_path: stdout_is_a_pty
241 .then(|| std::env::current_dir().log_err())
242 .flatten(),
243 cx,
244 }))
245 .log_err()
246 .map(Arc::new)
247 .unwrap_or_else(|| Arc::new(prompts::PromptBuilder::new(None).unwrap()));
248 register_slash_commands(Some(prompt_builder.clone()), cx);
249 register_tools(cx);
250 inline_assistant::init(
251 fs.clone(),
252 prompt_builder.clone(),
253 client.telemetry().clone(),
254 cx,
255 );
256 terminal_inline_assistant::init(
257 fs.clone(),
258 prompt_builder.clone(),
259 client.telemetry().clone(),
260 cx,
261 );
262 IndexedDocsRegistry::init_global(cx);
263
264 CommandPaletteFilter::update_global(cx, |filter, _cx| {
265 filter.hide_namespace(Assistant::NAMESPACE);
266 });
267 Assistant::update_global(cx, |assistant, cx| {
268 let settings = AssistantSettings::get_global(cx);
269
270 assistant.set_enabled(settings.enabled, cx);
271 });
272 cx.observe_global::<SettingsStore>(|cx| {
273 Assistant::update_global(cx, |assistant, cx| {
274 let settings = AssistantSettings::get_global(cx);
275 assistant.set_enabled(settings.enabled, cx);
276 });
277 })
278 .detach();
279
280 register_context_server_handlers(cx);
281
282 prompt_builder
283}
284
285fn register_context_server_handlers(cx: &mut AppContext) {
286 cx.subscribe(
287 &context_servers::manager::ContextServerManager::global(cx),
288 |manager, event, cx| match event {
289 context_servers::manager::Event::ServerStarted { server_id } => {
290 cx.update_model(
291 &manager,
292 |manager: &mut context_servers::manager::ContextServerManager, cx| {
293 let slash_command_registry = SlashCommandRegistry::global(cx);
294 let context_server_registry = ContextServerRegistry::global(cx);
295 if let Some(server) = manager.get_server(server_id) {
296 cx.spawn(|_, _| async move {
297 let Some(protocol) = server.client.read().clone() else {
298 return;
299 };
300
301 if protocol.capable(context_servers::protocol::ServerCapability::Prompts) {
302 if let Some(prompts) = protocol.list_prompts().await.log_err() {
303 for prompt in prompts
304 .into_iter()
305 .filter(context_server_command::acceptable_prompt)
306 {
307 log::info!(
308 "registering context server command: {:?}",
309 prompt.name
310 );
311 context_server_registry.register_command(
312 server.id.clone(),
313 prompt.name.as_str(),
314 );
315 slash_command_registry.register_command(
316 context_server_command::ContextServerSlashCommand::new(
317 &server, prompt,
318 ),
319 true,
320 );
321 }
322 }
323 }
324 })
325 .detach();
326 }
327 },
328 );
329
330 cx.update_model(
331 &manager,
332 |manager: &mut context_servers::manager::ContextServerManager, cx| {
333 let tool_registry = ToolRegistry::global(cx);
334 let context_server_registry = ContextServerRegistry::global(cx);
335 if let Some(server) = manager.get_server(server_id) {
336 cx.spawn(|_, _| async move {
337 let Some(protocol) = server.client.read().clone() else {
338 return;
339 };
340
341 if protocol.capable(context_servers::protocol::ServerCapability::Tools) {
342 if let Some(tools) = protocol.list_tools().await.log_err() {
343 for tool in tools.tools {
344 log::info!(
345 "registering context server tool: {:?}",
346 tool.name
347 );
348 context_server_registry.register_tool(
349 server.id.clone(),
350 tool.name.as_str(),
351 );
352 tool_registry.register_tool(
353 tools::context_server_tool::ContextServerTool::new(
354 server.id.clone(),
355 tool
356 ),
357 );
358 }
359 }
360 }
361 })
362 .detach();
363 }
364 },
365 );
366 }
367 context_servers::manager::Event::ServerStopped { server_id } => {
368 let slash_command_registry = SlashCommandRegistry::global(cx);
369 let context_server_registry = ContextServerRegistry::global(cx);
370 if let Some(commands) = context_server_registry.get_commands(server_id) {
371 for command_name in commands {
372 slash_command_registry.unregister_command_by_name(&command_name);
373 context_server_registry.unregister_command(&server_id, &command_name);
374 }
375 }
376
377 if let Some(tools) = context_server_registry.get_tools(server_id) {
378 let tool_registry = ToolRegistry::global(cx);
379 for tool_name in tools {
380 tool_registry.unregister_tool_by_name(&tool_name);
381 context_server_registry.unregister_tool(&server_id, &tool_name);
382 }
383 }
384 }
385 },
386 )
387 .detach();
388}
389
390fn init_language_model_settings(cx: &mut AppContext) {
391 update_active_language_model_from_settings(cx);
392
393 cx.observe_global::<SettingsStore>(update_active_language_model_from_settings)
394 .detach();
395 cx.subscribe(
396 &LanguageModelRegistry::global(cx),
397 |_, event: &language_model::Event, cx| match event {
398 language_model::Event::ProviderStateChanged
399 | language_model::Event::AddedProvider(_)
400 | language_model::Event::RemovedProvider(_) => {
401 update_active_language_model_from_settings(cx);
402 }
403 _ => {}
404 },
405 )
406 .detach();
407}
408
409fn update_active_language_model_from_settings(cx: &mut AppContext) {
410 let settings = AssistantSettings::get_global(cx);
411 let provider_name = LanguageModelProviderId::from(settings.default_model.provider.clone());
412 let model_id = LanguageModelId::from(settings.default_model.model.clone());
413 let inline_alternatives = settings
414 .inline_alternatives
415 .iter()
416 .map(|alternative| {
417 (
418 LanguageModelProviderId::from(alternative.provider.clone()),
419 LanguageModelId::from(alternative.model.clone()),
420 )
421 })
422 .collect::<Vec<_>>();
423 LanguageModelRegistry::global(cx).update(cx, |registry, cx| {
424 registry.select_active_model(&provider_name, &model_id, cx);
425 registry.select_inline_alternative_models(inline_alternatives, cx);
426 });
427}
428
429fn register_slash_commands(prompt_builder: Option<Arc<PromptBuilder>>, cx: &mut AppContext) {
430 let slash_command_registry = SlashCommandRegistry::global(cx);
431
432 slash_command_registry.register_command(file_command::FileSlashCommand, true);
433 slash_command_registry.register_command(delta_command::DeltaSlashCommand, true);
434 slash_command_registry.register_command(symbols_command::OutlineSlashCommand, true);
435 slash_command_registry.register_command(tab_command::TabSlashCommand, true);
436 slash_command_registry
437 .register_command(cargo_workspace_command::CargoWorkspaceSlashCommand, true);
438 slash_command_registry.register_command(prompt_command::PromptSlashCommand, true);
439 slash_command_registry.register_command(default_command::DefaultSlashCommand, false);
440 slash_command_registry.register_command(terminal_command::TerminalSlashCommand, true);
441 slash_command_registry.register_command(now_command::NowSlashCommand, false);
442 slash_command_registry.register_command(diagnostics_command::DiagnosticsSlashCommand, true);
443 slash_command_registry.register_command(fetch_command::FetchSlashCommand, false);
444 slash_command_registry.register_command(fetch_command::FetchSlashCommand, false);
445
446 if let Some(prompt_builder) = prompt_builder {
447 cx.observe_global::<SettingsStore>({
448 let slash_command_registry = slash_command_registry.clone();
449 let prompt_builder = prompt_builder.clone();
450 move |cx| {
451 if AssistantSettings::get_global(cx).are_live_diffs_enabled(cx) {
452 slash_command_registry.register_command(
453 workflow_command::WorkflowSlashCommand::new(prompt_builder.clone()),
454 true,
455 );
456 } else {
457 slash_command_registry.unregister_command_by_name(WorkflowSlashCommand::NAME);
458 }
459 }
460 })
461 .detach();
462
463 cx.observe_flag::<project_command::ProjectSlashCommandFeatureFlag, _>({
464 let slash_command_registry = slash_command_registry.clone();
465 move |is_enabled, _cx| {
466 if is_enabled {
467 slash_command_registry.register_command(
468 project_command::ProjectSlashCommand::new(prompt_builder.clone()),
469 true,
470 );
471 }
472 }
473 })
474 .detach();
475 }
476
477 cx.observe_flag::<auto_command::AutoSlashCommandFeatureFlag, _>({
478 let slash_command_registry = slash_command_registry.clone();
479 move |is_enabled, _cx| {
480 if is_enabled {
481 // [#auto-staff-ship] TODO remove this when /auto is no longer staff-shipped
482 slash_command_registry.register_command(auto_command::AutoCommand, true);
483 }
484 }
485 })
486 .detach();
487
488 update_slash_commands_from_settings(cx);
489 cx.observe_global::<SettingsStore>(update_slash_commands_from_settings)
490 .detach();
491
492 cx.observe_flag::<search_command::SearchSlashCommandFeatureFlag, _>({
493 let slash_command_registry = slash_command_registry.clone();
494 move |is_enabled, _cx| {
495 if is_enabled {
496 slash_command_registry.register_command(search_command::SearchSlashCommand, true);
497 }
498 }
499 })
500 .detach();
501}
502
503fn update_slash_commands_from_settings(cx: &mut AppContext) {
504 let slash_command_registry = SlashCommandRegistry::global(cx);
505 let settings = SlashCommandSettings::get_global(cx);
506
507 if settings.docs.enabled {
508 slash_command_registry.register_command(docs_command::DocsSlashCommand, true);
509 } else {
510 slash_command_registry.unregister_command(docs_command::DocsSlashCommand);
511 }
512
513 if settings.cargo_workspace.enabled {
514 slash_command_registry
515 .register_command(cargo_workspace_command::CargoWorkspaceSlashCommand, true);
516 } else {
517 slash_command_registry
518 .unregister_command(cargo_workspace_command::CargoWorkspaceSlashCommand);
519 }
520}
521
522fn register_tools(cx: &mut AppContext) {
523 let tool_registry = ToolRegistry::global(cx);
524 tool_registry.register_tool(tools::now_tool::NowTool);
525}
526
527pub fn humanize_token_count(count: usize) -> String {
528 match count {
529 0..=999 => count.to_string(),
530 1000..=9999 => {
531 let thousands = count / 1000;
532 let hundreds = (count % 1000 + 50) / 100;
533 if hundreds == 0 {
534 format!("{}k", thousands)
535 } else if hundreds == 10 {
536 format!("{}k", thousands + 1)
537 } else {
538 format!("{}.{}k", thousands, hundreds)
539 }
540 }
541 _ => format!("{}k", (count + 500) / 1000),
542 }
543}
544
545#[cfg(test)]
546#[ctor::ctor]
547fn init_logger() {
548 if std::env::var("RUST_LOG").is_ok() {
549 env_logger::init();
550 }
551}