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) } }