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