diff --git a/cmd/server/main.go b/cmd/server/main.go index a1e7e7e..4046bc8 100644 --- a/cmd/server/main.go +++ b/cmd/server/main.go @@ -215,7 +215,7 @@ func run() error { // reqtiming wraps outermost, outside the limiter and retries: a call // queued at the limiter is the engine's time, and every attempt of a // retried call is one wait on the provider. - llmClient := reqtiming.WrapClient(limit.Client(newLimiter("llm", cfg.Engine.LLM.Concurrency, logger))(registry.Client())) + llmClient := reqtiming.WrapClient(withModelRetry(limit.Client(newLimiter("llm", cfg.Engine.LLM.Concurrency, logger))(registry.Client()), modelRetry)) // judge is always the registry's proxy. Code that takes a different // path without a Judge resolves it per job or request through // models.ActiveJudge; judgewalk falls back to treewalk. @@ -740,7 +740,23 @@ func byokFactory(provider, apiKey, model string) (llmgate.Client, string, error) if err != nil { return nil, "", err } - return client, model, nil + return withModelRetry(client, modelRetry), model, nil +} + +// modelRetry is how the server's generative calls ride out a provider's +// refusals: 429s and 5xx, retried with backoff (Retry-After honoured, +// clamped to MaxDelay). Vertex's shared quota refuses in bursts; with no +// retry, 28 of 50 /v1/answer calls failed on their first 429 (HAL-3206). +var modelRetry = retry.Config{MaxRetries: 5, BaseDelay: time.Second, MaxDelay: 20 * time.Second} + +// withModelRetry wraps a generative client in modelRetry. It sits outside +// the limiter, so every attempt takes a slot and a refusal that narrows the +// limit is seen before the next attempt. +func withModelRetry(c llmgate.Client, cfg retry.Config) llmgate.Client { + if c == nil { + return nil + } + return retry.New(cfg)(c) } // buildStrategy constructs the retrieval strategy named by diff --git a/cmd/server/retry_test.go b/cmd/server/retry_test.go new file mode 100644 index 0000000..3d25d74 --- /dev/null +++ b/cmd/server/retry_test.go @@ -0,0 +1,56 @@ +package main + +import ( + "context" + "sync/atomic" + "testing" + "time" + + "github.com/hallelx2/llmgate" + "github.com/hallelx2/llmgate/middleware/retry" +) + +// A provider that refuses twice and then answers yields the answer: the +// server's model client retries 429s (HAL-3206). +func TestModelClientRidesOutRefusals(t *testing.T) { + var calls atomic.Int32 + inner := &llmgate.Mock{Respond: func(context.Context, llmgate.Request) (*llmgate.Response, error) { + if calls.Add(1) <= 2 { + return nil, llmgate.NewLLMError("gemini", 429, "vertex: Resource exhausted. Please try again later.", nil) + } + return &llmgate.Response{Content: "capex was $1,577 million [1]"}, nil + }} + fast := retry.Config{MaxRetries: modelRetry.MaxRetries, BaseDelay: time.Millisecond, MaxDelay: 5 * time.Millisecond} + resp, err := withModelRetry(inner, fast).Complete(context.Background(), llmgate.Request{Model: "m"}) + if err != nil { + t.Fatalf("err = %v after %d calls", err, calls.Load()) + } + if resp.Content != "capex was $1,577 million [1]" || calls.Load() != 3 { + t.Errorf("content %q after %d calls", resp.Content, calls.Load()) + } +} + +// A request the provider rejects as malformed is not retried. +func TestModelClientDoesNotRetryABadRequest(t *testing.T) { + var calls atomic.Int32 + inner := &llmgate.Mock{Respond: func(context.Context, llmgate.Request) (*llmgate.Response, error) { + calls.Add(1) + return nil, llmgate.NewLLMError("gemini", 400, "invalid argument", nil) + }} + fast := retry.Config{MaxRetries: 5, BaseDelay: time.Millisecond, MaxDelay: time.Millisecond} + if _, err := withModelRetry(inner, fast).Complete(context.Background(), llmgate.Request{}); err == nil { + t.Fatal("want the 400 back") + } + if calls.Load() != 1 { + t.Errorf("a 400 was tried %d times", calls.Load()) + } +} + +func TestModelRetryIsConfigured(t *testing.T) { + if modelRetry.MaxRetries < 3 || modelRetry.MaxDelay < 10*time.Second { + t.Errorf("modelRetry = %+v: too thin to ride out a shared-quota burst", modelRetry) + } + if withModelRetry(nil, modelRetry) != nil { + t.Error("a nil client must stay nil, so the endpoint can report no model") + } +}