llmgw/internal/agent/weights/fetch_test.go

package weights

import (
	"crypto/sha256"
	"encoding/hex"
	"errors"
	"net/http"
	"net/http/httptest"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"sync/atomic"
	"testing"
	"time"
)

func serveBytes(t *testing.T, body []byte, hits *atomic.Int64) *httptest.Server {
	t.Helper()
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if hits != nil {
			hits.Add(1)
		}
		http.ServeContent(w, r, "model.gguf", time.Time{}, strings.NewReader(string(body)))
	}))
	t.Cleanup(srv.Close)
	return srv
}

func TestALocalPathIsUsedNeverCopied(t *testing.T) {
	// A path is what an operator writes when the file is already there —
	// local disk, a Lustre mount, an image exported over the fabric. A
	// downloader that second-guessed that would copy a hundred gigabytes that
	// were already present.
	dir := t.TempDir()
	path := filepath.Join(dir, "model.gguf")
	if err := os.WriteFile(path, []byte("weights"), 0o644); err != nil {
		t.Fatal(err)
	}
	f := &Fetcher{CacheDir: t.TempDir()}

	got, err := f.Ensure(t.Context(), path, false)
	if err != nil {
		t.Fatal(err)
	}
	if got != path {
		t.Errorf("returned %q, want the original path", got)
	}
	if entries, _ := os.ReadDir(f.CacheDir); len(entries) != 0 {
		t.Errorf("the cache was written to for a local file: %v", entries)
	}
}

func TestAMissingLocalPathSaysWhatIsWrong(t *testing.T) {
	// The commonest setup mistake, and "no such file" from an engine twenty
	// minutes into a load is a poor way to learn it.
	f := &Fetcher{CacheDir: t.TempDir()}
	_, err := f.Ensure(t.Context(), filepath.Join(t.TempDir(), "absent.gguf"), false)
	if !errors.Is(err, ErrNotLocal) {
		t.Fatalf("err = %v, want ErrNotLocal", err)
	}
	if !strings.Contains(err.Error(), "absent.gguf") {
		t.Errorf("the error does not name the file: %v", err)
	}
}

func TestARemoteFileIsFetchedOnceAndReused(t *testing.T) {
	var hits atomic.Int64
	srv := serveBytes(t, []byte("the weights"), &hits)
	f := &Fetcher{CacheDir: t.TempDir()}

	first, err := f.Ensure(t.Context(), srv.URL+"/model.gguf", false)
	if err != nil {
		t.Fatal(err)
	}
	body, err := os.ReadFile(first)
	if err != nil || string(body) != "the weights" {
		t.Fatalf("fetched content = %q, err %v", body, err)
	}

	second, err := f.Ensure(t.Context(), srv.URL+"/model.gguf", false)
	if err != nil {
		t.Fatal(err)
	}
	if second != first {
		t.Errorf("the second call returned a different path: %q vs %q", second, first)
	}
	if hits.Load() != 1 {
		t.Errorf("the file was fetched %d times; the cache did not hold", hits.Load())
	}
}

func TestConcurrentRequestsProduceOneDownload(t *testing.T) {
	// Eight shards of a tensor-parallel replica start at the same instant and
	// want the same file. Eight downloads is a self-inflicted denial of service
	// on whatever is serving them.
	var hits atomic.Int64
	srv := serveBytes(t, []byte("shared weights"), &hits)
	f := &Fetcher{CacheDir: t.TempDir()}

	var wg sync.WaitGroup
	paths := make([]string, 8)
	errs := make([]error, 8)
	for i := range 8 {
		wg.Add(1)
		go func() {
			defer wg.Done()
			paths[i], errs[i] = f.Ensure(t.Context(), srv.URL+"/model.gguf", false)
		}()
	}
	wg.Wait()

	for i, err := range errs {
		if err != nil {
			t.Fatalf("caller %d failed: %v", i, err)
		}
		if paths[i] != paths[0] {
			t.Errorf("caller %d got %q, caller 0 got %q", i, paths[i], paths[0])
		}
	}
	if n := hits.Load(); n != 1 {
		t.Errorf("%d downloads for eight concurrent callers", n)
	}
}

func TestAPartialDownloadIsNeverVisibleAsTheModel(t *testing.T) {
	// A partial file under the real name would be found by the next scan,
	// reported as present, and handed to an engine that fails on it minutes
	// later — the exact failure all of this exists to prevent.
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Length", "1000")
		w.WriteHeader(http.StatusOK)
		_, _ = w.Write([]byte("only the first bytes"))
		// and then the connection dies without the rest
		if fl, ok := w.(http.Flusher); ok {
			fl.Flush()
		}
		panic(http.ErrAbortHandler)
	}))
	defer srv.Close()

	cache := t.TempDir()
	f := &Fetcher{CacheDir: cache}
	if _, err := f.Ensure(t.Context(), srv.URL+"/model.gguf", false); err == nil {
		t.Fatal("a truncated download was reported as success")
	}
	entries, _ := os.ReadDir(cache)
	for _, e := range entries {
		if !strings.HasSuffix(e.Name(), ".part") {
			t.Errorf("a non-part file was left behind: %q", e.Name())
		}
	}
}

func TestAnInterruptedDownloadResumes(t *testing.T) {
	// A forty gigabyte fetch that fails at ninety percent must not start again
	// from nothing.
	full := []byte(strings.Repeat("x", 4096))
	var ranges []string
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		ranges = append(ranges, r.Header.Get("Range"))
		http.ServeContent(w, r, "model.gguf", time.Time{}, strings.NewReader(string(full)))
	}))
	defer srv.Close()

	cache := t.TempDir()
	f := &Fetcher{CacheDir: cache}
	target, _ := f.cachePath(srv.URL + "/model.gguf")
	// Half a file already on disk, as an interrupted attempt would leave.
	if err := os.WriteFile(target+".part", full[:2048], 0o644); err != nil {
		t.Fatal(err)
	}

	got, err := f.Ensure(t.Context(), srv.URL+"/model.gguf", false)
	if err != nil {
		t.Fatal(err)
	}
	body, _ := os.ReadFile(got)
	if len(body) != len(full) {
		t.Errorf("resumed file is %d bytes, want %d", len(body), len(full))
	}
	if len(ranges) == 0 || !strings.Contains(ranges[0], "2048") {
		t.Errorf("no range request was made: %v", ranges)
	}
}

func TestAnHFSourceNamingARepositoryExplainsItself(t *testing.T) {
	// A whole repository of shards is deliberately not handled here: vLLM and
	// SGLang download those themselves and do it well. Saying so beats a
	// confusing 404.
	_, err := resolveURL("hf:Qwen/Qwen2.5-Coder-32B-Instruct")
	if err == nil {
		t.Fatal("a bare repository was accepted as a file")
	}
	for _, want := range []string{"model.gguf", "let the engine"} {
		if !strings.Contains(err.Error(), want) {
			t.Errorf("the error does not suggest what to do (%q): %v", want, err)
		}
	}
}

func TestAnHFFileResolvesToItsDownloadURL(t *testing.T) {
	got, err := resolveURL("hf:Qwen/Qwen2.5-Coder-3B-Instruct-GGUF/qwen2.5-coder-3b-instruct-q4_k_m.gguf")
	if err != nil {
		t.Fatal(err)
	}
	want := "https://huggingface.co/Qwen/Qwen2.5-Coder-3B-Instruct-GGUF/resolve/main/qwen2.5-coder-3b-instruct-q4_k_m.gguf"
	if got != want {
		t.Errorf("resolved to\n  %s\nwant\n  %s", got, want)
	}
}

func TestTwoSourcesWithTheSameFilenameDoNotCollide(t *testing.T) {
	f := &Fetcher{CacheDir: t.TempDir()}
	a, _ := f.cachePath("https://one.example/model.gguf")
	b, _ := f.cachePath("https://two.example/model.gguf")
	if a == b {
		t.Fatalf("both sources cache to %q", a)
	}
	// And a person looking at the directory can still tell what is in it.
	for _, p := range []string{a, b} {
		if !strings.HasSuffix(p, "model.gguf") {
			t.Errorf("cache path %q does not keep the filename", p)
		}
	}
}

func TestProgressIsReportedWithAPercentage(t *testing.T) {
	// A missing model is the difference between a replica that is slow to start
	// and one that is stuck, and without this they look identical.
	srv := serveBytes(t, []byte(strings.Repeat("y", 8192)), nil)
	var seen []Progress
	f := &Fetcher{
		CacheDir:   t.TempDir(),
		OnProgress: func(p Progress) { seen = append(seen, p) },
	}
	if _, err := f.Ensure(t.Context(), srv.URL+"/model.gguf", false); err != nil {
		t.Fatal(err)
	}
	if len(seen) == 0 {
		t.Fatal("no progress was reported")
	}
	last := seen[len(seen)-1]
	if last.Percent() != 100 {
		t.Errorf("final progress is %d%%, want 100 (%+v)", last.Percent(), last)
	}
}

func TestPercentIsUnknownRatherThanZeroWithoutATotal(t *testing.T) {
	// A server that sends no Content-Length is common, and reporting 0% for
	// a download that is nearly finished is worse than admitting ignorance.
	if got := (Progress{BytesDone: 500}).Percent(); got != -1 {
		t.Errorf("Percent with no total = %d, want -1", got)
	}
}

func TestNoCacheDirectoryIsAClearRefusal(t *testing.T) {
	f := &Fetcher{}
	_, err := f.Ensure(t.Context(), "https://example.invalid/model.gguf", false)
	if err == nil || !strings.Contains(err.Error(), "cache directory") {
		t.Errorf("err = %v, want a complaint about the cache directory", err)
	}
}

func TestABareModelIdIsLeftToTheEngine(t *testing.T) {
	// "Qwen/Qwen2.5-7B-Instruct" is a documented source that vLLM and SGLang
	// resolve themselves. Treating it as a missing file broke every existing
	// vLLM profile, which is how this was found.
	f := &Fetcher{CacheDir: t.TempDir()}
	for _, id := range []string{
		"Qwen/Qwen2.5-7B-Instruct",
		"meta-llama/Llama-3.3-70B-Instruct",
	} {
		got, err := f.Ensure(t.Context(), id, false)
		if err != nil {
			t.Errorf("%s was refused: %v", id, err)
		}
		if got != id {
			t.Errorf("%s was rewritten to %q", id, got)
		}
	}
}

func TestAnAbsolutePathThatIsMissingIsStillAnError(t *testing.T) {
	// The other half: an operator writing a path means a file, and a missing
	// one should be said plainly rather than handed to an engine.
	f := &Fetcher{CacheDir: t.TempDir()}
	if _, err := f.Ensure(t.Context(), "/models/definitely-not-here.gguf", false); !errors.Is(err, ErrNotLocal) {
		t.Errorf("err = %v, want ErrNotLocal", err)
	}
}

// A file whose sha256 the registry published is served only as that file:
// changed upstream after the model was added, it is refused with both
// checksums named, and a cached copy that is not it is fetched again.
func TestAFileThatIsNotWhatTheRegistryPublishedIsRefused(t *testing.T) {
	body := "the weights as published"
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		_, _ = w.Write([]byte(body))
	}))
	defer srv.Close()
	sum := sha256.Sum256([]byte(body))
	want := hex.EncodeToString(sum[:])
	f := &Fetcher{CacheDir: t.TempDir()}
	src := srv.URL + "/model.gguf"

	path, err := f.EnsureChecked(t.Context(), src, false, strings.ToUpper(want))
	if err != nil {
		t.Fatal(err)
	}
	if readChecksum(path) != want {
		t.Errorf("the checksum was not kept beside it: %q", readChecksum(path))
	}

	// The cached copy changed under it: fetched again, and the right one kept.
	if err := os.WriteFile(path, []byte("tampered"), 0o644); err != nil {
		t.Fatal(err)
	}
	_ = os.Remove(checksumPath(path))
	if _, err := f.EnsureChecked(t.Context(), src, false, want); err != nil {
		t.Fatal(err)
	}
	if b, _ := os.ReadFile(path); string(b) != body {
		t.Errorf("the cached copy that was not it was served: %q", b)
	}

	// Upstream changed after the model was added.
	body = "weights someone replaced"
	f2 := &Fetcher{CacheDir: t.TempDir()}
	_, err = f2.EnsureChecked(t.Context(), src, false, want)
	if err == nil || !strings.Contains(err.Error(), want) || !strings.Contains(err.Error(), "sha256") {
		t.Fatalf("not refused, or without both checksums: %v", err)
	}
	if _, err := os.Stat(CachePathFor(f2.CacheDir, src)); !os.IsNotExist(err) {
		t.Error("the file that was not it was kept")
	}
	// With nothing to check against, it is fetched as before.
	if _, err := f2.Ensure(t.Context(), src, false); err != nil {
		t.Error(err)
	}
}

// A cached file checked when it was fetched, that another model names with
// another checksum, is refused and left for the model it is.
func TestACheckedFileAnotherChecksumNamesIsRefusedAndLeft(t *testing.T) {
	body := "the weights"
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { _, _ = w.Write([]byte(body)) }))
	defer srv.Close()
	sum := sha256.Sum256([]byte(body))
	f := &Fetcher{CacheDir: t.TempDir()}
	src := srv.URL + "/m.gguf"
	path, err := f.EnsureChecked(t.Context(), src, false, hex.EncodeToString(sum[:]))
	if err != nil {
		t.Fatal(err)
	}
	wrong := strings.Repeat("0", 64)
	if _, err := f.EnsureChecked(t.Context(), src, false, wrong); err == nil || !strings.Contains(err.Error(), wrong) {
		t.Fatalf("not refused: %v", err)
	}
	if b, err := os.ReadFile(path); err != nil || string(b) != body {
		t.Errorf("the checked file was not left: %v", err)
	}
}