Coordinate retry attempts with backend permits

2026-09-23 16:48:36 +00:00
parent 65d5894648
commit ae4540822e
10 changed files with 541 additions and 4 deletions
--- a/internal/generation/coordinator.go
+++ b/internal/generation/coordinator.go
@@ -0,0 +1,186 @@
+// Package generation coordinates model attempts, backend permits, and bounded
+// built-in HTTP retries for one generation request.
+package generation
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"math/rand/v2"
+	"time"
+
+	"gitea.maximumdirect.net/eric/promptkit/internal/capacity"
+	"gitea.maximumdirect.net/eric/promptkit/internal/domain"
+	"gitea.maximumdirect.net/eric/promptkit/internal/llm"
+)
+
+const (
+	initialBackoff = 500 * time.Millisecond
+	maximumBackoff = 8 * time.Second
+)
+
+// Coordinator implements the internal model client using the existing backend
+// capacity manager. Exactly one of builtIn and injected is selected.
+type Coordinator struct {
+	manager  *capacity.Manager
+	builtIn  *llm.OpenAICompatibleClient
+	injected llm.Client
+	now      func() time.Time
+	jitter   func(time.Duration) time.Duration
+	wait     func(context.Context, time.Duration) error
+}
+
+// NewBuiltIn selects the built-in HTTP retry path.
+func NewBuiltIn(manager *capacity.Manager, client *llm.OpenAICompatibleClient) *Coordinator {
+	return &Coordinator{manager: manager, builtIn: client, now: time.Now, jitter: randomJitter, wait: waitContext}
+}
+
+// NewInjected selects one opaque client dispatch with no automatic retry.
+func NewInjected(manager *capacity.Manager, client llm.Client) *Coordinator {
+	return &Coordinator{manager: manager, injected: client, now: time.Now, jitter: randomJitter, wait: waitContext}
+}
+
+// ExhaustedError retains the final eligible status failure and the number of
+// attempts actually dispatched.
+type ExhaustedError struct {
+	attempts uint64
+	cause    error
+}
+
+func (e *ExhaustedError) Error() string {
+	if e == nil {
+		return "transient model retries exhausted"
+	}
+	return fmt.Sprintf("transient model retries exhausted after %d attempts", e.attempts)
+}
+
+func (e *ExhaustedError) Unwrap() error {
+	if e == nil {
+		return nil
+	}
+	return e.cause
+}
+
+func (e *ExhaustedError) Attempts() uint64 {
+	if e == nil {
+		return 0
+	}
+	return e.attempts
+}
+
+// Generate executes either one opaque client call or a frozen built-in HTTP
+// sequence with one permit per actual dispatch.
+func (c *Coordinator) Generate(ctx context.Context, req domain.GenerateRequest) (*domain.GenerateResponse, error) {
+	if c.injected != nil {
+		return c.generateInjected(ctx, req)
+	}
+	prepared, err := c.builtIn.PrepareGeneration(req)
+	if err != nil {
+		return nil, err
+	}
+
+	var attempts, retriesUsed uint64
+	for {
+		if err := ctx.Err(); err != nil {
+			return nil, err
+		}
+		outcome, dispatched := c.execute(ctx, req.Target.BackendID, prepared)
+		if dispatched {
+			attempts++
+		}
+		if outcome.Err == nil {
+			return outcome.Response, nil
+		}
+		if !dispatched || !eligibleStatus(outcome) {
+			return nil, outcome.Err
+		}
+		if retriesUsed == req.TransientRetries {
+			if retriesUsed == 0 {
+				return nil, outcome.Err
+			}
+			return nil, &ExhaustedError{attempts: attempts, cause: outcome.Err}
+		}
+		if err := ctx.Err(); err != nil {
+			return nil, err
+		}
+		retryOrdinal := retriesUsed + 1
+		delay := retryDelay(c.backoff(retryOrdinal), outcome.RetryNotBefore, c.now())
+		if err := c.wait(ctx, delay); err != nil {
+			return nil, err
+		}
+		if err := ctx.Err(); err != nil {
+			return nil, err
+		}
+		retriesUsed++
+	}
+}
+
+func (c *Coordinator) generateInjected(ctx context.Context, req domain.GenerateRequest) (*domain.GenerateResponse, error) {
+	release, err := c.manager.Acquire(ctx, req.Target.BackendID)
+	if err != nil {
+		return nil, err
+	}
+	defer release()
+	if err := ctx.Err(); err != nil {
+		return nil, err
+	}
+	return c.injected.Generate(ctx, req)
+}
+
+func (c *Coordinator) execute(ctx context.Context, backendID string, prepared *llm.PreparedGeneration) (outcome llm.AttemptOutcome, dispatched bool) {
+	release, err := c.manager.Acquire(ctx, backendID)
+	if err != nil {
+		return llm.AttemptOutcome{Err: err}, false
+	}
+	defer release()
+	if err := ctx.Err(); err != nil {
+		return llm.AttemptOutcome{Err: err}, false
+	}
+	return prepared.Execute(ctx), true
+}
+
+func eligibleStatus(outcome llm.AttemptOutcome) bool {
+	if !errors.Is(outcome.Err, llm.ErrUnexpectedStatus) {
+		return false
+	}
+	switch outcome.StatusCode {
+	case 429, 500, 502, 503, 504:
+		return true
+	default:
+		return false
+	}
+}
+
+func (c *Coordinator) backoff(ordinal uint64) time.Duration {
+	base := initialBackoff
+	for step := uint64(1); step < ordinal && base < maximumBackoff; step++ {
+		base *= 2
+	}
+	half := base / 2
+	return half + c.jitter(half)
+}
+
+func retryDelay(local time.Duration, notBefore, now time.Time) time.Duration {
+	if providerDelay := notBefore.Sub(now); providerDelay > local {
+		return providerDelay
+	}
+	return local
+}
+
+func randomJitter(max time.Duration) time.Duration {
+	return time.Duration(rand.Int64N(int64(max) + 1))
+}
+
+func waitContext(ctx context.Context, delay time.Duration) error {
+	if err := ctx.Err(); err != nil {
+		return err
+	}
+	timer := time.NewTimer(delay)
+	defer timer.Stop()
+	select {
+	case <-ctx.Done():
+		return ctx.Err()
+	case <-timer.C:
+		return ctx.Err()
+	}
+}
--- a/internal/generation/coordinator_test.go
+++ b/internal/generation/coordinator_test.go
@@ -0,0 +1,278 @@
+package generation
+
+import (
+	"context"
+	"errors"
+	"io"
+	"math"
+	"net/http"
+	"strings"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"gitea.maximumdirect.net/eric/promptkit/internal/capacity"
+	"gitea.maximumdirect.net/eric/promptkit/internal/domain"
+	"gitea.maximumdirect.net/eric/promptkit/internal/llm"
+)
+
+type transportFunc func(*http.Request) (*http.Response, error)
+
+func (f transportFunc) RoundTrip(req *http.Request) (*http.Response, error) { return f(req) }
+
+func response(status int, body string) *http.Response {
+	return &http.Response{StatusCode: status, Body: io.NopCloser(strings.NewReader(body)), Header: make(http.Header)}
+}
+
+func testRequest(retries uint64) domain.GenerateRequest {
+	return domain.GenerateRequest{
+		Prompt:           domain.RenderedPrompt{Messages: []domain.RenderedMessage{{Role: "user", Content: "hello"}}},
+		Target:           domain.ExecutionTarget{BackendID: "limited", Endpoint: "https://example.invalid/v1", Model: "model"},
+		TransientRetries: retries,
+	}
+}
+
+func testCoordinator(t *testing.T, transport transportFunc) *Coordinator {
+	t.Helper()
+	manager, err := capacity.NewManager(map[string]domain.BackendCapacityPolicy{
+		"limited": {ConcurrencyLimit: 1, QueueCapacity: 1},
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+	client, err := llm.NewOpenAICompatibleClient(llm.OpenAICompatibleConfig{HTTPClient: &http.Client{Transport: transport}})
+	if err != nil {
+		t.Fatal(err)
+	}
+	coordinator := NewBuiltIn(manager, client)
+	coordinator.jitter = func(time.Duration) time.Duration { return 0 }
+	coordinator.wait = func(ctx context.Context, _ time.Duration) error { return ctx.Err() }
+	return coordinator
+}
+
+func TestBuiltInRetryEligibilityAndExhaustion(t *testing.T) {
+	for _, tc := range []struct {
+		status   int
+		eligible bool
+	}{
+		{status: 429, eligible: true}, {status: 500, eligible: true},
+		{status: 502, eligible: true}, {status: 503, eligible: true},
+		{status: 504, eligible: true}, {status: 400}, {status: 501},
+	} {
+		t.Run(http.StatusText(tc.status), func(t *testing.T) {
+			calls := 0
+			coordinator := testCoordinator(t, func(*http.Request) (*http.Response, error) {
+				calls++
+				if calls == 1 {
+					return response(tc.status, `{"error":{}}`), nil
+				}
+				return response(200, `{"choices":[{"message":{"content":"ok"}}]}`), nil
+			})
+			result, err := coordinator.Generate(context.Background(), testRequest(1))
+			if tc.eligible {
+				if err != nil || result == nil || result.Content != "ok" || calls != 2 {
+					t.Fatalf("result=%+v, error=%v, calls=%d", result, err, calls)
+				}
+			} else if result != nil || err == nil || calls != 1 {
+				t.Fatalf("result=%+v, error=%v, calls=%d", result, err, calls)
+			}
+		})
+	}
+
+	for _, budget := range []uint64{0, 2} {
+		calls := 0
+		coordinator := testCoordinator(t, func(*http.Request) (*http.Response, error) {
+			calls++
+			return response(503, `{"error":{}}`), nil
+		})
+		_, err := coordinator.Generate(context.Background(), testRequest(budget))
+		var exhausted *ExhaustedError
+		if !errors.Is(err, llm.ErrUnexpectedStatus) || calls != int(budget+1) || errors.As(err, &exhausted) != (budget > 0) {
+			t.Fatalf("budget=%d calls=%d error=%v", budget, calls, err)
+		}
+		if budget > 0 && exhausted.Attempts() != budget+1 {
+			t.Fatalf("attempts=%d", exhausted.Attempts())
+		}
+	}
+
+	calls := 0
+	coordinator := testCoordinator(t, func(*http.Request) (*http.Response, error) {
+		calls++
+		if calls == 1 {
+			return response(503, `{"error":{}}`), nil
+		}
+		return response(400, `{"error":{}}`), nil
+	})
+	_, err := coordinator.Generate(context.Background(), testRequest(3))
+	var exhausted *ExhaustedError
+	if calls != 2 || !errors.Is(err, llm.ErrUnexpectedStatus) || errors.As(err, &exhausted) {
+		t.Fatalf("later permanent failure: calls=%d, error=%v", calls, err)
+	}
+}
+
+func TestBuiltInRetryWaitHonorsContextAndProviderMinimum(t *testing.T) {
+	when := time.Date(2026, 9, 23, 12, 0, 0, 0, time.UTC)
+	coordinator := testCoordinator(t, func(*http.Request) (*http.Response, error) { return response(503, `{"error":{}}`), nil })
+	if got := coordinator.backoff(1); got != 250*time.Millisecond {
+		t.Fatalf("first backoff=%v", got)
+	}
+	if got := coordinator.backoff(math.MaxUint64); got != 4*time.Second {
+		t.Fatalf("capped backoff=%v", got)
+	}
+	coordinator.jitter = func(max time.Duration) time.Duration { return max }
+	if got := coordinator.backoff(math.MaxUint64); got != 8*time.Second {
+		t.Fatalf("maximum equal jitter=%v", got)
+	}
+	if got := retryDelay(500*time.Millisecond, when.Add(3*time.Second), when); got != 3*time.Second {
+		t.Fatalf("provider minimum=%v", got)
+	}
+	if got := retryDelay(500*time.Millisecond, when.Add(-time.Second), when); got != 500*time.Millisecond {
+		t.Fatalf("past provider minimum=%v", got)
+	}
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	calls := 0
+	coordinator = testCoordinator(t, func(*http.Request) (*http.Response, error) {
+		calls++
+		return response(503, `{"error":{}}`), nil
+	})
+	coordinator.wait = func(ctx context.Context, delay time.Duration) error {
+		if delay < 250*time.Millisecond {
+			t.Fatalf("delay=%v", delay)
+		}
+		cancel()
+		return ctx.Err()
+	}
+	_, err := coordinator.Generate(ctx, testRequest(2))
+	if !errors.Is(err, context.Canceled) || calls != 1 {
+		t.Fatalf("canceled wait: error=%v calls=%d", err, calls)
+	}
+}
+
+func TestBuiltInRetryUsesSeparateAttemptDeadlines(t *testing.T) {
+	var contexts []context.Context
+	coordinator := testCoordinator(t, func(req *http.Request) (*http.Response, error) {
+		contexts = append(contexts, req.Context())
+		if len(contexts) == 1 {
+			return response(503, `{"error":{}}`), nil
+		}
+		return response(200, `{"choices":[{"message":{"content":"ok"}}]}`), nil
+	})
+	req := testRequest(1)
+	req.Target.TimeoutSeconds = 1
+	if _, err := coordinator.Generate(context.Background(), req); err != nil {
+		t.Fatal(err)
+	}
+	if len(contexts) != 2 || contexts[0] == contexts[1] {
+		t.Fatalf("attempt contexts=%d", len(contexts))
+	}
+	for _, ctx := range contexts {
+		if _, ok := ctx.Deadline(); !ok || !errors.Is(ctx.Err(), context.Canceled) {
+			t.Fatalf("attempt deadline/error=%t/%v", ok, ctx.Err())
+		}
+	}
+}
+
+func TestCanceledCapacityWaitDispatchesNothing(t *testing.T) {
+	calls := 0
+	coordinator := testCoordinator(t, func(*http.Request) (*http.Response, error) {
+		calls++
+		return response(200, `{"choices":[{"message":{"content":"ok"}}]}`), nil
+	})
+	release, err := coordinator.manager.Acquire(context.Background(), "limited")
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer release()
+	ctx, cancel := context.WithCancel(context.Background())
+	result := make(chan error, 1)
+	go func() { _, err := coordinator.Generate(ctx, testRequest(2)); result <- err }()
+	cancel()
+	if err := <-result; !errors.Is(err, context.Canceled) || calls != 0 {
+		t.Fatalf("capacity cancellation: error=%v calls=%d", err, calls)
+	}
+}
+
+func TestBuiltInRetryFreesPermitAndFreezesCredential(t *testing.T) {
+	const envName = "PROMPTKIT_GENERATION_RETRY_TEST_KEY"
+	t.Setenv(envName, "first-key")
+	var calls atomic.Int64
+	var firstAuth, retryAuth string
+	coordinator := testCoordinator(t, func(req *http.Request) (*http.Response, error) {
+		switch calls.Add(1) {
+		case 1:
+			firstAuth = req.Header.Get("Authorization")
+			return response(503, `{"error":{}}`), nil
+		case 2:
+			return response(200, `{"choices":[{"message":{"content":"other"}}]}`), nil
+		default:
+			retryAuth = req.Header.Get("Authorization")
+			return response(200, `{"choices":[{"message":{"content":"retry"}}]}`), nil
+		}
+	})
+	entered := make(chan struct{})
+	resume := make(chan struct{})
+	coordinator.wait = func(ctx context.Context, _ time.Duration) error {
+		close(entered)
+		select {
+		case <-resume:
+			return ctx.Err()
+		case <-ctx.Done():
+			return ctx.Err()
+		}
+	}
+	req := testRequest(1)
+	req.Target.APIKeyEnv = envName
+	first := make(chan error, 1)
+	go func() { _, err := coordinator.Generate(context.Background(), req); first <- err }()
+	<-entered
+	t.Setenv(envName, "later-key")
+	other, err := coordinator.Generate(context.Background(), testRequest(0))
+	if err != nil || other == nil || other.Content != "other" {
+		t.Fatalf("other call while backoff blocked = %+v, %v", other, err)
+	}
+	close(resume)
+	if err := <-first; err != nil {
+		t.Fatal(err)
+	}
+	if calls.Load() != 3 || firstAuth != "Bearer first-key" || retryAuth != firstAuth {
+		t.Fatalf("calls=%d, first auth=%q, retry auth=%q", calls.Load(), firstAuth, retryAuth)
+	}
+}
+
+func TestInjectedClientDispatchesOnceAndReleasesAfterPanic(t *testing.T) {
+	manager, err := capacity.NewManager(map[string]domain.BackendCapacityPolicy{"limited": {ConcurrencyLimit: 1}})
+	if err != nil {
+		t.Fatal(err)
+	}
+	client := &injectedClient{panicFirst: true}
+	coordinator := NewInjected(manager, client)
+	func() {
+		defer func() {
+			if recover() == nil {
+				t.Fatal("expected injected panic")
+			}
+		}()
+		_, _ = coordinator.Generate(context.Background(), testRequest(5))
+	}()
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	defer cancel()
+	result, err := coordinator.Generate(ctx, testRequest(5))
+	if err != nil || result == nil || client.calls != 2 {
+		t.Fatalf("after panic result=%+v error=%v calls=%d", result, err, client.calls)
+	}
+}
+
+type injectedClient struct {
+	calls      int
+	panicFirst bool
+}
+
+func (c *injectedClient) Generate(context.Context, domain.GenerateRequest) (*domain.GenerateResponse, error) {
+	c.calls++
+	if c.panicFirst && c.calls == 1 {
+		panic("injected panic")
+	}
+	return &domain.GenerateResponse{Content: "ok"}, nil
+}