feat: sync deployed v1.0.0-rc.38-custom.67

controller/responses_websocket_test.go

@@ -25,8 +25,10 @@ import (
 	"github.com/QuantumNous/new-api/relay"
 	"github.com/QuantumNous/new-api/relaykit/dto"
 	"github.com/QuantumNous/new-api/relaykit/types"
+	"github.com/QuantumNous/new-api/service"
 	"github.com/QuantumNous/new-api/setting"
 	"github.com/QuantumNous/new-api/setting/config"
+	"github.com/QuantumNous/new-api/setting/ratio_setting"
 	"github.com/alicebob/miniredis/v2"
 	"github.com/gin-gonic/gin"
 	"github.com/go-redis/redis/v8"
@@ -38,6 +40,84 @@ import (
 
 var responsesWSTestUserSequence atomic.Int64
 
+func TestResponsesHTTPPoolRetriesTwiceThenUsesBackup(t *testing.T) {
+	fixture := newResponsesWSBillingTest(t, `tier("request", fixed(0.002))`, func(*websocket.Conn, *http.Request) {})
+	oldRetries := common.RetryTimes
+	common.RetryTimes = 0
+	t.Cleanup(func() { common.RetryTimes = oldRetries })
+	oldGroups := setting.UserUsableGroups2JSONString()
+	require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(`{"default":"Default","vip":"VIP"}`))
+	t.Cleanup(func() { require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(oldGroups)) })
+	require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(`{"default":1,"vip":1}`))
+	if os.Getenv("TEST_PERSONAL_RATIO") == "1" {
+		require.NoError(t, model.SavePersonalGroupRatios(context.Background(), fixture.user.Id, map[string]float64{"default": 0.8, "vip": 0.2}))
+	}
+	service.ResetGroupAvailability()
+	t.Cleanup(service.ResetGroupAvailability)
+	var primaryCalls, backupCalls atomic.Int64
+	fixture.httpUpstream = func(w http.ResponseWriter, r *http.Request) {
+		primaryCalls.Add(1)
+		w.Header().Set("Content-Type", "application/json")
+		w.WriteHeader(500)
+		_, _ = fmt.Fprint(w, `{"error":{"type":"server_error","code":"server_error","message":"test failure"}}`)
+	}
+	backup := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		backupCalls.Add(1)
+		w.Header().Set("Content-Type", "application/json")
+		_, _ = fmt.Fprint(w, `{"id":"completed","status":"completed","usage":{"input_tokens":1000,"output_tokens":1,"total_tokens":1001}}`)
+	}))
+	t.Cleanup(backup.Close)
+	require.NoError(t, model.DB.Model(fixture.channel).Updates(map[string]any{"key": "single-primary-key", "channel_info": `{}`, "auto_ban": 0}).Error)
+	backupChannel := &model.Channel{Name: "backup-pool-test", Key: "test-only", Status: common.ChannelStatusEnabled, Type: constant.ChannelTypeOpenAI, Group: "vip", Models: "ws-billing", BaseURL: &backup.URL, AutoBan: common.GetPointer(0)}
+	require.NoError(t, model.DB.Create(backupChannel).Error)
+	require.NoError(t, model.DB.Create(&model.Ability{ChannelId: backupChannel.Id, Model: "ws-billing", Group: "vip", Enabled: true}).Error)
+	t.Cleanup(func() {
+		require.NoError(t, model.DB.Where("channel_id = ?", backupChannel.Id).Delete(&model.Ability{}).Error)
+		require.NoError(t, model.DB.Delete(backupChannel).Error)
+	})
+	require.NoError(t, model.DB.First(fixture.token, fixture.token.Id).Error)
+	body, err := common.Marshal(model.TokenRoutingPolicy{Version: 1, Mode: "custom", Groups: []model.TokenGroupEntry{{Group: "default", Priority: 0}, {Group: "vip", Priority: 1}}})
+	require.NoError(t, err)
+	fixture.token.RoutingPolicy = string(body)
+	fixture.token.Group = "default"
+	fixture.token.CrossGroupRetry = true
+	require.NoError(t, model.DB.Model(fixture.token).Updates(map[string]any{"routing_policy": string(body), "group": "default", "cross_group_retry": true}).Error)
+	require.NoError(t, fixture.token.SelectUpdate())
+	for range 2 {
+		req, err := http.NewRequest(http.MethodPost, fixture.gatewayURL+"/v1/responses", strings.NewReader(`{"model":"ws-billing","input":"hello"}`))
+		require.NoError(t, err)
+		req.Header.Set("Authorization", "Bearer sk-"+fixture.token.Key)
+		req.Header.Set("Content-Type", "application/json")
+		response, err := http.DefaultClient.Do(req)
+		require.NoError(t, err)
+		data, err := io.ReadAll(response.Body)
+		require.NoError(t, err)
+		require.NoError(t, response.Body.Close())
+		require.Equal(t, 200, response.StatusCode, string(data))
+		select {
+		case <-fixture.httpDone:
+		case <-time.After(3 * time.Second):
+			t.Fatal("HTTP request did not finish")
+		}
+	}
+	require.Equal(t, int64(2), primaryCalls.Load())
+	require.Equal(t, int64(2), backupCalls.Load())
+	var logs []model.Log
+	require.NoError(t, model.LOG_DB.Where("token_id = ? AND type = ?", fixture.token.Id, model.LogTypeConsume).Find(&logs).Error)
+	require.Len(t, logs, 2)
+	for _, log := range logs {
+		require.Equal(t, "vip", log.Group)
+		expected := 1000
+		if os.Getenv("TEST_PERSONAL_RATIO") == "1" {
+			expected = 200
+		}
+		require.Equal(t, expected, log.Quota)
+	}
+	if os.Getenv("TEST_PERSONAL_RATIO") == "1" {
+		assertResponsesWSAccounting(t, fixture, []int{200, 200})
+	}
+}
+
 func setupResponsesWSRequestTest(t *testing.T) (*model.User, *model.Token) {
 	t.Helper()
 	previousDB := model.DB
@@ -88,7 +168,7 @@ func setupResponsesWSRequestTest(t *testing.T) (*model.User, *model.Token) {
 		setting.ModelRequestRateLimitMutex.Unlock()
 		require.NoError(t, sqlDB.Close())
 	})
-	require.NoError(t, db.AutoMigrate(&model.User{}, &model.Token{}))
+	require.NoError(t, db.AutoMigrate(&model.User{}, &model.Token{}, &model.PersonalGroupRatio{}))
 	// The shared in-memory limiter outlives each database fixture. Give every
 	// user a separate quota bucket, including when the tests run with -count.
 	user := &model.User{Id: 5062000 + int(responsesWSTestUserSequence.Add(1)), Username: "responses-ws-user", Status: common.UserStatusEnabled, Group: "default", Quota: 1000, AuthVersion: 1}