diff --git a/crates/context_server/src/client.rs b/crates/context_server/src/client.rs index b8a321d01b..3761f19ec4 100644 --- a/crates/context_server/src/client.rs +++ b/crates/context_server/src/client.rs @@ -4,7 +4,7 @@ use futures::{FutureExt, StreamExt, channel::oneshot, future, select}; use futures_lite::future::yield_now; use gpui::{AppContext as _, AsyncApp, BackgroundExecutor, Task}; use parking_lot::Mutex; -use postage::barrier; +use postage::{barrier, prelude::Stream as _}; use serde::{Deserialize, Serialize, de::DeserializeOwned}; use serde_json::{Value, value::RawValue}; use slotmap::SlotMap; @@ -21,6 +21,7 @@ use std::{ use util::{ResultExt, TryFutureExt}; use crate::{ + oauth::WwwAuthenticate, transport::{StdioTransport, Transport}, types::{CancelledParams, ClientNotification, Notification as _, notifications::Cancelled}, }; @@ -56,19 +57,19 @@ pub(crate) struct Client { #[allow(clippy::type_complexity)] #[allow(dead_code)] io_tasks: Mutex>, Task>)>>, - #[allow(dead_code)] output_done_rx: Mutex>, executor: BackgroundExecutor, - #[allow(dead_code)] transport: Arc, request_timeout: Option, /// Single-slot side channel for the last transport-level error. When the /// output task encounters a send failure it stashes the error here and - /// exits; the next request to observe cancellation `.take()`s it so it can - /// propagate a typed error (e.g. `TransportError::AuthRequired`) instead - /// of a generic "cancelled". This works because `initialize` is the sole - /// in-flight request at startup, but would need rethinking if concurrent - /// requests are ever issued during that phase. + /// exits; the next request to observe cancellation `.take()`s it so it + /// can fail with the underlying cause (e.g. "connection refused") instead + /// of a generic "cancelled". This is best-effort diagnostics: with + /// concurrent requests in flight, a single arbitrary one receives the + /// stashed error. Nothing may depend on it for correctness — + /// authentication challenges are observed via [`Self::wait_for_shutdown`] + /// and [`Transport::auth_challenge`], which do not involve requests. last_transport_error: Arc>>, } @@ -348,6 +349,28 @@ impl Client { Ok(()) } + /// A future that resolves once the transport's output loop has terminated + /// — after a send failure, or when this client is dropped — yielding the + /// authentication challenge recorded by the transport if it shut down on a + /// `401 Unauthorized` response. + /// + /// Unlike `last_transport_error`, this does not require a request to be in + /// flight when the transport fails. Returns `None` if the shutdown signal + /// was already claimed: there is a single signal per client. + pub(crate) fn wait_for_shutdown( + &self, + ) -> Option>> { + let mut output_done = self.output_done_rx.lock().take()?; + let transport = self.transport.clone(); + Some( + async move { + output_done.recv().await; + transport.auth_challenge() + } + .boxed(), + ) + } + /// Sends a JSON-RPC request to the context server and waits for a response. /// This function handles serialization, deserialization, timeout, and error handling. pub async fn request( diff --git a/crates/context_server/src/context_server.rs b/crates/context_server/src/context_server.rs index 05a3451ea8..865f779b16 100644 --- a/crates/context_server/src/context_server.rs +++ b/crates/context_server/src/context_server.rs @@ -21,6 +21,7 @@ use parking_lot::RwLock; pub use settings::ContextServerCommand; use url::Url; +use crate::oauth::WwwAuthenticate; use crate::transport::HttpTransport; #[derive(Debug, Clone, PartialEq, Eq, Hash)] @@ -106,6 +107,16 @@ impl ContextServer { self.client.read().clone() } + /// The authentication challenge from the last `401 Unauthorized` response + /// this server's transport gave up on, if any. See + /// [`crate::transport::Transport::auth_challenge`]. + pub fn auth_challenge(&self) -> Option { + match &self.configuration { + ContextServerTransport::Stdio(..) => None, + ContextServerTransport::Custom(transport) => transport.auth_challenge(), + } + } + pub async fn start(&self, cx: &AsyncApp) -> Result<()> { self.initialize(self.new_client(cx)?).await } diff --git a/crates/context_server/src/protocol.rs b/crates/context_server/src/protocol.rs index 05082637c2..6eb4a84a46 100644 --- a/crates/context_server/src/protocol.rs +++ b/crates/context_server/src/protocol.rs @@ -8,11 +8,12 @@ use std::time::Duration; use anyhow::Result; -use futures::channel::oneshot; +use futures::{channel::oneshot, future::BoxFuture}; use gpui::AsyncApp; use serde_json::Value; use crate::client::{Client, NotificationSubscription}; +use crate::oauth::WwwAuthenticate; use crate::types::{self, Notification, Request}; pub struct ModelContextProtocol { @@ -122,6 +123,20 @@ impl InitializedContextServerProtocol { self.inner.notify(T::METHOD, params) } + /// A future that resolves once the underlying transport's output loop has + /// terminated — after a send failure, or when the client is dropped — + /// yielding the authentication challenge recorded by the transport if it + /// shut down on a `401 Unauthorized` response. + /// + /// Servers may accept `initialize` unauthenticated and only challenge a + /// later request or notification. Awaiting this is what lets the owner of + /// the connection notice such a challenge even when no request was in + /// flight to carry a typed error back. Returns `None` if the shutdown + /// signal was already claimed: there is a single signal per client. + pub fn wait_for_shutdown(&self) -> Option>> { + self.inner.wait_for_shutdown() + } + pub fn on_notification( &self, method: &'static str, diff --git a/crates/context_server/src/transport.rs b/crates/context_server/src/transport.rs index bffd7e4c4d..edb8a5e8da 100644 --- a/crates/context_server/src/transport.rs +++ b/crates/context_server/src/transport.rs @@ -6,6 +6,8 @@ use async_trait::async_trait; use futures::Stream; use std::pin::Pin; +use crate::oauth::WwwAuthenticate; + pub use http::*; pub use stdio_transport::*; @@ -19,4 +21,16 @@ pub trait Transport: Send + Sync { /// need the negotiated version (currently only HTTP, which must attach an /// `MCP-Protocol-Version` header from 2025-06-18 onward) can pick it up. fn set_protocol_version(&self, _version: &str) {} + + /// The authentication challenge from the last `401 Unauthorized` response + /// this transport gave up on, if any (currently only set by the HTTP + /// transport). + /// + /// The challenge is recorded right before the failed send tears down the + /// client's output loop. Observers of the client's shutdown read it from + /// here, so a 401 can initiate the OAuth flow even when it arrived on a + /// notification, with no request in flight to carry a typed error. + fn auth_challenge(&self) -> Option { + None + } } diff --git a/crates/context_server/src/transport/http.rs b/crates/context_server/src/transport/http.rs index bf374586b3..aa47a73fc3 100644 --- a/crates/context_server/src/transport/http.rs +++ b/crates/context_server/src/transport/http.rs @@ -58,6 +58,10 @@ pub struct HttpTransport { /// When set, the transport attaches `Authorization: Bearer` headers and /// handles 401 responses with token refresh + retry. token_provider: Option>, + /// The challenge from the last 401 this transport gave up on; cleared at + /// the start of each send so it always describes the most recent attempt. + /// See [`Transport::auth_challenge`]. + auth_challenge: SyncMutex>, } impl HttpTransport { @@ -92,6 +96,7 @@ impl HttpTransport { error_rx, headers, token_provider, + auth_challenge: SyncMutex::new(None), } } @@ -133,8 +138,21 @@ impl HttpTransport { Ok(request_builder.body(AsyncBody::from(message.to_vec()))?) } + /// Record the challenge so it remains observable after the failed send + /// tears down the client (see [`Transport::auth_challenge`]), and build + /// the typed error for the send itself. + fn auth_required(&self, www_authenticate: WwwAuthenticate) -> anyhow::Error { + *self.auth_challenge.lock() = Some(www_authenticate.clone()); + TransportError::AuthRequired { www_authenticate }.into() + } + /// Send a message and handle the response based on content type. async fn send_message(&self, message: String) -> Result<()> { + // The same server instance can be restarted over this transport; a + // challenge recorded by a previous client generation must not be + // observed by the current one. + *self.auth_challenge.lock() = None; + let is_notification = !message.contains("\"id\":") || message.contains("notifications/initialized"); @@ -174,13 +192,13 @@ impl HttpTransport { // If still 401 after refresh, give up. if response.status().as_u16() == 401 { - return Err(TransportError::AuthRequired { www_authenticate }.into()); + return Err(self.auth_required(www_authenticate)); } } else { - return Err(TransportError::AuthRequired { www_authenticate }.into()); + return Err(self.auth_required(www_authenticate)); } } else { - return Err(TransportError::AuthRequired { www_authenticate }.into()); + return Err(self.auth_required(www_authenticate)); } } @@ -335,6 +353,10 @@ impl Transport for HttpTransport { fn set_protocol_version(&self, version: &str) { *self.protocol_version.lock() = Some(version.to_string()); } + + fn auth_challenge(&self) -> Option { + self.auth_challenge.lock().clone() + } } impl Drop for HttpTransport { diff --git a/crates/project/src/context_server_store.rs b/crates/project/src/context_server_store.rs index ef5077d53a..bbd53e1e73 100644 --- a/crates/project/src/context_server_store.rs +++ b/crates/project/src/context_server_store.rs @@ -8,7 +8,7 @@ use std::time::Duration; use anyhow::{Context as _, Result}; use collections::{HashMap, HashSet}; use context_server::oauth::{self, McpOAuthTokenProvider, OAuthDiscovery, OAuthSession}; -use context_server::transport::{HttpTransport, TransportError}; +use context_server::transport::HttpTransport; use context_server::{ContextServer, ContextServerCommand, ContextServerId}; use credentials_provider::CredentialsProvider; use futures::future::Either; @@ -93,6 +93,9 @@ enum ContextServerState { Running { server: Arc, configuration: Arc, + /// Initiates the OAuth flow if the transport shuts down on an + /// authentication challenge; cancelled by any state transition. + _transport_watch: Task<()>, }, Stopped { server: Arc, @@ -708,9 +711,12 @@ impl ContextServerStore { let new_state = match server.clone().start(cx).await { Ok(_) => { debug_assert!(server.client().is_some()); + let _transport_watch = + Self::watch_transport_shutdown(this.clone(), server.clone(), cx); ContextServerState::Running { server, configuration, + _transport_watch, } } Err(err) => resolve_start_failure(&id, err, server, configuration, cx).await, @@ -733,6 +739,97 @@ impl ContextServerStore { ); } + /// Watches a running server's transport and initiates the OAuth flow if it + /// shuts down on an authentication challenge. + /// + /// MCP servers may accept `initialize` unauthenticated and only send a 401 + /// with a `WWW-Authenticate` challenge on a later request or notification. + /// The HTTP transport records the challenge, and the failed send tears + /// down the client's output loop. Observing that shutdown — rather than + /// relying on some request to carry a typed error back to a caller — is + /// what lets any post-initialize 401 move the server into `AuthRequired` + /// instead of leaving it `Running` with a dead client. + fn watch_transport_shutdown( + this: WeakEntity, + server: Arc, + cx: &mut AsyncApp, + ) -> Task<()> { + let Some(shutdown) = server + .client() + .and_then(|client| client.wait_for_shutdown()) + else { + return Task::ready(()); + }; + cx.spawn(async move |cx| { + let Some(www_authenticate) = shutdown.await else { + // Non-auth transport deaths leave the server state untouched, + // as they did before this watch existed. + return; + }; + this.update(cx, |this, cx| { + this.handle_auth_challenge(server, www_authenticate, cx); + }) + .log_err(); + }) + } + + fn handle_auth_challenge( + &mut self, + server: Arc, + www_authenticate: oauth::WwwAuthenticate, + cx: &mut Context, + ) { + let id = server.id(); + + // Act only if this exact server is still the one we consider running. + // If the state has changed since the challenge was recorded, whoever + // changed it owns the lifecycle now. + let Some(ContextServerState::Running { + server: running_server, + configuration, + .. + }) = self.servers.get(&id) + else { + return; + }; + if !Arc::ptr_eq(running_server, &server) { + return; + } + let configuration = configuration.clone(); + + log::info!("{id} received 401 after initialization; initiating OAuth authorization"); + + // The 401 already tore down the client's output loop. Stop the dead + // client, then resolve auth using the captured `WWW-Authenticate` — do + // not restart via `run_server`, as a fresh `initialize` would succeed + // and lose the challenge. + server.stop().log_err(); + + let task = cx.spawn({ + let id = id.clone(); + let server = server.clone(); + let configuration = configuration.clone(); + async move |this, cx| { + let new_state = + resolve_auth_required(&id, &www_authenticate, server, configuration, cx).await; + this.update(cx, |this, cx| { + this.update_server_state(id.clone(), new_state, cx) + }) + .log_err(); + } + }); + + self.update_server_state( + id, + ContextServerState::Starting { + configuration, + _task: task, + server, + }, + cx, + ); + } + fn remove_server(&mut self, id: &ContextServerId, cx: &mut Context) -> Result<()> { let state = self .servers @@ -1686,43 +1783,34 @@ async fn resolve_start_failure( configuration: Arc, cx: &AsyncApp, ) -> ContextServerState { - let www_authenticate = err.downcast_ref::().map(|e| match e { - TransportError::AuthRequired { www_authenticate } => www_authenticate.clone(), - }); - - if www_authenticate.is_some() && configuration.has_static_auth_header() { - log::warn!("{id} received 401 with a static Authorization header configured"); - return ContextServerState::Error { - configuration, - server, - error: "Server returned 401 Unauthorized. Check your configured Authorization header." - .into(), - }; - } - - let server_url = match configuration.as_ref() { - ContextServerConfiguration::Http { url, .. } if !configuration.has_static_auth_header() => { - url.clone() - } - _ => { - if www_authenticate.is_some() { - log::error!("{id} got OAuth 401 on a non-HTTP transport or with static auth"); - } else { - log::error!("{id} context server failed to start: {err}"); - } - return ContextServerState::Error { - configuration, - server, - error: err.to_string().into(), - }; - } - }; + // Read the challenge from the transport rather than downcasting `err`: it + // is recorded before the failed send's error propagates, so a 401 is + // recognized even when another error (e.g. the request timeout) wins the + // race to become the reported startup failure. + let www_authenticate = server.auth_challenge(); // When the error is NOT a 401 but there is a cached OAuth session in the // keychain, the session is likely stale/expired and caused the failure // (e.g. timeout because the server rejected the token silently). Clear it // so the next start attempt can get a clean 401 and trigger the auth flow. + // If there is no such session this is an ordinary startup error. if www_authenticate.is_none() { + let server_url = match configuration.as_ref() { + ContextServerConfiguration::Http { url, .. } + if !configuration.has_static_auth_header() => + { + url.clone() + } + _ => { + log::error!("{id} context server failed to start: {err}"); + return ContextServerState::Error { + configuration, + server, + error: err.to_string().into(), + }; + } + }; + let credentials_provider = cx.update(|cx| zed_credentials_provider::global(cx)); match ContextServerStore::load_session(&credentials_provider, &server_url, cx).await { Ok(Some(_)) => { @@ -1751,6 +1839,46 @@ async fn resolve_start_failure( let www_authenticate = www_authenticate .as_ref() .unwrap_or(&default_www_authenticate); + + resolve_auth_required(id, www_authenticate, server, configuration, cx).await +} + +/// Runs OAuth discovery for a server that returned a 401 and produces the +/// appropriate state (`AuthRequired`, `ClientSecretRequired`, or `Error`). +/// +/// Shared by the startup path ([`resolve_start_failure`]) and the +/// post-initialize path ([`ContextServerStore::handle_auth_challenge`]) so +/// that a 401 at any point — not only during `initialize` — can initiate the +/// OAuth flow. +async fn resolve_auth_required( + id: &ContextServerId, + www_authenticate: &oauth::WwwAuthenticate, + server: Arc, + configuration: Arc, + cx: &AsyncApp, +) -> ContextServerState { + if configuration.has_static_auth_header() { + log::warn!("{id} received 401 with a static Authorization header configured"); + return ContextServerState::Error { + configuration, + server, + error: "Server returned 401 Unauthorized. Check your configured Authorization header." + .into(), + }; + } + + let server_url = match configuration.as_ref() { + ContextServerConfiguration::Http { url, .. } => url.clone(), + _ => { + log::error!("{id} got OAuth 401 on a non-HTTP transport"); + return ContextServerState::Error { + configuration, + server, + error: "Server returned 401 Unauthorized on a non-HTTP transport".into(), + }; + } + }; + let http_client = cx.update(|cx| cx.http_client()); match context_server::oauth::discover(&http_client, &server_url, www_authenticate).await { diff --git a/crates/project/tests/integration/context_server_store.rs b/crates/project/tests/integration/context_server_store.rs index 090baacf03..a91cdd2ea8 100644 --- a/crates/project/tests/integration/context_server_store.rs +++ b/crates/project/tests/integration/context_server_store.rs @@ -914,6 +914,421 @@ async fn test_remote_context_server(cx: &mut TestAppContext) { cx.run_until_parked(); } +// A server may accept `initialize` unauthenticated (returns 200) yet only send +// `WWW-Authenticate` on a later request such as `tools/list` / `tools/call`. +// That post-initialize 401 must initiate the OAuth flow instead of surfacing as +// an opaque request failure while the server stays "Running". +#[gpui::test] +async fn test_http_server_authenticates_on_post_init_401(cx: &mut TestAppContext) { + use context_server::transport::TransportError; + + const SERVER_ID: &str = "auth-server"; + let server_id = ContextServerId(SERVER_ID.into()); + + set_fake_mcp_http_client(cx, |message| { + if message.contains("\"method\":\"initialize\"") { + Ok(initialize_response()) + } else if message.contains("notifications/initialized") { + Ok(notification_accepted_response()) + } else { + Ok(unauthorized_response()) + } + }); + + let (_fs, project) = setup_context_server_test(cx, json!({ "code.rs": "" }), vec![]).await; + let store = project.read_with(cx, |project, _| project.context_server_store()); + + set_http_context_server_configuration(&server_id, cx); + + { + let _server_events = assert_server_events( + &store, + vec![ + (server_id.clone(), ContextServerStatus::Starting), + (server_id.clone(), ContextServerStatus::Running), + ], + cx, + ); + cx.run_until_parked(); + } + + let client = store.read_with(cx, |store, _| { + store + .get_running_server(&server_id) + .expect("server should be running") + .client() + .expect("running server should have a client") + }); + + { + let _server_events = assert_server_events( + &store, + vec![ + (server_id.clone(), ContextServerStatus::Starting), + (server_id.clone(), ContextServerStatus::AuthRequired), + ], + cx, + ); + + // The 401 tears down the client, and the request that carried it fails + // with the typed error. + let error = client + .request::(()) + .await + .expect_err("request challenged with a 401 should fail"); + assert!( + matches!( + error.downcast_ref::(), + Some(TransportError::AuthRequired { .. }) + ), + "expected an AuthRequired error, got: {error}" + ); + + cx.run_until_parked(); + } + + cx.update(|cx| { + assert_eq!( + store.read(cx).status_for_server(&server_id), + Some(ContextServerStatus::AuthRequired), + "server should require authentication after a post-initialize 401" + ); + }); +} + +// The first 401 may even arrive on a notification (`notifications/initialized` +// here): the send fails with no request in flight to carry a typed error back, +// and the client is dead by the time anything notices. Watching the transport +// shutdown must still move the server into `AuthRequired`. +#[gpui::test] +async fn test_http_server_authenticates_on_notification_401(cx: &mut TestAppContext) { + const SERVER_ID: &str = "auth-server"; + let server_id = ContextServerId(SERVER_ID.into()); + + set_fake_mcp_http_client(cx, |message| { + if message.contains("\"method\":\"initialize\"") { + Ok(initialize_response()) + } else { + Ok(unauthorized_response()) + } + }); + + let (_fs, project) = setup_context_server_test(cx, json!({ "code.rs": "" }), vec![]).await; + let store = project.read_with(cx, |project, _| project.context_server_store()); + + set_http_context_server_configuration(&server_id, cx); + + { + let _server_events = assert_server_events( + &store, + vec![ + (server_id.clone(), ContextServerStatus::Starting), + (server_id.clone(), ContextServerStatus::Running), + (server_id.clone(), ContextServerStatus::Starting), + (server_id.clone(), ContextServerStatus::AuthRequired), + ], + cx, + ); + cx.run_until_parked(); + } + + cx.update(|cx| { + assert_eq!( + store.read(cx).status_for_server(&server_id), + Some(ContextServerStatus::AuthRequired), + "server should require authentication after a 401 on a notification" + ); + }); +} + +// A transport failure that is not an authentication challenge must not touch +// the server's state: no spurious auth flow, and (as before the transport +// watch existed) the server stays `Running`. +#[gpui::test] +async fn test_http_server_ignores_non_auth_transport_failure(cx: &mut TestAppContext) { + const SERVER_ID: &str = "flaky-server"; + let server_id = ContextServerId(SERVER_ID.into()); + + set_fake_mcp_http_client(cx, |message| { + if message.contains("\"method\":\"initialize\"") { + Ok(initialize_response()) + } else if message.contains("notifications/initialized") { + Ok(notification_accepted_response()) + } else { + Err(anyhow::anyhow!("connection reset")) + } + }); + + let (_fs, project) = setup_context_server_test(cx, json!({ "code.rs": "" }), vec![]).await; + let store = project.read_with(cx, |project, _| project.context_server_store()); + + set_http_context_server_configuration(&server_id, cx); + + { + let _server_events = assert_server_events( + &store, + vec![ + (server_id.clone(), ContextServerStatus::Starting), + (server_id.clone(), ContextServerStatus::Running), + ], + cx, + ); + cx.run_until_parked(); + + let client = store.read_with(cx, |store, _| { + store + .get_running_server(&server_id) + .expect("server should be running") + .client() + .expect("running server should have a client") + }); + + client + .request::(()) + .await + .expect_err("request should fail when the transport errors"); + + cx.run_until_parked(); + // Dropping the events guard asserts no further status change happened. + } + + cx.update(|cx| { + assert_eq!( + store.read(cx).status_for_server(&server_id), + Some(ContextServerStatus::Running), + "a non-auth transport failure should not change the server state" + ); + }); +} + +// A server may also require authentication on `initialize` itself. The +// challenge is read from the transport slot rather than the returned error, so +// the 401 is recognized even if another error (e.g. the request timeout) wins +// the race to become the reported startup failure. +#[gpui::test] +async fn test_http_server_authenticates_on_initialize_401(cx: &mut TestAppContext) { + const SERVER_ID: &str = "auth-server"; + let server_id = ContextServerId(SERVER_ID.into()); + + set_fake_mcp_http_client(cx, |_message| Ok(unauthorized_response())); + + let (_fs, project) = setup_context_server_test(cx, json!({ "code.rs": "" }), vec![]).await; + let store = project.read_with(cx, |project, _| project.context_server_store()); + + set_http_context_server_configuration(&server_id, cx); + + { + let _server_events = assert_server_events( + &store, + vec![ + (server_id.clone(), ContextServerStatus::Starting), + (server_id.clone(), ContextServerStatus::AuthRequired), + ], + cx, + ); + cx.run_until_parked(); + } + + cx.update(|cx| { + assert_eq!( + store.read(cx).status_for_server(&server_id), + Some(ContextServerStatus::AuthRequired), + "server should require authentication after a 401 on initialize" + ); + }); +} + +// Restarting a server reuses its transport (e.g. via the MCP settings page), +// so a challenge recorded by a previous client generation must not leak into +// the next one: after a successful restart, a non-auth transport failure must +// not trip a spurious auth flow on the stale challenge. +#[gpui::test] +async fn test_http_server_restart_clears_stale_auth_challenge(cx: &mut TestAppContext) { + use std::sync::atomic::{AtomicBool, Ordering}; + + const SERVER_ID: &str = "auth-server"; + let server_id = ContextServerId(SERVER_ID.into()); + + let restarted = Arc::new(AtomicBool::new(false)); + set_fake_mcp_http_client(cx, { + let restarted = restarted.clone(); + move |message| { + if message.contains("\"method\":\"initialize\"") { + Ok(initialize_response()) + } else if message.contains("notifications/initialized") { + Ok(notification_accepted_response()) + } else if restarted.load(Ordering::SeqCst) { + Err(anyhow::anyhow!("connection reset")) + } else { + Ok(unauthorized_response()) + } + } + }); + + let (_fs, project) = setup_context_server_test(cx, json!({ "code.rs": "" }), vec![]).await; + let store = project.read_with(cx, |project, _| project.context_server_store()); + + set_http_context_server_configuration(&server_id, cx); + cx.run_until_parked(); + + // A post-initialize 401 records a challenge on the transport and moves the + // server into AuthRequired. + let client = store.read_with(cx, |store, _| { + store + .get_running_server(&server_id) + .expect("server should be running") + .client() + .expect("running server should have a client") + }); + client + .request::(()) + .await + .expect_err("request challenged with a 401 should fail"); + // Drop our handle so the dead client fully goes away, as it does in + // production once the store has stopped it: a lingering client would + // compete with its successor for the reused transport's response channel. + drop(client); + cx.run_until_parked(); + cx.update(|cx| { + assert_eq!( + store.read(cx).status_for_server(&server_id), + Some(ContextServerStatus::AuthRequired), + ); + }); + + // The user restarts the same server instance instead of authenticating, + // and the server no longer challenges. + restarted.store(true, Ordering::SeqCst); + store.update(cx, |store, cx| { + let server = store.get_server(&server_id).expect("server should exist"); + store.start_server(server, cx); + }); + cx.run_until_parked(); + cx.update(|cx| { + assert_eq!( + store.read(cx).status_for_server(&server_id), + Some(ContextServerStatus::Running), + ); + }); + + let client = store.read_with(cx, |store, _| { + store + .get_running_server(&server_id) + .expect("server should be running") + .client() + .expect("running server should have a client") + }); + client + .request::(()) + .await + .expect_err("request should fail when the transport errors"); + cx.run_until_parked(); + + cx.update(|cx| { + assert_eq!( + store.read(cx).status_for_server(&server_id), + Some(ContextServerStatus::Running), + "a stale challenge from a previous client generation must not trigger auth" + ); + }); +} + +fn set_http_context_server_configuration(server_id: &ContextServerId, cx: &mut TestAppContext) { + set_context_server_configuration( + vec![( + server_id.0.clone(), + settings::ContextServerSettingsContent::Http { + enabled: true, + url: "https://mcp.example.com/mcp".to_string(), + headers: Default::default(), + timeout: None, + oauth: None, + }, + )], + cx, + ); +} + +/// A fake HTTP client that serves the OAuth discovery documents (CIMD-capable, +/// so no dynamic client registration is needed) and routes MCP endpoint POSTs +/// to `respond_to_mcp_message` by the JSON-RPC message in the request body. +fn set_fake_mcp_http_client( + cx: &mut TestAppContext, + respond_to_mcp_message: impl Fn(&str) -> Result> + + Send + + Sync + + 'static, +) { + let respond_to_mcp_message = Arc::new(respond_to_mcp_message); + let client = FakeHttpClient::create(move |request| { + let respond_to_mcp_message = respond_to_mcp_message.clone(); + async move { + let uri = request.uri().to_string(); + let discovery_document = if uri.contains("oauth-protected-resource") { + Some(json!({ + "resource": "https://mcp.example.com", + "authorization_servers": ["https://auth.example.com"], + "scopes_supported": ["mcp:read"] + })) + } else if uri.contains("oauth-authorization-server") { + Some(json!({ + "issuer": "https://auth.example.com", + "authorization_endpoint": "https://auth.example.com/authorize", + "token_endpoint": "https://auth.example.com/token", + "code_challenge_methods_supported": ["S256"], + "client_id_metadata_document_supported": true + })) + } else { + None + }; + if let Some(document) = discovery_document { + return Ok(json_response(document)); + } + + let mut body = request.into_body(); + let mut message = String::new(); + futures::AsyncReadExt::read_to_string(&mut body, &mut message).await?; + respond_to_mcp_message(&message) + } + }); + cx.update(|cx| cx.set_http_client(client)); +} + +fn json_response(body: serde_json::Value) -> Response { + Response::builder() + .status(200) + .header("Content-Type", "application/json") + .body(http_client::AsyncBody::from(body.to_string())) + .unwrap() +} + +fn initialize_response() -> Response { + json_response(json!({ + "jsonrpc": "2.0", + "id": 0, + "result": { + "protocolVersion": "2024-11-05", + "capabilities": {}, + "serverInfo": { "name": "test-server", "version": "1.0.0" } + } + })) +} + +fn notification_accepted_response() -> Response { + Response::builder() + .status(202) + .body(http_client::AsyncBody::empty()) + .unwrap() +} + +fn unauthorized_response() -> Response { + Response::builder() + .status(401) + .header("WWW-Authenticate", "Bearer") + .body(http_client::AsyncBody::empty()) + .unwrap() +} + struct ServerEvents { received_event_count: Rc>, expected_event_count: usize,