package helps import (	"bytes"	"context"	"errors"	"fmt"	"io"	"net/http"	"reflect"	"strings"	"sync"	"time" 	"github.com/gin-gonic/gin"	"github.com/router-for-me/CLIProxyAPI/v7/internal/clienterror"	internallogging "github.com/router-for-me/CLIProxyAPI/v7/internal/logging"	"github.com/router-for-me/CLIProxyAPI/v7/internal/thinking"	cliproxyauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth"	cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor"	"github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/usage"	"github.com/tidwall/gjson"	"github.com/tidwall/sjson") type UsageReporter struct {	provider            string	executorType        string	model               string	alias               string	authID              string	authIndex           string	authMu              sync.RWMutex	accessTokenHash     string	authType            string	apiKey              string	source              string	reasoning           string	serviceTier         string	generate            bool	stream              bool	requestedAt         time.Time	ttftMu              sync.RWMutex	ttft                time.Duration	firstPacketDuration time.Duration	firstPacketSet      bool	ttftStart           time.Time	ttftSet             bool	once                sync.Once} type usageExecutor interface {	Identifier() string} func NewExecutorUsageReporter(ctx context.Context, executor usageExecutor, model string, auth *cliproxyauth.Auth) *UsageReporter {	provider := ""	if executor != nil {		provider = executor.Identifier()	}	reporter := NewUsageReporter(ctx, provider, model, auth)	reporter.executorType = ExecutorTypeName(executor)	return reporter} func NewUsageReporter(ctx context.Context, provider, model string, auth *cliproxyauth.Auth) *UsageReporter {	apiKey := APIKeyFromContext(ctx)	alias := usage.RequestedModelAliasFromContext(ctx)	if alias == "" {		alias = model	}	reporter := &UsageReporter{		provider:    provider,		model:       model,		alias:       strings.TrimSpace(alias),		requestedAt: time.Now(),		apiKey:      apiKey,		source:      resolveUsageSource(auth, apiKey),		authType:    resolveUsageAuthType(auth),		reasoning:   usage.ReasoningEffortFromContext(ctx),		serviceTier: usage.ServiceTierFromContext(ctx),		generate:    usage.GenerateFromContext(ctx),		stream:      usage.StreamFromContext(ctx),	}	if auth != nil {		reporter.authID = auth.ID		reporter.authIndex = auth.EnsureIndex()		reporter.accessTokenHash = authAccessTokenSHA256(auth)	}	return reporter} // SetStream records whether the request was executed in streaming mode.func (r *UsageReporter) SetStream(stream bool) {	if r == nil {		return	}	r.stream = stream} // UpdateAccessTokenFingerprint records the token version actually used upstream.func (r *UsageReporter) UpdateAccessTokenFingerprint(auth *cliproxyauth.Auth) {	if r == nil {		return	}	r.authMu.Lock()	r.accessTokenHash = authAccessTokenSHA256(auth)	r.authMu.Unlock()} func (r *UsageReporter) accessTokenFingerprint() string {	if r == nil {		return ""	}	r.authMu.RLock()	defer r.authMu.RUnlock()	return r.accessTokenHash} func ExecutorTypeName(executor any) string {	if executor == nil {		return ""	}	executorType := reflect.TypeOf(executor)	for executorType.Kind() == reflect.Pointer {		executorType = executorType.Elem()	}	return strings.TrimSpace(executorType.Name())} func (r *UsageReporter) Publish(ctx context.Context, detail usage.Detail) {	r.publishWithOutcome(ctx, detail, false, usage.Failure{})} func (r *UsageReporter) PublishAdditionalModel(ctx context.Context, model string, detail usage.Detail) {	record, ok := r.buildAdditionalModelRecord(model, detail)	if !ok {		return	}	r.publishRecord(ctx, record)} func (r *UsageReporter) SetTranslatedReasoningEffort(payload []byte, format string) {	if r == nil {		return	}	r.reasoning = thinking.ExtractTranslatedReasoningEffort(payload, format)} func (r *UsageReporter) TrackHTTPClient(client *http.Client) *http.Client {	return r.trackHTTPClient(client, false)} // TrackHTTPClientRoundTripOnly records the TTFT start time upon sending the request// and captures first-packet arrival fallback on initial body reads, while keeping// effective TTFT unset. This allows protocol-aware streaming executors (like Codex SSE)// to mark effective TTFT explicitly upon receiving substantive token events, while// preserving first-packet fallback metrics for non-2xx error bodies or keepalive streams.func (r *UsageReporter) TrackHTTPClientRoundTripOnly(client *http.Client) *http.Client {	return r.trackHTTPClient(client, true)} func (r *UsageReporter) trackHTTPClient(client *http.Client, packetOnly bool) *http.Client {	if r == nil || client == nil {		return client	}	tracked := *client	transport := tracked.Transport	if transport == nil {		transport = http.DefaultTransport	}	tracked.Transport = usageTTFTRoundTripper{		base:       transport,		reporter:   r,		packetOnly: packetOnly,	}	return &tracked} func (r *UsageReporter) ObserveResponse(resp *http.Response) {	if r == nil || resp == nil || resp.Body == nil {		return	}	r.StartResponseTTFT()	resp.Body = &usageTTFTReadCloser{		ReadCloser: resp.Body,		mark: func() {			r.MarkFirstResponseByte()		},	}} func (r *UsageReporter) ObserveResponsePacketOnly(resp *http.Response) {	if r == nil || resp == nil || resp.Body == nil {		return	}	r.StartResponseTTFT()	resp.Body = &usageTTFTReadCloser{		ReadCloser: resp.Body,		mark: func() {			r.RecordFirstPacket()		},	}} func (r *UsageReporter) StartResponseTTFT() {	if r == nil {		return	}	r.ttftMu.Lock()	if !r.ttftSet && r.ttftStart.IsZero() {		r.ttftStart = time.Now()	}	r.ttftMu.Unlock()} func (r *UsageReporter) IsTTFTSet() bool {	if r == nil {		return false	}	r.ttftMu.RLock()	defer r.ttftMu.RUnlock()	return r.ttftSet} func (r *UsageReporter) IsFirstPacketSet() bool {	if r == nil {		return false	}	r.ttftMu.RLock()	defer r.ttftMu.RUnlock()	return r.firstPacketSet} // RecordFirstPacket records the arrival time of the first packet/chunk from upstream as a fallback.func (r *UsageReporter) RecordFirstPacket() {	r.ObserveTokenEvent(false)} // ObserveTokenEvent records the first packet fallback time on the first frame,// and if isToken is true, records the effective TTFT. It uses a fast-path// read check to return immediately with zero write lock contention once TTFT is set// or when subsequent non-token metadata frames arrive after the first packet.func (r *UsageReporter) ObserveTokenEvent(isToken bool) {	if r == nil {		return	}	r.ttftMu.RLock()	if r.ttftSet {		r.ttftMu.RUnlock()		return	}	start := r.ttftStart	alreadyRecordedPacket := r.firstPacketSet	r.ttftMu.RUnlock() 	if start.IsZero() {		return	} 	if !isToken && alreadyRecordedPacket {		return	} 	r.ttftMu.Lock()	defer r.ttftMu.Unlock()	if r.ttftSet {		return	}	if !r.firstPacketSet {		r.firstPacketDuration = time.Since(start)		r.firstPacketSet = true	}	if isToken {		r.ttft = time.Since(start)		r.ttftSet = true		r.ttftStart = time.Time{}	}} func (r *UsageReporter) MarkFirstResponseByte() {	if r == nil {		return	}	r.ttftMu.Lock()	if r.ttftSet {		r.ttftMu.Unlock()		return	}	start := r.ttftStart	r.ttftStart = time.Time{}	r.ttftMu.Unlock()	if start.IsZero() {		return	}	r.setTTFT(time.Since(start))} 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, r.provider, r.executorType)	if !hasNonZeroTokenUsage(detail) {		return usage.Record{}, false	}	return r.buildRecordForModel(model, detail, false, usage.Failure{}), true} func (r *UsageReporter) PublishFailure(ctx context.Context, errs ...error) {	r.publishWithOutcome(ctx, usage.Detail{}, true, failFromErrors(errs...))} func (r *UsageReporter) PublishFailureWithDetail(ctx context.Context, detail usage.Detail, errs ...error) {	r.publishWithOutcome(ctx, detail, true, failFromErrors(errs...))} func (r *UsageReporter) TrackFailure(ctx context.Context, errPtr *error) {	if r == nil || errPtr == nil {		return	}	if *errPtr != nil {		r.PublishFailure(ctx, *errPtr)	}} func (r *UsageReporter) publishWithOutcome(ctx context.Context, detail usage.Detail, failed bool, fail usage.Failure) {	if r == nil {		return	}	detail = normalizeUsageDetailTotal(detail, r.provider, r.executorType)	r.once.Do(func() {		r.publishRecord(ctx, r.buildRecord(detail, failed, fail))	})} func normalizeUsageDetailTotal(detail usage.Detail, provider, executorType string) usage.Detail {	return usage.EnsureTokenBreakdownForProvider(detail, provider, executorType)} func hasNonZeroTokenUsage(detail usage.Detail) bool {	return detail.InputTokens != 0 ||		detail.OutputTokens != 0 ||		detail.ReasoningTokens != 0 ||		detail.CachedTokens != 0 ||		detail.CacheReadTokens != 0 ||		detail.CacheCreationTokens != 0 ||		detail.TotalTokens != 0 ||		detail.TokenBreakdown.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() {		r.publishRecord(ctx, r.buildRecord(usage.Detail{}, false, usage.Failure{}))	})} func (r *UsageReporter) publishRecord(ctx context.Context, record usage.Record) {	record.ResponseHeaders = internallogging.GetResponseHeaders(ctx)	usage.PublishRecord(ctx, record)} func (r *UsageReporter) buildRecord(detail usage.Detail, failed bool, failures ...usage.Failure) usage.Record {	var fail usage.Failure	if len(failures) > 0 {		fail = failures[0]	}	if r == nil {		return usage.Record{Detail: detail, Failed: failed, Fail: fail, Generate: usage.GenerateFlag(true)}	}	return r.buildRecordForModel(r.model, detail, failed, fail)} func (r *UsageReporter) buildRecordForModel(model string, detail usage.Detail, failed bool, fail usage.Failure) usage.Record {	if r == nil {		return usage.Record{Model: model, Detail: detail, Failed: failed, Fail: fail, Generate: usage.GenerateFlag(true)}	}	return usage.Record{		Provider:            r.provider,		ExecutorType:        r.executorType,		Model:               model,		Alias:               r.alias,		Source:              r.source,		APIKey:              r.apiKey,		AuthID:              r.authID,		AuthIndex:           r.authIndex,		AccessTokenSHA256:   r.accessTokenFingerprint(),		AuthType:            r.authType,		ReasoningEffort:     r.reasoning,		ServiceTier:         r.serviceTier,		ResponseServiceTier: strings.TrimSpace(detail.ResponseServiceTier),		Generate:            usage.GenerateFlag(r.generate),		Stream:              r.stream,		RequestedAt:         r.requestedAt,		Latency:             r.latency(),		TTFT:                r.ttftDuration(),		Failed:              failed,		Fail:                fail,		Detail:              detail,	}} func failFromErrors(errs ...error) usage.Failure {	for _, err := range errs {		if err == nil {			continue		}		body := strings.TrimSpace(err.Error())		type responseBodyProvider interface {			ResponseBody() []byte		}		var responseErr responseBodyProvider		if errors.As(err, &responseErr) && responseErr != nil {			if responseBody := responseErr.ResponseBody(); len(responseBody) > 0 {				body = string(responseBody)			}		}		return usage.Failure{			Body:       body,			StatusCode: clienterror.HTTPStatusFromError(err),		}	}	return usage.Failure{}} 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 (r *UsageReporter) setTTFT(ttft time.Duration) {	if r == nil {		return	}	if ttft < 0 {		ttft = 0	}	r.ttftMu.Lock()	if r.ttftSet {		r.ttftMu.Unlock()		return	}	r.ttft = ttft	r.ttftSet = true	r.ttftStart = time.Time{}	r.ttftMu.Unlock()} func (r *UsageReporter) ttftDuration() time.Duration {	if r == nil {		return 0	}	r.ttftMu.RLock()	defer r.ttftMu.RUnlock()	if r.ttftSet {		return r.ttft	}	if r.firstPacketSet {		return r.firstPacketDuration	}	return 0} type usageTTFTRoundTripper struct {	base       http.RoundTripper	reporter   *UsageReporter	packetOnly bool} func (t usageTTFTRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) {	cliproxyexecutor.MarkUpstreamAttempt(req.Context())	t.reporter.StartResponseTTFT()	resp, errRoundTrip := t.base.RoundTrip(req)	if errRoundTrip != nil {		return resp, errRoundTrip	}	if t.packetOnly {		t.reporter.ObserveResponsePacketOnly(resp)	} else {		t.reporter.ObserveResponse(resp)	}	return resp, nil}

Lines 1–500 of 1298. Load more