feat: implement credential-based round-robin for gemini-cli virtual auths

Changes the RoundRobinSelector to use two-level round-robin when gemini-cli virtual auths are detected (via gemini_virtual_parent attr): - Level 1: cycle across credential groups (parent accounts) - Level 2: cycle within each group's project auths Credentials start from a random offset (rand.IntN) for fair distribution. Non-virtual auths and single-credential scenarios fall back to flat RR. Adds 3 test cases covering multi-credential grouping, single-parent fallback, and mixed virtual/non-virtual fallback.
2026-02-21 12:49:48 +08:00
parent 7b0eb41ebc
commit 5936f9895c
2 changed files with 201 additions and 6 deletions
@@ -402,3 +402,128 @@ func TestRoundRobinSelectorPick_CursorKeyCap(t *testing.T) {
 		t.Fatalf("selector.cursors missing key %q", "gemini:m3")
 	}
 }
+
+func TestRoundRobinSelectorPick_GeminiCLICredentialGrouping(t *testing.T) {
+	t.Parallel()
+
+	selector := &RoundRobinSelector{}
+
+	// Simulate two gemini-cli credentials, each with multiple projects:
+	// Credential A (parent = "cred-a.json") has 3 projects
+	// Credential B (parent = "cred-b.json") has 2 projects
+	auths := []*Auth{
+		{ID: "cred-a.json::proj-a1", Attributes: map[string]string{"gemini_virtual_parent": "cred-a.json"}},
+		{ID: "cred-a.json::proj-a2", Attributes: map[string]string{"gemini_virtual_parent": "cred-a.json"}},
+		{ID: "cred-a.json::proj-a3", Attributes: map[string]string{"gemini_virtual_parent": "cred-a.json"}},
+		{ID: "cred-b.json::proj-b1", Attributes: map[string]string{"gemini_virtual_parent": "cred-b.json"}},
+		{ID: "cred-b.json::proj-b2", Attributes: map[string]string{"gemini_virtual_parent": "cred-b.json"}},
+	}
+
+	// Two-level round-robin: consecutive picks must alternate between credentials.
+	// Credential group order is randomized, but within each call the group cursor
+	// advances by 1, so consecutive picks should cycle through different parents.
+	picks := make([]string, 6)
+	parents := make([]string, 6)
+	for i := 0; i < 6; i++ {
+		got, err := selector.Pick(context.Background(), "gemini-cli", "gemini-2.5-pro", cliproxyexecutor.Options{}, auths)
+		if err != nil {
+			t.Fatalf("Pick() #%d error = %v", i, err)
+		}
+		if got == nil {
+			t.Fatalf("Pick() #%d auth = nil", i)
+		}
+		picks[i] = got.ID
+		parents[i] = got.Attributes["gemini_virtual_parent"]
+	}
+
+	// Verify property: consecutive picks must alternate between credential groups.
+	for i := 1; i < len(parents); i++ {
+		if parents[i] == parents[i-1] {
+			t.Fatalf("Pick() #%d and #%d both from same parent %q (IDs: %q, %q); expected alternating credentials",
+				i-1, i, parents[i], picks[i-1], picks[i])
+		}
+	}
+
+	// Verify property: each credential's projects are picked in sequence (round-robin within group).
+	credPicks := map[string][]string{}
+	for i, id := range picks {
+		credPicks[parents[i]] = append(credPicks[parents[i]], id)
+	}
+	for parent, ids := range credPicks {
+		for i := 1; i < len(ids); i++ {
+			if ids[i] == ids[i-1] {
+				t.Fatalf("Credential %q picked same project %q twice in a row", parent, ids[i])
+			}
+		}
+	}
+}
+
+func TestRoundRobinSelectorPick_SingleParentFallsBackToFlat(t *testing.T) {
+	t.Parallel()
+
+	selector := &RoundRobinSelector{}
+
+	// All auths from the same parent - should fall back to flat round-robin
+	// because there's only one credential group (no benefit from two-level).
+	auths := []*Auth{
+		{ID: "cred-a.json::proj-a1", Attributes: map[string]string{"gemini_virtual_parent": "cred-a.json"}},
+		{ID: "cred-a.json::proj-a2", Attributes: map[string]string{"gemini_virtual_parent": "cred-a.json"}},
+		{ID: "cred-a.json::proj-a3", Attributes: map[string]string{"gemini_virtual_parent": "cred-a.json"}},
+	}
+
+	// With single parent group, parentOrder has length 1, so it uses flat round-robin.
+	// Sorted by ID: proj-a1, proj-a2, proj-a3
+	want := []string{
+		"cred-a.json::proj-a1",
+		"cred-a.json::proj-a2",
+		"cred-a.json::proj-a3",
+		"cred-a.json::proj-a1",
+	}
+
+	for i, expectedID := range want {
+		got, err := selector.Pick(context.Background(), "gemini-cli", "gemini-2.5-pro", cliproxyexecutor.Options{}, auths)
+		if err != nil {
+			t.Fatalf("Pick() #%d error = %v", i, err)
+		}
+		if got == nil {
+			t.Fatalf("Pick() #%d auth = nil", i)
+		}
+		if got.ID != expectedID {
+			t.Fatalf("Pick() #%d auth.ID = %q, want %q", i, got.ID, expectedID)
+		}
+	}
+}
+
+func TestRoundRobinSelectorPick_MixedVirtualAndNonVirtualFallsBackToFlat(t *testing.T) {
+	t.Parallel()
+
+	selector := &RoundRobinSelector{}
+
+	// Mix of virtual and non-virtual auths (e.g., a regular gemini-cli auth without projects
+	// alongside virtual ones). Should fall back to flat round-robin.
+	auths := []*Auth{
+		{ID: "cred-a.json::proj-a1", Attributes: map[string]string{"gemini_virtual_parent": "cred-a.json"}},
+		{ID: "cred-regular.json"}, // no gemini_virtual_parent
+	}
+
+	// groupByVirtualParent returns nil when any auth lacks the attribute,
+	// so flat round-robin is used. Sorted by ID: cred-a.json::proj-a1, cred-regular.json
+	want := []string{
+		"cred-a.json::proj-a1",
+		"cred-regular.json",
+		"cred-a.json::proj-a1",
+	}
+
+	for i, expectedID := range want {
+		got, err := selector.Pick(context.Background(), "gemini-cli", "", cliproxyexecutor.Options{}, auths)
+		if err != nil {
+			t.Fatalf("Pick() #%d error = %v", i, err)
+		}
+		if got == nil {
+			t.Fatalf("Pick() #%d auth = nil", i)
+		}
+		if got.ID != expectedID {
+			t.Fatalf("Pick() #%d auth.ID = %q, want %q", i, got.ID, expectedID)
+		}
+	}
+}