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

package openai

import (
	"context"
	"errors"

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

// 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",
		"model_group", "endpoint", "response_mode", "status", "usage_source",
	}
	usageTokenLabelNames = []string{
		"edge_id", "principal_ref", "principal_alias", "token_ref",
		"model_group", "endpoint", "response_mode", "token_type",
	}
	usageReasoningLabelNames = []string{
		"edge_id", "principal_ref", "principal_alias", "token_ref",
		"model_group", "endpoint", "response_mode",
	}
)

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. Never converted to a token estimate.",
	}, usageReasoningLabelNames)
)

// usageObservation is the metric-facing view of a request's token usage. Counts
// are provider-reported; a zero count means the provider did not report that
// token type and the emitter never estimates it.
type usageObservation struct {
	inputTokens       int
	outputTokens      int
	reasoningTokens   int
	cachedInputTokens int
	// reasoningChars counts observed reasoning text characters. It feeds the
	// reasoning-observed auxiliary metrics only and is never turned into a token
	// estimate (SDD S07).
	reasoningChars int
}

func (o usageObservation) hasTokens() bool {
	return o.inputTokens > 0 || o.outputTokens > 0 || o.reasoningTokens > 0 || o.cachedInputTokens > 0
}

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

// usageLabels carries the resolved low-cardinality label values for one request.
type usageLabels struct {
	edgeID         string
	principalRef   string
	principalAlias string
	tokenRef       string
	modelGroup     string
	endpoint       string
	responseMode   string
}

// usageLabelsFor builds the metric label values for a request from its
// authenticated principal context. Identity is taken only from the resolved
// principal (never from caller-supplied metadata.user).
func (s *Server) usageLabelsFor(ctx context.Context, modelGroup, endpoint, responseMode string) usageLabels {
	l := usageLabels{
		edgeID:       s.edgeIDValue(),
		modelGroup:   modelGroup,
		endpoint:     endpoint,
		responseMode: responseMode,
	}
	if p, ok := principalFromContext(ctx); ok {
		l.principalRef = p.PrincipalRef
		l.principalAlias = p.PrincipalAlias
		l.tokenRef = p.TokenRef
	}
	return l
}

func (l usageLabels) reasoningValues() []string {
	return []string{l.edgeID, l.principalRef, l.principalAlias, l.tokenRef, l.modelGroup, l.endpoint, l.responseMode}
}

// emitUsageMetrics increments the request counter for the terminal status and,
// when usage was observed, the per-token-type counters. Reasoning text observed
// without a provider-reported reasoning token count only increments the
// auxiliary reasoning counters; it is never added to the token counter.
func emitUsageMetrics(l usageLabels, status string, obs usageObservation) {
	source := usageSourceUnavailable
	if obs.hasTokens() {
		source = usageSourceProviderReported
	}
	openAIRequestsTotal.WithLabelValues(
		l.edgeID, l.principalRef, l.principalAlias, l.tokenRef,
		l.modelGroup, l.endpoint, l.responseMode, status, source,
	).Inc()

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

func addTokenCount(l usageLabels, tokenType string, count int) {
	if count <= 0 {
		return
	}
	openAIUsageTokensTotal.WithLabelValues(
		l.edgeID, l.principalRef, l.principalAlias, l.tokenRef,
		l.modelGroup, 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
}