llmgw/internal/web/admin_mirror.go

package web

import (
	"context"
	"errors"
	"net/http"
	"sort"
	"strconv"
	"strings"
	"time"

	"gitea.mynodes.xyz/ilker/llmgw/internal/mirror"
	"gitea.mynodes.xyz/ilker/llmgw/internal/store"
)

// MirrorAdmin is the mirror's management API as this page needs it.
//
// An interface rather than the client itself so the page can be tested without
// a mirror, and so a deployment without one leaves it nil and the page
// unmounted — the same shape the cluster dashboard uses.
type MirrorAdmin interface {
	List(ctx context.Context) (mirror.ModelsResponse, error)
	Fetch(ctx context.Context, req mirror.FetchRequest) (mirror.FetchResponse, error)
	Remove(ctx context.Context, model, version string) ([]string, error)
	// Addr is the mirror this points at, shown on the page because an operator
	// looking at an empty table needs to know which machine is empty.
	Addr() string
}

// mirrorRow is one model version as the page shows it.
type mirrorRow struct {
	mirror.Status
	// URL is what a catalog entry should say. The one string on this page
	// somebody will copy, so it is rendered rather than left to be assembled.
	URL     string
	Percent int
	// HasPercent is false when the upstream never sent a length. The bar is
	// then not drawn at all: a progress bar that has to invent its own total is
	// worse than a byte count that only goes up.
	HasPercent bool
}

type adminMirrorData struct {
	Addr string
	Base string
	Rows []mirrorRow

	Fetching int
	Present  int
	Failed   int

	// Unreachable separates "the mirror holds nothing" from "the mirror did not
	// answer". They render almost identically and mean opposite things: one is
	// a mirror waiting to be filled, the other is a machine to go and look at.
	Unreachable bool

	Message string
	Error   string
	CSRF    string
	Version string
	Now     string
	// RefreshSec is 0 when nothing is downloading. The page reloads only while
	// there is something to watch, because a reload while an operator is typing
	// a source URL into the form takes the URL with it.
	RefreshSec int
}

// RegisterAdminMirror mounts the model mirror under /admin/mirror.
//
// Fetching and removing, and nothing else. What the mirror holds is operator
// state rather than account state, but the reasoning is the same as
// /admin/users: these two are the work somebody does repeatedly and should not
// need a shell for, while anything that would let a stolen session reach
// further — pointing the control plane at a different mirror, say — stays in
// configuration where changing it means access to the host.
func (p *AccountPages) RegisterAdminMirror(mux *http.ServeMux, m MirrorAdmin) {
	p.hasMirror = true

	mux.HandleFunc("GET /admin/mirror", p.adminOnly(
		func(w http.ResponseWriter, r *http.Request, _ store.User) {
			p.renderMirror(w, r, m, "", "")
		}))

	mux.HandleFunc("POST /admin/mirror/fetch", p.mirrorAction(m,
		func(ctx context.Context, admin store.User, r *http.Request) (string, error) {
			req := mirror.FetchRequest{
				Source:  strings.TrimSpace(r.FormValue("source")),
				Model:   strings.TrimSpace(r.FormValue("model")),
				Version: strings.TrimSpace(r.FormValue("version")),
				File:    strings.TrimSpace(r.FormValue("file")),
			}
			if req.Source == "" {
				return "", errors.New("a source URL is required")
			}
			got, err := m.Fetch(ctx, req)
			if err != nil {
				return "", err
			}
			// Named as started, not as done. The mirror answered 202 and the
			// bytes are still moving; a message saying "added" would invite
			// somebody to point the catalog at a partial file.
			return "started downloading " + got.Status.Model + "/" + got.Status.Version, nil
		}))

	mux.HandleFunc("POST /admin/mirror/{model}/{version}/remove", p.mirrorAction(m,
		func(ctx context.Context, _ store.User, r *http.Request) (string, error) {
			model, ver := r.PathValue("model"), r.PathValue("version")
			removed, err := m.Remove(ctx, model, ver)
			if err != nil {
				return "", err
			}
			return "removed " + model + "/" + ver +
				" (" + plural(len(removed), "file", "files") + ")", nil
		}))
}

// mirrorAction wraps a change: admin session, CSRF, audit line, re-render.
//
// The audit line is not optional. Removing a version deletes weights every node
// may be about to fetch, and the only record that it happened on purpose is
// this log entry naming who clicked it.
func (p *AccountPages) mirrorAction(m MirrorAdmin,
	do func(ctx context.Context, admin store.User, r *http.Request) (string, error),
) http.HandlerFunc {
	return p.adminOnly(func(w http.ResponseWriter, r *http.Request, admin store.User) {
		sess, _, _ := p.session(r)
		if !p.accounts.CheckCSRF(sess, r.FormValue("csrf")) {
			http.Error(w, "bad request", http.StatusBadRequest)
			return
		}
		msg, err := do(r.Context(), admin, r)
		if err != nil {
			p.log.Warn("a mirror action was refused",
				"admin", admin.Email, "mirror", m.Addr(), "err", err)
			p.renderMirror(w, r, m, "", err.Error())
			return
		}
		p.log.Info("the mirror was changed",
			"admin", admin.Email, "mirror", m.Addr(), "what", msg)
		p.renderMirror(w, r, m, msg, "")
	})
}

func (p *AccountPages) renderMirror(w http.ResponseWriter, r *http.Request,
	m MirrorAdmin, msg, errMsg string,
) {
	d := adminMirrorData{
		Addr:    m.Addr(),
		Message: msg,
		Error:   errMsg,
		CSRF:    csrfFor(r),
		Version: p.version,
		Now:     time.Now().UTC().Format(time.RFC3339),
	}

	got, err := m.List(r.Context())
	if err != nil {
		p.log.Error("could not read the mirror", "mirror", m.Addr(), "err", err)
		d.Unreachable = true
		if d.Error == "" {
			d.Error = err.Error()
		}
		render(w, "admin_mirror", d)
		return
	}

	d.Base = got.Base
	for _, st := range got.Models {
		row := mirrorRow{
			Status: st,
			URL: mirror.Entry{Model: st.Model, Version: st.Version, File: st.File}.
				URL(got.Base),
		}
		if pc, ok := st.Percent(); ok {
			row.Percent, row.HasPercent = pc, true
		}
		switch st.State {
		case mirror.Fetching:
			d.Fetching++
		case mirror.Failed:
			d.Failed++
		default:
			d.Present++
		}
		d.Rows = append(d.Rows, row)
	}

	// Downloads first, then failures, then what is simply here. The two states
	// somebody opened this page for are at the top; the rest is a list they
	// scroll. Within a state the order is fixed so rows do not shuffle between
	// reloads while a byte count ticks.
	sort.SliceStable(d.Rows, func(i, j int) bool {
		a, b := d.Rows[i], d.Rows[j]
		if ra, rb := mirrorRank(a.State), mirrorRank(b.State); ra != rb {
			return ra < rb
		}
		if a.Model != b.Model {
			return a.Model < b.Model
		}
		if a.Version != b.Version {
			return a.Version < b.Version
		}
		return a.File < b.File
	})

	if d.Fetching > 0 {
		d.RefreshSec = 5
	}
	render(w, "admin_mirror", d)
}

func mirrorRank(s mirror.State) int {
	switch s {
	case mirror.Fetching:
		return 0
	case mirror.Failed:
		return 1
	default:
		return 2
	}
}

func plural(n int, one, many string) string {
	if n == 1 {
		return "1 " + one
	}
	return strconv.Itoa(n) + " " + many
}