llmgw/internal/mirror/client.go

package mirror

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"strings"
	"time"

	"gitea.mynodes.xyz/ilker/llmgw/internal/tlsconf"
)

// Client drives a mirror's management API from somewhere else.
//
// The mirror runs on its own machine with its own disk, so every operation the
// control plane wants — what is held, fetch this, how far along, remove that —
// is a request over the network rather than a function call. This is the half
// that makes the mirror something the control plane administers instead of
// something an operator has to SSH into.
type Client struct {
	base string
	hc   *http.Client
	// others are the mirrors kept in step with this one, asked in turn when
	// it cannot be reached: the same models, on another machine.
	others []string
}

// NewFailoverClient is NewClient over several mirrors kept in step: the
// first answers, and the next is asked when it cannot be reached.
func NewFailoverClient(addrs []string, tc tlsconf.Config, timeout time.Duration) (*Client, error) {
	if len(addrs) == 0 {
		return nil, fmt.Errorf("no mirror address")
	}
	c, err := NewClient(addrs[0], tc, timeout)
	if err != nil {
		return nil, err
	}
	for _, a := range addrs[1:] {
		u, err := url.Parse(a)
		if err != nil || u.Host == "" {
			return nil, fmt.Errorf("the mirror address %q is not a URL with a host", a)
		}
		c.others = append(c.others, strings.TrimSuffix(a, "/"))
	}
	return c, nil
}

// NewClient dials a mirror's management API, with mutual TLS when configured.
//
// NextProtos is cleared from what tlsconf builds, for the same reason the
// mirror's server clears it: that configuration is shaped for the agent stream,
// which needs HTTP/2 and pins "h2". A transport carrying a TLS config with a
// pinned protocol it is not itself set up to speak negotiates h2 and then talks
// HTTP/1.1 over it, which fails as a protocol error that mentions neither ALPN
// nor the mirror. Cleared here and ForceAttemptHTTP2 set, net/http negotiates
// whichever the two ends agree on and speaks it.
func NewClient(addr string, tc tlsconf.Config, timeout time.Duration) (*Client, error) {
	u, err := url.Parse(addr)
	if err != nil || u.Host == "" {
		return nil, fmt.Errorf("the mirror address %q is not a URL with a host", addr)
	}
	if timeout <= 0 {
		timeout = 15 * time.Second
	}
	hc := &http.Client{Timeout: timeout}

	if tc.Enabled() {
		cfg, err := tc.ClientConfig(u.Hostname())
		if err != nil {
			return nil, fmt.Errorf("the mirror's TLS configuration: %w", err)
		}
		cfg.NextProtos = nil
		hc.Transport = &http.Transport{TLSClientConfig: cfg, ForceAttemptHTTP2: true}
	}
	return &Client{base: strings.TrimSuffix(addr, "/"), hc: hc}, nil
}

// Addr is where this client points, for a page that has to say so.
func (c *Client) Addr() string { return c.base }

// APIError is a refusal the mirror explained.
//
// The status and the mirror's own sentence, both: a page showing "the mirror
// said 400" sends somebody to the wrong machine, while "the source URL returned
// 404" sends them to the right one. The status is kept alongside so a caller
// can still branch on it without parsing prose.
type APIError struct {
	Status int
	Msg    string
}

func (e *APIError) Error() string {
	if e.Msg == "" {
		return fmt.Sprintf("the mirror answered %d", e.Status)
	}
	return e.Msg
}

// Health reports whether the mirror is up and can read its own directory.
func (c *Client) Health(ctx context.Context) error {
	var out map[string]any
	return c.do(ctx, http.MethodGet, "/v1/health", nil, &out)
}

// List is everything the mirror holds or is fetching.
func (c *Client) List(ctx context.Context) (ModelsResponse, error) {
	var out ModelsResponse
	err := c.do(ctx, http.MethodGet, "/v1/models", nil, &out)
	return out, err
}

// Fetch asks the mirror to download a model.
//
// It returns as soon as the download is under way, not when it finishes —
// that is the mirror's contract and it is the caller's too. Poll Status.
func (c *Client) Fetch(ctx context.Context, req FetchRequest) (FetchResponse, error) {
	var out FetchResponse
	err := c.do(ctx, http.MethodPost, "/v1/models", req, &out)
	return out, err
}

// Upload sends a file for the mirror to keep at model/version/file, with
// where it came from; it answers the URL nodes fetch it from.
func (c *Client) Upload(ctx context.Context, model, version, file string, body []byte, source string) (FetchResponse, error) {
	var out FetchResponse
	req, err := http.NewRequestWithContext(ctx, http.MethodPut,
		c.base+"/v1/models/"+url.PathEscape(model)+"/"+url.PathEscape(version)+"/"+url.PathEscape(file), bytes.NewReader(body))
	if err != nil {
		return out, err
	}
	req.Header.Set("Content-Type", "application/octet-stream")
	req.Header.Set("X-Llmgw-Source", source)
	resp, err := c.hc.Do(req)
	if err != nil {
		return out, fmt.Errorf("reaching the mirror at %s: %w", c.base, err)
	}
	defer resp.Body.Close()
	raw, err := io.ReadAll(io.LimitReader(resp.Body, 1<<20))
	if err != nil {
		return out, err
	}
	if resp.StatusCode >= 300 {
		var e struct {
			Error string `json:"error"`
		}
		_ = json.Unmarshal(raw, &e)
		return out, &APIError{Status: resp.StatusCode, Msg: e.Error}
	}
	return out, json.Unmarshal(raw, &out)
}

// Status is where one model version stands.
func (c *Client) Status(ctx context.Context, model, version string) (FetchResponse, error) {
	var out FetchResponse
	err := c.do(ctx, http.MethodGet,
		"/v1/models/"+url.PathEscape(model)+"/"+url.PathEscape(version), nil, &out)
	return out, err
}

// Remove deletes a whole version and reports the files that went.
func (c *Client) Remove(ctx context.Context, model, version string) ([]string, error) {
	var out struct {
		Removed []string `json:"removed"`
	}
	err := c.do(ctx, http.MethodDelete,
		"/v1/models/"+url.PathEscape(model)+"/"+url.PathEscape(version), nil, &out)
	return out.Removed, err
}

func (c *Client) do(ctx context.Context, method, path string, body, out any) error {
	var b []byte
	if body != nil {
		var err error
		if b, err = json.Marshal(body); err != nil {
			return err
		}
	}
	// The mirror first, then those kept in step with it, while each cannot
	// be reached at all: one that answers, even with a refusal, is the one.
	var resp *http.Response
	var errs []error
	for _, base := range append([]string{c.base}, c.others...) {
		var rdr io.Reader
		if body != nil {
			rdr = bytes.NewReader(b)
		}
		req, err := http.NewRequestWithContext(ctx, method, base+path, rdr)
		if err != nil {
			return err
		}
		if body != nil {
			req.Header.Set("Content-Type", "application/json")
		}
		if resp, err = c.hc.Do(req); err == nil {
			break
		}
		// Wrapped with the address because the failure an operator sees on
		// a page is otherwise "connection refused" with no clue which of
		// the machines refused it.
		errs = append(errs, fmt.Errorf("reaching the mirror at %s: %w", base, err))
		if ctx.Err() != nil {
			break
		}
	}
	if resp == nil {
		return errors.Join(errs...)
	}
	defer resp.Body.Close()

	// Capped: the mirror is trusted but it is still a separate process over a
	// network, and a page that renders whatever it sends should not be able to
	// be handed a gigabyte.
	raw, err := io.ReadAll(io.LimitReader(resp.Body, 1<<20))
	if err != nil {
		return fmt.Errorf("reading the mirror's answer: %w", err)
	}

	if resp.StatusCode >= 300 {
		var e struct {
			Error string `json:"error"`
		}
		_ = json.Unmarshal(raw, &e)
		return &APIError{Status: resp.StatusCode, Msg: e.Error}
	}
	if out == nil {
		return nil
	}
	if err := json.Unmarshal(raw, out); err != nil {
		return fmt.Errorf("the mirror's answer was not the JSON expected: %w", err)
	}
	return nil
}