feat: migrate legacy hf.co tags to huggingface.co during client initialization (#645)

* feat: migrate legacy hf.co tags to huggingface.co during client initialization

* Update pkg/distribution/distribution/client_test.go

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

* fix format

---------

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

Author: ilopezluna

pkg/distribution/distribution/client.go

@@ -103,11 +103,37 @@ func NewClient(opts ...Option) (*Client, error) {
 	}
 
 	options.logger.Infoln("Successfully initialized store")
-	return &Client{
+	c := &Client{
 		store:    s,
 		log:      options.logger,
 		registry: registryClient,
-	}, nil
+	}
+
+	// Migrate any legacy hf.co tags to huggingface.co
+	if err := c.migrateHFTags(); err != nil {
+		options.logger.Warnf("Failed to migrate HuggingFace tags: %v", err)
+	}
+
+	return c, nil
+}
+
+// migrateHFTags normalizes legacy hf.co/ tags in the store to huggingface.co/.
+// This handles models that were pulled before the hf.co normalization was added,
+// ensuring they can be found by the cache check in PullModel.
+func (c *Client) migrateHFTags() error {
+	migrated, err := c.store.MigrateTags(func(tag string) string {
+		if rest, found := strings.CutPrefix(tag, "hf.co/"); found {
+			return "huggingface.co/" + rest
+		}
+		return tag
+	})
+	if err != nil {
+		return err
+	}
+	if migrated > 0 {
+		c.log.Infof("Migrated %d HuggingFace tag(s) from hf.co to huggingface.co", migrated)
+	}
+	return nil
 }
 
 // normalizeModelName adds the default organization prefix (ai/) and tag (:latest) if missing.

pkg/distribution/distribution/client_test.go

@@ -1141,6 +1141,94 @@ func randomFile(size int64) (string, error) {
 	return f.Name(), nil
 }
 
+func TestMigrateHFTagsOnClientInit(t *testing.T) {
+	testCases := []struct {
+		name          string
+		storedTag     string
+		lookupRef     string
+		shouldMigrate bool
+	}{
+		{
+			name:          "hf.co tag migrated to huggingface.co on init",
+			storedTag:     "hf.co/testorg/testmodel:latest",
+			lookupRef:     "hf.co/testorg/testmodel",
+			shouldMigrate: true,
+		},
+		{
+			name:          "hf.co tag with quantization migrated",
+			storedTag:     "hf.co/bartowski/llama-3.2-1b-instruct-gguf:Q4_K_M",
+			lookupRef:     "hf.co/bartowski/llama-3.2-1b-instruct-gguf:Q4_K_M",
+			shouldMigrate: true,
+		},
+		{
+			name:          "huggingface.co tag unchanged",
+			storedTag:     "huggingface.co/testorg/testmodel:latest",
+			lookupRef:     "huggingface.co/testorg/testmodel",
+			shouldMigrate: false,
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			tempDir := t.TempDir()
+
+			// Step 1: Create a client and write a model with the legacy tag
+			setupClient, err := newTestClient(tempDir)
+			if err != nil {
+				t.Fatalf("Failed to create setup client: %v", err)
+			}
+
+			model, err := gguf.NewModel(testGGUFFile)
+			if err != nil {
+				t.Fatalf("Failed to create model: %v", err)
+			}
+
+			if err := setupClient.store.Write(model, []string{tc.storedTag}, nil); err != nil {
+				t.Fatalf("Failed to write model to store: %v", err)
+			}
+
+			// Step 2: Create a NEW client (simulating restart) - migration should happen
+			client, err := newTestClient(tempDir)
+			if err != nil {
+				t.Fatalf("Failed to create client: %v", err)
+			}
+
+			// Step 3: Verify the model can be found using the reference
+			// (normalizeModelName converts hf.co -> huggingface.co, and migration should have updated the store)
+			foundModel, err := client.GetModel(tc.lookupRef)
+			if err != nil {
+				t.Fatalf("Failed to get model after migration: %v", err)
+			}
+
+			if foundModel == nil {
+				t.Fatal("Expected to find model after migration, got nil")
+			}
+
+			// Step 4: If the tag was hf.co, verify it was actually migrated in the store
+			if tc.shouldMigrate {
+				// The model should now have huggingface.co tag, not hf.co
+				tags := foundModel.Tags()
+				hasOldTag := false
+				hasNewTag := false
+				for _, tag := range tags {
+					if strings.HasPrefix(tag, "hf.co/") {
+						hasOldTag = true
+					}
+					if strings.HasPrefix(tag, "huggingface.co/") {
+						hasNewTag = true
+					}
+				}
+				if hasOldTag {
+					t.Errorf("Model still has old hf.co tag after migration: %v", tags)
+				}
+				if !hasNewTag {
+					t.Errorf("Model doesn't have huggingface.co tag after migration: %v", tags)
+				}
+			}
+		})
+	}
+}
+
 func TestPullHuggingFaceModelFromCache(t *testing.T) {
 	testCases := []struct {
 		name    string

pkg/distribution/internal/store/store.go

@@ -534,6 +534,39 @@ func (s *LocalStore) WriteLightweight(mdl oci.Image, tags []string) (err error)
 	return nil
 }
 
+// MigrateTags applies a transformation function to all tags in the index.
+// If the function returns a different string, the tag is updated.
+// Returns the number of tags that were migrated.
+func (s *LocalStore) MigrateTags(transform func(string) string) (int, error) {
+	index, err := s.readIndex()
+	if err != nil {
+		return 0, fmt.Errorf("reading index for migration: %w", err)
+	}
+
+	migrated := 0
+	changed := false
+	for i, entry := range index.Models {
+		newTags := make([]string, len(entry.Tags))
+		for j, tag := range entry.Tags {
+			newTag := transform(tag)
+			if newTag != tag {
+				migrated++
+				changed = true
+			}
+			newTags[j] = newTag
+		}
+		index.Models[i].Tags = newTags
+	}
+
+	if changed {
+		if err := s.writeIndex(index); err != nil {
+			return 0, fmt.Errorf("writing migrated index: %w", err)
+		}
+	}
+
+	return migrated, nil
+}
+
 // Read reads a model from the store by reference (either tag or ID)
 func (s *LocalStore) Read(reference string) (*Model, error) {
 	models, err := s.List()

pkg/distribution/internal/store/store_test.go

@@ -991,6 +991,82 @@ func TestResetStore(t *testing.T) {
 	}
 }
 
+func TestMigrateTags(t *testing.T) {
+	tempDir := t.TempDir()
+
+	storePath := filepath.Join(tempDir, "migrate-tags-store")
+	s, err := store.New(store.Options{
+		RootPath: storePath,
+	})
+	if err != nil {
+		t.Fatalf("Failed to create store: %v", err)
+	}
+
+	// Write a model with an hf.co tag (simulating pre-migration state)
+	mdl := newTestModel(t)
+	if err := s.Write(mdl, []string{"hf.co/testorg/testmodel:latest"}, nil); err != nil {
+		t.Fatalf("Write failed: %v", err)
+	}
+
+	// Write another model with a non-HF tag (should be unaffected)
+	mdl2Content := []byte("another model content")
+	mdl2Path := filepath.Join(tempDir, "other-model.gguf")
+	if err := os.WriteFile(mdl2Path, mdl2Content, 0644); err != nil {
+		t.Fatalf("Failed to write model file: %v", err)
+	}
+	mdl2, err := gguf.NewModel(mdl2Path)
+	if err != nil {
+		t.Fatalf("Failed to create model: %v", err)
+	}
+	if err := s.Write(mdl2, []string{"ai/some-model:latest"}, nil); err != nil {
+		t.Fatalf("Write failed: %v", err)
+	}
+
+	// Run migration: hf.co/ -> huggingface.co/
+	migrated, err := s.MigrateTags(func(tag string) string {
+		if rest, found := strings.CutPrefix(tag, "hf.co/"); found {
+			return "huggingface.co/" + rest
+		}
+		return tag
+	})
+	if err != nil {
+		t.Fatalf("MigrateTags failed: %v", err)
+	}
+
+	if migrated != 1 {
+		t.Errorf("Expected 1 migrated tag, got %d", migrated)
+	}
+
+	// Verify the model can be found with the new tag
+	if _, err := s.Read("huggingface.co/testorg/testmodel:latest"); err != nil {
+		t.Fatalf("Failed to read model with migrated tag: %v", err)
+	}
+
+	// Verify the old tag no longer works
+	if _, err := s.Read("hf.co/testorg/testmodel:latest"); !errors.Is(err, store.ErrModelNotFound) {
+		t.Errorf("Expected ErrModelNotFound for old tag, got: %v", err)
+	}
+
+	// Verify the non-HF model is unaffected
+	if _, err := s.Read("ai/some-model:latest"); err != nil {
+		t.Fatalf("Non-HF model should be unaffected: %v", err)
+	}
+
+	// Running migration again should be a no-op
+	migrated2, err := s.MigrateTags(func(tag string) string {
+		if rest, found := strings.CutPrefix(tag, "hf.co/"); found {
+			return "huggingface.co/" + rest
+		}
+		return tag
+	})
+	if err != nil {
+		t.Fatalf("Second MigrateTags failed: %v", err)
+	}
+	if migrated2 != 0 {
+		t.Errorf("Expected 0 migrated tags on second run, got %d", migrated2)
+	}
+}
+
 func TestWriteLightweight(t *testing.T) {
 	tempDir := t.TempDir()