feat: properly handle interrupt signal for proxy (#1206)

Co-authored-by: taskbot <[email protected]>
Co-authored-by: Copilot <[email protected]>
Co-authored-by: Don Browne <[email protected]>

Author: 4 people

cmd/thv/app/proxy.go

@@ -143,7 +143,8 @@ func init() {
 }
 
 func proxyCmdFunc(cmd *cobra.Command, args []string) error {
-	ctx := cmd.Context()
+	ctx, stopSignal := signal.NotifyContext(cmd.Context(), syscall.SIGINT, syscall.SIGTERM)
+	defer stopSignal()
 	// Get the server name
 	serverName := args[0]
 
@@ -223,21 +224,15 @@ func proxyCmdFunc(cmd *cobra.Command, args []string) error {
 		serverName, port, proxyTargetURI)
 	logger.Info("Press Ctrl+C to stop")
 
-	// Set up signal handling
-	sigCh := make(chan os.Signal, 1)
-	signal.Notify(sigCh, syscall.SIGINT, syscall.SIGTERM)
+	<-ctx.Done()
+	logger.Infof("Interrupt received, proxy is shutting down. Please wait for connections to close...")
 
-	// Wait for signal
-	sig := <-sigCh
-	logger.Infof("Received signal %s, stopping proxy...", sig)
-
-	// Stop the proxy
-	if err := proxy.Stop(ctx); err != nil {
-		logger.Warnf("Warning: Failed to stop proxy: %v", err)
+	if err := proxy.CloseListener(); err != nil {
+		logger.Warnf("Error closing proxy listener: %v", err)
 	}
-
-	logger.Infof("Proxy for server %s stopped", serverName)
-	return nil
+	shutdownCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	defer cancel()
+	return proxy.Stop(shutdownCtx)
 }
 
 // AuthInfo contains authentication information extracted from WWW-Authenticate header

pkg/transport/proxy/transparent/transparent_proxy.go

@@ -7,9 +7,11 @@ import (
 	"bytes"
 	"context"
 	"encoding/json"
+	"errors"
 	"fmt"
 	"io"
 	"mime"
+	"net"
 	"net/http"
 	"net/http/httputil"
 	"net/url"
@@ -61,6 +63,9 @@ type TransparentProxy struct {
 
 	// If mcp server has been initialized
 	IsServerInitialized bool
+
+	// Listener for the HTTP server
+	listener net.Listener
 }
 
 // NewTransparentProxy creates a new transparent proxy with optional middlewares.
@@ -129,10 +134,13 @@ func (t *tracingTransport) RoundTrip(req *http.Request) (*http.Response, error)
 
 	resp, err := t.forward(req)
 	if err != nil {
+		if errors.Is(err, context.Canceled) {
+			// Expected during shutdown or client disconnect—silently ignore
+			return nil, err
+		}
 		logger.Errorf("Failed to forward request: %v", err)
 		return nil, err
 	}
-
 	if resp.StatusCode == http.StatusOK {
 		// check if we saw a valid mcp header
 		ct := resp.Header.Get("Mcp-Session-Id")
@@ -286,6 +294,11 @@ func (p *TransparentProxy) Start(ctx context.Context) error {
 		mux.Handle("/metrics", p.prometheusHandler)
 		logger.Info("Prometheus metrics endpoint enabled at /metrics")
 	}
+	ln, err := net.Listen("tcp", fmt.Sprintf("%s:%d", p.host, p.port))
+	if err != nil {
+		return fmt.Errorf("failed to listen: %w", err)
+	}
+	p.listener = ln
 
 	// Create the server
 	p.server = &http.Server{
@@ -294,16 +307,17 @@ func (p *TransparentProxy) Start(ctx context.Context) error {
 		ReadHeaderTimeout: 10 * time.Second, // Prevent Slowloris attacks
 	}
 
-	// Start the server in a goroutine
 	go func() {
-		logger.Infof("Transparent proxy started for container %s on %s:%d -> %s",
-			p.containerName, p.host, p.port, p.targetURI)
-
-		if err := p.server.ListenAndServe(); err != nil && err != http.ErrServerClosed {
+		err := p.server.Serve(ln)
+		if err != nil && err != http.ErrServerClosed {
+			var opErr *net.OpError
+			if errors.As(err, &opErr) && opErr.Op == "accept" {
+				// Expected when listener is closed—silently return
+				return
+			}
 			logger.Errorf("Transparent proxy error: %v", err)
 		}
 	}()
-
 	// Start health-check monitoring only if health checker is enabled
 	if p.healthChecker != nil {
 		go p.monitorHealth(ctx)
@@ -312,6 +326,14 @@ func (p *TransparentProxy) Start(ctx context.Context) error {
 	return nil
 }
 
+// CloseListener closes the listener for the transparent proxy.
+func (p *TransparentProxy) CloseListener() error {
+	if p.listener != nil {
+		return p.listener.Close()
+	}
+	return nil
+}
+
 func (p *TransparentProxy) monitorHealth(parentCtx context.Context) {
 	ticker := time.NewTicker(10 * time.Second)
 	defer ticker.Stop()
@@ -352,7 +374,13 @@ func (p *TransparentProxy) Stop(ctx context.Context) error {
 
 	// Stop the HTTP server
 	if p.server != nil {
-		return p.server.Shutdown(ctx)
+		err := p.server.Shutdown(ctx)
+		if err != nil && err != http.ErrServerClosed && err != context.DeadlineExceeded {
+			logger.Warnf("Error during proxy shutdown: %v", err)
+			return err
+		}
+		logger.Infof("Server for %s stopped successfully", p.containerName)
+		p.server = nil
 	}
 
 	return nil