iop/apps/edge/internal/openai/usage_metrics.go

package openai

import (
	"context"
	"errors"
	"fmt"
	"strings"
	"sync"
	"sync/atomic"

	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promauto"

	edgeservice "iop/apps/edge/internal/service"
	"iop/packages/go/config"
)

// Usage metering is scoped to the OpenAI-compatible input surface only (SDD
// S04). CLI, A2A, and local-only adapter usage are intentionally excluded from
// these counters; this file is never imported by those surfaces.

// Terminal request status label values.
const (
	usageStatusSuccess = "success"
	usageStatusError   = "error"
	usageStatusCancel  = "cancel"
)

// usage_source label values on iop_openai_requests_total. provider_reported
// means at least one token type was reported for the request; unavailable means
// no usage was observed (status-only request count).
const (
	usageSourceProviderReported = "provider_reported"
	usageSourceUnavailable      = "unavailable"
)

// Endpoint label values distinguish the OpenAI-compatible route.
const (
	usageEndpointChatCompletions = "chat.completions"
	usageEndpointResponses       = "responses"
)

// Token type label values on iop_openai_usage_tokens_total.
const (
	tokenTypeInput       = "input"
	tokenTypeOutput      = "output"
	tokenTypeReasoning   = "reasoning"
	tokenTypeCachedInput = "cached_input"
)

// Label allowlists. These deliberately exclude request_id, session_id, raw
// bearer/provider tokens, and raw prompt/response text so the counters stay
// low-cardinality and never leak secrets (SDD S08). The label sets are exported
// as package vars so tests can assert the allowlist directly.
var (
	usageRequestLabelNames = []string{
		"edge_id", "principal_ref", "principal_alias", "token_ref",
		"route_model", "endpoint", "response_mode", "status", "usage_source",
	}
	usageTokenLabelNames = []string{
		"edge_id", "principal_ref", "principal_alias", "token_ref",
		"route_model", "usage_attribution", "provider_id", "served_model",
		"endpoint", "response_mode", "token_type",
	}
	usageReasoningLabelNames = []string{
		"edge_id", "principal_ref", "principal_alias", "token_ref",
		"route_model", "usage_attribution", "provider_id", "served_model",
		"endpoint", "response_mode",
	}
	usageReasoningEstimateLabelNames = []string{
		"edge_id", "principal_ref", "principal_alias", "token_ref",
		"route_model", "usage_attribution", "provider_id", "served_model",
		"endpoint", "response_mode", "estimation_method",
	}
)

var (
	openAIRequestsTotal = promauto.NewCounterVec(prometheus.CounterOpts{
		Name: "iop_openai_requests_total",
		Help: "OpenAI-compatible requests processed by the Edge input surface, by terminal status and usage source.",
	}, usageRequestLabelNames)

	openAIUsageTokensTotal = promauto.NewCounterVec(prometheus.CounterOpts{
		Name: "iop_openai_usage_tokens_total",
		Help: "Provider-reported OpenAI-compatible token usage by token type. Reasoning tokens are only counted when the provider reports them.",
	}, usageTokenLabelNames)

	openAIReasoningObservedTotal = promauto.NewCounterVec(prometheus.CounterOpts{
		Name: "iop_openai_reasoning_observed_total",
		Help: "Requests where reasoning text was observed but the provider did not report reasoning token usage.",
	}, usageReasoningLabelNames)

	openAIReasoningCharsTotal = promauto.NewCounterVec(prometheus.CounterOpts{
		Name: "iop_openai_reasoning_chars_total",
		Help: "Observed reasoning text characters for requests without provider-reported reasoning tokens. The token estimate is emitted separately via iop_openai_reasoning_estimated_tokens_total.",
	}, usageReasoningLabelNames)

	openAIReasoningEstimatedTokensTotal = promauto.NewCounterVec(prometheus.CounterOpts{
		Name: "iop_openai_reasoning_estimated_tokens_total",
		Help: "Estimated reasoning tokens for requests with observed reasoning text but without provider-reported reasoning token usage.",
	}, usageReasoningEstimateLabelNames)
)

// usageObservation is the metric-facing view of one attempt's token usage.
// providerReported distinguishes an observed zero-valued usage payload from an
// attempt where no provider usage fields were observed at all.
type usageObservation struct {
	inputTokens       int
	outputTokens      int
	reasoningTokens   int
	cachedInputTokens int
	providerReported  bool
	// reasoningChars counts observed reasoning text characters. It feeds the
	// reasoning-observed auxiliary metrics and the separate estimated-token
	// counter when the provider does not report reasoning token usage.
	reasoningChars int
}

func usageObservationFromOpenAIUsage(usage *openAIUsage, reasoningChars int) usageObservation {
	obs := usageObservation{reasoningChars: reasoningChars}
	if usage != nil {
		obs.providerReported = true
		obs.inputTokens = usage.PromptTokens
		obs.outputTokens = usage.CompletionTokens
		obs.reasoningTokens = usage.ReasoningTokens
		obs.cachedInputTokens = usage.CachedInputTokens
	}
	return obs
}

// openAIAttemptUsage is the request-local observation owned by one actual
// provider attempt. Event sources update it while the attempt controller alone
// claims and records the final immutable snapshot.
type openAIAttemptUsage struct {
	mu  sync.Mutex
	obs usageObservation
}

func (u *openAIAttemptUsage) observe(obs usageObservation) {
	if u == nil {
		return
	}
	u.mu.Lock()
	u.obs.inputTokens = obs.inputTokens
	u.obs.outputTokens = obs.outputTokens
	u.obs.reasoningTokens = obs.reasoningTokens
	u.obs.cachedInputTokens = obs.cachedInputTokens
	u.obs.providerReported = u.obs.providerReported || obs.providerReported
	if obs.reasoningChars > u.obs.reasoningChars {
		u.obs.reasoningChars = obs.reasoningChars
	}
	u.mu.Unlock()
}

func (u *openAIAttemptUsage) addReasoningChars(chars int) {
	if u == nil || chars <= 0 {
		return
	}
	u.mu.Lock()
	u.obs.reasoningChars += chars
	u.mu.Unlock()
}

func (u *openAIAttemptUsage) snapshot() usageObservation {
	if u == nil {
		return usageObservation{}
	}
	u.mu.Lock()
	defer u.mu.Unlock()
	return u.obs
}

// usageRequestLabels carries the stable caller and route values shared by the
// one request terminal and every provider attempt made on its behalf.
type usageRequestLabels struct {
	edgeID         string
	principalRef   string
	principalAlias string
	tokenRef       string
	routeModel     string
	endpoint       string
}

// usageDispatchBinding is the immutable metric binding for one actual provider
// attempt. attemptID and nodeID remain request-local deduplication/evidence
// fields and are deliberately absent from every public metric label vector.
type usageDispatchBinding struct {
	attemptID        string
	usageAttribution string
	providerID       string
	servedModel      string
	nodeID           string
	responseMode     string
}

// usageAttemptLabels is the public low-cardinality label subset for one actual
// provider attempt. route_model remains trace context while provider_id and
// served_model are the canonical usage attribution dimensions.
type usageAttemptLabels struct {
	usageRequestLabels
	usageAttribution string
	providerID       string
	servedModel      string
	responseMode     string
}

var openAIUsageAttemptSequence atomic.Uint64

func nextOpenAIUsageAttemptID() string {
	return fmt.Sprintf("usage-attempt-%d", openAIUsageAttemptSequence.Add(1))
}

// newUsageDispatchBinding derives the strict actual-provider binding supplied
// by service dispatch. Adapter and node values are never substituted for a
// missing provider id.
func newUsageDispatchBinding(dispatch edgeservice.RunDispatch, responseMode string) usageDispatchBinding {
	attribution := strings.TrimSpace(dispatch.UsageAttribution)
	if attribution == "" {
		attribution = config.UsageAttributionProvider
	}
	return usageDispatchBinding{
		attemptID:        nextOpenAIUsageAttemptID(),
		usageAttribution: attribution,
		providerID:       strings.TrimSpace(dispatch.ProviderID),
		servedModel:      strings.TrimSpace(dispatch.Target),
		nodeID:           strings.TrimSpace(dispatch.NodeID),
		responseMode:     strings.TrimSpace(responseMode),
	}
}

// openAIUsageRecorder owns request-terminal exactly-once state and
// per-attempt usage deduplication for one OpenAI-compatible request.
type openAIUsageRecorder struct {
	mu sync.Mutex

	request         usageRequestLabels
	attempts        map[string]usageDispatchBinding
	providerUsage   bool
	requestFinished bool
}

// newOpenAIUsageRecorder resolves authenticated identity once at ingress.
// Caller-supplied metadata.user and identity-like metadata never participate.
func (s *Server) newOpenAIUsageRecorder(ctx context.Context, routeModel, endpoint string) *openAIUsageRecorder {
	l := usageRequestLabels{
		edgeID:     s.edgeIDValue(),
		routeModel: strings.TrimSpace(routeModel),
		endpoint:   endpoint,
	}
	if p, ok := principalFromContext(ctx); ok {
		l.principalRef = p.PrincipalRef
		l.principalAlias = p.PrincipalAlias
		l.tokenRef = p.TokenRef
	}
	return &openAIUsageRecorder{request: l, attempts: make(map[string]usageDispatchBinding)}
}

func (l usageAttemptLabels) reasoningValues() []string {
	return []string{
		l.edgeID, l.principalRef, l.principalAlias, l.tokenRef,
		l.routeModel, l.usageAttribution, l.providerID, l.servedModel,
		l.endpoint, l.responseMode,
	}
}

// estimatedReasoningTokens converts observed reasoning text characters to a
// token estimate using the chars/4 rule of thumb. This is an operator-grade
// observation estimate only, never billing-grade. Caller reports a non-zero
// provider `reasoningTokens` value to skip this path entirely.
func estimatedReasoningTokens(chars int) int {
	return (chars + 3) / 4
}

// RecordAttempt emits one canonical provider series for one actual dispatch.
// Repeated finalization with the same internal attempt id is ignored. A missing
// strict provider/served-model binding emits no provider usage rather than
// falling back to adapter or node identity.
func (r *openAIUsageRecorder) RecordAttempt(binding usageDispatchBinding, obs usageObservation) {
	if r == nil {
		return
	}
	if strings.TrimSpace(binding.attemptID) == "" {
		binding.attemptID = nextOpenAIUsageAttemptID()
	}
	binding.usageAttribution = strings.TrimSpace(binding.usageAttribution)
	if binding.usageAttribution == "" {
		binding.usageAttribution = config.UsageAttributionProvider
	}
	binding.providerID = strings.TrimSpace(binding.providerID)
	binding.servedModel = strings.TrimSpace(binding.servedModel)
	binding.nodeID = strings.TrimSpace(binding.nodeID)
	binding.responseMode = strings.TrimSpace(binding.responseMode)

	r.mu.Lock()
	if _, duplicate := r.attempts[binding.attemptID]; duplicate {
		r.mu.Unlock()
		return
	}
	r.attempts[binding.attemptID] = binding
	if obs.providerReported && binding.providerID != "" && binding.servedModel != "" {
		r.providerUsage = true
	}
	request := r.request
	r.mu.Unlock()

	if binding.providerID == "" || binding.servedModel == "" {
		return
	}
	l := usageAttemptLabels{
		usageRequestLabels: request,
		usageAttribution:   binding.usageAttribution,
		providerID:         binding.providerID,
		servedModel:        binding.servedModel,
		responseMode:       binding.responseMode,
	}
	addTokenCount(l, tokenTypeInput, obs.inputTokens)
	addTokenCount(l, tokenTypeOutput, obs.outputTokens)
	addTokenCount(l, tokenTypeReasoning, obs.reasoningTokens)
	addTokenCount(l, tokenTypeCachedInput, obs.cachedInputTokens)

	if obs.reasoningChars > 0 && obs.reasoningTokens == 0 {
		openAIReasoningObservedTotal.WithLabelValues(l.reasoningValues()...).Inc()
		openAIReasoningCharsTotal.WithLabelValues(l.reasoningValues()...).Add(float64(obs.reasoningChars))
		estimate := estimatedReasoningTokens(obs.reasoningChars)
		openAIReasoningEstimatedTokensTotal.WithLabelValues(
			l.edgeID, l.principalRef, l.principalAlias, l.tokenRef,
			l.routeModel, l.usageAttribution, l.providerID, l.servedModel,
			l.endpoint, l.responseMode, "chars_div_4",
		).Add(float64(estimate))
	}
}

// FinishRequest records exactly one HTTP terminal using the final committed
// response mode. Reasoning-character observations alone do not classify the
// request as provider-reported usage.
func (r *openAIUsageRecorder) FinishRequest(status, responseMode string) {
	if r == nil {
		return
	}
	r.mu.Lock()
	if r.requestFinished {
		r.mu.Unlock()
		return
	}
	r.requestFinished = true
	request := r.request
	providerUsage := r.providerUsage
	r.mu.Unlock()

	source := usageSourceUnavailable
	if providerUsage {
		source = usageSourceProviderReported
	}
	openAIRequestsTotal.WithLabelValues(
		request.edgeID, request.principalRef, request.principalAlias, request.tokenRef,
		request.routeModel, request.endpoint, responseMode, status, source,
	).Inc()
}

func addTokenCount(l usageAttemptLabels, tokenType string, count int) {
	if count <= 0 {
		return
	}
	openAIUsageTokensTotal.WithLabelValues(
		l.edgeID, l.principalRef, l.principalAlias, l.tokenRef,
		l.routeModel, l.usageAttribution, l.providerID, l.servedModel,
		l.endpoint, l.responseMode, tokenType,
	).Add(float64(count))
}

// usageStatusForError classifies a post-dispatch run error into a terminal
// status label. Caller give-up (cancellation/timeout) is a cancel; everything
// else is an error.
func usageStatusForError(err error) string {
	if err == nil {
		return usageStatusSuccess
	}
	if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) || errors.Is(err, errRunTimedOut) {
		return usageStatusCancel
	}
	return usageStatusError
}

const (
	// responseModePassthrough and responseModeNormalized are internal execution
	// labels for the response_mode usage metric. They are derived from the
	// handler execution path, never from caller metadata: provider tunnel routes
	// report passthrough, normalized RunEvent routes report normalized. Callers
	// cannot select a response mode through OpenAI metadata.
	responseModePassthrough = "passthrough"
	responseModeNormalized  = "normalized"
)