new-api-custom/controller/detection_test.go

package controller

import (
	"bufio"
	"github.com/QuantumNous/new-api/common"
	"github.com/QuantumNous/new-api/model"
	"github.com/QuantumNous/new-api/pkg/billingexpr"
	fingerprint "github.com/QuantumNous/new-api/pkg/model_fingerprint"
	"github.com/QuantumNous/new-api/service"
	"github.com/gin-gonic/gin"
	"github.com/glebarez/sqlite"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"gorm.io/gorm"
	"io"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
	"time"
)

func TestDetectionProxyFlushesBeforeGenerationAndConsumesCredentialOnce(t *testing.T) {
	db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
	require.NoError(t, err)
	connection, err := db.DB()
	require.NoError(t, err)
	connection.SetMaxOpenConns(1)
	oldDB := model.DB
	model.DB = db
	t.Cleanup(func() { model.DB = oldDB; _ = connection.Close() })
	require.NoError(t, db.AutoMigrate(&model.Channel{}, &model.DetectionRun{}, &model.DetectionCall{}))
	release := make(chan struct{})
	upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		select {
		case <-release:
		case <-r.Context().Done():
			return
		}
		w.Header().Set("Content-Type", "application/json")
		_, _ = io.WriteString(w, `{"status":"completed","usage":{"input_tokens":10,"output_tokens":5},"output":[{"type":"message","content":[{"type":"output_text","text":"<svg></svg>"}]}]}`)
	}))
	defer upstream.Close()
	url := upstream.URL
	channel := model.Channel{Type: 1, Key: "test-only", BaseURL: &url, Status: common.ChannelStatusEnabled, Models: fingerprint.Model}
	require.NoError(t, db.Create(&channel).Error)
	price, _ := common.Marshal(service.DetectionPrice{Snapshot: billingexpr.BillingSnapshot{ExprString: "p+c", ExprHash: billingexpr.ExprHashString("p+c"), GroupRatio: 1, QuotaPerUnit: 500000, ExprVersion: 1}})
	run := model.DetectionRun{ID: "stream-test", RequestKey: "stream-test", Kind: "pelican", ChannelID: channel.Id, Status: "running", Deadline: time.Now().Unix() + 60, Reserved: 100, Price: string(price)}
	require.NoError(t, db.Create(&run).Error)
	call := model.DetectionCall{ID: "stream-call", RunID: run.ID, SecretHash: model.DetectionSecretHash("det_test"), State: "ready", ExpiresAt: time.Now().Unix() + 60}
	require.NoError(t, db.Create(&call).Error)
	router := gin.New()
	router.POST("/responses", DetectionProxy)
	gateway := httptest.NewServer(router)
	defer gateway.Close()
	client := http.Client{Timeout: 5 * time.Second}
	request, err := http.NewRequest(http.MethodPost, gateway.URL+"/responses", strings.NewReader(`{"model":"gpt-6-astra","input":"draw"}`))
	require.NoError(t, err)
	request.Header.Set("Authorization", "Bearer det_test")
	response, err := client.Do(request)
	close(release)
	require.NoError(t, err)
	defer response.Body.Close()
	reader := bufio.NewReader(response.Body)
	first, err := reader.ReadString('\n')
	require.NoError(t, err)
	assert.Equal(t, ": detection started\n", first)
	raw, err := io.ReadAll(reader)
	require.NoError(t, err)
	assert.Contains(t, string(raw), "response.completed")
	require.NoError(t, db.First(&call, "id = ?", call.ID).Error)
	assert.Equal(t, "finished", call.State)
	assert.Contains(t, call.Usage, "input_tokens")
	replay := httptest.NewRecorder()
	request = httptest.NewRequest(http.MethodPost, "/responses", strings.NewReader(`{"model":"gpt-6-astra","input":"draw"}`))
	request.Header.Set("Authorization", "Bearer det_test")
	router.ServeHTTP(replay, request)
	assert.Equal(t, 401, replay.Code)
}

func TestDetectionProxyRejectsEscapesAndEnforcesBounds(t *testing.T) {
	for _, body := range []string{
		`{"model":[],"input":"hi"}`,
		`{"model":"gpt-6-luna","input":"hi"}`,
		`{"model":"gpt-6-astra","input":[{"role":[],"content":"hi"}]}`,
		`{"model":"gpt-6-astra","input":"hi","previous_response_id":"x"}`,
		`{"model":"gpt-6-astra","input":[{"type":"input_image","image_url":"file:///etc/passwd"}]}`,
	} {
		_, err := detectionProxyRequest([]byte(body))
		require.Error(t, err)
	}
	request, err := detectionProxyRequest([]byte(`{"model":"gpt-6-astra","input":"draw a pelican","max_output_tokens":1000000,"tools":[{"type":"web_search"}]}`))
	require.NoError(t, err)
	require.NotNil(t, request.MaxOutputTokens)
	assert.EqualValues(t, 16384, *request.MaxOutputTokens)
	assert.Equal(t, "[]", string(request.Tools))
}

func TestDetectionOnlyBillsCompleteReportedUsage(t *testing.T) {
	answer, parseErr := fingerprint.ReadResponse(strings.NewReader(`{"status":"incomplete","usage":{"input_tokens":100,"output_tokens":10,"input_tokens_details":{"cached_tokens":40,"cache_creation_tokens":5}}}`), "application/json")
	require.ErrorIs(t, parseErr, fingerprint.ErrIncomplete)
	assert.True(t, answer.HasUsage)
	assert.EqualValues(t, 5, answer.CacheWriteTokens)
	assert.Empty(t, answer.Text)
	_, ok := detectionResponseUsage([]byte(`{"status":"completed"}`), "application/json")
	assert.False(t, ok)
	_, ok = detectionResponseUsage([]byte("data: {\"type\":\"response.output_text.delta\",\"delta\":\"hello\"}\n\n"), "text/event-stream")
	assert.False(t, ok)
	usage, ok := detectionResponseUsage([]byte("data: {\"type\":\"response.completed\",\"response\":{\"status\":\"completed\",\"usage\":{\"input_tokens\":100,\"output_tokens\":10}}}\n\n"), "text/event-stream")
	require.True(t, ok)
	assert.EqualValues(t, 100, usage.Get("input_tokens").Int())
	expression := `p * 2 + cr * 0.2 + c * 10`
	price := service.DetectionPrice{Snapshot: billingexpr.BillingSnapshot{ExprString: expression, ExprHash: billingexpr.ExprHashString(expression), GroupRatio: 0.18, QuotaPerUnit: 500000, ExprVersion: 1}}
	encoded, err := common.Marshal(price)
	require.NoError(t, err)
	quota, err := service.DetectionUsageQuota(string(encoded), 100, 10, 40)
	require.NoError(t, err)
	assert.Equal(t, 21, quota)
	_, err = service.DetectionUsageQuota(string(encoded), 10, 20, 11)
	assert.Error(t, err)
	at := time.Date(2026, 9, 24, 9, 0, 0, 0, time.UTC)
	price.Snapshot.ExprString = `hour("UTC") == 9 ? p * 2 : p * 100`
	price.Snapshot.ExprHash = billingexpr.ExprHashString(price.Snapshot.ExprString)
	price.Request.At = &at
	encoded, err = common.Marshal(price)
	require.NoError(t, err)
	quota, err = service.DetectionUsageQuota(string(encoded), 100, 0, 0)
	require.NoError(t, err)
	assert.Equal(t, 18, quota)
}

func TestDetectionArtworkCannotNavigate(t *testing.T) {
	cleaned := sanitizeDetectionHTML(`<meta http-equiv="refresh" content="0;url=https://evil.test"><script>alert(1)</script><a href="https://evil.test" onclick="alert(1)">x</a><iframe srcdoc="bad"></iframe><svg><circle cx="20" cy="20" r="10"/></svg><style>circle{fill:red}</style>`)
	for _, bad := range []string{"evil.test", "onclick", "<iframe", "http-equiv"} {
		assert.NotContains(t, cleaned, bad)
	}
	assert.Contains(t, cleaned, "<circle")
	assert.Contains(t, cleaned, "<style>")
}

func TestDetectionPriceBoundIncludesIntermediateTiers(t *testing.T) {
	bound, err := billingexpr.DetectionCostBound(`len < 100 ? p * 2 : len < 1000 ? p * 100 : p * 1`, 32768, 16384)
	require.NoError(t, err)
	assert.Equal(t, float64(3276800), bound)
	_, err = billingexpr.DetectionCostBound(`p / (c + 1)`, 32768, 16384)
	require.Error(t, err)
	bound, err = billingexpr.DetectionCostBound(`len <= 272000 ? tier("standard", p * 10 + c * 50 + cr * 1 + cc * 12) : tier("long", p * 20 + c * 100 + cr * 2 + cc * 25)`, 32768, 16384)
	require.NoError(t, err)
	assert.Greater(t, bound, float64(16384*100))
}

func TestDetectionLatestPerChannelStatus(t *testing.T) {
	now := time.Now().Unix()
	for _, tc := range []struct {
		name     string
		ids      []int
		evidence map[int]service.ChannelDetectionEvidence
		want     int
	}{
		{"successful fingerprint", []int{1}, map[int]service.ChannelDetectionEvidence{1: {At: now, Success: true}}, 1},
		{"failed availability", []int{1}, map[int]service.ChannelDetectionEvidence{1: {At: now, Success: false}}, 0},
		{"partial coverage", []int{1, 2}, map[int]service.ChannelDetectionEvidence{1: {At: now, Success: true}}, 2},
		{"failure and unknown", []int{1, 2}, map[int]service.ChannelDetectionEvidence{1: {At: now, Success: false}}, -1},
		{"expired success", []int{1}, map[int]service.ChannelDetectionEvidence{1: {At: now - 7201, Success: true}}, -1},
	} {
		t.Run(tc.name, func(t *testing.T) {
			status, _ := service.DetectionFallbackStatus(tc.ids, tc.evidence, now)
			assert.Equal(t, tc.want, status)
		})
	}
}