From d854bf028e1d6f937562fb6a0169284c469a2cab Mon Sep 17 00:00:00 2001 From: raysonmeng Date: Fri, 26 Jun 2026 20:25:39 +0800 Subject: [PATCH 1/2] =?UTF-8?q?feat(seatassign):=20shared=20WRH=20seat?= =?UTF-8?q?=E2=86=92account=20ranking=20module=20(master+proxy=20contract)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Weighted rendezvous hashing for deterministic seat→account assignment within a seat_group. Both control-master (seatgroup service + alloc-engine routing adapter) and aikey-proxy (offline group_resolve default) import it via go.mod replace so the ordering is byte-identical across repos — the proxy's offline default never disagrees with the master's assignment. Co-Authored-By: Claude Opus 4.8 (1M context) --- seatassign/go.mod | 3 + seatassign/seatassign.go | 121 +++++++++++++++++++++++++++ seatassign/seatassign_test.go | 149 ++++++++++++++++++++++++++++++++++ 3 files changed, 273 insertions(+) create mode 100644 seatassign/go.mod create mode 100644 seatassign/seatassign.go create mode 100644 seatassign/seatassign_test.go diff --git a/seatassign/go.mod b/seatassign/go.mod new file mode 100644 index 0000000..f935acc --- /dev/null +++ b/seatassign/go.mod @@ -0,0 +1,3 @@ +module github.com/AiKeyLabs/pkg/seatassign + +go 1.26.1 diff --git a/seatassign/seatassign.go b/seatassign/seatassign.go new file mode 100644 index 0000000..36c31ad --- /dev/null +++ b/seatassign/seatassign.go @@ -0,0 +1,121 @@ +// Package seatassign is the SINGLE SOURCE OF TRUTH for mapping a seat to its +// preferred order of provider accounts within a seat_group (weighted rendezvous +// / HRW hashing). +// +// Why a shared module (not duplicated in master + proxy): the master computes a +// seat's account assignment (the hash default, before any scheduler override) +// and the proxy computes the SAME hash default locally so it can route while +// offline or before the routing-override poll returns. If the two sides ranked +// accounts even slightly differently, the proxy's offline default would disagree +// with the master's assignment → inconsistent routing (a request could land on +// an account the master accounts to a different seat, breaking the per-account +// user caps and prompt-cache stickiness). Both repos import THIS package via +// go.mod replace (like pkg/usagehash), so the ordering is identical by +// construction. Any change to the algorithm flips the pinned-vector test on +// purpose — that is the signal to re-validate BOTH consumers. +// +// Algorithm: weighted rendezvous hashing (WRH). For each candidate account a, +// +// score(seat, a) = -weight / ln( h01(seat, a) ) +// +// (the standard WRH form: P(account is the max) ∝ weight) where h01 ∈ (0,1) is a +// uniform hash of (seatID, accountID). Highest score = +// the seat's primary; the full descending order is its fallback chain. WRH gives +// (1) even spread of seats across accounts, (2) adding/removing one account +// re-homes only ~1/N of seats — no mass reshuffle (unlike `hash mod N`), (3) +// per-account weights for capacity differences. Ties break deterministically by +// (priority, accountID) so the order is a total order with no hidden randomness. +package seatassign + +import ( + "crypto/sha256" + "encoding/binary" + "math" + "sort" +) + +// Account is one provider account a seat can be assigned to within a group. +type Account struct { + // AccountID is the stable identity hashed for ranking + // (seat_group_account.account_id). MUST be stable across master & proxy. + AccountID string + // Weight is the WRH capacity weight; <= 0 is treated as 1. A higher weight + // proportionally attracts more seats (for accounts with larger quota). + Weight float64 + // Priority is a deterministic tie-break when two accounts score equal + // (astronomically rare). Lower wins. + Priority int +} + +// Rank returns accounts ordered by the seat's preference, primary first. +// Deterministic and byte-identical across master & proxy (cross-repo contract). +// The input slice is not mutated. +func Rank(seatID string, accounts []Account) []Account { + out := make([]Account, len(accounts)) + copy(out, accounts) + scores := make(map[string]float64, len(out)) + for _, a := range out { + scores[a.AccountID] = score(seatID, a) + } + sort.SliceStable(out, func(i, j int) bool { + si, sj := scores[out[i].AccountID], scores[out[j].AccountID] + if si != sj { + return si > sj // higher score = preferred + } + if out[i].Priority != out[j].Priority { + return out[i].Priority < out[j].Priority + } + return out[i].AccountID < out[j].AccountID + }) + return out +} + +// Primary returns the seat's top-ranked account ID, or "" if accounts is empty. +func Primary(seatID string, accounts []Account) string { + if len(accounts) == 0 { + return "" + } + best := accounts[0] + bestScore := score(seatID, best) + for _, a := range accounts[1:] { + s := score(seatID, a) + if s > bestScore || (s == bestScore && less(a, best)) { + best, bestScore = a, s + } + } + return best.AccountID +} + +func less(a, b Account) bool { + if a.Priority != b.Priority { + return a.Priority < b.Priority + } + return a.AccountID < b.AccountID +} + +func score(seatID string, a Account) float64 { + w := a.Weight + if w <= 0 { + w = 1 + } + // Standard weighted rendezvous: -w/ln(h01). h01∈(0,1) ⇒ ln<0 ⇒ score>0; + // higher weight ⇒ higher score ⇒ proportionally more seats. + return -w / math.Log(h01(seatID, a.AccountID)) +} + +// h01 hashes (seatID, accountID) to a uniform float in (0,1). The seatID length +// prefix makes (seatID, accountID) unambiguous so ("ab","c") and ("a","bc") +// never collide. +func h01(seatID, accountID string) float64 { + h := sha256.New() + var n [8]byte + binary.BigEndian.PutUint64(n[:], uint64(len(seatID))) + _, _ = h.Write(n[:]) + _, _ = h.Write([]byte(seatID)) + _, _ = h.Write([]byte(accountID)) + sum := h.Sum(nil) + u := binary.BigEndian.Uint64(sum[:8]) + // Map to (0,1): +1 / +2 keeps the result strictly inside (0,1) so -ln is + // finite and positive. + return (float64(u) + 1) / (float64(math.MaxUint64) + 2) +} diff --git a/seatassign/seatassign_test.go b/seatassign/seatassign_test.go new file mode 100644 index 0000000..fb89e45 --- /dev/null +++ b/seatassign/seatassign_test.go @@ -0,0 +1,149 @@ +package seatassign + +import ( + "fmt" + "math" + "testing" +) + +func mkAccounts(n int) []Account { + a := make([]Account, n) + for i := 0; i < n; i++ { + a[i] = Account{AccountID: fmt.Sprintf("acc-%02d", i)} + } + return a +} + +func seats(n int) []string { + s := make([]string, n) + for i := 0; i < n; i++ { + s[i] = fmt.Sprintf("seat-%05d", i) + } + return s +} + +// TestDeterministic: Rank is a pure function — same inputs, same order, every +// call. This is the floor of the cross-repo byte-identical contract. +func TestDeterministic(t *testing.T) { + accs := mkAccounts(8) + for _, seat := range seats(50) { + first := Rank(seat, accs) + for r := 0; r < 5; r++ { + again := Rank(seat, accs) + for i := range first { + if first[i].AccountID != again[i].AccountID { + t.Fatalf("seat %s non-deterministic at %d: %s vs %s", + seat, i, first[i].AccountID, again[i].AccountID) + } + } + } + if Primary(seat, accs) != first[0].AccountID { + t.Fatalf("Primary != Rank[0] for %s", seat) + } + } +} + +// TestGoldenOrder freezes the ranking for a fixed input — the cross-repo +// contract. If WRH math/hashing changes, this fails on purpose: re-validate both +// master & proxy consumers, don't just paste new strings. +func TestGoldenOrder(t *testing.T) { + accs := mkAccounts(6) + got := Rank("seat-00042", accs) + order := make([]string, len(got)) + for i, a := range got { + order[i] = a.AccountID + } + want := []string{"acc-01", "acc-05", "acc-00", "acc-03", "acc-04", "acc-02"} + for i := range want { + if order[i] != want[i] { + t.Fatalf("golden order changed: got %v want %v", order, want) + } + } +} + +// TestEvenSpread: WRH spreads seats ~uniformly across accounts. +func TestEvenSpread(t *testing.T) { + const M, N = 10000, 10 + accs := mkAccounts(N) + counts := map[string]int{} + for _, seat := range seats(M) { + counts[Primary(seat, accs)]++ + } + exp := float64(M) / float64(N) + for id, c := range counts { + if dev := math.Abs(float64(c)-exp) / exp; dev > 0.15 { + t.Errorf("account %s got %d primaries, expected ~%.0f (dev %.1f%% > 15%%)", id, c, exp, dev*100) + } + } + if len(counts) != N { + t.Errorf("expected all %d accounts to receive seats, got %d", N, len(counts)) + } +} + +// TestAddAccountRehomesAboutOneOverN: adding one account re-homes only ~1/(N+1) +// of seats, and ONLY to the new account — existing accounts never trade seats +// among themselves (the core WRH property that avoids mass reshuffle). +func TestAddAccountRehomesAboutOneOverN(t *testing.T) { + const M, N = 10000, 10 + before := mkAccounts(N) + after := append(mkAccounts(N), Account{AccountID: "acc-NEW"}) + + moved, movedToNew := 0, 0 + for _, seat := range seats(M) { + b := Primary(seat, before) + a := Primary(seat, after) + if a != b { + moved++ + if a == "acc-NEW" { + movedToNew++ + } else { + t.Fatalf("seat %s moved between OLD accounts %s→%s (should never happen)", seat, b, a) + } + } + } + if moved != movedToNew { + t.Fatalf("all moves must go to the new account: moved=%d toNew=%d", moved, movedToNew) + } + exp := float64(M) / float64(N+1) + if dev := math.Abs(float64(moved)-exp) / exp; dev > 0.20 { + t.Errorf("re-homed %d seats, expected ~%.0f (1/(N+1)); dev %.1f%% > 20%%", moved, exp, dev*100) + } +} + +// TestWeightedFavorsHigh: a heavier-weight account attracts proportionally more +// seats. +func TestWeightedFavorsHigh(t *testing.T) { + const M = 10000 + accs := []Account{ + {AccountID: "w1-a"}, {AccountID: "w1-b"}, {AccountID: "w1-c"}, {AccountID: "w1-d"}, + {AccountID: "w3-heavy", Weight: 3}, + } + counts := map[string]int{} + for _, seat := range seats(M) { + counts[Primary(seat, accs)]++ + } + // Total weight 7; heavy share ≈ 3/7. Each w1 ≈ 1/7. Heavy should clearly + // dominate (> 2x any single light account). + heavy := counts["w3-heavy"] + for _, id := range []string{"w1-a", "w1-b", "w1-c", "w1-d"} { + if heavy < 2*counts[id] { + t.Errorf("heavy(%d) should be > 2x light %s(%d)", heavy, id, counts[id]) + } + } + if frac := float64(heavy) / M; frac < 0.33 || frac > 0.52 { + t.Errorf("heavy share %.2f outside expected ~3/7≈0.43", frac) + } +} + +func TestEmptyAndSingle(t *testing.T) { + if Primary("seat-x", nil) != "" { + t.Error("empty accounts must yield empty primary") + } + if len(Rank("seat-x", nil)) != 0 { + t.Error("empty accounts must yield empty rank") + } + one := []Account{{AccountID: "solo"}} + if Primary("seat-x", one) != "solo" { + t.Error("single account must be primary") + } +} From 140ec4123fe99f0612369d4c50541af617436219 Mon Sep 17 00:00:00 2001 From: merge-bot Date: Fri, 26 Jun 2026 21:42:37 +0800 Subject: [PATCH 2/2] chore(pkg): rename seat_group -> oauth_group (align to latest naming) --- seatassign/seatassign.go | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/seatassign/seatassign.go b/seatassign/seatassign.go index 36c31ad..9e89ac6 100644 --- a/seatassign/seatassign.go +++ b/seatassign/seatassign.go @@ -1,5 +1,5 @@ // Package seatassign is the SINGLE SOURCE OF TRUTH for mapping a seat to its -// preferred order of provider accounts within a seat_group (weighted rendezvous +// preferred order of provider accounts within a oauth_group (weighted rendezvous // / HRW hashing). // // Why a shared module (not duplicated in master + proxy): the master computes a @@ -37,7 +37,7 @@ import ( // Account is one provider account a seat can be assigned to within a group. type Account struct { // AccountID is the stable identity hashed for ranking - // (seat_group_account.account_id). MUST be stable across master & proxy. + // (oauth_group_account.account_id). MUST be stable across master & proxy. AccountID string // Weight is the WRH capacity weight; <= 0 is treated as 1. A higher weight // proportionally attracts more seats (for accounts with larger quota).