llmgw/internal/mirror/ollama_test.go

package mirror

import (
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"errors"
	"os"
	"path/filepath"
	"slices"
	"strings"
	"testing"

	"gitea.mynodes.xyz/ilker/llmgw/internal/agent/engines"
)

// ollamaFixture is an Ollama models directory laid out as Ollama writes it:
// blobs/sha256-<hex> named by their content's digest, and a manifest per
// model and tag under manifests/<host>/<namespace>/<model>/<tag>.
type ollamaFixture struct {
	t   *testing.T
	dir string
}

func newOllamaFixture(t *testing.T) *ollamaFixture {
	t.Helper()
	dir := filepath.Join(t.TempDir(), ".ollama", "models")
	for _, d := range []string{"blobs", "manifests"} {
		if err := os.MkdirAll(filepath.Join(dir, d), 0o755); err != nil {
			t.Fatal(err)
		}
	}
	return &ollamaFixture{t: t, dir: dir}
}

// blob stores content as Ollama does and returns the layer naming it.
func (f *ollamaFixture) blob(mediaType string, content []byte) OllamaLayer {
	f.t.Helper()
	sum := sha256.Sum256(content)
	hexsum := hex.EncodeToString(sum[:])
	if err := os.WriteFile(filepath.Join(f.dir, "blobs", "sha256-"+hexsum), content, 0o644); err != nil {
		f.t.Fatal(err)
	}
	return OllamaLayer{MediaType: mediaType, Digest: "sha256:" + hexsum, Size: int64(len(content))}
}

func (f *ollamaFixture) manifest(name string, layers ...OllamaLayer) {
	f.t.Helper()
	n, err := ParseOllamaName(name)
	if err != nil {
		f.t.Fatal(err)
	}
	config := f.blob("application/vnd.docker.container.image.v1+json",
		[]byte(`{"model_format":"gguf","model_family":"llama","file_type":"Q4_K_M"}`))
	b, _ := json.Marshal(map[string]any{
		"schemaVersion": 2,
		"mediaType":     "application/vnd.docker.distribution.manifest.v2+json",
		"config":        config,
		"layers":        layers,
	})
	p := n.manifestPath(f.dir)
	if err := os.MkdirAll(filepath.Dir(p), 0o755); err != nil {
		f.t.Fatal(err)
	}
	if err := os.WriteFile(p, b, 0o644); err != nil {
		f.t.Fatal(err)
	}
}

// gguf is a small file that begins as every GGUF does.
func gguf(body string) []byte { return []byte("GGUF\x03\x00\x00\x00" + body) }

// llama32 is llama3.2:3b as Ollama pulls it: the weights, a template, a
// license and parameters.
func (f *ollamaFixture) llama32() OllamaLayer {
	weights := f.blob(ollamaModelLayer, gguf("llama 3.2 3b weights"))
	f.manifest("llama3.2:3b",
		weights,
		f.blob("application/vnd.ollama.image.template", []byte(`<|start_header_id|>{{ .Prompt }}`)),
		f.blob("application/vnd.ollama.image.license", []byte("LLAMA 3.2 COMMUNITY LICENSE")),
		f.blob(ollamaParamsLayer, []byte(`{"num_ctx":8192,"stop":["<|eot_id|>"]}`)),
	)
	return weights
}

func TestOllamaNamesAreReadAsOllamaReadsThem(t *testing.T) {
	for in, want := range map[string]OllamaName{
		"llama3.2":                     {ollamaHost, ollamaNamespace, "llama3.2", "latest"},
		"llama3.2:3b":                  {ollamaHost, ollamaNamespace, "llama3.2", "3b"},
		"jmorgan/llava:7b-q4":          {ollamaHost, "jmorgan", "llava", "7b-q4"},
		"hf.co/bartowski/Qwen-GGUF:Q4": {"hf.co", "bartowski", "Qwen-GGUF", "Q4"},
		"localhost:5000/me/m:v1":       {"localhost:5000", "me", "m", "v1"},
		"qwen3:8b@sha256:abcd":         {ollamaHost, ollamaNamespace, "qwen3", "8b"},
	} {
		got, err := ParseOllamaName(in)
		if err != nil || got != want {
			t.Errorf("%q: got %+v, %v; want %+v", in, got, err, want)
		}
	}
	for _, bad := range []string{"", "a/b/c/d", "../etc:passwd", "x:.."} {
		if _, err := ParseOllamaName(bad); err == nil {
			t.Errorf("%q was taken for a name", bad)
		}
	}
	// Shown the way Ollama shows them: without its defaults.
	for in, want := range map[string]string{
		"llama3.2": "llama3.2:latest", "jmorgan/llava:7b": "jmorgan/llava:7b",
		"hf.co/bartowski/Qwen-GGUF:Q4": "hf.co/bartowski/Qwen-GGUF:Q4",
	} {
		n, _ := ParseOllamaName(in)
		if n.String() != want {
			t.Errorf("%q shows as %q", in, n.String())
		}
	}
}

func TestOllamasDirectoryIsFoundWhereOllamaKeepsIt(t *testing.T) {
	f := newOllamaFixture(t)
	// Given as the models directory, or as the ~/.ollama it is in.
	for _, given := range []string{f.dir, filepath.Dir(f.dir)} {
		if got, err := OllamaDir(given); err != nil || got != f.dir {
			t.Errorf("%s: got %q, %v", given, got, err)
		}
	}
	t.Setenv("OLLAMA_MODELS", f.dir)
	if got, err := OllamaDir(""); err != nil || got != f.dir {
		t.Errorf("OLLAMA_MODELS: got %q, %v", got, err)
	}
	if _, err := OllamaDir(t.TempDir()); err == nil {
		t.Error("a directory with no manifests was taken for Ollama's")
	}
}

func TestOllamasModelsAreListedFromItsManifests(t *testing.T) {
	f := newOllamaFixture(t)
	f.llama32()
	f.manifest("hf.co/bartowski/Qwen2.5-7B-Instruct-GGUF:Q4_K_M", f.blob(ollamaModelLayer, gguf("qwen")))
	// Ollama's own leftovers are not models.
	_ = os.WriteFile(filepath.Join(f.dir, "manifests", ".DS_Store"), nil, 0o644)

	got, err := ListOllama(f.dir)
	if err != nil {
		t.Fatal(err)
	}
	var names []string
	for _, n := range got {
		names = append(names, n.String())
	}
	if !slices.Equal(names, []string{"hf.co/bartowski/Qwen2.5-7B-Instruct-GGUF:Q4_K_M", "llama3.2:3b"}) {
		t.Errorf("listed %v", names)
	}
}

func TestAnOllamaManifestYieldsItsGGUFAndWhatIsLeftBehind(t *testing.T) {
	f := newOllamaFixture(t)
	weights := f.llama32()
	m, err := ReadOllama(f.dir, mustName(t, "llama3.2:3b"))
	if err != nil {
		t.Fatal(err)
	}
	if m.Model != weights {
		t.Errorf("model layer %+v, want %+v", m.Model, weights)
	}
	// What Ollama ran it with, to carry over as --ctx-size.
	if m.ContextLength != 8192 {
		t.Errorf("num_ctx=%d", m.ContextLength)
	}
	// The template is Ollama's own language; saying it is left is the honest
	// part, since llama.cpp will use the GGUF's.
	if len(m.Left) != 1 || !strings.Contains(m.Left[0], "template") {
		t.Errorf("left: %v", m.Left)
	}
}

func TestOllamaModelsLlamaCppCannotRunAreRefused(t *testing.T) {
	f := newOllamaFixture(t)
	// Ollama's newer layout keeps some models as separate tensors.
	f.manifest("gpt-oss:20b", f.blob(ollamaTensorLayer, []byte("tensor")), f.blob(ollamaTensorLayer, []byte("other")))
	// Two weight layers make a model this import does not join.
	f.manifest("split:1", f.blob(ollamaModelLayer, gguf("a")), f.blob(ollamaModelLayer, gguf("b")))
	f.manifest("empty:1", f.blob("application/vnd.ollama.image.license", []byte("MIT")))

	for name, want := range map[string]string{
		"gpt-oss:20b": "separate tensors", "split:1": "2 weight layers", "empty:1": "no GGUF layer",
		"absent:1": "no absent:1",
	} {
		_, err := ReadOllama(f.dir, mustName(t, name))
		if err == nil || !strings.Contains(err.Error(), want) {
			t.Errorf("%s: expected %q, got %v", name, want, err)
		}
	}
}

func TestAnOllamaModelIsImportedIntoTheMirrorWithItsChecksum(t *testing.T) {
	f := newOllamaFixture(t)
	weights := f.llama32()
	root := t.TempDir()
	m, err := ReadOllama(f.dir, mustName(t, "llama3.2:3b"))
	if err != nil {
		t.Fatal(err)
	}

	var reported []string
	added, err := ImportOllama(t.Context(), root, f.dir, m, "", "", func(file string, _, _ int64) {
		reported = append(reported, file)
	})
	if err != nil {
		t.Fatal(err)
	}
	// The tag alone moves in Ollama ("latest" is whatever was pulled last);
	// the start of the digest is what pins the version.
	sum := strings.TrimPrefix(weights.Digest, "sha256:")
	want := Entry{Model: "llama3.2", Version: "3b-" + sum[:12], File: "llama3.2-3b.gguf", Bytes: weights.Size}
	if len(added) != 1 || added[0].Entry != want {
		t.Fatalf("added %+v, want %+v", added, want)
	}
	if added[0].SHA256 != sum {
		t.Errorf("sha256 %s, want the digest Ollama named it by, %s", added[0].SHA256, sum)
	}
	if len(reported) == 0 || reported[0] != "llama3.2-3b.gguf" {
		t.Errorf("progress said %v", reported)
	}

	// What a node is given: the file, with the checksum it checks it against.
	mf, err := ManifestFor(root, want.Model, want.Version)
	if err != nil {
		t.Fatal(err)
	}
	if len(mf.Files) != 1 || mf.Files[0].SHA256 != sum {
		t.Errorf("manifest %+v", mf)
	}
	b, _ := os.ReadFile(want.Path(root))
	if string(b) != string(gguf("llama 3.2 3b weights")) {
		t.Error("the mirrored file is not the blob")
	}
	if got := readSource(want.Path(root)); got != "ollama://registry.ollama.ai/library/llama3.2:3b" {
		t.Errorf("source recorded as %q", got)
	}

	// The same model again is the same version, and the mirror never
	// replaces one.
	if _, err := ImportOllama(t.Context(), root, f.dir, m, "", "", nil); !errors.Is(err, ErrExists) {
		t.Errorf("a second import was not refused: %v", err)
	}
	// Under names the operator chose, it goes where they say.
	added, err = ImportOllama(t.Context(), root, f.dir, m, "llama-3.2-3b-instruct", "2026-10", nil)
	if err != nil || added[0].Entry.Path(root) != filepath.Join(root, "llama-3.2-3b-instruct", "2026-10", "llama3.2-3b.gguf") {
		t.Errorf("named import: %+v, %v", added, err)
	}
}

func TestAVisionModelsProjectorGoesBesideItForLlamaCpp(t *testing.T) {
	f := newOllamaFixture(t)
	f.manifest("llava:7b",
		f.blob(ollamaModelLayer, gguf("llava weights")),
		f.blob(ollamaProjectorLayer, gguf("clip projector")),
	)
	root := t.TempDir()
	m, err := ReadOllama(f.dir, mustName(t, "llava:7b"))
	if err != nil {
		t.Fatal(err)
	}
	added, err := ImportOllama(t.Context(), root, f.dir, m, "", "v1", nil)
	if err != nil {
		t.Fatal(err)
	}
	if len(added) != 2 || added[1].Entry.File != "mmproj-llava-7b.gguf" {
		t.Fatalf("added %+v", added)
	}

	// The version directory is the model's source, and llama.cpp's engine
	// finds in it the model to load and the projector to load beside it.
	dir := filepath.Join(root, "llava", "v1")
	s := engines.ShardSpec{HTTPPort: 20000, Profile: engines.Profile{Name: "cpu", Engine: "llamacpp", ModelPath: dir}}
	cmds, err := (&engines.LlamaCpp{}).Commands(s)
	if err != nil {
		t.Fatal(err)
	}
	args := cmds[0].Args
	if got := args[slices.Index(args, "--model")+1]; got != filepath.Join(dir, "llava-7b.gguf") {
		t.Errorf("--model %s", got)
	}
	if got := args[slices.Index(args, "--mmproj")+1]; got != filepath.Join(dir, "mmproj-llava-7b.gguf") {
		t.Errorf("--mmproj %s", got)
	}
}

func TestADamagedOllamaBlobIsNotMirrored(t *testing.T) {
	f := newOllamaFixture(t)
	f.manifest("llava:7b",
		f.blob(ollamaModelLayer, gguf("llava weights")),
		f.blob(ollamaProjectorLayer, gguf("clip projector")),
	)
	m, err := ReadOllama(f.dir, mustName(t, "llava:7b"))
	if err != nil {
		t.Fatal(err)
	}
	// A bit gone wrong on Ollama's disk: the projector no longer hashes to
	// the name it was stored under.
	path, _, _ := ollamaBlob(f.dir, m.Projectors[0].Digest)
	if err := os.WriteFile(path, gguf("clip projectoR"), 0o644); err != nil {
		t.Fatal(err)
	}

	root := t.TempDir()
	_, err = ImportOllama(t.Context(), root, f.dir, m, "", "v1", nil)
	if err == nil || !strings.Contains(err.Error(), "damaged") {
		t.Fatalf("expected the damage to be found, got %v", err)
	}
	// All or nothing: the weights copied before it are taken out again, so
	// no node is offered a vision model without its projector.
	if entries, _ := List(root); len(entries) != 0 {
		t.Errorf("left in the mirror: %+v", entries)
	}
	if _, err := os.Stat(filepath.Join(root, "llava")); !errors.Is(err, os.ErrNotExist) {
		t.Errorf("the version's directory was left behind: %v", err)
	}

	// A blob gone altogether says which.
	_ = os.Remove(path)
	if _, err := ImportOllama(t.Context(), root, f.dir, m, "", "v1", nil); err == nil ||
		!strings.Contains(err.Error(), m.Projectors[0].Digest) {
		t.Errorf("a missing blob: %v", err)
	}
}

func TestOnlyAGGUFIsImportedAsOne(t *testing.T) {
	// An HTML error page saved as a .gguf would reach every node otherwise.
	src := filepath.Join(t.TempDir(), "model.gguf")
	_ = os.WriteFile(src, []byte("<html>rate limited</html>"), 0o644)
	_, err := Import(t.Context(), t.TempDir(), src, Entry{Model: "m", Version: "v", File: "model.gguf"}, "", "file://"+src, nil)
	if err == nil || !strings.Contains(err.Error(), "not a GGUF") {
		t.Errorf("expected a refusal, got %v", err)
	}
}

func TestALocalFileIsImportedWithItsChecksum(t *testing.T) {
	src := filepath.Join(t.TempDir(), "q4_k_m.gguf")
	content := gguf("converted by hand")
	_ = os.WriteFile(src, content, 0o644)
	root := t.TempDir()
	e := Entry{Model: "qwen", Version: "2026-10-04", File: "q4_k_m.gguf"}

	a, err := Import(t.Context(), root, src, e, "", "file://"+src, nil)
	if err != nil {
		t.Fatal(err)
	}
	sum := sha256.Sum256(content)
	if a.SHA256 != hex.EncodeToString(sum[:]) || readChecksum(e.Path(root)) != a.SHA256 {
		t.Errorf("checksum %s, recorded %s", a.SHA256, readChecksum(e.Path(root)))
	}
	// A checksum the operator was given is checked too.
	e.Version = "again"
	if _, err := Import(t.Context(), root, src, e, strings.Repeat("0", 64), "file://"+src, nil); err == nil {
		t.Error("a file that does not hash to the sum given was imported")
	}
	if _, err := os.Stat(e.Path(root) + partSuffix); !errors.Is(err, os.ErrNotExist) {
		t.Error("a refused import left its partial file")
	}
}

func mustName(t *testing.T, s string) OllamaName {
	t.Helper()
	n, err := ParseOllamaName(s)
	if err != nil {
		t.Fatal(err)
	}
	return n
}