1use crate::{
2 AssistantContext, ContextEvent, ContextId, ContextOperation, ContextVersion, SavedContext,
3 SavedContextMetadata,
4};
5use anyhow::{Context as _, Result};
6use assistant_slash_command::{SlashCommandId, SlashCommandWorkingSet};
7use client::{Client, TypedEnvelope, proto, telemetry::Telemetry};
8use clock::ReplicaId;
9use collections::HashMap;
10use context_server::ContextServerId;
11use fs::{Fs, RemoveOptions};
12use futures::StreamExt;
13use fuzzy::StringMatchCandidate;
14use gpui::{App, AppContext as _, AsyncApp, Context, Entity, EventEmitter, Task, WeakEntity};
15use language::LanguageRegistry;
16use paths::contexts_dir;
17use project::{
18 Project,
19 context_server_store::{ContextServerStatus, ContextServerStore},
20};
21use prompt_store::PromptBuilder;
22use regex::Regex;
23use rpc::AnyProtoClient;
24use std::sync::LazyLock;
25use std::{cmp::Reverse, ffi::OsStr, mem, path::Path, sync::Arc, time::Duration};
26use util::{ResultExt, TryFutureExt};
27use zed_env_vars::ZED_STATELESS;
28
29pub(crate) fn init(client: &AnyProtoClient) {
30 client.add_entity_message_handler(ContextStore::handle_advertise_contexts);
31 client.add_entity_request_handler(ContextStore::handle_open_context);
32 client.add_entity_request_handler(ContextStore::handle_create_context);
33 client.add_entity_message_handler(ContextStore::handle_update_context);
34 client.add_entity_request_handler(ContextStore::handle_synchronize_contexts);
35}
36
37#[derive(Clone)]
38pub struct RemoteContextMetadata {
39 pub id: ContextId,
40 pub summary: Option<String>,
41}
42
43pub struct ContextStore {
44 contexts: Vec<ContextHandle>,
45 contexts_metadata: Vec<SavedContextMetadata>,
46 context_server_slash_command_ids: HashMap<ContextServerId, Vec<SlashCommandId>>,
47 host_contexts: Vec<RemoteContextMetadata>,
48 fs: Arc<dyn Fs>,
49 languages: Arc<LanguageRegistry>,
50 slash_commands: Arc<SlashCommandWorkingSet>,
51 telemetry: Arc<Telemetry>,
52 _watch_updates: Task<Option<()>>,
53 client: Arc<Client>,
54 project: Entity<Project>,
55 project_is_shared: bool,
56 client_subscription: Option<client::Subscription>,
57 _project_subscriptions: Vec<gpui::Subscription>,
58 prompt_builder: Arc<PromptBuilder>,
59}
60
61pub enum ContextStoreEvent {
62 ContextCreated(ContextId),
63}
64
65impl EventEmitter<ContextStoreEvent> for ContextStore {}
66
67enum ContextHandle {
68 Weak(WeakEntity<AssistantContext>),
69 Strong(Entity<AssistantContext>),
70}
71
72impl ContextHandle {
73 fn upgrade(&self) -> Option<Entity<AssistantContext>> {
74 match self {
75 ContextHandle::Weak(weak) => weak.upgrade(),
76 ContextHandle::Strong(strong) => Some(strong.clone()),
77 }
78 }
79
80 fn downgrade(&self) -> WeakEntity<AssistantContext> {
81 match self {
82 ContextHandle::Weak(weak) => weak.clone(),
83 ContextHandle::Strong(strong) => strong.downgrade(),
84 }
85 }
86}
87
88impl ContextStore {
89 pub fn new(
90 project: Entity<Project>,
91 prompt_builder: Arc<PromptBuilder>,
92 slash_commands: Arc<SlashCommandWorkingSet>,
93 cx: &mut App,
94 ) -> Task<Result<Entity<Self>>> {
95 let fs = project.read(cx).fs().clone();
96 let languages = project.read(cx).languages().clone();
97 let telemetry = project.read(cx).client().telemetry().clone();
98 cx.spawn(async move |cx| {
99 const CONTEXT_WATCH_DURATION: Duration = Duration::from_millis(100);
100 let (mut events, _) = fs.watch(contexts_dir(), CONTEXT_WATCH_DURATION).await;
101
102 let this = cx.new(|cx: &mut Context<Self>| {
103 let mut this = Self {
104 contexts: Vec::new(),
105 contexts_metadata: Vec::new(),
106 context_server_slash_command_ids: HashMap::default(),
107 host_contexts: Vec::new(),
108 fs,
109 languages,
110 slash_commands,
111 telemetry,
112 _watch_updates: cx.spawn(async move |this, cx| {
113 async move {
114 while events.next().await.is_some() {
115 this.update(cx, |this, cx| this.reload(cx))?.await.log_err();
116 }
117 anyhow::Ok(())
118 }
119 .log_err()
120 .await
121 }),
122 client_subscription: None,
123 _project_subscriptions: vec![
124 cx.subscribe(&project, Self::handle_project_event),
125 ],
126 project_is_shared: false,
127 client: project.read(cx).client(),
128 project: project.clone(),
129 prompt_builder,
130 };
131 this.handle_project_shared(project.clone(), cx);
132 this.synchronize_contexts(cx);
133 this.register_context_server_handlers(cx);
134 this.reload(cx).detach_and_log_err(cx);
135 this
136 })?;
137
138 Ok(this)
139 })
140 }
141
142 #[cfg(any(test, feature = "test-support"))]
143 pub fn fake(project: Entity<Project>, cx: &mut Context<Self>) -> Self {
144 Self {
145 contexts: Default::default(),
146 contexts_metadata: Default::default(),
147 context_server_slash_command_ids: Default::default(),
148 host_contexts: Default::default(),
149 fs: project.read(cx).fs().clone(),
150 languages: project.read(cx).languages().clone(),
151 slash_commands: Arc::default(),
152 telemetry: project.read(cx).client().telemetry().clone(),
153 _watch_updates: Task::ready(None),
154 client: project.read(cx).client(),
155 project,
156 project_is_shared: false,
157 client_subscription: None,
158 _project_subscriptions: Default::default(),
159 prompt_builder: Arc::new(PromptBuilder::new(None).unwrap()),
160 }
161 }
162
163 async fn handle_advertise_contexts(
164 this: Entity<Self>,
165 envelope: TypedEnvelope<proto::AdvertiseContexts>,
166 mut cx: AsyncApp,
167 ) -> Result<()> {
168 this.update(&mut cx, |this, cx| {
169 this.host_contexts = envelope
170 .payload
171 .contexts
172 .into_iter()
173 .map(|context| RemoteContextMetadata {
174 id: ContextId::from_proto(context.context_id),
175 summary: context.summary,
176 })
177 .collect();
178 cx.notify();
179 })
180 }
181
182 async fn handle_open_context(
183 this: Entity<Self>,
184 envelope: TypedEnvelope<proto::OpenContext>,
185 mut cx: AsyncApp,
186 ) -> Result<proto::OpenContextResponse> {
187 let context_id = ContextId::from_proto(envelope.payload.context_id);
188 let operations = this.update(&mut cx, |this, cx| {
189 anyhow::ensure!(
190 !this.project.read(cx).is_via_collab(),
191 "only the host contexts can be opened"
192 );
193
194 let context = this
195 .loaded_context_for_id(&context_id, cx)
196 .context("context not found")?;
197 anyhow::ensure!(
198 context.read(cx).replica_id() == ReplicaId::default(),
199 "context must be opened via the host"
200 );
201
202 anyhow::Ok(
203 context
204 .read(cx)
205 .serialize_ops(&ContextVersion::default(), cx),
206 )
207 })??;
208 let operations = operations.await;
209 Ok(proto::OpenContextResponse {
210 context: Some(proto::Context { operations }),
211 })
212 }
213
214 async fn handle_create_context(
215 this: Entity<Self>,
216 _: TypedEnvelope<proto::CreateContext>,
217 mut cx: AsyncApp,
218 ) -> Result<proto::CreateContextResponse> {
219 let (context_id, operations) = this.update(&mut cx, |this, cx| {
220 anyhow::ensure!(
221 !this.project.read(cx).is_via_collab(),
222 "can only create contexts as the host"
223 );
224
225 let context = this.create(cx);
226 let context_id = context.read(cx).id().clone();
227 cx.emit(ContextStoreEvent::ContextCreated(context_id.clone()));
228
229 anyhow::Ok((
230 context_id,
231 context
232 .read(cx)
233 .serialize_ops(&ContextVersion::default(), cx),
234 ))
235 })??;
236 let operations = operations.await;
237 Ok(proto::CreateContextResponse {
238 context_id: context_id.to_proto(),
239 context: Some(proto::Context { operations }),
240 })
241 }
242
243 async fn handle_update_context(
244 this: Entity<Self>,
245 envelope: TypedEnvelope<proto::UpdateContext>,
246 mut cx: AsyncApp,
247 ) -> Result<()> {
248 this.update(&mut cx, |this, cx| {
249 let context_id = ContextId::from_proto(envelope.payload.context_id);
250 if let Some(context) = this.loaded_context_for_id(&context_id, cx) {
251 let operation_proto = envelope.payload.operation.context("invalid operation")?;
252 let operation = ContextOperation::from_proto(operation_proto)?;
253 context.update(cx, |context, cx| context.apply_ops([operation], cx));
254 }
255 Ok(())
256 })?
257 }
258
259 async fn handle_synchronize_contexts(
260 this: Entity<Self>,
261 envelope: TypedEnvelope<proto::SynchronizeContexts>,
262 mut cx: AsyncApp,
263 ) -> Result<proto::SynchronizeContextsResponse> {
264 this.update(&mut cx, |this, cx| {
265 anyhow::ensure!(
266 !this.project.read(cx).is_via_collab(),
267 "only the host can synchronize contexts"
268 );
269
270 let mut local_versions = Vec::new();
271 for remote_version_proto in envelope.payload.contexts {
272 let remote_version = ContextVersion::from_proto(&remote_version_proto);
273 let context_id = ContextId::from_proto(remote_version_proto.context_id);
274 if let Some(context) = this.loaded_context_for_id(&context_id, cx) {
275 let context = context.read(cx);
276 let operations = context.serialize_ops(&remote_version, cx);
277 local_versions.push(context.version(cx).to_proto(context_id.clone()));
278 let client = this.client.clone();
279 let project_id = envelope.payload.project_id;
280 cx.background_spawn(async move {
281 let operations = operations.await;
282 for operation in operations {
283 client.send(proto::UpdateContext {
284 project_id,
285 context_id: context_id.to_proto(),
286 operation: Some(operation),
287 })?;
288 }
289 anyhow::Ok(())
290 })
291 .detach_and_log_err(cx);
292 }
293 }
294
295 this.advertise_contexts(cx);
296
297 anyhow::Ok(proto::SynchronizeContextsResponse {
298 contexts: local_versions,
299 })
300 })?
301 }
302
303 fn handle_project_shared(&mut self, _: Entity<Project>, cx: &mut Context<Self>) {
304 let is_shared = self.project.read(cx).is_shared();
305 let was_shared = mem::replace(&mut self.project_is_shared, is_shared);
306 if is_shared == was_shared {
307 return;
308 }
309
310 if is_shared {
311 self.contexts.retain_mut(|context| {
312 if let Some(strong_context) = context.upgrade() {
313 *context = ContextHandle::Strong(strong_context);
314 true
315 } else {
316 false
317 }
318 });
319 let remote_id = self.project.read(cx).remote_id().unwrap();
320 self.client_subscription = self
321 .client
322 .subscribe_to_entity(remote_id)
323 .log_err()
324 .map(|subscription| subscription.set_entity(&cx.entity(), &cx.to_async()));
325 self.advertise_contexts(cx);
326 } else {
327 self.client_subscription = None;
328 }
329 }
330
331 fn handle_project_event(
332 &mut self,
333 project: Entity<Project>,
334 event: &project::Event,
335 cx: &mut Context<Self>,
336 ) {
337 match event {
338 project::Event::RemoteIdChanged(_) => {
339 self.handle_project_shared(project, cx);
340 }
341 project::Event::Reshared => {
342 self.advertise_contexts(cx);
343 }
344 project::Event::HostReshared | project::Event::Rejoined => {
345 self.synchronize_contexts(cx);
346 }
347 project::Event::DisconnectedFromHost => {
348 self.contexts.retain_mut(|context| {
349 if let Some(strong_context) = context.upgrade() {
350 *context = ContextHandle::Weak(context.downgrade());
351 strong_context.update(cx, |context, cx| {
352 if context.replica_id() != ReplicaId::default() {
353 context.set_capability(language::Capability::ReadOnly, cx);
354 }
355 });
356 true
357 } else {
358 false
359 }
360 });
361 self.host_contexts.clear();
362 cx.notify();
363 }
364 _ => {}
365 }
366 }
367
368 pub fn unordered_contexts(&self) -> impl Iterator<Item = &SavedContextMetadata> {
369 self.contexts_metadata.iter()
370 }
371
372 pub fn create(&mut self, cx: &mut Context<Self>) -> Entity<AssistantContext> {
373 let context = cx.new(|cx| {
374 AssistantContext::local(
375 self.languages.clone(),
376 Some(self.project.clone()),
377 Some(self.telemetry.clone()),
378 self.prompt_builder.clone(),
379 self.slash_commands.clone(),
380 cx,
381 )
382 });
383 self.register_context(&context, cx);
384 context
385 }
386
387 pub fn create_remote_context(
388 &mut self,
389 cx: &mut Context<Self>,
390 ) -> Task<Result<Entity<AssistantContext>>> {
391 let project = self.project.read(cx);
392 let Some(project_id) = project.remote_id() else {
393 return Task::ready(Err(anyhow::anyhow!("project was not remote")));
394 };
395
396 let replica_id = project.replica_id();
397 let capability = project.capability();
398 let language_registry = self.languages.clone();
399 let project = self.project.clone();
400 let telemetry = self.telemetry.clone();
401 let prompt_builder = self.prompt_builder.clone();
402 let slash_commands = self.slash_commands.clone();
403 let request = self.client.request(proto::CreateContext { project_id });
404 cx.spawn(async move |this, cx| {
405 let response = request.await?;
406 let context_id = ContextId::from_proto(response.context_id);
407 let context_proto = response.context.context("invalid context")?;
408 let context = cx.new(|cx| {
409 AssistantContext::new(
410 context_id.clone(),
411 replica_id,
412 capability,
413 language_registry,
414 prompt_builder,
415 slash_commands,
416 Some(project),
417 Some(telemetry),
418 cx,
419 )
420 })?;
421 let operations = cx
422 .background_spawn(async move {
423 context_proto
424 .operations
425 .into_iter()
426 .map(ContextOperation::from_proto)
427 .collect::<Result<Vec<_>>>()
428 })
429 .await?;
430 context.update(cx, |context, cx| context.apply_ops(operations, cx))?;
431 this.update(cx, |this, cx| {
432 if let Some(existing_context) = this.loaded_context_for_id(&context_id, cx) {
433 existing_context
434 } else {
435 this.register_context(&context, cx);
436 this.synchronize_contexts(cx);
437 context
438 }
439 })
440 })
441 }
442
443 pub fn open_local_context(
444 &mut self,
445 path: Arc<Path>,
446 cx: &Context<Self>,
447 ) -> Task<Result<Entity<AssistantContext>>> {
448 if let Some(existing_context) = self.loaded_context_for_path(&path, cx) {
449 return Task::ready(Ok(existing_context));
450 }
451
452 let fs = self.fs.clone();
453 let languages = self.languages.clone();
454 let project = self.project.clone();
455 let telemetry = self.telemetry.clone();
456 let load = cx.background_spawn({
457 let path = path.clone();
458 async move {
459 let saved_context = fs.load(&path).await?;
460 SavedContext::from_json(&saved_context)
461 }
462 });
463 let prompt_builder = self.prompt_builder.clone();
464 let slash_commands = self.slash_commands.clone();
465
466 cx.spawn(async move |this, cx| {
467 let saved_context = load.await?;
468 let context = cx.new(|cx| {
469 AssistantContext::deserialize(
470 saved_context,
471 path.clone(),
472 languages,
473 prompt_builder,
474 slash_commands,
475 Some(project),
476 Some(telemetry),
477 cx,
478 )
479 })?;
480 this.update(cx, |this, cx| {
481 if let Some(existing_context) = this.loaded_context_for_path(&path, cx) {
482 existing_context
483 } else {
484 this.register_context(&context, cx);
485 context
486 }
487 })
488 })
489 }
490
491 pub fn delete_local_context(
492 &mut self,
493 path: Arc<Path>,
494 cx: &mut Context<Self>,
495 ) -> Task<Result<()>> {
496 let fs = self.fs.clone();
497
498 cx.spawn(async move |this, cx| {
499 fs.remove_file(
500 &path,
501 RemoveOptions {
502 recursive: false,
503 ignore_if_not_exists: true,
504 },
505 )
506 .await?;
507
508 this.update(cx, |this, cx| {
509 this.contexts.retain(|context| {
510 context
511 .upgrade()
512 .and_then(|context| context.read(cx).path())
513 != Some(&path)
514 });
515 this.contexts_metadata
516 .retain(|context| context.path.as_ref() != path.as_ref());
517 })?;
518
519 Ok(())
520 })
521 }
522
523 fn loaded_context_for_path(&self, path: &Path, cx: &App) -> Option<Entity<AssistantContext>> {
524 self.contexts.iter().find_map(|context| {
525 let context = context.upgrade()?;
526 if context.read(cx).path().map(Arc::as_ref) == Some(path) {
527 Some(context)
528 } else {
529 None
530 }
531 })
532 }
533
534 pub fn loaded_context_for_id(
535 &self,
536 id: &ContextId,
537 cx: &App,
538 ) -> Option<Entity<AssistantContext>> {
539 self.contexts.iter().find_map(|context| {
540 let context = context.upgrade()?;
541 if context.read(cx).id() == id {
542 Some(context)
543 } else {
544 None
545 }
546 })
547 }
548
549 pub fn open_remote_context(
550 &mut self,
551 context_id: ContextId,
552 cx: &mut Context<Self>,
553 ) -> Task<Result<Entity<AssistantContext>>> {
554 let project = self.project.read(cx);
555 let Some(project_id) = project.remote_id() else {
556 return Task::ready(Err(anyhow::anyhow!("project was not remote")));
557 };
558
559 if let Some(context) = self.loaded_context_for_id(&context_id, cx) {
560 return Task::ready(Ok(context));
561 }
562
563 let replica_id = project.replica_id();
564 let capability = project.capability();
565 let language_registry = self.languages.clone();
566 let project = self.project.clone();
567 let telemetry = self.telemetry.clone();
568 let request = self.client.request(proto::OpenContext {
569 project_id,
570 context_id: context_id.to_proto(),
571 });
572 let prompt_builder = self.prompt_builder.clone();
573 let slash_commands = self.slash_commands.clone();
574 cx.spawn(async move |this, cx| {
575 let response = request.await?;
576 let context_proto = response.context.context("invalid context")?;
577 let context = cx.new(|cx| {
578 AssistantContext::new(
579 context_id.clone(),
580 replica_id,
581 capability,
582 language_registry,
583 prompt_builder,
584 slash_commands,
585 Some(project),
586 Some(telemetry),
587 cx,
588 )
589 })?;
590 let operations = cx
591 .background_spawn(async move {
592 context_proto
593 .operations
594 .into_iter()
595 .map(ContextOperation::from_proto)
596 .collect::<Result<Vec<_>>>()
597 })
598 .await?;
599 context.update(cx, |context, cx| context.apply_ops(operations, cx))?;
600 this.update(cx, |this, cx| {
601 if let Some(existing_context) = this.loaded_context_for_id(&context_id, cx) {
602 existing_context
603 } else {
604 this.register_context(&context, cx);
605 this.synchronize_contexts(cx);
606 context
607 }
608 })
609 })
610 }
611
612 fn register_context(&mut self, context: &Entity<AssistantContext>, cx: &mut Context<Self>) {
613 let handle = if self.project_is_shared {
614 ContextHandle::Strong(context.clone())
615 } else {
616 ContextHandle::Weak(context.downgrade())
617 };
618 self.contexts.push(handle);
619 self.advertise_contexts(cx);
620 cx.subscribe(context, Self::handle_context_event).detach();
621 }
622
623 fn handle_context_event(
624 &mut self,
625 context: Entity<AssistantContext>,
626 event: &ContextEvent,
627 cx: &mut Context<Self>,
628 ) {
629 let Some(project_id) = self.project.read(cx).remote_id() else {
630 return;
631 };
632
633 match event {
634 ContextEvent::SummaryChanged => {
635 self.advertise_contexts(cx);
636 }
637 ContextEvent::PathChanged { old_path, new_path } => {
638 if let Some(old_path) = old_path.as_ref() {
639 for metadata in &mut self.contexts_metadata {
640 if &metadata.path == old_path {
641 metadata.path = new_path.clone();
642 break;
643 }
644 }
645 }
646 }
647 ContextEvent::Operation(operation) => {
648 let context_id = context.read(cx).id().to_proto();
649 let operation = operation.to_proto();
650 self.client
651 .send(proto::UpdateContext {
652 project_id,
653 context_id,
654 operation: Some(operation),
655 })
656 .log_err();
657 }
658 _ => {}
659 }
660 }
661
662 fn advertise_contexts(&self, cx: &App) {
663 let Some(project_id) = self.project.read(cx).remote_id() else {
664 return;
665 };
666
667 // For now, only the host can advertise their open contexts.
668 if self.project.read(cx).is_via_collab() {
669 return;
670 }
671
672 let contexts = self
673 .contexts
674 .iter()
675 .rev()
676 .filter_map(|context| {
677 let context = context.upgrade()?.read(cx);
678 if context.replica_id() == ReplicaId::default() {
679 Some(proto::ContextMetadata {
680 context_id: context.id().to_proto(),
681 summary: context
682 .summary()
683 .content()
684 .map(|summary| summary.text.clone()),
685 })
686 } else {
687 None
688 }
689 })
690 .collect();
691 self.client
692 .send(proto::AdvertiseContexts {
693 project_id,
694 contexts,
695 })
696 .ok();
697 }
698
699 fn synchronize_contexts(&mut self, cx: &mut Context<Self>) {
700 let Some(project_id) = self.project.read(cx).remote_id() else {
701 return;
702 };
703
704 let contexts = self
705 .contexts
706 .iter()
707 .filter_map(|context| {
708 let context = context.upgrade()?.read(cx);
709 if context.replica_id() != ReplicaId::default() {
710 Some(context.version(cx).to_proto(context.id().clone()))
711 } else {
712 None
713 }
714 })
715 .collect();
716
717 let client = self.client.clone();
718 let request = self.client.request(proto::SynchronizeContexts {
719 project_id,
720 contexts,
721 });
722 cx.spawn(async move |this, cx| {
723 let response = request.await?;
724
725 let mut context_ids = Vec::new();
726 let mut operations = Vec::new();
727 this.read_with(cx, |this, cx| {
728 for context_version_proto in response.contexts {
729 let context_version = ContextVersion::from_proto(&context_version_proto);
730 let context_id = ContextId::from_proto(context_version_proto.context_id);
731 if let Some(context) = this.loaded_context_for_id(&context_id, cx) {
732 context_ids.push(context_id);
733 operations.push(context.read(cx).serialize_ops(&context_version, cx));
734 }
735 }
736 })?;
737
738 let operations = futures::future::join_all(operations).await;
739 for (context_id, operations) in context_ids.into_iter().zip(operations) {
740 for operation in operations {
741 client.send(proto::UpdateContext {
742 project_id,
743 context_id: context_id.to_proto(),
744 operation: Some(operation),
745 })?;
746 }
747 }
748
749 anyhow::Ok(())
750 })
751 .detach_and_log_err(cx);
752 }
753
754 pub fn search(&self, query: String, cx: &App) -> Task<Vec<SavedContextMetadata>> {
755 let metadata = self.contexts_metadata.clone();
756 let executor = cx.background_executor().clone();
757 cx.background_spawn(async move {
758 if query.is_empty() {
759 metadata
760 } else {
761 let candidates = metadata
762 .iter()
763 .enumerate()
764 .map(|(id, metadata)| StringMatchCandidate::new(id, &metadata.title))
765 .collect::<Vec<_>>();
766 let matches = fuzzy::match_strings(
767 &candidates,
768 &query,
769 false,
770 true,
771 100,
772 &Default::default(),
773 executor,
774 )
775 .await;
776
777 matches
778 .into_iter()
779 .map(|mat| metadata[mat.candidate_id].clone())
780 .collect()
781 }
782 })
783 }
784
785 pub fn host_contexts(&self) -> &[RemoteContextMetadata] {
786 &self.host_contexts
787 }
788
789 fn reload(&mut self, cx: &mut Context<Self>) -> Task<Result<()>> {
790 let fs = self.fs.clone();
791 cx.spawn(async move |this, cx| {
792 if *ZED_STATELESS {
793 return Ok(());
794 }
795 fs.create_dir(contexts_dir()).await?;
796
797 let mut paths = fs.read_dir(contexts_dir()).await?;
798 let mut contexts = Vec::<SavedContextMetadata>::new();
799 while let Some(path) = paths.next().await {
800 let path = path?;
801 if path.extension() != Some(OsStr::new("json")) {
802 continue;
803 }
804
805 static ASSISTANT_CONTEXT_REGEX: LazyLock<Regex> =
806 LazyLock::new(|| Regex::new(r" - \d+.zed.json$").unwrap());
807
808 let metadata = fs.metadata(&path).await?;
809 if let Some((file_name, metadata)) = path
810 .file_name()
811 .and_then(|name| name.to_str())
812 .zip(metadata)
813 {
814 // This is used to filter out contexts saved by the new assistant.
815 if !ASSISTANT_CONTEXT_REGEX.is_match(file_name) {
816 continue;
817 }
818
819 if let Some(title) = ASSISTANT_CONTEXT_REGEX
820 .replace(file_name, "")
821 .lines()
822 .next()
823 {
824 contexts.push(SavedContextMetadata {
825 title: title.to_string().into(),
826 path: path.into(),
827 mtime: metadata.mtime.timestamp_for_user().into(),
828 });
829 }
830 }
831 }
832 contexts.sort_unstable_by_key(|context| Reverse(context.mtime));
833
834 this.update(cx, |this, cx| {
835 this.contexts_metadata = contexts;
836 cx.notify();
837 })
838 })
839 }
840
841 fn register_context_server_handlers(&self, cx: &mut Context<Self>) {
842 let context_server_store = self.project.read(cx).context_server_store();
843 cx.subscribe(&context_server_store, Self::handle_context_server_event)
844 .detach();
845
846 // Check for any servers that were already running before the handler was registered
847 for server in context_server_store.read(cx).running_servers() {
848 self.load_context_server_slash_commands(server.id(), context_server_store.clone(), cx);
849 }
850 }
851
852 fn handle_context_server_event(
853 &mut self,
854 context_server_store: Entity<ContextServerStore>,
855 event: &project::context_server_store::Event,
856 cx: &mut Context<Self>,
857 ) {
858 match event {
859 project::context_server_store::Event::ServerStatusChanged { server_id, status } => {
860 match status {
861 ContextServerStatus::Running => {
862 self.load_context_server_slash_commands(
863 server_id.clone(),
864 context_server_store,
865 cx,
866 );
867 }
868 ContextServerStatus::Stopped | ContextServerStatus::Error(_) => {
869 if let Some(slash_command_ids) =
870 self.context_server_slash_command_ids.remove(server_id)
871 {
872 self.slash_commands.remove(&slash_command_ids);
873 }
874 }
875 _ => {}
876 }
877 }
878 }
879 }
880
881 fn load_context_server_slash_commands(
882 &self,
883 server_id: ContextServerId,
884 context_server_store: Entity<ContextServerStore>,
885 cx: &mut Context<Self>,
886 ) {
887 let Some(server) = context_server_store.read(cx).get_running_server(&server_id) else {
888 return;
889 };
890 let slash_command_working_set = self.slash_commands.clone();
891 cx.spawn(async move |this, cx| {
892 let Some(protocol) = server.client() else {
893 return;
894 };
895
896 if protocol.capable(context_server::protocol::ServerCapability::Prompts)
897 && let Some(response) = protocol
898 .request::<context_server::types::requests::PromptsList>(())
899 .await
900 .log_err()
901 {
902 let slash_command_ids = response
903 .prompts
904 .into_iter()
905 .filter(assistant_slash_commands::acceptable_prompt)
906 .map(|prompt| {
907 log::info!("registering context server command: {:?}", prompt.name);
908 slash_command_working_set.insert(Arc::new(
909 assistant_slash_commands::ContextServerSlashCommand::new(
910 context_server_store.clone(),
911 server.id(),
912 prompt,
913 ),
914 ))
915 })
916 .collect::<Vec<_>>();
917
918 this.update(cx, |this, _cx| {
919 this.context_server_slash_command_ids
920 .insert(server_id.clone(), slash_command_ids);
921 })
922 .log_err();
923 }
924 })
925 .detach();
926 }
927}