ai-gateway-go/internal/platform/httpserver/server.go

package httpserver

import (
	"crypto/rand"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"log/slog"
	"net/http"
	"runtime/debug"
	"sync/atomic"
	"time"

	"aigateway.local/core/internal/gateway"
	"aigateway.local/core/internal/platform/config"
	"aigateway.local/core/internal/platform/health"
)

type Dependencies struct {
	Config        config.Config
	Logger        *slog.Logger
	Checker       health.Checker
	Gateway       http.Handler
	Control       http.Handler
	Version       string
	StartedAt     time.Time
	BootstrapUses func() uint64
	ExtraMetrics  func() string
}

type metrics struct {
	requests atomic.Uint64
	panics   atomic.Uint64
}

func New(dependencies Dependencies) *http.Server {
	mux := http.NewServeMux()
	stats := &metrics{}
	mux.HandleFunc("GET /healthz", func(writer http.ResponseWriter, _ *http.Request) {
		writeJSON(writer, http.StatusOK, map[string]any{"status": "ok", "version": dependencies.Version})
	})
	mux.HandleFunc("GET /readyz", func(writer http.ResponseWriter, request *http.Request) {
		report := dependencies.Checker.Readiness(request.Context())
		status := http.StatusOK
		if !report.Ready {
			status = http.StatusServiceUnavailable
		}
		writeJSON(writer, status, report)
	})
	mux.HandleFunc("GET /metrics", func(writer http.ResponseWriter, _ *http.Request) {
		writer.Header().Set("Content-Type", "text/plain; version=0.0.4")
		bootstrapUses := uint64(0)
		if dependencies.BootstrapUses != nil {
			bootstrapUses = dependencies.BootstrapUses()
		}
		_, _ = fmt.Fprintf(writer, "gateway_http_requests_total %d\ngateway_http_panics_total %d\ngateway_bootstrap_api_key_uses_total %d\ngateway_uptime_seconds %.0f\n", stats.requests.Load(), stats.panics.Load(), bootstrapUses, time.Since(dependencies.StartedAt).Seconds())
		if dependencies.ExtraMetrics != nil {
			_, _ = fmt.Fprint(writer, dependencies.ExtraMetrics())
		}
	})
	mux.Handle("/v1/", dependencies.Gateway)
	if dependencies.Control != nil {
		mux.Handle("/api/", dependencies.Control)
	}

	handler := recoverMiddleware(dependencies.Logger, stats, requestIDMiddleware(stats, mux))
	return &http.Server{
		Addr:              dependencies.Config.Server.Address,
		Handler:           handler,
		ReadHeaderTimeout: dependencies.Config.Server.ReadHeaderTimeout,
		IdleTimeout:       dependencies.Config.Server.IdleTimeout,
		// WriteTimeout intentionally remains zero: SSE responses may be long lived.
	}
}

func requestIDMiddleware(stats *metrics, next http.Handler) http.Handler {
	return http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
		stats.requests.Add(1)
		requestID := request.Header.Get("X-Request-ID")
		if requestID == "" || len(requestID) > 128 {
			requestID = newRequestID()
		}
		writer.Header().Set("X-Request-ID", requestID)
		next.ServeHTTP(writer, request.WithContext(gateway.WithRequestID(request.Context(), requestID)))
	})
}

func recoverMiddleware(logger *slog.Logger, stats *metrics, next http.Handler) http.Handler {
	return http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
		defer func() {
			if recovered := recover(); recovered != nil {
				stats.panics.Add(1)
				logger.Error("http handler panic", "request_id", gateway.RequestID(request.Context()), "panic", recovered, "stack", string(debug.Stack()))
				writeJSON(writer, http.StatusInternalServerError, map[string]string{"error": "internal server error"})
			}
		}()
		next.ServeHTTP(writer, request)
	})
}

func writeJSON(writer http.ResponseWriter, status int, value any) {
	writer.Header().Set("Content-Type", "application/json")
	writer.WriteHeader(status)
	_ = json.NewEncoder(writer).Encode(value)
}

func newRequestID() string {
	buffer := make([]byte, 12)
	if _, err := rand.Read(buffer); err != nil {
		return fmt.Sprintf("req_%x", time.Now().UnixNano())
	}
	return "req_" + hex.EncodeToString(buffer)
}