ModelRouter/internal/gateway/server.go

// Package gateway exposes an OpenAI-compatible HTTP API over the provider
// registry: POST /v1/chat/completions (SDK + SSE), POST /v1/images/generations,
// GET /v1/models, protected by shared gateway API keys, plus a web UI and
// management API for adapters and sources.
package gateway

import (
	"embed"
	"encoding/json"
	"io/fs"
	"log"
	"net/http"
	"strings"

	"llmsproxy/internal/core"
)

//go:embed ui/*
var uiFS embed.FS

// Gateway is the HTTP handler for the OpenAI-compatible endpoint + web UI.
type Gateway struct {
	core    *core.Core
	apiKeys map[string]bool
	ui      http.Handler
}

func New(c *core.Core, gatewayKeys []string) (*Gateway, error) {
	keys := map[string]bool{}
	for _, k := range gatewayKeys {
		if k != "" {
			keys[k] = true
		}
	}
	sub, err := fs.Sub(uiFS, "ui")
	if err != nil {
		return nil, err
	}
	return &Gateway{
		core:    c,
		apiKeys: keys,
		ui:      http.FileServer(http.FS(sub)),
	}, nil
}

func (g *Gateway) Handler() http.Handler {
	return g.auth(http.HandlerFunc(g.routes))
}

func (g *Gateway) routes(w http.ResponseWriter, r *http.Request) {
	switch {
	case r.URL.Path == "/v1/chat/completions":
		g.handleChat(w, r)
	case r.URL.Path == "/v1/images/generations":
		g.handleImage(w, r)
	case r.URL.Path == "/v1/models":
		g.handleModels(w, r)
	case r.URL.Path == "/api/adapters" || strings.HasPrefix(r.URL.Path, "/api/adapters/"):
		g.handleAdaptersAPI(w, r)
	case r.URL.Path == "/api/sources" || strings.HasPrefix(r.URL.Path, "/api/sources/"):
		g.handleSourcesAPI(w, r)
	case r.URL.Path == "/api/status":
		g.handleStatusAPI(w, r)
	default:
		g.serveUI(w, r)
	}
}

func (g *Gateway) serveUI(w http.ResponseWriter, r *http.Request) {
	// serve index.html directly for the root path (FileServer would 301 it)
	if r.URL.Path == "/" || r.URL.Path == "/ui" {
		data, err := uiFS.ReadFile("ui/index.html")
		if err != nil {
			http.Error(w, "ui missing", http.StatusInternalServerError)
			return
		}
		w.Header().Set("Content-Type", "text/html; charset=utf-8")
		w.Write(data)
		return
	}
	g.ui.ServeHTTP(w, r)
}

func (g *Gateway) auth(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if len(g.apiKeys) == 0 {
			next.ServeHTTP(w, r)
			return
		}
		key := ""
		if h := r.Header.Get("Authorization"); h != "" {
			parts := strings.SplitN(h, " ", 2)
			if len(parts) == 2 && strings.EqualFold(parts[0], "Bearer") {
				key = parts[1]
			}
		}
		if key == "" {
			key = r.URL.Query().Get("api_key")
		}
		if !g.apiKeys[key] {
			writeError(w, http.StatusUnauthorized, "invalid_api_key", "invalid gateway api key")
			return
		}
		next.ServeHTTP(w, r)
	})
}

func (g *Gateway) handleModels(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodGet {
		writeError(w, http.StatusMethodNotAllowed, "method_not_allowed", "use GET")
		return
	}
	models := g.core.Registry().ModelList()
	type modelObj struct {
		ID     string `json:"id"`
		Object string `json:"object"`
	}
	objs := make([]modelObj, 0, len(models))
	for _, m := range models {
		objs = append(objs, modelObj{ID: m, Object: "model"})
	}
	writeJSON(w, http.StatusOK, map[string]interface{}{
		"object": "list",
		"data":   objs,
	})
}

func (g *Gateway) handleStatusAPI(w http.ResponseWriter, r *http.Request) {
	writeJSON(w, http.StatusOK, map[string]interface{}{
		"default_model": g.core.DefaultModel(),
		"models":        g.core.Registry().ModelList(),
		"sources":       g.core.Registry().Status(),
		"adapters":      g.core.ListAdapters(),
	})
}

func writeError(w http.ResponseWriter, code int, errType, msg string) {
	writeJSON(w, code, map[string]interface{}{
		"error": map[string]interface{}{"type": errType, "message": msg},
	})
}

func writeJSON(w http.ResponseWriter, code int, v interface{}) {
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(code)
	if err := json.NewEncoder(w).Encode(v); err != nil {
		log.Printf("[gateway] write json: %v", err)
	}
}