neural-hunt/internal/artifact/budget.go

package artifact

import (
	"context"
	"errors"
	"fmt"
	"sync"
	"time"

	"neuralhunt/internal/settings"
)

var ErrOpenAIBudgetExceeded = errors.New("OpenAI cost circuit breaker active")

type OpenAIBudgetError struct {
	Reason     string
	RetryAfter time.Duration
}

func (e *OpenAIBudgetError) Error() string {
	if e == nil {
		return ErrOpenAIBudgetExceeded.Error()
	}
	return fmt.Sprintf("%s: %s", ErrOpenAIBudgetExceeded, e.Reason)
}

func (e *OpenAIBudgetError) Unwrap() error { return ErrOpenAIBudgetExceeded }

type openAIBudgetSnapshot struct {
	Calls1H   int64
	Calls24H  int64
	Cost24H   float64
	Oldest1H  int64
	Oldest24H int64
}

var openAIBudgetMu sync.Mutex

func (w *Worker) openAIBudgetSnapshot(ctx context.Context) (openAIBudgetSnapshot, error) {
	now := time.Now().UTC()
	oneHour := now.Add(-time.Hour).UnixMilli()
	day := now.Add(-24 * time.Hour).UnixMilli()
	var x openAIBudgetSnapshot
	if err := w.db.QueryRowContext(ctx, `SELECT
		COALESCE(sum(CASE WHEN created_at>=? THEN 1 ELSE 0 END),0),
		count(*),COALESCE(sum(estimated_cost_usd),0),
		COALESCE(min(CASE WHEN created_at>=? THEN created_at END),0),
		COALESCE(min(created_at),0)
		FROM artifact_api_usage
		WHERE provider='openai' AND created_at>=?`, oneHour, oneHour, day).
		Scan(&x.Calls1H, &x.Calls24H, &x.Cost24H, &x.Oldest1H, &x.Oldest24H); err != nil {
		return x, err
	}
	return x, nil
}

func retryFromOldest(oldestMS int64, window time.Duration) time.Duration {
	if oldestMS <= 0 {
		return time.Minute
	}
	d := time.Until(time.UnixMilli(oldestMS).UTC().Add(window))
	if d < time.Second {
		d = time.Second
	}
	return d
}

func (w *Worker) checkOpenAIBudget(ctx context.Context, cfg settings.Runtime) error {
	if cfg.OpenAIMaxCalls1H <= 0 && cfg.OpenAIMaxCalls24H <= 0 && cfg.OpenAIMaxCost24HUSD <= 0 {
		return nil
	}
	if cfg.OpenAIMaxCost24HUSD > 0 {
		if _, ok := openAIImagePricing(cfg.ArtifactModel); !ok {
			return &OpenAIBudgetError{Reason: "daily USD budget is enabled but this image model has no pinned local price", RetryAfter: 10 * time.Minute}
		}
	}
	x, err := w.openAIBudgetSnapshot(ctx)
	if err != nil {
		return err
	}
	var wait time.Duration
	var reasons []string
	if cfg.OpenAIMaxCalls1H > 0 && x.Calls1H >= int64(cfg.OpenAIMaxCalls1H) {
		reasons = append(reasons, fmt.Sprintf("%d/%d OpenAI image calls in the last hour", x.Calls1H, cfg.OpenAIMaxCalls1H))
		wait = maxDuration(wait, retryFromOldest(x.Oldest1H, time.Hour))
	}
	if cfg.OpenAIMaxCalls24H > 0 && x.Calls24H >= int64(cfg.OpenAIMaxCalls24H) {
		reasons = append(reasons, fmt.Sprintf("%d/%d OpenAI image calls in the last 24h", x.Calls24H, cfg.OpenAIMaxCalls24H))
		wait = maxDuration(wait, retryFromOldest(x.Oldest24H, 24*time.Hour))
	}
	if cfg.OpenAIMaxCost24HUSD > 0 && x.Cost24H+cfg.OpenAIBudgetReserveUSD > cfg.OpenAIMaxCost24HUSD {
		reasons = append(reasons, fmt.Sprintf("estimated cost $%.4f + $%.4f safety reserve exceeds $%.4f/24h", x.Cost24H, cfg.OpenAIBudgetReserveUSD, cfg.OpenAIMaxCost24HUSD))
		// Re-evaluate periodically because several old calls may age out together.
		wait = maxDuration(wait, 5*time.Minute)
	}
	if len(reasons) > 0 {
		return &OpenAIBudgetError{Reason: joinBudgetReasons(reasons), RetryAfter: wait}
	}
	return nil
}

func joinBudgetReasons(parts []string) string {
	out := ""
	for i, p := range parts {
		if i > 0 {
			out += "; "
		}
		out += p
	}
	return out
}

func maxDuration(a, b time.Duration) time.Duration {
	if b > a {
		return b
	}
	return a
}

// openAITrackedRequest serializes budget check -> provider call -> usage log.
// That prevents two concurrent admin/worker requests from both observing the
// same remaining budget before either successful call is recorded.
func (w *Worker) openAITrackedRequest(ctx context.Context, cfg settings.Runtime, taskID, kind, prompt string, refs []referenceImage) (imageResult, error) {
	openAIBudgetMu.Lock()
	defer openAIBudgetMu.Unlock()
	if err := w.checkOpenAIBudget(ctx, cfg); err != nil {
		return imageResult{}, err
	}
	res, err := w.openAIRequest(ctx, cfg, prompt, refs)
	if err != nil {
		return imageResult{}, err
	}
	if err := w.recordOpenAIUsage(ctx, taskID, kind, res); err != nil {
		if res.Meta == nil {
			res.Meta = map[string]any{}
		}
		// Never retry a successfully generated image merely because telemetry
		// failed; that could double-spend. Surface the logging failure in metadata.
		res.Meta["usage_log_error"] = err.Error()
	}
	return res, nil
}