From 5216ab2bbdb2e61e7c1a822322aa9bb0b62a1035 Mon Sep 17 00:00:00 2001
From: Jack Amadeo <jackamadeo@block.xyz>
Date: Tue, 23 Sep 2025 15:58:00 -0400
Subject: [PATCH] fix: return auth errors (#451)

---
 .../common/reqwest/streamable_http_client.rs  | 21 ++++++++++--
 .../src/transport/streamable_http_client.rs   | 32 ++++++++++++++-----
 2 files changed, 43 insertions(+), 10 deletions(-)

diff --git a/crates/rmcp/src/transport/common/reqwest/streamable_http_client.rs b/crates/rmcp/src/transport/common/reqwest/streamable_http_client.rs
index a3f963e..bea1a71 100644
--- a/crates/rmcp/src/transport/common/reqwest/streamable_http_client.rs
+++ b/crates/rmcp/src/transport/common/reqwest/streamable_http_client.rs
@@ -1,6 +1,7 @@
-use std::sync::Arc;
+use std::{borrow::Cow, sync::Arc};
 
 use futures::{StreamExt, stream::BoxStream};
+use http::header::WWW_AUTHENTICATE;
 use reqwest::header::ACCEPT;
 use sse_stream::{Sse, SseStream};
 
@@ -101,7 +102,23 @@ impl StreamableHttpClient for reqwest::Client {
         if let Some(session_id) = session_id {
             request = request.header(HEADER_SESSION_ID, session_id.as_ref());
         }
-        let response = request.json(&message).send().await?.error_for_status()?;
+        let response = request.json(&message).send().await?;
+        if response.status() == reqwest::StatusCode::UNAUTHORIZED {
+            if let Some(header) = response.headers().get(WWW_AUTHENTICATE) {
+                let header = header
+                    .to_str()
+                    .map_err(|_| {
+                        StreamableHttpError::UnexpectedServerResponse(Cow::from(
+                            "invalid www-authenticate header value",
+                        ))
+                    })?
+                    .to_string();
+                return Err(StreamableHttpError::AuthRequired(AuthRequiredError {
+                    www_authenticate_header: header,
+                }));
+            }
+        }
+        let response = response.error_for_status()?;
         if response.status() == reqwest::StatusCode::ACCEPTED {
             return Ok(StreamableHttpPostResponse::Accepted);
         }
diff --git a/crates/rmcp/src/transport/streamable_http_client.rs b/crates/rmcp/src/transport/streamable_http_client.rs
index e25e57d..ad451b1 100644
--- a/crates/rmcp/src/transport/streamable_http_client.rs
+++ b/crates/rmcp/src/transport/streamable_http_client.rs
@@ -18,6 +18,11 @@ use crate::{
 
 type BoxedSseStream = BoxStream<'static, Result<Sse, SseError>>;
 
+#[derive(Debug)]
+pub struct AuthRequiredError {
+    pub www_authenticate_header: String,
+}
+
 #[derive(Error, Debug)]
 pub enum StreamableHttpError<E: std::error::Error + Send + Sync + 'static> {
     #[error("SSE error: {0}")]
@@ -48,6 +53,8 @@ pub enum StreamableHttpError<E: std::error::Error + Send + Sync + 'static> {
     #[cfg_attr(docsrs, doc(cfg(feature = "auth")))]
     #[error("Auth error: {0}")]
     Auth(#[from] crate::transport::auth::AuthError),
+    #[error("Auth required")]
+    AuthRequired(AuthRequiredError),
 }
 
 #[derive(Debug, Clone, Error)]
@@ -274,8 +281,7 @@ impl<C: StreamableHttpClient> Worker for StreamableHttpClientWorker<C> {
             responder,
             message: initialize_request,
         } = context.recv_from_handler().await?;
-        let _ = responder.send(Ok(()));
-        let (message, session_id) = self
+        let (message, session_id) = match self
             .client
             .post_message(
                 config.uri.clone(),
@@ -284,12 +290,22 @@ impl<C: StreamableHttpClient> Worker for StreamableHttpClientWorker<C> {
                 self.config.auth_header,
             )
             .await
-            .map_err(WorkerQuitReason::fatal_context("send initialize request"))?
-            .expect_initialized::<C::Error>()
-            .await
-            .map_err(WorkerQuitReason::fatal_context(
-                "process initialize response",
-            ))?;
+        {
+            Ok(res) => {
+                let _ = responder.send(Ok(()));
+                res.expect_initialized::<C::Error>().await.map_err(
+                    WorkerQuitReason::fatal_context("process initialize response"),
+                )?
+            }
+            Err(err) => {
+                let msg = format!("{:?}", err);
+                let _ = responder.send(Err(err));
+                return Err(WorkerQuitReason::fatal(
+                    StreamableHttpError::TransportChannelClosed,
+                    msg,
+                ));
+            }
+        };
         let session_id: Option<Arc<str>> = if let Some(session_id) = session_id {
             Some(session_id.into())
         } else {