refactor: make oauth discovery reactive (#1052)

This commit is contained in:
Dale Seo 2026-07-27 10:43:06 -04:00 • committed by GitHub
commit c284607794
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 581 additions and 241 deletions

View file

@ -195,25 +195,73 @@ const REDIRECT_URI: &str = "http://localhost:3000/callback";
const SCOPE_STEP_UP_INITIAL_SCOPES: &[&str] = &["mcp:basic"];
const SCOPE_STEP_UP_ESCALATED_SCOPES: &[&str] = &["mcp:basic", "mcp:write"];
/// Perform the headless OAuth authorization-code flow.
/// Attempt the real connection unauthenticated and return the server's
/// `WWW-Authenticate` challenge from the 401 — the reactive discovery
/// trigger.
///
/// 1. Discover metadata, register (or use CIMD), get auth URL
/// 2. Fetch the auth URL with redirect:manual → extract code from Location header
/// 3. Exchange code for token
/// 4. Return an `AuthClient` wrapping `reqwest::Client`
/// `None` (server accepted the unauthenticated connection, which is then
/// closed cleanly) is a legitimate outcome, not an error: the scope-step-up
/// and scope-retry-limit mocks allow unauthenticated `initialize` and only
/// enforce authorization on tool calls.
async fn initialize_challenge(
server_url: &str,
lifecycle: ClientLifecycleMode,
) -> anyhow::Result<Option<String>> {
let transport = StreamableHttpClientTransport::from_uri(server_url);
match BasicClientHandler
.serve_with_lifecycle(transport, lifecycle)
.await
{
Ok(client) => {
client.cancel().await.ok();
Ok(None)
}
Err(error) => match error.auth_challenge() {
Some(challenge) => Ok(Some(challenge.to_string())),
None => Err(error.into()),
},
}
}
fn with_optional_challenge(
request: AuthorizationRequest,
challenge: Option<String>,
) -> AuthorizationRequest {
match challenge {
Some(challenge) => request.with_challenge(challenge),
None => request,
}
}
/// Perform the headless OAuth authorization-code flow, reactively:
///
/// 1. Attempt the real connection; take the 401's WWW-Authenticate challenge
/// 2. Discover from the challenge, register (or use CIMD), get auth URL
/// 3. Fetch the auth URL with redirect:manual → extract code from Location header
/// 4. Exchange code for token
/// 5. Return an `AuthClient` wrapping `reqwest::Client`
async fn perform_oauth_flow(
server_url: &str,
_ctx: &ConformanceContext,
) -> anyhow::Result<AuthClient<reqwest::Client>> {
// Always the discover lifecycle here (not `conformance_lifecycle()`):
// this flow serves `run_auth_client`, whose 2026-07-28 auth mocks require
// the per-request MCP-Protocol-Version negotiation.
let challenge = initialize_challenge(
server_url,
ClientLifecycleMode::Discover {
preferred_versions: preferred_protocol_versions(),
},
)
.await?;
let mut oauth = OAuthState::new(server_url, None).await?;
// Discover + register + get auth URL
let request = AuthorizationRequest::new(REDIRECT_URI)
.with_client_name("conformance-client")
.with_client_metadata_url(CIMD_CLIENT_METADATA_URL);
oauth
.start_authorization(
AuthorizationRequest::new(REDIRECT_URI)
.with_client_name("conformance-client")
.with_client_metadata_url(CIMD_CLIENT_METADATA_URL),
)
.start_authorization(with_optional_challenge(request, challenge))
.await?;
let auth_url = oauth.get_authorization_url().await?;
@ -255,14 +303,14 @@ async fn perform_oauth_flow_preregistered(
client_id: &str,
client_secret: &str,
) -> anyhow::Result<AuthClient<reqwest::Client>> {
let challenge = initialize_challenge(server_url, conformance_lifecycle()).await?;
let mut oauth = OAuthState::new(server_url, None).await?;
let request = AuthorizationRequest::new(REDIRECT_URI)
.with_preregistered_client(client_id)
.with_client_secret(client_secret);
oauth
.start_authorization(
AuthorizationRequest::new(REDIRECT_URI)
.with_preregistered_client(client_id)
.with_client_secret(client_secret),
)
.start_authorization(with_optional_challenge(request, challenge))
.await?;
let auth_url = oauth.get_authorization_url().await?;
@ -325,14 +373,14 @@ async fn run_auth_scope_step_up_client(
server_url: &str,
_ctx: &ConformanceContext,
) -> anyhow::Result<()> {
let challenge = initialize_challenge(server_url, conformance_lifecycle()).await?;
let mut oauth = OAuthState::new(server_url, None).await?;
let request = AuthorizationRequest::new(REDIRECT_URI)
.with_scopes(SCOPE_STEP_UP_INITIAL_SCOPES.iter().copied())
.with_client_name("conformance-client")
.with_client_metadata_url(CIMD_CLIENT_METADATA_URL);
oauth
.start_authorization(
AuthorizationRequest::new(REDIRECT_URI)
.with_scopes(SCOPE_STEP_UP_INITIAL_SCOPES.iter().copied())
.with_client_name("conformance-client")
.with_client_metadata_url(CIMD_CLIENT_METADATA_URL),
)
.start_authorization(with_optional_challenge(request, challenge))
.await?;
let auth_url = oauth.get_authorization_url().await?;
@ -427,15 +475,15 @@ async fn run_auth_scope_retry_limit_client(
) -> anyhow::Result<()> {
let max_retries = 3u32;
let mut attempt = 0u32;
let challenge = initialize_challenge(server_url, conformance_lifecycle()).await?;
loop {
let mut oauth = OAuthState::new(server_url, None).await?;
let request = AuthorizationRequest::new(REDIRECT_URI)
.with_client_name("conformance-client")
.with_client_metadata_url(CIMD_CLIENT_METADATA_URL);
oauth
.start_authorization(
AuthorizationRequest::new(REDIRECT_URI)
.with_client_name("conformance-client")
.with_client_metadata_url(CIMD_CLIENT_METADATA_URL),
)
.start_authorization(with_optional_challenge(request, challenge.clone()))
.await?;
let auth_url = oauth.get_authorization_url().await?;
let callback = headless_authorize(&auth_url).await?;
@ -509,7 +557,10 @@ async fn migration_token(
return Ok(manager.get_access_token().await?);
}
let resolution = manager.resolve_metadata().await?;
let challenge = initialize_challenge(server_url, conformance_lifecycle()).await?;
let resolution = manager
.resolve_metadata_from_challenge(challenge.as_deref())
.await?;
manager.set_metadata(resolution.metadata);
manager
.register_client("conformance-client", REDIRECT_URI, &[])
@ -609,7 +660,10 @@ async fn run_client_credentials_basic(
.unwrap_or("conformance-test-secret");
let mut manager = AuthorizationManager::new(server_url).await?;
let resolution = manager.resolve_metadata().await?;
let challenge = initialize_challenge(server_url, conformance_lifecycle()).await?;
let resolution = manager
.resolve_metadata_from_challenge(challenge.as_deref())
.await?;
let token_endpoint = resolution.metadata.token_endpoint.clone();
manager.set_metadata(resolution.metadata);

View file

@ -86,6 +86,31 @@ impl ClientInitializeError {
context: context.into(),
}
}
/// The `WWW-Authenticate` challenge from the 401/403 the transport hit
/// during initialization, if that is why initialization failed.
///
/// This is the trigger of the reactive OAuth flow: feed the challenge to
/// `AuthorizationRequest::with_challenge` to authorize, then reconnect.
#[cfg(feature = "transport-streamable-http-client")]
pub fn auth_challenge(&self) -> Option<&str> {
use crate::transport::streamable_http_client::{AuthRequiredError, InsufficientScopeError};
let Self::TransportError { error, .. } = self else {
return None;
};
let mut source: Option<&(dyn std::error::Error + 'static)> = Some(error.error.as_ref());
while let Some(current) = source {
if let Some(auth_required) = current.downcast_ref::<AuthRequiredError>() {
return Some(&auth_required.www_authenticate_header);
}
if let Some(insufficient_scope) = current.downcast_ref::<InsufficientScopeError>() {
return Some(&insufficient_scope.www_authenticate_header);
}
source = current.source();
}
None
}
}
/// Helper function to get the next message from the stream

View file

@ -31,12 +31,6 @@ use crate::transport::common::http_header::HEADER_MCP_PROTOCOL_VERSION;
const DEFAULT_HTTP_TIMEOUT: Duration = Duration::from_secs(30);
const MAX_OAUTH_HTTP_RESPONSE_BODY_BYTES: usize = 1024 * 1024;
const MAX_OAUTH_DISCOVERY_REDIRECTS: usize = 10;
const RESOURCE_METADATA_POST_PROBE_BODY: &str = concat!(
r#"{"jsonrpc":"2.0","id":"auth-discovery","method":"initialize","params":{"#,
r#""protocolVersion":"2024-11-05","capabilities":{},"#,
r#""clientInfo":{"name":"rmcp-auth-discovery","version":"0.0.0"}}"#,
r#"}"#
);
const CLOUD_METADATA_HOSTS: &[&str] = &[
"metadata",
"metadata.google.internal",
@ -634,6 +628,12 @@ pub struct WWWAuthenticateParams {
}
impl WWWAuthenticateParams {
/// Parse a `WWW-Authenticate` header value, resolving a relative
/// `resource_metadata` URL against `base_url`.
pub fn parse(header: &str, base_url: &Url) -> Self {
AuthorizationManager::extract_www_authenticate_params(header, base_url)
}
/// check if this is an insufficient_scope error
pub fn is_insufficient_scope(&self) -> bool {
self.error.as_deref() == Some("insufficient_scope")
@ -736,6 +736,11 @@ pub struct AuthorizationRequest {
/// OIDC Dynamic Client Registration `application_type` (SEP-837),
/// e.g. `"native"` or `"web"`.
pub application_type: Option<String>,
/// `WWW-Authenticate` header value from a real request's 401 response.
/// When set, discovery is seeded from the challenge (its
/// `resource_metadata` URL and `scope` hint) instead of probing the
/// server — the reactive discovery path.
pub challenge: Option<String>,
}
impl AuthorizationRequest {
@ -750,6 +755,7 @@ impl AuthorizationRequest {
client_secret: None,
client_metadata_url: None,
application_type: None,
challenge: None,
}
}
@ -805,6 +811,13 @@ impl AuthorizationRequest {
self.application_type = Some(application_type.into());
self
}
/// Seed discovery from the `WWW-Authenticate` header value of a real
/// request's 401 response instead of probing the server.
pub fn with_challenge(mut self, www_authenticate: impl Into<String>) -> Self {
self.challenge = Some(www_authenticate.into());
self
}
}
// add type aliases for oauth2 types
@ -1419,6 +1432,52 @@ impl AuthorizationManager {
})
}
/// Resolve authorization server metadata starting from the
/// `WWW-Authenticate` challenge of a real request's 401 response — the
/// reactive discovery path, matching the TypeScript and Python SDKs.
///
/// Seeds scope selection with the challenge's `scope` hint and prefers
/// the challenge's `resource_metadata` URL; falls back to
/// [`resolve_metadata`](Self::resolve_metadata) when the challenge is
/// `None` or carries no usable metadata pointer.
pub async fn resolve_metadata_from_challenge(
&self,
www_authenticate: Option<&str>,
) -> Result<AuthorizationMetadataResolution, AuthError> {
let Some(www_authenticate) = www_authenticate else {
return self.resolve_metadata().await;
};
let params = WWWAuthenticateParams::parse(www_authenticate, &self.base_url);
self.record_challenge_scope(&params).await;
if let Some(resource_metadata_url) = &params.resource_metadata_url
&& let Some(metadata) = self
.discover_oauth_server_from_resource_metadata_url(resource_metadata_url)
.await?
{
return Ok(AuthorizationMetadataResolution {
metadata,
source: AuthorizationMetadataSource::ProtectedResourceMetadata,
});
}
self.resolve_metadata().await
}
/// Store a challenge's `scope` hint for later scope selection.
async fn record_challenge_scope(&self, params: &WWWAuthenticateParams) {
let Some(scope) = &params.scope else {
return;
};
let scopes: Vec<String> = scope.split_whitespace().map(str::to_string).collect();
if scopes.is_empty() {
return;
}
debug!("WWW-Authenticate challenge contains scope: {scope}");
*self.www_auth_scopes.write().await = scopes;
}
fn legacy_authorization_metadata(base_url: &Url) -> AuthorizationMetadata {
let endpoint = |path: &str| {
let mut url = base_url.clone();
@ -2056,7 +2115,7 @@ impl AuthorizationManager {
/// refresh token or the server rejected it, return `AuthorizationRequired`
/// so the caller can re-prompt the user. Infrastructure errors (e.g. store
/// I/O failures, misconfigured client) are propagated as-is.
async fn try_refresh_or_reauth(&self) -> Result<String, AuthError> {
pub(crate) async fn try_refresh_or_reauth(&self) -> Result<String, AuthError> {
match self.refresh_token().await {
Ok(new_creds) => {
tracing::info!("Refreshed access token.");
@ -2328,12 +2387,19 @@ impl AuthorizationManager {
async fn discover_oauth_server_via_resource_metadata(
&self,
) -> Result<Option<AuthorizationMetadata>, AuthError> {
let Some(resource_metadata_url) = self.discover_resource_metadata_url().await? else {
let Some(resource_metadata_url) = self.discover_resource_metadata_url().await else {
return Ok(None);
};
self.discover_oauth_server_from_resource_metadata_url(&resource_metadata_url)
.await
}
async fn discover_oauth_server_from_resource_metadata_url(
&self,
resource_metadata_url: &Url,
) -> Result<Option<AuthorizationMetadata>, AuthError> {
let Some(resource_metadata) = self
.fetch_resource_metadata_from_url(&resource_metadata_url)
.fetch_resource_metadata_from_url(resource_metadata_url)
.await?
else {
return Ok(None);
@ -2443,11 +2509,10 @@ impl AuthorizationManager {
&& Self::is_same_origin(&root_resource, &path_resource)
}
async fn discover_resource_metadata_url(&self) -> Result<Option<Url>, AuthError> {
if let Ok(Some(resource_metadata_url)) =
self.fetch_resource_metadata_url(&self.base_url, true).await
async fn discover_resource_metadata_url(&self) -> Option<Url> {
if let Some(resource_metadata_url) = self.probe_resource_metadata_url(&self.base_url).await
{
return Ok(Some(resource_metadata_url));
return Some(resource_metadata_url);
}
// If the primary URL doesn't use WWW-Authenticate, try oauth-protected-resource discovery.
@ -2459,86 +2524,42 @@ impl AuthorizationManager {
discovery_url.set_query(None);
discovery_url.set_fragment(None);
discovery_url.set_path(&candidate_path);
if let Ok(Some(resource_metadata_url)) = self
.fetch_resource_metadata_url(&discovery_url, false)
.await
if let Some(resource_metadata_url) =
self.probe_resource_metadata_url(&discovery_url).await
{
return Ok(Some(resource_metadata_url));
return Some(resource_metadata_url);
}
}
Ok(None)
None
}
/// Extract the resource metadata url from the WWW-Authenticate header value.
/// Probe `url` with a GET, extracting the resource metadata url from a
/// 200 (the url itself is the metadata document) or from a 401's
/// WWW-Authenticate header value.
/// https://www.rfc-editor.org/rfc/rfc9728.html#name-use-of-www-authenticate-for
async fn fetch_resource_metadata_url(
&self,
url: &Url,
allow_post_probe: bool,
) -> Result<Option<Url>, AuthError> {
async fn probe_resource_metadata_url(&self, url: &Url) -> Option<Url> {
let response = match self.discovery_get(url).await {
Ok(r) => r,
Err(e) => {
debug!("resource metadata probe failed: {}", e);
return Ok(None);
return None;
}
};
match response.status() {
StatusCode::OK => Ok(Some(url.clone())),
StatusCode::UNAUTHORIZED => Ok(self
.extract_resource_metadata_url_from_www_authenticate(&response)
.await),
StatusCode::NOT_FOUND | StatusCode::METHOD_NOT_ALLOWED if allow_post_probe => {
self.fetch_resource_metadata_url_with_post_probe(url).await
StatusCode::OK => Some(url.clone()),
StatusCode::UNAUTHORIZED => {
self.extract_resource_metadata_url_from_www_authenticate(&response)
.await
}
status => {
debug!("resource metadata probe returned unexpected status: {status}");
Ok(None)
None
}
}
}
async fn fetch_resource_metadata_url_with_post_probe(
&self,
url: &Url,
) -> Result<Option<Url>, AuthError> {
let request = oauth2::http::Request::builder()
.method("POST")
.uri(url.as_str())
.header(HEADER_MCP_PROTOCOL_VERSION, "2024-11-05")
.header(CONTENT_TYPE, "application/json")
.body(RESOURCE_METADATA_POST_PROBE_BODY.as_bytes().to_vec())
.map_err(|error| AuthError::InternalError(error.to_string()))?;
let response = match self
.http_client
.execute(OAuthHttpRequest::new(
request,
OAuthHttpRedirectPolicy::Stop,
))
.await
{
Ok(response) => response,
Err(error) => {
debug!("resource metadata POST probe failed: {}", error);
return Ok(None);
}
};
if response.status() != StatusCode::UNAUTHORIZED {
debug!(
"resource metadata POST probe returned unexpected status: {}",
response.status()
);
return Ok(None);
}
Ok(self
.extract_resource_metadata_url_from_www_authenticate(&response)
.await)
}
async fn extract_resource_metadata_url_from_www_authenticate(
&self,
response: &HttpResponse,
@ -2549,14 +2570,9 @@ impl AuthorizationManager {
continue;
};
let params = Self::extract_www_authenticate_params(value_str, &self.base_url);
if let Some(url) = params.resource_metadata_url {
if let Some(scope) = &params.scope {
debug!("WWW-Authenticate header contains scope: {}", scope);
let scopes: Vec<String> =
scope.split_whitespace().map(|s| s.to_string()).collect();
*self.www_auth_scopes.write().await = scopes;
}
parsed_url = Some(url);
if params.resource_metadata_url.is_some() {
self.record_challenge_scope(&params).await;
parsed_url = params.resource_metadata_url;
break;
}
}
@ -3432,7 +3448,7 @@ impl OAuthState {
)
}
async fn placeholder(&self) -> Result<Self, AuthError> {
async fn placeholder_state(&self) -> Result<Self, AuthError> {
let (http_client, refresh_redirect_policy) = self.oauth_http_client_config();
Ok(OAuthState::Unauthorized(
AuthorizationManager::new_inner(
@ -3548,7 +3564,7 @@ impl OAuthState {
&mut self,
request: AuthorizationRequest,
) -> Result<(), AuthError> {
let placeholder = self.placeholder().await?;
let placeholder = self.placeholder_state().await?;
let old = std::mem::replace(self, placeholder);
let OAuthState::Unauthorized(mut manager) = old else {
*self = old;
@ -3557,7 +3573,10 @@ impl OAuthState {
));
};
debug!("start discovery");
let metadata = match manager.resolve_metadata().await {
let resolution = manager
.resolve_metadata_from_challenge(request.challenge.as_deref())
.await;
let metadata = match resolution {
Ok(resolution) => resolution.metadata,
Err(e) => {
*self = OAuthState::Unauthorized(manager);
@ -3580,7 +3599,7 @@ impl OAuthState {
/// complete authorization
pub async fn complete_authorization(&mut self) -> Result<(), AuthError> {
let placeholder = self.placeholder().await?;
let placeholder = self.placeholder_state().await?;
if let OAuthState::Session(session) = std::mem::replace(self, placeholder) {
*self = OAuthState::Authorized(session.auth_manager);
Ok(())
@ -3590,7 +3609,7 @@ impl OAuthState {
}
/// covert to authorized http client
pub async fn to_authorized_http_client(&mut self) -> Result<(), AuthError> {
let placeholder = self.placeholder().await?;
let placeholder = self.placeholder_state().await?;
if let OAuthState::Authorized(manager) = std::mem::replace(self, placeholder) {
*self = OAuthState::AuthorizedHttpClient(AuthorizedHttpClient::new(
Arc::new(manager),
@ -3610,7 +3629,7 @@ impl OAuthState {
required_scope: &str,
redirect_uri: &str,
) -> Result<String, AuthError> {
let placeholder = self.placeholder().await?;
let placeholder = self.placeholder_state().await?;
let old = std::mem::replace(self, placeholder);
let OAuthState::Authorized(manager) = old else {
*self = old;
@ -3742,7 +3761,7 @@ impl OAuthState {
&mut self,
config: ClientCredentialsConfig,
) -> Result<(), AuthError> {
let placeholder = self.placeholder().await?;
let placeholder = self.placeholder_state().await?;
let OAuthState::Unauthorized(mut manager) = std::mem::replace(self, placeholder) else {
return Err(AuthError::InternalError(
"Client credentials flow requires Unauthorized state".to_string(),
@ -4137,7 +4156,6 @@ mod tests {
.body(Vec::new())
.unwrap();
let client = RecordingOAuthHttpClient::with_responses(vec![
empty_response(404),
challenge,
http_response(
200,
@ -4178,7 +4196,6 @@ mod tests {
(
"https://auth.example.com/tenant1/token",
vec![
"https://mcp.example.com/mcp",
"https://mcp.example.com/mcp",
"https://mcp.example.com/custom/metadata/location.json",
"https://auth.example.com/.well-known/oauth-authorization-server/tenant1",
@ -4187,14 +4204,12 @@ mod tests {
],
)
);
assert_eq!(
assert!(
client
.requests()
.iter()
.take(2)
.map(|request| request.method.as_str())
.collect::<Vec<_>>(),
vec!["GET", "POST"]
.all(|request| request.method == "GET"),
"discovery must not send non-GET requests"
);
}
@ -4205,7 +4220,6 @@ mod tests {
empty_response(404),
empty_response(404),
empty_response(404),
empty_response(404),
]);
let manager = AuthorizationManager::new_with_oauth_http_client(
"https://legacy.example.com/",
@ -4234,7 +4248,6 @@ mod tests {
"https://legacy.example.com/token",
Some("https://legacy.example.com/register"),
vec![
"https://legacy.example.com/",
"https://legacy.example.com/",
"https://legacy.example.com/.well-known/oauth-protected-resource",
"https://legacy.example.com/.well-known/oauth-authorization-server",
@ -4293,6 +4306,85 @@ mod tests {
);
}
#[tokio::test]
async fn resolve_metadata_from_challenge_uses_challenge_pointer_and_scope() {
let client = RecordingOAuthHttpClient::with_responses(vec![
http_response(
200,
serde_json::json!({
"resource": "https://mcp.example.com/mcp",
"authorization_servers": ["https://auth.example.com"]
}),
),
http_response(
200,
serde_json::json!({
"issuer": "https://auth.example.com",
"authorization_endpoint": "https://auth.example.com/authorize",
"token_endpoint": "https://auth.example.com/token"
}),
),
]);
let manager = AuthorizationManager::new_with_oauth_http_client(
"https://mcp.example.com/mcp",
Arc::new(client.clone()),
)
.await
.unwrap();
let resolution = manager
.resolve_metadata_from_challenge(Some(
r#"Bearer resource_metadata="https://mcp.example.com/custom/prm.json", scope="mcp:read mcp:write""#,
))
.await
.unwrap();
assert_eq!(
(
resolution.source,
resolution.metadata.token_endpoint.as_str(),
manager.select_scopes(None, &[]),
client.requests().first().map(|request| request.uri.clone()),
),
(
AuthorizationMetadataSource::ProtectedResourceMetadata,
"https://auth.example.com/token",
vec!["mcp:read".to_string(), "mcp:write".to_string()],
// discovery starts at the challenge's pointer: no probing
Some("https://mcp.example.com/custom/prm.json".to_string()),
)
);
}
#[tokio::test]
async fn resolve_metadata_from_challenge_falls_back_without_metadata_pointer() {
let client = RecordingOAuthHttpClient::with_responses(vec![
empty_response(404),
empty_response(404),
empty_response(404),
empty_response(404),
]);
let manager = AuthorizationManager::new_with_oauth_http_client(
"https://legacy.example.com/",
Arc::new(client),
)
.await
.unwrap();
let resolution = manager
.resolve_metadata_from_challenge(Some(r#"Bearer scope="mcp:basic""#))
.await
.unwrap();
assert_eq!(
(resolution.source, manager.select_scopes(None, &[]),),
(
AuthorizationMetadataSource::LegacyEndpointFallback,
vec!["mcp:basic".to_string()],
)
);
}
#[tokio::test]
async fn resolve_metadata_reports_authorization_server_metadata() {
let client = RecordingOAuthHttpClient::with_responses(vec![

View file

@ -1,11 +1,70 @@
use std::collections::HashMap;
use http::{HeaderName, HeaderValue};
use tracing::debug;
use crate::transport::{
auth::AuthClient,
auth::{AuthClient, AuthError},
streamable_http_client::{StreamableHttpClient, StreamableHttpError},
};
impl<C> AuthClient<C>
where
C: StreamableHttpClient + Send + Sync,
{
/// Run `call` with a token when one is available, reacting to the
/// server's auth verdict instead of requiring credentials up front:
///
/// - no usable credentials → the request goes out unauthenticated, and a
/// 401 propagates as [`StreamableHttpError::AuthRequired`] carrying the
/// `WWW-Authenticate` challenge for the caller to authorize with;
/// - a token the server rejects (e.g. revoked) → one silent refresh, one
/// retry, then the challenge propagates.
async fn call_reacting_to_challenges<T, F, Fut>(
&self,
auth_token: Option<String>,
call: F,
) -> Result<T, StreamableHttpError<C::Error>>
where
F: Fn(Option<String>) -> Fut,
Fut: Future<Output = Result<T, StreamableHttpError<C::Error>>>,
{
// Missing credentials are not an error in the reactive model: the
// request goes out unauthenticated and the server's 401 challenge
// drives authorization.
let auth_token = match auth_token {
None => match self.get_access_token().await {
Ok(token) => Some(token),
Err(AuthError::AuthorizationRequired) => None,
Err(error) => return Err(error.into()),
},
token => token,
};
match call(auth_token.clone()).await {
Err(StreamableHttpError::AuthRequired(challenge)) => {
// One silent recovery attempt: refresh the rejected token and
// retry only when the refresh actually produced a new one.
let Some(sent_token) = auth_token else {
return Err(StreamableHttpError::AuthRequired(challenge));
};
let refreshed = {
let manager = self.auth_manager.lock().await;
manager.try_refresh_or_reauth().await
};
match refreshed {
Ok(fresh_token) if fresh_token != sent_token => call(Some(fresh_token)).await,
Ok(_) => Err(StreamableHttpError::AuthRequired(challenge)),
Err(error) => {
debug!("token refresh after server rejection failed: {error}");
Err(StreamableHttpError::AuthRequired(challenge))
}
}
}
result => result,
}
}
}
impl<C> StreamableHttpClient for AuthClient<C>
where
C: StreamableHttpClient + Send + Sync,
@ -16,16 +75,21 @@ where
&self,
uri: std::sync::Arc<str>,
session_id: std::sync::Arc<str>,
mut auth_token: Option<String>,
auth_token: Option<String>,
custom_headers: HashMap<HeaderName, HeaderValue>,
) -> Result<(), crate::transport::streamable_http_client::StreamableHttpError<Self::Error>>
{
if auth_token.is_none() {
auth_token = Some(self.get_access_token().await?);
}
self.http_client
.delete_session(uri, session_id, auth_token, custom_headers)
.await
self.call_reacting_to_challenges(auth_token, |token| {
let uri = uri.clone();
let session_id = session_id.clone();
let custom_headers = custom_headers.clone();
async move {
self.http_client
.delete_session(uri, session_id, token, custom_headers)
.await
}
})
.await
}
async fn get_stream(
@ -33,18 +97,24 @@ where
uri: std::sync::Arc<str>,
session_id: Option<std::sync::Arc<str>>,
last_event_id: Option<String>,
mut auth_token: Option<String>,
auth_token: Option<String>,
custom_headers: HashMap<HeaderName, HeaderValue>,
) -> Result<
futures::stream::BoxStream<'static, Result<sse_stream::Sse, sse_stream::Error>>,
crate::transport::streamable_http_client::StreamableHttpError<Self::Error>,
> {
if auth_token.is_none() {
auth_token = Some(self.get_access_token().await?);
}
self.http_client
.get_stream(uri, session_id, last_event_id, auth_token, custom_headers)
.await
self.call_reacting_to_challenges(auth_token, |token| {
let uri = uri.clone();
let session_id = session_id.clone();
let last_event_id = last_event_id.clone();
let custom_headers = custom_headers.clone();
async move {
self.http_client
.get_stream(uri, session_id, last_event_id, token, custom_headers)
.await
}
})
.await
}
async fn get_stream_with_max_sse_event_size(
@ -52,26 +122,32 @@ where
uri: std::sync::Arc<str>,
session_id: Option<std::sync::Arc<str>>,
last_event_id: Option<String>,
mut auth_token: Option<String>,
auth_token: Option<String>,
custom_headers: HashMap<HeaderName, HeaderValue>,
max_sse_event_size: usize,
) -> Result<
futures::stream::BoxStream<'static, Result<sse_stream::Sse, sse_stream::Error>>,
crate::transport::streamable_http_client::StreamableHttpError<Self::Error>,
> {
if auth_token.is_none() {
auth_token = Some(self.get_access_token().await?);
}
self.http_client
.get_stream_with_max_sse_event_size(
uri,
session_id,
last_event_id,
auth_token,
custom_headers,
max_sse_event_size,
)
.await
self.call_reacting_to_challenges(auth_token, |token| {
let uri = uri.clone();
let session_id = session_id.clone();
let last_event_id = last_event_id.clone();
let custom_headers = custom_headers.clone();
async move {
self.http_client
.get_stream_with_max_sse_event_size(
uri,
session_id,
last_event_id,
token,
custom_headers,
max_sse_event_size,
)
.await
}
})
.await
}
async fn post_message(
@ -79,18 +155,24 @@ where
uri: std::sync::Arc<str>,
message: crate::model::ClientJsonRpcMessage,
session_id: Option<std::sync::Arc<str>>,
mut auth_token: Option<String>,
auth_token: Option<String>,
custom_headers: HashMap<HeaderName, HeaderValue>,
) -> Result<
crate::transport::streamable_http_client::StreamableHttpPostResponse,
StreamableHttpError<Self::Error>,
> {
if auth_token.is_none() {
auth_token = Some(self.get_access_token().await?);
}
self.http_client
.post_message(uri, message, session_id, auth_token, custom_headers)
.await
self.call_reacting_to_challenges(auth_token, |token| {
let uri = uri.clone();
let message = message.clone();
let session_id = session_id.clone();
let custom_headers = custom_headers.clone();
async move {
self.http_client
.post_message(uri, message, session_id, token, custom_headers)
.await
}
})
.await
}
async fn post_message_with_max_sse_event_size(
@ -98,25 +180,31 @@ where
uri: std::sync::Arc<str>,
message: crate::model::ClientJsonRpcMessage,
session_id: Option<std::sync::Arc<str>>,
mut auth_token: Option<String>,
auth_token: Option<String>,
custom_headers: HashMap<HeaderName, HeaderValue>,
max_sse_event_size: usize,
) -> Result<
crate::transport::streamable_http_client::StreamableHttpPostResponse,
StreamableHttpError<Self::Error>,
> {
if auth_token.is_none() {
auth_token = Some(self.get_access_token().await?);
}
self.http_client
.post_message_with_max_sse_event_size(
uri,
message,
session_id,
auth_token,
custom_headers,
max_sse_event_size,
)
.await
self.call_reacting_to_challenges(auth_token, |token| {
let uri = uri.clone();
let message = message.clone();
let session_id = session_id.clone();
let custom_headers = custom_headers.clone();
async move {
self.http_client
.post_message_with_max_sse_event_size(
uri,
message,
session_id,
token,
custom_headers,
max_sse_event_size,
)
.await
}
})
.await
}
}

View file

@ -124,7 +124,8 @@ fn negotiate_version_headers(
(version, headers)
}
#[derive(Debug)]
#[derive(Debug, Error)]
#[error("authorization required: {www_authenticate_header}")]
#[non_exhaustive]
pub struct AuthRequiredError {
pub www_authenticate_header: String,
@ -139,7 +140,8 @@ impl AuthRequiredError {
}
}
#[derive(Debug)]
#[derive(Debug, Error)]
#[error("insufficient scope: {www_authenticate_header}")]
#[non_exhaustive]
pub struct InsufficientScopeError {
pub www_authenticate_header: String,
@ -197,15 +199,31 @@ pub enum StreamableHttpError<E: std::error::Error + Send + Sync + 'static> {
#[error("Auth error: {0}")]
Auth(#[from] crate::transport::auth::AuthError),
#[error("Auth required")]
AuthRequired(AuthRequiredError),
AuthRequired(#[source] AuthRequiredError),
#[error("Insufficient scope")]
InsufficientScope(InsufficientScopeError),
InsufficientScope(#[source] InsufficientScopeError),
#[error("Header name '{0}' is reserved and conflicts with default headers")]
ReservedHeaderConflict(String),
#[error("Session expired (HTTP 404)")]
SessionExpired,
}
impl<E: std::error::Error + Send + Sync + 'static> StreamableHttpError<E> {
/// The `WWW-Authenticate` challenge carried by this error, when the
/// server answered 401 ([`AuthRequired`](Self::AuthRequired)) or 403
/// ([`InsufficientScope`](Self::InsufficientScope)). Feed it to
/// [`AuthorizationRequest::with_challenge`](crate::transport::auth::AuthorizationRequest::with_challenge)
/// to authorize reactively.
#[cfg(feature = "auth")]
pub fn auth_challenge(&self) -> Option<&str> {
match self {
Self::AuthRequired(error) => Some(&error.www_authenticate_header),
Self::InsufficientScope(error) => Some(&error.www_authenticate_header),
_ => None,
}
}
}
#[derive(Debug, Clone, Error)]
#[non_exhaustive]
pub enum StreamableHttpProtocolError {

View file

@ -8,8 +8,9 @@ use axum::{
routing::get,
};
use rmcp::{
ServiceExt,
RoleClient, ServiceExt,
model::ClientInfo,
service::RunningService,
transport::{
StreamableHttpClientTransport,
auth::{AuthClient, AuthorizationRequest, OAuthState},
@ -55,6 +56,109 @@ async fn callback_handler(
Html(CALLBACK_HTML.to_string())
}
enum ConnectOutcome {
/// The server accepted the unauthenticated connection.
Connected(RunningService<RoleClient, ClientInfo>),
/// The server answered 401; authorize with this `WWW-Authenticate`
/// challenge and reconnect.
AuthRequired(String),
}
/// Attempt the real connection unauthenticated — the reactive discovery
/// trigger (matching the TypeScript and Python SDKs). The server's 401
/// challenge, not a probe, tells us whether and how to authorize.
async fn try_connect(http_client: reqwest::Client, server_url: &str) -> Result<ConnectOutcome> {
let transport = StreamableHttpClientTransport::with_client(
http_client,
StreamableHttpClientTransportConfig::with_uri(server_url),
);
match ClientInfo::default().serve(transport).await {
Ok(client) => Ok(ConnectOutcome::Connected(client)),
Err(error) => match error.auth_challenge() {
Some(challenge) => Ok(ConnectOutcome::AuthRequired(challenge.to_string())),
None => Err(error.into()),
},
}
}
/// Run the browser OAuth flow seeded by the server's challenge, then
/// reconnect with the authorized transport.
async fn authorize_and_connect(
challenge: String,
oauth_http_client: reqwest::Client,
server_url: &str,
client_metadata_url: &str,
code_receiver: oneshot::Receiver<CallbackParams>,
output: &mut BufWriter<tokio::io::Stdout>,
) -> Result<RunningService<RoleClient, ClientInfo>> {
tracing::info!("Server requires authorization: {challenge}");
// initialize oauth state machine
let mut oauth_state = OAuthState::new(server_url, Some(oauth_http_client))
.await
.context("Failed to initialize oauth state machine")?;
// Seed discovery from the server's challenge, and use CIMD (SEP-991)
// with client metadata URL. Passing no scopes lets the SDK auto-select
// from the challenge's scope hint, Protected Resource Metadata, or AS
// metadata.
oauth_state
.start_authorization(
AuthorizationRequest::new(MCP_REDIRECT_URI)
.with_client_name("Test MCP Client")
.with_client_metadata_url(client_metadata_url)
.with_challenge(challenge),
)
.await
.context("Failed to start authorization")?;
// Output authorization URL to user
output
.write_all(b"Please open the following URL in your browser to authorize:\n\n")
.await?;
output
.write_all(oauth_state.get_authorization_url().await?.as_bytes())
.await?;
output
.write_all(b"\n\nWaiting for browser callback, please do not close this window...\n")
.await?;
output.flush().await?;
// Wait for authorization code
tracing::info!("Waiting for authorization code...");
let CallbackParams {
code: auth_code,
state: csrf_token,
iss,
} = code_receiver
.await
.context("Failed to get authorization code")?;
tracing::info!("Received authorization code: {}", auth_code);
// Exchange code for access token
tracing::info!("Exchanging authorization code for access token...");
oauth_state
.handle_callback_with_issuer(&auth_code, &csrf_token, iss.as_deref())
.await
.context("Failed to handle callback")?;
tracing::info!("Successfully obtained access token");
output
.write_all(b"\nAuthorization successful! Access token obtained.\n\n")
.await?;
output.flush().await?;
// Reconnect with the authorized transport
tracing::info!("Establishing authorized connection to MCP server...");
let am = oauth_state
.into_authorization_manager()
.ok_or_else(|| anyhow::anyhow!("Failed to get authorization manager"))?;
let auth_client = AuthClient::new(reqwest::Client::default(), am);
let transport = StreamableHttpClientTransport::with_client(
auth_client,
StreamableHttpClientTransportConfig::with_uri(server_url),
);
Ok(ClientInfo::default().serve(transport).await?)
}
#[tokio::main]
async fn main() -> Result<()> {
// Initialize logging
@ -123,73 +227,32 @@ async fn main() -> Result<()> {
.build()
.context("Failed to build OAuth HTTP client")?;
// initialize oauth state machine
let mut oauth_state = OAuthState::new(&server_url, Some(oauth_http_client))
.await
.context("Failed to initialize oauth state machine")?;
// use CIMD (SEP-991) with client metadata URL.
// passing no scopes lets the SDK auto-select from the server's
// WWW-Authenticate header, Protected Resource Metadata, or AS metadata.
oauth_state
.start_authorization(
AuthorizationRequest::new(MCP_REDIRECT_URI)
.with_client_name("Test MCP Client")
.with_client_metadata_url(&client_metadata_url),
)
.await
.context("Failed to start authorization")?;
// Output authorization URL to user
let mut output = BufWriter::new(tokio::io::stdout());
output.write_all(b"\n=== MCP OAuth Client ===\n\n").await?;
output
.write_all(b"Please open the following URL in your browser to authorize:\n\n")
.await?;
output
.write_all(oauth_state.get_authorization_url().await?.as_bytes())
.await?;
output
.write_all(b"\n\nWaiting for browser callback, please do not close this window...\n")
.await?;
output.flush().await?;
// Wait for authorization code
tracing::info!("Waiting for authorization code...");
let CallbackParams {
code: auth_code,
state: csrf_token,
iss,
} = code_receiver
.await
.context("Failed to get authorization code")?;
tracing::info!("Received authorization code: {}", auth_code);
// Exchange code for access token
tracing::info!("Exchanging authorization code for access token...");
oauth_state
.handle_callback_with_issuer(&auth_code, &csrf_token, iss.as_deref())
.await
.context("Failed to handle callback")?;
tracing::info!("Successfully obtained access token");
output
.write_all(b"\nAuthorization successful! Access token obtained.\n\n")
.await?;
output.flush().await?;
// Create authorized transport, this transport is authorized by the oauth state machine
tracing::info!("Establishing authorized connection to MCP server...");
let am = oauth_state
.into_authorization_manager()
.ok_or_else(|| anyhow::anyhow!("Failed to get authorization manager"))?;
let client = AuthClient::new(reqwest::Client::default(), am);
let transport = StreamableHttpClientTransport::with_client(
client,
StreamableHttpClientTransportConfig::with_uri(server_url.as_str()),
);
// Create client and connect to MCP server
let client_service = ClientInfo::default();
let client = client_service.serve(transport).await?;
// Reactive discovery: attempt the real connection first. The server's
// 401 challenge — not a probe — tells us whether and how to authorize.
// The transport gets a default client: `oauth_http_client`'s request
// timeout would cut long-lived SSE streams short.
tracing::info!("Attempting connection to MCP server...");
let client = match try_connect(reqwest::Client::default(), &server_url).await? {
ConnectOutcome::Connected(client) => {
tracing::info!("Server accepted the connection without authorization");
client
}
ConnectOutcome::AuthRequired(challenge) => {
authorize_and_connect(
challenge,
oauth_http_client,
&server_url,
&client_metadata_url,
code_receiver,
&mut output,
)
.await?
}
};
tracing::info!("Successfully connected to MCP server");
// Test API requests