homelab-frontend/internal/proxy/router.go

// Package proxy provides request routing and forwarding.
package proxy

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"net/url"

	"forgejo.riotpiao.com/rock/homelab-frontend/internal/config"
)

// RouteRequest determines which upstream should handle the request.
// For /v1/chat/completions, /v1/embeddings, and /v1/rerank, it uses body-based dispatch.
// For other routes, it looks up by path in the configured routes.
func (h *Handler) RouteRequest(r *http.Request) (*Route, error) {
	// For /v1/chat/completions, /v1/embeddings, /v1/rerank use body-based dispatch
	if r.Method == "POST" && (r.URL.Path == "/v1/chat/completions" || r.URL.Path == "/v1/embeddings" || r.URL.Path == "/v1/rerank") {
		return h.routeByModel(r, r.URL.Path)
	}

	// For other paths, try to find a matching route by path
	// Look for exact path match or path prefix match
	for _, route := range h.routes {
		// Check if route's pathRewrite matches request path
		if route.Upstream.PathRewrite != "" && route.Upstream.PathRewrite == r.URL.Path {
			return route, nil
		}
	}

	// If no routes configured, return error
	if len(h.routes) == 0 {
		return nil, fmt.Errorf("no routes configured")
	}

	// Return the first configured route as fallback
	for _, route := range h.routes {
		return route, nil
	}

	return nil, fmt.Errorf("no route available")
}

// routeByModel reads the request body to find the "model" field and routes accordingly.
// The body is preserved for forwarding to the upstream.
// Returns a modelValidationError for client errors (invalid JSON, missing/unknown model).
// The path parameter indicates which endpoint is being called (/v1/chat/completions, /v1/embeddings, /v1/rerank)
func (h *Handler) routeByModel(r *http.Request, path string) (*Route, error) {
	// If there's no body, we can't determine the model
	if r.Body == nil {
		return nil, &modelValidationError{
			Kind:    "missing_model",
			Message: "request body required",
		}
	}

	// Read the body to extract the model name
	// We need to be careful to preserve the body for the upstream
	bodyBytes, err := io.ReadAll(r.Body)
	if err != nil {
		return nil, &modelValidationError{
			Kind:    "invalid_request",
			Message: fmt.Sprintf("failed to read request body: %v", err),
		}
	}

	// Restore the body so it can be read again by the upstream
	r.Body = io.NopCloser(bytes.NewReader(bodyBytes))

	// Extract only the top-level "model" value with a token walk: the common
	// path must not materialise the whole body (messages and all) as a map.
	modelRaw, hasModel, err := extractTopLevelModel(bodyBytes)
	if err != nil {
		return nil, &modelValidationError{
			Kind:    "invalid_json",
			Message: "request body is not valid JSON",
		}
	}

	// Extract the model name
	if !hasModel {
		return nil, &modelValidationError{
			Kind:    "missing_model",
			Message: "'model' field is missing",
		}
	}

	// Handle null model
	if bytes.Equal(bytes.TrimSpace(modelRaw), []byte("null")) {
		return nil, &modelValidationError{
			Kind:    "null_model",
			Message: "'model' field is null",
		}
	}

	// Extract as string
	var modelName string
	if err := json.Unmarshal(modelRaw, &modelName); err != nil {
		return nil, &modelValidationError{
			Kind:    "missing_model",
			Message: "'model' field must be a string",
		}
	}

	// Handle empty string
	if modelName == "" {
		return nil, &modelValidationError{
			Kind:    "empty_model",
			Message: "'model' field cannot be empty",
		}
	}

	// Look up the model in the registry
	modelUpstream := h.config.LookupModel(modelName)
	if modelUpstream == nil {
		return nil, &modelValidationError{
			Kind:    "unknown_model",
			Message: fmt.Sprintf("unknown model: %q", modelName),
			Model:   modelName,
		}
	}

	// If upstream expects a different model name, rewrite the body. This path is
	// uncommon, so it may decode the body as raw fields; the common path above
	// does not build a full map.
	if modelUpstream.UpstreamModel != "" && modelUpstream.UpstreamModel != modelName {
		var fields map[string]json.RawMessage
		if err := json.Unmarshal(bodyBytes, &fields); err != nil {
			return nil, &modelValidationError{
				Kind:    "invalid_request",
				Message: fmt.Sprintf("failed to rewrite model name: %v", err),
			}
		}
		newModel, err := json.Marshal(modelUpstream.UpstreamModel)
		if err != nil {
			return nil, &modelValidationError{
				Kind:    "invalid_request",
				Message: fmt.Sprintf("failed to rewrite model name: %v", err),
			}
		}
		fields["model"] = newModel
		newBody, err := json.Marshal(fields)
		if err != nil {
			return nil, &modelValidationError{
				Kind:    "invalid_request",
				Message: fmt.Sprintf("failed to rewrite model name: %v", err),
			}
		}
		r.Body = io.NopCloser(bytes.NewReader(newBody))
		r.ContentLength = int64(len(newBody))
	}

	// Determine the upstream path based on the request path
	upstreamPath := path
	if path == "/v1/rerank" {
		// Rerank endpoint uses /rerank path on upstream
		upstreamPath = "/rerank"
	}

	// Create a route for this model with appropriate timeouts
	// These are sensible defaults for LLM models
	upstreamCfg := &config.Upstream{
		Address:        modelUpstream.Address,
		PathRewrite:    upstreamPath,
		ConnectTimeout: h.defaultConnectTimeout,
		ReadTimeout:    h.defaultReadTimeout,
		WriteTimeout:   h.defaultWriteTimeout,
		MaxBodySize:    h.maxBodyBytes,
		AuthRequired:   false,
	}

	targetURL := &url.URL{Scheme: "http", Host: modelUpstream.Address}

	route := &Route{
		Name:      "v1-chat-" + modelName,
		ModelName: modelName,
		Upstream:  upstreamCfg,
		Transport: h.getOrCreateTransport(modelUpstream.Address, upstreamCfg),
		Director: func(req *http.Request) {
			directorFunc(req, targetURL, upstreamCfg)
		},
	}

	return route, nil
}

// extractTopLevelModel walks a JSON object body with a streaming decoder and
// returns the raw value of the top-level "model" field. It never builds the
// full document as a map, so a large "messages" array is skipped by token.
// hasModel is false when the key is absent or the document is a bare null.
func extractTopLevelModel(body []byte) (model json.RawMessage, hasModel bool, err error) {
	dec := json.NewDecoder(bytes.NewReader(body))
	tok, err := dec.Token()
	if err != nil {
		return nil, false, err
	}
	if tok == nil {
		// A bare `null` body decoded to a nil map before, i.e. no model field.
		if _, err := dec.Token(); err != io.EOF {
			if err == nil {
				return nil, false, fmt.Errorf("unexpected data after JSON null")
			}
			return nil, false, err
		}
		return nil, false, nil
	}
	if delim, ok := tok.(json.Delim); !ok || delim != '{' {
		return nil, false, fmt.Errorf("request body is not a JSON object")
	}
	for dec.More() {
		keyTok, err := dec.Token()
		if err != nil {
			return nil, false, err
		}
		key, _ := keyTok.(string)
		if key == "model" {
			if err := dec.Decode(&model); err != nil {
				return nil, false, err
			}
			hasModel = true
			continue
		}
		if err := skipJSONValue(dec); err != nil {
			return nil, false, err
		}
	}
	// Consume the closing brace and reject trailing data, preserving the
	// invalid_json answer the old whole-body decode produced.
	if _, err := dec.Token(); err != nil {
		return nil, false, err
	}
	if _, err := dec.Token(); err != io.EOF {
		if err == nil {
			return nil, false, fmt.Errorf("unexpected data after JSON object")
		}
		return nil, false, err
	}
	return model, hasModel, nil
}

// skipJSONValue advances dec past the next JSON value without materialising it.
func skipJSONValue(dec *json.Decoder) error {
	tok, err := dec.Token()
	if err != nil {
		return err
	}
	delim, ok := tok.(json.Delim)
	if !ok {
		return nil // scalar
	}
	if delim != '{' && delim != '[' {
		return fmt.Errorf("unexpected delimiter %v", delim)
	}
	depth := 1
	for depth > 0 {
		tok, err := dec.Token()
		if err != nil {
			return err
		}
		if d, ok := tok.(json.Delim); ok {
			if d == '{' || d == '[' {
				depth++
			} else {
				depth--
			}
		}
	}
	return nil
}