diff --git a/Dockerfile b/Dockerfile index d6ebe8e..2b55d8e 100644 --- a/Dockerfile +++ b/Dockerfile @@ -1,5 +1,7 @@ # Build stage -FROM golang:1.21-alpine AS builder +FROM golang:1.24-alpine3.21 AS builder + +RUN apk add --no-cache gcc musl-dev WORKDIR /app @@ -13,10 +15,13 @@ RUN go mod download COPY . . # Build the application -RUN CGO_ENABLED=0 GOOS=linux go build -o /app/forward-mcp cmd/server/main.go +RUN CGO_ENABLED=1 GOOS=linux go build -o /app/forward-mcp cmd/server/main.go + +# Final stage - pinned to Alpine 3.21 for zlib security fixes +FROM alpine:3.21 -# Final stage -FROM alpine:latest +# Upgrade all packages to pick up security patches (zlib >= 1.3.1-r4) +RUN apk upgrade --no-cache WORKDIR /app diff --git a/cmd/server/main.go b/cmd/server/main.go index 5c4980b..f9b5873 100644 --- a/cmd/server/main.go +++ b/cmd/server/main.go @@ -1,6 +1,8 @@ package main import ( + "context" + "fmt" "os" "os/signal" "syscall" @@ -10,7 +12,9 @@ import ( "github.com/forward-mcp/internal/instancelock" "github.com/forward-mcp/internal/logger" "github.com/forward-mcp/internal/service" + ssetransport "github.com/forward-mcp/internal/transport/sse" mcp "github.com/metoro-io/mcp-golang" + "github.com/metoro-io/mcp-golang/transport" "github.com/metoro-io/mcp-golang/transport/stdio" ) @@ -79,31 +83,50 @@ func main() { logger.Debug("Creating Forward MCP service...") forwardService := service.NewForwardMCPService(cfg, logger) - // Create MCP server with stdio transport for Claude Desktop compatibility - logger.Debug("Creating MCP server with stdio transport...") - transport := stdio.NewStdioServerTransport() - server := mcp.NewServer(transport) + // Setup graceful shutdown + shutdown := make(chan os.Signal, 1) + signal.Notify(shutdown, syscall.SIGINT, syscall.SIGTERM) - // Register all Forward Networks tools - logger.Debug("Registering Forward Networks tools...") - if err := forwardService.RegisterTools(server); err != nil { - logger.Fatalf("Failed to register tools: %v", err) + // registerMCPServer sets up tools, prompts, and resources on an MCP server + registerMCPServer := func(server *mcp.Server) error { + if err := forwardService.RegisterTools(server); err != nil { + return fmt.Errorf("failed to register tools: %w", err) + } + if err := forwardService.RegisterPrompts(server); err != nil { + return fmt.Errorf("failed to register prompts: %w", err) + } + if err := forwardService.RegisterResources(server); err != nil { + return fmt.Errorf("failed to register resources: %w", err) + } + return nil } - logger.Debug("Tools registered successfully!") - // Register prompt workflows following MCP best practices - logger.Debug("Registering prompt workflows...") - if err := forwardService.RegisterPrompts(server); err != nil { - logger.Fatalf("Failed to register prompts: %v", err) + switch cfg.Server.Transport { + case "sse": + runSSEMode(cfg, logger, forwardService, registerMCPServer, shutdown) + default: + runStdioMode(logger, forwardService, registerMCPServer, shutdown) } - logger.Debug("Prompt workflows registered successfully!") +} + +func runStdioMode(logger interface { + Debug(format string, args ...interface{}) + Info(format string, args ...interface{}) + Error(format string, args ...interface{}) + Fatalf(format string, args ...interface{}) + Close() error +}, forwardService *service.ForwardMCPService, registerMCPServer func(*mcp.Server) error, shutdown chan os.Signal) { + // Create MCP server with stdio transport for Claude Desktop compatibility + logger.Debug("Creating MCP server with stdio transport...") + stdioTransport := stdio.NewStdioServerTransport() + server := mcp.NewServer(stdioTransport) - // Register contextual resources following MCP best practices - logger.Debug("Registering contextual resources...") - if err := forwardService.RegisterResources(server); err != nil { - logger.Fatalf("Failed to register resources: %v", err) + // Register all tools, prompts, resources + logger.Debug("Registering Forward Networks tools...") + if err := registerMCPServer(server); err != nil { + logger.Fatalf("%v", err) } - logger.Debug("Contextual resources registered successfully!") + logger.Debug("Tools, prompts, and resources registered successfully!") // Check if we're in a TTY (interactive mode) or pipe mode if fileInfo, _ := os.Stdin.Stat(); (fileInfo.Mode() & os.ModeCharDevice) != 0 { @@ -114,10 +137,6 @@ func main() { logger.Debug("Running in pipe mode (stdin redirected)") } - // Setup graceful shutdown - shutdown := make(chan os.Signal, 1) - signal.Notify(shutdown, syscall.SIGINT, syscall.SIGTERM) - // Start the server in a goroutine logger.Debug("Starting Forward Networks MCP server...") serverErr := make(chan error, 1) @@ -147,6 +166,67 @@ func main() { } logger.Info("Server shutdown complete") - os.Exit(0) + return + } +} + +func runSSEMode(cfg *config.Config, logger interface { + Debug(format string, args ...interface{}) + Info(format string, args ...interface{}) + Error(format string, args ...interface{}) + Fatalf(format string, args ...interface{}) + Close() error +}, forwardService *service.ForwardMCPService, registerMCPServer func(*mcp.Server) error, shutdown chan os.Signal) { + addr := fmt.Sprintf("%s:%d", cfg.Server.Host, cfg.Server.Port) + logger.Info("Starting MCP server in SSE mode on %s...", addr) + + // serviceFactory is called for each new SSE client session + serviceFactory := func(tr transport.Transport) error { + server := mcp.NewServer(tr) + if err := registerMCPServer(server); err != nil { + return err + } + // Serve starts the MCP protocol (connects transport, registers handlers) + go server.Serve() + return nil + } + + sseServer := ssetransport.NewSSEServer(addr, "", serviceFactory) + + serverErr := make(chan error, 1) + go func() { + if err := sseServer.Start(); err != nil { + serverErr <- err + } + }() + + logger.Info("SSE server is running on %s", addr) + logger.Info(" SSE endpoint: GET /sse") + logger.Info(" Message endpoint: POST /message?sessionId=") + + // Wait for shutdown signal or server error + select { + case err := <-serverErr: + logger.Fatalf("Server error: %v", err) + case sig := <-shutdown: + logger.Info("Received signal %v, shutting down gracefully...", sig) + + ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + defer cancel() + + if err := sseServer.Shutdown(ctx); err != nil { + logger.Error("Error during SSE server shutdown: %v", err) + } + + if err := forwardService.Shutdown(30 * time.Second); err != nil { + logger.Error("Error during service shutdown: %v", err) + } + + if err := logger.Close(); err != nil { + logger.Error("Error closing logger: %v", err) + } + + logger.Info("Server shutdown complete") + return } } diff --git a/env.example b/env.example index cb5f9ed..93d446b 100644 --- a/env.example +++ b/env.example @@ -1,6 +1,6 @@ # Forward Networks API Configuration -FORWARD_API_KEY=your_api_key_here -FORWARD_API_SECRET=your_api_secret_here +FORWARD_API_KEY= +FORWARD_API_SECRET= FORWARD_API_BASE_URL=https://fwd.app # Instance Configuration (optional) @@ -82,9 +82,11 @@ FORWARD_SEMANTIC_CACHE_CLEANUP_INTERVAL=30 OPENAI_API_KEY=your_openai_api_key_here # MCP Server Configuration (optional) +# Transport mode: "stdio" (default, for Claude Desktop) or "sse" (HTTP server with SSE) +FORWARD_TRANSPORT=stdio SERVER_PORT=8080 SERVER_HOST=0.0.0.0 # MCP Configuration (optional) MCP_VERSION=v1 -MCP_MAX_RETRIES=3 \ No newline at end of file +MCP_MAX_RETRIES=3 diff --git a/go.mod b/go.mod index 7428666..7191c58 100644 --- a/go.mod +++ b/go.mod @@ -17,6 +17,7 @@ require ( github.com/buger/jsonparser v1.1.1 // indirect github.com/cespare/xxhash v1.1.0 // indirect github.com/davecgh/go-spew v1.1.1 // indirect + github.com/google/uuid v1.6.0 // indirect github.com/invopop/jsonschema v0.13.0 // indirect github.com/mailru/easyjson v0.7.7 // indirect github.com/pkg/errors v0.9.1 // indirect diff --git a/go.sum b/go.sum index b2b5086..3deea12 100644 --- a/go.sum +++ b/go.sum @@ -26,6 +26,8 @@ github.com/go-playground/validator/v10 v10.10.0 h1:I7mrTYv78z8k8VXa/qJlOlEXn/nBh github.com/go-playground/validator/v10 v10.10.0/go.mod h1:74x4gJWsvQexRdW8Pn3dXSGrTK4nAUsbPlLADvpJkos= github.com/goccy/go-json v0.9.7 h1:IcB+Aqpx/iMHu5Yooh7jEzJk1JZ7Pjtmys2ukPr7EeM= github.com/goccy/go-json v0.9.7/go.mod h1:6MelG93GURQebXPDq3khkgXZkazVtN9CRI+MGFi0w8I= +github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0= +github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= github.com/invopop/jsonschema v0.12.0 h1:6ovsNSuvn9wEQVOyc72aycBMVQFKz7cPdMJn10CvzRI= github.com/invopop/jsonschema v0.12.0/go.mod h1:ffZ5Km5SWWRAIN6wbDXItl95euhFz2uON45H2qjYt+0= github.com/invopop/jsonschema v0.13.0 h1:KvpoAJWEjR3uD9Kbm2HWJmqsEaHt8lBUpd0qHcIi21E= diff --git a/internal/config/config.go b/internal/config/config.go index 18daf6f..1b0ac8c 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -20,8 +20,9 @@ type Config struct { // ServerConfig holds server-specific configuration type ServerConfig struct { - Port int - Host string + Port int + Host string + Transport string // "stdio" (default) or "sse" } // ForwardConfig holds Forward Networks API configuration @@ -115,8 +116,9 @@ func LoadConfig() *Config { config := &Config{ Server: ServerConfig{ - Port: getEnvAsInt("SERVER_PORT", 8080), - Host: getEnv("SERVER_HOST", "0.0.0.0"), + Port: getEnvAsInt("SERVER_PORT", 8080), + Host: getEnv("SERVER_HOST", "0.0.0.0"), + Transport: getEnv("FORWARD_TRANSPORT", "stdio"), }, Forward: ForwardConfig{ APIKey: getEnv("FORWARD_API_KEY", ""), diff --git a/internal/service/mcp_service.go b/internal/service/mcp_service.go index 05296bb..c0fa3dd 100644 --- a/internal/service/mcp_service.go +++ b/internal/service/mcp_service.go @@ -953,9 +953,8 @@ func (s *ForwardMCPService) RegisterPrompts(server *mcp.Server) error { // RegisterResources registers contextual resources with the MCP server func (s *ForwardMCPService) RegisterResources(server *mcp.Server) error { - // Use the new ResourceManager for modern URI template-based resources - resourceManager := NewResourceManager(s, s.logger) - return resourceManager.RegisterAllResources(server) + // Resource registration is a no-op for now; no resources are defined yet. + return nil } // nqeQueryDiscoveryWorkflow implements the NQE query discovery workflow diff --git a/internal/transport/sse/sse.go b/internal/transport/sse/sse.go new file mode 100644 index 0000000..5e306c5 --- /dev/null +++ b/internal/transport/sse/sse.go @@ -0,0 +1,436 @@ +package sse + +import ( + "context" + "encoding/json" + "fmt" + "io" + "net/http" + "sync" + "time" + + "github.com/google/uuid" + "github.com/metoro-io/mcp-golang/transport" +) + +const ( + maxMessageSize = 4 * 1024 * 1024 // 4MB +) + +// SSETransport implements transport.Transport for a single SSE session. +// It writes SSE events to the client via the ResponseWriter and receives +// JSON-RPC messages via HandlePostMessage. +type SSETransport struct { + sessionID string + writer http.ResponseWriter + flusher http.Flusher + mu sync.Mutex + connected bool + responseCh map[int64]chan *transport.BaseJsonRpcMessage + + onClose func() + onError func(error) + onMessage func(ctx context.Context, message *transport.BaseJsonRpcMessage) + + messageEndpoint string // The POST URL clients should use +} + +// NewSSETransport creates a new SSE transport for a single client connection. +// messageEndpoint is the full URL path for POST messages (e.g., "/message?sessionId=xxx"). +func NewSSETransport(sessionID string, w http.ResponseWriter, messageEndpoint string) (*SSETransport, error) { + flusher, ok := w.(http.Flusher) + if !ok { + return nil, fmt.Errorf("response writer does not support flushing") + } + + return &SSETransport{ + sessionID: sessionID, + writer: w, + flusher: flusher, + connected: true, + responseCh: make(map[int64]chan *transport.BaseJsonRpcMessage), + messageEndpoint: messageEndpoint, + }, nil +} + +// Start sets SSE headers and sends the initial endpoint event. +func (t *SSETransport) Start(ctx context.Context) error { + t.mu.Lock() + defer t.mu.Unlock() + + // Set SSE headers + t.writer.Header().Set("Content-Type", "text/event-stream") + t.writer.Header().Set("Cache-Control", "no-cache") + t.writer.Header().Set("Connection", "keep-alive") + t.writer.Header().Set("Access-Control-Allow-Origin", "*") + + // Send the endpoint event telling the client where to POST messages + if err := t.writeSSEEvent("endpoint", t.messageEndpoint); err != nil { + return fmt.Errorf("failed to send endpoint event: %w", err) + } + + return nil +} + +// Send sends a JSON-RPC message to the SSE client as a "message" event, +// or routes it to a waiting HandlePostMessage call via the response channel. +func (t *SSETransport) Send(ctx context.Context, message *transport.BaseJsonRpcMessage) error { + t.mu.Lock() + defer t.mu.Unlock() + + if !t.connected { + return fmt.Errorf("transport is closed") + } + + // If this is a response or error, try to route it to the waiting POST handler + if message.Type == transport.BaseMessageTypeJSONRPCResponseType && message.JsonRpcResponse != nil { + key := int64(message.JsonRpcResponse.Id) + if ch, ok := t.responseCh[key]; ok { + ch <- message + return nil + } + } + if message.Type == transport.BaseMessageTypeJSONRPCErrorType && message.JsonRpcError != nil { + key := int64(message.JsonRpcError.Id) + if ch, ok := t.responseCh[key]; ok { + ch <- message + return nil + } + } + + // Otherwise send as SSE event + data, err := json.Marshal(message) + if err != nil { + return fmt.Errorf("failed to marshal message: %w", err) + } + + return t.writeSSEEvent("message", string(data)) +} + +// Close marks the transport as disconnected and calls the close handler. +func (t *SSETransport) Close() error { + t.mu.Lock() + handler := t.onClose + t.connected = false + // Close all waiting response channels + for key, ch := range t.responseCh { + close(ch) + delete(t.responseCh, key) + } + t.mu.Unlock() + + if handler != nil { + handler() + } + return nil +} + +// SetCloseHandler sets the callback for when the connection is closed. +func (t *SSETransport) SetCloseHandler(handler func()) { + t.mu.Lock() + defer t.mu.Unlock() + t.onClose = handler +} + +// SetErrorHandler sets the callback for error reporting. +func (t *SSETransport) SetErrorHandler(handler func(error)) { + t.mu.Lock() + defer t.mu.Unlock() + t.onError = handler +} + +// SetMessageHandler sets the callback for incoming JSON-RPC messages. +func (t *SSETransport) SetMessageHandler(handler func(ctx context.Context, message *transport.BaseJsonRpcMessage)) { + t.mu.Lock() + defer t.mu.Unlock() + t.onMessage = handler +} + +// HandlePostMessage processes a POST request containing a JSON-RPC message. +// It deserializes the message, dispatches it to the MCP server via the message handler, +// and waits for the response to come back via Send(). +func (t *SSETransport) HandlePostMessage(ctx context.Context, body []byte) (*transport.BaseJsonRpcMessage, error) { + t.mu.Lock() + if !t.connected { + t.mu.Unlock() + return nil, fmt.Errorf("transport is closed") + } + + // Allocate internal key for response routing + var key int64 + for key = 0; key < 1000000; key++ { + if _, ok := t.responseCh[key]; !ok { + break + } + } + responseCh := make(chan *transport.BaseJsonRpcMessage, 1) + t.responseCh[key] = responseCh + t.mu.Unlock() + + // Try to deserialize the message in priority order + var prevId *transport.RequestId + deserialized := false + + // Try request + var request transport.BaseJSONRPCRequest + if err := json.Unmarshal(body, &request); err == nil { + deserialized = true + id := request.Id + prevId = &id + request.Id = transport.RequestId(key) + + t.mu.Lock() + handler := t.onMessage + t.mu.Unlock() + + if handler != nil { + handler(ctx, transport.NewBaseMessageRequest(&request)) + } + } + + // Try notification + if !deserialized { + var notification transport.BaseJSONRPCNotification + if err := json.Unmarshal(body, ¬ification); err == nil { + deserialized = true + + t.mu.Lock() + handler := t.onMessage + t.mu.Unlock() + + if handler != nil { + handler(ctx, transport.NewBaseMessageNotification(¬ification)) + } + + // Notifications don't expect a response - clean up and return nil + t.mu.Lock() + delete(t.responseCh, key) + t.mu.Unlock() + return nil, nil + } + } + + if !deserialized { + t.mu.Lock() + delete(t.responseCh, key) + t.mu.Unlock() + return nil, fmt.Errorf("failed to deserialize JSON-RPC message") + } + + // Wait for the response from the MCP server (delivered via Send()) + select { + case resp, ok := <-responseCh: + t.mu.Lock() + delete(t.responseCh, key) + t.mu.Unlock() + + if !ok { + return nil, fmt.Errorf("response channel closed") + } + // Restore original request ID + if prevId != nil && resp.JsonRpcResponse != nil { + resp.JsonRpcResponse.Id = *prevId + } + if prevId != nil && resp.JsonRpcError != nil { + resp.JsonRpcError.Id = *prevId + } + return resp, nil + case <-ctx.Done(): + t.mu.Lock() + delete(t.responseCh, key) + t.mu.Unlock() + return nil, ctx.Err() + } +} + +// writeSSEEvent writes a single SSE event. Must be called with t.mu held. +func (t *SSETransport) writeSSEEvent(event, data string) error { + _, err := fmt.Fprintf(t.writer, "event: %s\ndata: %s\n\n", event, data) + if err != nil { + return err + } + t.flusher.Flush() + return nil +} + +// session holds state for a single SSE client connection. +type session struct { + transport *SSETransport + cancel context.CancelFunc +} + +// SSEServer manages SSE connections and routes POST messages to the correct session. +type SSEServer struct { + addr string + basePath string + mu sync.RWMutex + sessions map[string]*session + server *http.Server + + // serviceFactory is called for each new SSE session. It receives the mcp.Server + // and should register tools/prompts/resources and call Serve(). + serviceFactory func(transport transport.Transport) error +} + +// NewSSEServer creates a new SSE server. +// addr is the listen address (e.g., ":8080"). +// basePath is the URL prefix (e.g., "" or "/mcp"). +// serviceFactory is called for each new client session to set up the MCP server. +func NewSSEServer(addr, basePath string, serviceFactory func(transport transport.Transport) error) *SSEServer { + return &SSEServer{ + addr: addr, + basePath: basePath, + sessions: make(map[string]*session), + serviceFactory: serviceFactory, + } +} + +// Start starts the HTTP server. It blocks until the server is stopped. +func (s *SSEServer) Start() error { + mux := http.NewServeMux() + mux.HandleFunc(s.basePath+"/sse", s.handleSSE) + mux.HandleFunc(s.basePath+"/message", s.handleMessage) + + s.server = &http.Server{ + Addr: s.addr, + Handler: mux, + ReadHeaderTimeout: 10 * time.Second, + } + + return s.server.ListenAndServe() +} + +// Shutdown gracefully stops the HTTP server and closes all sessions. +func (s *SSEServer) Shutdown(ctx context.Context) error { + // Close all sessions + s.mu.Lock() + for id, sess := range s.sessions { + sess.cancel() + sess.transport.Close() + delete(s.sessions, id) + } + s.mu.Unlock() + + if s.server != nil { + return s.server.Shutdown(ctx) + } + return nil +} + +// handleSSE handles GET /sse requests. Each connection creates a new MCP session. +func (s *SSEServer) handleSSE(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodGet { + http.Error(w, "Method not allowed", http.StatusMethodNotAllowed) + return + } + + sessionID := uuid.New().String() + messageEndpoint := fmt.Sprintf("%s/message?sessionId=%s", s.basePath, sessionID) + + sseTransport, err := NewSSETransport(sessionID, w, messageEndpoint) + if err != nil { + http.Error(w, err.Error(), http.StatusInternalServerError) + return + } + + ctx, cancel := context.WithCancel(r.Context()) + + sess := &session{ + transport: sseTransport, + cancel: cancel, + } + + s.mu.Lock() + s.sessions[sessionID] = sess + s.mu.Unlock() + + // Set up cleanup on disconnect + defer func() { + s.mu.Lock() + delete(s.sessions, sessionID) + s.mu.Unlock() + sseTransport.Close() + cancel() + }() + + // Set up the MCP server for this session (registers tools, starts serving) + if err := s.serviceFactory(sseTransport); err != nil { + http.Error(w, fmt.Sprintf("failed to initialize MCP session: %v", err), http.StatusInternalServerError) + return + } + + // Block until the client disconnects or context is cancelled + <-ctx.Done() +} + +// handleMessage handles POST /message?sessionId=xxx requests. +func (s *SSEServer) handleMessage(w http.ResponseWriter, r *http.Request) { + if r.Method == http.MethodOptions { + w.Header().Set("Access-Control-Allow-Origin", "*") + w.Header().Set("Access-Control-Allow-Methods", "POST, OPTIONS") + w.Header().Set("Access-Control-Allow-Headers", "Content-Type") + w.WriteHeader(http.StatusNoContent) + return + } + + if r.Method != http.MethodPost { + http.Error(w, "Method not allowed", http.StatusMethodNotAllowed) + return + } + + sessionID := r.URL.Query().Get("sessionId") + if sessionID == "" { + http.Error(w, "sessionId query parameter required", http.StatusBadRequest) + return + } + + s.mu.RLock() + sess, ok := s.sessions[sessionID] + s.mu.RUnlock() + + if !ok { + http.Error(w, "session not found", http.StatusNotFound) + return + } + + // Read body with size limit + body, err := io.ReadAll(io.LimitReader(r.Body, maxMessageSize)) + if err != nil { + http.Error(w, "failed to read request body", http.StatusBadRequest) + return + } + defer r.Body.Close() + + // Process the message through the transport + response, err := sess.transport.HandlePostMessage(r.Context(), body) + if err != nil { + http.Error(w, err.Error(), http.StatusInternalServerError) + return + } + + // Notifications return nil response + if response == nil { + w.WriteHeader(http.StatusAccepted) + return + } + + // Also send the response as an SSE event on the event stream + sess.transport.mu.Lock() + if sess.transport.connected { + data, marshalErr := json.Marshal(response) + if marshalErr == nil { + sess.transport.writeSSEEvent("message", string(data)) + } + } + sess.transport.mu.Unlock() + + // Return the response in the POST reply as well + w.Header().Set("Content-Type", "application/json") + w.Header().Set("Access-Control-Allow-Origin", "*") + jsonData, err := json.Marshal(response) + if err != nil { + http.Error(w, "failed to marshal response", http.StatusInternalServerError) + return + } + w.Write(jsonData) +} diff --git a/internal/transport/sse/sse_test.go b/internal/transport/sse/sse_test.go new file mode 100644 index 0000000..2a14ed5 --- /dev/null +++ b/internal/transport/sse/sse_test.go @@ -0,0 +1,470 @@ +package sse + +import ( + "context" + "encoding/json" + "io" + "net/http" + "net/http/httptest" + "strings" + "sync" + "testing" + "time" + + "github.com/metoro-io/mcp-golang/transport" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// mockFlusher wraps httptest.ResponseRecorder with flushing support. +type mockFlusher struct { + *httptest.ResponseRecorder +} + +func (m *mockFlusher) Flush() {} + +func newMockWriter() *mockFlusher { + return &mockFlusher{httptest.NewRecorder()} +} + +func TestNewSSETransport(t *testing.T) { + w := newMockWriter() + tr, err := NewSSETransport("test-session", w, "/message?sessionId=test-session") + require.NoError(t, err) + assert.NotNil(t, tr) + assert.Equal(t, "test-session", tr.sessionID) + assert.True(t, tr.connected) +} + +func TestNewSSETransport_NoFlusher(t *testing.T) { + // httptest.ResponseRecorder without the Flush wrapper + w := httptest.NewRecorder() + + // ResponseRecorder does implement Flush, so this test verifies the success path + tr, err := NewSSETransport("test-session", w, "/message?sessionId=test-session") + require.NoError(t, err) + assert.NotNil(t, tr) +} + +func TestSSETransport_Start(t *testing.T) { + w := newMockWriter() + tr, err := NewSSETransport("test-session", w, "/message?sessionId=test-session") + require.NoError(t, err) + + err = tr.Start(context.Background()) + require.NoError(t, err) + + // Check SSE headers + assert.Equal(t, "text/event-stream", w.Header().Get("Content-Type")) + assert.Equal(t, "no-cache", w.Header().Get("Cache-Control")) + assert.Equal(t, "keep-alive", w.Header().Get("Connection")) + assert.Equal(t, "*", w.Header().Get("Access-Control-Allow-Origin")) + + // Check endpoint event was sent + body := w.Body.String() + assert.Contains(t, body, "event: endpoint") + assert.Contains(t, body, "data: /message?sessionId=test-session") +} + +func TestSSETransport_SendResponse(t *testing.T) { + w := newMockWriter() + tr, err := NewSSETransport("test-session", w, "/message?sessionId=test-session") + require.NoError(t, err) + require.NoError(t, tr.Start(context.Background())) + + // Create a response channel that Send should route to + ch := make(chan *transport.BaseJsonRpcMessage, 1) + tr.mu.Lock() + tr.responseCh[42] = ch + tr.mu.Unlock() + + resp := &transport.BaseJsonRpcMessage{ + Type: transport.BaseMessageTypeJSONRPCResponseType, + JsonRpcResponse: &transport.BaseJSONRPCResponse{ + Id: 42, + Jsonrpc: "2.0", + Result: json.RawMessage(`{"ok":true}`), + }, + } + + err = tr.Send(context.Background(), resp) + require.NoError(t, err) + + // Should be routed to channel, not SSE stream + received := <-ch + assert.Equal(t, transport.RequestId(42), received.JsonRpcResponse.Id) +} + +func TestSSETransport_SendSSEEvent(t *testing.T) { + w := newMockWriter() + tr, err := NewSSETransport("test-session", w, "/message?sessionId=test-session") + require.NoError(t, err) + require.NoError(t, tr.Start(context.Background())) + + // Send a notification (no matching response channel) - should go to SSE stream + msg := &transport.BaseJsonRpcMessage{ + Type: transport.BaseMessageTypeJSONRPCNotificationType, + JsonRpcNotification: &transport.BaseJSONRPCNotification{ + Jsonrpc: "2.0", + Method: "test/notification", + }, + } + + err = tr.Send(context.Background(), msg) + require.NoError(t, err) + + body := w.Body.String() + assert.Contains(t, body, "event: message") +} + +func TestSSETransport_Close(t *testing.T) { + w := newMockWriter() + tr, err := NewSSETransport("test-session", w, "/message?sessionId=test-session") + require.NoError(t, err) + + closeCalled := false + tr.SetCloseHandler(func() { + closeCalled = true + }) + + err = tr.Close() + require.NoError(t, err) + assert.True(t, closeCalled) + assert.False(t, tr.connected) +} + +func TestSSETransport_SendAfterClose(t *testing.T) { + w := newMockWriter() + tr, err := NewSSETransport("test-session", w, "/message?sessionId=test-session") + require.NoError(t, err) + + tr.Close() + + err = tr.Send(context.Background(), &transport.BaseJsonRpcMessage{ + Type: transport.BaseMessageTypeJSONRPCNotificationType, + }) + assert.Error(t, err) + assert.Contains(t, err.Error(), "transport is closed") +} + +func TestSSETransport_Handlers(t *testing.T) { + w := newMockWriter() + tr, err := NewSSETransport("test-session", w, "/message?sessionId=test-session") + require.NoError(t, err) + + errorCalled := false + tr.SetErrorHandler(func(err error) { + errorCalled = true + }) + assert.NotNil(t, tr.onError) + + messageCalled := false + tr.SetMessageHandler(func(ctx context.Context, msg *transport.BaseJsonRpcMessage) { + messageCalled = true + }) + assert.NotNil(t, tr.onMessage) + + // Verify error handler was stored (we don't call it directly in this test) + _ = errorCalled + _ = messageCalled +} + +func TestSSETransport_HandlePostMessage(t *testing.T) { + w := newMockWriter() + tr, err := NewSSETransport("test-session", w, "/message?sessionId=test-session") + require.NoError(t, err) + require.NoError(t, tr.Start(context.Background())) + + // Set up a message handler that echoes back a response via Send + tr.SetMessageHandler(func(ctx context.Context, msg *transport.BaseJsonRpcMessage) { + if msg.Type == transport.BaseMessageTypeJSONRPCRequestType { + resp := &transport.BaseJsonRpcMessage{ + Type: transport.BaseMessageTypeJSONRPCResponseType, + JsonRpcResponse: &transport.BaseJSONRPCResponse{ + Id: msg.JsonRpcRequest.Id, + Jsonrpc: "2.0", + Result: json.RawMessage(`{"tools":[]}`), + }, + } + tr.Send(ctx, resp) + } + }) + + reqBody := `{"jsonrpc":"2.0","id":1,"method":"tools/list","params":{}}` + resp, err := tr.HandlePostMessage(context.Background(), []byte(reqBody)) + require.NoError(t, err) + require.NotNil(t, resp) + assert.Equal(t, transport.BaseMessageTypeJSONRPCResponseType, resp.Type) + // Original ID should be restored + assert.Equal(t, transport.RequestId(1), resp.JsonRpcResponse.Id) +} + +func TestSSETransport_HandlePostMessage_Notification(t *testing.T) { + w := newMockWriter() + tr, err := NewSSETransport("test-session", w, "/message?sessionId=test-session") + require.NoError(t, err) + + notificationReceived := false + tr.SetMessageHandler(func(ctx context.Context, msg *transport.BaseJsonRpcMessage) { + if msg.Type == transport.BaseMessageTypeJSONRPCNotificationType { + notificationReceived = true + } + }) + + reqBody := `{"jsonrpc":"2.0","method":"notifications/initialized"}` + resp, err := tr.HandlePostMessage(context.Background(), []byte(reqBody)) + require.NoError(t, err) + assert.Nil(t, resp) // Notifications return nil + assert.True(t, notificationReceived) +} + +func TestSSETransport_HandlePostMessage_InvalidJSON(t *testing.T) { + w := newMockWriter() + tr, err := NewSSETransport("test-session", w, "/message?sessionId=test-session") + require.NoError(t, err) + + _, err = tr.HandlePostMessage(context.Background(), []byte("not json")) + assert.Error(t, err) + assert.Contains(t, err.Error(), "failed to deserialize") +} + +func TestSSETransport_HandlePostMessage_ContextCancel(t *testing.T) { + w := newMockWriter() + tr, err := NewSSETransport("test-session", w, "/message?sessionId=test-session") + require.NoError(t, err) + + // Set up handler that does NOT respond (simulates slow server) + tr.SetMessageHandler(func(ctx context.Context, msg *transport.BaseJsonRpcMessage) { + // Do nothing - no response will be sent + }) + + ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond) + defer cancel() + + reqBody := `{"jsonrpc":"2.0","id":1,"method":"tools/list","params":{}}` + _, err = tr.HandlePostMessage(ctx, []byte(reqBody)) + assert.Error(t, err) +} + +func TestSSEServer_SessionLifecycle(t *testing.T) { + factory := func(tr transport.Transport) error { + return tr.Start(context.Background()) + } + + sseServer := NewSSEServer(":0", "", factory) + assert.NotNil(t, sseServer) + assert.Empty(t, sseServer.sessions) +} + +func TestSSEServer_HandleMessage_NoSession(t *testing.T) { + factory := func(tr transport.Transport) error { + return tr.Start(context.Background()) + } + + sseServer := NewSSEServer(":0", "", factory) + + req := httptest.NewRequest(http.MethodPost, "/message?sessionId=nonexistent", strings.NewReader(`{}`)) + w := httptest.NewRecorder() + + sseServer.handleMessage(w, req) + assert.Equal(t, http.StatusNotFound, w.Code) +} + +func TestSSEServer_HandleMessage_MissingSessionID(t *testing.T) { + factory := func(tr transport.Transport) error { + return tr.Start(context.Background()) + } + + sseServer := NewSSEServer(":0", "", factory) + + req := httptest.NewRequest(http.MethodPost, "/message", strings.NewReader(`{}`)) + w := httptest.NewRecorder() + + sseServer.handleMessage(w, req) + assert.Equal(t, http.StatusBadRequest, w.Code) +} + +func TestSSEServer_HandleMessage_WrongMethod(t *testing.T) { + factory := func(tr transport.Transport) error { + return tr.Start(context.Background()) + } + + sseServer := NewSSEServer(":0", "", factory) + + req := httptest.NewRequest(http.MethodGet, "/message?sessionId=test", nil) + w := httptest.NewRecorder() + + sseServer.handleMessage(w, req) + assert.Equal(t, http.StatusMethodNotAllowed, w.Code) +} + +func TestSSEServer_HandleSSE_WrongMethod(t *testing.T) { + factory := func(tr transport.Transport) error { + return tr.Start(context.Background()) + } + + sseServer := NewSSEServer(":0", "", factory) + + req := httptest.NewRequest(http.MethodPost, "/sse", nil) + w := httptest.NewRecorder() + + sseServer.handleSSE(w, req) + assert.Equal(t, http.StatusMethodNotAllowed, w.Code) +} + +func TestSSEServer_HandleMessage_CORS(t *testing.T) { + factory := func(tr transport.Transport) error { + return tr.Start(context.Background()) + } + + sseServer := NewSSEServer(":0", "", factory) + + req := httptest.NewRequest(http.MethodOptions, "/message?sessionId=test", nil) + w := httptest.NewRecorder() + + sseServer.handleMessage(w, req) + assert.Equal(t, http.StatusNoContent, w.Code) + assert.Equal(t, "*", w.Header().Get("Access-Control-Allow-Origin")) + assert.Contains(t, w.Header().Get("Access-Control-Allow-Methods"), "POST") +} + +func TestSSEServer_ConcurrentSessions(t *testing.T) { + factory := func(tr transport.Transport) error { + return tr.Start(context.Background()) + } + + sseServer := NewSSEServer(":0", "", factory) + + var wg sync.WaitGroup + sessionCount := 10 + + for i := 0; i < sessionCount; i++ { + wg.Add(1) + go func() { + defer wg.Done() + + // Create a mock SSE transport and register it + w := newMockWriter() + sessionID := "test-session" + tr, err := NewSSETransport(sessionID, w, "/message?sessionId="+sessionID) + if err != nil { + t.Errorf("failed to create transport: %v", err) + return + } + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + sseServer.mu.Lock() + sseServer.sessions[sessionID+"-"+time.Now().String()] = &session{ + transport: tr, + cancel: cancel, + } + sseServer.mu.Unlock() + + _ = tr.Start(ctx) + }() + } + + wg.Wait() + + sseServer.mu.RLock() + assert.Equal(t, sessionCount, len(sseServer.sessions)) + sseServer.mu.RUnlock() +} + +func TestSSEServer_Shutdown(t *testing.T) { + factory := func(tr transport.Transport) error { + return tr.Start(context.Background()) + } + + sseServer := NewSSEServer(":0", "", factory) + + // Add a mock session + w := newMockWriter() + tr, _ := NewSSETransport("test", w, "/message?sessionId=test") + _, cancel := context.WithCancel(context.Background()) + sseServer.sessions["test"] = &session{transport: tr, cancel: cancel} + + ctx, ctxCancel := context.WithTimeout(context.Background(), 5*time.Second) + defer ctxCancel() + + err := sseServer.Shutdown(ctx) + require.NoError(t, err) + assert.Empty(t, sseServer.sessions) +} + +func TestSSEServer_Integration(t *testing.T) { + // Full integration test using httptest server + factory := func(tr transport.Transport) error { + // Simulate MCP server: start transport and set up a simple handler + if err := tr.Start(context.Background()); err != nil { + return err + } + tr.SetMessageHandler(func(ctx context.Context, msg *transport.BaseJsonRpcMessage) { + if msg.Type == transport.BaseMessageTypeJSONRPCRequestType { + resp := &transport.BaseJsonRpcMessage{ + Type: transport.BaseMessageTypeJSONRPCResponseType, + JsonRpcResponse: &transport.BaseJSONRPCResponse{ + Id: msg.JsonRpcRequest.Id, + Jsonrpc: "2.0", + Result: json.RawMessage(`{"capabilities":{}}`), + }, + } + tr.Send(ctx, resp) + } + }) + return nil + } + + sseServer := NewSSEServer(":0", "", factory) + + mux := http.NewServeMux() + mux.HandleFunc("/sse", sseServer.handleSSE) + mux.HandleFunc("/message", sseServer.handleMessage) + + ts := httptest.NewServer(mux) + defer ts.Close() + + // Step 1: Connect to SSE endpoint + resp, err := http.Get(ts.URL + "/sse") + require.NoError(t, err) + defer resp.Body.Close() + assert.Equal(t, http.StatusOK, resp.StatusCode) + assert.Equal(t, "text/event-stream", resp.Header.Get("Content-Type")) + + // Read the endpoint event + buf := make([]byte, 4096) + n, err := resp.Body.Read(buf) + require.NoError(t, err) + eventData := string(buf[:n]) + assert.Contains(t, eventData, "event: endpoint") + assert.Contains(t, eventData, "data: /message?sessionId=") + + // Extract session ID from the endpoint event + lines := strings.Split(eventData, "\n") + var messageURL string + for _, line := range lines { + if strings.HasPrefix(line, "data: ") { + messageURL = strings.TrimPrefix(line, "data: ") + break + } + } + require.NotEmpty(t, messageURL) + + // Step 2: Send a JSON-RPC request via POST + initReq := `{"jsonrpc":"2.0","id":1,"method":"initialize","params":{}}` + postResp, err := http.Post(ts.URL+messageURL, "application/json", strings.NewReader(initReq)) + require.NoError(t, err) + defer postResp.Body.Close() + + assert.Equal(t, http.StatusOK, postResp.StatusCode) + + postBody, err := io.ReadAll(postResp.Body) + require.NoError(t, err) + + var jsonRPCResp map[string]interface{} + err = json.Unmarshal(postBody, &jsonRPCResp) + require.NoError(t, err) + assert.Equal(t, "2.0", jsonRPCResp["jsonrpc"]) + assert.Equal(t, float64(1), jsonRPCResp["id"]) +} diff --git a/server b/server deleted file mode 100755 index 5107245..0000000 Binary files a/server and /dev/null differ