fix SSEOption (#63)

* fix SSEOption

* fix SSEOption

* fix SSEOption

* fix SSEOption

* fix SSEOption

---------

Co-authored-by: Buf Generate <[email protected]>

Author: a67793581

server/sse.go

@@ -6,6 +6,7 @@ import (
 	"fmt"
 	"net/http"
 	"net/http/httptest"
+	"net/url"
 	"strings"
 	"sync"
 
@@ -57,7 +58,23 @@ type SSEOption func(*SSEServer)
 // WithBaseURL sets the base URL for the SSE server
 func WithBaseURL(baseURL string) SSEOption {
 	return func(s *SSEServer) {
-		s.baseURL = baseURL
+		if baseURL != "" {
+			u, err := url.Parse(baseURL)
+			if err != nil {
+				return
+			}
+			if u.Scheme != "http" && u.Scheme != "https" {
+				return
+			}
+			// Check if the host is empty or only contains a port
+			if u.Host == "" || strings.HasPrefix(u.Host, ":") {
+				return
+			}
+			if len(u.Query()) > 0 {
+				return
+			}
+		}
+		s.baseURL = strings.TrimSuffix(baseURL, "/")
 	}
 }
 
@@ -69,7 +86,6 @@ func WithBasePath(basePath string) SSEOption {
 			basePath = "/" + basePath
 		}
 		s.basePath = strings.TrimSuffix(basePath, "/")
-		s.baseURL = s.baseURL + s.basePath
 	}
 }
 
@@ -108,7 +124,6 @@ func NewSSEServer(server *MCPServer, opts ...SSEOption) *SSEServer {
 		server:          server,
 		sseEndpoint:     "/sse",
 		messageEndpoint: "/message",
-		basePath:        "",
 	}
 
 	// Apply all options
@@ -219,12 +234,7 @@ func (s *SSEServer) handleSSE(w http.ResponseWriter, r *http.Request) {
 		}
 	}()
 
-	messageEndpoint := fmt.Sprintf(
-		"%s%s?sessionId=%s",
-		s.baseURL,
-		s.messageEndpoint,
-		sessionID,
-	)
+	messageEndpoint := fmt.Sprintf("%s?sessionId=%s", s.CompleteMessageEndpoint(), sessionID)
 
 	// Send the initial endpoint event
 	fmt.Fprintf(w, "event: endpoint\ndata: %s\r\n\r\n", messageEndpoint)
@@ -345,22 +355,47 @@ func (s *SSEServer) SendEventToSession(
 		return fmt.Errorf("event queue full")
 	}
 }
+func (s *SSEServer) GetUrlPath(input string) (string, error) {
+	parse, err := url.Parse(input)
+	if err != nil {
+		return "", fmt.Errorf("failed to parse URL %s: %w", input, err)
+	}
+	return parse.Path, nil
+}
+
+func (s *SSEServer) CompleteSseEndpoint() string {
+	return s.baseURL + s.basePath + s.sseEndpoint
+}
+func (s *SSEServer) CompleteSsePath() string {
+	path, err := s.GetUrlPath(s.CompleteSseEndpoint())
+	if err != nil {
+		return s.basePath + s.sseEndpoint
+	}
+	return path
+}
+
+func (s *SSEServer) CompleteMessageEndpoint() string {
+	return s.baseURL + s.basePath + s.messageEndpoint
+}
+func (s *SSEServer) CompleteMessagePath() string {
+	path, err := s.GetUrlPath(s.CompleteMessageEndpoint())
+	if err != nil {
+		return s.basePath + s.messageEndpoint
+	}
+	return path
+}
 
 // ServeHTTP implements the http.Handler interface.
 func (s *SSEServer) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	path := r.URL.Path
-
-	// Construct the full SSE and message paths
-	ssePath := s.basePath + s.sseEndpoint
-	messagePath := s.basePath + s.messageEndpoint
-
 	// Use exact path matching rather than Contains
-	if path == ssePath {
+	ssePath := s.CompleteSsePath()
+	if ssePath != "" && path == ssePath {
 		s.handleSSE(w, r)
 		return
 	}
-
-	if path == messagePath {
+	messagePath := s.CompleteMessagePath()
+	if messagePath != "" && path == messagePath {
 		s.handleMessage(w, r)
 		return
 	}

server/sse_test.go

@@ -5,6 +5,7 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
+	"math/rand"
 	"net/http"
 	"net/http/httptest"
 	"strings"
@@ -29,9 +30,9 @@ func TestSSEServer(t *testing.T) {
 		if sseServer.server == nil {
 			t.Error("MCPServer should not be nil")
 		}
-		if sseServer.baseURL != "http://localhost:8080/mcp" {
+		if sseServer.baseURL != "http://localhost:8080" {
 			t.Errorf(
-				"Expected baseURL http://localhost:8080/mcp, got %s",
+				"Expected baseURL http://localhost:8080, got %s",
 				sseServer.baseURL,
 			)
 		}
@@ -350,7 +351,7 @@ func TestSSEServer(t *testing.T) {
 		sseServer := NewSSEServer(mcpServer)
 
 		mux := http.NewServeMux()
-		mux.Handle("/mcp/", http.StripPrefix("/mcp", sseServer))
+		mux.Handle("/mcp/", sseServer)
 
 		ts := httptest.NewServer(mux)
 		defer ts.Close()
@@ -614,4 +615,48 @@ func TestSSEServer(t *testing.T) {
 		}
 	})
 
+	t.Run("SSEOption should not have negative effects when used repeatedly but should always remain consistent.", func(t *testing.T) {
+		mcpServer := NewMCPServer("test", "1.0.0")
+		basePath := "/mcp-test"
+		baseURL := "http://localhost:8080/test"
+		messageEndpoint := "/message-test"
+		sseEndpoint := "/sse-test"
+		srv := &http.Server{}
+		rands := []SSEOption{
+			WithBasePath(basePath),
+			WithBaseURL(baseURL),
+			WithMessageEndpoint(messageEndpoint),
+			WithSSEEndpoint(sseEndpoint),
+			WithHTTPServer(srv),
+		}
+		for i := 0; i < 100; i++ {
+
+			var options []SSEOption
+			for i2 := 0; i2 < 100; i2++ {
+				index := rand.Intn(len(rands))
+				options = append(options, rands[index])
+			}
+			sseServer := NewSSEServer(mcpServer, options...)
+
+			if sseServer.basePath != basePath {
+				t.Fatalf("basePath %v, got: %v", basePath, sseServer.basePath)
+			}
+
+			if sseServer.baseURL != baseURL {
+				t.Fatalf("baseURL %v, got: %v", baseURL, sseServer.baseURL)
+			}
+
+			if sseServer.sseEndpoint != sseEndpoint {
+				t.Fatalf("sseEndpoint %v, got: %v", sseEndpoint, sseServer.sseEndpoint)
+			}
+
+			if sseServer.messageEndpoint != messageEndpoint {
+				t.Fatalf("messageEndpoint  %v, got: %v", messageEndpoint, sseServer.messageEndpoint)
+			}
+
+			if sseServer.srv != srv {
+				t.Fatalf("srv  %v, got: %v", srv, sseServer.srv)
+			}
+		}
+	})
 }