tenant.openagents/omega

1pub mod extension;
2pub mod registry;
3
4use std::path::Path;
5use std::sync::Arc;
6use std::time::Duration;
7
8use anyhow::{Context as _, Result};
9use collections::{HashMap, HashSet};
10use context_server::oauth::{self, McpOAuthTokenProvider, OAuthDiscovery, OAuthSession};
11use context_server::transport::HttpTransport;
12use context_server::{ContextServer, ContextServerCommand, ContextServerId};
13use credentials_provider::CredentialsProvider;
14use futures::future::Either;
15use futures::{FutureExt as _, StreamExt as _, future::join_all};
16use gpui::{
17    App, AsyncApp, Context, Entity, EventEmitter, Subscription, Task, TaskExt, WeakEntity, actions,
18};
19use http_client::HttpClient;
20use itertools::Itertools;
21use rand::Rng as _;
22use registry::ContextServerDescriptorRegistry;
23use remote::{Interactive, RemoteClient};
24use rpc::{AnyProtoClient, TypedEnvelope, proto};
25use settings::{Settings as _, SettingsLocation, SettingsStore, WorktreeId};
26use util::{ResultExt as _, rel_path::RelPath};
27
28use crate::{
29    DisableAiSettings, Project,
30    project_settings::{ContextServerSettings, OAuthClientSettings, ProjectSettings},
31    worktree_store::{WorktreeStore, WorktreeStoreEvent},
32};
33
34/// Maximum timeout for context server requests
35/// Prevents extremely large timeout values from tying up resources indefinitely.
36const MAX_TIMEOUT_SECS: u64 = 600; // 10 minutes
37
38pub fn init(cx: &mut App) {
39    extension::init(cx);
40}
41
42actions!(
43    context_server,
44    [
45        /// Restarts the context server.
46        Restart
47    ]
48);
49
50#[derive(Debug, Clone, PartialEq, Eq, Hash)]
51pub enum ContextServerStatus {
52    Starting,
53    Running,
54    Stopped,
55    Error(Arc<str>),
56    /// The server returned 401 and OAuth authorization is needed. The UI
57    /// should show an "Authenticate" button.
58    AuthRequired,
59    /// The server has a pre-registered OAuth client_id, but a client_secret
60    /// is needed and not available in settings or the keychain.
61    ClientSecretRequired {
62        error: Option<Arc<str>>,
63    },
64    /// The OAuth browser flow is in progress — the user has been redirected
65    /// to the authorization server and we're waiting for the callback.
66    Authenticating,
67}
68
69impl ContextServerStatus {
70    fn from_state(state: &ContextServerState) -> Self {
71        match state {
72            ContextServerState::Starting { .. } => ContextServerStatus::Starting,
73            ContextServerState::Running { .. } => ContextServerStatus::Running,
74            ContextServerState::Stopped { .. } => ContextServerStatus::Stopped,
75            ContextServerState::Error { error, .. } => ContextServerStatus::Error(error.clone()),
76            ContextServerState::AuthRequired { .. } => ContextServerStatus::AuthRequired,
77            ContextServerState::ClientSecretRequired { error, .. } => {
78                ContextServerStatus::ClientSecretRequired {
79                    error: error.clone(),
80                }
81            }
82            ContextServerState::Authenticating { .. } => ContextServerStatus::Authenticating,
83        }
84    }
85}
86
87enum ContextServerState {
88    Starting {
89        server: Arc<ContextServer>,
90        configuration: Arc<ContextServerConfiguration>,
91        _task: Task<()>,
92    },
93    Running {
94        server: Arc<ContextServer>,
95        configuration: Arc<ContextServerConfiguration>,
96        /// Initiates the OAuth flow if the transport shuts down on an
97        /// authentication challenge; cancelled by any state transition.
98        _transport_watch: Task<()>,
99    },
100    Stopped {
101        server: Arc<ContextServer>,
102        configuration: Arc<ContextServerConfiguration>,
103    },
104    Error {
105        server: Arc<ContextServer>,
106        configuration: Arc<ContextServerConfiguration>,
107        error: Arc<str>,
108    },
109    /// The server requires OAuth authorization before it can be used. The
110    /// `OAuthDiscovery` holds everything needed to start the browser flow.
111    AuthRequired {
112        server: Arc<ContextServer>,
113        configuration: Arc<ContextServerConfiguration>,
114        discovery: Arc<OAuthDiscovery>,
115    },
116    /// A pre-registered client_id is configured but no client_secret was found
117    /// in settings or the keychain.
118    ClientSecretRequired {
119        server: Arc<ContextServer>,
120        configuration: Arc<ContextServerConfiguration>,
121        discovery: Arc<OAuthDiscovery>,
122        error: Option<Arc<str>>,
123    },
124    /// The OAuth browser flow is in progress. The user has been redirected
125    /// to the authorization server and we're waiting for the callback.
126    Authenticating {
127        server: Arc<ContextServer>,
128        configuration: Arc<ContextServerConfiguration>,
129        _task: Task<()>,
130    },
131}
132
133impl ContextServerState {
134    pub fn server(&self) -> Arc<ContextServer> {
135        match self {
136            ContextServerState::Starting { server, .. }
137            | ContextServerState::Running { server, .. }
138            | ContextServerState::Stopped { server, .. }
139            | ContextServerState::Error { server, .. }
140            | ContextServerState::AuthRequired { server, .. }
141            | ContextServerState::ClientSecretRequired { server, .. }
142            | ContextServerState::Authenticating { server, .. } => server.clone(),
143        }
144    }
145
146    pub fn configuration(&self) -> Arc<ContextServerConfiguration> {
147        match self {
148            ContextServerState::Starting { configuration, .. }
149            | ContextServerState::Running { configuration, .. }
150            | ContextServerState::Stopped { configuration, .. }
151            | ContextServerState::Error { configuration, .. }
152            | ContextServerState::AuthRequired { configuration, .. }
153            | ContextServerState::ClientSecretRequired { configuration, .. }
154            | ContextServerState::Authenticating { configuration, .. } => configuration.clone(),
155        }
156    }
157}
158
159#[derive(Debug, PartialEq, Eq)]
160pub enum ContextServerConfiguration {
161    Custom {
162        command: ContextServerCommand,
163        remote: bool,
164    },
165    Extension {
166        command: ContextServerCommand,
167        settings: serde_json::Value,
168        remote: bool,
169    },
170    Http {
171        url: url::Url,
172        headers: HashMap<String, String>,
173        timeout: Option<u64>,
174        oauth: Option<OAuthClientSettings>,
175    },
176}
177
178impl ContextServerConfiguration {
179    pub fn command(&self) -> Option<&ContextServerCommand> {
180        match self {
181            ContextServerConfiguration::Custom { command, .. } => Some(command),
182            ContextServerConfiguration::Extension { command, .. } => Some(command),
183            ContextServerConfiguration::Http { .. } => None,
184        }
185    }
186
187    pub fn has_static_auth_header(&self) -> bool {
188        match self {
189            ContextServerConfiguration::Http { headers, .. } => headers
190                .keys()
191                .any(|k| k.eq_ignore_ascii_case("authorization")),
192            _ => false,
193        }
194    }
195
196    pub fn remote(&self) -> bool {
197        match self {
198            ContextServerConfiguration::Custom { remote, .. } => *remote,
199            ContextServerConfiguration::Extension { remote, .. } => *remote,
200            ContextServerConfiguration::Http { .. } => false,
201        }
202    }
203
204    pub async fn from_settings(
205        settings: ContextServerSettings,
206        id: ContextServerId,
207        registry: Entity<ContextServerDescriptorRegistry>,
208        worktree_store: Entity<WorktreeStore>,
209        cx: &AsyncApp,
210    ) -> Option<Self> {
211        const EXTENSION_COMMAND_TIMEOUT: Duration = Duration::from_secs(30);
212
213        match settings {
214            ContextServerSettings::Stdio {
215                enabled: _,
216                command,
217                remote,
218            } => Some(ContextServerConfiguration::Custom { command, remote }),
219            ContextServerSettings::Extension {
220                enabled: _,
221                settings,
222                remote,
223            } => {
224                let descriptor =
225                    cx.update(|cx| registry.read(cx).context_server_descriptor(&id.0))?;
226
227                let command_future = descriptor.command(worktree_store, cx);
228                let timeout_future = cx.background_executor().timer(EXTENSION_COMMAND_TIMEOUT);
229
230                match futures::future::select(command_future, timeout_future).await {
231                    Either::Left((Ok(command), _)) => Some(ContextServerConfiguration::Extension {
232                        command,
233                        settings,
234                        remote,
235                    }),
236                    Either::Left((Err(e), _)) => {
237                        log::error!(
238                            "Failed to create context server configuration from settings: {e:#}"
239                        );
240                        None
241                    }
242                    Either::Right(_) => {
243                        log::error!(
244                            "Timed out resolving command for extension context server {id}"
245                        );
246                        None
247                    }
248                }
249            }
250            ContextServerSettings::Http {
251                enabled: _,
252                url,
253                headers: auth,
254                timeout,
255                oauth,
256            } => {
257                let url = url::Url::parse(&url).log_err()?;
258                Some(ContextServerConfiguration::Http {
259                    url,
260                    headers: auth,
261                    timeout,
262                    oauth,
263                })
264            }
265        }
266    }
267}
268
269pub type ContextServerFactory =
270    Box<dyn Fn(ContextServerId, Arc<ContextServerConfiguration>) -> Arc<ContextServer>>;
271
272enum ContextServerStoreState {
273    Local {
274        downstream_client: Option<(u64, AnyProtoClient)>,
275        is_headless: bool,
276    },
277    Remote {
278        project_id: u64,
279        upstream_client: Entity<RemoteClient>,
280    },
281}
282
283#[derive(Clone, PartialEq)]
284struct ContextServerSettingsEntry {
285    worktree_id: Option<WorktreeId>,
286    settings: ContextServerSettings,
287}
288
289pub struct ContextServerStore {
290    state: ContextServerStoreState,
291    context_server_settings: HashMap<Arc<str>, ContextServerSettingsEntry>,
292    servers: HashMap<ContextServerId, ContextServerState>,
293    server_ids: Vec<ContextServerId>,
294    worktree_store: Entity<WorktreeStore>,
295    project: Option<WeakEntity<Project>>,
296    registry: Entity<ContextServerDescriptorRegistry>,
297    update_servers_task: Option<Task<Result<()>>>,
298    context_server_factory: Option<ContextServerFactory>,
299    /// The working directory each server was last started with. The working
300    /// directory of a stdio server depends on the resolved project root, which
301    /// can only become available after the server has already started (e.g. a
302    /// worktree is added moments after launch). Tracking it lets
303    /// `maintain_servers` restart a server when its working directory changes,
304    /// since the working directory is not part of `ContextServerConfiguration`.
305    server_working_directories: HashMap<ContextServerId, Option<Arc<Path>>>,
306    needs_server_update: bool,
307    ai_disabled: bool,
308    _subscriptions: Vec<Subscription>,
309}
310
311pub struct ServerStatusChangedEvent {
312    pub server_id: ContextServerId,
313    pub status: ContextServerStatus,
314}
315
316impl EventEmitter<ServerStatusChangedEvent> for ContextServerStore {}
317
318impl ContextServerStore {
319    pub fn local(
320        worktree_store: Entity<WorktreeStore>,
321        weak_project: Option<WeakEntity<Project>>,
322        headless: bool,
323        cx: &mut Context<Self>,
324    ) -> Self {
325        Self::new_internal(
326            !headless,
327            None,
328            ContextServerDescriptorRegistry::default_global(cx),
329            worktree_store,
330            weak_project,
331            ContextServerStoreState::Local {
332                downstream_client: None,
333                is_headless: headless,
334            },
335            cx,
336        )
337    }
338
339    pub fn remote(
340        project_id: u64,
341        upstream_client: Entity<RemoteClient>,
342        worktree_store: Entity<WorktreeStore>,
343        weak_project: Option<WeakEntity<Project>>,
344        cx: &mut Context<Self>,
345    ) -> Self {
346        Self::new_internal(
347            true,
348            None,
349            ContextServerDescriptorRegistry::default_global(cx),
350            worktree_store,
351            weak_project,
352            ContextServerStoreState::Remote {
353                project_id,
354                upstream_client,
355            },
356            cx,
357        )
358    }
359
360    pub fn init_headless(session: &AnyProtoClient) {
361        session.add_entity_request_handler(Self::handle_get_context_server_command);
362    }
363
364    pub fn shared(&mut self, project_id: u64, client: AnyProtoClient) {
365        if let ContextServerStoreState::Local {
366            downstream_client, ..
367        } = &mut self.state
368        {
369            *downstream_client = Some((project_id, client));
370        }
371    }
372
373    pub fn is_remote_project(&self) -> bool {
374        matches!(self.state, ContextServerStoreState::Remote { .. })
375    }
376
377    /// Returns all configured context server ids, excluding the ones that are disabled
378    pub fn configured_server_ids(&self) -> Vec<ContextServerId> {
379        self.context_server_settings
380            .iter()
381            .filter(|(_, entry)| entry.settings.enabled())
382            .map(|(id, _)| ContextServerId(id.clone()))
383            .collect()
384    }
385
386    #[cfg(feature = "test-support")]
387    pub fn test(
388        registry: Entity<ContextServerDescriptorRegistry>,
389        worktree_store: Entity<WorktreeStore>,
390        weak_project: Option<WeakEntity<Project>>,
391        cx: &mut Context<Self>,
392    ) -> Self {
393        Self::new_internal(
394            false,
395            None,
396            registry,
397            worktree_store,
398            weak_project,
399            ContextServerStoreState::Local {
400                downstream_client: None,
401                is_headless: false,
402            },
403            cx,
404        )
405    }
406
407    #[cfg(feature = "test-support")]
408    pub fn test_maintain_server_loop(
409        context_server_factory: Option<ContextServerFactory>,
410        registry: Entity<ContextServerDescriptorRegistry>,
411        worktree_store: Entity<WorktreeStore>,
412        weak_project: Option<WeakEntity<Project>>,
413        cx: &mut Context<Self>,
414    ) -> Self {
415        Self::new_internal(
416            true,
417            context_server_factory,
418            registry,
419            worktree_store,
420            weak_project,
421            ContextServerStoreState::Local {
422                downstream_client: None,
423                is_headless: false,
424            },
425            cx,
426        )
427    }
428
429    #[cfg(feature = "test-support")]
430    pub fn set_context_server_factory(&mut self, factory: ContextServerFactory) {
431        self.context_server_factory = Some(factory);
432    }
433
434    #[cfg(feature = "test-support")]
435    pub fn registry(&self) -> &Entity<ContextServerDescriptorRegistry> {
436        &self.registry
437    }
438
439    #[cfg(feature = "test-support")]
440    pub fn test_start_server(&mut self, server: Arc<ContextServer>, cx: &mut Context<Self>) {
441        let configuration = Arc::new(ContextServerConfiguration::Custom {
442            command: ContextServerCommand {
443                path: "test".into(),
444                args: vec![],
445                env: None,
446                timeout: None,
447            },
448            remote: false,
449        });
450        self.run_server(server, configuration, cx);
451    }
452
453    fn new_internal(
454        maintain_server_loop: bool,
455        context_server_factory: Option<ContextServerFactory>,
456        registry: Entity<ContextServerDescriptorRegistry>,
457        worktree_store: Entity<WorktreeStore>,
458        weak_project: Option<WeakEntity<Project>>,
459        state: ContextServerStoreState,
460        cx: &mut Context<Self>,
461    ) -> Self {
462        let mut subscriptions = vec![cx.observe_global::<SettingsStore>(move |this, cx| {
463            let ai_disabled = DisableAiSettings::get_global(cx).disable_ai;
464            let ai_was_disabled = this.ai_disabled;
465            this.ai_disabled = ai_disabled;
466
467            let settings = Self::resolve_all_context_server_settings(&this.worktree_store, cx);
468            let settings_changed = this.context_server_settings != settings;
469
470            if settings_changed {
471                this.context_server_settings = settings;
472            }
473
474            // When AI is disabled, stop all running servers
475            if ai_disabled {
476                let server_ids: Vec<_> = this.servers.keys().cloned().collect();
477                for id in server_ids {
478                    this.stop_server(&id, cx).log_err();
479                }
480                return;
481            }
482
483            // Trigger updates if AI was re-enabled or settings changed
484            if maintain_server_loop && (ai_was_disabled || settings_changed) {
485                this.available_context_servers_changed(cx);
486            }
487        })];
488
489        if maintain_server_loop {
490            subscriptions.push(cx.observe(&registry, |this, _registry, cx| {
491                if !DisableAiSettings::get_global(cx).disable_ai {
492                    this.available_context_servers_changed(cx);
493                }
494            }));
495            subscriptions.push(cx.subscribe(&worktree_store, |this, _store, event, cx| {
496                if matches!(
497                    event,
498                    WorktreeStoreEvent::WorktreeAdded(_)
499                        | WorktreeStoreEvent::WorktreeRemoved(_, _)
500                ) && !DisableAiSettings::get_global(cx).disable_ai
501                {
502                    this.available_context_servers_changed(cx);
503                }
504            }));
505        }
506
507        let ai_disabled = DisableAiSettings::get_global(cx).disable_ai;
508        let mut this = Self {
509            state,
510            _subscriptions: subscriptions,
511            context_server_settings: Self::resolve_all_context_server_settings(&worktree_store, cx),
512            worktree_store,
513            project: weak_project,
514            registry,
515            needs_server_update: false,
516            ai_disabled,
517            servers: HashMap::default(),
518            server_ids: Default::default(),
519            update_servers_task: None,
520            context_server_factory,
521            server_working_directories: HashMap::default(),
522        };
523        if maintain_server_loop && !DisableAiSettings::get_global(cx).disable_ai {
524            this.available_context_servers_changed(cx);
525        }
526        this
527    }
528
529    pub fn get_server(&self, id: &ContextServerId) -> Option<Arc<ContextServer>> {
530        self.servers.get(id).map(|state| state.server())
531    }
532
533    pub fn get_running_server(&self, id: &ContextServerId) -> Option<Arc<ContextServer>> {
534        if let Some(ContextServerState::Running { server, .. }) = self.servers.get(id) {
535            Some(server.clone())
536        } else {
537            None
538        }
539    }
540
541    pub fn status_for_server(&self, id: &ContextServerId) -> Option<ContextServerStatus> {
542        self.servers.get(id).map(ContextServerStatus::from_state)
543    }
544
545    pub fn configuration_for_server(
546        &self,
547        id: &ContextServerId,
548    ) -> Option<Arc<ContextServerConfiguration>> {
549        self.servers.get(id).map(|state| state.configuration())
550    }
551
552    /// Returns the configured settings for a server, if it is present in the user
553    /// or project settings. This is available regardless of whether the server is
554    /// currently running, unlike [`Self::configuration_for_server`].
555    pub fn settings_for_server(&self, id: &ContextServerId) -> Option<&ContextServerSettings> {
556        self.context_server_settings
557            .get(&id.0)
558            .map(|entry| &entry.settings)
559    }
560
561    /// Returns whether a server is provided by an extension (as opposed to a
562    /// custom Stdio/HTTP server configured directly in settings).
563    ///
564    /// This is derived from the configured settings rather than the runtime
565    /// configuration, so it stays correct even when a custom server is disabled
566    /// or has not been started yet (in which case it has no runtime state).
567    pub fn is_extension_provided(&self, id: &ContextServerId, cx: &App) -> bool {
568        match self.settings_for_server(id) {
569            Some(ContextServerSettings::Stdio { .. } | ContextServerSettings::Http { .. }) => false,
570            Some(ContextServerSettings::Extension { .. }) => true,
571            // No custom settings entry: the server can only originate from an
572            // extension descriptor in the registry.
573            None => self
574                .registry
575                .read(cx)
576                .context_server_descriptor(&id.0)
577                .is_some(),
578        }
579    }
580
581    /// Returns whether a server is enabled.
582    /// Servers with no settings entry only originate from an extension
583    /// descriptor in the registry, and those are enabled by default
584    /// ([`ContextServerSettings::default_extension`]).
585    pub fn is_server_enabled(&self, id: &ContextServerId, cx: &App) -> bool {
586        match self.settings_for_server(id) {
587            Some(settings) => settings.enabled(),
588            None => self
589                .registry
590                .read(cx)
591                .context_server_descriptor(&id.0)
592                .is_some(),
593        }
594    }
595
596    /// Returns a sorted slice of available unique context server IDs. Within the
597    /// slice, context servers which have `mcp-server-` as a prefix in their ID will
598    /// appear after servers that do not have this prefix in their ID.
599    pub fn server_ids(&self) -> &[ContextServerId] {
600        self.server_ids.as_slice()
601    }
602
603    fn populate_server_ids(&mut self, cx: &App) {
604        self.server_ids = self
605            .servers
606            .keys()
607            .cloned()
608            .chain(
609                self.registry
610                    .read(cx)
611                    .context_server_descriptors()
612                    .into_iter()
613                    .map(|(id, _)| ContextServerId(id)),
614            )
615            .chain(
616                self.context_server_settings
617                    .keys()
618                    .map(|id| ContextServerId(id.clone())),
619            )
620            .unique()
621            .sorted_unstable_by(
622                // Sort context servers: ones without mcp-server- prefix first, then prefixed ones
623                |a, b| {
624                    const MCP_PREFIX: &str = "mcp-server-";
625                    match (a.0.strip_prefix(MCP_PREFIX), b.0.strip_prefix(MCP_PREFIX)) {
626                        // If one has mcp-server- prefix and other doesn't, non-mcp comes first
627                        (Some(_), None) => std::cmp::Ordering::Greater,
628                        (None, Some(_)) => std::cmp::Ordering::Less,
629                        // If both have same prefix status, sort by appropriate key
630                        (Some(a), Some(b)) => a.cmp(b),
631                        (None, None) => a.0.cmp(&b.0),
632                    }
633                },
634            )
635            .collect();
636    }
637
638    pub fn running_servers(&self) -> Vec<Arc<ContextServer>> {
639        self.servers
640            .values()
641            .filter_map(|state| {
642                if let ContextServerState::Running { server, .. } = state {
643                    Some(server.clone())
644                } else {
645                    None
646                }
647            })
648            .collect()
649    }
650
651    pub fn start_server(&mut self, server: Arc<ContextServer>, cx: &mut Context<Self>) {
652        cx.spawn(async move |this, cx| {
653            let this = this.upgrade().context("Context server store dropped")?;
654            let id = server.id();
655            let settings_entry = this
656                .update(cx, |this, _| {
657                    this.context_server_settings.get(&id.0).cloned()
658                })
659                .context("Failed to get context server settings")?;
660
661            if !settings_entry.settings.enabled() {
662                return anyhow::Ok(());
663            }
664
665            let (registry, worktree_store) = this.update(cx, |this, _| {
666                (this.registry.clone(), this.worktree_store.clone())
667            });
668            let configuration = ContextServerConfiguration::from_settings(
669                settings_entry.settings,
670                id.clone(),
671                registry,
672                worktree_store,
673                cx,
674            )
675            .await
676            .context("Failed to create context server configuration")?;
677
678            this.update(cx, |this, cx| {
679                this.run_server(server, Arc::new(configuration), cx)
680            });
681            Ok(())
682        })
683        .detach_and_log_err(cx);
684    }
685
686    pub fn stop_server(&mut self, id: &ContextServerId, cx: &mut Context<Self>) -> Result<()> {
687        if matches!(
688            self.servers.get(id),
689            Some(ContextServerState::Stopped { .. })
690        ) {
691            return Ok(());
692        }
693
694        let state = self
695            .servers
696            .remove(id)
697            .context("Context server not found")?;
698
699        let server = state.server();
700        let configuration = state.configuration();
701        let result = server.stop();
702        drop(state);
703
704        self.update_server_state(
705            id.clone(),
706            ContextServerState::Stopped {
707                configuration,
708                server,
709            },
710            cx,
711        );
712
713        result
714    }
715
716    fn run_server(
717        &mut self,
718        server: Arc<ContextServer>,
719        configuration: Arc<ContextServerConfiguration>,
720        cx: &mut Context<Self>,
721    ) {
722        let id = server.id();
723        if matches!(
724            self.servers.get(&id),
725            Some(
726                ContextServerState::Starting { .. }
727                    | ContextServerState::Running { .. }
728                    | ContextServerState::Authenticating { .. },
729            )
730        ) {
731            self.stop_server(&id, cx).log_err();
732        }
733        let task = cx.spawn({
734            let id = server.id();
735            let server = server.clone();
736            let configuration = configuration.clone();
737
738            async move |this, cx| {
739                let new_state = match server.clone().start(cx).await {
740                    Ok(_) => {
741                        debug_assert!(server.client().is_some());
742                        let _transport_watch =
743                            Self::watch_transport_shutdown(this.clone(), server.clone(), cx);
744                        ContextServerState::Running {
745                            server,
746                            configuration,
747                            _transport_watch,
748                        }
749                    }
750                    Err(err) => resolve_start_failure(&id, err, server, configuration, cx).await,
751                };
752                this.update(cx, |this, cx| {
753                    this.update_server_state(id.clone(), new_state, cx)
754                })
755                .log_err();
756            }
757        });
758
759        self.update_server_state(
760            id.clone(),
761            ContextServerState::Starting {
762                configuration,
763                _task: task,
764                server,
765            },
766            cx,
767        );
768    }
769
770    /// Watches a running server's transport and initiates the OAuth flow if it
771    /// shuts down on an authentication challenge.
772    ///
773    /// MCP servers may accept `initialize` unauthenticated and only send a 401
774    /// with a `WWW-Authenticate` challenge on a later request or notification.
775    /// The HTTP transport records the challenge, and the failed send tears
776    /// down the client's output loop. Observing that shutdown — rather than
777    /// relying on some request to carry a typed error back to a caller — is
778    /// what lets any post-initialize 401 move the server into `AuthRequired`
779    /// instead of leaving it `Running` with a dead client.
780    fn watch_transport_shutdown(
781        this: WeakEntity<Self>,
782        server: Arc<ContextServer>,
783        cx: &mut AsyncApp,
784    ) -> Task<()> {
785        let Some(shutdown) = server
786            .client()
787            .and_then(|client| client.wait_for_shutdown())
788        else {
789            return Task::ready(());
790        };
791        cx.spawn(async move |cx| {
792            let Some(www_authenticate) = shutdown.await else {
793                // Non-auth transport deaths leave the server state untouched,
794                // as they did before this watch existed.
795                return;
796            };
797            this.update(cx, |this, cx| {
798                this.handle_auth_challenge(server, www_authenticate, cx);
799            })
800            .log_err();
801        })
802    }
803
804    fn handle_auth_challenge(
805        &mut self,
806        server: Arc<ContextServer>,
807        www_authenticate: oauth::WwwAuthenticate,
808        cx: &mut Context<Self>,
809    ) {
810        let id = server.id();
811
812        // Act only if this exact server is still the one we consider running.
813        // If the state has changed since the challenge was recorded, whoever
814        // changed it owns the lifecycle now.
815        let Some(ContextServerState::Running {
816            server: running_server,
817            configuration,
818            ..
819        }) = self.servers.get(&id)
820        else {
821            return;
822        };
823        if !Arc::ptr_eq(running_server, &server) {
824            return;
825        }
826        let configuration = configuration.clone();
827
828        log::info!("{id} received 401 after initialization; initiating OAuth authorization");
829
830        // The 401 already tore down the client's output loop. Stop the dead
831        // client, then resolve auth using the captured `WWW-Authenticate` — do
832        // not restart via `run_server`, as a fresh `initialize` would succeed
833        // and lose the challenge.
834        server.stop().log_err();
835
836        let task = cx.spawn({
837            let id = id.clone();
838            let server = server.clone();
839            let configuration = configuration.clone();
840            async move |this, cx| {
841                let new_state =
842                    resolve_auth_required(&id, &www_authenticate, server, configuration, cx).await;
843                this.update(cx, |this, cx| {
844                    this.update_server_state(id.clone(), new_state, cx)
845                })
846                .log_err();
847            }
848        });
849
850        self.update_server_state(
851            id,
852            ContextServerState::Starting {
853                configuration,
854                _task: task,
855                server,
856            },
857            cx,
858        );
859    }
860
861    fn remove_server(&mut self, id: &ContextServerId, cx: &mut Context<Self>) -> Result<()> {
862        let state = self
863            .servers
864            .remove(id)
865            .context("Context server not found")?;
866        self.server_working_directories.remove(id);
867
868        if let ContextServerConfiguration::Http { url, .. } = state.configuration().as_ref() {
869            let server_url = url.clone();
870            let id = id.clone();
871            cx.spawn(async move |_this, cx| {
872                let credentials_provider = cx.update(|cx| zed_credentials_provider::global(cx));
873                if let Err(err) = Self::clear_session(&credentials_provider, &server_url, &cx).await
874                {
875                    log::warn!("{} failed to clear OAuth session on removal: {}", id, err);
876                }
877            })
878            .detach();
879        }
880
881        drop(state);
882        cx.emit(ServerStatusChangedEvent {
883            server_id: id.clone(),
884            status: ContextServerStatus::Stopped,
885        });
886        cx.notify();
887        Ok(())
888    }
889
890    /// The project root a locally-spawned stdio server should use as its working
891    /// directory: the active project directory, falling back to the first visible
892    /// worktree. Resolves to `None` before any worktree is available.
893    fn resolve_root_path(&self, cx: &App) -> Option<Arc<Path>> {
894        self.project
895            .as_ref()
896            .and_then(|project| {
897                project
898                    .read_with(cx, |project, cx| project.active_project_directory(cx))
899                    .ok()
900                    .flatten()
901            })
902            .or_else(|| {
903                self.worktree_store.read_with(cx, |store, cx| {
904                    store.visible_worktrees(cx).fold(None, |acc, item| {
905                        if acc.is_none() {
906                            item.read(cx).root_dir()
907                        } else {
908                            acc
909                        }
910                    })
911                })
912            })
913    }
914
915    pub async fn create_context_server(
916        this: WeakEntity<Self>,
917        id: ContextServerId,
918        configuration: Arc<ContextServerConfiguration>,
919        cx: &mut AsyncApp,
920    ) -> Result<(Arc<ContextServer>, Arc<ContextServerConfiguration>)> {
921        let remote = configuration.remote();
922        let needs_remote_command = match configuration.as_ref() {
923            ContextServerConfiguration::Custom { .. }
924            | ContextServerConfiguration::Extension { .. } => remote,
925            ContextServerConfiguration::Http { .. } => false,
926        };
927
928        let (remote_state, is_remote_project) = this.update(cx, |this, _| {
929            let remote_state = match &this.state {
930                ContextServerStoreState::Remote {
931                    project_id,
932                    upstream_client,
933                } if needs_remote_command => Some((*project_id, upstream_client.clone())),
934                _ => None,
935            };
936            (remote_state, this.is_remote_project())
937        })?;
938
939        let root_path: Option<Arc<Path>> =
940            this.update(cx, |this, cx| this.resolve_root_path(cx))?;
941
942        let configuration = if let Some((project_id, upstream_client)) = remote_state {
943            let root_dir = root_path.as_ref().map(|p| p.display().to_string());
944
945            let response = upstream_client
946                .update(cx, |client, _| {
947                    client
948                        .proto_client()
949                        .request(proto::GetContextServerCommand {
950                            project_id,
951                            server_id: id.0.to_string(),
952                            root_dir: root_dir.clone(),
953                        })
954                })
955                .await?;
956
957            let remote_command = upstream_client.update(cx, |client, _| {
958                client.build_command(
959                    Some(response.path),
960                    &response.args,
961                    &response.env.into_iter().collect(),
962                    root_dir,
963                    None,
964                    Interactive::Yes,
965                )
966            })?;
967
968            let command = ContextServerCommand {
969                path: remote_command.program.into(),
970                args: remote_command.args,
971                env: Some(remote_command.env.into_iter().collect()),
972                timeout: None,
973            };
974
975            Arc::new(ContextServerConfiguration::Custom { command, remote })
976        } else {
977            configuration
978        };
979
980        if let Some(server) = this.update(cx, |this, _| {
981            this.context_server_factory
982                .as_ref()
983                .map(|factory| factory(id.clone(), configuration.clone()))
984        })? {
985            return Ok((server, configuration));
986        }
987
988        let cached_token_provider: Option<Arc<dyn oauth::OAuthTokenProvider>> =
989            if let ContextServerConfiguration::Http { url, .. } = configuration.as_ref() {
990                if configuration.has_static_auth_header() {
991                    None
992                } else {
993                    let credentials_provider = cx.update(|cx| zed_credentials_provider::global(cx));
994                    let http_client = cx.update(|cx| cx.http_client());
995
996                    match Self::load_session(&credentials_provider, url, &cx).await {
997                        Ok(Some(session)) => {
998                            log::info!("{} loaded cached OAuth session from keychain", id);
999                            Some(Self::create_oauth_token_provider(
1000                                &id,
1001                                url,
1002                                session,
1003                                http_client,
1004                                credentials_provider,
1005                                cx,
1006                            ))
1007                        }
1008                        Ok(None) => None,
1009                        Err(err) => {
1010                            log::warn!("{} failed to load cached OAuth session: {}", id, err);
1011                            None
1012                        }
1013                    }
1014                }
1015            } else {
1016                None
1017            };
1018
1019        let server: Arc<ContextServer> = this.update(cx, |this, cx| {
1020            let global_timeout = this.timeout_for_server(&id, cx);
1021
1022            match configuration.as_ref() {
1023                ContextServerConfiguration::Http {
1024                    url,
1025                    headers,
1026                    timeout,
1027                    oauth: _,
1028                } => {
1029                    let transport = HttpTransport::new_with_token_provider(
1030                        cx.http_client(),
1031                        url.to_string(),
1032                        headers.clone(),
1033                        cx.background_executor().clone(),
1034                        cached_token_provider.clone(),
1035                    );
1036                    anyhow::Ok(Arc::new(ContextServer::new_with_timeout(
1037                        id,
1038                        Arc::new(transport),
1039                        Some(Duration::from_secs(
1040                            timeout.unwrap_or(global_timeout).min(MAX_TIMEOUT_SECS),
1041                        )),
1042                    )))
1043                }
1044                _ => {
1045                    let mut command = configuration
1046                        .command()
1047                        .context("Missing command configuration for stdio context server")?
1048                        .clone();
1049                    command.timeout = Some(
1050                        command
1051                            .timeout
1052                            .unwrap_or(global_timeout)
1053                            .min(MAX_TIMEOUT_SECS),
1054                    );
1055
1056                    // Don't pass remote paths as working directory for locally-spawned processes
1057                    let working_directory = if is_remote_project { None } else { root_path };
1058                    anyhow::Ok(Arc::new(ContextServer::stdio(
1059                        id,
1060                        command,
1061                        working_directory,
1062                    )))
1063                }
1064            }
1065        })??;
1066
1067        Ok((server, configuration))
1068    }
1069
1070    async fn handle_get_context_server_command(
1071        this: Entity<Self>,
1072        envelope: TypedEnvelope<proto::GetContextServerCommand>,
1073        mut cx: AsyncApp,
1074    ) -> Result<proto::ContextServerCommand> {
1075        let server_id = ContextServerId(envelope.payload.server_id.into());
1076
1077        let (settings_entry, registry, worktree_store) =
1078            this.update(&mut cx, |this, inner_cx| {
1079                let ContextServerStoreState::Local {
1080                    is_headless: true, ..
1081                } = &this.state
1082                else {
1083                    anyhow::bail!(
1084                        "unexpected GetContextServerCommand request in a non-local project"
1085                    );
1086                };
1087
1088                let settings = this
1089                    .context_server_settings
1090                    .get(&server_id.0)
1091                    .cloned()
1092                    .or_else(|| {
1093                        this.registry
1094                            .read(inner_cx)
1095                            .context_server_descriptor(&server_id.0)
1096                            .map(|_| ContextServerSettingsEntry {
1097                                worktree_id: None,
1098                                settings: ContextServerSettings::default_extension(),
1099                            })
1100                    })
1101                    .with_context(|| format!("context server `{}` not found", server_id))?;
1102
1103                anyhow::Ok((settings, this.registry.clone(), this.worktree_store.clone()))
1104            })?;
1105
1106        let configuration = ContextServerConfiguration::from_settings(
1107            settings_entry.settings,
1108            server_id.clone(),
1109            registry,
1110            worktree_store,
1111            &cx,
1112        )
1113        .await
1114        .with_context(|| format!("failed to build configuration for `{}`", server_id))?;
1115
1116        let command = configuration
1117            .command()
1118            .context("context server has no command (HTTP servers don't need RPC)")?;
1119
1120        Ok(proto::ContextServerCommand {
1121            path: command.path.display().to_string(),
1122            args: command.args.clone(),
1123            env: command
1124                .env
1125                .clone()
1126                .map(|env| env.into_iter().collect())
1127                .unwrap_or_default(),
1128        })
1129    }
1130
1131    /// Merges context server settings from all visible worktrees so that servers defined
1132    /// in any project folder in a multi-root workspace are picked up.
1133    fn resolve_all_context_server_settings(
1134        worktree_store: &Entity<WorktreeStore>,
1135        cx: &App,
1136    ) -> HashMap<Arc<str>, ContextServerSettingsEntry> {
1137        let mut merged = HashMap::default();
1138        for worktree in worktree_store.read(cx).visible_worktrees(cx) {
1139            let worktree_id = worktree.read(cx).id();
1140            let location = settings::SettingsLocation {
1141                worktree_id,
1142                path: RelPath::empty(),
1143            };
1144            for (id, settings) in &ProjectSettings::get(Some(location), cx).context_servers {
1145                merged
1146                    .entry(id.clone())
1147                    .or_insert_with(|| ContextServerSettingsEntry {
1148                        worktree_id: Some(worktree_id),
1149                        settings: settings.clone(),
1150                    });
1151            }
1152        }
1153        merged
1154    }
1155
1156    fn create_oauth_token_provider(
1157        id: &ContextServerId,
1158        server_url: &url::Url,
1159        session: OAuthSession,
1160        http_client: Arc<dyn HttpClient>,
1161        credentials_provider: Arc<dyn CredentialsProvider>,
1162        cx: &mut AsyncApp,
1163    ) -> Arc<dyn oauth::OAuthTokenProvider> {
1164        let (token_refresh_tx, mut token_refresh_rx) = futures::channel::mpsc::unbounded();
1165        let id = id.clone();
1166        let server_url = server_url.clone();
1167
1168        cx.spawn(async move |cx| {
1169            while let Some(refreshed_session) = token_refresh_rx.next().await {
1170                if let Err(err) =
1171                    Self::store_session(&credentials_provider, &server_url, &refreshed_session, &cx)
1172                        .await
1173                {
1174                    log::warn!("{} failed to persist refreshed OAuth session: {}", id, err);
1175                }
1176            }
1177            log::debug!("{} OAuth session persistence task ended", id);
1178        })
1179        .detach();
1180
1181        Arc::new(McpOAuthTokenProvider::new(
1182            session,
1183            http_client,
1184            Some(token_refresh_tx),
1185        ))
1186    }
1187
1188    fn timeout_for_server(&self, id: &ContextServerId, cx: &App) -> u64 {
1189        let worktree_id = self
1190            .context_server_settings
1191            .get(&id.0)
1192            .as_ref()
1193            .and_then(|entry| entry.worktree_id);
1194
1195        ProjectSettings::get(
1196            worktree_id.map(|id| SettingsLocation {
1197                worktree_id: id,
1198                path: &RelPath::empty(),
1199            }),
1200            cx,
1201        )
1202        .context_server_timeout
1203    }
1204
1205    /// Initiate the OAuth browser flow for a server in the `AuthRequired` state.
1206    ///
1207    /// This starts a loopback HTTP callback server on an ephemeral port, builds
1208    /// the authorization URL, opens the user's browser, waits for the callback,
1209    /// exchanges the code for tokens, persists them in the keychain, and restarts
1210    /// the server with the new token provider.
1211    pub fn authenticate_server(
1212        &mut self,
1213        id: &ContextServerId,
1214        cx: &mut Context<Self>,
1215    ) -> Result<()> {
1216        let state = self.servers.get(id).context("Context server not found")?;
1217        let global_timeout = self.timeout_for_server(id, cx);
1218
1219        let (discovery, server, configuration) = match state {
1220            ContextServerState::AuthRequired {
1221                discovery,
1222                server,
1223                configuration,
1224            } => (discovery.clone(), server.clone(), configuration.clone()),
1225            _ => anyhow::bail!("Server is not in AuthRequired state"),
1226        };
1227
1228        let needs_keychain_check = match configuration.as_ref() {
1229            ContextServerConfiguration::Http {
1230                url,
1231                oauth: Some(oauth_settings),
1232                ..
1233            } if oauth_settings.client_secret.is_none() => Some(url.clone()),
1234            _ => None,
1235        };
1236
1237        let id = id.clone();
1238
1239        let task = cx.spawn({
1240            let id = id.clone();
1241            let server = server.clone();
1242            let configuration = configuration.clone();
1243            async move |this, cx| {
1244                if let Some(server_url) = needs_keychain_check {
1245                    let credentials_provider = cx.update(|cx| zed_credentials_provider::global(cx));
1246                    let has_keychain_secret =
1247                        Self::load_client_secret(&credentials_provider, &server_url, cx)
1248                            .await
1249                            .ok()
1250                            .flatten()
1251                            .is_some();
1252
1253                    if !has_keychain_secret {
1254                        this.update(cx, |this, cx| {
1255                            this.update_server_state(
1256                                id.clone(),
1257                                ContextServerState::ClientSecretRequired {
1258                                    server,
1259                                    configuration,
1260                                    discovery,
1261                                    error: None,
1262                                },
1263                                cx,
1264                            );
1265                        })
1266                        .log_err();
1267                        return;
1268                    }
1269                }
1270
1271                let result = Self::run_oauth_flow(
1272                    this.clone(),
1273                    id.clone(),
1274                    discovery.clone(),
1275                    configuration.clone(),
1276                    global_timeout,
1277                    cx,
1278                )
1279                .await;
1280
1281                if let Err(err) = &result {
1282                    log::error!("{} OAuth authentication failed: {:?}", id, err);
1283                    this.update(cx, |this, cx| {
1284                        this.update_server_state(
1285                            id.clone(),
1286                            ContextServerState::Error {
1287                                server,
1288                                configuration,
1289                                error: format!("{err:#}").into(),
1290                            },
1291                            cx,
1292                        )
1293                    })
1294                    .log_err();
1295                }
1296            }
1297        });
1298
1299        self.update_server_state(
1300            id,
1301            ContextServerState::Authenticating {
1302                server,
1303                configuration,
1304                _task: task,
1305            },
1306            cx,
1307        );
1308
1309        Ok(())
1310    }
1311
1312    /// Store the client secret and proceed with authentication.
1313    pub fn submit_client_secret(
1314        &mut self,
1315        id: &ContextServerId,
1316        secret: String,
1317        cx: &mut Context<Self>,
1318    ) -> Result<()> {
1319        let state = self.servers.get(id).context("Context server not found")?;
1320        let global_timeout = self.timeout_for_server(id, cx);
1321
1322        let (server, configuration, discovery) = match state {
1323            ContextServerState::ClientSecretRequired {
1324                server,
1325                configuration,
1326                discovery,
1327                ..
1328            } => (server.clone(), configuration.clone(), discovery.clone()),
1329            _ => anyhow::bail!("Server is not in ClientSecretRequired state"),
1330        };
1331
1332        let server_url = match configuration.as_ref() {
1333            ContextServerConfiguration::Http { url, .. } => url.clone(),
1334            _ => anyhow::bail!("OAuth only supported for HTTP servers"),
1335        };
1336
1337        let id = id.clone();
1338
1339        let task = cx.spawn({
1340            let id = id.clone();
1341            let server = server.clone();
1342            let configuration = configuration.clone();
1343            async move |this, cx| {
1344                // Store the secret if non-empty (empty means public client / skip).
1345                if !secret.is_empty() {
1346                    let credentials_provider = cx.update(|cx| zed_credentials_provider::global(cx));
1347                    if let Err(err) =
1348                        Self::store_client_secret(&credentials_provider, &server_url, &secret, cx)
1349                            .await
1350                    {
1351                        log::error!(
1352                            "{} failed to store client secret in keychain: {:?}",
1353                            id,
1354                            err
1355                        );
1356                    }
1357                }
1358
1359                let result = Self::run_oauth_flow(
1360                    this.clone(),
1361                    id.clone(),
1362                    discovery.clone(),
1363                    configuration.clone(),
1364                    global_timeout,
1365                    cx,
1366                )
1367                .await;
1368
1369                if let Err(err) = &result {
1370                    log::error!("{} OAuth authentication failed: {:?}", id, err);
1371
1372                    let is_bad_client_credentials = err
1373                        .downcast_ref::<oauth::OAuthTokenError>()
1374                        .is_some_and(|e| e.error == "unauthorized_client");
1375
1376                    if is_bad_client_credentials {
1377                        // Clear the bad secret from the keychain so the user
1378                        // gets a fresh prompt.
1379                        let credentials_provider =
1380                            cx.update(|cx| zed_credentials_provider::global(cx));
1381                        Self::clear_client_secret(&credentials_provider, &server_url, cx)
1382                            .await
1383                            .log_err();
1384
1385                        this.update(cx, |this, cx| {
1386                            this.update_server_state(
1387                                id.clone(),
1388                                ContextServerState::ClientSecretRequired {
1389                                    server,
1390                                    configuration,
1391                                    discovery,
1392                                    error: Some(format!("{err:#}").into()),
1393                                },
1394                                cx,
1395                            );
1396                        })
1397                        .log_err();
1398                    } else {
1399                        this.update(cx, |this, cx| {
1400                            this.update_server_state(
1401                                id.clone(),
1402                                ContextServerState::Error {
1403                                    server,
1404                                    configuration,
1405                                    error: format!("{err:#}").into(),
1406                                },
1407                                cx,
1408                            )
1409                        })
1410                        .log_err();
1411                    }
1412                }
1413            }
1414        });
1415
1416        self.update_server_state(
1417            id,
1418            ContextServerState::Authenticating {
1419                server,
1420                configuration,
1421                _task: task,
1422            },
1423            cx,
1424        );
1425
1426        Ok(())
1427    }
1428
1429    async fn run_oauth_flow(
1430        this: WeakEntity<Self>,
1431        id: ContextServerId,
1432        discovery: Arc<OAuthDiscovery>,
1433        configuration: Arc<ContextServerConfiguration>,
1434        global_timeout: u64,
1435        cx: &mut AsyncApp,
1436    ) -> Result<()> {
1437        let resource = oauth::canonical_server_uri(&discovery.resource_metadata.resource);
1438        let pkce = oauth::generate_pkce_challenge();
1439
1440        let mut state_bytes = [0u8; 32];
1441        rand::rng().fill(&mut state_bytes);
1442        let state_param: String = state_bytes.iter().map(|b| format!("{:02x}", b)).collect();
1443
1444        // Start a loopback HTTP server on an ephemeral port. The redirect URI
1445        // includes this port so the browser sends the callback directly to our
1446        // process.
1447        let (redirect_uri, callback_rx) =
1448            oauth::start_callback_server().context("Failed to start OAuth callback server")?;
1449
1450        let http_client = cx.update(|cx| cx.http_client());
1451        let credentials_provider = cx.update(|cx| zed_credentials_provider::global(cx));
1452        let server_url = match configuration.as_ref() {
1453            ContextServerConfiguration::Http { url, .. } => url.clone(),
1454            _ => anyhow::bail!("OAuth authentication only supported for HTTP servers"),
1455        };
1456
1457        let client_registration = match configuration.as_ref() {
1458            ContextServerConfiguration::Http {
1459                url,
1460                oauth: Some(oauth_settings),
1461                ..
1462            } => {
1463                // Pre-registered client. Resolve the secret from settings, then keychain.
1464                let client_secret = if oauth_settings.client_secret.is_some() {
1465                    oauth_settings.client_secret.clone()
1466                } else {
1467                    Self::load_client_secret(&credentials_provider, url, cx)
1468                        .await
1469                        .ok()
1470                        .flatten()
1471                };
1472                oauth::OAuthClientRegistration {
1473                    client_id: oauth_settings.client_id.clone(),
1474                    client_secret,
1475                }
1476            }
1477            _ => oauth::resolve_client_registration(&http_client, &discovery, &redirect_uri)
1478                .await
1479                .context("Failed to resolve OAuth client registration")?,
1480        };
1481
1482        let auth_url = oauth::build_authorization_url(
1483            &discovery.auth_server_metadata,
1484            &client_registration.client_id,
1485            &redirect_uri,
1486            &discovery.scopes,
1487            &resource,
1488            &pkce,
1489            &state_param,
1490        );
1491
1492        cx.update(|cx| cx.open_url(auth_url.as_str()));
1493
1494        let callback = callback_rx
1495            .await
1496            .context("OAuth callback server received an invalid request")?;
1497
1498        if callback.state != state_param {
1499            anyhow::bail!("OAuth state parameter mismatch (possible CSRF)");
1500        }
1501
1502        let tokens = oauth::exchange_code(
1503            &http_client,
1504            &discovery.auth_server_metadata,
1505            &callback.code,
1506            &client_registration.client_id,
1507            &redirect_uri,
1508            &pkce.verifier,
1509            &resource,
1510            client_registration.client_secret.as_deref(),
1511        )
1512        .await
1513        .context("Failed to exchange authorization code for tokens")?;
1514
1515        let session = OAuthSession {
1516            token_endpoint: discovery.auth_server_metadata.token_endpoint.clone(),
1517            resource: discovery.resource_metadata.resource.clone(),
1518            client_registration,
1519            tokens,
1520        };
1521
1522        Self::store_session(&credentials_provider, &server_url, &session, cx)
1523            .await
1524            .context("Failed to persist OAuth session in keychain")?;
1525
1526        let token_provider = Self::create_oauth_token_provider(
1527            &id,
1528            &server_url,
1529            session,
1530            http_client.clone(),
1531            credentials_provider,
1532            cx,
1533        );
1534
1535        let new_server = this.update(cx, |_this, cx| match configuration.as_ref() {
1536            ContextServerConfiguration::Http {
1537                url,
1538                headers,
1539                timeout,
1540                oauth: _,
1541            } => {
1542                let transport = HttpTransport::new_with_token_provider(
1543                    http_client.clone(),
1544                    url.to_string(),
1545                    headers.clone(),
1546                    cx.background_executor().clone(),
1547                    Some(token_provider.clone()),
1548                );
1549                Ok(Arc::new(ContextServer::new_with_timeout(
1550                    id.clone(),
1551                    Arc::new(transport),
1552                    Some(Duration::from_secs(
1553                        timeout.unwrap_or(global_timeout).min(MAX_TIMEOUT_SECS),
1554                    )),
1555                )))
1556            }
1557            _ => anyhow::bail!("OAuth authentication only supported for HTTP servers"),
1558        })??;
1559
1560        this.update(cx, |this, cx| {
1561            this.run_server(new_server, configuration, cx);
1562        })?;
1563
1564        Ok(())
1565    }
1566
1567    /// Store the full OAuth session in the system keychain, keyed by the
1568    /// server's canonical URI.
1569    async fn store_session(
1570        credentials_provider: &Arc<dyn CredentialsProvider>,
1571        server_url: &url::Url,
1572        session: &OAuthSession,
1573        cx: &AsyncApp,
1574    ) -> Result<()> {
1575        let key = Self::keychain_key(server_url);
1576        let json = serde_json::to_string(session)?;
1577        credentials_provider
1578            .write_credentials(&key, "mcp-oauth", json.as_bytes(), cx)
1579            .await
1580    }
1581
1582    /// Load the full OAuth session from the system keychain for the given
1583    /// server URL.
1584    async fn load_session(
1585        credentials_provider: &Arc<dyn CredentialsProvider>,
1586        server_url: &url::Url,
1587        cx: &AsyncApp,
1588    ) -> Result<Option<OAuthSession>> {
1589        let key = Self::keychain_key(server_url);
1590        match credentials_provider.read_credentials(&key, cx).await? {
1591            Some((_username, password_bytes)) => {
1592                let session: OAuthSession = serde_json::from_slice(&password_bytes)?;
1593                Ok(Some(session))
1594            }
1595            None => Ok(None),
1596        }
1597    }
1598
1599    /// Clear the stored OAuth session from the system keychain.
1600    async fn clear_session(
1601        credentials_provider: &Arc<dyn CredentialsProvider>,
1602        server_url: &url::Url,
1603        cx: &AsyncApp,
1604    ) -> Result<()> {
1605        let key = Self::keychain_key(server_url);
1606        credentials_provider.delete_credentials(&key, cx).await
1607    }
1608
1609    fn keychain_key(server_url: &url::Url) -> String {
1610        format!("mcp-oauth:{}", oauth::canonical_server_uri(server_url))
1611    }
1612
1613    fn client_secret_keychain_key(server_url: &url::Url) -> String {
1614        format!(
1615            "mcp-oauth-client-secret:{}",
1616            oauth::canonical_server_uri(server_url)
1617        )
1618    }
1619
1620    async fn load_client_secret(
1621        credentials_provider: &Arc<dyn CredentialsProvider>,
1622        server_url: &url::Url,
1623        cx: &AsyncApp,
1624    ) -> Result<Option<String>> {
1625        let key = Self::client_secret_keychain_key(server_url);
1626        match credentials_provider.read_credentials(&key, cx).await? {
1627            Some((_username, secret_bytes)) => Ok(Some(String::from_utf8(secret_bytes)?)),
1628            None => Ok(None),
1629        }
1630    }
1631
1632    pub async fn store_client_secret(
1633        credentials_provider: &Arc<dyn CredentialsProvider>,
1634        server_url: &url::Url,
1635        secret: &str,
1636        cx: &AsyncApp,
1637    ) -> Result<()> {
1638        let key = Self::client_secret_keychain_key(server_url);
1639        credentials_provider
1640            .write_credentials(&key, "mcp-oauth-client-secret", secret.as_bytes(), cx)
1641            .await
1642    }
1643
1644    async fn clear_client_secret(
1645        credentials_provider: &Arc<dyn CredentialsProvider>,
1646        server_url: &url::Url,
1647        cx: &AsyncApp,
1648    ) -> Result<()> {
1649        let key = Self::client_secret_keychain_key(server_url);
1650        credentials_provider.delete_credentials(&key, cx).await
1651    }
1652
1653    /// Log out of an OAuth-authenticated MCP server: clear the stored OAuth
1654    /// session from the keychain and stop the server.
1655    pub fn logout_server(&mut self, id: &ContextServerId, cx: &mut Context<Self>) -> Result<()> {
1656        let state = self.servers.get(id).context("Context server not found")?;
1657        let configuration = state.configuration();
1658
1659        let server_url = match configuration.as_ref() {
1660            ContextServerConfiguration::Http { url, .. } => url.clone(),
1661            _ => anyhow::bail!("logout only applies to HTTP servers with OAuth"),
1662        };
1663
1664        let id = id.clone();
1665        self.stop_server(&id, cx)?;
1666
1667        cx.spawn(async move |this, cx| {
1668            let credentials_provider = cx.update(|cx| zed_credentials_provider::global(cx));
1669            if let Err(err) = Self::clear_session(&credentials_provider, &server_url, &cx).await {
1670                log::error!("{} failed to clear OAuth session: {}", id, err);
1671            }
1672            // Also clear any client secret so the user gets a fresh prompt on
1673            // the next authentication attempt.
1674            Self::clear_client_secret(&credentials_provider, &server_url, &cx)
1675                .await
1676                .log_err();
1677            // Trigger server recreation so the next start uses a fresh
1678            // transport without the old (now-invalidated) token provider.
1679            this.update(cx, |this, cx| {
1680                this.available_context_servers_changed(cx);
1681            })
1682            .log_err();
1683        })
1684        .detach();
1685
1686        Ok(())
1687    }
1688
1689    fn update_server_state(
1690        &mut self,
1691        id: ContextServerId,
1692        state: ContextServerState,
1693        cx: &mut Context<Self>,
1694    ) {
1695        let status = ContextServerStatus::from_state(&state);
1696        self.servers.insert(id.clone(), state);
1697        cx.emit(ServerStatusChangedEvent {
1698            server_id: id,
1699            status,
1700        });
1701        cx.notify();
1702    }
1703
1704    fn available_context_servers_changed(&mut self, cx: &mut Context<Self>) {
1705        if self.update_servers_task.is_some() {
1706            self.needs_server_update = true;
1707        } else {
1708            self.needs_server_update = false;
1709            self.update_servers_task = Some(cx.spawn(async move |this, cx| {
1710                if let Err(err) = Self::maintain_servers(this.clone(), cx).await {
1711                    log::error!("Error maintaining context servers: {}", err);
1712                }
1713
1714                this.update(cx, |this, cx| {
1715                    this.populate_server_ids(cx);
1716                    cx.notify();
1717                    this.update_servers_task.take();
1718                    if this.needs_server_update {
1719                        this.available_context_servers_changed(cx);
1720                    }
1721                })?;
1722
1723                Ok(())
1724            }));
1725        }
1726    }
1727
1728    async fn maintain_servers(this: WeakEntity<Self>, cx: &mut AsyncApp) -> Result<()> {
1729        // Don't start context servers if AI is disabled
1730        let ai_disabled = this.update(cx, |_, cx| DisableAiSettings::get_global(cx).disable_ai)?;
1731        if ai_disabled {
1732            // Stop all running servers when AI is disabled
1733            this.update(cx, |this, cx| {
1734                let server_ids: Vec<_> = this.servers.keys().cloned().collect();
1735                for id in server_ids {
1736                    let _ = this.stop_server(&id, cx);
1737                }
1738            })?;
1739            return Ok(());
1740        }
1741
1742        let (mut configured_servers, registry, worktree_store) = this.update(cx, |this, _| {
1743            (
1744                this.context_server_settings.clone(),
1745                this.registry.clone(),
1746                this.worktree_store.clone(),
1747            )
1748        })?;
1749
1750        for (id, _) in registry.read_with(cx, |registry, _| registry.context_server_descriptors()) {
1751            configured_servers
1752                .entry(id)
1753                .or_insert(ContextServerSettingsEntry {
1754                    worktree_id: None,
1755                    settings: ContextServerSettings::default_extension(),
1756                });
1757        }
1758
1759        let (enabled_servers, disabled_servers): (HashMap<_, _>, HashMap<_, _>) =
1760            configured_servers
1761                .into_iter()
1762                .partition(|(_, entry)| entry.settings.enabled());
1763
1764        let configured_servers =
1765            join_all(enabled_servers.into_iter().map(|(id, settings_entry)| {
1766                let id = ContextServerId(id);
1767                ContextServerConfiguration::from_settings(
1768                    settings_entry.settings,
1769                    id.clone(),
1770                    registry.clone(),
1771                    worktree_store.clone(),
1772                    cx,
1773                )
1774                .map(move |config| (id, config))
1775            }))
1776            .await
1777            .into_iter()
1778            .filter_map(|(id, config)| config.map(|config| (id, config)))
1779            .collect::<HashMap<_, _>>();
1780
1781        let mut servers_to_start = Vec::new();
1782        let mut servers_to_remove = HashSet::default();
1783        let mut servers_to_stop = HashSet::default();
1784
1785        this.update(cx, |this, cx| {
1786            for server_id in this.servers.keys() {
1787                // All servers that are not in desired_servers should be removed from the store.
1788                // This can happen if the user removed a server from the context server settings.
1789                if !configured_servers.contains_key(server_id) {
1790                    if disabled_servers.contains_key(&server_id.0) {
1791                        servers_to_stop.insert(server_id.clone());
1792                    } else {
1793                        servers_to_remove.insert(server_id.clone());
1794                    }
1795                }
1796            }
1797
1798            let is_remote_project = this.is_remote_project();
1799            let root_path = this.resolve_root_path(cx);
1800
1801            for (id, config) in configured_servers {
1802                let state = this.servers.get(&id);
1803                let is_stopped = matches!(state, Some(ContextServerState::Stopped { .. }));
1804                let existing_config = state.as_ref().map(|state| state.configuration());
1805                let working_directory =
1806                    working_directory_for(&config, root_path.clone(), is_remote_project);
1807                // A running server that was started before the project root became
1808                // available keeps its stale working directory, since the working
1809                // directory is not part of `ContextServerConfiguration`. Restart it
1810                // when the resolved working directory no longer matches.
1811                let working_directory_changed = state.is_some()
1812                    && !is_stopped
1813                    && this.server_working_directories.get(&id) != Some(&working_directory);
1814                if existing_config.as_deref() != Some(&config)
1815                    || is_stopped
1816                    || working_directory_changed
1817                {
1818                    let config = Arc::new(config);
1819                    servers_to_start.push((id.clone(), config, working_directory));
1820                    if this.servers.contains_key(&id) {
1821                        servers_to_stop.insert(id);
1822                    }
1823                }
1824            }
1825
1826            anyhow::Ok(())
1827        })??;
1828
1829        this.update(cx, |this, inner_cx| {
1830            for id in servers_to_stop {
1831                this.stop_server(&id, inner_cx)?;
1832            }
1833            for id in servers_to_remove {
1834                this.remove_server(&id, inner_cx)?;
1835            }
1836            anyhow::Ok(())
1837        })??;
1838
1839        for (id, config, working_directory) in servers_to_start {
1840            match Self::create_context_server(this.clone(), id.clone(), config, cx).await {
1841                Ok((server, config)) => {
1842                    this.update(cx, |this, cx| {
1843                        this.server_working_directories
1844                            .insert(id.clone(), working_directory);
1845                        this.run_server(server, config, cx);
1846                    })?;
1847                }
1848                Err(err) => {
1849                    log::error!("{id} context server failed to create: {err:#}");
1850                    this.update(cx, |_this, cx| {
1851                        cx.emit(ServerStatusChangedEvent {
1852                            server_id: id,
1853                            status: ContextServerStatus::Error(err.to_string().into()),
1854                        });
1855                        cx.notify();
1856                    })?;
1857                }
1858            }
1859        }
1860
1861        Ok(())
1862    }
1863}
1864
1865/// The working directory a server will be spawned with, mirroring the choice
1866/// made in [`ContextServerStore::create_context_server`]: only locally-spawned
1867/// stdio servers use the project root; HTTP and remote servers use none.
1868fn working_directory_for(
1869    configuration: &ContextServerConfiguration,
1870    root_path: Option<Arc<Path>>,
1871    is_remote_project: bool,
1872) -> Option<Arc<Path>> {
1873    match configuration {
1874        ContextServerConfiguration::Http { .. } => None,
1875        _ if is_remote_project => None,
1876        _ => root_path,
1877    }
1878}
1879
1880/// Determines the appropriate server state after a start attempt fails.
1881///
1882/// When the error is an HTTP 401 with no static auth header configured,
1883/// attempts OAuth discovery so the UI can offer an authentication flow.
1884async fn resolve_start_failure(
1885    id: &ContextServerId,
1886    err: anyhow::Error,
1887    server: Arc<ContextServer>,
1888    configuration: Arc<ContextServerConfiguration>,
1889    cx: &AsyncApp,
1890) -> ContextServerState {
1891    // Read the challenge from the transport rather than downcasting `err`: it
1892    // is recorded before the failed send's error propagates, so a 401 is
1893    // recognized even when another error (e.g. the request timeout) wins the
1894    // race to become the reported startup failure.
1895    let www_authenticate = server.auth_challenge();
1896
1897    // When the error is NOT a 401 but there is a cached OAuth session in the
1898    // keychain, the session is likely stale/expired and caused the failure
1899    // (e.g. timeout because the server rejected the token silently). Clear it
1900    // so the next start attempt can get a clean 401 and trigger the auth flow.
1901    // If there is no such session this is an ordinary startup error.
1902    if www_authenticate.is_none() {
1903        let server_url = match configuration.as_ref() {
1904            ContextServerConfiguration::Http { url, .. }
1905                if !configuration.has_static_auth_header() =>
1906            {
1907                url.clone()
1908            }
1909            _ => {
1910                log::error!("{id} context server failed to start: {err}");
1911                return ContextServerState::Error {
1912                    configuration,
1913                    server,
1914                    error: err.to_string().into(),
1915                };
1916            }
1917        };
1918
1919        let credentials_provider = cx.update(|cx| zed_credentials_provider::global(cx));
1920        match ContextServerStore::load_session(&credentials_provider, &server_url, cx).await {
1921            Ok(Some(_)) => {
1922                log::info!("{id} start failed with a cached OAuth session present; clearing it");
1923                ContextServerStore::clear_session(&credentials_provider, &server_url, cx)
1924                    .await
1925                    .log_err();
1926            }
1927            _ => {
1928                log::error!("{id} context server failed to start: {err}");
1929                return ContextServerState::Error {
1930                    configuration,
1931                    server,
1932                    error: err.to_string().into(),
1933                };
1934            }
1935        }
1936    }
1937
1938    let default_www_authenticate = oauth::WwwAuthenticate {
1939        resource_metadata: None,
1940        scope: None,
1941        error: None,
1942        error_description: None,
1943    };
1944    let www_authenticate = www_authenticate
1945        .as_ref()
1946        .unwrap_or(&default_www_authenticate);
1947
1948    resolve_auth_required(id, www_authenticate, server, configuration, cx).await
1949}
1950
1951/// Runs OAuth discovery for a server that returned a 401 and produces the
1952/// appropriate state (`AuthRequired`, `ClientSecretRequired`, or `Error`).
1953///
1954/// Shared by the startup path ([`resolve_start_failure`]) and the
1955/// post-initialize path ([`ContextServerStore::handle_auth_challenge`]) so
1956/// that a 401 at any point — not only during `initialize` — can initiate the
1957/// OAuth flow.
1958async fn resolve_auth_required(
1959    id: &ContextServerId,
1960    www_authenticate: &oauth::WwwAuthenticate,
1961    server: Arc<ContextServer>,
1962    configuration: Arc<ContextServerConfiguration>,
1963    cx: &AsyncApp,
1964) -> ContextServerState {
1965    if configuration.has_static_auth_header() {
1966        log::warn!("{id} received 401 with a static Authorization header configured");
1967        return ContextServerState::Error {
1968            configuration,
1969            server,
1970            error: "Server returned 401 Unauthorized. Check your configured Authorization header."
1971                .into(),
1972        };
1973    }
1974
1975    let server_url = match configuration.as_ref() {
1976        ContextServerConfiguration::Http { url, .. } => url.clone(),
1977        _ => {
1978            log::error!("{id} got OAuth 401 on a non-HTTP transport");
1979            return ContextServerState::Error {
1980                configuration,
1981                server,
1982                error: "Server returned 401 Unauthorized on a non-HTTP transport".into(),
1983            };
1984        }
1985    };
1986
1987    let http_client = cx.update(|cx| cx.http_client());
1988
1989    match context_server::oauth::discover(&http_client, &server_url, www_authenticate).await {
1990        Ok(discovery) => {
1991            use context_server::oauth::{
1992                ClientRegistrationStrategy, determine_registration_strategy,
1993            };
1994
1995            let has_preregistered_client_id = matches!(
1996                configuration.as_ref(),
1997                ContextServerConfiguration::Http { oauth: Some(_), .. }
1998            );
1999
2000            let strategy = determine_registration_strategy(&discovery.auth_server_metadata);
2001
2002            if matches!(strategy, ClientRegistrationStrategy::Unavailable)
2003                && !has_preregistered_client_id
2004            {
2005                log::error!(
2006                    "{id} authorization server supports neither CIMD nor DCR, \
2007                     and no pre-registered client_id is configured"
2008                );
2009                return ContextServerState::Error {
2010                    configuration,
2011                    server,
2012                    error: "Authorization server supports neither CIMD nor DCR. \
2013                            Configure a pre-registered client_id in your settings \
2014                            under the \"oauth\" key."
2015                        .into(),
2016                };
2017            }
2018
2019            log::info!(
2020                "{id} requires OAuth authorization (auth server: {})",
2021                discovery.auth_server_metadata.issuer,
2022            );
2023            ContextServerState::AuthRequired {
2024                server,
2025                configuration,
2026                discovery: Arc::new(discovery),
2027            }
2028        }
2029        Err(discovery_err) => {
2030            log::error!("{id} OAuth discovery failed: {discovery_err}");
2031            ContextServerState::Error {
2032                configuration,
2033                server,
2034                error: format!("OAuth discovery failed: {discovery_err}").into(),
2035            }
2036        }
2037    }
2038}
2039