From ec4e9f7131f9fd7702f56808b6b101106c2f8997 Mon Sep 17 00:00:00 2001
From: hallelx2 <halleluyaholudele@gmail.com>
Date: Sat, 17 Oct 2026 00:59:42 +0100
Subject: [PATCH] fix(server): retry the model client on 429 and 5xx

The deployed server built its generative client from the limiter and the
model registry, with no retry. The comment above it said it sat "outside the
limiter and retries", but nothing retried. The Judge had retry.NewJudge;
the answer model and BYOK clients had nothing. On vlb-engine, 28 of the
pilot's 50 /v1/answer calls returned 500 on their first Vertex 429.

Both now go through retry.New: 5 retries, exponential backoff from 1 s,
capped at 20 s, Retry-After honoured. It sits outside the limiter, so every
attempt takes a slot. 400s and the other non-transient classes are not
retried.
---
 cmd/server/main.go       | 20 ++++++++++++--
 cmd/server/retry_test.go | 56 ++++++++++++++++++++++++++++++++++++++++
 2 files changed, 74 insertions(+), 2 deletions(-)
 create mode 100644 cmd/server/retry_test.go

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")
+	}
+}