From bd2c9729ff69b691ed22554c13d8eb53376469a4 Mon Sep 17 00:00:00 2001 From: N Rohith Reddy Date: Thu, 16 Jul 2026 15:27:55 -0400 Subject: [PATCH 1/2] fix(bandit): repair PairedModel when Thompson sampling switches model Mirror banditexplore.repairBandPair so serving the scorer's runner-up cannot collapse the session-pin band pair and silently disable Stage 2 band swap. Fixes #753 Signed-off-by: N Rohith Reddy --- internal/router/bandit/bandit.go | 45 ++++++++++++++++++++- internal/router/bandit/bandit_test.go | 56 +++++++++++++++++++++++++++ 2 files changed, 99 insertions(+), 2 deletions(-) diff --git a/internal/router/bandit/bandit.go b/internal/router/bandit/bandit.go index a2ee6188f..2ee1b200b 100644 --- a/internal/router/bandit/bandit.go +++ b/internal/router/bandit/bandit.go @@ -109,8 +109,13 @@ func (b *Router) Route(ctx context.Context, req router.Request) (router.Decision argmaxModel := dec.Model pick := candidates[0] b.annotate(&dec, pick.model, pick.provider, b.propensity(clusterIDs, candidates, pick.model), pick.sample) - if pick.model != argmaxModel && md.EffectiveKnobsHash != 0 { - md.EffectiveKnobsHash = mixModel(md.EffectiveKnobsHash, pick.model) + if pick.model != argmaxModel { + if md.EffectiveKnobsHash != 0 { + md.EffectiveKnobsHash = mixModel(md.EffectiveKnobsHash, pick.model) + } + // The scorer's runner-up was computed against the argmax and may now + // equal the served peer; recompute it against the served model. + repairBandPair(&md.PairedModel, &md.PairedProvider, &md.PairedScore, md.CandidateScores, candidates, pick.model) } return dec, nil } @@ -185,6 +190,42 @@ func (b *Router) annotate(dec *router.Decision, model, provider string, propensi } } +// repairBandPair recomputes the band pair's runner-up against the served +// model, writing it back through the metadata pointers. Picks the +// highest-scoring servable peer other than served (ties broken by name); +// clears the pair if none remains. Providers come from the live candidate +// list (resolved before annotate) so the prior argmax stays eligible even +// when it was only bound via Decision.Provider. +func repairBandPair(outModel, outProvider *string, outScore *float32, scores map[string]float32, cands []candidate, served string) { + providers := make(map[string]string, len(cands)) + for _, c := range cands { + providers[c.model] = c.provider + } + models := make([]string, 0, len(scores)) + for m := range scores { + models = append(models, m) + } + sort.Strings(models) + + bestModel, bestProvider := "", "" + var bestScore float32 + for _, m := range models { + if m == served { + continue + } + sc := scores[m] + if bestModel != "" && sc <= bestScore { + continue + } + provider, ok := providers[m] + if !ok || provider == "" { + continue + } + bestModel, bestProvider, bestScore = m, provider, sc + } + *outModel, *outProvider, *outScore = bestModel, bestProvider, bestScore +} + func mixModel(h uint64, model string) uint64 { f := fnv.New64a() var b [8]byte diff --git a/internal/router/bandit/bandit_test.go b/internal/router/bandit/bandit_test.go index 4eeda5b79..f87d8d57f 100644 --- a/internal/router/bandit/bandit_test.go +++ b/internal/router/bandit/bandit_test.go @@ -133,6 +133,11 @@ func TestRoute_PropensityReflectsPosteriorOverlap(t *testing.T) { func TestRoute_KeepsArgmaxWhenSameModel(t *testing.T) { scores := map[string]float32{"claude-haiku-4-5": 0.9, "claude-sonnet-4-6": 0.5} inner := clusterDecision(scores, []int{0}, "claude-haiku-4-5", "anthropic") + // Scorer's runner-up; keeping the argmax must leave this pair untouched + // rather than recomputing it. + inner.Metadata.PairedModel = "claude-sonnet-4-6" + inner.Metadata.PairedProvider = "anthropic" + inner.Metadata.PairedScore = 0.5 post, err := LoadPosterior("testdata/ts_posterior.json") if err != nil { t.Fatal(err) @@ -149,6 +154,57 @@ func TestRoute_KeepsArgmaxWhenSameModel(t *testing.T) { if dec.Metadata.EffectiveKnobsHash != 99 { t.Fatal("keeping argmax must preserve cache key") } + if dec.Metadata.PairedModel != "claude-sonnet-4-6" { + t.Fatalf("argmax keep must preserve scorer runner-up sonnet, got %q", dec.Metadata.PairedModel) + } + if dec.Metadata.PairedScore != 0.5 { + t.Fatalf("argmax keep must preserve scorer paired score, got %v", dec.Metadata.PairedScore) + } +} + +func TestRoute_RepairsBandPairWhenServingRunnerUp(t *testing.T) { + // Scorer picks sonnet (argmax 0.90) with haiku (0.85) as the frozen + // runner-up. Zero-noise TS prefers haiku's higher posterior mean, so the + // served model becomes the former runner-up. The pin pair must not + // collapse to {haiku, haiku}: it recomputes to the best remaining peer, + // sonnet (the prior argmax). + scores := map[string]float32{ + "claude-sonnet-4-6": 0.90, + "claude-haiku-4-5": 0.85, + } + inner := clusterDecision(scores, []int{0}, "claude-sonnet-4-6", "anthropic") + inner.Metadata.PairedModel = "claude-haiku-4-5" + inner.Metadata.PairedProvider = "anthropic" + inner.Metadata.PairedScore = 0.85 + post := &Posterior{cells: map[int]map[string]Arm{ + 0: { + "claude-haiku-4-5": {Mean: 0.9, Variance: 0}, + "claude-sonnet-4-6": {Mean: 0.1, Variance: 0}, + }, + }} + b := New(&fakeInner{dec: inner}, post) + b.norm = func() float64 { return 0 } + b.trials = 0 + + dec, err := b.Route(context.Background(), router.Request{}) + if err != nil { + t.Fatal(err) + } + if dec.Model != "claude-haiku-4-5" { + t.Fatalf("expected Thompson pick haiku (runner-up), got %q", dec.Model) + } + if dec.Metadata.PairedModel == dec.Model { + t.Fatalf("band pair collapsed: Model and PairedModel are both %q", dec.Model) + } + if dec.Metadata.PairedModel != "claude-sonnet-4-6" { + t.Fatalf("expected runner-up recomputed to sonnet, got %q", dec.Metadata.PairedModel) + } + if dec.Metadata.PairedProvider != "anthropic" { + t.Fatalf("expected paired provider anthropic, got %q", dec.Metadata.PairedProvider) + } + if dec.Metadata.PairedScore != 0.90 { + t.Fatalf("expected paired score 0.90, got %v", dec.Metadata.PairedScore) + } } func TestRoute_UsesPerRequestProviderBinding(t *testing.T) { From febd91611a3b676ce38dfcbf98a14a0afe5c97fc Mon Sep 17 00:00:00 2001 From: N Rohith Reddy Date: Thu, 16 Jul 2026 15:33:51 -0400 Subject: [PATCH 2/2] style(bandit): keep test comments to one-line behavior asserts Match AGENTS.md: concise comments that state the asserted behavior, without scenario narration. Signed-off-by: N Rohith Reddy --- internal/router/bandit/bandit.go | 8 +++----- internal/router/bandit/bandit_test.go | 9 ++------- 2 files changed, 5 insertions(+), 12 deletions(-) diff --git a/internal/router/bandit/bandit.go b/internal/router/bandit/bandit.go index 2ee1b200b..75e88cec5 100644 --- a/internal/router/bandit/bandit.go +++ b/internal/router/bandit/bandit.go @@ -191,11 +191,9 @@ func (b *Router) annotate(dec *router.Decision, model, provider string, propensi } // repairBandPair recomputes the band pair's runner-up against the served -// model, writing it back through the metadata pointers. Picks the -// highest-scoring servable peer other than served (ties broken by name); -// clears the pair if none remains. Providers come from the live candidate -// list (resolved before annotate) so the prior argmax stays eligible even -// when it was only bound via Decision.Provider. +// model. Picks the highest-scoring servable peer other than served (ties +// broken by name); clears the pair if none remains. Providers come from the +// live candidate list so peers resolved only via Decision.Provider stay eligible. func repairBandPair(outModel, outProvider *string, outScore *float32, scores map[string]float32, cands []candidate, served string) { providers := make(map[string]string, len(cands)) for _, c := range cands { diff --git a/internal/router/bandit/bandit_test.go b/internal/router/bandit/bandit_test.go index f87d8d57f..03b82086e 100644 --- a/internal/router/bandit/bandit_test.go +++ b/internal/router/bandit/bandit_test.go @@ -131,10 +131,9 @@ func TestRoute_PropensityReflectsPosteriorOverlap(t *testing.T) { } func TestRoute_KeepsArgmaxWhenSameModel(t *testing.T) { + // Keeping the argmax must leave the scorer's PairedModel untouched. scores := map[string]float32{"claude-haiku-4-5": 0.9, "claude-sonnet-4-6": 0.5} inner := clusterDecision(scores, []int{0}, "claude-haiku-4-5", "anthropic") - // Scorer's runner-up; keeping the argmax must leave this pair untouched - // rather than recomputing it. inner.Metadata.PairedModel = "claude-sonnet-4-6" inner.Metadata.PairedProvider = "anthropic" inner.Metadata.PairedScore = 0.5 @@ -163,11 +162,7 @@ func TestRoute_KeepsArgmaxWhenSameModel(t *testing.T) { } func TestRoute_RepairsBandPairWhenServingRunnerUp(t *testing.T) { - // Scorer picks sonnet (argmax 0.90) with haiku (0.85) as the frozen - // runner-up. Zero-noise TS prefers haiku's higher posterior mean, so the - // served model becomes the former runner-up. The pin pair must not - // collapse to {haiku, haiku}: it recomputes to the best remaining peer, - // sonnet (the prior argmax). + // Serving the former runner-up must recompute PairedModel to the prior argmax. scores := map[string]float32{ "claude-sonnet-4-6": 0.90, "claude-haiku-4-5": 0.85,