Skip to content

Add option to StreamableHTTPServer to allow custom http server instance #347

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 2 commits into from
May 30, 2025
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 3 additions & 1 deletion server/sse.go
Original file line number Diff line number Diff line change
@@ -227,7 +227,9 @@ func WithSSEEndpoint(endpoint string) SSEOption {
}
}

// WithHTTPServer sets the HTTP server instance
// WithHTTPServer sets the HTTP server instance.
// NOTE: When providing a custom HTTP server, you must handle routing yourself
// If routing is not set up, the server will start but won't handle any MCP requests.
func WithHTTPServer(srv *http.Server) SSEOption {
return func(s *SSEServer) {
s.srv = srv
30 changes: 24 additions & 6 deletions server/streamable_http.go
Original file line number Diff line number Diff line change
@@ -73,6 +73,15 @@ func WithHTTPContextFunc(fn HTTPContextFunc) StreamableHTTPOption {
}
}

// WithStreamableHTTPServer sets the HTTP server instance for StreamableHTTPServer.
// NOTE: When providing a custom HTTP server, you must handle routing yourself
// If routing is not set up, the server will start but won't handle any MCP requests.
func WithStreamableHTTPServer(srv *http.Server) StreamableHTTPOption {
return func(s *StreamableHTTPServer) {
s.httpServer = srv
}
}

// WithLogger sets the logger for the server
func WithLogger(logger util.Logger) StreamableHTTPOption {
return func(s *StreamableHTTPServer) {
@@ -155,15 +164,24 @@ func (s *StreamableHTTPServer) ServeHTTP(w http.ResponseWriter, r *http.Request)
// s.Start(":8080")
func (s *StreamableHTTPServer) Start(addr string) error {
s.mu.Lock()
mux := http.NewServeMux()
mux.Handle(s.endpointPath, s)
s.httpServer = &http.Server{
Addr: addr,
Handler: mux,
if s.httpServer == nil {
mux := http.NewServeMux()
mux.Handle(s.endpointPath, s)
s.httpServer = &http.Server{
Addr: addr,
Handler: mux,
}
} else {
if s.httpServer.Addr == "" {
s.httpServer.Addr = addr
} else if s.httpServer.Addr != addr {
return fmt.Errorf("conflicting listen address: WithStreamableHTTPServer(%q) vs Start(%q)", s.httpServer.Addr, addr)
}
Comment on lines +174 to +179
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nit: Should we document in this scenario that routing must be handled by the user? In the default/previous case, routing was handled by the Start method, but for the custom HTTP server instance, it is not. Maybe document it in a user-facing way with the WithStreamableHTTPServer function?

It can be a brief note because we can reasonably expect someone who wants to use a custom HTTP server to do this, regardless. But if they don't, the server would still start, but it wouldn't be able to do anything, right?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

okay, makes sense, will update for both sse and streamable http.

}
srv := s.httpServer
s.mu.Unlock()

return s.httpServer.ListenAndServe()
return srv.ListenAndServe()
}

// Shutdown gracefully stops the server, closing all active sessions
50 changes: 50 additions & 0 deletions server/streamable_http_test.go
Original file line number Diff line number Diff line change
@@ -670,6 +670,56 @@ func TestStreamableHTTP_SessionWithTools(t *testing.T) {
})
}

func TestStreamableHTTPServer_WithOptions(t *testing.T) {
t.Run("WithStreamableHTTPServer sets httpServer field", func(t *testing.T) {
mcpServer := NewMCPServer("test", "1.0.0")
customServer := &http.Server{Addr: ":9999"}
httpServer := NewStreamableHTTPServer(mcpServer, WithStreamableHTTPServer(customServer))

if httpServer.httpServer != customServer {
t.Errorf("Expected httpServer to be set to custom server instance, got %v", httpServer.httpServer)
}
})

t.Run("Start with conflicting address returns error", func(t *testing.T) {
mcpServer := NewMCPServer("test", "1.0.0")
customServer := &http.Server{Addr: ":9999"}
httpServer := NewStreamableHTTPServer(mcpServer, WithStreamableHTTPServer(customServer))

err := httpServer.Start(":8888")
if err == nil {
t.Error("Expected error for conflicting address, got nil")
} else if !strings.Contains(err.Error(), "conflicting listen address") {
t.Errorf("Expected error message to contain 'conflicting listen address', got '%s'", err.Error())
}
})

t.Run("Options consistency test", func(t *testing.T) {
mcpServer := NewMCPServer("test", "1.0.0")
endpointPath := "/test-mcp"
customServer := &http.Server{}

// Options to test
options := []StreamableHTTPOption{
WithEndpointPath(endpointPath),
WithStreamableHTTPServer(customServer),
}

// Apply options multiple times and verify consistency
for i := 0; i < 10; i++ {
server := NewStreamableHTTPServer(mcpServer, options...)

if server.endpointPath != endpointPath {
t.Errorf("Expected endpointPath %s, got %s", endpointPath, server.endpointPath)
}

if server.httpServer != customServer {
t.Errorf("Expected httpServer to match, got %v", server.httpServer)
}
}
})
}
Comment on lines +673 to +721
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🛠️ Refactor suggestion

Good test coverage but missing some important edge cases.

The current tests cover the basic functionality well:

  • ✅ Custom server field assignment
  • ✅ Address conflict detection
  • ✅ Options consistency

However, consider adding tests for these scenarios to improve coverage:

Add these missing test cases:

+	t.Run("Start with custom server and no address sets address", func(t *testing.T) {
+		mcpServer := NewMCPServer("test", "1.0.0")
+		customServer := &http.Server{} // No address set
+		httpServer := NewStreamableHTTPServer(mcpServer, WithStreamableHTTPServer(customServer))
+
+		// This should work and set the address
+		go func() {
+			time.Sleep(10 * time.Millisecond)
+			httpServer.Shutdown(context.Background())
+		}()
+
+		err := httpServer.Start(":0") // Use :0 for random port
+		if err != nil && !strings.Contains(err.Error(), "use of closed network connection") {
+			t.Errorf("Expected no error or connection closed error, got: %v", err)
+		}
+
+		if customServer.Addr != ":0" {
+			t.Errorf("Expected custom server address to be set to ':0', got '%s'", customServer.Addr)
+		}
+	})
+
+	t.Run("Start with custom server and matching address succeeds", func(t *testing.T) {
+		mcpServer := NewMCPServer("test", "1.0.0")
+		customServer := &http.Server{Addr: ":0"}
+		httpServer := NewStreamableHTTPServer(mcpServer, WithStreamableHTTPServer(customServer))
+
+		go func() {
+			time.Sleep(10 * time.Millisecond)
+			httpServer.Shutdown(context.Background())
+		}()
+
+		err := httpServer.Start(":0") // Same address
+		if err != nil && !strings.Contains(err.Error(), "use of closed network connection") {
+			t.Errorf("Expected no error or connection closed error, got: %v", err)
+		}
+	})
🤖 Prompt for AI Agents
In server/streamable_http_test.go around lines 673 to 721, the existing tests
cover basic functionality but lack edge case coverage. Add tests for scenarios
such as starting the server with an empty or invalid address, verifying behavior
when no options are provided, testing the default endpoint path, and ensuring
proper error handling when the underlying HTTP server fails to start. Implement
these additional test cases to improve overall test coverage and robustness.


func postJSON(url string, bodyObject any) (*http.Response, error) {
jsonBody, _ := json.Marshal(bodyObject)
req, _ := http.NewRequest(http.MethodPost, url, bytes.NewBuffer(jsonBody))