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"

	"go.maximumdirect.net/promptkit/internal/capacity"
	"go.maximumdirect.net/promptkit/internal/domain"
	"go.maximumdirect.net/promptkit/internal/llm"
	"go.maximumdirect.net/promptkit/internal/observation"
)

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) {
	response, err, _ := c.generate(ctx, req, false)
	return response, err
}

// GenerateWithEvidence runs the built-in retry path and returns its factual
// attempt evidence alongside any response retained before cancellation.
func (c *Coordinator) GenerateWithEvidence(
	ctx context.Context,
	req domain.GenerateRequest,
) (*domain.GenerateResponse, error, GenerationEvidence) {
	return c.generate(ctx, req, true)
}

// PreparationError identifies a built-in request preparation failure.
type PreparationError struct{ Cause error }

func (e *PreparationError) Error() string { return e.Cause.Error() }
func (e *PreparationError) Unwrap() error { return e.Cause }

func (c *Coordinator) generate(
	ctx context.Context,
	req domain.GenerateRequest,
	captureEvidence bool,
) (*domain.GenerateResponse, error, GenerationEvidence) {
	if c.injected != nil {
		return c.generateInjected(ctx, req, captureEvidence)
	}
	prepared, err := c.builtIn.PrepareGeneration(req)
	if err != nil {
		return nil, &PreparationError{Cause: err}, GenerationEvidence{PreparationFailed: true}
	}

	evidence := GenerationEvidence{}
	if captureEvidence {
		evidence.InitialEndpoint = prepared.InitialEndpoint()
	}
	var usage observation.UsageAccumulator
	finish := func(response *domain.GenerateResponse, err error) (*domain.GenerateResponse, error, GenerationEvidence) {
		if captureEvidence {
			evidence.Usage = usage.Observed()
		}
		return response, err, evidence
	}
	var attempts, retriesUsed uint64
	recorder := observation.FromContext(ctx)
	generation := req.Generation
	if generation == 0 {
		generation = 1
	}
	for {
		if err := ctx.Err(); err != nil {
			return finish(nil, err)
		}
		ordinal := attempts + 1
		result := c.execute(ctx, req.Target.BackendID, prepared, captureEvidence)
		if result.Dispatched {
			attempts++
		}
		recordAttempt(result, false, captureEvidence, &evidence, &usage, recorder, generation, ordinal)
		if result.Outcome.Err == nil {
			if err := ctx.Err(); err != nil {
				if captureEvidence {
					return finish(result.Outcome.Response, err)
				}
				return finish(nil, err)
			}
			return finish(result.Outcome.Response, nil)
		}
		if !result.Dispatched || !eligibleStatus(result.Outcome) {
			if captureEvidence {
				return finish(result.Outcome.Response, result.Outcome.Err)
			}
			return finish(nil, result.Outcome.Err)
		}
		if retriesUsed == req.TransientRetries {
			if retriesUsed == 0 {
				return finish(nil, result.Outcome.Err)
			}
			return finish(nil, &ExhaustedError{attempts: attempts, cause: result.Outcome.Err})
		}
		if err := ctx.Err(); err != nil {
			return finish(nil, err)
		}
		retryOrdinal := retriesUsed + 1
		delay := retryDelay(c.backoff(retryOrdinal), result.Outcome.RetryNotBefore, c.now())
		waitStart := c.now()
		waitErr := c.wait(ctx, delay)
		waitDuration := elapsed(waitStart, c.now())
		if captureEvidence {
			evidence.RetryDelays = append(evidence.RetryDelays, waitDuration)
		}
		if recorder != nil {
			recorder.Delay(observation.DelayData{Generation: generation, Attempt: attempts,
				Elapsed: waitDuration, Failure: ClassifyFailure(waitErr)})
		}
		if waitErr != nil {
			return finish(nil, waitErr)
		}
		if err := ctx.Err(); err != nil {
			return finish(nil, err)
		}
		retriesUsed++
	}
}

func (c *Coordinator) generateInjected(ctx context.Context, req domain.GenerateRequest, captureEvidence bool) (*domain.GenerateResponse, error, GenerationEvidence) {
	evidence := GenerationEvidence{}
	var usage observation.UsageAccumulator
	generation := req.Generation
	if generation == 0 {
		generation = 1
	}
	result := c.executeInjected(ctx, req)
	recordAttempt(result, true, captureEvidence, &evidence, &usage, observation.FromContext(ctx), generation, 1)
	if captureEvidence {
		evidence.Usage = usage.Observed()
	}
	if result.Outcome.Err != nil {
		return nil, result.Outcome.Err, evidence
	}
	if err := ctx.Err(); err != nil {
		if captureEvidence {
			return result.Outcome.Response, err, evidence
		}
		return nil, err, evidence
	}
	return result.Outcome.Response, nil, evidence
}

func (c *Coordinator) executeInjected(ctx context.Context, req domain.GenerateRequest) attemptResult {
	result := attemptResult{StartedAt: c.now()}
	release, err := c.manager.Acquire(ctx, req.Target.BackendID)
	result.CapacityWait = c.capacityWait(req.Target.BackendID, result.StartedAt)
	if err != nil {
		result.Outcome.Err = err
		return result
	}
	defer release()
	if err := ctx.Err(); err != nil {
		result.Outcome.Err = err
		return result
	}
	start := c.now()
	result.Dispatched = true
	requested := req.StructuredOutput != nil
	result.Outcome.StructuredOutputRequested = &requested
	result.Outcome.Response, result.Outcome.Err = c.injected.Generate(ctx, req)
	result.RequestDuration = elapsed(start, c.now())
	if result.Outcome.Err == nil && result.Outcome.Response == nil {
		result.Outcome.Err = llm.ErrMalformedResponse
	}
	if result.Outcome.Err == nil {
		result.Outcome.Usage = result.Outcome.Response.Usage
		result.Outcome.UsagePresence = result.Outcome.Response.UsagePresence
	}
	return result
}

type attemptResult struct {
	Outcome         llm.AttemptOutcome
	Dispatched      bool
	CapacityWait    time.Duration
	RequestDuration time.Duration
	StartedAt       time.Time
}

// recordAttempt runs after the attempt has released its permit, including when
// the recorder calls application code. Injected clients expose no HTTP facts.
func recordAttempt(result attemptResult, injected, capture bool, evidence *GenerationEvidence, usage *observation.UsageAccumulator, recorder *observation.Recorder, generation, ordinal uint64) {
	outcome := result.Outcome
	failure := ClassifyFailure(outcome.Err)
	if capture {
		attempt := AttemptEvidence{
			StartedAt: result.StartedAt.UTC(), Dispatched: result.Dispatched,
			CapacityWait: result.CapacityWait, RequestDuration: result.RequestDuration,
			Failure: failure,
		}
		if injected {
			attempt.ResponseReceived = outcome.Response != nil
			attempt.GenerationCompleted = outcome.Err == nil && outcome.Response != nil
			if result.Dispatched {
				attempt.CapacityWait = 0
			}
		} else {
			attempt.ResponseReceived = outcome.Evidence.ResponseReceived
			attempt.HTTP = cloneHTTPEvidence(outcome.Evidence)
			attempt.ResponseDecoded = outcome.ResponseDecoded
			attempt.GenerationCompleted = outcome.GenerationCompleted
		}
		if outcome.StructuredOutputRequested != nil {
			requested := *outcome.StructuredOutputRequested
			attempt.StructuredOutputRequested = &requested
		}
		if result.Dispatched {
			usage.Add(outcome.Usage, outcome.UsagePresence)
			var single observation.UsageAccumulator
			single.Add(outcome.Usage, outcome.UsagePresence)
			attempt.Usage = single.Observed()
		}
		evidence.Attempts = append(evidence.Attempts, attempt)
	}
	if recorder != nil {
		recorder.Attempt(observation.AttemptData{Generation: generation, Attempt: ordinal,
			Dispatched: result.Dispatched, CapacityWait: result.CapacityWait, GenerationDuration: result.RequestDuration,
			Outcome: attemptOutcome(outcome.Err), Failure: failure, StatusCode: outcome.StatusCode,
			Usage: outcome.Usage, UsagePresence: outcome.UsagePresence})
	}
}

func (c *Coordinator) execute(ctx context.Context, backendID string, prepared *llm.PreparedGeneration, captureEvidence bool) attemptResult {
	result := attemptResult{StartedAt: c.now()}
	release, err := c.manager.Acquire(ctx, backendID)
	result.CapacityWait = c.capacityWait(backendID, result.StartedAt)
	if err != nil {
		result.Outcome.Err = err
		return result
	}
	defer release()
	if err := ctx.Err(); err != nil {
		result.Outcome.Err = err
		return result
	}
	dispatchStart := c.now()
	if captureEvidence {
		result.Outcome = prepared.ExecuteWithEvidence(ctx)
		result.Dispatched = result.Outcome.Evidence.DispatchInvoked
	} else {
		result.Outcome = prepared.Execute(ctx)
		result.Dispatched = true
	}
	result.RequestDuration = elapsed(dispatchStart, c.now())
	return result
}

func (c *Coordinator) capacityWait(backendID string, start time.Time) time.Duration {
	if !c.manager.Limited(backendID) {
		return 0
	}
	return elapsed(start, c.now())
}

func elapsed(start, end time.Time) time.Duration {
	if d := end.Sub(start); d > 0 {
		return d
	}
	return 0
}

func attemptOutcome(err error) domain.ExecutionOutcome {
	if err == nil {
		return domain.ExecutionOutcomeSuccess
	}
	return domain.ExecutionOutcomeError
}

// ClassifyFailure maps an execution error to the shared content-free category.
func ClassifyFailure(err error) domain.ExecutionFailureKind {
	if err == nil {
		return domain.ExecutionFailureNone
	}
	if errors.Is(err, context.Canceled) {
		return domain.ExecutionFailureCanceled
	}
	if errors.Is(err, context.DeadlineExceeded) {
		return domain.ExecutionFailureDeadline
	}
	if errors.Is(err, llm.ErrUnexpectedStatus) {
		return domain.ExecutionFailureProviderStatus
	}
	if errors.Is(err, llm.ErrTransport) {
		return domain.ExecutionFailureTransport
	}
	if errors.Is(err, llm.ErrMalformedResponse) {
		return domain.ExecutionFailureMalformedResponse
	}
	if errors.Is(err, llm.ErrInvalidRequest) {
		return domain.ExecutionFailureInvalidRequest
	}
	return domain.ExecutionFailureGeneration
}

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