vecna/pkg/server/server.go

package server

import (
	"fmt"
	"net/http"
	"strings"
	"time"

	"github.com/prometheus/client_golang/prometheus/promhttp"
	"github.com/uptrace/bunrouter"
	"go.uber.org/zap"

	"github.com/Warky-Devs/vecna.git/pkg/adapter"
	"github.com/Warky-Devs/vecna.git/pkg/config"
	"github.com/Warky-Devs/vecna.git/pkg/embedclient"
	"github.com/Warky-Devs/vecna.git/pkg/forwardproxy"
	"github.com/Warky-Devs/vecna.git/pkg/metrics"
	"github.com/Warky-Devs/vecna.git/pkg/server/spec"
)

// New builds and returns a configured bunrouter.Router.
// Returns an error if route registration panics (e.g. conflicting routes).
func New(
	cfg *config.Config,
	clients map[string]embedclient.Client,
	adp adapter.Adapter,
	extraMaps map[string]ExtraMap,
	forwarders map[string]*forwardproxy.Router,
	reg *metrics.Registry,
	logger *zap.Logger,
) (router *bunrouter.Router, err error) {
	defer func() {
		if r := recover(); r != nil {
			logger.Error("panic during router setup", zap.Any("recover", r))
			err = fmt.Errorf("router setup panic: %v", r)
		}
	}()

	router = bunrouter.New(
		bunrouter.WithMiddleware(recoverMiddleware(logger)),
		bunrouter.WithMiddleware(traceMiddleware()),
		bunrouter.WithMiddleware(metricsMiddleware(reg, adp)),
		bunrouter.WithMiddleware(loggingMiddleware(logger)),
	)

	h := &handler{cfg: cfg, clients: clients, adapter: adp, extraMaps: extraMaps, forwarders: forwarders, logger: logger}

	// Public routes — no authentication required.
	router.GET("/", spec.DocsHandler())
	router.GET("/docs", spec.DocsHandler())
	router.GET("/openapi.yaml", spec.SpecHandler())

	// All API routes require authentication.
	authed := router.NewGroup("", bunrouter.WithMiddleware(authMiddleware(cfg.Server.APIKeys)))

	authed.POST("/v1/embeddings", h.openAIEmbeddings)
	authed.GET("/v1/embeddings", h.openAIEmbeddingsGet)
	// Google API uses a literal colon as a method separator (e.g. /v1/models/foo:embedContent).
	// bunrouter can't distinguish two routes with the same :param prefix, so a single wildcard
	// captures the full "model:action" segment and dispatches internally.
	authed.POST("/v1/models/*modelaction", h.googleDispatch)

	// Extra-map routes: /map/:mapping/v1/... uses the adapter configured under extra_maps[mapping].
	authed.POST("/map/:mapping/v1/embeddings", h.openAIEmbeddingsMapped)
	authed.GET("/map/:mapping/v1/embeddings", h.openAIEmbeddingsGetMapped)
	authed.POST("/map/:mapping/v1/models/*modelaction", h.googleDispatchMapped)

	// Generic OpenAI-compatible passthrough — forwarded verbatim to the resolved
	// target with no vecna-specific processing (no dimension adaptation). The
	// target is chosen the same way as embeddings: the body's "model" field must
	// name a configured forward target, falling back to forward.default.
	registerPassthroughRoutes(authed, h)

	// Metrics — only when enabled.
	// If metrics.api_key is set, routes are guarded by that key (on the authed group).
	// If metrics.api_key is blank, routes are public — no auth headers checked at all.
	if cfg.Metrics.Enabled {
		metricsHandler := promhttp.HandlerFor(reg.Prometheus(), promhttp.HandlerOpts{})
		path := cfg.Metrics.Path
		if path == "" {
			path = "/metrics"
		}
		dash := dashboardHandler(reg)
		if cfg.Metrics.APIKey != "" {
			authed.GET(path, metricsAuthHandler(cfg.Metrics.APIKey, metricsHandler))
			authed.GET("/dashboard", metricsKeyMiddleware(cfg.Metrics.APIKey, dash))
		} else {
			router.GET(path, func(w http.ResponseWriter, req bunrouter.Request) error {
				metricsHandler.ServeHTTP(w, req.Request)
				return nil
			})
			router.GET("/dashboard", dash)
		}
	}

	return router, nil
}

// recoverMiddleware catches panics from any handler or middleware below it.
// Without this, an unrecovered panic kills the request goroutine mid-response
// and the client sees a raw connection reset instead of an error; with it,
// the client gets a normal 500 JSON error (or, if the panic happened after
// the response was already committed — e.g. mid-stream — the connection is
// still cut, but the panic is at least logged instead of silently crashing).
func recoverMiddleware(logger *zap.Logger) bunrouter.MiddlewareFunc {
	return func(next bunrouter.HandlerFunc) bunrouter.HandlerFunc {
		return func(w http.ResponseWriter, req bunrouter.Request) (err error) {
			rw := &recoverWriter{ResponseWriter: w}
			defer func() {
				r := recover()
				if r == nil {
					return
				}
				logger.Error("panic recovered",
					zap.Any("recover", r),
					zap.String("method", req.Method),
					zap.String("path", req.URL.Path),
					zap.Stack("stack"),
				)
				if rw.wroteHeader {
					err = fmt.Errorf("panic after response started: %v", r)
					return
				}
				err = writeError(rw, http.StatusInternalServerError, "internal server error")
			}()
			return next(rw, req)
		}
	}
}

// recoverWriter tracks whether a response has already been committed, so
// recoverMiddleware knows whether it's still safe to write an error body.
type recoverWriter struct {
	http.ResponseWriter
	wroteHeader bool
}

func (rw *recoverWriter) WriteHeader(code int) {
	rw.wroteHeader = true
	rw.ResponseWriter.WriteHeader(code)
}

func (rw *recoverWriter) Write(b []byte) (int, error) {
	rw.wroteHeader = true
	return rw.ResponseWriter.Write(b)
}

func (rw *recoverWriter) Flush() {
	if f, ok := rw.ResponseWriter.(http.Flusher); ok {
		f.Flush()
	}
}

// authMiddleware rejects requests without a valid API key when api_keys is
// configured. The key may be supplied as a Bearer token (OpenAI-style), or
// via the X-Api-Key or Api-Key headers (Anthropic/Azure OpenAI-style), so
// clients built against any of those conventions work unmodified.
func authMiddleware(apiKeys []string) bunrouter.MiddlewareFunc {
	if len(apiKeys) == 0 {
		return func(next bunrouter.HandlerFunc) bunrouter.HandlerFunc { return next }
	}
	keySet := make(map[string]struct{}, len(apiKeys))
	for _, k := range apiKeys {
		keySet[k] = struct{}{}
	}
	return func(next bunrouter.HandlerFunc) bunrouter.HandlerFunc {
		return func(w http.ResponseWriter, req bunrouter.Request) error {
			token := requestAPIKey(req.Request)
			if _, ok := keySet[token]; !ok || token == "" {
				return writeError(w, http.StatusUnauthorized, "unauthorized")
			}
			return next(w, req)
		}
	}
}

// requestAPIKey extracts the client-supplied API key, checking (in order)
// the Authorization Bearer scheme, X-Api-Key, and Api-Key headers.
func requestAPIKey(req *http.Request) string {
	if auth := req.Header.Get("Authorization"); auth != "" {
		return strings.TrimPrefix(auth, "Bearer ")
	}
	if key := req.Header.Get("X-Api-Key"); key != "" {
		return key
	}
	return req.Header.Get("Api-Key")
}

// traceMiddleware injects a *RequestTrace into every request context.
func traceMiddleware() bunrouter.MiddlewareFunc {
	return func(next bunrouter.HandlerFunc) bunrouter.HandlerFunc {
		return func(w http.ResponseWriter, req bunrouter.Request) error {
			ctx := WithTrace(req.Context())
			return next(w, req.WithContext(ctx))
		}
	}
}

// metricsMiddleware records Prometheus observations after the handler returns.
func metricsMiddleware(reg *metrics.Registry, adp adapter.Adapter) bunrouter.MiddlewareFunc {
	if reg == nil {
		return func(next bunrouter.HandlerFunc) bunrouter.HandlerFunc { return next }
	}
	adpType := fmt.Sprintf("%T", adp)
	return reg.Middleware(func(req bunrouter.Request) metrics.TraceSnapshot {
		t := TraceFromContext(req.Context())
		total := time.Since(t.Start)
		return metrics.TraceSnapshot{
			TotalSeconds:     total.Seconds(),
			ForwardSeconds:   t.ForwardDuration.Seconds(),
			TranslateSeconds: t.TranslateDuration.Seconds(),
			ForwardTarget:    t.ForwardTarget,
			ForwardURL:       t.ForwardURL,
			ForwardModel:     t.ForwardModel,
			AdapterType:      adpType,
			PromptTokens:     t.PromptTokens,
			TotalTokens:      t.TotalTokens,
		}
	})
}

// loggingMiddleware logs method, path, status, and timing via zap.
func loggingMiddleware(logger *zap.Logger) bunrouter.MiddlewareFunc {
	return func(next bunrouter.HandlerFunc) bunrouter.HandlerFunc {
		return func(w http.ResponseWriter, req bunrouter.Request) error {
			sw := &statusWriter{ResponseWriter: w, status: http.StatusOK}
			err := next(sw, req)
			t := TraceFromContext(req.Context())
			total := time.Since(t.Start)

			logger.Info("request",
				zap.String("method", req.Method),
				zap.String("path", req.URL.Path),
				zap.Int("status", sw.status),
				zap.Int64("total_ms", total.Milliseconds()),
				zap.Int64("forward_ms", t.ForwardDuration.Milliseconds()),
				zap.Int64("translate_ms", t.TranslateDuration.Milliseconds()),
			)
			return err
		}
	}
}

// metricsKeyMiddleware guards a bunrouter.HandlerFunc with a dedicated API key
// (Bearer, X-Api-Key, or Api-Key — see requestAPIKey).
func metricsKeyMiddleware(apiKey string, h bunrouter.HandlerFunc) bunrouter.HandlerFunc {
	return func(w http.ResponseWriter, req bunrouter.Request) error {
		if requestAPIKey(req.Request) != apiKey {
			return writeError(w, http.StatusUnauthorized, "unauthorized")
		}
		return h(w, req)
	}
}

// metricsAuthHandler wraps a standard http.Handler with the same API key auth.
func metricsAuthHandler(apiKey string, h http.Handler) bunrouter.HandlerFunc {
	return func(w http.ResponseWriter, req bunrouter.Request) error {
		if requestAPIKey(req.Request) != apiKey {
			return writeError(w, http.StatusUnauthorized, "unauthorized")
		}
		h.ServeHTTP(w, req.Request)
		return nil
	}
}

// statusWriter captures the HTTP status code written by a handler.
type statusWriter struct {
	http.ResponseWriter
	status int
}

func (sw *statusWriter) WriteHeader(code int) {
	sw.status = code
	sw.ResponseWriter.WriteHeader(code)
}

func (sw *statusWriter) Flush() {
	if f, ok := sw.ResponseWriter.(http.Flusher); ok {
		f.Flush()
	}
}