llmgw/internal/agent/supervisor/health.go

package supervisor

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"time"
)

// HTTPHealth checks an engine over HTTP.
//
// Readiness is two questions, not one (plan §5.4): /health says the process is
// up, and /v1/models says the weights finished loading and the model is being
// served. An engine answers /health long before it can serve a request, so
// routing on /health alone sends traffic into a model that is still loading.
type HTTPHealth struct {
	client      *http.Client
	healthPath  string
	modelsPath  string
	requireName bool
}

// HTTPHealthOption customizes the checker.
type HTTPHealthOption func(*HTTPHealth)

// WithHealthPath overrides the readiness path.
func WithHealthPath(p string) HTTPHealthOption {
	return func(h *HTTPHealth) { h.healthPath = p }
}

// WithoutModelCheck skips the /v1/models step, for engines that do not expose it.
func WithoutModelCheck() HTTPHealthOption {
	return func(h *HTTPHealth) { h.requireName = false }
}

// NewHTTPHealth builds a checker.
func NewHTTPHealth(timeout time.Duration, opts ...HTTPHealthOption) *HTTPHealth {
	if timeout <= 0 {
		timeout = 3 * time.Second
	}
	h := &HTTPHealth{
		client:      &http.Client{Timeout: timeout},
		healthPath:  "/health",
		modelsPath:  "/v1/models",
		requireName: true,
	}
	for _, o := range opts {
		o(h)
	}
	return h
}

// Ready reports whether the engine at baseURL is serving model.
func (h *HTTPHealth) Ready(ctx context.Context, baseURL, model, apiKey string) error {
	var body healthBody
	if err := h.get(ctx, baseURL+h.healthPath, apiKey, &body); err != nil {
		return fmt.Errorf("health: %w", err)
	}
	if body.notReady() {
		return fmt.Errorf("health: the engine says %q", body.Status)
	}
	if !h.requireName || model == "" {
		return nil
	}

	var list struct {
		Data []struct {
			ID string `json:"id"`
		} `json:"data"`
	}
	if err := h.get(ctx, baseURL+h.modelsPath, apiKey, &list); err != nil {
		return fmt.Errorf("models: %w", err)
	}
	for _, d := range list.Data {
		if d.ID == model {
			return nil
		}
	}
	return fmt.Errorf("engine is up but does not serve %q yet", model)
}

func (h *HTTPHealth) get(ctx context.Context, url, apiKey string, out any) error {
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
	if err != nil {
		return err
	}
	if apiKey != "" {
		req.Header.Set("Authorization", "Bearer "+apiKey)
	}
	resp, err := h.client.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		_, _ = io.Copy(io.Discard, io.LimitReader(resp.Body, 1<<10))
		return fmt.Errorf("status %s", resp.Status)
	}
	if out == nil {
		return nil
	}
	if hb, ok := out.(*healthBody); ok {
		// Most engines answer /health with nothing, or with text: only a
		// JSON status says anything, and a body that is not one is not a
		// failure.
		_ = json.NewDecoder(io.LimitReader(resp.Body, 64<<10)).Decode(hb)
		return nil
	}
	return json.NewDecoder(resp.Body).Decode(out)
}

// healthBody is what an engine's /health may say in its body. Some answer
// 200 from the moment their HTTP server is up and say in the body that the
// model is still loading — FreeToken's {"status":"loading","phase":...} —
// and a model that lists itself in /v1/models meanwhile would be routed to.
type healthBody struct {
	Status string `json:"status"`
}

// notReady is a body that says the engine cannot serve yet, or has failed.
func (b healthBody) notReady() bool {
	switch b.Status {
	case "loading", "starting", "initializing", "warming_up", "error", "unhealthy":
		return true
	}
	return false
}