llmgw/internal/agent/removeweights_test.go

package agent

import (
	"log/slog"
	"os"
	"path/filepath"
	"testing"

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

// agentWithCache is an agent that only knows where its weights live, which is
// all removeWeights needs.
func agentWithCache(t *testing.T, dir string) *Agent {
	t.Helper()
	return &Agent{
		cfg: Config{ModelDirs: []string{dir}},
		log: slog.New(slog.DiscardHandler),
	}
}

func write(t *testing.T, path string) string {
	t.Helper()
	if err := os.WriteFile(path, []byte("weights"), 0o644); err != nil {
		t.Fatal(err)
	}
	return path
}

func TestTheAgentRefusesToDeleteWhatItDidNotFetch(t *testing.T) {
	// The control plane computes these paths, but the node holds the files. A
	// control plane that is wrong, newer than this agent, or tampered with must
	// not be able to talk it into deleting a model an operator spent a week
	// copying here. So the agent checks again, on its own.
	dir := t.TempDir()
	a := agentWithCache(t, dir)

	operators := write(t, filepath.Join(dir, "qwen2.5-coder-14b-instruct-q8_0.gguf"))
	fetched := write(t, weights.CachePathFor(dir, "https://models.internal/gone.gguf"))

	a.removeWeights([]string{operators, fetched})

	if _, err := os.Stat(operators); err != nil {
		t.Errorf("a file llmgw never fetched was deleted: %v", err)
	}
	if _, err := os.Stat(fetched); !os.IsNotExist(err) {
		t.Errorf("the fetched file was not removed: %v", err)
	}
}

func TestTheAgentRefusesPathsOutsideItsCache(t *testing.T) {
	// Even with a name that looks like one of ours. The cache directory is the
	// boundary, and traversal out of it is the attack this closes.
	dir := t.TempDir()
	elsewhere := t.TempDir()
	a := agentWithCache(t, dir)

	// A correctly named cache file, in the wrong directory.
	outside := write(t, weights.CachePathFor(elsewhere, "https://models.internal/x.gguf"))
	// And the same thing reached by traversal from inside the cache.
	traversal := filepath.Join(dir, "..", filepath.Base(elsewhere),
		filepath.Base(outside))

	a.removeWeights([]string{outside, traversal})

	if _, err := os.Stat(outside); err != nil {
		t.Errorf("a file outside the cache was deleted: %v", err)
	}
}

func TestRemovingAFileThatIsAlreadyGoneIsNotAFailure(t *testing.T) {
	// The expected outcome of a retry. It must not stop the rest of the batch.
	dir := t.TempDir()
	a := agentWithCache(t, dir)

	missing := weights.CachePathFor(dir, "https://models.internal/never.gguf")
	present := write(t, weights.CachePathFor(dir, "https://models.internal/here.gguf"))

	a.removeWeights([]string{missing, present})

	if _, err := os.Stat(present); !os.IsNotExist(err) {
		t.Errorf("a missing file stopped the batch: %v", err)
	}
}

func TestWithNoCacheConfiguredNothingIsDeleted(t *testing.T) {
	dir := t.TempDir()
	a := &Agent{cfg: Config{}, log: slog.New(slog.DiscardHandler)}
	kept := write(t, weights.CachePathFor(dir, "https://models.internal/x.gguf"))

	a.removeWeights([]string{kept})

	if _, err := os.Stat(kept); err != nil {
		t.Errorf("an agent with no cache directory deleted something: %v", err)
	}
}

// A model of several files is fetched into a directory of its own, and an
// unused one goes whole: it was refused as not empty, and nothing was freed.
func TestAnUnusedSetIsRemovedWhole(t *testing.T) {
	dir := t.TempDir()
	a := agentWithCache(t, dir)
	set := weights.CachePathFor(dir, "https://mirror/coder/v1")
	if err := os.MkdirAll(set, 0o755); err != nil {
		t.Fatal(err)
	}
	write(t, filepath.Join(set, "model-00001-of-00002.gguf"))
	write(t, filepath.Join(set, ".complete"))
	// A link in the cache is taken away, not what it points at.
	target := t.TempDir()
	write(t, filepath.Join(target, "keep.gguf"))
	link := weights.CachePathFor(dir, "https://mirror/linked/v1")
	if err := os.Symlink(target, link); err != nil {
		t.Fatal(err)
	}

	a.removeWeights([]string{set, link})

	if _, err := os.Stat(set); !os.IsNotExist(err) {
		t.Errorf("the set is still there: %v", err)
	}
	if _, err := os.Lstat(link); !os.IsNotExist(err) {
		t.Errorf("the link is still there: %v", err)
	}
	if _, err := os.Stat(filepath.Join(target, "keep.gguf")); err != nil {
		t.Errorf("what the link pointed at was deleted: %v", err)
	}
}