ModelRouter/internal/gateway/keys.go

package gateway

import (
	"context"
	"encoding/json"
	"net/http"
	"strings"

	"llmsproxy/internal/config"
	"llmsproxy/internal/scheduler"
)

// handleKeysAPI manages gateway keys: GET /api/keys (admin: all keys),
// GET /api/keys/me (own key for any role), POST /api/keys (admin: create),
// PUT /api/keys/{key} (admin: update), DELETE /api/keys/{key} (admin: remove).
func (g *Gateway) handleKeysAPI(w http.ResponseWriter, r *http.Request) {
	path := strings.TrimPrefix(r.URL.Path, "/api/keys")
	path = strings.Trim(path, "/")
	role := reqRole(r.Context())

	// /api/keys/me/auto — a user's own chain, readable without admin rights.
	// Must be matched before the generic {key}/auto case below so it does not
	// ask admin rights of a plain user asking about their own chain.
	if path == "me/auto" {
		g.handleKeyMeAuto(w, r)
		return
	}
	if path == "me" {
		g.handleKeyMe(w, r)
		return
	}
	// /api/keys/{key}/auto — the per-key AUTO chain editor.
	if i := strings.LastIndex(path, "/auto"); i > 0 && path[i+len("/auto"):] == "" {
		g.handleKeyAutoAPI(w, r, path[:i])
		return
	}
	if role != "admin" {
		writeError(w, http.StatusForbidden, "forbidden", "admin role required")
		return
	}

	switch r.Method {
	case http.MethodGet:
		if path != "" {
			// A GET on a specific key is almost always a client that meant to
			// DELETE or PUT it but let fetch default to GET. Name the verbs
			// instead of only pointing back at the collection endpoint.
			writeError(w, http.StatusNotFound, "not_found",
				"no such endpoint; use GET /api/keys to list, DELETE /api/keys/{key} to remove, PUT /api/keys/{key} to update")
			return
		}
		writeJSON(w, http.StatusOK, map[string]interface{}{"keys": g.core.ListKeys()})
	case http.MethodPost:
		var body struct {
			Name   string              `json:"name"`
			Role   string              `json:"role"`
			Models []config.ModelScope `json:"models"`
			Note   string              `json:"note"`
		}
		if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
			writeError(w, http.StatusBadRequest, "invalid_request", "invalid json: "+err.Error())
			return
		}
		body.Role = config.NormalizeRole(body.Role)
		rec, err := g.core.CreateKey(body.Name, body.Role, body.Models, body.Note)
		if err != nil {
			writeError(w, http.StatusBadRequest, "key_error", err.Error())
			return
		}
		writeJSON(w, http.StatusOK, map[string]interface{}{"ok": true, "key": rec})
	case http.MethodPut, http.MethodPatch:
		if path == "" {
			writeError(w, http.StatusBadRequest, "invalid_request", "key required")
			return
		}
		var body struct {
			Name   string              `json:"name"`
			Role   string              `json:"role"`
			Models []config.ModelScope `json:"models"`
			Note   string              `json:"note"`
		}
		if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
			writeError(w, http.StatusBadRequest, "invalid_request", "invalid json: "+err.Error())
			return
		}
		// Each scope entry carries its own token/request caps, so replacing the
		// scope replaces the budgets with it — there is no separate key-wide
		// quota that could drift out of sync with the models.
		rec, err := g.core.UpdateKey(path, body.Name, body.Role, body.Models, body.Note)
		if err != nil {
			writeError(w, http.StatusBadRequest, "key_error", err.Error())
			return
		}
		writeJSON(w, http.StatusOK, map[string]interface{}{"ok": true, "key": rec})
	case http.MethodDelete:
		if path == "" {
			writeError(w, http.StatusBadRequest, "invalid_request", "key required")
			return
		}
		if path == reqKey(r.Context()) {
			writeError(w, http.StatusBadRequest, "invalid_request", "cannot delete the key you are logged in with")
			return
		}
		ok, err := g.core.DeleteKey(path)
		if err != nil {
			writeError(w, http.StatusBadRequest, "key_error", err.Error())
			return
		}
		if !ok {
			writeError(w, http.StatusNotFound, "not_found", "key not found")
			return
		}
		writeJSON(w, http.StatusOK, map[string]interface{}{"ok": true})
	default:
		writeError(w, http.StatusMethodNotAllowed, "method_not_allowed", "")
	}
}

// handleKeyMe returns the authenticated key's own record (users see only
// themselves; admins can use this as a convenience too).
func (g *Gateway) handleKeyMe(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodGet {
		writeError(w, http.StatusMethodNotAllowed, "method_not_allowed", "use GET")
		return
	}
	rec, ok := g.core.FindKey(reqKey(r.Context()))
	if !ok {
		writeError(w, http.StatusUnauthorized, "invalid_api_key", "key not found")
		return
	}
	if !rec.Seed {
		for _, s := range g.core.GatewayKeys() {
			if s == rec.Key {
				rec.Seed = true
				break
			}
		}
	}
	writeJSON(w, http.StatusOK, map[string]interface{}{"key": rec})
}

// allowedModels returns the model scope for the request's key; nil means
// unrestricted (admin keys and user keys without an explicit scope).
func (g *Gateway) allowedModels(ctx context.Context) []config.ModelScope {
	if reqRole(ctx) == "admin" {
		return nil
	}
	rec, ok := g.core.FindKey(reqKey(ctx))
	if !ok {
		return nil
	}
	return rec.Models
}

// handleKeyAutoAPI manages one key's own AUTO chain:
// GET /api/keys/{key}/auto returns it, PUT replaces it, DELETE clears the
// override so the key falls back to the global chain. Admin only — a user can
// inspect their own chain through GET /api/keys/me/auto.
//
// An empty chain and a cleared chain are deliberately the same state: "inherit
// the global chain". That keeps the WebUI's "use global" toggle a plain DELETE
// with no sentinel value to carry through config.yaml.
func (g *Gateway) handleKeyAutoAPI(w http.ResponseWriter, r *http.Request, key string) {
	if reqRole(r.Context()) != "admin" {
		writeError(w, http.StatusForbidden, "forbidden", "admin role required")
		return
	}
	rec, ok := g.core.FindKey(key)
	if !ok {
		writeError(w, http.StatusNotFound, "not_found", "key not found")
		return
	}
	switch r.Method {
	case http.MethodGet:
		writeJSON(w, http.StatusOK, map[string]interface{}{
			"key":      key,
			"auto":     rec.Auto,
			"inherits": !rec.HasOwnAuto(),
		})
	case http.MethodPut, http.MethodPatch:
		var body struct {
			Auto []config.ModelScope `json:"auto"`
		}
		if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
			writeError(w, http.StatusBadRequest, "invalid_request", "invalid json: "+err.Error())
			return
		}
		if len(body.Auto) == 0 {
			// Treat an empty PUT as "give up the override" so the endpoint
			// cannot persist a chain that would 503 with no slots.
			if err := g.core.SaveKeyAuto(key, nil); err != nil {
				writeError(w, http.StatusBadRequest, "key_error", err.Error())
				return
			}
			writeJSON(w, http.StatusOK, map[string]interface{}{"ok": true, "inherits": true})
			return
		}
		if err := g.core.SaveKeyAuto(key, body.Auto); err != nil {
			writeError(w, http.StatusBadRequest, "key_error", err.Error())
			return
		}
		// Report whether the chain actually resolved to usable slots. An
		// admin who typed a model that no longer exists should learn it here,
		// not from the user's next 503.
		chain, _ := g.core.AutoChainFor(key)
		writeJSON(w, http.StatusOK, map[string]interface{}{
			"ok": true, "inherits": false, "slots": chainSlots(chain),
		})
	case http.MethodDelete:
		if err := g.core.SaveKeyAuto(key, nil); err != nil {
			writeError(w, http.StatusBadRequest, "key_error", err.Error())
			return
		}
		writeJSON(w, http.StatusOK, map[string]interface{}{"ok": true, "inherits": true})
	default:
		writeError(w, http.StatusMethodNotAllowed, "method_not_allowed", "")
	}
}

// chainSlots counts the usable slots in a chain, for API feedback after an
// edit. A chain that compiled to zero slots is reported so the caller can warn
// instead of letting the next request fail with no_provider.
func chainSlots(ch *scheduler.Chain) int {
	if ch == nil {
		return 0
	}
	n := 0
	for _, tn := range ch.Tiers {
		n += len(tn.Slots)
	}
	return n
}

// handleKeyMeAuto lets a user read their own AUTO chain without admin rights.
func (g *Gateway) handleKeyMeAuto(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodGet {
		writeError(w, http.StatusMethodNotAllowed, "method_not_allowed", "use GET")
		return
	}
	rec, ok := g.core.FindKey(reqKey(r.Context()))
	if !ok {
		writeError(w, http.StatusUnauthorized, "invalid_api_key", "key not found")
		return
	}
	writeJSON(w, http.StatusOK, map[string]interface{}{
		"key": rec.Key, "auto": rec.Auto, "inherits": !rec.HasOwnAuto(),
	})
}

// handleAutoAPI manages the AUTO scheduling slots: GET /api/auto returns the
// current rules; PUT /api/auto replaces them (admin only).
func (g *Gateway) handleAutoAPI(w http.ResponseWriter, r *http.Request) {
	if r.Method == http.MethodGet {
		writeJSON(w, http.StatusOK, map[string]interface{}{
			"rules":       g.core.AutoRules(),
			"image_rules": g.core.AutoImageRules(),
			"states":      g.core.AutoSlotStates(),
		})
		return
	}
	if reqRole(r.Context()) != "admin" {
		writeError(w, http.StatusForbidden, "forbidden", "admin role required")
		return
	}
	switch r.Method {
	case http.MethodPut, http.MethodPost:
		var body struct {
			Rules      []config.ModelScope `json:"rules"`
			ImageRules []config.ModelScope `json:"image_rules"`
		}
		if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
			writeError(w, http.StatusBadRequest, "invalid_request", "invalid json: "+err.Error())
			return
		}
		if body.Rules != nil {
			if err := g.core.SaveAutoRules(body.Rules); err != nil {
				writeError(w, http.StatusBadRequest, "auto_error", err.Error())
				return
			}
		}
		if body.ImageRules != nil {
			if err := g.core.SaveAutoImageRules(body.ImageRules); err != nil {
				writeError(w, http.StatusBadRequest, "auto_error", err.Error())
				return
			}
		}
		writeJSON(w, http.StatusOK, map[string]interface{}{
			"ok":          true,
			"rules":       g.core.AutoRules(),
			"image_rules": g.core.AutoImageRules(),
		})
	default:
		writeError(w, http.StatusMethodNotAllowed, "method_not_allowed", "")
	}
}