fix: harden stream billing and affinity handling

.gitignore

@@ -35,3 +35,14 @@ data/
 .test
 token_estimator_test.go
 skills-lock.json
+.cla
+
+# repo-root screenshot artifacts and openapi dumps (not nested)
+/*.png
+/*.openapi.json
+/tmp
+
+# scripts/ top-level: research notes & probe scripts (cf-worker/ subdir stays tracked)
+/scripts/*.md
+/scripts/*.py
+/scripts/*.sh

common/constants.go

@@ -170,6 +170,8 @@ var BatchUpdateInterval int
 
 var RelayTimeout int // unit is second
 
+var RelayResponseHeaderTimeout int // unit is second; timeout for receiving response headers from upstream
+
 var RelayMaxIdleConns int
 var RelayMaxIdleConnsPerHost int
 

common/custom-event.go

@@ -63,9 +63,13 @@ func encode(writer io.Writer, event CustomEvent) error {
 }
 
 func writeData(w stringWriter, data interface{}) error {
-	dataReplacer.WriteString(w, fmt.Sprint(data))
+	if _, err := dataReplacer.WriteString(w, fmt.Sprint(data)); err != nil {
+		return err
+	}
 	if strings.HasPrefix(data.(string), "data") {
-		w.writeString("\n\n")
+		if _, err := w.writeString("\n\n"); err != nil {
+			return err
+		}
 	}
 	return nil
 }

common/init.go

@@ -102,6 +102,7 @@ func InitEnv() {
 	SyncFrequency = GetEnvOrDefault("SYNC_FREQUENCY", 60)
 	BatchUpdateInterval = GetEnvOrDefault("BATCH_UPDATE_INTERVAL", 5)
 	RelayTimeout = GetEnvOrDefault("RELAY_TIMEOUT", 0)
+	RelayResponseHeaderTimeout = GetEnvOrDefault("RELAY_RESPONSE_HEADER_TIMEOUT", 60)
 	RelayMaxIdleConns = GetEnvOrDefault("RELAY_MAX_IDLE_CONNS", 500)
 	RelayMaxIdleConnsPerHost = GetEnvOrDefault("RELAY_MAX_IDLE_CONNS_PER_HOST", 100)
 

common/quota_math.go

@@ -0,0 +1,95 @@
+package common
+
+import (
+	"fmt"
+	"math"
+	"strconv"
+
+	"github.com/shopspring/decimal"
+)
+
+const (
+	MaxQuota = math.MaxInt32
+	MinQuota = math.MinInt32
+)
+
+type QuotaClampKind string
+
+const (
+	QuotaClampOverflow  QuotaClampKind = "overflow"
+	QuotaClampUnderflow QuotaClampKind = "underflow"
+	QuotaClampNaN       QuotaClampKind = "nan"
+)
+
+type QuotaClamp struct {
+	Op       string         `json:"op"`
+	Kind     QuotaClampKind `json:"kind"`
+	Original string         `json:"original"`
+	Clamped  int            `json:"clamped"`
+}
+
+func (c *QuotaClamp) Error() string {
+	if c == nil {
+		return ""
+	}
+	return fmt.Sprintf("quota conversion (%s) %s: original=%s, clamped=%d", c.Op, c.Kind, c.Original, c.Clamped)
+}
+
+func saturateQuota(value float64, op string) (int, *QuotaClamp) {
+	var clamp *QuotaClamp
+	switch {
+	case math.IsNaN(value):
+		clamp = &QuotaClamp{Op: op, Kind: QuotaClampNaN, Original: "NaN", Clamped: 0}
+	case value > MaxQuota:
+		clamp = &QuotaClamp{Op: op, Kind: QuotaClampOverflow, Original: strconv.FormatFloat(value, 'g', -1, 64), Clamped: MaxQuota}
+	case value < MinQuota:
+		clamp = &QuotaClamp{Op: op, Kind: QuotaClampUnderflow, Original: strconv.FormatFloat(value, 'g', -1, 64), Clamped: MinQuota}
+	default:
+		return int(value), nil
+	}
+	return clamp.Clamped, clamp
+}
+
+func QuotaFromFloat(value float64) int {
+	quota, _ := QuotaFromFloatChecked(value)
+	return quota
+}
+
+func QuotaFromFloatChecked(value float64) (int, *QuotaClamp) {
+	return saturateQuota(value, "QuotaFromFloat")
+}
+
+func QuotaFromFloatStrict(value float64) (int, error) {
+	quota, clamp := QuotaFromFloatChecked(value)
+	if clamp != nil {
+		return 0, clamp
+	}
+	return quota, nil
+}
+
+func QuotaRound(value float64) int {
+	quota, _ := QuotaRoundChecked(value)
+	return quota
+}
+
+func QuotaRoundChecked(value float64) (int, *QuotaClamp) {
+	return saturateQuota(math.Round(value), "QuotaRound")
+}
+
+func QuotaRoundStrict(value float64) (int, error) {
+	quota, clamp := QuotaRoundChecked(value)
+	if clamp != nil {
+		return 0, clamp
+	}
+	return quota, nil
+}
+
+func QuotaFromDecimal(value decimal.Decimal) int {
+	quota, _ := QuotaFromDecimalChecked(value)
+	return quota
+}
+
+func QuotaFromDecimalChecked(value decimal.Decimal) (int, *QuotaClamp) {
+	rounded, _ := value.Round(0).Float64()
+	return saturateQuota(rounded, "QuotaFromDecimal")
+}

common/quota_math_test.go

@@ -0,0 +1,64 @@
+package common
+
+import (
+	"math"
+	"testing"
+
+	"github.com/shopspring/decimal"
+	"github.com/stretchr/testify/require"
+)
+
+func TestQuotaFromFloatSaturatesOutOfRangeValues(t *testing.T) {
+	tests := []struct {
+		name string
+		in   float64
+		want int
+	}{
+		{name: "overflow", in: math.MaxFloat64, want: MaxQuota},
+		{name: "underflow", in: -math.MaxFloat64, want: MinQuota},
+		{name: "nan", in: math.NaN(), want: 0},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got := QuotaFromFloat(tt.in)
+			require.Equal(t, tt.want, got)
+		})
+	}
+}
+
+func TestQuotaRoundStrictRejectsSaturation(t *testing.T) {
+	quota, err := QuotaRoundStrict(float64(MaxQuota) + 1)
+
+	require.Error(t, err)
+	require.Zero(t, quota)
+	clamp, ok := err.(*QuotaClamp)
+	require.True(t, ok)
+	require.Equal(t, QuotaClampOverflow, clamp.Kind)
+}
+
+func TestQuotaRoundStrictAcceptsIntegerBounds(t *testing.T) {
+	maxQuota, maxErr := QuotaRoundStrict(float64(MaxQuota))
+	minQuota, minErr := QuotaRoundStrict(float64(MinQuota))
+
+	require.NoError(t, maxErr)
+	require.NoError(t, minErr)
+	require.Equal(t, MaxQuota, maxQuota)
+	require.Equal(t, MinQuota, minQuota)
+}
+
+func TestQuotaClampNaNIsJSONSafe(t *testing.T) {
+	_, clamp := QuotaFromFloatChecked(math.NaN())
+	require.NotNil(t, clamp)
+
+	data, err := Marshal(clamp)
+
+	require.NoError(t, err)
+	require.Contains(t, string(data), `"original":"NaN"`)
+}
+
+func TestQuotaFromDecimalRoundsBeforeSaturating(t *testing.T) {
+	quota := QuotaFromDecimal(decimal.NewFromFloat(1.5))
+
+	require.Equal(t, 2, quota)
+}

controller/channel-test.go

@@ -132,7 +132,7 @@ func testChannel(channel *model.Channel, testModel string, endpointType string,
 			requestPath = "/v1/responses/compact"
 		}
 	}
-	if strings.HasPrefix(requestPath, "/v1/responses/compact") {
+	if strings.HasPrefix(requestPath, "/v1/responses/compact") && !ratio_setting.CompactUseBaseModel() {
 		testModel = ratio_setting.WithCompactModelSuffix(testModel)
 	}
 

controller/channel.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
+	"io"
 	"net/http"
 	"strconv"
 	"strings"
@@ -12,11 +13,13 @@ import (
 	"github.com/QuantumNous/new-api/common"
 	"github.com/QuantumNous/new-api/constant"
 	"github.com/QuantumNous/new-api/dto"
+	"github.com/QuantumNous/new-api/logger"
 	"github.com/QuantumNous/new-api/model"
 	relaychannel "github.com/QuantumNous/new-api/relay/channel"
 	"github.com/QuantumNous/new-api/relay/channel/gemini"
 	"github.com/QuantumNous/new-api/relay/channel/ollama"
 	"github.com/QuantumNous/new-api/service"
+	"github.com/QuantumNous/new-api/setting/system_setting"
 
 	"github.com/gin-gonic/gin"
 	"gorm.io/gorm"
@@ -69,6 +72,20 @@ func clearChannelInfo(channel *model.Channel) {
 	}
 }
 
+// fillChannelCooldown annotates a channel with its current in-memory cooldown
+// status (reason + expiry) so the admin UI can show why and for how long a
+// channel was temporarily taken out of selection.
+func fillChannelCooldown(channel *model.Channel) {
+	if channel == nil {
+		return
+	}
+	if reason, expires, cooling := model.GetChannelCooldown(channel.Id); cooling {
+		channel.CoolingDown = true
+		channel.CooldownReason = reason
+		channel.CooldownExpires = expires
+	}
+}
+
 func applyChannelStatusFilter(query *gorm.DB, statusFilter int) *gorm.DB {
 	if statusFilter == common.ChannelStatusEnabled {
 		return query.Where("status = ?", common.ChannelStatusEnabled)
@@ -159,6 +176,7 @@ func GetAllChannels(c *gin.Context) {
 
 	for _, datum := range channelData {
 		clearChannelInfo(datum)
+		fillChannelCooldown(datum)
 	}
 
 	countQuery := buildChannelListQuery(groupFilter, statusFilter, -1)
@@ -365,6 +383,7 @@ func SearchChannels(c *gin.Context) {
 
 	for _, datum := range pagedData {
 		clearChannelInfo(datum)
+		fillChannelCooldown(datum)
 	}
 
 	c.JSON(http.StatusOK, gin.H{
@@ -392,6 +411,7 @@ func GetChannel(c *gin.Context) {
 	}
 	if channel != nil {
 		clearChannelInfo(channel)
+		fillChannelCooldown(channel)
 	}
 	c.JSON(http.StatusOK, gin.H{
 		"success": true,
@@ -1054,10 +1074,27 @@ func FetchModels(c *gin.Context) {
 		return
 	}
 
-	client := &http.Client{}
+	client := service.GetHttpClient()
+	if client == nil {
+		c.JSON(http.StatusServiceUnavailable, gin.H{
+			"success": false,
+			"message": "HTTP client is not initialized",
+		})
+		return
+	}
 	url := fmt.Sprintf("%s/v1/models", baseURL)
+	fetchSetting := system_setting.GetFetchSetting()
+	if err := common.ValidateURLWithFetchSetting(url, fetchSetting.EnableSSRFProtection, fetchSetting.AllowPrivateIp, fetchSetting.DomainFilterMode, fetchSetting.IpFilterMode, fetchSetting.DomainList, fetchSetting.IpList, fetchSetting.AllowedPorts, fetchSetting.ApplyIPFilterForDomain); err != nil {
+		c.JSON(http.StatusBadRequest, gin.H{
+			"success": false,
+			"message": fmt.Sprintf("Invalid models URL: %s", err.Error()),
+		})
+		return
+	}
 
-	request, err := http.NewRequest("GET", url, nil)
+	requestCtx, cancel := context.WithTimeout(c.Request.Context(), 30*time.Second)
+	defer cancel()
+	request, err := http.NewRequestWithContext(requestCtx, http.MethodGet, url, nil)
 	if err != nil {
 		c.JSON(http.StatusInternalServerError, gin.H{
 			"success": false,
@@ -1070,29 +1107,31 @@ func FetchModels(c *gin.Context) {
 
 	response, err := client.Do(request)
 	if err != nil {
-		c.JSON(http.StatusInternalServerError, gin.H{
+		logger.LogError(c, "failed to fetch models: "+err.Error())
+		c.JSON(http.StatusBadGateway, gin.H{
 			"success": false,
-			"message": err.Error(),
+			"message": "Failed to fetch models",
 		})
 		return
 	}
-	//check status code
+	defer response.Body.Close()
 	if response.StatusCode != http.StatusOK {
-		c.JSON(http.StatusInternalServerError, gin.H{
+		c.JSON(http.StatusBadGateway, gin.H{
 			"success": false,
-			"message": "Failed to fetch models",
+			"message": fmt.Sprintf("Failed to fetch models: upstream status %d", response.StatusCode),
 		})
 		return
 	}
-	defer response.Body.Close()
 
+	const maxModelsResponseBytes = 5 << 20
+	limitedBody := io.LimitReader(response.Body, maxModelsResponseBytes+1)
 	var result struct {
 		Data []struct {
 			ID string `json:"id"`
 		} `json:"data"`
 	}
 
-	if err := json.NewDecoder(response.Body).Decode(&result); err != nil {
+	if err := json.NewDecoder(limitedBody).Decode(&result); err != nil {
 		c.JSON(http.StatusInternalServerError, gin.H{
 			"success": false,
 			"message": err.Error(),