new-api-custom/service/relay_attempt.go

package service

import (
	"fmt"
	"net/http"
	"strings"
	"time"

	"github.com/QuantumNous/new-api/common"
	"github.com/QuantumNous/new-api/constant"
	"github.com/QuantumNous/new-api/model"
	relaycommon "github.com/QuantumNous/new-api/relay/common"
	"github.com/QuantumNous/new-api/relaykit/dto"
	"github.com/QuantumNous/new-api/relaykit/types"
	"github.com/gin-gonic/gin"
)

func BeginRelayAttempt(c *gin.Context, info *relaycommon.RelayInfo) *types.NewAPIError {
	if !info.DurableAttemptBilling {
		return nil
	}
	reserved := info.PriceData.QuotaToPreConsume
	if snap := info.TieredBillingSnapshot; snap != nil {
		reserved = snap.EstimatedQuotaAfterGroup
	}
	price, err := common.Marshal(struct {
		Price      any
		Expression any
	}{info.PriceData, info.TieredBillingSnapshot})
	if err != nil {
		return types.NewError(err, types.ErrorCodeModelPriceError)
	}
	a := &model.RelayAttempt{AttemptKey: fmt.Sprintf("%s:%d", info.RequestId, info.UpstreamRequests+1), RequestID: info.RequestId, Number: info.UpstreamRequests + 1, UserID: info.UserId, TokenID: info.TokenId, ChannelID: c.GetInt("channel_id"), GroupName: info.UsingGroup, RequestedModel: info.OriginModelName, Reserved: reserved, EstimateQuota: reserved, PriceJSON: string(price), Deadline: time.Now().Add(10 * time.Minute).Unix(), GuardResult: "unchecked"}
	if err := model.ReserveRelayAttempt(a, info.UserSetting.BillingPreference); err != nil {
		return types.NewErrorWithStatusCode(err, types.ErrorCodeInsufficientUserQuota, http.StatusForbidden, types.ErrOptionWithSkipRetry())
	}
	info.DurableAttemptID = a.ID
	info.AttemptDetached = false
	info.Billing = nil
	info.BillingSource = a.Funding
	info.SubscriptionId = a.SubscriptionID
	info.FinalPreConsumedQuota = a.Reserved
	c.Set("relay_attempt_number", a.Number)
	c.Set("attempt_guard_result", "")
	c.Set("attempt_returned_model", "")
	c.Set("attempt_usage_source", "")
	requestContext := c.Request.Context()
	go func(id int64) {
		ticker := time.NewTicker(30 * time.Second)
		defer ticker.Stop()
		for {
			select {
			case <-requestContext.Done():
				return
			case <-ticker.C:
				r := model.DB.Model(&model.RelayAttempt{}).Where("id = ? AND state <> ?", id, "settled").Update("deadline", time.Now().Add(10*time.Minute).Unix())
				if r.Error == nil && r.RowsAffected == 0 {
					return
				}
			}
		}
	}(a.ID)
	return nil
}

// Log and account through one durable transaction instead of the ordinary
// single-generation session. A checkpoint calculates price but never debits.
func resolveAttemptGuardResult(reason string, info *relaycommon.RelayInfo) string {
	if reason != "" {
		return reason
	}
	if info != nil && info.ResponseModel != nil && info.ResponseModel.ReturnedModel != "" {
		if info.ResponseModel.Mismatch {
			return "mismatch_allowed"
		}
		return "accepted"
	}
	if info != nil && info.AstraAutoSwitch {
		return "accepted"
	}
	return "unchecked"
}

func recordDurableAttempt(c *gin.Context, info *relaycommon.RelayInfo, usage *dto.Usage, params model.RecordConsumeLogParams) {
	number := c.GetInt("relay_attempt_number")
	reason := resolveAttemptGuardResult(c.GetString("attempt_guard_result"), info)
	if reason == "model_mismatch" {
		reason = "mismatch_rejected"
	}
	returned := c.GetString("attempt_returned_model")
	if returned == "" && info.ResponseModel != nil {
		returned = info.ResponseModel.ReturnedModel
	}
	source := c.GetString("attempt_usage_source")
	if source == "" {
		source = "upstream"
		if usage != nil && usage.UsageSource == "estimated" {
			source = "estimated"
		}
	}
	status := "verified"
	if reason == "unchecked" {
		status = "unchecked"
	} else if reason == "mismatch_allowed" {
		status = "mismatch_allowed"
	} else if reason == "transport_error" {
		status = "transport_error"
	} else if reason != "accepted" {
		status = "rejected"
		if params.Quota > 0 {
			status = "charged_rejection"
		}
	} else if number > 1 {
		status = "switched_success"
	}
	params.Other.SetPublic("routing_status", status)
	params.Other.SetPublic("attempt_usage_source", source)
	params.Other.SetPublic("routing_attempt", map[string]any{"request_id": info.RequestId, "number": number, "guard_result": reason, "usage_source": source, "requested_model": info.OriginModelName, "returned_model": returned})
	params.Other.SetPublic("response_model", map[string]any{"requested_model": info.OriginModelName, "upstream_model": info.UpstreamModelName, "returned_model": returned, "mismatch": returned != "" && returned != info.OriginModelName})
	if reason != "accepted" && reason != "unchecked" {
		params.Other.SetPublic("model_guard_result", reason)
	}
	log := &model.Log{UserId: info.UserId, Username: common.GetContextKeyString(c, constant.ContextKeyUserName), CreatedAt: time.Now().Unix(), Type: model.LogTypeConsume, Content: params.Content, TokenName: params.TokenName, ModelName: params.ModelName, Quota: params.Quota, PromptTokens: params.PromptTokens, CompletionTokens: params.CompletionTokens, UseTime: params.UseTimeSeconds, IsStream: params.IsStream, ChannelId: params.ChannelId, TokenId: params.TokenId, Group: params.Group, RequestId: info.RequestId, Other: params.Other.JSONString(), CorrelationKey: c.GetString("log_correlation_key")}
	billable := params.Quota > 0 || params.PromptTokens > 0 || params.CompletionTokens > 0
	usageBody, _ := common.Marshal(usage)
	logBody, _ := common.Marshal(log)
	deadline := time.Now().Add(90 * time.Second).Unix()
	if c.GetBool("attempt_checkpoint") {
		deadline = time.Now().Add(10 * time.Minute).Unix()
	}
	if c.GetBool("attempt_checkpoint") {
		if err := model.CheckpointRelayAttempt(info.DurableAttemptID, params.Quota, string(usageBody), string(logBody), returned, reason, source, billable, deadline); err != nil {
			common.SysError("attempt checkpoint: " + err.Error())
		}
		return
	}
	details := model.RelayAttempt{UsageJSON: string(usageBody), ReturnedModel: returned, GuardResult: reason, UsageSource: source, Billable: billable}
	if err := model.SettleRelayAttempt(info.DurableAttemptID, params.Quota, log, details); err != nil {
		common.SysError("attempt settlement: " + err.Error())
		info.AttemptDetached = true
		// Retain the authoritative usage in memory during transient DB outages;
		// restart recovery can independently use the last durable checkpoint.
		id := info.DurableAttemptID
		quota := params.Quota
		go func() {
			for range 60 {
				time.Sleep(time.Second)
				if model.SettleRelayAttempt(id, quota, log, details) == nil {
					return
				}
			}
		}()
		return
	}
	if err := model.DeliverRelayAttemptLogs(); err != nil {
		common.SysError("attempt log outbox: " + err.Error())
	}
}

type GuardUsageCollector struct {
	ctx            *gin.Context
	info           relaycommon.RelayInfo
	usage          dto.Usage
	outputTokens   int
	observed       bool
	hasUsage       bool
	lastCheckpoint time.Time
}

func NewGuardUsageCollector(c *gin.Context, info *relaycommon.RelayInfo) *GuardUsageCollector {
	copyInfo := *info
	if info.ChannelMeta != nil {
		meta := *info.ChannelMeta
		copyInfo.ChannelMeta = &meta
	}
	if info.TieredBillingSnapshot != nil {
		snap := *info.TieredBillingSnapshot
		copyInfo.TieredBillingSnapshot = &snap
	}
	copyInfo.ResponseModel = nil
	copyInfo.ResponseGuard = nil
	copyInfo.ResponsesUsageInfo = nil
	copyInfo.StreamStatus = relaycommon.NewStreamStatus()
	copyContext := c.Copy()
	state := RequestPolicy(c)
	copyContext.Set(requestPolicyContextKey, &RequestPolicyState{StartedAt: state.StartedAt, Attempts: state.Attempts, SelectedGroup: state.SelectedGroup, SessionMode: state.SessionMode, SessionModeSource: state.SessionModeSource, RuleName: state.RuleName, events: state.Events()})
	return &GuardUsageCollector{ctx: copyContext, info: copyInfo}
}

func (collector *GuardUsageCollector) Observe(payload []byte, stream bool) {
	var event dto.ResponsesStreamResponse
	if stream {
		if common.Unmarshal(payload, &event) != nil {
			return
		}
	} else {
		var response dto.OpenAIResponsesResponse
		if common.Unmarshal(payload, &response) != nil {
			return
		}
		event.Type = "response.completed"
		event.Response = &response
	}
	if event.Response != nil {
		collector.info.ObserveResponseModel(event.Response.Model)
		if event.Response.Model != "" {
			collector.ctx.Set("attempt_returned_model", event.Response.Model)
			if event.Response.Model != collector.info.OriginModelName {
				collector.ctx.Set("attempt_guard_result", "model_mismatch")
			}
		}
		if event.Response.Usage != nil && (event.Type == "response.completed" || event.Type == "response.incomplete" || event.Type == "response.failed" || event.Type == "response.cancelled") {
			ApplyResponsesUsage(&collector.usage, event.Response.Usage)
			collector.hasUsage = true
		}
		if collector.outputTokens == 0 {
			text := ExtractOutputTextFromResponses(event.Response)
			if text != "" {
				collector.outputTokens = CountTextToken(text, collector.info.OriginModelName)
				collector.observed = true
			}
		}
	}
	if strings.HasSuffix(event.Type, ".delta") && event.Delta != "" {
		collector.observed = true
		collector.outputTokens += CountTextToken(event.Delta, collector.info.OriginModelName)
	}
	if collector.hasUsage || collector.observed || collector.ctx.GetString("attempt_returned_model") != "" && collector.ctx.GetString("attempt_returned_model") != collector.info.OriginModelName {
		if time.Since(collector.lastCheckpoint) > time.Second || collector.hasUsage {
			collector.persist(true)
			collector.lastCheckpoint = time.Now()
		}
	}
}

func (collector *GuardUsageCollector) persist(checkpoint bool) {
	usage := collector.usage
	if !collector.hasUsage {
		usage.PromptTokens = collector.info.GetEstimatePromptTokens()
		usage.CompletionTokens = collector.outputTokens
		usage.TotalTokens = usage.PromptTokens + usage.CompletionTokens
		collector.ctx.Set("attempt_usage_source", "estimated")
	} else {
		collector.ctx.Set("attempt_usage_source", "upstream")
	}
	collector.ctx.Set("attempt_checkpoint", checkpoint)
	PostTextConsumeQuota(collector.ctx, &collector.info, &usage, nil)
}

// Flush the most recent observed output when a stream ends between periodic
// checkpoints, including disconnects and late guard rejection.
func (collector *GuardUsageCollector) CheckpointFinal(reason string) {
	if !collector.hasUsage && !collector.observed && reason != "model_mismatch" {
		return
	}
	if reason != "accepted" && reason != "pending" {
		collector.ctx.Set("attempt_guard_result", reason)
	}
	collector.persist(true)
}

func (collector *GuardUsageCollector) FinishRejected(info *relaycommon.RelayInfo, guard *relaycommon.ResponseGuardAttempt, resp *http.Response) bool {
	if info.DurableAttemptID == 0 {
		return false
	}
	collector.ctx.Set("attempt_guard_result", guard.Reason)
	if guard.Reason != "model_mismatch" && !collector.hasUsage && !collector.observed {
		return false
	}
	info.AttemptDetached = true
	collector.persist(true)
	guard.Detach()
	go func() {
		defer func() {
			if recovered := recover(); recovered != nil {
				common.SysError(fmt.Sprintf("attempt usage collection panic: %v", recovered))
			}
		}()
		defer guard.Close()
		if resp != nil && resp.Body != nil {
			defer resp.Body.Close()
		}
		defer collector.persist(false)
		if guard.Drain != nil {
			guard.Drain()
		}
	}()
	return true
}

func EndFailedRelayAttempt(c *gin.Context, info *relaycommon.RelayInfo) {
	if info.DurableAttemptID == 0 || info.AttemptDetached {
		return
	}
	var attempt model.RelayAttempt
	if err := model.DB.First(&attempt, info.DurableAttemptID).Error; err != nil {
		common.SysError(err.Error())
		return
	}
	if attempt.State == "settled" {
		return
	}
	if !attempt.Billable {
		reason := "failed"
		if guard := info.ResponseGuard.Last(); guard != nil {
			reason = guard.Reason
		}
		_ = model.CheckpointRelayAttempt(attempt.ID, 0, "{}", "", attempt.ReturnedModel, reason, "", false, time.Now().Unix())
	}
	quota := 0
	var log *model.Log
	if attempt.Billable {
		quota = attempt.ObservedQuota
		log = &model.Log{}
		if common.UnmarshalJsonStr(attempt.LogJSON, log) != nil {
			log = nil
		}
	}
	if err := model.SettleRelayAttempt(attempt.ID, quota, log); err != nil {
		common.SysError(err.Error())
	}
}

func StartRelayAttemptRecovery() {
	go func() {
		for {
			var attempts []model.RelayAttempt
			if err := model.DB.Where("state <> ? AND deadline < ?", "settled", time.Now().Unix()).Limit(100).Find(&attempts).Error; err == nil {
				for _, attempt := range attempts {
					quota := 0
					var log *model.Log
					if attempt.Billable {
						quota = attempt.ObservedQuota
						log = &model.Log{}
						if common.UnmarshalJsonStr(attempt.LogJSON, log) != nil {
							continue
						}
					}
					if err := model.SettleRelayAttempt(attempt.ID, quota, log); err != nil {
						common.SysError("attempt recovery: " + err.Error())
					}
				}
			}
			if err := model.DeliverRelayAttemptLogs(); err != nil {
				common.SysError("attempt outbox: " + err.Error())
			}
			time.Sleep(5 * time.Second)
		}
	}()
}