llmgw/terraform-provider-llmgw/internal/fakeapi/fakeapi.go

// Package fakeapi is an admin API in memory, for testing the client and the
// provider without a cluster.
//
// It does what the real one does at the edges the provider depends on: it
// demands the token, refuses a change not sent as JSON, fills in the defaults
// the catalog fills in, resolves a scope named by path to its ID, and answers
// with the same shapes and statuses. It is not the gateway; the OpenAPI
// contract test is what ties the client to that.
package fakeapi

import (
	"encoding/json"
	"fmt"
	"net/http"
	"slices"
	"sort"
	"strings"
	"sync"
	"time"

	"gitea.mynodes.xyz/ilker/llmgw/terraform-provider-llmgw/internal/client"
)

// API is the fake. Its fields may be read by a test while it serves.
type API struct {
	Token string

	mu          sync.Mutex
	models      map[string]client.Model
	order       []string
	quotas      map[string]client.QuotaView
	attachments []client.Attachment
	nodes       []client.NodeView
	actors      []string
	// scopes maps scope:ID to the path people know it by, rebuilt from
	// the tree whenever it changes.
	scopes map[string]string
	tenancy
}

// New builds a fake with a small tenancy tree — Acme, Acme/ml,
// Acme/ml/training — one person, ayse@example.com (u1), two guard policies
// and the nodes given.
func New(token string, nodes ...client.NodeView) *API {
	a := &API{
		Token:  token,
		models: map[string]client.Model{},
		quotas: map[string]client.QuotaView{},
		nodes:  nodes,
		tenancy: tenancy{
			levels: map[string][]node{
				client.LevelOrg:     {{ID: "o1", Name: "Acme"}},
				client.LevelTeam:    {{ID: "t1", Name: "ml", In: "o1"}},
				client.LevelProject: {{ID: "p1", Name: "training", In: "t1"}},
			},
			keys:  map[string]client.KeyView{},
			users: map[string]*client.UserView{"u1": {ID: "u1", Email: "ayse@example.com", Memberships: []client.MembershipView{}}},
		},
	}
	a.rebuild()
	return a
}

// Actors are who each change said it was made by, in order.
func (a *API) Actors() []string {
	a.mu.Lock()
	defer a.mu.Unlock()
	return slices.Clone(a.actors)
}

// Model is a model as the fake holds it, for a test to look at.
func (a *API) Model(name string) (client.Model, bool) {
	a.mu.Lock()
	defer a.mu.Unlock()
	m, ok := a.models[name]
	return m, ok
}

// Quota is a quota as the fake holds it, by scope:ID.
func (a *API) Quota(scope string) (client.QuotaView, bool) {
	a.mu.Lock()
	defer a.mu.Unlock()
	q, ok := a.quotas[scope]
	return q, ok
}

// Attachments are the guard attachments as the fake holds them.
func (a *API) Attachments() []client.Attachment {
	a.mu.Lock()
	defer a.mu.Unlock()
	return slices.Clone(a.attachments)
}

// ScaleBehindTerraform changes a model's scaling as someone at a terminal
// would, for a test of drift.
func (a *API) ScaleBehindTerraform(name string, minimum, maximum int64) {
	a.mu.Lock()
	defer a.mu.Unlock()
	m := a.models[name]
	m.Scaling.Min, m.Scaling.Max = minimum, maximum
	a.models[name] = m
}

// RemoveBehindTerraform removes a model as someone at a terminal would.
func (a *API) RemoveBehindTerraform(name string) {
	a.mu.Lock()
	defer a.mu.Unlock()
	delete(a.models, name)
	a.order = slices.DeleteFunc(a.order, func(n string) bool { return n == name })
}

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

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

// ServeHTTP checks what the admin listener checks, then routes.
func (a *API) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	if r.Header.Get("Authorization") != "Bearer "+a.Token {
		writeError(w, http.StatusUnauthorized, "the admin API needs its token")
		return
	}
	if r.Method != http.MethodGet {
		if ct, _, _ := strings.Cut(r.Header.Get("Content-Type"), ";"); ct != "application/json" {
			http.Error(w, "a change to the admin API must be sent as application/json", http.StatusUnsupportedMediaType)
			return
		}
		a.mu.Lock()
		a.actors = append(a.actors, r.Header.Get("X-Llmgw-Actor"))
		a.mu.Unlock()
	}
	a.mux().ServeHTTP(w, r)
}

func (a *API) mux() *http.ServeMux {
	mux := http.NewServeMux()
	mux.HandleFunc("GET /api/v1/models", a.listModels)
	mux.HandleFunc("POST /api/v1/models", a.addModel)
	mux.HandleFunc("GET /api/v1/models/{name}", a.getModel)
	mux.HandleFunc("DELETE /api/v1/models/{name}", a.removeModel)
	mux.HandleFunc("POST /api/v1/models/{name}/scale", a.scale)
	mux.HandleFunc("POST /api/v1/models/{name}/class", a.class)
	mux.HandleFunc("POST /api/v1/models/{name}/owner", a.owner)
	mux.HandleFunc("POST /api/v1/models/{name}/adapters", a.adapters)
	mux.HandleFunc("GET /api/v1/nodes", func(w http.ResponseWriter, _ *http.Request) {
		a.mu.Lock()
		defer a.mu.Unlock()
		out := slices.Clone(a.nodes)
		if out == nil {
			out = []client.NodeView{}
		}
		writeJSON(w, http.StatusOK, out)
	})
	mux.HandleFunc("GET /api/v1/quotas", a.listQuotas)
	mux.HandleFunc("PUT /api/v1/quotas", a.putQuota)
	mux.HandleFunc("DELETE /api/v1/quotas/{scope}", a.deleteQuota)
	mux.HandleFunc("GET /api/v1/guard", a.guard)
	mux.HandleFunc("POST /api/v1/guard/attachments", a.attach(true))
	mux.HandleFunc("DELETE /api/v1/guard/attachments", a.attach(false))
	a.tenancyRoutes(mux)
	return mux
}

// normalize fills in what the catalog fills in.
func normalize(m client.Model) client.Model {
	if m.Scaling.Max == 0 && m.Scaling.Min > 0 {
		m.Scaling.Max = m.Scaling.Min
	}
	for i := range m.Profiles {
		p := &m.Profiles[i]
		if p.Engine == "" {
			p.Engine = "vllm"
		}
		if p.TP <= 0 {
			p.TP = 1
		}
		if p.Nodes <= 0 {
			p.Nodes = 1
		}
		if p.StartupTimeout <= 0 {
			p.StartupTimeout = 20 * time.Minute
		}
		if p.Weight <= 0 {
			p.Weight = 1
		}
	}
	return m
}

func (a *API) listModels(w http.ResponseWriter, _ *http.Request) {
	a.mu.Lock()
	defer a.mu.Unlock()
	out := []client.ModelView{}
	for _, n := range a.order {
		m := a.models[n]
		v := client.ModelView{Name: m.Name, Description: m.Description, Aliases: m.Aliases, Source: m.Source,
			Min: m.Scaling.Min, Max: m.Scaling.Max, Desired: m.Scaling.Min, Class: m.Class, Owner: m.Owner}
		if v.Class == "" {
			v.Class = "standard"
		}
		for _, p := range m.Profiles {
			v.Profiles = append(v.Profiles, p.Name)
		}
		if m.External != nil {
			v.External = m.External.URL + " · " + m.External.Model
		}
		out = append(out, v)
	}
	writeJSON(w, http.StatusOK, out)
}

func (a *API) addModel(w http.ResponseWriter, r *http.Request) {
	var m client.Model
	if err := json.NewDecoder(r.Body).Decode(&m); err != nil {
		writeError(w, http.StatusBadRequest, "the model definition is not valid JSON: "+err.Error())
		return
	}
	a.mu.Lock()
	defer a.mu.Unlock()
	if m.Name == "" || (m.Source == "" && m.External == nil) {
		writeError(w, http.StatusBadRequest, "a model needs a name and a source")
		return
	}
	if _, ok := a.models[m.Name]; ok {
		writeError(w, http.StatusBadRequest, fmt.Sprintf("model %q is already in the catalog", m.Name))
		return
	}
	a.models[m.Name] = normalize(m)
	a.order = append(a.order, m.Name)
	writeJSON(w, http.StatusOK, map[string]string{"added": m.Name})
}

func (a *API) getModel(w http.ResponseWriter, r *http.Request) {
	a.mu.Lock()
	defer a.mu.Unlock()
	m, ok := a.models[r.PathValue("name")]
	if !ok {
		writeError(w, http.StatusNotFound, "no model named "+r.PathValue("name")+" is in the catalog")
		return
	}
	writeJSON(w, http.StatusOK, m)
}

func (a *API) removeModel(w http.ResponseWriter, r *http.Request) {
	a.mu.Lock()
	defer a.mu.Unlock()
	name := r.PathValue("name")
	if _, ok := a.models[name]; !ok {
		// The real API says this with a 400: the model was not added at
		// runtime, or is not there at all.
		writeError(w, http.StatusBadRequest, fmt.Sprintf("model %q was not added at runtime", name))
		return
	}
	delete(a.models, name)
	a.order = slices.DeleteFunc(a.order, func(n string) bool { return n == name })
	writeJSON(w, http.StatusOK, map[string]string{"removed": name})
}

// change applies f to a model, answering 400 when there is none.
func (a *API) change(w http.ResponseWriter, r *http.Request, body any, f func(*client.Model) error) (client.Model, bool) {
	if err := json.NewDecoder(r.Body).Decode(body); err != nil {
		writeError(w, http.StatusBadRequest, "the request is not valid JSON: "+err.Error())
		return client.Model{}, false
	}
	a.mu.Lock()
	defer a.mu.Unlock()
	m, ok := a.models[r.PathValue("name")]
	if !ok {
		writeError(w, http.StatusBadRequest, fmt.Sprintf("model %q was not added at runtime", r.PathValue("name")))
		return client.Model{}, false
	}
	if err := f(&m); err != nil {
		writeError(w, http.StatusBadRequest, err.Error())
		return client.Model{}, false
	}
	m = normalize(m)
	a.models[m.Name] = m
	return m, true
}

func (a *API) scale(w http.ResponseWriter, r *http.Request) {
	var req client.ScaleRequest
	m, ok := a.change(w, r, &req, func(m *client.Model) error {
		if req.Min < 0 || req.Max < 0 {
			return fmt.Errorf("a replica count cannot be negative")
		}
		m.Scaling = client.Scaling{Min: req.Min, Max: req.Max, IdleTimeout: req.IdleTimeout}
		return nil
	})
	if ok {
		writeJSON(w, http.StatusOK, map[string]any{"model": m.Name, "min": m.Scaling.Min, "max": m.Scaling.Max,
			"idle_timeout": m.Scaling.IdleTimeout, "on_demand": m.Scaling.Min == 0 && m.Scaling.Max > 0})
	}
}

func (a *API) class(w http.ResponseWriter, r *http.Request) {
	var req client.ClassRequest
	m, ok := a.change(w, r, &req, func(m *client.Model) error {
		switch req.Class {
		case "", "critical", "standard", "opportunistic":
			m.Class = req.Class
			return nil
		}
		return fmt.Errorf("class %q is not critical, standard or opportunistic", req.Class)
	})
	if ok {
		writeJSON(w, http.StatusOK, map[string]string{"model": m.Name, "class": m.Class})
	}
}

func (a *API) owner(w http.ResponseWriter, r *http.Request) {
	var req client.OwnerRequest
	m, ok := a.change(w, r, &req, func(m *client.Model) error {
		if req.Owner == "" {
			m.Owner = ""
			return nil
		}
		key, err := a.resolve(req.Owner)
		m.Owner = key
		return err
	})
	if ok {
		writeJSON(w, http.StatusOK, client.OwnerView{Model: m.Name, Owner: m.Owner})
	}
}

func (a *API) adapters(w http.ResponseWriter, r *http.Request) {
	var req client.AdaptersRequest
	m, ok := a.change(w, r, &req, func(m *client.Model) error {
		m.Adapters = req.Adapters
		if len(m.Adapters) == 0 {
			m.Adapters = nil
		}
		return nil
	})
	if ok {
		writeJSON(w, http.StatusOK, client.AdaptersRequest{Adapters: m.Adapters})
	}
}

// resolve reads org:X, team:X or project:X, X an ID or a path, as scope:ID.
// Must be called holding mu.
func (a *API) resolve(s string) (string, error) {
	kind, what, ok := strings.Cut(s, ":")
	if !ok || what == "" {
		return "", fmt.Errorf("%q is not org:NAME, team:ORG/TEAM or project:ORG/TEAM/PROJECT", s)
	}
	if _, ok := a.scopes[s]; ok {
		return s, nil
	}
	norm := func(p string) string {
		parts := strings.Split(p, "/")
		for i := range parts {
			parts[i] = strings.ToLower(strings.TrimSpace(parts[i]))
		}
		return strings.Join(parts, "/")
	}
	for key, name := range a.scopes {
		if strings.HasPrefix(key, kind+":") && norm(name) == norm(what) {
			return key, nil
		}
	}
	return "", fmt.Errorf("there is no %s %q", kind, what)
}

func (a *API) listQuotas(w http.ResponseWriter, _ *http.Request) {
	a.mu.Lock()
	defer a.mu.Unlock()
	out := []client.QuotaView{}
	for _, q := range a.quotas {
		out = append(out, q)
	}
	sort.Slice(out, func(i, j int) bool { return out[i].Scope < out[j].Scope })
	writeJSON(w, http.StatusOK, out)
}

func (a *API) putQuota(w http.ResponseWriter, r *http.Request) {
	var req client.QuotaRequest
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		writeError(w, http.StatusBadRequest, "the request is not valid JSON: "+err.Error())
		return
	}
	a.mu.Lock()
	defer a.mu.Unlock()
	key, err := a.resolve(req.Scope)
	if err != nil {
		writeError(w, http.StatusBadRequest, err.Error())
		return
	}
	empty := func(r client.Resources) bool { return r.GPUs == nil && r.CPUCores == nil && r.MemoryGB == nil }
	if empty(req.Guarantee) && empty(req.Limit) {
		writeError(w, http.StatusBadRequest, "a quota needs a guarantee or a limit; to remove one, clear it")
		return
	}
	a.quotas[key] = client.QuotaView{Scope: key, Name: a.scopes[key], Guarantee: req.Guarantee, Limit: req.Limit}
	writeJSON(w, http.StatusOK, client.QuotaSet{Scope: key})
}

func (a *API) deleteQuota(w http.ResponseWriter, r *http.Request) {
	a.mu.Lock()
	defer a.mu.Unlock()
	key, err := a.resolve(r.PathValue("scope"))
	if err != nil {
		writeError(w, http.StatusBadRequest, err.Error())
		return
	}
	delete(a.quotas, key)
	writeJSON(w, http.StatusOK, map[string]string{"removed": key})
}

func (a *API) guard(w http.ResponseWriter, _ *http.Request) {
	a.mu.Lock()
	defer a.mu.Unlock()
	type view struct {
		client.Attachment
		Name string `json:"name,omitempty"`
	}
	atts := []view{}
	for _, at := range a.attachments {
		atts = append(atts, view{Attachment: at, Name: a.scopes[at.Scope+":"+at.ScopeID]})
	}
	writeJSON(w, http.StatusOK, map[string]any{
		"policies": []string{"baseline", "finance"}, "default": []string{"baseline"}, "attachments": atts,
	})
}

func (a *API) attach(on bool) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		var req client.GuardAttachmentRequest
		if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
			writeError(w, http.StatusBadRequest, err.Error())
			return
		}
		a.mu.Lock()
		defer a.mu.Unlock()
		kind, what, _ := strings.Cut(req.Target, ":")
		at := client.Attachment{Scope: kind, ScopeID: what, Policy: req.Policy}
		if kind == "org" || kind == "team" || kind == "project" {
			key, err := a.resolve(req.Target)
			if err != nil {
				writeError(w, http.StatusBadRequest, err.Error())
				return
			}
			_, at.ScopeID, _ = strings.Cut(key, ":")
		} else if kind != "key" && kind != "user" {
			writeError(w, http.StatusBadRequest, fmt.Sprintf("%q: attach to org, team, project, key or user", kind))
			return
		}
		i := slices.Index(a.attachments, at)
		switch {
		case on && req.Policy != "baseline" && req.Policy != "finance":
			writeError(w, http.StatusBadRequest, fmt.Sprintf("no policy %q in server.yaml's guard.policies", req.Policy))
			return
		case on && i < 0:
			a.attachments = append(a.attachments, at)
		case !on && i < 0:
			writeError(w, http.StatusNotFound, "no such attachment")
			return
		case !on:
			a.attachments = slices.Delete(a.attachments, i, i+1)
		}
		writeJSON(w, http.StatusOK, at)
	}
}