refactor: make oauth discovery reactive

Author: DaleSeo

conformance/src/bin/client.rs

@@ -1,14 +1,14 @@
 use rmcp::{
     ClientHandler, ClientLifecycleMode, ClientServiceExt, ErrorData, RoleClient, ServiceExt,
     model::*,
-    service::RequestContext,
+    service::{ClientInitializeError, RequestContext},
     transport::{
         AuthClient, AuthorizationManager, StreamableHttpClientTransport,
         auth::{
             AuthorizationCallback, AuthorizationRequest, ClientCredentialsConfig,
             InMemoryCredentialStore, JwtSigningAlgorithm, OAuthState,
         },
-        streamable_http_client::StreamableHttpClientTransportConfig,
+        streamable_http_client::{StreamableHttpClientTransportConfig, StreamableHttpError},
     },
 };
 use serde_json::{Value, json};
@@ -195,26 +195,69 @@ 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. Returns `None` when the server accepts unauthenticated clients
+/// (the connection is closed cleanly in that case).
+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 auth_challenge_from_init_error(&error) {
+            Some(challenge) => Ok(Some(challenge)),
+            None => Err(error.into()),
+        },
+    }
+}
+
+fn auth_challenge_from_init_error(error: &ClientInitializeError) -> Option<String> {
+    let ClientInitializeError::TransportError { error, .. } = error else {
+        return None;
+    };
+    error
+        .error
+        .downcast_ref::<StreamableHttpError<reqwest::Error>>()?
+        .auth_challenge()
+        .map(str::to_string)
+}
+
+/// Perform the headless OAuth authorization-code flow, reactively:
 ///
-/// 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`
+/// 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>> {
+    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
-    oauth
-        .start_authorization(
-            AuthorizationRequest::new(REDIRECT_URI)
-                .with_client_name("conformance-client")
-                .with_client_metadata_url(CIMD_CLIENT_METADATA_URL),
-        )
-        .await?;
+    let mut request = AuthorizationRequest::new(REDIRECT_URI)
+        .with_client_name("conformance-client")
+        .with_client_metadata_url(CIMD_CLIENT_METADATA_URL);
+    if let Some(challenge) = challenge {
+        request = request.with_challenge(challenge);
+    }
+    oauth.start_authorization(request).await?;
 
     let auth_url = oauth.get_authorization_url().await?;
     tracing::debug!("Authorization URL: {}", auth_url);
@@ -255,15 +298,16 @@ 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?;
 
-    oauth
-        .start_authorization(
-            AuthorizationRequest::new(REDIRECT_URI)
-                .with_preregistered_client(client_id)
-                .with_client_secret(client_secret),
-        )
-        .await?;
+    let mut request = AuthorizationRequest::new(REDIRECT_URI)
+        .with_preregistered_client(client_id)
+        .with_client_secret(client_secret);
+    if let Some(challenge) = challenge {
+        request = request.with_challenge(challenge);
+    }
+    oauth.start_authorization(request).await?;
 
     let auth_url = oauth.get_authorization_url().await?;
     let callback = headless_authorize(&auth_url).await?;
@@ -325,15 +369,16 @@ 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?;
-    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),
-        )
-        .await?;
+    let mut 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);
+    if let Some(challenge) = challenge {
+        request = request.with_challenge(challenge);
+    }
+    oauth.start_authorization(request).await?;
 
     let auth_url = oauth.get_authorization_url().await?;
     let callback = headless_authorize(&auth_url).await?;
@@ -427,16 +472,17 @@ 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?;
-        oauth
-            .start_authorization(
-                AuthorizationRequest::new(REDIRECT_URI)
-                    .with_client_name("conformance-client")
-                    .with_client_metadata_url(CIMD_CLIENT_METADATA_URL),
-            )
-            .await?;
+        let mut request = AuthorizationRequest::new(REDIRECT_URI)
+            .with_client_name("conformance-client")
+            .with_client_metadata_url(CIMD_CLIENT_METADATA_URL);
+        if let Some(challenge) = &challenge {
+            request = request.with_challenge(challenge.clone());
+        }
+        oauth.start_authorization(request).await?;
         let auth_url = oauth.get_authorization_url().await?;
         let callback = headless_authorize(&auth_url).await?;
         oauth
@@ -509,7 +555,10 @@ async fn migration_token(
         return Ok(manager.get_access_token().await?);
     }
 
-    let resolution = manager.resolve_metadata().await?;
+    let resolution = match initialize_challenge(server_url, conformance_lifecycle()).await? {
+        Some(challenge) => manager.resolve_metadata_from_challenge(&challenge).await?,
+        None => manager.resolve_metadata().await?,
+    };
     manager.set_metadata(resolution.metadata);
     manager
         .register_client("conformance-client", REDIRECT_URI, &[])
@@ -609,7 +658,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 resolution = match initialize_challenge(server_url, conformance_lifecycle()).await? {
+        Some(challenge) => manager.resolve_metadata_from_challenge(&challenge).await?,
+        None => manager.resolve_metadata().await?,
+    };
     let token_endpoint = resolution.metadata.token_endpoint.clone();
     manager.set_metadata(resolution.metadata);