llmgw/internal/mirror/api.go

package mirror

import (
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"os"
	"strings"
)

// API is the mirror's management interface.
//
// Management only. The weights themselves are served by an ordinary file
// handler on a separate listener, and that split is deliberate: nodes fetching
// forty gigabytes should not have to authenticate against a management API, and
// putting a client certificate in front of the byte path would mean
// re-certifying the agent's fetcher to change how models are administered.
//
// So: bytes over plain HTTP where every node can reach them, control over an
// API that can be locked down without touching them.
type API struct {
	m *Manager
	// base is how a catalog should refer to files here, reported alongside each
	// model so the caller does not have to reconstruct the URL and get it
	// subtly wrong.
	base string
}

// NewAPI wraps a manager in handlers.
func NewAPI(m *Manager, base string) *API { return &API{m: m, base: base} }

// Handler routes the management endpoints.
func (a *API) Handler() http.Handler {
	mux := http.NewServeMux()
	mux.HandleFunc("GET /v1/models", a.list)
	mux.HandleFunc("POST /v1/models", a.fetch)
	mux.HandleFunc("GET /v1/models/{model}/{version}", a.status)
	mux.HandleFunc("DELETE /v1/models/{model}/{version}", a.remove)
	mux.HandleFunc("PUT /v1/models/{model}/{version}/{file}", a.upload)
	mux.HandleFunc("GET /v1/health", a.health)
	mux.HandleFunc("PUT /v1/firewall", a.putFirewall)
	return mux
}

// ModelsResponse is what the mirror holds, and where to point a catalog.
type ModelsResponse struct {
	// Base is the address nodes fetch from. Returned rather than assumed
	// because the caller writing a catalog entry needs the same spelling the
	// mirror would use, and a mismatch there is what makes a file look unused
	// when it is not.
	Base   string   `json:"base"`
	Models []Status `json:"models"`
	// Removed are the versions removed here, model/version to when (ms),
	// for a peer mirror to remove too.
	Removed map[string]int64 `json:"removed,omitempty"`
}

func (a *API) list(w http.ResponseWriter, _ *http.Request) {
	all, err := a.m.All()
	if err != nil {
		fail(w, http.StatusInternalServerError, err)
		return
	}
	writeJSON(w, http.StatusOK, ModelsResponse{Base: a.base, Models: all, Removed: a.m.Tombstones()})
}

// FetchRequest asks the mirror to download a model.
type FetchRequest struct {
	Source  string `json:"source"`
	Model   string `json:"model"`
	Version string `json:"version"`
	// File is optional; the last element of the source URL is used when it is
	// empty, which is right almost always and wrong visibly rather than
	// quietly when it is not.
	File string `json:"file"`
	// SHA256 is what the registry published for the file: a download that
	// is not it is deleted and the fetch fails. Empty checks nothing.
	SHA256 string `json:"sha256,omitempty"`
}

// FetchResponse reports where a download stands, and where it will be.
type FetchResponse struct {
	Status Status `json:"status"`
	// URL is what a catalog entry should say once the download finishes. Given
	// now rather than later so the operator can write the catalog while the
	// bytes are still moving.
	URL string `json:"url"`
}

func (a *API) fetch(w http.ResponseWriter, r *http.Request) {
	var req FetchRequest
	if err := json.NewDecoder(http.MaxBytesReader(w, r.Body, 1<<16)).Decode(&req); err != nil {
		fail(w, http.StatusBadRequest, fmt.Errorf("reading the request: %w", err))
		return
	}
	if req.File == "" {
		req.File = LastPathElement(req.Source)
	}
	if req.Model == "" || req.Version == "" || req.File == "" {
		fail(w, http.StatusBadRequest,
			errors.New("model, version and a filename are all required; the filename can come from the source URL"))
		return
	}

	e := Entry{Model: req.Model, Version: req.Version, File: req.File}
	st, err := a.m.StartChecked(req.Source, e, strings.ToLower(req.SHA256))
	switch {
	case errors.Is(err, ErrAlreadyFetching):
		// Not an error the caller needs to handle: a control plane that asks
		// twice for the same model wanted it once. Reporting the download
		// already running says the same thing as starting it would have.
		writeJSON(w, http.StatusAccepted, FetchResponse{Status: st, URL: e.URL(a.base)})
		return
	case err != nil:
		fail(w, http.StatusBadRequest, err)
		return
	}
	// 202, not 200: the bytes are not here yet. A caller that treats this as
	// "done" would point a catalog at a file that is still arriving, and every
	// node would then fetch a partial one.
	writeJSON(w, http.StatusAccepted, FetchResponse{Status: st, URL: e.URL(a.base)})
}

func (a *API) status(w http.ResponseWriter, r *http.Request) {
	model, version := r.PathValue("model"), r.PathValue("version")
	st, ok := a.m.Status(model, version)
	if !ok {
		fail(w, http.StatusNotFound,
			fmt.Errorf("the mirror holds nothing for %s/%s", model, version))
		return
	}
	writeJSON(w, http.StatusOK, FetchResponse{
		Status: st,
		URL:    Entry{Model: st.Model, Version: st.Version, File: st.File}.URL(a.base),
	})
}

func (a *API) remove(w http.ResponseWriter, r *http.Request) {
	model, version := r.PathValue("model"), r.PathValue("version")
	if err := errors.Join(ValidName(model), ValidName(version)); err != nil {
		fail(w, http.StatusBadRequest, err)
		return
	}

	// A download in flight is not deleted out from under itself: the fetch
	// would carry on writing and the file would reappear, which looks like the
	// delete silently failed.
	if st, ok := a.m.Status(model, version); ok && st.State == Fetching {
		fail(w, http.StatusConflict,
			fmt.Errorf("%s/%s is still downloading; nothing was removed", model, version))
		return
	}

	removed, err := a.m.RemoveVersion(model, version)
	switch {
	case errors.Is(err, os.ErrNotExist):
		fail(w, http.StatusNotFound, fmt.Errorf("the mirror holds nothing for %s/%s", model, version))
		return
	case err != nil:
		fail(w, http.StatusInternalServerError, err)
		return
	}
	writeJSON(w, http.StatusOK, map[string]any{"removed": removed})
}

// health says the mirror is up and can read its own directory.
//
// Reading the directory rather than only answering: a mirror whose disk has
// gone away still accepts connections, and a check that only proves the process
// is running would call that healthy right up until a node tried to fetch.
func (a *API) health(w http.ResponseWriter, _ *http.Request) {
	entries, err := List(a.m.root)
	if err != nil {
		fail(w, http.StatusServiceUnavailable, err)
		return
	}
	var bytes int64
	for _, e := range entries {
		bytes += e.Bytes
	}
	writeJSON(w, http.StatusOK, map[string]any{
		"status": "ok", "files": len(entries), "bytes": bytes, "base": a.base,
	})
}

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

// fail answers with the reason in the body.
//
// In the body because the caller is a program that will log it and a person who
// will read that log. A bare status code turns "the source URL 404s" into "the
// mirror said 400", and the next hour goes into finding out which.
func fail(w http.ResponseWriter, code int, err error) {
	writeJSON(w, code, map[string]string{"error": err.Error()})
}

// LastPathElement is the filename a URL ends in.
func LastPathElement(u string) string {
	u = strings.SplitN(u, "?", 2)[0]
	if i := strings.LastIndex(u, "/"); i >= 0 && i < len(u)-1 {
		return u[i+1:]
	}
	return u
}

// maxUpload bounds a file sent to the mirror rather than fetched by it.
const maxUpload = 8 << 30

// upload keeps a file the caller sends — an adapter a node trained, which
// is on no registry to fetch from — with the same discipline as an import:
// its checksum and source beside it, and an existing version never
// replaced. X-Llmgw-Source says where it came from.
func (a *API) upload(w http.ResponseWriter, r *http.Request) {
	e := Entry{Model: r.PathValue("model"), Version: r.PathValue("version"), File: r.PathValue("file")}
	for _, n := range []string{e.Model, e.Version, e.File} {
		if err := ValidName(n); err != nil {
			fail(w, http.StatusBadRequest, err)
			return
		}
	}
	if _, err := os.Stat(e.Path(a.m.root)); err == nil {
		fail(w, http.StatusConflict, fmt.Errorf("%s/%s/%s: %w", e.Model, e.Version, e.File, ErrExists))
		return
	}
	tmp, err := os.CreateTemp(a.m.root, ".upload-*")
	if err != nil {
		fail(w, http.StatusInternalServerError, err)
		return
	}
	defer os.Remove(tmp.Name())
	n, err := io.Copy(tmp, http.MaxBytesReader(w, r.Body, maxUpload))
	if cerr := tmp.Close(); err == nil {
		err = cerr
	}
	if err != nil {
		fail(w, http.StatusBadRequest, fmt.Errorf("receiving the file: %w", err))
		return
	}
	if n == 0 {
		fail(w, http.StatusBadRequest, errors.New("the file is empty"))
		return
	}
	source := strings.TrimSpace(r.Header.Get("X-Llmgw-Source"))
	if source == "" {
		source = "uploaded"
	}
	added, err := Import(r.Context(), a.m.root, tmp.Name(), e, "", source, nil)
	switch {
	case errors.Is(err, ErrExists):
		fail(w, http.StatusConflict, err)
		return
	case err != nil:
		fail(w, http.StatusBadRequest, err)
		return
	}
	writeJSON(w, http.StatusCreated, FetchResponse{
		Status: Status{Model: e.Model, Version: e.Version, File: e.File, State: Present, BytesDone: added.Entry.Bytes, BytesTotal: added.Entry.Bytes},
		URL:    e.URL(a.base),
	})
}