xinmi/internal/models/requestbudget/context.go

// Package requestbudget carries a durable request reservation to provider transports.
package requestbudget

import (
	"context"
	"errors"
	"fmt"
	"net/http"
)

var ErrStopped = errors.New("wiki recovery budget stopped")

type Guard interface {
	Key() string
	Check(context.Context) error
	Reserve(context.Context, string, string) error
}

type guardKey struct{}
type phaseKey struct{}
type modelKey struct{}

func WithGuard(ctx context.Context, guard Guard) context.Context {
	return context.WithValue(ctx, guardKey{}, guard)
}

func FromContext(ctx context.Context) Guard {
	guard, _ := ctx.Value(guardKey{}).(Guard)
	return guard
}

func WithPhase(ctx context.Context, phase string) context.Context {
	return context.WithValue(ctx, phaseKey{}, phase)
}

func Phase(ctx context.Context) string {
	phase, _ := ctx.Value(phaseKey{}).(string)
	return phase
}

func WithModel(ctx context.Context, modelID string) context.Context {
	return context.WithValue(ctx, modelKey{}, modelID)
}

func Check(ctx context.Context) error {
	if guard := FromContext(ctx); guard != nil {
		return guard.Check(ctx)
	}
	return ctx.Err()
}

// RequireSupported prevents a controlled request from reaching an adapter
// whose actual sends are not instrumented. Ordinary requests are unaffected.
func RequireSupported(ctx context.Context, supported bool) error {
	if FromContext(ctx) != nil && !supported {
		return fmt.Errorf("%w: unsupported provider transport", ErrStopped)
	}
	return nil
}

// Transport reserves each RoundTrip, including redirects and SDK fallbacks.
// Failed or cancelled sends intentionally consume their reservation.
type Transport struct{ Base http.RoundTripper }

func (t Transport) RoundTrip(req *http.Request) (*http.Response, error) {
	if err := req.Context().Err(); err != nil {
		return nil, err
	}
	if guard := FromContext(req.Context()); guard != nil {
		modelID, _ := req.Context().Value(modelKey{}).(string)
		if modelID == "" {
			return nil, fmt.Errorf("%w: missing provider model identity", ErrStopped)
		}
		if err := guard.Reserve(req.Context(), modelID, Phase(req.Context())); err != nil {
			return nil, fmt.Errorf("%w: %w", ErrStopped, err)
		}
		// A replayable POST with an idempotency header can otherwise be retried
		// inside http.Transport without reserving again. Clone so Client.Do
		// keeps the original GetBody for separately accounted 307/308 redirects.
		req = req.Clone(req.Context())
		req.GetBody = nil
	}
	base := t.Base
	if base == nil {
		base = http.DefaultTransport
	}
	return base.RoundTrip(req)
}