fix(authentication): Allow connecting to Oauth servers that use protected-resource fallback instead of the WWW-authenticate header (#8148)

Signed-off-by: Cameron Yick <cameron.yick@datadoghq.com>
Signed-off-by: Douwe Osinga <douwe@squareup.com>
Co-authored-by: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com>
Co-authored-by: Douwe Osinga <douwe@squareup.com>

crates/goose/src/agents/extension_manager.rs

@@ -6,7 +6,7 @@ use futures::{future, FutureExt};
 use once_cell::sync::Lazy;
 use rmcp::service::{ClientInitializeError, ServiceError};
 use rmcp::transport::streamable_http_client::{
-    AuthRequiredError, StreamableHttpClientTransportConfig, StreamableHttpError,
+    StreamableHttpClientTransportConfig, StreamableHttpError,
 };
 use rmcp::transport::{
     ConfigureCommandExt, DynamicTransportError, StreamableHttpClientTransport, TokioChildProcess,
@@ -296,25 +296,26 @@ async fn child_process_client(
     }
 }
 
-fn extract_auth_error(
+/// Retry with OAuth for typed auth challenges and wrapped bare HTTP 401 responses.
-    res: &Result<McpClient, ClientInitializeError>,
+fn should_attempt_oauth_fallback(res: &Result<McpClient, ClientInitializeError>) -> bool {
-) -> Option<&AuthRequiredError> {
+    let Err(ClientInitializeError::TransportError {
-    match res {
+        error: DynamicTransportError { error, .. },
-        Ok(_) => None,
+        ..
-        Err(err) => match err {
+    }) = res
-            ClientInitializeError::TransportError {
+    else {
-                error: DynamicTransportError { error, .. },
+        return false;
-                ..
+    };
-            } => error
+
-                .downcast_ref::<StreamableHttpError<reqwest::Error>>()
+    if let Some(http_err) = error.downcast_ref::<StreamableHttpError<reqwest::Error>>() {
-                .and_then(|auth_error| match auth_error {
+        match http_err {
-                    StreamableHttpError::AuthRequired(auth_required_error) => {
+            StreamableHttpError::AuthRequired(_) => true,
-                        Some(auth_required_error)
+            StreamableHttpError::UnexpectedServerResponse(body) => body.starts_with("HTTP 401"),
-                    }
+            _ => false,
-                    _ => None,
+        }
-                }),
+    } else {
-            _ => None,
+        error
-        },
+            .to_string()
+            .contains("unexpected server response: HTTP 401")
     }
 }
 
@@ -453,37 +454,39 @@ async fn create_streamable_http_client(
     )
     .await;
 
-    if extract_auth_error(&client_res).is_some() {
+    if should_attempt_oauth_fallback(&client_res) {
-        let auth_manager = oauth_flow(&uri.to_string(), &name.to_string())
+        match oauth_flow(&uri.to_string(), &name.to_string()).await {
-            .await
+            Ok(auth_manager) => {
-            .map_err(|_| ExtensionError::SetupError("auth error".to_string()))?;
+                let mut auth_headers = HeaderMap::new();
-        let mut auth_headers = HeaderMap::new();
+                auth_headers.insert(reqwest::header::USER_AGENT, GOOSE_USER_AGENT);
-        auth_headers.insert(reqwest::header::USER_AGENT, GOOSE_USER_AGENT);
+                let auth_http_client = reqwest::Client::builder()
-        let auth_http_client = reqwest::Client::builder()
+                    .default_headers(auth_headers)
-            .default_headers(auth_headers)
+                    .build()
-            .build()
+                    .map_err(|_| {
-            .map_err(|_| {
+                        ExtensionError::ConfigError("could not construct http client".to_string())
-                ExtensionError::ConfigError("could not construct http client".to_string())
+                    })?;
-            })?;
+                let auth_client = AuthClient::new(auth_http_client, auth_manager);
-        let auth_client = AuthClient::new(auth_http_client, auth_manager);
+                let transport = StreamableHttpClientTransport::with_client(
-        let transport = StreamableHttpClientTransport::with_client(
+                    auth_client,
-            auth_client,
+                    StreamableHttpClientTransportConfig {
-            StreamableHttpClientTransportConfig {
+                        uri: uri.into(),
-                uri: uri.into(),
+                        ..Default::default()
-                ..Default::default()
+                    },
-            },
+                );
-        );
+                Ok(Box::new(
-        Ok(Box::new(
+                    McpClient::connect(
-            McpClient::connect(
+                        transport,
-                transport,
+                        timeout_duration,
-                timeout_duration,
+                        provider,
-                provider,
+                        client_name,
-                client_name,
+                        capabilities,
-                capabilities,
+                        roots_dir.to_path_buf(),
-                roots_dir.to_path_buf(),
+                    )
-            )
+                    .await?,
-            .await?,
+                ))
-        ))
+            }
+            Err(_) => Ok(Box::new(client_res?)),
+        }
     } else {
         Ok(Box::new(client_res?))
     }
@@ -2305,4 +2308,43 @@ mod tests {
             "old extension must be preserved when replacement client creation fails"
         );
     }
+
+    fn transport_err(error: Box<dyn std::error::Error + Send + Sync>) -> ClientInitializeError {
+        ClientInitializeError::TransportError {
+            error: rmcp::transport::DynamicTransportError {
+                transport_name: "test".into(),
+                transport_type_id: std::any::TypeId::of::<()>(),
+                error,
+            },
+            context: "test context".into(),
+        }
+    }
+
+    fn streamable_err(
+        e: rmcp::transport::streamable_http_client::StreamableHttpError<reqwest::Error>,
+    ) -> ClientInitializeError {
+        transport_err(Box::new(e))
+    }
+
+    #[test]
+    fn test_oauth_fallback_on_typed_auth_required() {
+        let err = streamable_err(
+            rmcp::transport::streamable_http_client::StreamableHttpError::AuthRequired(
+                rmcp::transport::streamable_http_client::AuthRequiredError {
+                    www_authenticate_header: "Bearer realm=\"test\"".to_string(),
+                },
+            ),
+        );
+        assert!(should_attempt_oauth_fallback(&Err(err)));
+    }
+
+    #[test]
+    fn test_oauth_fallback_on_unexpected_response_http_401_prefix() {
+        let err = streamable_err(
+            rmcp::transport::streamable_http_client::StreamableHttpError::UnexpectedServerResponse(
+                std::borrow::Cow::Borrowed("HTTP 401 Unauthorized"),
+            ),
+        );
+        assert!(should_attempt_oauth_fallback(&Err(err)));
+    }
 }