bring mcp server inline with warp and api template pattern

Author: stuartwk

mcp/logging_middleware.go

@@ -2,10 +2,10 @@ package main
 
 import (
 	"context"
-	"log"
 	"time"
 
 	"github.com/modelcontextprotocol/go-sdk/mcp"
+	"github.com/urnetwork/glog"
 )
 
 func createLoggingMiddleware() mcp.Middleware {
@@ -19,7 +19,7 @@ func createLoggingMiddleware() mcp.Middleware {
 			sessionID := req.GetSession().ID()
 
 			// Log request details.
-			log.Printf("[REQUEST] Session: %s | Method: %s",
+			glog.Infof("[REQUEST] Session: %s | Method: %s",
 				sessionID,
 				method)
 
@@ -30,13 +30,13 @@ func createLoggingMiddleware() mcp.Middleware {
 			duration := time.Since(start)
 
 			if err != nil {
-				log.Printf("[RESPONSE] Session: %s | Method: %s | Status: ERROR | Duration: %v | Error: %v",
+				glog.Infof("[RESPONSE] Session: %s | Method: %s | Status: ERROR | Duration: %v | Error: %v",
 					sessionID,
 					method,
 					duration,
 					err)
 			} else {
-				log.Printf("[RESPONSE] Session: %s | Method: %s | Status: OK | Duration: %v",
+				glog.Infof("[RESPONSE] Session: %s | Method: %s | Status: OK | Duration: %v",
 					sessionID,
 					method,
 					duration)

mcp/main.go

@@ -1,23 +1,76 @@
 package main
 
 import (
+	"context"
 	"flag"
-	"fmt"
-	"log"
+	"runtime"
+	"syscall"
+	"time"
+
+	"github.com/urnetwork/glog"
+	"github.com/urnetwork/server"
 )
 
+const DrainTimeout = 60 * time.Second
+
 var (
 	host = flag.String("host", "0.0.0.0", "host to listen on")
-	port = flag.Int("port", 8000, "port number to listen on")
+	port = flag.Int("port", 80, "port number to listen on")
+	p    = flag.Int("p", 80, "port number to listen on (short form)")
 )
 
 func main() {
 	flag.Parse()
 
-	addr := fmt.Sprintf("%s:%d", *host, *port)
-	log.Printf("Starting MCP server on %s", addr)
+	quitEvent := server.NewEventWithContext(context.Background())
+	defer quitEvent.Set()
+
+	closeFn := quitEvent.SetOnSignals(syscall.SIGQUIT, syscall.SIGTERM)
+	defer closeFn()
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	// drain on sigterm
+	go server.HandleError(func() {
+		defer cancel()
+		select {
+		case <-ctx.Done():
+			return
+		case <-quitEvent.Ctx.Done():
+			select {
+			case <-ctx.Done():
+				return
+			case <-time.After(DrainTimeout):
+			}
+		}
+	})
+
+	// Debugging - log goroutine count
+	go server.HandleError(func() {
+		for {
+			select {
+			case <-ctx.Done():
+				return
+			case <-time.After(30 * time.Second):
+			}
+
+			glog.Infof("[mcp]goroutines=%d/%d\n", runtime.NumGoroutine(), runtime.GOMAXPROCS(0))
+		}
+	})
+
+	server.Warmup()
+
+	glog.Infof(
+		"[mcp]serving %s %s on %s:%d\n",
+		server.RequireEnv(),
+		server.RequireVersion(),
+		*host,
+		*port,
+	)
 
-	if err := runServer(addr); err != nil {
-		log.Fatalf("Server failed: %v", err)
+	if err := runServer(ctx, *port); err != nil {
+		glog.Fatalf("[mcp]Server failed: %v", err)
 	}
+	glog.Infof("[mcp]close\n")
 }

mcp/server.go

@@ -1,29 +1,62 @@
 package main
 
 import (
-	"log"
+	"context"
+	"net"
 	"net/http"
+	"strconv"
+	"time"
 
 	"github.com/modelcontextprotocol/go-sdk/mcp"
+	"github.com/urnetwork/glog"
+	"github.com/urnetwork/server"
+	"github.com/urnetwork/server/router"
 )
 
-func runServer(addr string) error {
-	server := mcp.NewServer(&mcp.Implementation{
+func runServer(ctx context.Context, port int) error {
+	mcpServer := mcp.NewServer(&mcp.Implementation{
 		Name:    "urnetwork-mcp",
 		Version: "0.0.1-beta.1",
 	}, nil)
 
 	// add logging
-	server.AddReceivingMiddleware(createLoggingMiddleware())
+	mcpServer.AddReceivingMiddleware(createLoggingMiddleware())
 
 	// register tools
-	registerTools(server)
+	registerTools(mcpServer)
 
-	handler := mcp.NewStreamableHTTPHandler(func(req *http.Request) *mcp.Server {
-		return server
+	mcpHandler := mcp.NewStreamableHTTPHandler(func(req *http.Request) *mcp.Server {
+		return mcpServer
 	}, nil)
 
-	log.Printf("MCP server listening on %s", addr)
+	// Wrap MCP handler to match http.HandlerFunc signature
+	mcpHandlerFunc := func(w http.ResponseWriter, r *http.Request) {
+		mcpHandler.ServeHTTP(w, r)
+	}
 
-	return http.ListenAndServe(addr, handler)
+	// Define routes using the router package
+	routes := []*router.Route{
+		router.NewRoute("GET", "/status", router.WarpStatus),
+		router.NewRoute("*", ".*", mcpHandlerFunc),
+	}
+
+	listenIpv4, _, listenPort := server.RequireListenIpPort(port)
+
+	reusePort := false
+
+	httpServerOptions := server.HttpServerOptions{
+		ReadTimeout:  15 * time.Second,
+		WriteTimeout: 30 * time.Second,
+		IdleTimeout:  5 * time.Minute,
+	}
+
+	glog.Infof("MCP server listening on %s:%d", listenIpv4, listenPort)
+
+	return server.HttpListenAndServeWithReusePort(
+		ctx,
+		net.JoinHostPort(listenIpv4, strconv.Itoa(listenPort)),
+		router.NewRouter(ctx, routes),
+		reusePort,
+		httpServerOptions,
+	)
 }