diff --git a/router/deployment_router.go b/router/deployment_router.go index a220b89..5beda83 100644 --- a/router/deployment_router.go +++ b/router/deployment_router.go @@ -146,8 +146,12 @@ func (r *DeploymentRouter) Chat(ctx context.Context, messages []client.EyrieMess if attempts < 1 { attempts = 1 } + // Track the deployment that just failed this stage so the next attempt + // prefers a different deployment when one is available, instead of + // re-selecting the same dead endpoint up to stage.Retries times. + recentlyFailed := "" for attempt := 0; attempt < attempts; attempt++ { - choice := selectDeploymentChoice(choices) + choice := selectDeploymentChoice(choices, recentlyFailed) resp, err := r.chatWithDeployment(ctx, messages, opts, target, choice.DeploymentID) if err == nil { r.recordSuccess(choice.DeploymentID) @@ -161,6 +165,7 @@ func (r *DeploymentRouter) Chat(ctx context.Context, messages []client.EyrieMess } return nil, err } + recentlyFailed = choice.DeploymentID } } if lastErr == nil { @@ -189,8 +194,12 @@ func (r *DeploymentRouter) StreamChat(ctx context.Context, messages []client.Eyr if attempts < 1 { attempts = 1 } + // Prefer a different deployment than the one that just failed + // this stage, instead of re-selecting the same dead endpoint + // up to stage.Retries times. + recentlyFailed := "" for attempt := 0; attempt < attempts; attempt++ { - choice := selectDeploymentChoice(choices) + choice := selectDeploymentChoice(choices, recentlyFailed) fallback, err := r.streamWithDeployment(streamCtx, out, messages, opts, target, choice.DeploymentID) if err == nil { r.recordSuccess(choice.DeploymentID) @@ -215,6 +224,7 @@ func (r *DeploymentRouter) StreamChat(ctx context.Context, messages []client.Eyr } return } + recentlyFailed = choice.DeploymentID } } if lastErr == nil { @@ -560,25 +570,41 @@ func offeringSupportsTools(offering catalog.ModelOffering, tools []string) bool return true } -func selectDeploymentChoice(choices []DeploymentChoice) DeploymentChoice { +// selectDeploymentChoice picks a deployment from choices using weighted random +// selection. The deployment whose ID equals exclude is skipped when more than +// one option is available, so a retry after a failure prefers a different +// endpoint instead of hammering the same dead one. +func selectDeploymentChoice(choices []DeploymentChoice, exclude string) DeploymentChoice { if len(choices) == 1 { return choices[0] } + alternatives := choices + if exclude != "" { + filtered := make([]DeploymentChoice, 0, len(choices)) + for _, c := range choices { + if c.DeploymentID != exclude { + filtered = append(filtered, c) + } + } + if len(filtered) > 0 { + alternatives = filtered + } + } total := 0 - for _, choice := range choices { + for _, choice := range alternatives { total += choice.Weight } if total <= 0 { - return choices[0] + return alternatives[0] } n := rand.IntN(total) // #nosec G404 -- non-cryptographic weighted load-balancing choice, not a security decision - for _, choice := range choices { + for _, choice := range alternatives { n -= choice.Weight if n < 0 { return choice } } - return choices[len(choices)-1] + return alternatives[len(alternatives)-1] } func isOutputEvent(event client.EyrieStreamEvent) bool { diff --git a/router/deployment_router_test.go b/router/deployment_router_test.go index 9d89c20..d29b117 100644 --- a/router/deployment_router_test.go +++ b/router/deployment_router_test.go @@ -16,11 +16,13 @@ type deploymentMockProvider struct { lastModel string lastTools []client.EyrieTool streamDone bool + callCount int } func (m *deploymentMockProvider) Chat(_ context.Context, _ []client.EyrieMessage, opts client.ChatOptions) (*client.EyrieResponse, error) { m.lastModel = opts.Model m.lastTools = opts.Tools + m.callCount++ if m.err != nil { return nil, m.err } @@ -360,3 +362,50 @@ func TestDeploymentRouterNativeMimoUsesConfiguredXiaomiDeployment(t *testing.T) t.Fatalf("native model = %q", mimo.lastModel) } } + +// TestDeploymentRouterRetriesPreferDifferentEndpoint verifies the fix for +// "deployment retry can re-select the same dead deployment": when a stage +// has multiple deployments and the first choice fails transiently, the next +// attempt should prefer a different deployment (and the healthy one is +// reached) instead of retrying the same dead endpoint up to stage.Retries. +func TestDeploymentRouterRetriesPreferDifferentEndpoint(t *testing.T) { + t.Parallel() + dead := &deploymentMockProvider{name: "direct", err: fmt.Errorf("HTTP 503 unavailable")} + healthy := &deploymentMockProvider{name: "vertex"} + r, err := NewDeploymentRouter(DeploymentRouterOptions{ + Catalog: testCompiledCatalog(t), + Deployments: map[string]DeploymentAdapter{ + "anthropic-direct": {Provider: dead}, + "anthropic-vertex": {Provider: healthy}, + }, + Routing: RoutingPolicy{Providers: map[string][]RoutingStage{"anthropic": {{ + Deployments: []DeploymentChoice{ + {DeploymentID: "anthropic-direct", Weight: 100}, + {DeploymentID: "anthropic-vertex", Weight: 1}, + }, + Retries: 3, + }}}}, + }) + if err != nil { + t.Fatal(err) + } + + resp, err := r.Chat(context.Background(), + []client.EyrieMessage{{Role: "user", Content: "hi"}}, + client.ChatOptions{Model: "anthropic/claude-sonnet-4-6"}) + if err != nil { + t.Fatalf("Chat: %v", err) + } + if resp.Content != "from vertex" { + t.Fatalf("expected vertex (healthy) deployment after direct (dead) failed, got %q", resp.Content) + } + // The dead provider is tried at most once (to discover the failure); the + // retry must prefer the healthy endpoint instead of re-selecting the same + // dead one up to stage.Retries times. + if dead.callCount > 1 { + t.Fatalf("dead deployment retried %d times; want at most 1", dead.callCount) + } + if healthy.callCount != 1 { + t.Fatalf("healthy deployment called %d times; want 1", healthy.callCount) + } +}