new-api/relay/relay_task_test.go

package relay

import (
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"

	"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/types"
	"github.com/gin-gonic/gin"
	"github.com/glebarez/sqlite"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"gorm.io/gorm"
)

func setupRelayTaskTestDB(t *testing.T) {
	t.Helper()

	db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
	require.NoError(t, err)
	require.NoError(t, db.AutoMigrate(&model.Task{}, &model.Channel{}))
	sqlDB, err := db.DB()
	require.NoError(t, err)
	sqlDB.SetMaxOpenConns(1)

	origDB := model.DB
	origUsingSQLite := common.UsingSQLite
	origRedisEnabled := common.RedisEnabled
	model.DB = db
	common.UsingSQLite = true
	common.RedisEnabled = false

	t.Cleanup(func() {
		model.DB = origDB
		common.UsingSQLite = origUsingSQLite
		common.RedisEnabled = origRedisEnabled
		require.NoError(t, sqlDB.Close())
	})
}

func TestResolveOriginTask_RestoresMatrixOriginPricing(t *testing.T) {
	setupRelayTaskTestDB(t)
	require.NoError(t, model.DB.Create(&model.Channel{
		Id:     11,
		Type:   constant.ChannelTypeMiniMax,
		Key:    "test-key",
		Status: common.ChannelStatusEnabled,
		Name:   "hailuo",
		Group:  "default",
		Models: "matrix-video",
	}).Error)
	require.NoError(t, model.DB.Create(&model.Task{
		TaskID:    "task_origin",
		UserId:    7,
		ChannelId: 11,
		Properties: model.Properties{
			OriginModelName: "matrix-video",
		},
		PrivateData: model.TaskPrivateData{
			BillingContext: &model.TaskBillingContext{
				BillingMode:       types.BillingModeMatrix,
				BillingUnit:       types.BillingUnitPer1MTokens,
				TokenUnitPriceUSD: 0.5,
				PricingSnapshot: map[string]any{
					"resolution":   "720P",
					"price":        0.5,
					"billing_unit": types.BillingUnitPer1MTokens,
					"pricing_mode": string(types.PricingModeMatched),
				},
				PerCallBilling: false,
			},
		},
	}).Error)

	c := newRelayTaskRemixContext(t, "task_origin")
	info := &relaycommon.RelayInfo{UserId: 7, TaskRelayInfo: &relaycommon.TaskRelayInfo{}}

	taskErr := ResolveOriginTask(c, info)
	require.Nil(t, taskErr)
	require.NotNil(t, info.OriginPricing)
	require.NotNil(t, info.LockedChannel)
	assert.Equal(t, constant.TaskActionRemix, info.Action)
	assert.Equal(t, "matrix-video", info.OriginModelName)
	assert.Equal(t, types.BillingModeMatrix, info.OriginPricing.BillingMode)
	assert.Equal(t, types.BillingUnitPer1MTokens, info.OriginPricing.BillingUnit)
	assert.Equal(t, 0.5, info.OriginPricing.TokenUnitPriceUSD)
	assert.Equal(t, "720P", info.OriginPricing.Snapshot["resolution"])
	lockedChannel, ok := info.LockedChannel.(*model.Channel)
	require.True(t, ok)
	assert.Equal(t, 11, lockedChannel.Id)
}

func TestResolveOriginTask_RestoresLegacyOriginPricing(t *testing.T) {
	setupRelayTaskTestDB(t)
	require.NoError(t, model.DB.Create(&model.Channel{
		Id:     12,
		Type:   constant.ChannelTypeOpenAI,
		Key:    "test-key",
		Status: common.ChannelStatusEnabled,
		Name:   "sora",
		Group:  "default",
		Models: "legacy-video",
	}).Error)
	require.NoError(t, model.DB.Create(&model.Task{
		TaskID:    "task_legacy",
		UserId:    7,
		ChannelId: 12,
		Properties: model.Properties{
			OriginModelName: "legacy-video",
		},
		PrivateData: model.TaskPrivateData{
			BillingContext: &model.TaskBillingContext{
				ModelPrice:     0.2,
				OtherRatios:    map[string]float64{"seconds": 10, "size": 1.666667},
				PerCallBilling: true,
			},
		},
	}).Error)

	c := newRelayTaskRemixContext(t, "task_legacy")
	info := &relaycommon.RelayInfo{UserId: 7, TaskRelayInfo: &relaycommon.TaskRelayInfo{}}

	taskErr := ResolveOriginTask(c, info)
	require.Nil(t, taskErr)
	require.NotNil(t, info.OriginPricing)
	require.NotNil(t, info.LockedChannel)
	assert.Equal(t, "legacy-video", info.OriginModelName)
	assert.Equal(t, 0.2, info.OriginPricing.PriceUSD)
	assert.Equal(t, map[string]float64{"seconds": 10, "size": 1.666667}, info.OriginPricing.OtherRatios)
	assert.True(t, info.OriginPricing.PerCallBilling)
	lockedChannel, ok := info.LockedChannel.(*model.Channel)
	require.True(t, ok)
	assert.Equal(t, 12, lockedChannel.Id)
}

func TestBuildUpstreamRequestSnapshotRedactsDataURIsAndKeepsDebugFields(t *testing.T) {
	raw := []byte(`{
		"model":"doubao-seedance-2-0-260128",
		"content":[
			{"type":"text","text":"prompt"},
			{"type":"image_url","image_url":{"url":"data:image/png;base64,QUJDREVGRw=="}},
			{"type":"video_url","video_url":{"url":"asset://video-1"}},
			{"type":"audio_url","audio_url":{"url":"https://example.test/audio.wav"}}
		]
	}`)

	snapshot := buildUpstreamRequestSnapshot(raw)

	require.NotNil(t, snapshot)
	require.Equal(t, len(raw), snapshot.RawBytes)
	require.NotEmpty(t, snapshot.SHA256)
	require.True(t, snapshot.Redacted)
	require.False(t, snapshot.Truncated)
	data, err := common.Marshal(snapshot)
	require.NoError(t, err)
	jsonBody := string(data)
	require.Contains(t, jsonBody, `"model":"doubao-seedance-2-0-260128"`)
	require.Contains(t, jsonBody, `"asset://video-1"`)
	require.Contains(t, jsonBody, `"https://example.test/audio.wav"`)
	require.Contains(t, jsonBody, `"kind":"data_uri"`)
	require.Contains(t, jsonBody, `"media_type":"image/png"`)
	require.NotContains(t, jsonBody, "QUJDREVGRw==")
	require.Equal(t, len(data), snapshot.StoredBytes)
}

func TestBuildUpstreamRequestSnapshotTruncatesOversizedSnapshot(t *testing.T) {
	raw := []byte(`{"model":"m","prompt":"` + strings.Repeat("x", maxUpstreamRequestSnapshotBytes+1024) + `"}`)

	snapshot := buildUpstreamRequestSnapshot(raw)

	require.NotNil(t, snapshot)
	require.True(t, snapshot.Truncated)
	require.NotEmpty(t, snapshot.SHA256)
	data, err := common.Marshal(snapshot)
	require.NoError(t, err)
	require.LessOrEqual(t, len(data), maxUpstreamRequestSnapshotBytes)
	require.Equal(t, len(data), snapshot.StoredBytes)
}

func TestTaskModel2DtoDoesNotExposeUpstreamRequestSnapshot(t *testing.T) {
	task := &model.Task{
		TaskID: "task_public",
		PrivateData: model.TaskPrivateData{
			UpstreamRequest: &model.TaskUpstreamRequestSnapshot{
				Body: map[string]any{"secret": "hidden"},
			},
		},
	}

	dtoTask := TaskModel2Dto(task)
	data, err := common.Marshal(dtoTask)

	require.NoError(t, err)
	require.NotContains(t, string(data), "upstream_request")
	require.NotContains(t, string(data), "hidden")
}

func TestTaskModel2DtoSanitizesAipingInternalIDFromData(t *testing.T) {
	task := &model.Task{
		TaskID: "task_public",
		Data: []byte(`{
			"id":"cgt-upstream",
			"aiping_id":"b7ee658f-acac-433d-897d-5c9ed55ef402",
			"status":"succeeded"
		}`),
	}

	dtoTask := TaskModel2Dto(task)
	data, err := common.Marshal(dtoTask)

	require.NoError(t, err)
	require.NotContains(t, string(data), "aiping_id")
	require.NotContains(t, string(data), "b7ee658f-acac-433d-897d-5c9ed55ef402")
	require.Contains(t, string(data), `"id":"task_public"`)
}

func newRelayTaskRemixContext(t *testing.T, originTaskID string) *gin.Context {
	t.Helper()
	gin.SetMode(gin.TestMode)
	w := httptest.NewRecorder()
	c, _ := gin.CreateTestContext(w)
	c.Request = httptest.NewRequest(http.MethodPost, "/v1/videos/"+originTaskID+"/remix", nil)
	c.Params = gin.Params{{Key: "video_id", Value: originTaskID}}
	return c
}