Option mode (#39)

* Option mode

* fix test

Author: a67793581

.gitignore

@@ -1,2 +1,3 @@
 .aider*
 .env
+.idea

examples/everything/main.go

@@ -321,10 +321,6 @@ func handleSendNotification(
 	}, nil
 }
 
-func ServeSSE(mcpServer *server.MCPServer, addr string) *server.SSEServer {
-	return server.NewSSEServer(mcpServer, fmt.Sprintf("http://%s", addr))
-}
-
 func handleLongRunningOperationTool(
 	ctx context.Context,
 	request mcp.CallToolRequest,
@@ -418,19 +414,14 @@ func handleNotification(
 func main() {
 	var transport string
 	flag.StringVar(&transport, "t", "stdio", "Transport type (stdio or sse)")
-	flag.StringVar(
-		&transport,
-		"transport",
-		"stdio",
-		"Transport type (stdio or sse)",
-	)
+	flag.StringVar(&transport, "transport", "stdio", "Transport type (stdio or sse)")
 	flag.Parse()
 
 	mcpServer := NewMCPServer()
 
 	// Only check for "sse" since stdio is the default
 	if transport == "sse" {
-		sseServer := ServeSSE(mcpServer, "localhost:8080")
+		sseServer := server.NewSSEServer(mcpServer)
 		log.Printf("SSE server listening on :8080")
 		if err := sseServer.Start(":8080"); err != nil {
 			log.Fatalf("Server error: %v", err)

server/sse.go

@@ -12,15 +12,6 @@ import (
 	"github.com/mark3labs/mcp-go/mcp"
 )
 
-// SSEServer implements a Server-Sent Events (SSE) based MCP server.
-// It provides real-time communication capabilities over HTTP using the SSE protocol.
-type SSEServer struct {
-	server   *MCPServer
-	baseURL  string
-	sessions sync.Map
-	srv      *http.Server
-}
-
 // sseSession represents an active SSE connection.
 type sseSession struct {
 	writer     http.ResponseWriter
@@ -29,19 +20,67 @@ type sseSession struct {
 	eventQueue chan string // Channel for queuing events
 }
 
-// NewSSEServer creates a new SSE server instance with the given MCP server and base URL.
-func NewSSEServer(server *MCPServer, baseURL string) *SSEServer {
-	return &SSEServer{
-		server:  server,
-		baseURL: baseURL,
+// SSEServer implements a Server-Sent Events (SSE) based MCP server.
+// It provides real-time communication capabilities over HTTP using the SSE protocol.
+type SSEServer struct {
+	server          *MCPServer
+	baseURL         string
+	messageEndpoint string
+	sseEndpoint     string
+	sessions        sync.Map
+	srv             *http.Server
+}
+
+// Option defines a function type for configuring SSEServer
+type Option func(*SSEServer)
+
+// WithBaseURL sets the base URL for the SSE server
+func WithBaseURL(baseURL string) Option {
+	return func(s *SSEServer) {
+		s.baseURL = baseURL
+	}
+}
+
+// WithMessageEndpoint sets the message endpoint path
+func WithMessageEndpoint(endpoint string) Option {
+	return func(s *SSEServer) {
+		s.messageEndpoint = endpoint
 	}
 }
 
+// WithSSEEndpoint sets the SSE endpoint path
+func WithSSEEndpoint(endpoint string) Option {
+	return func(s *SSEServer) {
+		s.sseEndpoint = endpoint
+	}
+}
+
+// WithHTTPServer sets the HTTP server instance
+func WithHTTPServer(srv *http.Server) Option {
+	return func(s *SSEServer) {
+		s.srv = srv
+	}
+}
+
+// NewSSEServer creates a new SSE server instance with the given MCP server and options.
+func NewSSEServer(server *MCPServer, opts ...Option) *SSEServer {
+	s := &SSEServer{
+		server:          server,
+		sseEndpoint:     "/sse",
+		messageEndpoint: "/message",
+	}
+
+	// Apply all options
+	for _, opt := range opts {
+		opt(s)
+	}
+
+	return s
+}
+
 // NewTestServer creates a test server for testing purposes
 func NewTestServer(server *MCPServer) *httptest.Server {
-	sseServer := &SSEServer{
-		server: server,
-	}
+	sseServer := NewSSEServer(server)
 
 	testServer := httptest.NewServer(sseServer)
 	sseServer.baseURL = testServer.URL
@@ -132,8 +171,9 @@ func (s *SSEServer) handleSSE(w http.ResponseWriter, r *http.Request) {
 	}()
 
 	messageEndpoint := fmt.Sprintf(
-		"%s/message?sessionId=%s",
+		"%s%s?sessionId=%s",
 		s.baseURL,
+		s.messageEndpoint,
 		sessionID,
 	)
 
@@ -260,9 +300,9 @@ func (s *SSEServer) SendEventToSession(
 // ServeHTTP implements the http.Handler interface.
 func (s *SSEServer) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	switch r.URL.Path {
-	case "/sse":
+	case s.sseEndpoint:
 		s.handleSSE(w, r)
-	case "/message":
+	case s.messageEndpoint:
 		s.handleMessage(w, r)
 	default:
 		http.NotFound(w, r)

server/sse_test.go

@@ -16,7 +16,7 @@ import (
 func TestSSEServer(t *testing.T) {
 	t.Run("Can instantiate", func(t *testing.T) {
 		mcpServer := NewMCPServer("test", "1.0.0")
-		sseServer := NewSSEServer(mcpServer, "http://localhost:8080")
+		sseServer := NewSSEServer(mcpServer, WithBaseURL("http://localhost:8080"))
 
 		if sseServer == nil {
 			t.Error("SSEServer should not be nil")
@@ -234,7 +234,7 @@ func TestSSEServer(t *testing.T) {
 
 	t.Run("Can be used as http.Handler", func(t *testing.T) {
 		mcpServer := NewMCPServer("test", "1.0.0")
-		sseServer := NewSSEServer(mcpServer, "http://localhost:8080")
+		sseServer := NewSSEServer(mcpServer, WithBaseURL("http://localhost:8080"))
 
 		ts := httptest.NewServer(sseServer)
 		defer ts.Close()
@@ -263,7 +263,7 @@ func TestSSEServer(t *testing.T) {
 			t.Fatalf("Failed to connect to SSE endpoint: %v", err)
 		}
 		defer resp.Body.Close()
-		
+
 		if resp.StatusCode != http.StatusOK {
 			t.Errorf("Expected status 200, got %d", resp.StatusCode)
 		}
@@ -287,7 +287,7 @@ func TestSSEServer(t *testing.T) {
 
 	t.Run("Works with middleware", func(t *testing.T) {
 		mcpServer := NewMCPServer("test", "1.0.0")
-		sseServer := NewSSEServer(mcpServer, "http://localhost:8080")
+		sseServer := NewSSEServer(mcpServer, WithBaseURL("http://localhost:8080"))
 
 		middleware := func(next http.Handler) http.Handler {
 			return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
@@ -336,7 +336,7 @@ func TestSSEServer(t *testing.T) {
 
 	t.Run("Works with custom mux", func(t *testing.T) {
 		mcpServer := NewMCPServer("test", "1.0.0")
-		sseServer := NewSSEServer(mcpServer, "")
+		sseServer := NewSSEServer(mcpServer)
 
 		mux := http.NewServeMux()
 		mux.Handle("/mcp/", http.StripPrefix("/mcp", sseServer))
@@ -397,7 +397,7 @@ func TestSSEServer(t *testing.T) {
 			t.Fatalf("Failed to send message: %v", err)
 		}
 		defer resp.Body.Close()
-		
+
 		if resp.StatusCode != http.StatusAccepted {
 			t.Errorf("Expected status 202, got %d", resp.StatusCode)
 		}