promptkit/internal/generation/coordinator.go

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