From 407cebf92114de5d3054096f6c1d1bdd69f507c7 Mon Sep 17 00:00:00 2001
From: Salman Mohammed <smohammed@squareup.com>
Date: Wed, 29 Jan 2025 16:48:04 -0500
Subject: [PATCH] fix: post endpoint url on sse endpoint event (#900)

---
 crates/mcp-client/examples/sse.rs      |  2 +-
 crates/mcp-client/src/transport/sse.rs | 12 +++++++-----
 2 files changed, 8 insertions(+), 6 deletions(-)

diff --git a/crates/mcp-client/examples/sse.rs b/crates/mcp-client/examples/sse.rs
index 8b93b372c1..360a2bbcd7 100644
--- a/crates/mcp-client/examples/sse.rs
+++ b/crates/mcp-client/examples/sse.rs
@@ -43,7 +43,7 @@ async fn main() -> Result<()> {
     println!("Connected to server: {server_info:?}\n");
 
     // Sleep for 100ms to allow the server to start - surprisingly this is required!
-    tokio::time::sleep(Duration::from_millis(100)).await;
+    tokio::time::sleep(Duration::from_millis(500)).await;
 
     // List tools
     let tools = client.list_tools(None).await?;
diff --git a/crates/mcp-client/src/transport/sse.rs b/crates/mcp-client/src/transport/sse.rs
index a5f05ed62e..e11325859d 100644
--- a/crates/mcp-client/src/transport/sse.rs
+++ b/crates/mcp-client/src/transport/sse.rs
@@ -9,6 +9,7 @@ use std::sync::Arc;
 use tokio::sync::{mpsc, RwLock};
 use tokio::time::{timeout, Duration};
 use tracing::warn;
+use url::Url;
 
 use super::{send_message, Transport, TransportHandle};
 
@@ -90,12 +91,13 @@ impl SseActor {
             match event {
                 SSE::Event(e) if e.event_type == "endpoint" => {
                     // SSE server uses the "endpoint" event to tell us the POST URL
-                    let base_url = sse_url.trim_end_matches('/').trim_end_matches("sse");
-                    let endpoint_path = e.data.trim_start_matches('/');
-                    let post_url = format!("{}{}", base_url, endpoint_path);
+                    let base_url = Url::parse(&sse_url).expect("Invalid base URL");
+                    let post_url = base_url
+                        .join(&e.data)
+                        .expect("Failed to resolve endpoint URL");
 
-                    println!("Discovered SSE POST endpoint: {post_url}");
-                    *post_endpoint.write().await = Some(post_url);
+                    println!("Discovered SSE POST endpoint: {}", post_url);
+                    *post_endpoint.write().await = Some(post_url.to_string());
                     break;
                 }
                 _ => continue,