Introduce wrapper middleware to intercept http requests

Author: ardaguclu

pkg/kubernetes-mcp-server/cmd/root.go

@@ -9,6 +9,8 @@ import (
 	"strconv"
 	"strings"
 
+	"github.com/manusa/kubernetes-mcp-server/pkg/middleware"
+
 	"github.com/spf13/cobra"
 
 	"k8s.io/cli-runtime/pkg/genericiooptions"
@@ -206,9 +208,11 @@ func (m *MCPServerOptions) Run() error {
 
 	if m.StaticConfig.Port != "" {
 		mux := http.NewServeMux()
+		wrappedMux := middleware.RequestMiddleware(mux)
+
 		httpServer := &http.Server{
 			Addr:    ":" + m.StaticConfig.Port,
-			Handler: mux,
+			Handler: wrappedMux,
 		}
 
 		sseServer := mcpServer.ServeSse(m.SSEBaseUrl, httpServer)

pkg/mcp/mcp.go

@@ -98,6 +98,7 @@ func (s *Server) ServeHTTP(httpServer *http.Server) *server.StreamableHTTPServer
 	options := []server.StreamableHTTPOption{
 		server.WithHTTPContextFunc(contextFunc),
 		server.WithStreamableHTTPServer(httpServer),
+		server.WithStateLess(true),
 	}
 	return server.NewStreamableHTTPServer(s.server, options...)
 }

pkg/middleware/middleware.go

@@ -0,0 +1,61 @@
+package middleware
+
+import (
+	"bufio"
+	"net"
+	"net/http"
+	"time"
+
+	"k8s.io/klog/v2"
+)
+
+func RequestMiddleware(next http.Handler) http.Handler {
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		start := time.Now()
+
+		lrw := &loggingResponseWriter{
+			ResponseWriter: w,
+			statusCode:     http.StatusOK,
+		}
+
+		next.ServeHTTP(lrw, r)
+
+		duration := time.Since(start)
+		klog.V(5).Infof("%s %s %d %v", r.Method, r.URL.Path, lrw.statusCode, duration)
+	})
+}
+
+type loggingResponseWriter struct {
+	http.ResponseWriter
+	statusCode    int
+	headerWritten bool
+}
+
+func (lrw *loggingResponseWriter) WriteHeader(code int) {
+	if !lrw.headerWritten {
+		lrw.statusCode = code
+		lrw.headerWritten = true
+		lrw.ResponseWriter.WriteHeader(code)
+	}
+}
+
+func (lrw *loggingResponseWriter) Write(b []byte) (int, error) {
+	if !lrw.headerWritten {
+		lrw.statusCode = http.StatusOK
+		lrw.headerWritten = true
+	}
+	return lrw.ResponseWriter.Write(b)
+}
+
+func (lrw *loggingResponseWriter) Flush() {
+	if flusher, ok := lrw.ResponseWriter.(http.Flusher); ok {
+		flusher.Flush()
+	}
+}
+
+func (lrw *loggingResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
+	if hijacker, ok := lrw.ResponseWriter.(http.Hijacker); ok {
+		return hijacker.Hijack()
+	}
+	return nil, nil, http.ErrNotSupported
+}