new-api-custom/relay/relay_attempt_test.go

package relay

import (
	"context"
	"fmt"
	"github.com/QuantumNous/new-api/common"
	"github.com/QuantumNous/new-api/constant"
	"github.com/QuantumNous/new-api/model"
	"github.com/QuantumNous/new-api/pkg/billingexpr"
	"github.com/QuantumNous/new-api/relay/channel"
	"github.com/QuantumNous/new-api/relay/channel/openai"
	relaycommon "github.com/QuantumNous/new-api/relay/common"
	"github.com/QuantumNous/new-api/relaykit/dto"
	"github.com/QuantumNous/new-api/service"
	hosttypes "github.com/QuantumNous/new-api/types"
	"github.com/gin-gonic/gin"
	"github.com/glebarez/sqlite"
	"github.com/stretchr/testify/require"
	"gorm.io/gorm"
	"io"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
	"time"
)

func TestRelayAttemptMismatchCollectsUsageAndChargesRequestedPrice(t *testing.T) {
	service.InitTokenEncoders()
	db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
	require.NoError(t, err)
	sqlDB, err := db.DB()
	require.NoError(t, err)
	sqlDB.SetMaxOpenConns(1)
	previousDB, previousLog, previousRedis := model.DB, model.LOG_DB, common.RedisEnabled
	model.DB = db
	model.LOG_DB = db
	common.RedisEnabled = false
	t.Cleanup(func() {
		model.DB = previousDB
		model.LOG_DB = previousLog
		common.RedisEnabled = previousRedis
		require.NoError(t, sqlDB.Close())
	})
	require.NoError(t, db.AutoMigrate(&model.User{}, &model.Token{}, &model.Channel{}, &model.RelayAttempt{}, &model.Log{}, &model.QuotaData{}))
	user := model.User{Username: "attempt-relay", Quota: 10000}
	require.NoError(t, db.Create(&user).Error)
	token := model.Token{UserId: user.Id, Key: "attempt-relay", Status: common.TokenStatusEnabled, RemainQuota: 10000}
	require.NoError(t, db.Create(&token).Error)
	expression := `p * 2 + c * 4 + cr * 0.5`
	for i, ratio := range []float64{1, 3, 1} {
		selected := model.Channel{Name: fmt.Sprint(i), Status: common.ChannelStatusEnabled}
		require.NoError(t, db.Create(&selected).Error)
		// The final usage is buffered behind the rejected prefix: Drain must reuse
		// the same reader rather than lose bytes already read from the socket.
		body := "data: {\"type\":\"response.created\",\"response\":{\"model\":\"gpt-5.6-luna\"}}\n\n" +
			"data: {\"type\":\"response.output_text.delta\",\"delta\":\"hidden generated text\"}\n\n" +
			"data: {\"type\":\"response.completed\",\"response\":{\"model\":\"gpt-5.6-luna\",\"usage\":{\"input_tokens\":10,\"output_tokens\":2,\"total_tokens\":12,\"input_tokens_details\":{\"cached_tokens\":4}}}}\n\ndata: [DONE]\n\n"
		server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			w.Header().Set("Content-Type", "text/event-stream")
			_, _ = io.WriteString(w, body)
		}))
		ctx, cancel := context.WithCancel(context.Background())
		c, writer := gin.CreateTestContext(httptest.NewRecorder())
		_ = writer
		c.Request = httptest.NewRequest("POST", "/v1/responses", strings.NewReader(`{"model":"gpt-6-astra","input":"secret prompt"}`)).WithContext(ctx)
		c.Set("channel_id", selected.Id)
		c.Set("id", user.Id)
		c.Set("token_id", token.Id)
		common.SetContextKey(c, constant.ContextKeyUserName, user.Username)
		info := &relaycommon.RelayInfo{RequestId: fmt.Sprintf("attempt-request-%d", i), UserId: user.Id, TokenId: token.Id, TokenKey: token.Key, OriginModelName: "gpt-6-astra", UsingGroup: fmt.Sprint(ratio), UserSetting: dto.UserSetting{BillingPreference: "wallet_only"}, StartTime: time.Now(), DurableAttemptBilling: true, ResponseGuard: &relaycommon.ResponseModelGuard{}, ChannelMeta: &relaycommon.ChannelMeta{ChannelId: selected.Id, UpstreamModelName: "gpt-6-astra"}, PriceData: hosttypes.PriceData{GroupRatioInfo: hosttypes.GroupRatioInfo{GroupRatio: ratio}}, TieredBillingSnapshot: &billingexpr.BillingSnapshot{BillingMode: "tiered_expr", ExprString: expression, ExprHash: billingexpr.ExprHashString(expression), QuotaPerUnit: common.QuotaPerUnit, GroupRatio: ratio, EstimatedQuotaAfterGroup: 100}}
		info.SetEstimatePromptTokens(30)
		info.AstraAutoSwitch = i < 2
		if !info.AstraAutoSwitch {
			info.ResponseGuard = nil
		}
		require.Nil(t, service.BeginRelayAttempt(c, info))
		req, err := http.NewRequest("POST", server.URL, strings.NewReader("{}"))
		require.NoError(t, err)
		resp, err := channel.DoRequest(c, req, info)
		if info.AstraAutoSwitch {
			require.Error(t, err)
			require.Nil(t, resp)
			require.False(t, c.Writer.Written())
			require.True(t, info.AttemptDetached)
			require.False(t, service.RequestPolicy(c).Successful, "collector cannot mark the client request successful")
		} else {
			require.NoError(t, err)
			usage, apiErr := openai.OaiResponsesHandler(c, info, resp)
			require.Nil(t, apiErr)
			require.True(t, c.Writer.Written())
			require.True(t, info.ResponseModel.Mismatch)
			ConsumeResponsesQuota(c, info, usage)
		}
		var attempt model.RelayAttempt
		require.Eventually(t, func() bool {
			return db.First(&attempt, info.DurableAttemptID).Error == nil && attempt.State == "settled"
		}, time.Second*5, time.Millisecond*10)
		require.Equal(t, int(11*ratio), attempt.Quota)
		if info.AstraAutoSwitch {
			require.Equal(t, "mismatch_rejected", attempt.GuardResult)
		} else {
			require.Equal(t, "mismatch_allowed", attempt.GuardResult)
			require.Contains(t, attempt.LogJSON, "mismatch_allowed")
		}
		require.Equal(t, "upstream", attempt.UsageSource)
		require.NotContains(t, attempt.LogJSON, "secret prompt")
		require.NotContains(t, attempt.LogJSON, "hidden generated text")
		require.NoError(t, model.DeliverRelayAttemptLogs())
		service.EndFailedRelayAttempt(c, info)
		cancel()
		server.Close()
	}
	require.NoError(t, db.First(&user, user.Id).Error)
	require.Equal(t, 9945, user.Quota)
	require.NoError(t, db.First(&token, token.Id).Error)
	require.Equal(t, 55, token.UsedQuota)
	var logs int64
	require.NoError(t, db.Model(&model.Log{}).Count(&logs).Error)
	require.EqualValues(t, 3, logs)
}