Expose SSE server as HttpHandler

Author: ezynda3

server/sse.go

@@ -244,3 +244,15 @@ func (s *SSEServer) SendEventToSession(
 		return nil
 	}
 }
+
+// ServeHTTP implements the http.Handler interface.
+func (s *SSEServer) ServeHTTP(w http.ResponseWriter, r *http.Request) {
+	switch r.URL.Path {
+	case "/sse":
+		s.handleSSE(w, r)
+	case "/message":
+		s.handleMessage(w, r)
+	default:
+		http.NotFound(w, r)
+	}
+}

server/sse_test.go

@@ -2,9 +2,11 @@ package server
 
 import (
 	"bytes"
+	"context"
 	"encoding/json"
 	"fmt"
 	"net/http"
+	"net/http/httptest"
 	"strings"
 	"sync"
 	"testing"
@@ -229,4 +231,178 @@ func TestSSEServer(t *testing.T) {
 			t.Fatal("Timeout waiting for sessions to complete")
 		}
 	})
+
+	t.Run("Can be used as http.Handler", func(t *testing.T) {
+		mcpServer := NewMCPServer("test", "1.0.0")
+		sseServer := NewSSEServer(mcpServer, "http://localhost:8080")
+
+		ts := httptest.NewServer(sseServer)
+		defer ts.Close()
+
+		// Test 404 for unknown path first (simpler case)
+		resp, err := http.Get(fmt.Sprintf("%s/unknown", ts.URL))
+		if err != nil {
+			t.Fatalf("Failed to make request: %v", err)
+		}
+		defer resp.Body.Close()
+		if resp.StatusCode != http.StatusNotFound {
+			t.Errorf("Expected status 404, got %d", resp.StatusCode)
+		}
+
+		// Test SSE endpoint with proper cleanup
+		ctx, cancel := context.WithCancel(context.Background())
+		defer cancel()
+
+		req, err := http.NewRequestWithContext(ctx, "GET", fmt.Sprintf("%s/sse", ts.URL), nil)
+		if err != nil {
+			t.Fatalf("Failed to create request: %v", err)
+		}
+
+		resp, err = http.DefaultClient.Do(req)
+		if err != nil {
+			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)
+		}
+
+		// Read initial message in goroutine
+		done := make(chan struct{})
+		go func() {
+			defer close(done)
+			buf := make([]byte, 1024)
+			_, err := resp.Body.Read(buf)
+			if err != nil && err.Error() != "context canceled" {
+				t.Errorf("Failed to read from SSE stream: %v", err)
+			}
+		}()
+
+		// Wait briefly for initial response then cancel
+		time.Sleep(100 * time.Millisecond)
+		cancel()
+		<-done
+	})
+
+	t.Run("Works with middleware", func(t *testing.T) {
+		mcpServer := NewMCPServer("test", "1.0.0")
+		sseServer := NewSSEServer(mcpServer, "http://localhost:8080")
+
+		middleware := func(next http.Handler) http.Handler {
+			return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+				w.Header().Set("X-Test", "middleware")
+				next.ServeHTTP(w, r)
+			})
+		}
+
+		ts := httptest.NewServer(middleware(sseServer))
+		defer ts.Close()
+
+		ctx, cancel := context.WithCancel(context.Background())
+		defer cancel()
+
+		req, err := http.NewRequestWithContext(ctx, "GET", fmt.Sprintf("%s/sse", ts.URL), nil)
+		if err != nil {
+			t.Fatalf("Failed to create request: %v", err)
+		}
+
+		resp, err := http.DefaultClient.Do(req)
+		if err != nil {
+			t.Fatalf("Failed to connect to SSE endpoint: %v", err)
+		}
+		defer resp.Body.Close()
+
+		if resp.Header.Get("X-Test") != "middleware" {
+			t.Error("Middleware header not found")
+		}
+
+		// Read initial message in goroutine
+		done := make(chan struct{})
+		go func() {
+			defer close(done)
+			buf := make([]byte, 1024)
+			_, err := resp.Body.Read(buf)
+			if err != nil && err.Error() != "context canceled" {
+				t.Errorf("Failed to read from SSE stream: %v", err)
+			}
+		}()
+
+		// Wait briefly then cancel
+		time.Sleep(100 * time.Millisecond)
+		cancel()
+		<-done
+	})
+
+	t.Run("Works with custom mux", func(t *testing.T) {
+		mcpServer := NewMCPServer("test", "1.0.0")
+		sseServer := NewSSEServer(mcpServer, "")
+
+		mux := http.NewServeMux()
+		mux.Handle("/mcp/", http.StripPrefix("/mcp", sseServer))
+
+		ts := httptest.NewServer(mux)
+		defer ts.Close()
+
+		sseServer.baseURL = ts.URL + "/mcp"
+
+		ctx, cancel := context.WithCancel(context.Background())
+		defer cancel()
+
+		req, err := http.NewRequestWithContext(ctx, "GET", fmt.Sprintf("%s/mcp/sse", ts.URL), nil)
+		if err != nil {
+			t.Fatalf("Failed to create request: %v", err)
+		}
+
+		resp, err := http.DefaultClient.Do(req)
+		if err != nil {
+			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)
+		}
+
+		// Read the endpoint event
+		buf := make([]byte, 1024)
+		n, err := resp.Body.Read(buf)
+		if err != nil {
+			t.Fatalf("Failed to read SSE response: %v", err)
+		}
+
+		endpointEvent := string(buf[:n])
+		messageURL := strings.TrimSpace(
+			strings.Split(strings.Split(endpointEvent, "data: ")[1], "\n")[0],
+		)
+
+		// The messageURL should already be correct since we set the baseURL correctly
+		// Test message endpoint
+		initRequest := map[string]interface{}{
+			"jsonrpc": "2.0",
+			"id":      1,
+			"method":  "initialize",
+			"params": map[string]interface{}{
+				"protocolVersion": "2024-11-05",
+				"clientInfo": map[string]interface{}{
+					"name":    "test-client",
+					"version": "1.0.0",
+				},
+			},
+		}
+		requestBody, _ := json.Marshal(initRequest)
+
+		resp, err = http.Post(messageURL, "application/json", bytes.NewBuffer(requestBody))
+		if err != nil {
+			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)
+		}
+
+		// Clean up SSE connection
+		cancel()
+	})
 }