package helps import (	"bytes"	"context"	"fmt"	"strings"	"sync"	"time" 	"github.com/gin-gonic/gin"	cliproxyauth "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/auth"	"github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/usage"	"github.com/tidwall/gjson"	"github.com/tidwall/sjson") type UsageReporter struct {	provider    string	model       string	authID      string	authIndex   string	authType    string	apiKey      string	source      string	requestedAt time.Time	once        sync.Once} func NewUsageReporter(ctx context.Context, provider, model string, auth *cliproxyauth.Auth) *UsageReporter {	apiKey := APIKeyFromContext(ctx)	reporter := &UsageReporter{		provider:    provider,		model:       model,		requestedAt: time.Now(),		apiKey:      apiKey,		source:      resolveUsageSource(auth, apiKey),		authType:    resolveUsageAuthType(auth),	}	if auth != nil {		reporter.authID = auth.ID		reporter.authIndex = auth.EnsureIndex()	}	return reporter} func (r *UsageReporter) Publish(ctx context.Context, detail usage.Detail) {	r.publishWithOutcome(ctx, detail, false)} func (r *UsageReporter) PublishAdditionalModel(ctx context.Context, model string, detail usage.Detail) {	record, ok := r.buildAdditionalModelRecord(model, detail)	if !ok {		return	}	usage.PublishRecord(ctx, record)} func (r *UsageReporter) buildAdditionalModelRecord(model string, detail usage.Detail) (usage.Record, bool) {	if r == nil {		return usage.Record{}, false	}	model = strings.TrimSpace(model)	if model == "" {		return usage.Record{}, false	}	detail = normalizeUsageDetailTotal(detail)	if !hasNonZeroTokenUsage(detail) {		return usage.Record{}, false	}	return r.buildRecordForModel(model, detail, false), true} func (r *UsageReporter) PublishFailure(ctx context.Context) {	r.publishWithOutcome(ctx, usage.Detail{}, true)} func (r *UsageReporter) TrackFailure(ctx context.Context, errPtr *error) {	if r == nil || errPtr == nil {		return	}	if *errPtr != nil {		r.PublishFailure(ctx)	}} func (r *UsageReporter) publishWithOutcome(ctx context.Context, detail usage.Detail, failed bool) {	if r == nil {		return	}	detail = normalizeUsageDetailTotal(detail)	r.once.Do(func() {		usage.PublishRecord(ctx, r.buildRecord(detail, failed))	})} func normalizeUsageDetailTotal(detail usage.Detail) usage.Detail {	if detail.TotalTokens == 0 {		total := detail.InputTokens + detail.OutputTokens + detail.ReasoningTokens		if total > 0 {			detail.TotalTokens = total		}	}	return detail} func hasNonZeroTokenUsage(detail usage.Detail) bool {	return detail.InputTokens != 0 ||		detail.OutputTokens != 0 ||		detail.ReasoningTokens != 0 ||		detail.CachedTokens != 0 ||		detail.TotalTokens != 0} // ensurePublished guarantees that a usage record is emitted exactly once.// It is safe to call multiple times; only the first call wins due to once.Do.// This is used to ensure request counting even when upstream responses do not// include any usage fields (tokens), especially for streaming paths.func (r *UsageReporter) EnsurePublished(ctx context.Context) {	if r == nil {		return	}	r.once.Do(func() {		usage.PublishRecord(ctx, r.buildRecord(usage.Detail{}, false))	})} func (r *UsageReporter) buildRecord(detail usage.Detail, failed bool) usage.Record {	if r == nil {		return usage.Record{Detail: detail, Failed: failed}	}	return r.buildRecordForModel(r.model, detail, failed)} func (r *UsageReporter) buildRecordForModel(model string, detail usage.Detail, failed bool) usage.Record {	if r == nil {		return usage.Record{Model: model, Detail: detail, Failed: failed}	}	return usage.Record{		Provider:    r.provider,		Model:       model,		Source:      r.source,		APIKey:      r.apiKey,		AuthID:      r.authID,		AuthIndex:   r.authIndex,		AuthType:    r.authType,		RequestedAt: r.requestedAt,		Latency:     r.latency(),		Failed:      failed,		Detail:      detail,	}} func (r *UsageReporter) latency() time.Duration {	if r == nil || r.requestedAt.IsZero() {		return 0	}	latency := time.Since(r.requestedAt)	if latency < 0 {		return 0	}	return latency} func APIKeyFromContext(ctx context.Context) string {	if ctx == nil {		return ""	}	ginCtx, ok := ctx.Value("gin").(*gin.Context)	if !ok || ginCtx == nil {		return ""	}	if v, exists := ginCtx.Get("apiKey"); exists {		switch value := v.(type) {		case string:			return value		case fmt.Stringer:			return value.String()		default:			return fmt.Sprintf("%v", value)		}	}	return ""} func resolveUsageSource(auth *cliproxyauth.Auth, ctxAPIKey string) string {	if auth != nil {		provider := strings.TrimSpace(auth.Provider)		if strings.EqualFold(provider, "gemini-cli") {			if id := strings.TrimSpace(auth.ID); id != "" {				return id			}		}		if strings.EqualFold(provider, "vertex") {			if auth.Metadata != nil {				if projectID, ok := auth.Metadata["project_id"].(string); ok {					if trimmed := strings.TrimSpace(projectID); trimmed != "" {						return trimmed					}				}				if project, ok := auth.Metadata["project"].(string); ok {					if trimmed := strings.TrimSpace(project); trimmed != "" {						return trimmed					}				}			}		}		if _, value := auth.AccountInfo(); value != "" {			return strings.TrimSpace(value)		}		if auth.Metadata != nil {			if email, ok := auth.Metadata["email"].(string); ok {				if trimmed := strings.TrimSpace(email); trimmed != "" {					return trimmed				}			}		}		if auth.Attributes != nil {			if key := strings.TrimSpace(auth.Attributes["api_key"]); key != "" {				return key			}		}	}	if trimmed := strings.TrimSpace(ctxAPIKey); trimmed != "" {		return trimmed	}	return ""} func resolveUsageAuthType(auth *cliproxyauth.Auth) string {	if auth == nil {		return ""	}	kind, _ := auth.AccountInfo()	kind = strings.TrimSpace(kind)	if kind == "api_key" {		return "apikey"	}	return kind} func ParseCodexUsage(data []byte) (usage.Detail, bool) {	usageNode := gjson.ParseBytes(data).Get("response.usage")	if !usageNode.Exists() {		return usage.Detail{}, false	}	return parseOpenAIStyleUsageNode(usageNode), true} func ParseCodexImageToolUsage(data []byte) (usage.Detail, bool) {	usageNode := gjson.ParseBytes(data).Get("response.tool_usage.image_gen")	if !usageNode.Exists() || !usageNode.IsObject() {		return usage.Detail{}, false	}	return parseOpenAIStyleUsageNode(usageNode), true} func ParseOpenAIUsage(data []byte) usage.Detail {	usageNode := gjson.ParseBytes(data).Get("usage")	if !usageNode.Exists() {		return usage.Detail{}	}	return parseOpenAIStyleUsageNode(usageNode)} func parseOpenAIStyleUsageNode(usageNode gjson.Result) usage.Detail {	inputNode := usageNode.Get("prompt_tokens")	if !inputNode.Exists() {		inputNode = usageNode.Get("input_tokens")	}	outputNode := usageNode.Get("completion_tokens")	if !outputNode.Exists() {		outputNode = usageNode.Get("output_tokens")	}	detail := usage.Detail{		InputTokens:  inputNode.Int(),		OutputTokens: outputNode.Int(),		TotalTokens:  usageNode.Get("total_tokens").Int(),	}	cached := usageNode.Get("prompt_tokens_details.cached_tokens")	if !cached.Exists() {		cached = usageNode.Get("input_tokens_details.cached_tokens")	}	if cached.Exists() {		detail.CachedTokens = cached.Int()	}	reasoning := usageNode.Get("completion_tokens_details.reasoning_tokens")	if !reasoning.Exists() {		reasoning = usageNode.Get("output_tokens_details.reasoning_tokens")	}	if reasoning.Exists() {		detail.ReasoningTokens = reasoning.Int()	}	return detail} func ParseOpenAIStreamUsage(line []byte) (usage.Detail, bool) {	payload := jsonPayload(line)	if len(payload) == 0 || !gjson.ValidBytes(payload) {		return usage.Detail{}, false	}	usageNode := gjson.GetBytes(payload, "usage")	if !usageNode.Exists() {		return usage.Detail{}, false	}	detail := usage.Detail{		InputTokens:  usageNode.Get("prompt_tokens").Int(),		OutputTokens: usageNode.Get("completion_tokens").Int(),		TotalTokens:  usageNode.Get("total_tokens").Int(),	}	if cached := usageNode.Get("prompt_tokens_details.cached_tokens"); cached.Exists() {		detail.CachedTokens = cached.Int()	}	if reasoning := usageNode.Get("completion_tokens_details.reasoning_tokens"); reasoning.Exists() {		detail.ReasoningTokens = reasoning.Int()	}	return detail, true} func ParseClaudeUsage(data []byte) usage.Detail {	usageNode := gjson.ParseBytes(data).Get("usage")	if !usageNode.Exists() {		return usage.Detail{}	}	detail := usage.Detail{		InputTokens:  usageNode.Get("input_tokens").Int(),		OutputTokens: usageNode.Get("output_tokens").Int(),		CachedTokens: usageNode.Get("cache_read_input_tokens").Int(),	}	if detail.CachedTokens == 0 {		// fall back to creation tokens when read tokens are absent		detail.CachedTokens = usageNode.Get("cache_creation_input_tokens").Int()	}	detail.TotalTokens = detail.InputTokens + detail.OutputTokens	return detail} func ParseClaudeStreamUsage(line []byte) (usage.Detail, bool) {	payload := jsonPayload(line)	if len(payload) == 0 || !gjson.ValidBytes(payload) {		return usage.Detail{}, false	}	usageNode := gjson.GetBytes(payload, "usage")	if !usageNode.Exists() {		return usage.Detail{}, false	}	detail := usage.Detail{		InputTokens:  usageNode.Get("input_tokens").Int(),		OutputTokens: usageNode.Get("output_tokens").Int(),		CachedTokens: usageNode.Get("cache_read_input_tokens").Int(),	}	if detail.CachedTokens == 0 {		detail.CachedTokens = usageNode.Get("cache_creation_input_tokens").Int()	}	detail.TotalTokens = detail.InputTokens + detail.OutputTokens	return detail, true} func parseGeminiFamilyUsageDetail(node gjson.Result) usage.Detail {	detail := usage.Detail{		InputTokens:     node.Get("promptTokenCount").Int(),		OutputTokens:    node.Get("candidatesTokenCount").Int(),		ReasoningTokens: node.Get("thoughtsTokenCount").Int(),		TotalTokens:     node.Get("totalTokenCount").Int(),		CachedTokens:    node.Get("cachedContentTokenCount").Int(),	}	if detail.TotalTokens == 0 {		detail.TotalTokens = detail.InputTokens + detail.OutputTokens + detail.ReasoningTokens	}	return detail} func hasGeminiFamilyUsageTokenFields(node gjson.Result) bool {	return node.Get("promptTokenCount").Exists() ||		node.Get("candidatesTokenCount").Exists() ||		node.Get("thoughtsTokenCount").Exists() ||		node.Get("totalTokenCount").Exists() ||		node.Get("cachedContentTokenCount").Exists()} func ParseGeminiCLIUsage(data []byte) usage.Detail {	usageNode := gjson.ParseBytes(data)	node := firstExistingUsageNode(usageNode,		"response.usageMetadata",		"response.usage_metadata",		"usageMetadata",		"usage_metadata",	)	if !node.Exists() {		return usage.Detail{}	}	return parseGeminiFamilyUsageDetail(node)} func ParseGeminiUsage(data []byte) usage.Detail {	usageNode := gjson.ParseBytes(data)	node := usageNode.Get("usageMetadata")	if !node.Exists() {		node = usageNode.Get("usage_metadata")	}	if !node.Exists() {		return usage.Detail{}	}	return parseGeminiFamilyUsageDetail(node)} func ParseGeminiStreamUsage(line []byte) (usage.Detail, bool) {	payload := jsonPayload(line)	if len(payload) == 0 || !gjson.ValidBytes(payload) {		return usage.Detail{}, false	}	node := gjson.GetBytes(payload, "usageMetadata")	if !node.Exists() {		node = gjson.GetBytes(payload, "usage_metadata")	}	if !node.Exists() {		return usage.Detail{}, false	}	return parseGeminiFamilyUsageDetail(node), true} func ParseGeminiCLIStreamUsage(line []byte) (usage.Detail, bool) {	payload := jsonPayload(line)	if len(payload) == 0 || !gjson.ValidBytes(payload) {		return usage.Detail{}, false	}	root := gjson.ParseBytes(payload)	node := firstExistingUsageNode(root,		"response.usageMetadata",		"response.usage_metadata",		"usageMetadata",		"usage_metadata",	)	if !node.Exists() {		return usage.Detail{}, false	}	if !hasGeminiFamilyUsageTokenFields(node) {		return usage.Detail{}, false	}	return parseGeminiFamilyUsageDetail(node), true} func firstExistingUsageNode(root gjson.Result, paths ...string) gjson.Result {	for _, path := range paths {		node := root.Get(path)		if node.Exists() {			return node		}	}	return gjson.Result{}} func ParseAntigravityUsage(data []byte) usage.Detail {	usageNode := gjson.ParseBytes(data)	node := usageNode.Get("response.usageMetadata")	if !node.Exists() {		node = usageNode.Get("usageMetadata")	}	if !node.Exists() {		node = usageNode.Get("usage_metadata")	}	if !node.Exists() {		return usage.Detail{}	}	return parseGeminiFamilyUsageDetail(node)} func ParseAntigravityStreamUsage(line []byte) (usage.Detail, bool) {	payload := jsonPayload(line)	if len(payload) == 0 || !gjson.ValidBytes(payload) {		return usage.Detail{}, false	}	node := gjson.GetBytes(payload, "response.usageMetadata")	if !node.Exists() {		node = gjson.GetBytes(payload, "usageMetadata")	}	if !node.Exists() {		node = gjson.GetBytes(payload, "usage_metadata")	}	if !node.Exists() {		return usage.Detail{}, false	}	return parseGeminiFamilyUsageDetail(node), true} var stopChunkWithoutUsage sync.Map func rememberStopWithoutUsage(traceID string) {	stopChunkWithoutUsage.Store(traceID, struct{}{})	time.AfterFunc(10*time.Minute, func() { stopChunkWithoutUsage.Delete(traceID) })} // FilterSSEUsageMetadata removes usageMetadata from SSE events that are not// terminal (finishReason != "stop"). Stop chunks are left untouched. This// function is shared between aistudio and antigravity executors.func FilterSSEUsageMetadata(payload []byte) []byte {	if len(payload) == 0 {		return payload	} 

Lines 1–500 of 659. Load more