Skip to content

Refactor: Consolidate shutdown middleware, centralize session utilities, remove deprecated wrappers #483

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.

Sign up for GitHub

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
lpcox merged 5 commits into from
Jan 27, 2026
Merged

Refactor: Consolidate shutdown middleware, centralize session utilities, remove deprecated wrappers #483

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
fef300b
Initial plan
Copilot Jan 26, 2026
c0c8122
Phase 1: Consolidate duplicate shutdown middleware functions
Copilot Jan 26, 2026
9a66ae5
Phase 2: Move session truncation utility to auth package
Copilot Jan 26, 2026
e24def0
Phase 3: Remove deprecated extractSessionFromAuth wrapper
Copilot Jan 26, 2026
94fb1e0
Fix formatting: Remove extra blank lines
Copilot Jan 26, 2026
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
13 changes: 13 additions & 0 deletions internal/auth/header.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -165,3 +165,16 @@ func ExtractSessionID(authHeader string) string {
log.Print("Using plain API key as session ID")
return authHeader
}

// TruncateSessionID returns a truncated session ID for safe logging (first 8 chars).
// Returns "(none)" for empty session IDs, and appends "..." for truncated values.
// This is useful for logging session IDs without exposing sensitive information.
func TruncateSessionID(sessionID string) string {
if sessionID == "" {
return "(none)"
}
if len(sessionID) <= 8 {
return sessionID
}
return sessionID[:8] + "..."
}
11 changes: 0 additions & 11 deletions internal/server/auth.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,6 @@ import (
"net/http"
"time"

"github.com/githubnext/gh-aw-mcpg/internal/auth"
"github.com/githubnext/gh-aw-mcpg/internal/logger"
)

Expand Down Expand Up @@ -42,16 +41,6 @@ func authMiddleware(apiKey string, next http.HandlerFunc) http.HandlerFunc {
}
}

// extractSessionFromAuth extracts session ID from Authorization header.
// This function delegates to auth.ExtractSessionID for consistent session ID extraction.
// Per spec 7.1: When API key is configured, Authorization contains plain API key.
// When API key is not configured, supports Bearer token for backward compatibility.
//
// Deprecated: Use auth.ExtractSessionID directly instead.
func extractSessionFromAuth(authHeader string) string {
return auth.ExtractSessionID(authHeader)
}

// logRuntimeError logs runtime errors to stdout per spec section 9.2
func logRuntimeError(errorType, detail string, r *http.Request, serverName *string) {
timestamp := time.Now().UTC().Format(time.RFC3339)
Expand Down
11 changes: 7 additions & 4 deletions internal/server/routed.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@ import (
"net/http"
"sync"

"github.com/githubnext/gh-aw-mcpg/internal/auth"
"github.com/githubnext/gh-aw-mcpg/internal/logger"
"github.com/githubnext/gh-aw-mcpg/internal/mcp"
sdk "github.com/modelcontextprotocol/go-sdk/mcp"
Expand All @@ -18,10 +19,12 @@ var logRouted = logger.New("server:routed")

// rejectIfShutdown is a middleware that rejects requests with HTTP 503 when gateway is shutting down
// Per spec 5.1.3: "Immediately reject any new RPC requests to /mcp/{server-name} endpoints with HTTP 503"
func rejectIfShutdown(unifiedServer *UnifiedServer, next http.Handler) http.Handler {
// The logNamespace parameter is used to create a logger for debug output specific to the call site.
func rejectIfShutdown(unifiedServer *UnifiedServer, next http.Handler, logNamespace string) http.Handler {
log := logger.New(logNamespace)
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if unifiedServer.IsShutdown() {
logRouted.Printf("Rejecting request during shutdown: remote=%s, method=%s, path=%s", r.RemoteAddr, r.Method, r.URL.Path)
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)
Expand Down Expand Up @@ -109,7 +112,7 @@ func CreateHTTPServerForRoutedMode(addr string, unifiedServer *UnifiedServer, ap
routeHandler := sdk.NewStreamableHTTPHandler(func(r *http.Request) *sdk.Server {
// Extract session ID from Authorization header
authHeader := r.Header.Get("Authorization")
sessionID := extractSessionFromAuth(authHeader)
sessionID := auth.ExtractSessionID(authHeader)

// Reject requests without Authorization header
if sessionID == "" {
Expand Down Expand Up @@ -156,7 +159,7 @@ func CreateHTTPServerForRoutedMode(addr string, unifiedServer *UnifiedServer, ap

// Apply shutdown check middleware (spec 5.1.3)
// This must come before auth to ensure shutdown takes precedence
shutdownHandler := rejectIfShutdown(unifiedServer, loggedHandler)
shutdownHandler := rejectIfShutdown(unifiedServer, loggedHandler, "server:routed")

// Apply auth middleware if API key is configured (spec 7.1)
finalHandler := shutdownHandler
Expand Down
26 changes: 8 additions & 18 deletions internal/server/sdk_logging.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@ import (
"strings"
"time"

"github.com/githubnext/gh-aw-mcpg/internal/auth"
"github.com/githubnext/gh-aw-mcpg/internal/logger"
)

Expand Down Expand Up @@ -45,12 +46,12 @@ func WithSDKLogging(handler http.Handler, mode string) http.Handler {

// Extract session info for logging context
authHeader := r.Header.Get("Authorization")
sessionID := extractSessionFromAuth(authHeader)
sessionID := auth.ExtractSessionID(authHeader)
mcpSessionID := r.Header.Get("Mcp-Session-Id")

// Log incoming request
logSDK.Printf(">>> SDK Request [%s] session=%s mcp-session=%s method=%s path=%s",
mode, truncateSession(sessionID), truncateSession(mcpSessionID), r.Method, r.URL.Path)
mode, auth.TruncateSessionID(sessionID), auth.TruncateSessionID(mcpSessionID), r.Method, r.URL.Path)

// Capture and log request body for POST requests
var requestBody []byte
Expand All @@ -66,7 +67,7 @@ func WithSDKLogging(handler http.Handler, mode string) http.Handler {
if err := json.Unmarshal(requestBody, &jsonrpcReq); err == nil {
logSDK.Printf(" JSON-RPC Request: method=%s id=%v", jsonrpcReq.Method, jsonrpcReq.ID)
logger.LogDebug("sdk-frontend", "JSON-RPC request parsed: mode=%s, method=%s, id=%v, session=%s",
mode, jsonrpcReq.Method, jsonrpcReq.ID, truncateSession(sessionID))
mode, jsonrpcReq.Method, jsonrpcReq.ID, auth.TruncateSessionID(sessionID))
} else {
logSDK.Printf(" Failed to parse JSON-RPC request: %v", err)
logSDK.Printf(" Raw body: %s", string(requestBody))
Expand Down Expand Up @@ -100,14 +101,14 @@ func WithSDKLogging(handler http.Handler, mode string) http.Handler {
strings.Contains(jsonrpcResp.Error.Message, "invalid during") {
logSDK.Printf(" ⚠️ PROTOCOL STATE ERROR DETECTED")
logSDK.Printf(" Request method was: %s", jsonrpcReq.Method)
logSDK.Printf(" Session ID: %s", truncateSession(sessionID))
logSDK.Printf(" MCP-Session-Id header: %s", truncateSession(mcpSessionID))
logSDK.Printf(" Session ID: %s", auth.TruncateSessionID(sessionID))
logSDK.Printf(" MCP-Session-Id header: %s", auth.TruncateSessionID(mcpSessionID))
logSDK.Printf(" This error indicates SDK's StreamableHTTPHandler created fresh protocol state")

logger.LogWarn("sdk-frontend",
"Protocol state error: mode=%s, method=%s, session=%s, mcp_session=%s, error=%q",
mode, jsonrpcReq.Method, truncateSession(sessionID),
truncateSession(mcpSessionID), jsonrpcResp.Error.Message)
mode, jsonrpcReq.Method, auth.TruncateSessionID(sessionID),
auth.TruncateSessionID(mcpSessionID), jsonrpcResp.Error.Message)
} else {
logger.LogError("sdk-frontend",
"JSON-RPC error: mode=%s, method=%s, code=%d, message=%q",
Expand Down Expand Up @@ -140,14 +141,3 @@ func WithSDKLogging(handler http.Handler, mode string) http.Handler {
}
})
}

// truncateSession returns a truncated session ID for logging (first 8 chars)
func truncateSession(s string) string {
if s == "" {
return "(none)"
}
if len(s) <= 8 {
return s
}
return s[:8] + "..."
}
21 changes: 3 additions & 18 deletions internal/server/transport.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,28 +7,13 @@ import (
"log"
"net/http"

"github.com/githubnext/gh-aw-mcpg/internal/auth"
"github.com/githubnext/gh-aw-mcpg/internal/logger"
sdk "github.com/modelcontextprotocol/go-sdk/mcp"
)

var logTransport = logger.New("server:transport")

// rejectIfShutdownUnified is a middleware that rejects requests with HTTP 503 when gateway is shutting down
// Per spec 5.1.3: "Immediately reject any new RPC requests to /mcp/{server-name} endpoints with HTTP 503"
func rejectIfShutdownUnified(unifiedServer *UnifiedServer, next http.Handler) http.Handler {
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if unifiedServer.IsShutdown() {
logTransport.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))
return
}
next.ServeHTTP(w, r)
})
}

// HTTPTransport wraps the SDK's HTTP transport
type HTTPTransport struct {
Addr string
Expand Down Expand Up @@ -91,7 +76,7 @@ func CreateHTTPServerForMCP(addr string, unifiedServer *UnifiedServer, apiKey st

// Extract session ID from Authorization header
authHeader := r.Header.Get("Authorization")
sessionID := extractSessionFromAuth(authHeader)
sessionID := auth.ExtractSessionID(authHeader)

// Reject requests without Authorization header
if sessionID == "" {
Expand Down Expand Up @@ -138,7 +123,7 @@ func CreateHTTPServerForMCP(addr string, unifiedServer *UnifiedServer, apiKey st

// Apply shutdown check middleware (spec 5.1.3)
// This must come before auth to ensure shutdown takes precedence
shutdownHandler := rejectIfShutdownUnified(unifiedServer, loggedHandler)
shutdownHandler := rejectIfShutdown(unifiedServer, loggedHandler, "server:transport")

// Apply auth middleware if API key is configured (spec 7.1)
finalHandler := shutdownHandler
Expand Down
Loading