yenru0/zed
1
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

MCP remote server OAuth authentication (#51768)

Browse Source 
Closes #43162

Implements the OAuth 2.0 Authorization Code + PKCE authentication flow
for remote MCP servers using Streamable HTTP transport, as specified by
the [MCP auth
specification](https://modelcontextprotocol.io/specification/2025-03-26/basic/authorization).

Previously, connecting to a remote MCP server that required OAuth would
silently fail with a timeout — the server's 401 response was never
handled. Now, Zed detects the 401, performs OAuth discovery, and guides
the user through browser-based authentication.

Step-up authentication and pre-registered clients are not in scope for
this PR, but will be done as follow-ups.

## Overview

- **401 detection** — When the HTTP transport receives a 401 during
server startup, it surfaces a typed `TransportError::AuthRequired` with
parsed `WWW-Authenticate` header info.
- **OAuth discovery** — Protected Resource Metadata (RFC 9728) and
Authorization Server Metadata (RFC 8414) are fetched to locate the
authorization and token endpoints.
- **Client registration** — Zed first tries CIMD (Client ID Metadata
Document) hosted at `zed.dev`. If the server doesn't support CIMD, falls
back to Dynamic Client Registration (DCR).
- **Browser flow** — A loopback HTTP callback server starts on a
preferred fixed port (27523, listed in the CIMD), the user's browser
opens to the authorization URL, and Zed waits for the callback with the
authorization code.
- **Token exchange & persistence** — The code is exchanged for
access/refresh tokens using PKCE. The session is persisted in the system
keychain so subsequent startups restore it without another browser flow.
- **Automatic refresh** — The HTTP transport transparently refreshes
expired tokens using the refresh token, and persists the updated session
to the keychain.

## UI changes

- Servers requiring auth show a warning indicator with an
**"Authenticate"** button
- During auth, a spinner and **"Waiting for authorization..."** message
are shown
- A **"Log Out"** option is available in the server settings menu for
OAuth-authenticated servers
- The configure server modal handles the auth flow inline when
configuring a new server that needs authentication.

Release Notes:

- Added OAuth authentication support for remote MCP servers. Servers
requiring OAuth now show an "Authenticate" button when they need you to
log in. You will be redirected in your browser to the authorization
server of the MCP server to go through the authorization flow.

---------

Co-authored-by: Danilo Leal <daniloleal09@gmail.com>
This commit is contained in:
Tom HouléandDanilo Leal authored and GitHub committed 2026-03-23 12:54:08 +01:00
1 parent 968a13d22b
commit 302aa859f7
13 files changed
+4240 -162
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

Cargo.lock
Generated
+5
View File
@@ -3572,6 +3572,7 @@ version = "0.1.0"
dependencies = [
"anyhow",
"async-trait",
"base64 0.22.1",
"collections",
"futures 0.3.31",
"gpui",
@@ -3580,14 +3581,17 @@ dependencies = [
"net",
"parking_lot",
"postage",
"rand 0.9.2",
"schemars",
"serde",
"serde_json",
"settings",
"sha2",
"slotmap",
"smol",
"tempfile",
"terminal",
"tiny_http",
"url",
"util",
]
@@ -13189,6 +13193,7 @@ dependencies = [
"clock",
"collections",
"context_server",
"credentials_provider",
"dap",
"encoding_rs",
"extension",
crates/agent/src/tools/context_server_registry.rs
+4 -2
View File
@@ -253,12 +253,14 @@ impl ContextServerRegistry {
let project::context_server_store::ServerStatusChangedEvent { server_id, status } = event;
match status {
ContextServerStatus::Starting => {}
ContextServerStatus::Starting | ContextServerStatus::Authenticating => {}
ContextServerStatus::Running => {
self.reload_tools_for_server(server_id.clone(), cx);
self.reload_prompts_for_server(server_id.clone(), cx);
}
ContextServerStatus::Stopped | ContextServerStatus::Error(_) => {
ContextServerStatus::Stopped
| ContextServerStatus::Error(_)
| ContextServerStatus::AuthRequired => {
if let Some(registered_server) = self.registered_servers.remove(server_id) {
if !registered_server.tools.is_empty() {
cx.emit(ContextServerRegistryEvent::ToolsChanged);
crates/agent_ui/src/agent_configuration.rs
+140 -23
View File
@@ -517,11 +517,7 @@ impl AgentConfiguration {
}
}
fn render_context_servers_section(
&mut self,
window: &mut Window,
cx: &mut Context<Self>,
) -> impl IntoElement {
fn render_context_servers_section(&mut self, cx: &mut Context<Self>) -> impl IntoElement {
let context_server_ids = self.context_server_store.read(cx).server_ids();
let add_server_popover = PopoverMenu::new("add-server-popover")
@@ -601,7 +597,7 @@ impl AgentConfiguration {
} else {
parent.children(itertools::intersperse_with(
context_server_ids.iter().cloned().map(|context_server_id| {
self.render_context_server(context_server_id, window, cx)
self.render_context_server(context_server_id, cx)
.into_any_element()
}),
|| {
@@ -618,7 +614,6 @@ impl AgentConfiguration {
fn render_context_server(
&self,
context_server_id: ContextServerId,
window: &mut Window,
cx: &Context<Self>,
) -> impl use<> + IntoElement {
let server_status = self
@@ -646,6 +641,9 @@ impl AgentConfiguration {
} else {
None
};
let auth_required = matches!(server_status, ContextServerStatus::AuthRequired);
let authenticating = matches!(server_status, ContextServerStatus::Authenticating);
let context_server_store = self.context_server_store.clone();
let tool_count = self
.context_server_registry
@@ -689,11 +687,33 @@ impl AgentConfiguration {
Indicator::dot().color(Color::Muted).into_any_element(),
"Server is stopped.",
),
ContextServerStatus::AuthRequired => (
Indicator::dot().color(Color::Warning).into_any_element(),
"Authentication required.",
),
ContextServerStatus::Authenticating => (
Icon::new(IconName::LoadCircle)
.size(IconSize::XSmall)
.color(Color::Accent)
.with_keyed_rotate_animation(
SharedString::from(format!("{}-authenticating", context_server_id.0)),
3,
)
.into_any_element(),
"Waiting for authorization...",
),
};
let is_remote = server_configuration
.as_ref()
.map(|config| matches!(config.as_ref(), ContextServerConfiguration::Http { .. }))
.unwrap_or(false);
let should_show_logout_button = server_configuration.as_ref().is_some_and(|config| {
matches!(config.as_ref(), ContextServerConfiguration::Http { .. })
&& !config.has_static_auth_header()
});
let context_server_configuration_menu = PopoverMenu::new("context-server-config-menu")
.trigger_with_tooltip(
IconButton::new("context-server-config-menu", IconName::Settings)
@@ -708,6 +728,7 @@ impl AgentConfiguration {
let language_registry = self.language_registry.clone();
let workspace = self.workspace.clone();
let context_server_registry = self.context_server_registry.clone();
let context_server_store = context_server_store.clone();
move |window, cx| {
Some(ContextMenu::build(window, cx, |menu, _window, _cx| {
@@ -754,6 +775,17 @@ impl AgentConfiguration {
.ok();
}
}))
.when(should_show_logout_button, |this| {
this.entry("Log Out", None, {
let context_server_store = context_server_store.clone();
let context_server_id = context_server_id.clone();
move |_window, cx| {
context_server_store.update(cx, |store, cx| {
store.logout_server(&context_server_id, cx).log_err();
});
}
})
})
.separator()
.entry("Uninstall", None, {
let fs = fs.clone();
@@ -810,6 +842,9 @@ impl AgentConfiguration {
}
});
let feedback_base_container =
|| h_flex().py_1().min_w_0().w_full().gap_1().justify_between();
v_flex()
.min_w_0()
.id(item_id.clone())
@@ -868,6 +903,7 @@ impl AgentConfiguration {
.on_click({
let context_server_manager = self.context_server_store.clone();
let fs = self.fs.clone();
let context_server_id = context_server_id.clone();
move |state, _window, cx| {
let is_enabled = match state {
@@ -915,32 +951,113 @@ impl AgentConfiguration {
)
.map(|parent| {
if let Some(error) = error {
return parent
.child(
feedback_base_container()
.child(
h_flex()
.pr_4()
.min_w_0()
.w_full()
.gap_2()
.child(
Icon::new(IconName::XCircle)
.size(IconSize::XSmall)
.color(Color::Error),
)
.child(
div().min_w_0().flex_1().child(
Label::new(error)
.color(Color::Muted)
.size(LabelSize::Small),
),
),
)
.when(should_show_logout_button, |this| {
this.child(
Button::new("error-logout-server", "Log Out")
.style(ButtonStyle::Outlined)
.label_size(LabelSize::Small)
.on_click({
let context_server_store =
context_server_store.clone();
let context_server_id =
context_server_id.clone();
move |_event, _window, cx| {
context_server_store.update(
cx,
|store, cx| {
store
.logout_server(
&context_server_id,
cx,
)
.log_err();
},
);
}
}),
)
}),
);
}
if auth_required {
return parent.child(
h_flex()
.gap_2()
.pr_4()
.items_start()
feedback_base_container()
.child(
h_flex()
.flex_none()
.h(window.line_height() / 1.6_f32)
.justify_center()
.pr_4()
.min_w_0()
.w_full()
.gap_2()
.child(
Icon::new(IconName::XCircle)
Icon::new(IconName::Info)
.size(IconSize::XSmall)
.color(Color::Error),
.color(Color::Muted),
)
.child(
Label::new("Authenticate to connect this server")
.color(Color::Muted)
.size(LabelSize::Small),
),
)
.child(
div().w_full().child(
Label::new(error)
.buffer_font(cx)
.color(Color::Muted)
.size(LabelSize::Small),
),
Button::new("error-logout-server", "Authenticate")
.style(ButtonStyle::Outlined)
.label_size(LabelSize::Small)
.on_click({
let context_server_store = context_server_store.clone();
let context_server_id = context_server_id.clone();
move |_event, _window, cx| {
context_server_store.update(cx, |store, cx| {
store
.authenticate_server(&context_server_id, cx)
.log_err();
});
}
}),
),
);
}
if authenticating {
return parent.child(
h_flex()
.mt_1()
.pr_4()
.min_w_0()
.w_full()
.gap_2()
.child(
div().size_3().flex_shrink_0(), // Alignment Div
)
.child(
Label::new("Authenticating…")
.color(Color::Muted)
.size(LabelSize::Small),
),
);
}
parent
})
}
@@ -1234,7 +1351,7 @@ impl Render for AgentConfiguration {
.min_w_0()
.overflow_y_scroll()
.child(self.render_agent_servers_section(cx))
.child(self.render_context_servers_section(window, cx))
.child(self.render_context_servers_section(cx))
.child(self.render_provider_configuration_section(cx)),
)
.vertical_scrollbar_for(&self.scroll_handle, window, cx),
crates/agent_ui/src/agent_configuration/configure_context_server_modal.rs
+203 -65
View File
@@ -1,25 +1,27 @@
use std::sync::{Arc, Mutex};
use anyhow::{Context as _, Result};
use collections::HashMap;
use context_server::{ContextServerCommand, ContextServerId};
use editor::{Editor, EditorElement, EditorStyle};
use gpui::{
AsyncWindowContext, DismissEvent, Entity, EventEmitter, FocusHandle, Focusable, ScrollHandle,
Task, TextStyle, TextStyleRefinement, UnderlineStyle, WeakEntity, prelude::*,
Subscription, Task, TextStyle, TextStyleRefinement, UnderlineStyle, WeakEntity, prelude::*,
};
use language::{Language, LanguageRegistry};
use markdown::{Markdown, MarkdownElement, MarkdownStyle};
use notifications::status_toast::{StatusToast, ToastIcon};
use parking_lot::Mutex;
use project::{
context_server_store::{
ContextServerStatus, ContextServerStore, registry::ContextServerDescriptorRegistry,
ContextServerStatus, ContextServerStore, ServerStatusChangedEvent,
registry::ContextServerDescriptorRegistry,
},
project_settings::{ContextServerSettings, ProjectSettings},
worktree_store::WorktreeStore,
};
use serde::Deserialize;
use settings::{Settings as _, update_settings_file};
use std::sync::Arc;
use theme::ThemeSettings;
use ui::{
CommonAnimationExt, KeyBinding, Modal, ModalFooter, ModalHeader, Section, Tooltip,
@@ -237,6 +239,8 @@ fn context_server_input(existing: Option<(ContextServerId, ContextServerCommand)
format!(
r#"{{
/// Configure an MCP server that runs locally via stdin/stdout
///
/// The name of your MCP server
"{name}": {{
/// The command which runs the MCP server
@@ -280,6 +284,8 @@ fn context_server_http_input(
format!(
r#"{{
/// Configure an MCP server that you connect to over HTTP
///
/// The name of your remote MCP server
"{name}": {{
/// The URL of the remote MCP server
@@ -342,6 +348,8 @@ fn resolve_context_server_extension(
enum State {
Idle,
Waiting,
AuthRequired { server_id: ContextServerId },
Authenticating { _server_id: ContextServerId },
Error(SharedString),
}
@@ -352,6 +360,7 @@ pub struct ConfigureContextServerModal {
state: State,
original_server_id: Option<ContextServerId>,
scroll_handle: ScrollHandle,
_auth_subscription: Option<Subscription>,
}
impl ConfigureContextServerModal {
@@ -475,6 +484,7 @@ impl ConfigureContextServerModal {
cx,
),
scroll_handle: ScrollHandle::new(),
_auth_subscription: None,
})
})
})
@@ -486,6 +496,13 @@ impl ConfigureContextServerModal {
}
fn confirm(&mut self, _: &menu::Confirm, cx: &mut Context<Self>) {
if matches!(
self.state,
State::Waiting | State::AuthRequired { .. } | State::Authenticating { .. }
) {
return;
}
self.state = State::Idle;
let Some(workspace) = self.workspace.upgrade() else {
return;
@@ -515,14 +532,19 @@ impl ConfigureContextServerModal {
async move |this, cx| {
let result = wait_for_context_server_task.await;
this.update(cx, |this, cx| match result {
Ok(_) => {
Ok(ContextServerStatus::Running) => {
this.state = State::Idle;
this.show_configured_context_server_toast(id, cx);
cx.emit(DismissEvent);
}
Ok(ContextServerStatus::AuthRequired) => {
this.state = State::AuthRequired { server_id: id };
cx.notify();
}
Err(err) => {
this.set_error(err, cx);
}
Ok(_) => {}
})
}
})
@@ -558,6 +580,49 @@ impl ConfigureContextServerModal {
cx.emit(DismissEvent);
}
fn authenticate(&mut self, server_id: ContextServerId, cx: &mut Context<Self>) {
self.context_server_store.update(cx, |store, cx| {
store.authenticate_server(&server_id, cx).log_err();
});
self.state = State::Authenticating {
_server_id: server_id.clone(),
};
self._auth_subscription = Some(cx.subscribe(
&self.context_server_store,
move |this, _, event: &ServerStatusChangedEvent, cx| {
if event.server_id != server_id {
return;
}
match &event.status {
ContextServerStatus::Running => {
this._auth_subscription = None;
this.state = State::Idle;
this.show_configured_context_server_toast(event.server_id.clone(), cx);
cx.emit(DismissEvent);
}
ContextServerStatus::AuthRequired => {
this._auth_subscription = None;
this.state = State::AuthRequired {
server_id: event.server_id.clone(),
};
cx.notify();
}
ContextServerStatus::Error(error) => {
this._auth_subscription = None;
this.set_error(error.clone(), cx);
}
ContextServerStatus::Authenticating
| ContextServerStatus::Starting
| ContextServerStatus::Stopped => {}
}
},
));
cx.notify();
}
fn show_configured_context_server_toast(&self, id: ContextServerId, cx: &mut App) {
self.workspace
.update(cx, {
@@ -615,7 +680,8 @@ impl ConfigureContextServerModal {
}
fn render_modal_description(&self, window: &mut Window, cx: &mut Context<Self>) -> AnyElement {
const MODAL_DESCRIPTION: &str = "Visit the MCP server configuration docs to find all necessary arguments and environment variables.";
const MODAL_DESCRIPTION: &str =
"Check the server docs for required arguments and environment variables.";
if let ConfigurationSource::Extension {
installation_instructions: Some(installation_instructions),
@@ -637,6 +703,67 @@ impl ConfigureContextServerModal {
}
}
fn render_tab_bar(&self, cx: &mut Context<Self>) -> Option<AnyElement> {
let is_http = match &self.source {
ConfigurationSource::New { is_http, .. } => *is_http,
_ => return None,
};
let tab = |label: &'static str, active: bool| {
div()
.id(label)
.cursor_pointer()
.p_1()
.text_sm()
.border_b_1()
.when(active, |this| {
this.border_color(cx.theme().colors().border_focused)
})
.when(!active, |this| {
this.border_color(gpui::transparent_black())
.text_color(cx.theme().colors().text_muted)
.hover(|s| s.text_color(cx.theme().colors().text))
})
.child(label)
};
Some(
h_flex()
.pt_1()
.mb_2p5()
.gap_1()
.border_b_1()
.border_color(cx.theme().colors().border.opacity(0.5))
.child(
tab("Local", !is_http).on_click(cx.listener(|this, _, window, cx| {
if let ConfigurationSource::New { editor, is_http } = &mut this.source {
if *is_http {
*is_http = false;
let new_text = context_server_input(None);
editor.update(cx, |editor, cx| {
editor.set_text(new_text, window, cx);
});
}
}
})),
)
.child(
tab("Remote", is_http).on_click(cx.listener(|this, _, window, cx| {
if let ConfigurationSource::New { editor, is_http } = &mut this.source {
if !*is_http {
*is_http = true;
let new_text = context_server_http_input(None);
editor.update(cx, |editor, cx| {
editor.set_text(new_text, window, cx);
});
}
}
})),
)
.into_any_element(),
)
}
fn render_modal_content(&self, cx: &App) -> AnyElement {
let editor = match &self.source {
ConfigurationSource::New { editor, .. } => editor,
@@ -682,7 +809,10 @@ impl ConfigureContextServerModal {
fn render_modal_footer(&self, cx: &mut Context<Self>) -> ModalFooter {
let focus_handle = self.focus_handle(cx);
let is_connecting = matches!(self.state, State::Waiting);
let is_busy = matches!(
self.state,
State::Waiting | State::AuthRequired { .. } | State::Authenticating { .. }
);
ModalFooter::new()
.start_slot::<Button>(
@@ -714,36 +844,6 @@ impl ConfigureContextServerModal {
move |_, _, cx| cx.open_url(&repository_url)
}),
)
} else if let ConfigurationSource::New { is_http, .. } = &self.source {
let label = if *is_http {
"Configure Local"
} else {
"Configure Remote"
};
let tooltip = if *is_http {
"Configure an MCP server that runs on stdin/stdout."
} else {
"Configure an MCP server that you connect to over HTTP"
};
Some(
Button::new("toggle-kind", label)
.tooltip(Tooltip::text(tooltip))
.on_click(cx.listener(|this, _, window, cx| match &mut this.source {
ConfigurationSource::New { editor, is_http } => {
*is_http = !*is_http;
let new_text = if *is_http {
context_server_http_input(None)
} else {
context_server_input(None)
};
editor.update(cx, |editor, cx| {
editor.set_text(new_text, window, cx);
})
}
_ => {}
})),
)
} else {
None
},
@@ -777,7 +877,7 @@ impl ConfigureContextServerModal {
"Configure Server"
},
)
.disabled(is_connecting)
.disabled(is_busy)
.key_binding(
KeyBinding::for_action_in(&menu::Confirm, &focus_handle, cx)
.map(|kb| kb.size(rems_from_px(12.))),
@@ -791,29 +891,62 @@ impl ConfigureContextServerModal {
)
}
fn render_waiting_for_context_server() -> Div {
fn render_loading(&self, label: impl Into<SharedString>) -> Div {
h_flex()
.gap_2()
.h_8()
.gap_1p5()
.justify_center()
.child(
Icon::new(IconName::ArrowCircle)
Icon::new(IconName::LoadCircle)
.size(IconSize::XSmall)
.color(Color::Info)
.with_rotate_animation(2)
.into_any_element(),
.color(Color::Muted)
.with_rotate_animation(3),
)
.child(Label::new(label).size(LabelSize::Small).color(Color::Muted))
}
fn render_auth_required(&self, server_id: &ContextServerId, cx: &mut Context<Self>) -> Div {
h_flex()
.h_8()
.min_w_0()
.w_full()
.gap_2()
.justify_center()
.child(
h_flex()
.gap_1p5()
.child(
Icon::new(IconName::Info)
.size(IconSize::Small)
.color(Color::Muted),
)
.child(
Label::new("Authenticate to connect this server")
.size(LabelSize::Small)
.color(Color::Muted),
),
)
.child(
Label::new("Waiting for Context Server")
.size(LabelSize::Small)
.color(Color::Muted),
Button::new("authenticate-server", "Authenticate")
.style(ButtonStyle::Outlined)
.label_size(LabelSize::Small)
.on_click({
let server_id = server_id.clone();
cx.listener(move |this, _event, _window, cx| {
this.authenticate(server_id.clone(), cx);
})
}),
)
}
fn render_modal_error(error: SharedString) -> Div {
h_flex()
.gap_2()
.h_8()
.gap_1p5()
.justify_center()
.child(
Icon::new(IconName::Warning)
.size(IconSize::XSmall)
.size(IconSize::Small)
.color(Color::Warning),
)
.child(
@@ -828,7 +961,7 @@ impl Render for ConfigureContextServerModal {
fn render(&mut self, window: &mut Window, cx: &mut Context<Self>) -> impl IntoElement {
div()
.elevation_3(cx)
.w(rems(34.))
.w(rems(40.))
.key_context("ConfigureContextServerModal")
.on_action(
cx.listener(|this, _: &menu::Cancel, _window, cx| this.cancel(&menu::Cancel, cx)),
@@ -855,11 +988,18 @@ impl Render for ConfigureContextServerModal {
.overflow_y_scroll()
.track_scroll(&self.scroll_handle)
.child(self.render_modal_description(window, cx))
.children(self.render_tab_bar(cx))
.child(self.render_modal_content(cx))
.child(match &self.state {
State::Idle => div(),
State::Waiting => {
Self::render_waiting_for_context_server()
self.render_loading("Connecting Server…")
}
State::AuthRequired { server_id } => {
self.render_auth_required(&server_id.clone(), cx)
}
State::Authenticating { .. } => {
self.render_loading("Authenticating…")
}
State::Error(error) => {
Self::render_modal_error(error.clone())
@@ -878,7 +1018,7 @@ fn wait_for_context_server(
context_server_store: &Entity<ContextServerStore>,
context_server_id: ContextServerId,
cx: &mut App,
) -> Task<Result<(), Arc<str>>> {
) -> Task<Result<ContextServerStatus, Arc<str>>> {
use std::time::Duration;
const WAIT_TIMEOUT: Duration = Duration::from_secs(120);
@@ -888,31 +1028,29 @@ fn wait_for_context_server(
let context_server_id_for_timeout = context_server_id.clone();
let subscription = cx.subscribe(context_server_store, move |_, event, _cx| {
let project::context_server_store::ServerStatusChangedEvent { server_id, status } = event;
let ServerStatusChangedEvent { server_id, status } = event;
if server_id != &context_server_id {
return;
}
match status {
ContextServerStatus::Running => {
if server_id == &context_server_id
&& let Some(tx) = tx.lock().unwrap().take()
{
let _ = tx.send(Ok(()));
ContextServerStatus::Running | ContextServerStatus::AuthRequired => {
if let Some(tx) = tx.lock().take() {
let _ = tx.send(Ok(status.clone()));
}
}
ContextServerStatus::Stopped => {
if server_id == &context_server_id
&& let Some(tx) = tx.lock().unwrap().take()
{
if let Some(tx) = tx.lock().take() {
let _ = tx.send(Err("Context server stopped running".into()));
}
}
ContextServerStatus::Error(error) => {
if server_id == &context_server_id
&& let Some(tx) = tx.lock().unwrap().take()
{
if let Some(tx) = tx.lock().take() {
let _ = tx.send(Err(error.clone()));
}
}
_ => {}
ContextServerStatus::Starting | ContextServerStatus::Authenticating => {}
}
});
crates/assistant_text_thread/src/text_thread_store.rs
+4 -2
View File
@@ -901,14 +901,16 @@ impl TextThreadStore {
cx,
);
}
ContextServerStatus::Stopped | ContextServerStatus::Error(_) => {
ContextServerStatus::Stopped
| ContextServerStatus::Error(_)
| ContextServerStatus::AuthRequired => {
if let Some(slash_command_ids) =
self.context_server_slash_command_ids.remove(server_id)
{
self.slash_commands.remove(&slash_command_ids);
}
}
_ => {}
ContextServerStatus::Starting | ContextServerStatus::Authenticating => {}
}
}
crates/context_server/Cargo.toml
+4
View File
@@ -17,6 +17,7 @@ test-support = ["gpui/test-support"]
[dependencies]
anyhow.workspace = true
async-trait.workspace = true
base64.workspace = true
collections.workspace = true
futures.workspace = true
gpui.workspace = true
@@ -24,14 +25,17 @@ http_client = { workspace = true, features = ["test-support"] }
log.workspace = true
net.workspace = true
parking_lot.workspace = true
rand.workspace = true
postage.workspace = true
schemars.workspace = true
serde_json.workspace = true
serde.workspace = true
settings.workspace = true
sha2.workspace = true
slotmap.workspace = true
smol.workspace = true
tempfile.workspace = true
tiny_http.workspace = true
url = { workspace = true, features = ["serde"] }
util.workspace = true
terminal.workspace = true
crates/context_server/src/client.rs
+27 -5
View File
@@ -35,7 +35,7 @@ pub const METHOD_NOT_FOUND: i32 = -32601;
pub const INVALID_PARAMS: i32 = -32602;
pub const INTERNAL_ERROR: i32 = -32603;
type ResponseHandler = Box<dyn Send + FnOnce(Result<String, Error>)>;
type ResponseHandler = Box<dyn Send + FnOnce(String)>;
type NotificationHandler = Box<dyn Send + FnMut(Value, AsyncApp)>;
type RequestHandler = Box<dyn Send + FnMut(RequestId, &RawValue, AsyncApp)>;
@@ -62,6 +62,14 @@ pub(crate) struct Client {
#[allow(dead_code)]
transport: Arc<dyn Transport>,
request_timeout: Option<Duration>,
/// 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.
last_transport_error: Arc<Mutex<Option<anyhow::Error>>>,
}
#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
@@ -223,13 +231,16 @@ impl Client {
input.or(err)
});
let last_transport_error: Arc<Mutex<Option<anyhow::Error>>> = Arc::new(Mutex::new(None));
let output_task = cx.background_spawn({
let transport = transport.clone();
let last_transport_error = last_transport_error.clone();
Self::handle_output(
transport,
outbound_rx,
output_done_tx,
response_handlers.clone(),
last_transport_error,
)
.log_err()
});
@@ -246,6 +257,7 @@ impl Client {
output_done_rx: Mutex::new(Some(output_done_rx)),
transport,
request_timeout,
last_transport_error,
})
}
@@ -279,7 +291,7 @@ impl Client {
if let Some(handlers) = response_handlers.lock().as_mut()
&& let Some(handler) = handlers.remove(&response.id)
{
handler(Ok(message.to_string()));
handler(message.to_string());
}
} else if let Ok(notification) = serde_json::from_str::<AnyNotification>(&message) {
subscription_set.lock().notify(
@@ -315,6 +327,7 @@ impl Client {
outbound_rx: channel::Receiver<String>,
output_done_tx: barrier::Sender,
response_handlers: Arc<Mutex<Option<HashMap<RequestId, ResponseHandler>>>>,
last_transport_error: Arc<Mutex<Option<anyhow::Error>>>,
) -> anyhow::Result<()> {
let _clear_response_handlers = util::defer({
let response_handlers = response_handlers.clone();
@@ -324,7 +337,11 @@ impl Client {
});
while let Ok(message) = outbound_rx.recv().await {
log::trace!("outgoing message: {}", message);
transport.send(message).await?;
if let Err(err) = transport.send(message).await {
log::debug!("transport send failed: {:#}", err);
*last_transport_error.lock() = Some(err);
return Ok(());
}
}
drop(output_done_tx);
Ok(())
@@ -408,7 +425,7 @@ impl Client {
response = rx.fuse() => {
let elapsed = started.elapsed();
log::trace!("took {elapsed:?} to receive response to {method:?} id {id}");
match response? {
match response {
Ok(response) => {
let parsed: AnyResponse = serde_json::from_str(&response)?;
if let Some(error) = parsed.error {
@@ -419,7 +436,12 @@ impl Client {
anyhow::bail!("Invalid response: no result or error");
}
}
Err(_) => anyhow::bail!("cancelled")
Err(_canceled) => {
if let Some(err) = self.last_transport_error.lock().take() {
return Err(err);
}
anyhow::bail!("cancelled")
}
}
}
_ = cancel_fut => {
crates/context_server/src/context_server.rs
+1
View File
@@ -1,5 +1,6 @@
pub mod client;
pub mod listener;
pub mod oauth;
pub mod protocol;
#[cfg(any(test, feature = "test-support"))]
pub mod test;
crates/context_server/src/oauth.rs
+2800
View File
File diff suppressed because it is too large. Load diff
crates/context_server/src/transport/http.rs
+507 -10
View File
@@ -8,8 +8,30 @@ use parking_lot::Mutex as SyncMutex;
use smol::channel;
use std::{pin::Pin, sync::Arc};
use crate::oauth::{self, OAuthTokenProvider, WwwAuthenticate};
use crate::transport::Transport;
/// Typed errors returned by the HTTP transport that callers can downcast from
/// `anyhow::Error` to handle specific failure modes.
#[derive(Debug)]
pub enum TransportError {
/// The server returned 401 and token refresh either wasn't possible or
/// failed. The caller should initiate the OAuth authorization flow.
AuthRequired { www_authenticate: WwwAuthenticate },
}
impl std::fmt::Display for TransportError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
TransportError::AuthRequired { .. } => {
write!(f, "OAuth authorization required")
}
}
}
}
impl std::error::Error for TransportError {}
// Constants from MCP spec
const HEADER_SESSION_ID: &str = "Mcp-Session-Id";
const EVENT_STREAM_MIME_TYPE: &str = "text/event-stream";
@@ -25,8 +47,11 @@ pub struct HttpTransport {
response_rx: channel::Receiver<String>,
error_tx: channel::Sender<String>,
error_rx: channel::Receiver<String>,
// Authentication headers to include in requests
/// Static headers to include in every request (e.g. from server config).
headers: HashMap<String, String>,
/// When set, the transport attaches `Authorization: Bearer` headers and
/// handles 401 responses with token refresh + retry.
token_provider: Option<Arc<dyn OAuthTokenProvider>>,
}
impl HttpTransport {
@@ -35,6 +60,16 @@ impl HttpTransport {
endpoint: String,
headers: HashMap<String, String>,
executor: BackgroundExecutor,
) -> Self {
Self::new_with_token_provider(http_client, endpoint, headers, executor, None)
}
pub fn new_with_token_provider(
http_client: Arc<dyn HttpClient>,
endpoint: String,
headers: HashMap<String, String>,
executor: BackgroundExecutor,
token_provider: Option<Arc<dyn OAuthTokenProvider>>,
) -> Self {
let (response_tx, response_rx) = channel::unbounded();
let (error_tx, error_rx) = channel::unbounded();
@@ -49,14 +84,14 @@ impl HttpTransport {
error_tx,
error_rx,
headers,
token_provider,
}
}
/// Send a message and handle the response based on content type
async fn send_message(&self, message: String) -> Result<()> {
let is_notification =
!message.contains("\"id\":") || message.contains("notifications/initialized");
/// Build a POST request for the given message body, attaching all standard
/// headers (content-type, accept, session ID, static headers, and bearer
/// token if available).
fn build_request(&self, message: &[u8]) -> Result<http_client::Request<AsyncBody>> {
let mut request_builder = Request::builder()
.method(Method::POST)
.uri(&self.endpoint)
@@ -70,15 +105,71 @@ impl HttpTransport {
request_builder = request_builder.header(key.as_str(), value.as_str());
}
// Add session ID if we have one (except for initialize)
// Attach bearer token when a token provider is present.
if let Some(token) = self.token_provider.as_ref().and_then(|p| p.access_token()) {
request_builder = request_builder.header("Authorization", format!("Bearer {}", token));
}
// Add session ID if we have one (except for initialize).
if let Some(ref session_id) = *self.session_id.lock() {
request_builder = request_builder.header(HEADER_SESSION_ID, session_id.as_str());
}
let request = request_builder.body(AsyncBody::from(message.into_bytes()))?;
Ok(request_builder.body(AsyncBody::from(message.to_vec()))?)
}
/// Send a message and handle the response based on content type.
async fn send_message(&self, message: String) -> Result<()> {
let is_notification =
!message.contains("\"id\":") || message.contains("notifications/initialized");
// If we currently have no access token, try refreshing before sending
// the request so restored but expired sessions do not need an initial
// 401 round-trip before they can recover.
if let Some(ref provider) = self.token_provider {
if provider.access_token().is_none() {
provider.try_refresh().await.unwrap_or(false);
}
}
let request = self.build_request(message.as_bytes())?;
let mut response = self.http_client.send(request).await?;
// Handle different response types based on status and content-type
// On 401, try refreshing the token and retry once.
if response.status().as_u16() == 401 {
let www_auth_header = response
.headers()
.get("www-authenticate")
.and_then(|v| v.to_str().ok())
.unwrap_or("Bearer");
let www_authenticate =
oauth::parse_www_authenticate(www_auth_header).unwrap_or(WwwAuthenticate {
resource_metadata: None,
scope: None,
error: None,
error_description: None,
});
if let Some(ref provider) = self.token_provider {
if provider.try_refresh().await.unwrap_or(false) {
// Retry with the refreshed token.
let retry_request = self.build_request(message.as_bytes())?;
response = self.http_client.send(retry_request).await?;
// If still 401 after refresh, give up.
if response.status().as_u16() == 401 {
return Err(TransportError::AuthRequired { www_authenticate }.into());
}
} else {
return Err(TransportError::AuthRequired { www_authenticate }.into());
}
} else {
return Err(TransportError::AuthRequired { www_authenticate }.into());
}
}
// Handle different response types based on status and content-type.
match response.status() {
status if status.is_success() => {
// Check content type
@@ -233,6 +324,7 @@ impl Drop for HttpTransport {
let endpoint = self.endpoint.clone();
let session_id = self.session_id.lock().clone();
let headers = self.headers.clone();
let access_token = self.token_provider.as_ref().and_then(|p| p.access_token());
if let Some(session_id) = session_id {
self.executor
@@ -242,11 +334,17 @@ impl Drop for HttpTransport {
.uri(&endpoint)
.header(HEADER_SESSION_ID, &session_id);
// Add authentication headers if present
// Add static authentication headers.
for (key, value) in headers {
request_builder = request_builder.header(key.as_str(), value.as_str());
}
// Attach bearer token if available.
if let Some(token) = access_token {
request_builder =
request_builder.header("Authorization", format!("Bearer {}", token));
}
let request = request_builder.body(AsyncBody::empty());
if let Ok(request) = request {
@@ -257,3 +355,402 @@ impl Drop for HttpTransport {
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use async_trait::async_trait;
use gpui::TestAppContext;
use parking_lot::Mutex as SyncMutex;
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
/// A mock token provider that returns a configurable token and tracks
/// refresh attempts.
struct FakeTokenProvider {
token: SyncMutex<Option<String>>,
refreshed_token: SyncMutex<Option<String>>,
refresh_succeeds: AtomicBool,
refresh_count: AtomicUsize,
}
impl FakeTokenProvider {
fn new(token: Option<&str>, refresh_succeeds: bool) -> Arc<Self> {
Self::with_refreshed_token(token, None, refresh_succeeds)
}
fn with_refreshed_token(
token: Option<&str>,
refreshed_token: Option<&str>,
refresh_succeeds: bool,
) -> Arc<Self> {
Arc::new(Self {
token: SyncMutex::new(token.map(String::from)),
refreshed_token: SyncMutex::new(refreshed_token.map(String::from)),
refresh_succeeds: AtomicBool::new(refresh_succeeds),
refresh_count: AtomicUsize::new(0),
})
}
fn set_token(&self, token: &str) {
*self.token.lock() = Some(token.to_string());
}
fn refresh_count(&self) -> usize {
self.refresh_count.load(Ordering::SeqCst)
}
}
#[async_trait]
impl OAuthTokenProvider for FakeTokenProvider {
fn access_token(&self) -> Option<String> {
self.token.lock().clone()
}
async fn try_refresh(&self) -> Result<bool> {
self.refresh_count.fetch_add(1, Ordering::SeqCst);
let refresh_succeeds = self.refresh_succeeds.load(Ordering::SeqCst);
if refresh_succeeds {
if let Some(token) = self.refreshed_token.lock().clone() {
*self.token.lock() = Some(token);
}
}
Ok(refresh_succeeds)
}
}
fn make_fake_http_client(
handler: impl Fn(
http_client::Request<AsyncBody>,
) -> std::pin::Pin<
Box<dyn std::future::Future<Output = anyhow::Result<Response<AsyncBody>>> + Send>,
> + Send
+ Sync
+ 'static,
) -> Arc<dyn HttpClient> {
http_client::FakeHttpClient::create(handler) as Arc<dyn HttpClient>
}
fn json_response(status: u16, body: &str) -> anyhow::Result<Response<AsyncBody>> {
Ok(Response::builder()
.status(status)
.header("Content-Type", "application/json")
.body(AsyncBody::from(body.as_bytes().to_vec()))
.unwrap())
}
#[gpui::test]
async fn test_bearer_token_attached_to_requests(cx: &mut TestAppContext) {
// Capture the Authorization header from the request.
let captured_auth = Arc::new(SyncMutex::new(None::<String>));
let captured_auth_clone = captured_auth.clone();
let client = make_fake_http_client(move |req| {
let auth = req
.headers()
.get("Authorization")
.map(|v| v.to_str().unwrap().to_string());
*captured_auth_clone.lock() = auth;
Box::pin(async { json_response(200, r#"{"jsonrpc":"2.0","id":1,"result":{}}"#) })
});
let provider = FakeTokenProvider::new(Some("test-access-token"), false);
let transport = HttpTransport::new_with_token_provider(
client,
"http://mcp.example.com/mcp".to_string(),
HashMap::default(),
cx.background_executor.clone(),
Some(provider),
);
transport
.send(r#"{"jsonrpc":"2.0","id":1,"method":"ping"}"#.to_string())
.await
.expect("send should succeed");
assert_eq!(
captured_auth.lock().as_deref(),
Some("Bearer test-access-token"),
);
}
#[gpui::test]
async fn test_no_bearer_token_without_provider(cx: &mut TestAppContext) {
let captured_auth = Arc::new(SyncMutex::new(None::<String>));
let captured_auth_clone = captured_auth.clone();
let client = make_fake_http_client(move |req| {
let auth = req
.headers()
.get("Authorization")
.map(|v| v.to_str().unwrap().to_string());
*captured_auth_clone.lock() = auth;
Box::pin(async { json_response(200, r#"{"jsonrpc":"2.0","id":1,"result":{}}"#) })
});
let transport = HttpTransport::new(
client,
"http://mcp.example.com/mcp".to_string(),
HashMap::default(),
cx.background_executor.clone(),
);
transport
.send(r#"{"jsonrpc":"2.0","id":1,"method":"ping"}"#.to_string())
.await
.expect("send should succeed");
assert!(captured_auth.lock().is_none());
}
#[gpui::test]
async fn test_missing_token_triggers_refresh_before_first_request(cx: &mut TestAppContext) {
let captured_auth = Arc::new(SyncMutex::new(None::<String>));
let captured_auth_clone = captured_auth.clone();
let client = make_fake_http_client(move |req| {
let auth = req
.headers()
.get("Authorization")
.map(|v| v.to_str().unwrap().to_string());
*captured_auth_clone.lock() = auth;
Box::pin(async { json_response(200, r#"{"jsonrpc":"2.0","id":1,"result":{}}"#) })
});
let provider = FakeTokenProvider::with_refreshed_token(None, Some("refreshed-token"), true);
let transport = HttpTransport::new_with_token_provider(
client,
"http://mcp.example.com/mcp".to_string(),
HashMap::default(),
cx.background_executor.clone(),
Some(provider.clone()),
);
transport
.send(r#"{"jsonrpc":"2.0","id":1,"method":"ping"}"#.to_string())
.await
.expect("send should succeed after proactive refresh");
assert_eq!(provider.refresh_count(), 1);
assert_eq!(
captured_auth.lock().as_deref(),
Some("Bearer refreshed-token"),
);
}
#[gpui::test]
async fn test_invalid_token_still_triggers_refresh_and_retry(cx: &mut TestAppContext) {
let request_count = Arc::new(AtomicUsize::new(0));
let request_count_clone = request_count.clone();
let client = make_fake_http_client(move |_req| {
let count = request_count_clone.fetch_add(1, Ordering::SeqCst);
Box::pin(async move {
if count == 0 {
Ok(Response::builder()
.status(401)
.header(
"WWW-Authenticate",
r#"Bearer error="invalid_token", resource_metadata="https://mcp.example.com/.well-known/oauth-protected-resource""#,
)
.body(AsyncBody::from(b"Unauthorized".to_vec()))
.unwrap())
} else {
json_response(200, r#"{"jsonrpc":"2.0","id":1,"result":{}}"#)
}
})
});
let provider = FakeTokenProvider::with_refreshed_token(
Some("old-token"),
Some("refreshed-token"),
true,
);
let transport = HttpTransport::new_with_token_provider(
client,
"http://mcp.example.com/mcp".to_string(),
HashMap::default(),
cx.background_executor.clone(),
Some(provider.clone()),
);
transport
.send(r#"{"jsonrpc":"2.0","id":1,"method":"ping"}"#.to_string())
.await
.expect("send should succeed after refresh");
assert_eq!(provider.refresh_count(), 1);
assert_eq!(request_count.load(Ordering::SeqCst), 2);
}
#[gpui::test]
async fn test_401_triggers_refresh_and_retry(cx: &mut TestAppContext) {
let request_count = Arc::new(AtomicUsize::new(0));
let request_count_clone = request_count.clone();
let client = make_fake_http_client(move |_req| {
let count = request_count_clone.fetch_add(1, Ordering::SeqCst);
Box::pin(async move {
if count == 0 {
// First request: 401.
Ok(Response::builder()
.status(401)
.header(
"WWW-Authenticate",
r#"Bearer resource_metadata="https://mcp.example.com/.well-known/oauth-protected-resource""#,
)
.body(AsyncBody::from(b"Unauthorized".to_vec()))
.unwrap())
} else {
// Retry after refresh: 200.
json_response(200, r#"{"jsonrpc":"2.0","id":1,"result":{}}"#)
}
})
});
let provider = FakeTokenProvider::new(Some("old-token"), true);
// Simulate the refresh updating the token.
let provider_ref = provider.clone();
let transport = HttpTransport::new_with_token_provider(
client,
"http://mcp.example.com/mcp".to_string(),
HashMap::default(),
cx.background_executor.clone(),
Some(provider.clone()),
);
// Set the new token that will be used on retry.
provider_ref.set_token("refreshed-token");
transport
.send(r#"{"jsonrpc":"2.0","id":1,"method":"ping"}"#.to_string())
.await
.expect("send should succeed after refresh");
assert_eq!(provider_ref.refresh_count(), 1);
assert_eq!(request_count.load(Ordering::SeqCst), 2);
}
#[gpui::test]
async fn test_401_returns_auth_required_when_refresh_fails(cx: &mut TestAppContext) {
let client = make_fake_http_client(|_req| {
Box::pin(async {
Ok(Response::builder()
.status(401)
.header(
"WWW-Authenticate",
r#"Bearer resource_metadata="https://mcp.example.com/.well-known/oauth-protected-resource", scope="read write""#,
)
.body(AsyncBody::from(b"Unauthorized".to_vec()))
.unwrap())
})
});
// Refresh returns false — no new token available.
let provider = FakeTokenProvider::new(Some("stale-token"), false);
let transport = HttpTransport::new_with_token_provider(
client,
"http://mcp.example.com/mcp".to_string(),
HashMap::default(),
cx.background_executor.clone(),
Some(provider.clone()),
);
let err = transport
.send(r#"{"jsonrpc":"2.0","id":1,"method":"ping"}"#.to_string())
.await
.unwrap_err();
let transport_err = err
.downcast_ref::<TransportError>()
.expect("error should be TransportError");
match transport_err {
TransportError::AuthRequired { www_authenticate } => {
assert_eq!(
www_authenticate
.resource_metadata
.as_ref()
.map(|u| u.as_str()),
Some("https://mcp.example.com/.well-known/oauth-protected-resource"),
);
assert_eq!(
www_authenticate.scope,
Some(vec!["read".to_string(), "write".to_string()]),
);
}
}
assert_eq!(provider.refresh_count(), 1);
}
#[gpui::test]
async fn test_401_returns_auth_required_without_provider(cx: &mut TestAppContext) {
let client = make_fake_http_client(|_req| {
Box::pin(async {
Ok(Response::builder()
.status(401)
.header("WWW-Authenticate", "Bearer")
.body(AsyncBody::from(b"Unauthorized".to_vec()))
.unwrap())
})
});
// No token provider at all.
let transport = HttpTransport::new(
client,
"http://mcp.example.com/mcp".to_string(),
HashMap::default(),
cx.background_executor.clone(),
);
let err = transport
.send(r#"{"jsonrpc":"2.0","id":1,"method":"ping"}"#.to_string())
.await
.unwrap_err();
let transport_err = err
.downcast_ref::<TransportError>()
.expect("error should be TransportError");
match transport_err {
TransportError::AuthRequired { www_authenticate } => {
assert!(www_authenticate.resource_metadata.is_none());
assert!(www_authenticate.scope.is_none());
}
}
}
#[gpui::test]
async fn test_401_after_successful_refresh_still_returns_auth_required(
cx: &mut TestAppContext,
) {
// Both requests return 401 — the server rejects the refreshed token too.
let client = make_fake_http_client(|_req| {
Box::pin(async {
Ok(Response::builder()
.status(401)
.header("WWW-Authenticate", "Bearer")
.body(AsyncBody::from(b"Unauthorized".to_vec()))
.unwrap())
})
});
let provider = FakeTokenProvider::new(Some("token"), true);
let transport = HttpTransport::new_with_token_provider(
client,
"http://mcp.example.com/mcp".to_string(),
HashMap::default(),
cx.background_executor.clone(),
Some(provider.clone()),
);
let err = transport
.send(r#"{"jsonrpc":"2.0","id":1,"method":"ping"}"#.to_string())
.await
.unwrap_err();
err.downcast_ref::<TransportError>()
.expect("error should be TransportError");
// Refresh was attempted exactly once.
assert_eq!(provider.refresh_count(), 1);
}
}
crates/project/Cargo.toml
+1
View File
@@ -45,6 +45,7 @@ client.workspace = true
clock.workspace = true
collections.workspace = true
context_server.workspace = true
credentials_provider.workspace = true
dap.workspace = true
extension.workspace = true
fancy-regex.workspace = true
crates/project/src/context_server_store.rs
+543 -54
View File
@@ -7,10 +7,16 @@ 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::{ContextServer, ContextServerCommand, ContextServerId};
use futures::{FutureExt as _, future::Either, future::join_all};
use credentials_provider::CredentialsProvider;
use futures::future::Either;
use futures::{FutureExt as _, StreamExt as _, future::join_all};
use gpui::{App, AsyncApp, Context, Entity, EventEmitter, Subscription, Task, WeakEntity, actions};
use http_client::HttpClient;
use itertools::Itertools;
use rand::Rng as _;
use registry::ContextServerDescriptorRegistry;
use remote::RemoteClient;
use rpc::{AnyProtoClient, TypedEnvelope, proto};
@@ -45,6 +51,12 @@ pub enum ContextServerStatus {
Running,
Stopped,
Error(Arc<str>),
/// The server returned 401 and OAuth authorization is needed. The UI
/// should show an "Authenticate" button.
AuthRequired,
/// The OAuth browser flow is in progress — the user has been redirected
/// to the authorization server and we're waiting for the callback.
Authenticating,
}
impl ContextServerStatus {
@@ -54,6 +66,8 @@ impl ContextServerStatus {
ContextServerState::Running { .. } => ContextServerStatus::Running,
ContextServerState::Stopped { .. } => ContextServerStatus::Stopped,
ContextServerState::Error { error, .. } => ContextServerStatus::Error(error.clone()),
ContextServerState::AuthRequired { .. } => ContextServerStatus::AuthRequired,
ContextServerState::Authenticating { .. } => ContextServerStatus::Authenticating,
}
}
}
@@ -77,24 +91,42 @@ enum ContextServerState {
configuration: Arc<ContextServerConfiguration>,
error: Arc<str>,
},
/// The server requires OAuth authorization before it can be used. The
/// `OAuthDiscovery` holds everything needed to start the browser flow.
AuthRequired {
server: Arc<ContextServer>,
configuration: Arc<ContextServerConfiguration>,
discovery: Arc<OAuthDiscovery>,
},
/// The OAuth browser flow is in progress. The user has been redirected
/// to the authorization server and we're waiting for the callback.
Authenticating {
server: Arc<ContextServer>,
configuration: Arc<ContextServerConfiguration>,
_task: Task<()>,
},
}
impl ContextServerState {
pub fn server(&self) -> Arc<ContextServer> {
match self {
ContextServerState::Starting { server, .. } => server.clone(),
ContextServerState::Running { server, .. } => server.clone(),
ContextServerState::Stopped { server, .. } => server.clone(),
ContextServerState::Error { server, .. } => server.clone(),
ContextServerState::Starting { server, .. }
| ContextServerState::Running { server, .. }
| ContextServerState::Stopped { server, .. }
| ContextServerState::Error { server, .. }
| ContextServerState::AuthRequired { server, .. }
| ContextServerState::Authenticating { server, .. } => server.clone(),
}
}
pub fn configuration(&self) -> Arc<ContextServerConfiguration> {
match self {
ContextServerState::Starting { configuration, .. } => configuration.clone(),
ContextServerState::Running { configuration, .. } => configuration.clone(),
ContextServerState::Stopped { configuration, .. } => configuration.clone(),
ContextServerState::Error { configuration, .. } => configuration.clone(),
ContextServerState::Starting { configuration, .. }
| ContextServerState::Running { configuration, .. }
| ContextServerState::Stopped { configuration, .. }
| ContextServerState::Error { configuration, .. }
| ContextServerState::AuthRequired { configuration, .. }
| ContextServerState::Authenticating { configuration, .. } => configuration.clone(),
}
}
}
@@ -126,6 +158,15 @@ impl ContextServerConfiguration {
}
}
pub fn has_static_auth_header(&self) -> bool {
match self {
ContextServerConfiguration::Http { headers, .. } => headers
.keys()
.any(|k| k.eq_ignore_ascii_case("authorization")),
_ => false,
}
}
pub fn remote(&self) -> bool {
match self {
ContextServerConfiguration::Custom { remote, .. } => *remote,
@@ -517,9 +558,10 @@ impl ContextServerStore {
pub fn start_server(&mut self, server: Arc<ContextServer>, cx: &mut Context<Self>) {
cx.spawn(async move |this, cx| {
let this = this.upgrade().context("Context server store dropped")?;
let id = server.id();
let settings = this
.update(cx, |this, _| {
this.context_server_settings.get(&server.id().0).cloned()
this.context_server_settings.get(&id.0).cloned()
})
.context("Failed to get context server settings")?;
@@ -532,7 +574,7 @@ impl ContextServerStore {
});
let configuration = ContextServerConfiguration::from_settings(
settings,
server.id(),
id.clone(),
registry,
worktree_store,
cx,
@@ -590,7 +632,11 @@ impl ContextServerStore {
let id = server.id();
if matches!(
self.servers.get(&id),
Some(ContextServerState::Starting { .. } | ContextServerState::Running { .. })
Some(
ContextServerState::Starting { .. }
| ContextServerState::Running { .. }
| ContextServerState::Authenticating { .. },
)
) {
self.stop_server(&id, cx).log_err();
}
@@ -600,38 +646,20 @@ impl ContextServerStore {
let configuration = configuration.clone();
async move |this, cx| {
match server.clone().start(cx).await {
let new_state = match server.clone().start(cx).await {
Ok(_) => {
debug_assert!(server.client().is_some());
this.update(cx, |this, cx| {
this.update_server_state(
id.clone(),
ContextServerState::Running {
server,
configuration,
},
cx,
)
})
.log_err()
}
Err(err) => {
log::error!("{} context server failed to start: {}", id, err);
this.update(cx, |this, cx| {
this.update_server_state(
id.clone(),
ContextServerState::Error {
configuration,
server,
error: err.to_string().into(),
},
cx,
)
})
.log_err()
ContextServerState::Running {
server,
configuration,
}
}
Err(err) => resolve_start_failure(&id, err, server, configuration, cx).await,
};
this.update(cx, |this, cx| {
this.update_server_state(id.clone(), new_state, cx)
})
.log_err();
}
});
@@ -651,6 +679,20 @@ impl ContextServerStore {
.servers
.remove(id)
.context("Context server not found")?;
if let ContextServerConfiguration::Http { url, .. } = state.configuration().as_ref() {
let server_url = url.clone();
let id = id.clone();
cx.spawn(async move |_this, cx| {
let credentials_provider = cx.update(|cx| <dyn CredentialsProvider>::global(cx));
if let Err(err) = Self::clear_session(&credentials_provider, &server_url, &cx).await
{
log::warn!("{} failed to clear OAuth session on removal: {}", id, err);
}
})
.detach();
}
drop(state);
cx.emit(ServerStatusChangedEvent {
server_id: id.clone(),
@@ -742,29 +784,71 @@ impl ContextServerStore {
configuration
};
if let Some(server) = this.update(cx, |this, _| {
this.context_server_factory
.as_ref()
.map(|factory| factory(id.clone(), configuration.clone()))
})? {
return Ok((server, configuration));
}
let cached_token_provider: Option<Arc<dyn oauth::OAuthTokenProvider>> =
if let ContextServerConfiguration::Http { url, .. } = configuration.as_ref() {
if configuration.has_static_auth_header() {
None
} else {
let credentials_provider =
cx.update(|cx| <dyn CredentialsProvider>::global(cx));
let http_client = cx.update(|cx| cx.http_client());
match Self::load_session(&credentials_provider, url, &cx).await {
Ok(Some(session)) => {
log::info!("{} loaded cached OAuth session from keychain", id);
Some(Self::create_oauth_token_provider(
&id,
url,
session,
http_client,
credentials_provider,
cx,
))
}
Ok(None) => None,
Err(err) => {
log::warn!("{} failed to load cached OAuth session: {}", id, err);
None
}
}
}
} else {
None
};
let server: Arc<ContextServer> = this.update(cx, |this, cx| {
let global_timeout =
Self::resolve_project_settings(&this.worktree_store, cx).context_server_timeout;
if let Some(factory) = this.context_server_factory.as_ref() {
return anyhow::Ok(factory(id.clone(), configuration.clone()));
}
match configuration.as_ref() {
ContextServerConfiguration::Http {
url,
headers,
timeout,
} => anyhow::Ok(Arc::new(ContextServer::http(
id,
url,
headers.clone(),
cx.http_client(),
cx.background_executor().clone(),
Some(Duration::from_secs(
timeout.unwrap_or(global_timeout).min(MAX_TIMEOUT_SECS),
)),
)?)),
} => {
let transport = HttpTransport::new_with_token_provider(
cx.http_client(),
url.to_string(),
headers.clone(),
cx.background_executor().clone(),
cached_token_provider.clone(),
);
anyhow::Ok(Arc::new(ContextServer::new_with_timeout(
id,
Arc::new(transport),
Some(Duration::from_secs(
timeout.unwrap_or(global_timeout).min(MAX_TIMEOUT_SECS),
)),
)))
}
_ => {
let mut command = configuration
.command()
@@ -861,6 +945,310 @@ impl ContextServerStore {
ProjectSettings::get(location, cx)
}
fn create_oauth_token_provider(
id: &ContextServerId,
server_url: &url::Url,
session: OAuthSession,
http_client: Arc<dyn HttpClient>,
credentials_provider: Arc<dyn CredentialsProvider>,
cx: &mut AsyncApp,
) -> Arc<dyn oauth::OAuthTokenProvider> {
let (token_refresh_tx, mut token_refresh_rx) = futures::channel::mpsc::unbounded();
let id = id.clone();
let server_url = server_url.clone();
cx.spawn(async move |cx| {
while let Some(refreshed_session) = token_refresh_rx.next().await {
if let Err(err) =
Self::store_session(&credentials_provider, &server_url, &refreshed_session, &cx)
.await
{
log::warn!("{} failed to persist refreshed OAuth session: {}", id, err);
}
}
log::debug!("{} OAuth session persistence task ended", id);
})
.detach();
Arc::new(McpOAuthTokenProvider::new(
session,
http_client,
Some(token_refresh_tx),
))
}
/// Initiate the OAuth browser flow for a server in the `AuthRequired` state.
///
/// This starts a loopback HTTP callback server on an ephemeral port, builds
/// the authorization URL, opens the user's browser, waits for the callback,
/// exchanges the code for tokens, persists them in the keychain, and restarts
/// the server with the new token provider.
pub fn authenticate_server(
&mut self,
id: &ContextServerId,
cx: &mut Context<Self>,
) -> Result<()> {
let state = self.servers.get(id).context("Context server not found")?;
let (discovery, server, configuration) = match state {
ContextServerState::AuthRequired {
discovery,
server,
configuration,
} => (discovery.clone(), server.clone(), configuration.clone()),
_ => anyhow::bail!("Server is not in AuthRequired state"),
};
let id = id.clone();
let task = cx.spawn({
let id = id.clone();
let server = server.clone();
let configuration = configuration.clone();
async move |this, cx| {
let result = Self::run_oauth_flow(
this.clone(),
id.clone(),
discovery.clone(),
configuration.clone(),
cx,
)
.await;
if let Err(err) = &result {
log::error!("{} OAuth authentication failed: {:?}", id, err);
// Transition back to AuthRequired so the user can retry
// rather than landing in a terminal Error state.
this.update(cx, |this, cx| {
this.update_server_state(
id.clone(),
ContextServerState::AuthRequired {
server,
configuration,
discovery,
},
cx,
)
})
.log_err();
}
}
});
self.update_server_state(
id,
ContextServerState::Authenticating {
server,
configuration,
_task: task,
},
cx,
);
Ok(())
}
async fn run_oauth_flow(
this: WeakEntity<Self>,
id: ContextServerId,
discovery: Arc<OAuthDiscovery>,
configuration: Arc<ContextServerConfiguration>,
cx: &mut AsyncApp,
) -> Result<()> {
let resource = oauth::canonical_server_uri(&discovery.resource_metadata.resource);
let pkce = oauth::generate_pkce_challenge();
let mut state_bytes = [0u8; 32];
rand::rng().fill(&mut state_bytes);
let state_param: String = state_bytes.iter().map(|b| format!("{:02x}", b)).collect();
// Start a loopback HTTP server on an ephemeral port. The redirect URI
// includes this port so the browser sends the callback directly to our
// process.
let (redirect_uri, callback_rx) = oauth::start_callback_server()
.await
.context("Failed to start OAuth callback server")?;
let http_client = cx.update(|cx| cx.http_client());
let credentials_provider = cx.update(|cx| <dyn CredentialsProvider>::global(cx));
let server_url = match configuration.as_ref() {
ContextServerConfiguration::Http { url, .. } => url.clone(),
_ => anyhow::bail!("OAuth authentication only supported for HTTP servers"),
};
let client_registration =
oauth::resolve_client_registration(&http_client, &discovery, &redirect_uri)
.await
.context("Failed to resolve OAuth client registration")?;
let auth_url = oauth::build_authorization_url(
&discovery.auth_server_metadata,
&client_registration.client_id,
&redirect_uri,
&discovery.scopes,
&resource,
&pkce,
&state_param,
);
cx.update(|cx| cx.open_url(auth_url.as_str()));
let callback = callback_rx
.await
.map_err(|_| {
anyhow::anyhow!("OAuth callback server was shut down before receiving a response")
})?
.context("OAuth callback server received an invalid request")?;
if callback.state != state_param {
anyhow::bail!("OAuth state parameter mismatch (possible CSRF)");
}
let tokens = oauth::exchange_code(
&http_client,
&discovery.auth_server_metadata,
&callback.code,
&client_registration.client_id,
&redirect_uri,
&pkce.verifier,
&resource,
)
.await
.context("Failed to exchange authorization code for tokens")?;
let session = OAuthSession {
token_endpoint: discovery.auth_server_metadata.token_endpoint.clone(),
resource: discovery.resource_metadata.resource.clone(),
client_registration,
tokens,
};
Self::store_session(&credentials_provider, &server_url, &session, cx)
.await
.context("Failed to persist OAuth session in keychain")?;
let token_provider = Self::create_oauth_token_provider(
&id,
&server_url,
session,
http_client.clone(),
credentials_provider,
cx,
);
let new_server = this.update(cx, |this, cx| {
let global_timeout =
Self::resolve_project_settings(&this.worktree_store, cx).context_server_timeout;
match configuration.as_ref() {
ContextServerConfiguration::Http {
url,
headers,
timeout,
} => {
let transport = HttpTransport::new_with_token_provider(
http_client.clone(),
url.to_string(),
headers.clone(),
cx.background_executor().clone(),
Some(token_provider.clone()),
);
Ok(Arc::new(ContextServer::new_with_timeout(
id.clone(),
Arc::new(transport),
Some(Duration::from_secs(
timeout.unwrap_or(global_timeout).min(MAX_TIMEOUT_SECS),
)),
)))
}
_ => anyhow::bail!("OAuth authentication only supported for HTTP servers"),
}
})??;
this.update(cx, |this, cx| {
this.run_server(new_server, configuration, cx);
})?;
Ok(())
}
/// Store the full OAuth session in the system keychain, keyed by the
/// server's canonical URI.
async fn store_session(
credentials_provider: &Arc<dyn CredentialsProvider>,
server_url: &url::Url,
session: &OAuthSession,
cx: &AsyncApp,
) -> Result<()> {
let key = Self::keychain_key(server_url);
let json = serde_json::to_string(session)?;
credentials_provider
.write_credentials(&key, "mcp-oauth", json.as_bytes(), cx)
.await
}
/// Load the full OAuth session from the system keychain for the given
/// server URL.
async fn load_session(
credentials_provider: &Arc<dyn CredentialsProvider>,
server_url: &url::Url,
cx: &AsyncApp,
) -> Result<Option<OAuthSession>> {
let key = Self::keychain_key(server_url);
match credentials_provider.read_credentials(&key, cx).await? {
Some((_username, password_bytes)) => {
let session: OAuthSession = serde_json::from_slice(&password_bytes)?;
Ok(Some(session))
}
None => Ok(None),
}
}
/// Clear the stored OAuth session from the system keychain.
async fn clear_session(
credentials_provider: &Arc<dyn CredentialsProvider>,
server_url: &url::Url,
cx: &AsyncApp,
) -> Result<()> {
let key = Self::keychain_key(server_url);
credentials_provider.delete_credentials(&key, cx).await
}
fn keychain_key(server_url: &url::Url) -> String {
format!("mcp-oauth:{}", oauth::canonical_server_uri(server_url))
}
/// Log out of an OAuth-authenticated MCP server: clear the stored OAuth
/// session from the keychain and stop the server.
pub fn logout_server(&mut self, id: &ContextServerId, cx: &mut Context<Self>) -> Result<()> {
let state = self.servers.get(id).context("Context server not found")?;
let configuration = state.configuration();
let server_url = match configuration.as_ref() {
ContextServerConfiguration::Http { url, .. } => url.clone(),
_ => anyhow::bail!("logout only applies to HTTP servers with OAuth"),
};
let id = id.clone();
self.stop_server(&id, cx)?;
cx.spawn(async move |this, cx| {
let credentials_provider = cx.update(|cx| <dyn CredentialsProvider>::global(cx));
if let Err(err) = Self::clear_session(&credentials_provider, &server_url, &cx).await {
log::error!("{} failed to clear OAuth session: {}", id, err);
}
// Trigger server recreation so the next start uses a fresh
// transport without the old (now-invalidated) token provider.
this.update(cx, |this, cx| {
this.available_context_servers_changed(cx);
})
.log_err();
})
.detach();
Ok(())
}
fn update_server_state(
&mut self,
id: ContextServerId,
@@ -1014,3 +1402,104 @@ impl ContextServerStore {
Ok(())
}
}
/// Determines the appropriate server state after a start attempt fails.
///
/// When the error is an HTTP 401 with no static auth header configured,
/// attempts OAuth discovery so the UI can offer an authentication flow.
async fn resolve_start_failure(
id: &ContextServerId,
err: anyhow::Error,
server: Arc<ContextServer>,
configuration: Arc<ContextServerConfiguration>,
cx: &AsyncApp,
) -> ContextServerState {
let www_authenticate = err.downcast_ref::<TransportError>().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(),
};
}
};
// 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 www_authenticate.is_none() {
let credentials_provider = cx.update(|cx| <dyn CredentialsProvider>::global(cx));
match ContextServerStore::load_session(&credentials_provider, &server_url, cx).await {
Ok(Some(_)) => {
log::info!("{id} start failed with a cached OAuth session present; clearing it");
ContextServerStore::clear_session(&credentials_provider, &server_url, cx)
.await
.log_err();
}
_ => {
log::error!("{id} context server failed to start: {err}");
return ContextServerState::Error {
configuration,
server,
error: err.to_string().into(),
};
}
}
}
let default_www_authenticate = oauth::WwwAuthenticate {
resource_metadata: None,
scope: None,
error: None,
error_description: None,
};
let www_authenticate = www_authenticate
.as_ref()
.unwrap_or(&default_www_authenticate);
let http_client = cx.update(|cx| cx.http_client());
match context_server::oauth::discover(&http_client, &server_url, www_authenticate).await {
Ok(discovery) => {
log::info!(
"{id} requires OAuth authorization (auth server: {})",
discovery.auth_server_metadata.issuer,
);
ContextServerState::AuthRequired {
server,
configuration,
discovery: Arc::new(discovery),
}
}
Err(discovery_err) => {
log::error!("{id} OAuth discovery failed: {discovery_err}");
ContextServerState::Error {
configuration,
server,
error: format!("OAuth discovery failed: {discovery_err}").into(),
}
}
}
}
crates/ui/src/components/modal.rs
+1 -1
View File
@@ -162,7 +162,7 @@ impl RenderOnce for ModalHeader {
children.insert(
0,
Headline::new(headline)
.size(HeadlineSize::XSmall)
.size(HeadlineSize::Small)
.color(Color::Muted)
.into_any_element(),
);