diff --git a/backend/windmill-mcp/src/client/redacting.rs b/backend/windmill-mcp/src/client/redacting.rs index a14f862be5..0024ed6162 100644 --- a/backend/windmill-mcp/src/client/redacting.rs +++ b/backend/windmill-mcp/src/client/redacting.rs @@ -1,4 +1,4 @@ -use std::{collections::HashMap, future::Future, sync::Arc}; +use std::{borrow::Cow, collections::HashMap, future::Future, sync::Arc}; use futures::stream::BoxStream; use reqwest::header::{HeaderName, HeaderValue}; @@ -20,35 +20,19 @@ type Error = StreamableHttpError; #[derive(Clone)] pub(super) struct UrlRedactingClient(pub(super) reqwest::Client); -async fn redacted( - uri: Arc, - request: impl Future>, -) -> Result { +async fn redacted(request: impl Future>) -> Result { match request.with_subscriber(Dispatch::none()).await { Err(StreamableHttpError::Client(e)) => Err(StreamableHttpError::Client(e.without_url())), - Err(StreamableHttpError::UnexpectedServerResponse(body)) => { - Err(StreamableHttpError::UnexpectedServerResponse( - without_request_target(&uri, &body).into(), - )) + // An owned message is `HTTP : `, and the body is the server's + // to write: a common one echoes the request target (`Cannot POST /mcp/`). + Err(StreamableHttpError::UnexpectedServerResponse(Cow::Owned(message))) => { + let status = message.split(": ").next().unwrap_or_default().to_string(); + Err(StreamableHttpError::UnexpectedServerResponse(status.into())) } r => r, } } -/// The body of an error response is the server's to write, and a common one -/// echoes the request target (`Cannot POST /mcp/`). -fn without_request_target(uri: &str, body: &str) -> String { - let mut body = body.replace(uri, ""); - if let Ok(url) = reqwest::Url::parse(uri) { - for part in [Some(url.path()).filter(|p| p.len() > 1), url.query()] { - if let Some(part) = part.filter(|p| !p.is_empty()) { - body = body.replace(part, ""); - } - } - } - body -} - impl StreamableHttpClient for UrlRedactingClient { type Error = reqwest::Error; @@ -61,7 +45,6 @@ impl StreamableHttpClient for UrlRedactingClient { custom_headers: HashMap, ) -> Result { redacted( - uri.clone(), self.0 .post_message(uri, message, session_id, auth_header, custom_headers), ) @@ -77,17 +60,14 @@ impl StreamableHttpClient for UrlRedactingClient { custom_headers: HashMap, max_sse_event_size: usize, ) -> Result { - redacted( - uri.clone(), - self.0.post_message_with_max_sse_event_size( - uri, - message, - session_id, - auth_header, - custom_headers, - max_sse_event_size, - ), - ) + redacted(self.0.post_message_with_max_sse_event_size( + uri, + message, + session_id, + auth_header, + custom_headers, + max_sse_event_size, + )) .await } @@ -99,7 +79,6 @@ impl StreamableHttpClient for UrlRedactingClient { custom_headers: HashMap, ) -> Result<(), Error> { redacted( - uri.clone(), self.0 .delete_session(uri, session_id, auth_header, custom_headers), ) @@ -115,7 +94,6 @@ impl StreamableHttpClient for UrlRedactingClient { custom_headers: HashMap, ) -> Result>, Error> { redacted( - uri.clone(), self.0 .get_stream(uri, session_id, last_event_id, auth_header, custom_headers), ) @@ -131,42 +109,14 @@ impl StreamableHttpClient for UrlRedactingClient { custom_headers: HashMap, max_sse_event_size: usize, ) -> Result>, Error> { - redacted( - uri.clone(), - self.0.get_stream_with_max_sse_event_size( - uri, - session_id, - last_event_id, - auth_header, - custom_headers, - max_sse_event_size, - ), - ) + redacted(self.0.get_stream_with_max_sse_event_size( + uri, + session_id, + last_event_id, + auth_header, + custom_headers, + max_sse_event_size, + )) .await } } - -#[cfg(test)] -mod tests { - use super::without_request_target; - - #[test] - fn error_body_loses_the_request_target() { - let uri = "https://mcp.example.com/mcp/KEY?token=T0K"; - for body in [ - "Cannot POST /mcp/KEY", - "no route for /mcp/KEY?token=T0K", - uri, - ] { - let scrubbed = without_request_target(uri, &format!("HTTP 404: {body}")); - assert!( - !scrubbed.contains("KEY") && !scrubbed.contains("T0K"), - "{scrubbed}" - ); - } - assert_eq!( - without_request_target("https://mcp.example.com/", "a/b"), - "a/b" - ); - } -}