internal/hub/download_test.go
Rama main · réplica de lectura del repositorio original.
package hub
import (
"bytes"
"context"
"fmt"
"io"
"log"
"net/http"
"net/http/httptest"
"os"
"path/filepath"
"strconv"
"strings"
"sync/atomic"
"testing"
"time"
)
// fakeSource is a Source backed by an httptest server, so the download manager
// can be exercised end to end without touching Hugging Face.
type fakeSource struct {
base string
files []File
}
func (f *fakeSource) ID() string { return "fake" }
func (f *fakeSource) Name() string { return "Fake" }
func (f *fakeSource) Search(context.Context, Query) (Page, error) {
return Page{Models: []Model{{ID: "o/r", Source: "fake"}}}, nil
}
func (f *fakeSource) Detail(_ context.Context, repo string) (ModelDetail, error) {
return ModelDetail{Model: Model{ID: repo, Source: "fake"}, Files: f.files}, nil
}
func (f *fakeSource) FileURL(repo, path string) (string, error) {
return f.base + "/" + repo + "/" + path, nil
}
func (f *fakeSource) CheckURL(raw string) error {
if !strings.HasPrefix(raw, f.base) {
return fmt.Errorf("refusing %s", raw)
}
return nil
}
func (f *fakeSource) Authorize(*http.Request) {}
// serveBlob answers a GET with Range support, optionally cutting the response
// short after cutAfter bytes to simulate a dropped connection.
func serveBlob(w http.ResponseWriter, r *http.Request, blob []byte, cutAfter int) {
start := 0
if rng := r.Header.Get("Range"); strings.HasPrefix(rng, "bytes=") {
n, err := strconv.Atoi(strings.TrimSuffix(strings.TrimPrefix(rng, "bytes="), "-"))
if err == nil {
start = n
}
if start >= len(blob) {
w.WriteHeader(http.StatusRequestedRangeNotSatisfiable)
return
}
w.Header().Set("Content-Range", fmt.Sprintf("bytes %d-%d/%d", start, len(blob)-1, len(blob)))
w.WriteHeader(http.StatusPartialContent)
}
body := blob[start:]
if cutAfter > 0 && cutAfter < len(body) {
body = body[:cutAfter]
w.Write(body)
// No Content-Length was set, so a short write plus a flush and a panic
// closes the connection mid-body the way a dropped line does.
if fl, ok := w.(http.Flusher); ok {
fl.Flush()
}
panic(http.ErrAbortHandler)
}
w.Write(body)
}
// discardLogger silences httptest's "connection hijacked/aborted" noise for the
// tests that deliberately drop a connection.
func discardLogger() *log.Logger { return log.New(io.Discard, "", 0) }
func waitJob(t *testing.T, m *Manager, id string) Job {
t.Helper()
deadline := time.Now().Add(20 * time.Second)
for time.Now().Before(deadline) {
j, ok := m.Job(id)
if ok && j.Done() {
return j
}
time.Sleep(10 * time.Millisecond)
}
j, _ := m.Job(id)
t.Fatalf("job %s did not finish, stuck in phase %q", id, j.Phase)
return Job{}
}
func blobOf(n int) []byte {
b := make([]byte, n)
for i := range b {
b[i] = byte(i % 251)
}
return b
}
func TestManager_DownloadsEveryFile(t *testing.T) {
a, b := blobOf(4096), blobOf(1500)
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if strings.HasSuffix(r.URL.Path, "b.gguf") {
serveBlob(w, r, b, 0)
return
}
serveBlob(w, r, a, 0)
}))
defer srv.Close()
root := t.TempDir()
src := &fakeSource{base: srv.URL, files: []File{
{Path: "a-00001-of-00002.gguf", SizeBytes: int64(len(a))},
{Path: "sub/b.gguf", SizeBytes: int64(len(b))},
}}
var completed atomic.Int32
m := NewManager(func() string { return root }, nil, src)
m.OnComplete = func(Job) error { completed.Add(1); return nil }
id, err := m.Start(context.Background(), StartRequest{
Source: "fake", Repo: "o/r", Files: []string{"a-00001-of-00002.gguf", "sub/b.gguf"},
})
if err != nil {
t.Fatal(err)
}
job := waitJob(t, m, id)
if job.Phase != PhaseDone {
t.Fatalf("phase = %q, err = %q", job.Phase, job.Err)
}
dir := filepath.Join(root, "o", "r")
if got, err := os.ReadFile(filepath.Join(dir, "a-00001-of-00002.gguf")); err != nil || !bytes.Equal(got, a) {
t.Errorf("shard 1 mismatch: err %v, %d bytes", err, len(got))
}
if got, err := os.ReadFile(filepath.Join(dir, "sub", "b.gguf")); err != nil || !bytes.Equal(got, b) {
t.Errorf("nested file mismatch: err %v, %d bytes", err, len(got))
}
if job.Downloaded != job.Total {
t.Errorf("progress ended at %d/%d", job.Downloaded, job.Total)
}
if completed.Load() != 1 {
t.Errorf("OnComplete ran %d times, want 1", completed.Load())
}
// Nothing partial may survive: model discovery walks this folder.
if _, err := os.Stat(filepath.Join(dir, "a-00001-of-00002.gguf"+partSuffix)); !os.IsNotExist(err) {
t.Error(".part file left behind after a successful download")
}
// Nor may the journal: there is nothing left to resume.
if _, err := os.Stat(filepath.Join(dir, journalName)); !os.IsNotExist(err) {
t.Error("the download journal survived a finished job")
}
}
func TestManager_ResumesAfterDroppedConnection(t *testing.T) {
blob := blobOf(8192)
var hits atomic.Int32
var sawRange atomic.Bool
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.Header.Get("Range") != "" {
sawRange.Store(true)
}
// Cut the first attempt short; serve the rest normally.
cut := 0
if hits.Add(1) == 1 {
cut = 3000
}
serveBlob(w, r, blob, cut)
}))
defer srv.Close()
srv.Config.ErrorLog = discardLogger()
root := t.TempDir()
src := &fakeSource{base: srv.URL, files: []File{{Path: "m.gguf", SizeBytes: int64(len(blob))}}}
m := NewManager(func() string { return root }, nil, src)
id, err := m.Start(context.Background(), StartRequest{Source: "fake", Repo: "o/r", Files: []string{"m.gguf"}})
if err != nil {
t.Fatal(err)
}
job := waitJob(t, m, id)
if job.Phase != PhaseDone {
t.Fatalf("phase = %q, err = %q", job.Phase, job.Err)
}
if !sawRange.Load() {
t.Error("retry did not send a Range header — the transfer restarted from zero")
}
got, err := os.ReadFile(filepath.Join(root, "o", "r", "m.gguf"))
if err != nil {
t.Fatal(err)
}
if !bytes.Equal(got, blob) {
t.Errorf("resumed file is %d bytes and does not match the source", len(got))
}
}
func TestManager_SkipsFileAlreadyOnDisk(t *testing.T) {
blob := blobOf(2048)
var hits atomic.Int32
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
hits.Add(1)
serveBlob(w, r, blob, 0)
}))
defer srv.Close()
root := t.TempDir()
dir := filepath.Join(root, "o", "r")
if err := os.MkdirAll(dir, 0o755); err != nil {
t.Fatal(err)
}
if err := os.WriteFile(filepath.Join(dir, "m.gguf"), blob, 0o644); err != nil {
t.Fatal(err)
}
src := &fakeSource{base: srv.URL, files: []File{{Path: "m.gguf", SizeBytes: int64(len(blob))}}}
m := NewManager(func() string { return root }, nil, src)
id, err := m.Start(context.Background(), StartRequest{Source: "fake", Repo: "o/r", Files: []string{"m.gguf"}})
if err != nil {
t.Fatal(err)
}
job := waitJob(t, m, id)
if job.Phase != PhaseDone {
t.Fatalf("phase = %q, err = %q", job.Phase, job.Err)
}
if hits.Load() != 0 {
t.Errorf("refetched a file already on disk (%d requests)", hits.Load())
}
if !job.Files[0].Skipped || job.Downloaded != job.Total {
t.Errorf("skipped file not accounted for: %+v (%d/%d)", job.Files[0], job.Downloaded, job.Total)
}
}
func TestManager_GatedRepoFailsFastAsAuthError(t *testing.T) {
var hits atomic.Int32
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
hits.Add(1)
w.WriteHeader(http.StatusForbidden)
w.Write([]byte("access to model o/r is restricted"))
}))
defer srv.Close()
root := t.TempDir()
src := &fakeSource{base: srv.URL, files: []File{{Path: "m.gguf", SizeBytes: 10}}}
m := NewManager(func() string { return root }, nil, src)
id, err := m.Start(context.Background(), StartRequest{Source: "fake", Repo: "o/r", Files: []string{"m.gguf"}})
if err != nil {
t.Fatal(err)
}
job := waitJob(t, m, id)
if job.Phase != PhaseError || !job.Gated {
t.Fatalf("phase = %q gated = %v, want error + gated", job.Phase, job.Gated)
}
// A license wall is terminal: retrying it five times only delays telling
// the user to accept the license.
if hits.Load() != 1 {
t.Errorf("retried a 403 %d times", hits.Load())
}
if !strings.Contains(job.Err, "accept the license") {
t.Errorf("error does not tell the user what to do: %q", job.Err)
}
}
// heldServer serves the first `head` bytes of blob and then holds the
// connection open until the test lets go, so a transfer can be caught mid-flight
// and interrupted deterministically.
func heldServer(t *testing.T, blob []byte, head int) (*httptest.Server, func()) {
t.Helper()
release := make(chan struct{})
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Write(blob[:head])
if fl, ok := w.(http.Flusher); ok {
fl.Flush()
}
<-release
}))
srv.Config.ErrorLog = discardLogger()
return srv, func() { close(release); srv.Close() }
}
func waitPhase(t *testing.T, m *Manager, id, phase string) Job {
t.Helper()
deadline := time.Now().Add(20 * time.Second)
for time.Now().Before(deadline) {
if j, ok := m.Job(id); ok && j.Phase == phase {
return j
}
time.Sleep(10 * time.Millisecond)
}
j, _ := m.Job(id)
t.Fatalf("job %s never reached %q, stuck in %q", id, phase, j.Phase)
return Job{}
}
// waitPartial waits until there are real bytes in the .part file. Neither of
// the job's own signals can say that: PhaseDownloading is set before the
// connection is even made (so the UI says "downloading" rather than sitting on
// "resolving" through DNS and TLS), and Downloaded is reported at most once per
// progressEvery, so a transfer smaller than a MiB never bumps it off zero. The
// file on disk is the thing this test is actually about.
func waitPartial(t *testing.T, part string) {
t.Helper()
deadline := time.Now().Add(20 * time.Second)
for time.Now().Before(deadline) {
if st, err := os.Stat(part); err == nil && st.Size() > 0 {
return
}
time.Sleep(10 * time.Millisecond)
}
t.Fatalf("no bytes landed in %s", part)
}
func TestManager_PauseKeepsPartialAndResumeFinishes(t *testing.T) {
blob := blobOf(1 << 20)
srv, stop := heldServer(t, blob, 4096)
root := t.TempDir()
src := &fakeSource{base: srv.URL, files: []File{{Path: "m.gguf", SizeBytes: int64(len(blob))}}}
m := NewManager(func() string { return root }, nil, src)
id, err := m.Start(context.Background(), StartRequest{Source: "fake", Repo: "o/r", Files: []string{"m.gguf"}})
if err != nil {
t.Fatal(err)
}
part := filepath.Join(root, "o", "r", "m.gguf"+partSuffix)
waitPartial(t, part)
if err := m.Pause(id); err != nil {
t.Fatal(err)
}
waitPhase(t, m, id, PhasePaused)
// The partial must survive — that is the entire difference from cancel.
st, err := os.Stat(part)
if err != nil {
t.Fatalf("pause discarded the partial file: %v", err)
}
if st.Size() == 0 {
t.Fatal("pause left an empty partial")
}
if _, err := os.Stat(filepath.Join(root, "o", "r", "m.gguf")); !os.IsNotExist(err) {
t.Error("a paused download left a file that looks complete")
}
// Resume against a server that now answers in full: the job keeps its id and
// finishes from the bytes already on disk.
stop()
full := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
serveBlob(w, r, blob, 0)
}))
defer full.Close()
src.base = full.URL
if err := m.Resume(id); err != nil {
t.Fatal(err)
}
job := waitJob(t, m, id)
if job.Phase != PhaseDone {
t.Fatalf("phase = %q, err = %q", job.Phase, job.Err)
}
got, err := os.ReadFile(filepath.Join(root, "o", "r", "m.gguf"))
if err != nil || !bytes.Equal(got, blob) {
t.Errorf("resumed file mismatch: err %v, %d bytes", err, len(got))
}
}
func TestManager_CancelDiscardsPartial(t *testing.T) {
blob := blobOf(1 << 20)
srv, stop := heldServer(t, blob, 4096)
defer stop()
root := t.TempDir()
src := &fakeSource{base: srv.URL, files: []File{{Path: "m.gguf", SizeBytes: int64(len(blob))}}}
m := NewManager(func() string { return root }, nil, src)
id, err := m.Start(context.Background(), StartRequest{Source: "fake", Repo: "o/r", Files: []string{"m.gguf"}})
if err != nil {
t.Fatal(err)
}
waitPhase(t, m, id, PhaseDownloading)
if err := m.Cancel(id); err != nil {
t.Fatal(err)
}
job := waitJob(t, m, id)
if job.Phase != PhaseCanceled {
t.Fatalf("phase = %q, want canceled (err %q)", job.Phase, job.Err)
}
// Cancel means discard: an abandoned .part is a multi-GB orphan nothing ever
// mentions again.
if _, err := os.Stat(filepath.Join(root, "o", "r", "m.gguf"+partSuffix)); !os.IsNotExist(err) {
t.Errorf("cancel kept the partial file: %v", err)
}
if _, err := os.Stat(filepath.Join(root, "o", "r", "m.gguf")); !os.IsNotExist(err) {
t.Error("a canceled download left a file that looks complete")
}
// The empty folders it created go with it.
if _, err := os.Stat(filepath.Join(root, "o")); !os.IsNotExist(err) {
t.Error("cancel left the publisher folder behind")
}
}
func TestManager_CancelDiscardsFilesTheJobFinished(t *testing.T) {
// Two files: one lands complete, the second is still running when the job is
// canceled. A lone shard is not a model, so the finished one must go too —
// but a file that was already there before the job (Skipped) must not.
done, running := blobOf(2048), blobOf(1<<20)
release := make(chan struct{})
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if strings.HasSuffix(r.URL.Path, "a.gguf") {
serveBlob(w, r, done, 0)
return
}
w.Write(running[:4096])
if fl, ok := w.(http.Flusher); ok {
fl.Flush()
}
<-release
}))
defer srv.Close()
defer close(release)
srv.Config.ErrorLog = discardLogger()
root := t.TempDir()
dir := filepath.Join(root, "o", "r")
if err := os.MkdirAll(dir, 0o755); err != nil {
t.Fatal(err)
}
pre := blobOf(512)
if err := os.WriteFile(filepath.Join(dir, "pre.gguf"), pre, 0o644); err != nil {
t.Fatal(err)
}
src := &fakeSource{base: srv.URL, files: []File{
{Path: "pre.gguf", SizeBytes: int64(len(pre))},
{Path: "a.gguf", SizeBytes: int64(len(done))},
{Path: "b.gguf", SizeBytes: int64(len(running))},
}}
m := NewManager(func() string { return root }, nil, src)
id, err := m.Start(context.Background(), StartRequest{
Source: "fake", Repo: "o/r", Files: []string{"pre.gguf", "a.gguf", "b.gguf"},
})
if err != nil {
t.Fatal(err)
}
// b.gguf is the one that hangs, so wait until a.gguf has landed.
deadline := time.Now().Add(20 * time.Second)
for time.Now().Before(deadline) {
if _, err := os.Stat(filepath.Join(dir, "a.gguf")); err == nil {
break
}
time.Sleep(10 * time.Millisecond)
}
if err := m.Cancel(id); err != nil {
t.Fatal(err)
}
waitJob(t, m, id)
if _, err := os.Stat(filepath.Join(dir, "a.gguf")); !os.IsNotExist(err) {
t.Error("cancel kept a file the job itself downloaded")
}
if got, err := os.ReadFile(filepath.Join(dir, "pre.gguf")); err != nil || !bytes.Equal(got, pre) {
t.Errorf("cancel deleted a file that predated the job: err %v", err)
}
}
func TestSweepPartials_AgeGated(t *testing.T) {
root := t.TempDir()
dir := filepath.Join(root, "o", "r")
if err := os.MkdirAll(dir, 0o755); err != nil {
t.Fatal(err)
}
old := filepath.Join(dir, "old.gguf"+partSuffix)
fresh := filepath.Join(dir, "fresh.gguf"+partSuffix)
keep := filepath.Join(dir, "model.gguf")
for _, p := range []string{old, fresh, keep} {
if err := os.WriteFile(p, blobOf(1024), 0o644); err != nil {
t.Fatal(err)
}
}
stale := time.Now().Add(-48 * time.Hour)
if err := os.Chtimes(old, stale, stale); err != nil {
t.Fatal(err)
}
n, freed := SweepPartials(root, 24*time.Hour, nil)
if n != 1 || freed != 1024 {
t.Errorf("swept %d files / %d bytes, want 1 / 1024", n, freed)
}
if _, err := os.Stat(old); !os.IsNotExist(err) {
t.Error("the stale partial survived the sweep")
}
// A partial is resumable, so a recent one is work in progress, not garbage.
if _, err := os.Stat(fresh); err != nil {
t.Error("the sweep ate a fresh partial")
}
if _, err := os.Stat(keep); err != nil {
t.Error("the sweep touched a real model file")
}
}
func TestManager_QueuesSecondJobForSameRepo(t *testing.T) {
release := make(chan struct{})
var hits atomic.Int32
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
hits.Add(1)
<-release
_, _ = w.Write(make([]byte, 1<<20))
}))
defer srv.Close()
root := t.TempDir()
src := &fakeSource{base: srv.URL, files: []File{
{Path: "a.gguf", SizeBytes: 1 << 20},
{Path: "b.gguf", SizeBytes: 1 << 20},
}}
m := NewManager(func() string { return root }, nil, src)
first, err := m.Start(context.Background(), StartRequest{Source: "fake", Repo: "o/r", Files: []string{"a.gguf"}})
if err != nil {
t.Fatal(err)
}
// Picking a second file from the same repo is the ordinary case (two quants,
// or a model and its projector) — it waits its turn instead of erroring.
second, err := m.Start(context.Background(), StartRequest{Source: "fake", Repo: "o/r", Files: []string{"b.gguf"}})
if err != nil {
t.Fatalf("second job for the same repo was refused: %v", err)
}
if j, _ := m.Job(second); j.Phase != PhaseQueued {
t.Errorf("queued job is in phase %q, want %q", j.Phase, PhaseQueued)
}
if got := hits.Load(); got > 1 {
t.Errorf("%d concurrent transfers for one repo, want 1", got)
}
close(release)
waitJob(t, m, first)
waitJob(t, m, second)
if j, _ := m.Job(second); j.Phase != PhaseDone {
t.Errorf("queued job ended in %q (%s), want %q", j.Phase, j.Err, PhaseDone)
}
if _, err := os.Stat(filepath.Join(root, "o", "r", "b.gguf")); err != nil {
t.Errorf("the queued job never ran: %v", err)
}
}
// A job canceled before its turn takes nothing with it: the bytes under that
// repo belong to whatever is running there.
func TestManager_CancelQueuedJob(t *testing.T) {
release := make(chan struct{})
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
<-release
}))
defer srv.Close()
defer close(release)
root := t.TempDir()
src := &fakeSource{base: srv.URL, files: []File{
{Path: "a.gguf", SizeBytes: 1 << 20},
{Path: "b.gguf", SizeBytes: 1 << 20},
}}
m := NewManager(func() string { return root }, nil, src)
first, err := m.Start(context.Background(), StartRequest{Source: "fake", Repo: "o/r", Files: []string{"a.gguf"}})
if err != nil {
t.Fatal(err)
}
second, err := m.Start(context.Background(), StartRequest{Source: "fake", Repo: "o/r", Files: []string{"b.gguf"}})
if err != nil {
t.Fatal(err)
}
if err := m.Cancel(second); err != nil {
t.Fatalf("canceling a queued job: %v", err)
}
if j, _ := m.Job(second); j.Phase != PhaseCanceled {
t.Errorf("canceled queued job is in phase %q", j.Phase)
}
// The running job is untouched by the other job's cancel.
if j, _ := m.Job(first); j.Done() {
t.Errorf("canceling the queued job stopped the running one (phase %q)", j.Phase)
}
m.Cancel(first)
waitJob(t, m, first)
}
// waitForBytes waits until size(path) is positive, so a test can pause a held
// transfer knowing some bytes really landed rather than racing the cancel.
func waitForBytes(t *testing.T, path string) int64 {
t.Helper()
deadline := time.Now().Add(20 * time.Second)
for time.Now().Before(deadline) {
if n := partSize(path); n > 0 {
return n
}
time.Sleep(10 * time.Millisecond)
}
t.Fatalf("no bytes ever landed in %s", path)
return 0
}
// A download interrupted by a process kill leaves its `.part` and a journal
// record beside it. Restore has to turn that back into a paused job whose
// progress is read off the disk, and the existing Resume verb has to finish it
// — this is the whole point of the journal.
func TestManager_RestoreResumesPartial(t *testing.T) {
blob := blobOf(1 << 20)
srv, stop := heldServer(t, blob, 4096)
root := t.TempDir()
src := &fakeSource{base: srv.URL, files: []File{{Path: "m.gguf", SizeBytes: int64(len(blob))}}}
m := NewManager(func() string { return root }, nil, src)
id, err := m.Start(context.Background(), StartRequest{
Source: "fake", Repo: "o/r", Files: []string{"m.gguf"}, Label: "m.gguf",
})
if err != nil {
t.Fatal(err)
}
part := filepath.Join(root, "o", "r", "m.gguf"+partSuffix)
waitForBytes(t, part)
if err := m.Pause(id); err != nil {
t.Fatal(err)
}
waitPhase(t, m, id, PhasePaused)
// A fresh Manager over the same folder stands in for a restarted process:
// the old job list is gone, only the disk survived.
m2 := NewManager(func() string { return root }, nil, src)
if n := m2.Restore(24 * time.Hour); n != 1 {
t.Fatalf("Restore returned %d jobs, want 1", n)
}
job, ok := m2.Job(id)
if !ok {
t.Fatalf("the journaled job %s was not restored", id)
}
if job.Phase != PhasePaused {
t.Errorf("restored job is %q, want paused", job.Phase)
}
if job.Label != "m.gguf" || job.Repo != "o/r" || job.Source != "fake" {
t.Errorf("restored job lost its identity: %+v", job)
}
if job.Downloaded <= 0 || job.Downloaded >= job.Total {
t.Errorf("restored progress = %d/%d, want a partial count", job.Downloaded, job.Total)
}
if job.Files[0].Done != job.Downloaded {
t.Errorf("file progress %d disagrees with job progress %d", job.Files[0].Done, job.Downloaded)
}
// Resume against a server that answers in full; the bytes already on disk
// must carry over rather than being fetched again.
stop()
full := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.Header.Get("Range") == "" {
t.Error("resume did not send a Range header")
}
serveBlob(w, r, blob, 0)
}))
defer full.Close()
src.base = full.URL
if err := m2.Resume(id); err != nil {
t.Fatal(err)
}
got := waitJob(t, m2, id)
if got.Phase != PhaseDone {
t.Fatalf("resumed job ended %q (%s), want done", got.Phase, got.Err)
}
b, err := os.ReadFile(filepath.Join(root, "o", "r", "m.gguf"))
if err != nil || !bytes.Equal(b, blob) {
t.Errorf("resumed file mismatch: err %v, %d bytes", err, len(b))
}
}
// Restore must rebuild per-file progress: a shard that finished before the kill
// counts complete, the one that was running counts its `.part`, and the two add
// up to the row's progress.
func TestManager_RestoreKeepsFinishedFiles(t *testing.T) {
finished, running := blobOf(2048), blobOf(1<<20)
release := make(chan struct{})
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if strings.HasSuffix(r.URL.Path, "a.gguf") {
serveBlob(w, r, finished, 0)
return
}
w.Write(running[:4096])
if fl, ok := w.(http.Flusher); ok {
fl.Flush()
}
<-release
}))
defer srv.Close()
defer close(release)
srv.Config.ErrorLog = discardLogger()
root := t.TempDir()
dir := filepath.Join(root, "o", "r")
src := &fakeSource{base: srv.URL, files: []File{
{Path: "a.gguf", SizeBytes: int64(len(finished))},
{Path: "b.gguf", SizeBytes: int64(len(running))},
}}
m := NewManager(func() string { return root }, nil, src)
id, err := m.Start(context.Background(), StartRequest{
Source: "fake", Repo: "o/r", Files: []string{"a.gguf", "b.gguf"},
})
if err != nil {
t.Fatal(err)
}
// Wait for a.gguf to land and b.gguf's partial to start, then stop mid-flight.
waitForBytes(t, filepath.Join(dir, "b.gguf"+partSuffix))
deadline := time.Now().Add(20 * time.Second)
for time.Now().Before(deadline) && existingSize(filepath.Join(dir, "a.gguf")) < int64(len(finished)) {
time.Sleep(10 * time.Millisecond)
}
if existingSize(filepath.Join(dir, "a.gguf")) != int64(len(finished)) {
t.Fatal("the first shard never landed")
}
if err := m.Pause(id); err != nil {
t.Fatal(err)
}
waitPhase(t, m, id, PhasePaused)
m2 := NewManager(func() string { return root }, nil, src)
if n := m2.Restore(24 * time.Hour); n != 1 {
t.Fatalf("Restore returned %d jobs, want 1", n)
}
job, _ := m2.Job(id)
if !job.Files[0].Skipped || job.Files[0].Done != int64(len(finished)) {
t.Errorf("finished shard not restored as complete: %+v", job.Files[0])
}
if job.Files[1].Done <= 0 || job.Files[1].Done >= int64(len(running)) {
t.Errorf("partial shard progress = %d, want a partial count", job.Files[1].Done)
}
if want := int64(len(finished)) + job.Files[1].Done; job.Downloaded != want {
t.Errorf("job progress = %d, want %d", job.Downloaded, want)
}
if job.Total != int64(len(finished)+len(running)) {
t.Errorf("job total = %d, want %d", job.Total, len(finished)+len(running))
}
// The finished shard is removed while the job is parked. Resume must refetch
// it rather than trust the restored Skipped flag — that flag described the
// disk at restore time, not now.
if err := os.Remove(filepath.Join(dir, "a.gguf")); err != nil {
t.Fatal(err)
}
full := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if strings.HasSuffix(r.URL.Path, "a.gguf") {
serveBlob(w, r, finished, 0)
return
}
serveBlob(w, r, running, 0)
}))
defer full.Close()
src.base = full.URL
if err := m2.Resume(id); err != nil {
t.Fatal(err)
}
got := waitJob(t, m2, id)
if got.Phase != PhaseDone {
t.Fatalf("resumed job ended %q (%s), want done", got.Phase, got.Err)
}
for name, want := range map[string][]byte{"a.gguf": finished, "b.gguf": running} {
b, err := os.ReadFile(filepath.Join(dir, name))
if err != nil || !bytes.Equal(b, want) {
t.Errorf("%s after resume: err %v, %d bytes", name, err, len(b))
}
}
}
// Two jobs share one repo folder and therefore one journal file, and both have
// to come back after a restart.
func TestManager_RestoreKeepsBothJobsForOneRepo(t *testing.T) {
blob := blobOf(1 << 20)
srv, stop := heldServer(t, blob, 4096)
defer stop()
root := t.TempDir()
src := &fakeSource{base: srv.URL, files: []File{
{Path: "a.gguf", SizeBytes: int64(len(blob))},
{Path: "b.gguf", SizeBytes: int64(len(blob))},
}}
m := NewManager(func() string { return root }, nil, src)
first, err := m.Start(context.Background(), StartRequest{Source: "fake", Repo: "o/r", Files: []string{"a.gguf"}, Label: "a"})
if err != nil {
t.Fatal(err)
}
second, err := m.Start(context.Background(), StartRequest{Source: "fake", Repo: "o/r", Files: []string{"b.gguf"}, Label: "b"})
if err != nil {
t.Fatal(err)
}
waitForBytes(t, filepath.Join(root, "o", "r", "a.gguf"+partSuffix))
if err := m.Pause(first); err != nil {
t.Fatal(err)
}
waitPhase(t, m, first, PhasePaused)
// The pause frees the repo, so the queued job starts on its own.
waitForBytes(t, filepath.Join(root, "o", "r", "b.gguf"+partSuffix))
if err := m.Pause(second); err != nil {
t.Fatal(err)
}
waitPhase(t, m, second, PhasePaused)
m2 := NewManager(func() string { return root }, nil, src)
if n := m2.Restore(24 * time.Hour); n != 2 {
t.Fatalf("Restore returned %d jobs, want 2", n)
}
for _, id := range []string{first, second} {
if _, ok := m2.Job(id); !ok {
t.Errorf("job %s was not restored", id)
}
}
}
// A record whose files are all on disk at their recorded revision is work that
// already finished — the process died between the last rename and the journal
// clear. It must not come back as a paused row.
func TestManager_RestoreDropsFinishedRecord(t *testing.T) {
root := t.TempDir()
dir := filepath.Join(root, "o", "r")
if err := os.MkdirAll(dir, 0o755); err != nil {
t.Fatal(err)
}
blob := blobOf(100)
if err := os.WriteFile(filepath.Join(dir, "m.gguf"), blob, 0o644); err != nil {
t.Fatal(err)
}
m := NewManager(func() string { return root }, nil)
m.recordManifest(dir, "m.gguf", manifestEntry{Size: int64(len(blob)), OID: "oid-1"})
m.recordDownload(dir, downloadRecord{
ID: "dl-finished", Source: "fake", Repo: "o/r", Started: time.Now(),
Files: []recordFile{{Path: "m.gguf", Size: int64(len(blob)), OID: "oid-1"}},
})
if n := m.Restore(24 * time.Hour); n != 0 {
t.Errorf("restored %d jobs for a finished download", n)
}
if _, err := os.Stat(filepath.Join(dir, journalName)); !os.IsNotExist(err) {
t.Error("the finished job's journal record survived restore")
}
}
// A record with no bytes anywhere has nothing to resume: a job killed before
// its first byte, or a `.part` the sweep took. It is cleared, not rendered as
// an empty paused row for the life of the install.
func TestManager_RestoreDropsEmptyRecord(t *testing.T) {
root := t.TempDir()
dir := filepath.Join(root, "o", "r")
m := NewManager(func() string { return root }, nil)
m.recordDownload(dir, downloadRecord{
ID: "dl-empty", Source: "fake", Repo: "o/r", Started: time.Now(),
Files: []recordFile{{Path: "m.gguf", Size: 4096}},
})
if n := m.Restore(24 * time.Hour); n != 0 {
t.Errorf("restored %d jobs with no bytes on disk", n)
}
if _, err := os.Stat(filepath.Join(dir, journalName)); !os.IsNotExist(err) {
t.Error("an empty job's journal record survived restore")
}
}
// Restore and SweepPartials share the age gate: a `.part` old enough for the
// sweep is not progress, or the row would promise bytes about to be deleted.
func TestManager_RestoreIgnoresAgedPartial(t *testing.T) {
root := t.TempDir()
dir := filepath.Join(root, "o", "r")
if err := os.MkdirAll(dir, 0o755); err != nil {
t.Fatal(err)
}
part := filepath.Join(dir, "m.gguf"+partSuffix)
if err := os.WriteFile(part, blobOf(100), 0o644); err != nil {
t.Fatal(err)
}
old := time.Now().Add(-48 * time.Hour)
if err := os.Chtimes(part, old, old); err != nil {
t.Fatal(err)
}
m := NewManager(func() string { return root }, nil)
m.recordDownload(dir, downloadRecord{
ID: "dl-old", Source: "fake", Repo: "o/r", Started: time.Now(),
Files: []recordFile{{Path: "m.gguf", Size: 4096}},
})
if n := m.Restore(24 * time.Hour); n != 0 {
t.Errorf("restored %d job(s) off a partial the sweep will delete", n)
}
}
// LocalFiles is what puts "downloaded" on a picker row, so it has to agree with
// what the transfer actually leaves behind: a finished file counts, a `.part`
// does not.
func TestManager_LocalFiles(t *testing.T) {
root := t.TempDir()
dir := filepath.Join(root, "o", "r", "sub")
if err := os.MkdirAll(dir, 0o755); err != nil {
t.Fatal(err)
}
for name, size := range map[string]int{
filepath.Join(root, "o", "r", "done.gguf"): 2048,
filepath.Join(root, "o", "r", "half.gguf"+partSuffix): 10,
filepath.Join(dir, "nested.gguf"): 7,
} {
if err := os.WriteFile(name, make([]byte, size), 0o644); err != nil {
t.Fatal(err)
}
}
m := NewManager(func() string { return root }, nil)
// The journal is bookkeeping like the manifest: it must not read as a repo
// file, or the picker would show a row for it.
m.recordDownload(filepath.Join(root, "o", "r"), downloadRecord{
ID: "dl-x", Source: "fake", Repo: "o/r", Started: time.Now(),
Files: []recordFile{{Path: "done.gguf", Size: 2048}},
})
got := m.LocalFiles("o/r")
if got["done.gguf"].Size != 2048 {
t.Errorf("done.gguf = %d, want 2048", got["done.gguf"].Size)
}
if got["sub/nested.gguf"].Size != 7 {
t.Errorf("nested file missing or wrong size: %v", got)
}
if _, ok := got["half.gguf"]; ok {
t.Error("a .part was reported as a local file")
}
if _, ok := got[journalName]; ok {
t.Error("the download journal was reported as a local file")
}
// Nothing recorded these, so they carry no identity — which is what keeps
// a hand-copied file from ever being called stale.
if got["done.gguf"].OID != "" {
t.Errorf("an unrecorded file reported an identity: %q", got["done.gguf"].OID)
}
if len(m.LocalFiles("nobody/here")) != 0 {
t.Error("an unknown repo reported local files")
}
}
// A manifest is only believed while it still describes the file next to it.
// Recording an id and then letting the bytes change behind our back must drop
// the id, not keep vouching for a file that is gone.
func TestManager_LocalFilesManifest(t *testing.T) {
root := t.TempDir()
dir := filepath.Join(root, "o", "r")
if err := os.MkdirAll(dir, 0o755); err != nil {
t.Fatal(err)
}
if err := os.WriteFile(filepath.Join(dir, "a.gguf"), make([]byte, 100), 0o644); err != nil {
t.Fatal(err)
}
m := NewManager(func() string { return root }, nil)
m.recordManifest(dir, "a.gguf", manifestEntry{Size: 100, OID: "sha-one"})
if got := m.LocalFiles("o/r"); got["a.gguf"].OID != "sha-one" {
t.Errorf("recorded identity not reported: %+v", got["a.gguf"])
}
// The manifest itself is our bookkeeping, not a repo file.
if _, ok := m.LocalFiles("o/r")[manifestName]; ok {
t.Error("the manifest was reported as a downloaded file")
}
// Same name, different bytes, written by something that is not us.
if err := os.WriteFile(filepath.Join(dir, "a.gguf"), make([]byte, 140), 0o644); err != nil {
t.Fatal(err)
}
if got := m.LocalFiles("o/r"); got["a.gguf"].OID != "" {
t.Errorf("a replaced file kept the old identity: %+v", got["a.gguf"])
}
}
// The bug this whole path exists for: a repo re-uploads a quant under the same
// name, at a size a byte comparison cannot tell apart. Skipping it is what used
// to force the user to rename or delete the file by hand.
func TestManager_HaveCurrent(t *testing.T) {
root := t.TempDir()
dir := filepath.Join(root, "o", "r")
if err := os.MkdirAll(dir, 0o755); err != nil {
t.Fatal(err)
}
if err := os.WriteFile(filepath.Join(dir, "a.gguf"), make([]byte, 100), 0o644); err != nil {
t.Fatal(err)
}
m := NewManager(func() string { return root }, nil)
m.recordManifest(dir, "a.gguf", manifestEntry{Size: 100, OID: "old"})
man := readManifest(dir)
cases := []struct {
name string
f JobFile
force bool
want bool
}{
{"same revision is skipped", JobFile{Path: "a.gguf", Size: 100, OID: "old"}, false, true},
{"replaced upstream at the same size is refetched", JobFile{Path: "a.gguf", Size: 100, OID: "new"}, false, false},
{"force refetches regardless", JobFile{Path: "a.gguf", Size: 100, OID: "old"}, true, false},
{"a hub stating no id falls back to size", JobFile{Path: "a.gguf", Size: 100}, false, true},
{"a truncated copy is refetched", JobFile{Path: "a.gguf", Size: 4096, OID: "old"}, false, false},
{"a file we do not have is fetched", JobFile{Path: "b.gguf", Size: 100, OID: "new"}, false, false},
}
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
if got := m.haveCurrent(dir, tc.f, man, tc.force); got != tc.want {
t.Errorf("haveCurrent = %v, want %v", got, tc.want)
}
})
}
// No manifest at all — the hand-copied case. There is nothing to compare
// against, so the size test stands and the file is left alone.
if !m.haveCurrent(dir, JobFile{Path: "a.gguf", Size: 100, OID: "new"}, manifest{Files: map[string]manifestEntry{}}, false) {
t.Error("an unrecorded file was refetched on an id we never had")
}
}
func TestManager_RejectsFileNotInRepo(t *testing.T) {
root := t.TempDir()
src := &fakeSource{base: "http://127.0.0.1:1", files: []File{{Path: "m.gguf", SizeBytes: 10}}}
m := NewManager(func() string { return root }, nil, src)
// A path the hub never listed has no size, so neither the disk check nor
// the progress bar would have a denominator — and it is how a crafted
// request would try to reach an arbitrary file.
if _, err := m.Start(context.Background(), StartRequest{Source: "fake", Repo: "o/r", Files: []string{"../../etc/passwd"}}); err == nil {
t.Error("Start accepted a file that is not in the repo listing")
}
}
func TestManager_NoModelsRootIsAClearError(t *testing.T) {
src := &fakeSource{base: "http://127.0.0.1:1", files: []File{{Path: "m.gguf", SizeBytes: 10}}}
m := NewManager(func() string { return "" }, nil, src)
_, err := m.Start(context.Background(), StartRequest{Source: "fake", Repo: "o/r", Files: []string{"m.gguf"}})
if err == nil || !strings.Contains(err.Error(), "models folder") {
t.Errorf("err = %v, want a models-folder message", err)
}
}
// Clearing the panel's history drops the terminal rows and nothing else: a
// paused pull is outstanding work, and losing it from the list would be the
// same as losing the download.
func TestManager_ClearFinishedKeepsUnfinishedJobs(t *testing.T) {
m := NewManager(func() string { return t.TempDir() }, nil)
seed := func(id, phase string) {
m.jobs[id] = &Job{ID: id, Phase: phase}
m.order = append(m.order, id)
}
seed("a", PhaseDone)
seed("b", PhaseDownloading)
seed("c", PhaseCanceled)
seed("d", PhasePaused)
seed("e", PhaseError)
seed("f", PhaseQueued)
if n := m.ClearFinished(); n != 3 {
t.Errorf("cleared %d, want the 3 terminal jobs", n)
}
var left []string
for _, j := range m.Jobs() {
left = append(left, j.ID)
}
want := []string{"b", "d", "f"}
if len(left) != len(want) {
t.Fatalf("left = %v, want %v", left, want)
}
for i, id := range want {
if left[i] != id {
t.Fatalf("left = %v, want %v (order preserved)", left, want)
}
}
// A second pass has nothing to do rather than something to break.
if n := m.ClearFinished(); n != 0 {
t.Errorf("second clear removed %d, want 0", n)
}
}