fix: eliminate 3 duplicate code patterns (launcher goroutine, API key extraction, JSON response)

Co-authored-by: lpcox <15877973+lpcox@users.noreply.github.com>

Author: Copilot

internal/cmd/root.go

@@ -285,10 +285,7 @@ func run(cmd *cobra.Command, args []string) error {
 		logger.LogInfoMd("startup", "Routes: /mcp/<server> for servers: %v", unifiedServer.GetServerIDs())
 
 		// Extract API key from gateway config (spec 7.1)
-		apiKey := ""
-		if cfg.Gateway != nil {
-			apiKey = cfg.Gateway.APIKey
-		}
+		apiKey := cfg.GetAPIKey()
 
 		httpServer = server.CreateHTTPServerForRoutedMode(listenAddr, unifiedServer, apiKey)
 	} else {
@@ -298,10 +295,7 @@ func run(cmd *cobra.Command, args []string) error {
 		logger.LogInfoMd("startup", "Endpoint: /mcp")
 
 		// Extract API key from gateway config (spec 7.1)
-		apiKey := ""
-		if cfg.Gateway != nil {
-			apiKey = cfg.Gateway.APIKey
-		}
+		apiKey := cfg.GetAPIKey()
 
 		httpServer = server.CreateHTTPServerForMCP(listenAddr, unifiedServer, apiKey)
 	}
@@ -362,10 +356,7 @@ func writeGatewayConfig(cfg *config.Config, listenAddr, mode string, w io.Writer
 	debugLog.Printf("Resolved gateway address: host=%s, port=%s", host, port)
 
 	// Extract API key from gateway config (per spec section 7.1)
-	apiKey := ""
-	if cfg.Gateway != nil {
-		apiKey = cfg.Gateway.APIKey
-	}
+	apiKey := cfg.GetAPIKey()
 	debugLog.Printf("Gateway config: auth_enabled=%v", apiKey != "")
 
 	debugLog.Printf("Gateway auth: apiKeyConfigured=%v", apiKey != "")

internal/config/config_core.go

@@ -91,6 +91,14 @@ type GatewayConfig struct {
 	PayloadSizeThreshold int `toml:"payload_size_threshold" json:"payload_size_threshold,omitempty"`
 }
 
+// GetAPIKey returns the gateway API key, handling a nil Gateway safely.
+func (c *Config) GetAPIKey() string {
+	if c.Gateway == nil {
+		return ""
+	}
+	return c.Gateway.APIKey
+}
+
 // ServerConfig represents an individual MCP server configuration.
 type ServerConfig struct {
 	// Type is the server type: "stdio" or "http"

internal/config/config_test.go

@@ -1609,3 +1609,21 @@ func TestApplyGatewayDefaults_OtherFieldsUnaffected(t *testing.T) {
 	// Defaults applied for the zero fields
 	assert.Equal(t, DefaultPort, cfg.Port)
 }
+
+// TestGetAPIKey verifies that GetAPIKey handles nil Gateway and returns the key when set.
+func TestGetAPIKey(t *testing.T) {
+	t.Run("nil Gateway returns empty string", func(t *testing.T) {
+		cfg := &Config{}
+		assert.Equal(t, "", cfg.GetAPIKey())
+	})
+
+	t.Run("Gateway with no key returns empty string", func(t *testing.T) {
+		cfg := &Config{Gateway: &GatewayConfig{}}
+		assert.Equal(t, "", cfg.GetAPIKey())
+	})
+
+	t.Run("Gateway with key returns key", func(t *testing.T) {
+		cfg := &Config{Gateway: &GatewayConfig{APIKey: "my-secret-key"}}
+		assert.Equal(t, "my-secret-key", cfg.GetAPIKey())
+	})
+}

internal/launcher/launcher.go

@@ -123,56 +123,12 @@ func GetOrLaunch(l *Launcher, serverID string) (*mcp.Connection, error) {
 	}
 
 	// stdio backends from this point
-	// Warn if using direct command in a container
-	isDirectCommand := serverCfg.Command != "docker"
-	if l.runningInContainer && isDirectCommand {
-		l.logSecurityWarning(serverID, serverCfg)
-	}
-
-	// Log the command being executed
-	l.logLaunchStart(serverID, "", serverCfg, isDirectCommand)
-
-	// Check for environment variable passthrough (only check args after -e flags)
-	l.logEnvPassthrough(serverCfg.Args)
-
-	if len(serverCfg.Env) > 0 {
-		log.Printf("[LAUNCHER] Additional env vars: %v", sanitize.TruncateSecretMap(serverCfg.Env))
-	}
-
-	log.Printf("[LAUNCHER] Starting server with %v timeout", l.startupTimeout)
-	logLauncher.Printf("Starting server with timeout: serverID=%s, timeout=%v", serverID, l.startupTimeout)
-
-	// Create a buffered channel to receive connection result
-	// Buffer size of 1 prevents goroutine leak if timeout occurs before connection completes
-	resultChan := make(chan connectionResult, 1)
-
-	logLauncher.Printf("Starting connection goroutine: serverID=%s", serverID)
-	// Launch connection in a goroutine
-	go func() {
-		conn, err := mcp.NewConnection(l.ctx, serverID, serverCfg.Command, serverCfg.Args, serverCfg.Env)
-		resultChan <- connectionResult{conn, err}
-	}()
-
-	// Wait for connection with timeout
-	select {
-	case result := <-resultChan:
-		conn, err := result.conn, result.err
-		if err != nil {
-			// Enhanced error logging for command-based servers
-			l.logLaunchError(serverID, "", err, serverCfg, isDirectCommand)
-			return nil, fmt.Errorf("failed to create connection: %w", err)
-		}
-
-		l.logLaunchSuccess(serverID, "")
-
-		l.connections[serverID] = conn
-		return conn, nil
-
-	case <-time.After(l.startupTimeout):
-		// Timeout occurred
-		l.logTimeoutError(serverID, "")
-		return nil, fmt.Errorf("server startup timeout after %v", l.startupTimeout)
+	conn, err := l.launchStdioConnection(serverID, "", serverCfg)
+	if err != nil {
+		return nil, err
 	}
+	l.connections[serverID] = conn
+	return conn, nil
 }
 
 // GetOrLaunchForSession returns a session-aware connection or launches a new one
@@ -219,6 +175,22 @@ func GetOrLaunchForSession(l *Launcher, serverID, sessionID string) (*mcp.Connec
 		return conn, nil
 	}
 
+	conn, err := l.launchStdioConnection(serverID, sessionID, serverCfg)
+	if err != nil {
+		// Record error in session pool
+		l.sessionPool.RecordError(serverID, sessionID)
+		return nil, err
+	}
+
+	// Add to session pool
+	l.sessionPool.Set(serverID, sessionID, conn)
+	return conn, nil
+}
+
+// launchStdioConnection creates a new stdio connection using a goroutine+timeout pattern.
+// It handles the security warning, launch logging, and the buffered-channel/select timeout.
+// Returns the raw *mcp.Connection on success; the caller is responsible for storing it.
+func (l *Launcher) launchStdioConnection(serverID, sessionID string, serverCfg *config.ServerConfig) (*mcp.Connection, error) {
 	// Warn if using direct command in a container
 	isDirectCommand := serverCfg.Command != "docker"
 	if l.runningInContainer && isDirectCommand {
@@ -235,13 +207,14 @@ func GetOrLaunchForSession(l *Launcher, serverID, sessionID string) (*mcp.Connec
 		log.Printf("[LAUNCHER] Additional env vars: %v", sanitize.TruncateSecretMap(serverCfg.Env))
 	}
 
-	log.Printf("[LAUNCHER] Starting server for session with %v timeout", l.startupTimeout)
-	logLauncher.Printf("Starting session server with timeout: serverID=%s, sessionID=%s, timeout=%v", serverID, sessionID, l.startupTimeout)
+	log.Printf("[LAUNCHER] Starting server with %v timeout", l.startupTimeout)
+	logLauncher.Printf("Starting server with timeout: serverID=%s, sessionID=%s, timeout=%v", serverID, sessionID, l.startupTimeout)
 
 	// Create a buffered channel to receive connection result
 	// Buffer size of 1 prevents goroutine leak if timeout occurs before connection completes
 	resultChan := make(chan connectionResult, 1)
 
+	logLauncher.Printf("Starting connection goroutine: serverID=%s", serverID)
 	// Launch connection in a goroutine
 	go func() {
 		conn, err := mcp.NewConnection(l.ctx, serverID, serverCfg.Command, serverCfg.Args, serverCfg.Env)
@@ -254,24 +227,13 @@ func GetOrLaunchForSession(l *Launcher, serverID, sessionID string) (*mcp.Connec
 		conn, err := result.conn, result.err
 		if err != nil {
 			l.logLaunchError(serverID, sessionID, err, serverCfg, isDirectCommand)
-
-			// Record error in session pool
-			l.sessionPool.RecordError(serverID, sessionID)
-
 			return nil, fmt.Errorf("failed to create connection: %w", err)
 		}
-
 		l.logLaunchSuccess(serverID, sessionID)
-
-		// Add to session pool
-		l.sessionPool.Set(serverID, sessionID, conn)
 		return conn, nil
 
 	case <-time.After(l.startupTimeout):
-		// Timeout occurred
 		l.logTimeoutError(serverID, sessionID)
-		// Record error in session pool before returning
-		l.sessionPool.RecordError(serverID, sessionID)
 		return nil, fmt.Errorf("server startup timeout after %v", l.startupTimeout)
 	}
 }

internal/server/handlers.go

@@ -2,7 +2,6 @@ package server
 
 import (
 	"context"
-	"encoding/json"
 	"log"
 	"net/http"
 	"os"
@@ -51,9 +50,7 @@ func handleClose(unifiedServer *UnifiedServer) http.Handler {
 		if unifiedServer.IsShutdown() {
 			logHandlers.Print("Gateway already shutdown, returning 410 Gone")
 			logger.LogWarn("shutdown", "Close endpoint called but gateway already closed, remote=%s", r.RemoteAddr)
-			w.Header().Set("Content-Type", "application/json")
-			w.WriteHeader(http.StatusGone) // 410 Gone
-			json.NewEncoder(w).Encode(map[string]interface{}{
+			writeJSONResponse(w, http.StatusGone, map[string]interface{}{
 				"error": "Gateway has already been closed",
 			})
 			return
@@ -65,14 +62,11 @@ func handleClose(unifiedServer *UnifiedServer) http.Handler {
 		logHandlers.Printf("Shutdown completed: servers_terminated=%d", serversTerminated)
 
 		// Return success response (spec 5.1.3)
-		w.Header().Set("Content-Type", "application/json")
-		w.WriteHeader(http.StatusOK)
-		response := map[string]interface{}{
+		writeJSONResponse(w, http.StatusOK, map[string]interface{}{
 			"status":            "closed",
 			"message":           "Gateway shutdown initiated",
 			"serversTerminated": serversTerminated,
-		}
-		json.NewEncoder(w).Encode(response)
+		})
 
 		logger.LogInfo("shutdown", "Close endpoint response sent, servers_terminated=%d", serversTerminated)
 		log.Printf("Gateway shutdown initiated. Terminated %d server(s)", serversTerminated)

internal/server/health.go

@@ -1,7 +1,6 @@
 package server
 
 import (
-	"encoding/json"
 	"net/http"
 
 	"github.com/github/gh-aw-mcpg/internal/logger"
@@ -52,11 +51,8 @@ func HandleHealth(unifiedServer *UnifiedServer) http.HandlerFunc {
 	return func(w http.ResponseWriter, r *http.Request) {
 		logHealth.Printf("Health check request: method=%s, remote=%s", r.Method, r.RemoteAddr)
 
-		w.Header().Set("Content-Type", "application/json")
-		w.WriteHeader(http.StatusOK)
-
 		response := BuildHealthResponse(unifiedServer)
 		logHealth.Printf("Health response: status=%s, servers=%d", response.Status, len(response.Servers))
-		json.NewEncoder(w).Encode(response)
+		writeJSONResponse(w, http.StatusOK, response)
 	}
 }

internal/server/http_helpers.go

@@ -3,6 +3,7 @@ package server
 import (
 	"bytes"
 	"context"
+	"encoding/json"
 	"io"
 	"log"
 	"net/http"
@@ -15,6 +16,14 @@ import (
 
 var logHelpers = logger.New("server:helpers")
 
+// writeJSONResponse sets the Content-Type header, writes the status code, and encodes
+// body as JSON. It centralises the three-line pattern used across HTTP handlers.
+func writeJSONResponse(w http.ResponseWriter, statusCode int, body interface{}) {
+	w.Header().Set("Content-Type", "application/json")
+	w.WriteHeader(statusCode)
+	json.NewEncoder(w).Encode(body)
+}
+
 // withResponseLogging wraps an http.Handler to log response bodies
 func withResponseLogging(handler http.Handler) http.Handler {
 	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {

internal/server/routed.go

@@ -22,9 +22,7 @@ func rejectIfShutdown(unifiedServer *UnifiedServer, next http.Handler, logNamesp
 		if unifiedServer.IsShutdown() {
 			log.Printf("Rejecting request during shutdown: remote=%s, method=%s, path=%s", r.RemoteAddr, r.Method, r.URL.Path)
 			logger.LogWarn("shutdown", "Request rejected during shutdown, remote=%s, path=%s", r.RemoteAddr, r.URL.Path)
-			w.Header().Set("Content-Type", "application/json")
-			w.WriteHeader(http.StatusServiceUnavailable)
-			w.Write([]byte(shutdownErrorJSON))
+			writeJSONResponse(w, http.StatusServiceUnavailable, map[string]string{"error": "Gateway is shutting down"})
 			return
 		}
 		next.ServeHTTP(w, r)