controller/responses_websocket_test.go
Rama main · réplica de lectura del repositorio original.
package controller
import (
"bufio"
"context"
"errors"
"fmt"
"io"
"net/http"
"net/http/httptest"
"os"
"path/filepath"
"strings"
"sync"
"sync/atomic"
"testing"
"time"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/i18n"
"github.com/QuantumNous/new-api/middleware"
"github.com/QuantumNous/new-api/model"
"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/alicebob/miniredis/v2"
"github.com/gin-gonic/gin"
"github.com/go-redis/redis/v8"
"github.com/gorilla/websocket"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"gorm.io/gorm"
)
var responsesWSTestUserSequence atomic.Int64
func setupResponsesWSRequestTest(t *testing.T) (*model.User, *model.Token) {
t.Helper()
previousDB := model.DB
previousLogDB := model.LOG_DB
previousType := common.MainDatabaseType()
previousLogType := common.LogDatabaseType()
previousRedis := common.RedisEnabled
previousRDB := common.RDB
previousMaster, previousMemory, previousSQLite := common.IsMasterNode, common.MemoryCacheEnabled, common.SQLitePath
previousEnabled := setting.ModelRequestRateLimitEnabled
previousDuration := setting.ModelRequestRateLimitDurationMinutes
previousTotal := setting.ModelRequestRateLimitCount
previousSuccess := setting.ModelRequestRateLimitSuccessCount
setting.ModelRequestRateLimitMutex.Lock()
previousGroups := setting.ModelRequestRateLimitGroup
setting.ModelRequestRateLimitGroup = nil
setting.ModelRequestRateLimitMutex.Unlock()
t.Setenv("SQL_DSN", os.Getenv("TEST_RESPONSES_SQL_DSN"))
t.Setenv("LOG_SQL_DSN", os.Getenv("TEST_RESPONSES_LOG_SQL_DSN"))
common.IsMasterNode, common.MemoryCacheEnabled, common.RedisEnabled = false, false, false
common.SQLitePath = filepath.Join(t.TempDir(), "responses.db")
require.NoError(t, model.InitDB())
db := model.DB
sqlDB, err := db.DB()
require.NoError(t, err)
sqlDB.SetMaxOpenConns(1)
require.NoError(t, model.InitLogDB())
if model.LOG_DB != db {
logSQL, err := model.LOG_DB.DB()
require.NoError(t, err)
logSQL.SetMaxOpenConns(1)
t.Cleanup(func() { require.NoError(t, logSQL.Close()) })
}
setting.ModelRequestRateLimitEnabled = false
t.Cleanup(func() {
model.DB = previousDB
model.LOG_DB = previousLogDB
common.SetDatabaseTypes(previousType, previousLogType)
common.RedisEnabled = previousRedis
common.RDB = previousRDB
common.IsMasterNode, common.MemoryCacheEnabled, common.SQLitePath = previousMaster, previousMemory, previousSQLite
setting.ModelRequestRateLimitEnabled = previousEnabled
setting.ModelRequestRateLimitDurationMinutes = previousDuration
setting.ModelRequestRateLimitCount = previousTotal
setting.ModelRequestRateLimitSuccessCount = previousSuccess
setting.ModelRequestRateLimitMutex.Lock()
setting.ModelRequestRateLimitGroup = previousGroups
setting.ModelRequestRateLimitMutex.Unlock()
require.NoError(t, sqlDB.Close())
})
require.NoError(t, db.AutoMigrate(&model.User{}, &model.Token{}))
// 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}
require.NoError(t, db.Create(user).Error)
token := &model.Token{UserId: user.Id, Key: "responseswstoken", Status: common.TokenStatusEnabled, ExpiredTime: -1, RemainQuota: 100}
require.NoError(t, db.Create(token).Error)
t.Cleanup(func() {
require.NoError(t, db.Unscoped().Delete(token).Error)
require.NoError(t, db.Unscoped().Delete(user).Error)
})
return user, token
}
func newResponsesWSTestRunner(t *testing.T, token *model.Token) (relay.ResponsesWSRequestRunner, *http.Request) {
t.Helper()
var runner relay.ResponsesWSRequestRunner
engine := gin.New()
require.NoError(t, engine.SetTrustedProxies([]string{"127.0.0.1"}))
engine.GET("/v1/responses", middleware.TokenAuth(), func(c *gin.Context) {
runner = newResponsesWSRequestRunner(c)
})
request := httptest.NewRequest(http.MethodGet, "/v1/responses", nil)
request.RemoteAddr = "127.0.0.1:12345"
request.Header.Set("Authorization", "Bearer sk-"+token.Key)
request.Header.Set("X-Forwarded-For", "203.0.113.8")
request.Header.Set("Connection", "Upgrade")
request.Header.Set("Upgrade", "websocket")
request.Header.Set("Sec-WebSocket-Key", "handshake-key")
request.Header.Set("Sec-WebSocket-Version", "13")
request.Header.Set("Sec-WebSocket-Extensions", "permessage-deflate")
request.Header.Set("Sec-WebSocket-Protocol", "responses")
request.Header.Set("Content-Length", "99")
request.Header.Set("Content-Encoding", "gzip")
request.Header.Set("X-Application-Header", "preserved")
response := httptest.NewRecorder()
engine.ServeHTTP(response, request)
require.Equal(t, http.StatusOK, response.Code)
require.NotNil(t, runner)
return runner, request
}
func TestResponsesWSRequestRunnerRefreshesBillingContextAndCleansBody(t *testing.T) {
user, token := setupResponsesWSRequestTest(t)
allowedIP := "203.0.113.8"
require.NoError(t, model.DB.Model(token).Update("allow_ips", allowedIP).Error)
runner, handshake := newResponsesWSTestRunner(t, token)
// A later upstream header mutation must not replace the client's credential.
handshake.Header.Set("Authorization", "Bearer upstream-secret")
handshake.Header.Set("X-Forwarded-For", "198.51.100.99")
var storage common.BodyStorage
request := httptest.NewRequest(http.MethodPost, "/v1/responses", strings.NewReader(`{"input":"first"}`))
require.Nil(t, runner(request, "ws-first", func(c *gin.Context) *types.NewAPIError {
assert.Equal(t, user.Id, c.GetInt("id"))
assert.Equal(t, 100, c.GetInt("token_quota"))
assert.Equal(t, "203.0.113.8", c.ClientIP())
assert.Equal(t, "ws-first", c.GetString(common.RequestIdKey))
assert.Equal(t, "ws-first", c.Request.Context().Value(common.RequestIdKey))
for _, header := range []string{"Connection", "Upgrade", "Sec-WebSocket-Key", "Sec-WebSocket-Version", "Sec-WebSocket-Extensions", "Sec-WebSocket-Protocol", "Content-Length", "Content-Encoding"} {
assert.Empty(t, c.GetHeader(header))
}
assert.Equal(t, "preserved", c.GetHeader("X-Application-Header"))
c.Set("previous_turn", true)
var err error
storage, err = common.GetBodyStorage(c)
require.NoError(t, err)
data, err := storage.Bytes()
require.NoError(t, err)
assert.JSONEq(t, `{"input":"first"}`, string(data))
return nil
}))
_, err := storage.Bytes()
assert.ErrorIs(t, err, common.ErrStorageClosed)
require.NoError(t, model.DB.Model(token).Updates(map[string]any{
"remain_quota": 2, "model_limits_enabled": true, "model_limits": "gpt-5.1",
}).Error)
require.NoError(t, model.DB.Model(user).Updates(map[string]any{
"group": "updated", "setting": `{"billing_preference":"wallet_only"}`,
}).Error)
request = httptest.NewRequest(http.MethodPost, "/v1/responses", strings.NewReader(`{"input":"second"}`))
require.Nil(t, runner(request, "ws-second", func(c *gin.Context) *types.NewAPIError {
assert.Equal(t, 2, c.GetInt("token_quota"))
assert.True(t, c.GetBool("token_model_limit_enabled"))
modelLimits, _ := c.Get("token_model_limit")
assert.Equal(t, map[string]bool{"gpt-5.1": true}, modelLimits)
assert.Equal(t, "updated", common.GetContextKeyString(c, constant.ContextKeyUsingGroup))
userSetting, ok := common.GetContextKeyType[dto.UserSetting](c, constant.ContextKeyUserSetting)
require.True(t, ok)
assert.Equal(t, "wallet_only", userSetting.BillingPreference)
assert.Equal(t, "ws-second", c.GetString(common.RequestIdKey))
_, leaked := c.Get("previous_turn")
assert.False(t, leaked)
_, leaked = c.Get(common.KeyBodyStorage)
assert.False(t, leaked)
return nil
}))
}
func TestResponsesWSRequestRunnerRejectsRevokedCredentials(t *testing.T) {
for _, tc := range []struct {
name string
status int
update func(*testing.T, *model.User, *model.Token)
}{
{name: "disabled token", status: http.StatusUnauthorized, update: func(t *testing.T, _ *model.User, token *model.Token) {
require.NoError(t, model.DB.Model(token).Update("status", common.TokenStatusDisabled).Error)
}},
{name: "expired token", status: http.StatusUnauthorized, update: func(t *testing.T, _ *model.User, token *model.Token) {
require.NoError(t, model.DB.Model(token).Update("expired_time", time.Now().Add(-time.Minute).Unix()).Error)
}},
{name: "exhausted token", status: http.StatusUnauthorized, update: func(t *testing.T, _ *model.User, token *model.Token) {
require.NoError(t, model.DB.Model(token).Update("remain_quota", 0).Error)
}},
{name: "disabled user", status: http.StatusForbidden, update: func(t *testing.T, user *model.User, _ *model.Token) {
require.NoError(t, model.DB.Model(user).Update("status", common.UserStatusDisabled).Error)
}},
{name: "deleted token", status: http.StatusUnauthorized, update: func(t *testing.T, _ *model.User, token *model.Token) {
require.NoError(t, model.DB.Delete(token).Error)
}},
{name: "changed IP restriction", status: http.StatusForbidden, update: func(t *testing.T, _ *model.User, token *model.Token) {
require.NoError(t, model.DB.Model(token).Update("allow_ips", "198.51.100.0/24").Error)
}},
} {
t.Run(tc.name, func(t *testing.T) {
user, token := setupResponsesWSRequestTest(t)
runner, _ := newResponsesWSTestRunner(t, token)
require.Nil(t, runner(httptest.NewRequest(http.MethodPost, "/v1/responses", nil), "before", func(*gin.Context) *types.NewAPIError { return nil }))
tc.update(t, user, token)
called := false
apiError := runner(httptest.NewRequest(http.MethodPost, "/v1/responses", nil), "after", func(*gin.Context) *types.NewAPIError {
called = true
return nil
})
require.NotNil(t, apiError)
assert.Equal(t, tc.status, apiError.StatusCode)
assert.False(t, called)
assert.NotContains(t, apiError.Error(), token.Key)
})
}
}
func TestResponsesWSRequestRunnerUsesExistingMemoryRateLimit(t *testing.T) {
require.NoError(t, i18n.Init())
_, token := setupResponsesWSRequestTest(t)
setting.ModelRequestRateLimitEnabled = true
setting.ModelRequestRateLimitDurationMinutes = 1
setting.ModelRequestRateLimitCount = 1
setting.ModelRequestRateLimitSuccessCount = 100
runner, _ := newResponsesWSTestRunner(t, token)
called := 0
handle := func(*gin.Context) *types.NewAPIError { called++; return nil }
require.Nil(t, runner(httptest.NewRequest(http.MethodPost, "/v1/responses", nil), "first", handle))
apiError := runner(httptest.NewRequest(http.MethodPost, "/v1/responses", nil), "limited", handle)
require.NotNil(t, apiError)
assert.Equal(t, http.StatusTooManyRequests, apiError.StatusCode)
message := i18n.Translate("", i18n.MsgRateLimitTotalReached, map[string]any{"Minutes": 1, "Max": 1})
assert.Equal(t, common.MessageWithRequestId(message, "limited"), apiError.Error())
assert.Equal(t, 1, called)
}
func TestResponsesWSRequestRunnerSharesRedisSuccessLimitWithHTTP(t *testing.T) {
_, token := setupResponsesWSRequestTest(t)
redisServer := miniredis.RunT(t)
client := redis.NewClient(&redis.Options{Addr: redisServer.Addr()})
t.Cleanup(func() { require.NoError(t, client.Close()) })
common.RedisEnabled = true
common.RDB = client
setting.ModelRequestRateLimitEnabled = true
setting.ModelRequestRateLimitDurationMinutes = 1
setting.ModelRequestRateLimitCount = 0
setting.ModelRequestRateLimitSuccessCount = 1
runner, _ := newResponsesWSTestRunner(t, token)
failed := types.NewErrorWithStatusCode(errors.New("invalid input"), types.ErrorCodeInvalidRequest, http.StatusBadRequest)
assert.Same(t, failed, runner(httptest.NewRequest(http.MethodPost, "/v1/responses", nil), "failed", func(*gin.Context) *types.NewAPIError { return failed }))
require.Nil(t, runner(httptest.NewRequest(http.MethodPost, "/v1/responses", nil), "successful", func(*gin.Context) *types.NewAPIError { return nil }))
engine := gin.New()
called := false
engine.POST("/v1/responses", middleware.TokenAuth(), middleware.ModelRequestRateLimit(), func(c *gin.Context) { called = true })
request := httptest.NewRequest(http.MethodPost, "/v1/responses", nil)
request.Header.Set("Authorization", "Bearer sk-"+token.Key)
response := httptest.NewRecorder()
engine.ServeHTTP(response, request)
assert.Equal(t, http.StatusTooManyRequests, response.Code)
assert.False(t, called)
}
type responsesWSBillingTest struct {
httpUpstream http.HandlerFunc
httpDone chan struct{}
user *model.User
token *model.Token
channel *model.Channel
client *websocket.Conn
done chan struct{}
upstreamDone chan struct{}
connections atomic.Int32
gatewayURL string
}
func (fixture *responsesWSBillingTest) closeAndWait(t *testing.T) {
t.Helper()
_ = fixture.client.Close()
select {
case <-fixture.done:
case <-time.After(3 * time.Second):
t.Error("gateway handler did not stop after closing its client")
}
if fixture.connections.Load() > 0 {
select {
case <-fixture.upstreamDone:
case <-time.After(3 * time.Second):
t.Error("upstream connection was not closed")
}
}
// Health classification is synchronous at the request boundary; only the
// Redis write is asynchronous, see waitPerfCounters.
}
// waitPerfCounters waits for the asynchronous health samples of the fixture
// model to reach wantRequests and returns the request and success counters.
func waitPerfCounters(t *testing.T, wantRequests int64) (requests, successes int64) {
t.Helper()
require.Eventually(t, func() bool {
requests, successes = 0, 0
keys, err := common.RDB.Keys(context.Background(), "perf:ws-billing:*").Result()
if err != nil {
return false
}
for _, key := range keys {
n, _ := common.RDB.HGet(context.Background(), key, "req").Int64()
requests += n
n, _ = common.RDB.HGet(context.Background(), key, "ok").Int64()
successes += n
}
return requests >= wantRequests
}, 3*time.Second, 10*time.Millisecond)
return requests, successes
}
func newResponsesWSBillingTest(t *testing.T, expression string, handle func(*websocket.Conn, *http.Request), httpEvents ...string) *responsesWSBillingTest {
t.Helper()
user, token := setupResponsesWSRequestTest(t)
previousTimeout := constant.StreamingTimeout
constant.StreamingTimeout = 30
t.Cleanup(func() { constant.StreamingTimeout = previousTimeout })
previousBatch, previousLogs, previousCount, previousQuota := common.BatchUpdateEnabled, common.LogConsumeEnabled, constant.CountToken, common.QuotaPerUnit
common.BatchUpdateEnabled, common.LogConsumeEnabled, constant.CountToken, common.QuotaPerUnit = false, true, false, 500000
saved := map[string]string{}
require.NoError(t, config.GlobalConfig.SaveToDB(func(key, value string) error {
switch key {
case "billing_setting.billing_mode", "billing_setting.billing_expr", "group_ratio_setting.group_ratio", "perf_metrics_setting.enabled":
saved[key] = value
}
return nil
}))
t.Cleanup(func() {
common.BatchUpdateEnabled, common.LogConsumeEnabled, constant.CountToken, common.QuotaPerUnit = previousBatch, previousLogs, previousCount, previousQuota
require.NoError(t, config.GlobalConfig.LoadFromDB(saved))
})
expressions, err := common.Marshal(map[string]string{"ws-billing": expression})
require.NoError(t, err)
require.NoError(t, config.GlobalConfig.LoadFromDB(map[string]string{
"billing_setting.billing_mode": `{"ws-billing":"tiered_expr"}`,
"billing_setting.billing_expr": string(expressions),
"group_ratio_setting.group_ratio": `{"default":1}`,
"perf_metrics_setting.enabled": "true",
}))
require.NoError(t, model.DB.AutoMigrate(&model.Channel{}, &model.Ability{}))
require.NoError(t, model.LOG_DB.AutoMigrate(&model.Log{}))
require.NoError(t, model.DB.Model(user).Updates(map[string]any{"quota": 100000, "setting": `{"billing_preference":"wallet_only"}`}).Error)
require.NoError(t, model.DB.Model(token).Update("remain_quota", 3000).Error)
fixture := &responsesWSBillingTest{user: user, token: token, done: make(chan struct{}), upstreamDone: make(chan struct{}), httpDone: make(chan struct{}, 1)}
redisServer := miniredis.RunT(t)
redisClient := redis.NewClient(&redis.Options{Addr: redisServer.Addr()})
common.RDB, common.RedisEnabled = redisClient, true
t.Cleanup(func() { require.NoError(t, redisClient.Close()) })
var upstreamClosed sync.Once
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if !websocket.IsWebSocketUpgrade(r) {
if fixture.httpUpstream != nil {
fixture.httpUpstream(w, r)
return
}
w.Header().Set("Content-Type", "text/event-stream")
for _, event := range httpEvents {
_, _ = fmt.Fprintf(w, "data: %s\n\n", event)
}
_, _ = fmt.Fprint(w, "data: [DONE]\n\n")
return
}
ws, err := (&websocket.Upgrader{}).Upgrade(w, r, nil)
if !assert.NoError(t, err) {
return
}
fixture.connections.Add(1)
defer ws.Close()
defer upstreamClosed.Do(func() { close(fixture.upstreamDone) })
handle(ws, r)
}))
t.Cleanup(upstream.Close)
channel := &model.Channel{Name: "responses-ws-upstream", Key: "upstream-first\nupstream-second", Status: common.ChannelStatusEnabled, Type: constant.ChannelTypeOpenAI, Group: "default", Models: "ws-billing", BaseURL: &upstream.URL,
ChannelInfo: model.ChannelInfo{IsMultiKey: true, MultiKeySize: 2, MultiKeyMode: constant.MultiKeyModePolling}}
modelMapping := `{"ws-billing":"gpt-4o"}`
channel.ModelMapping = &modelMapping
channel.SetSetting(dto.ChannelSettings{ResponsesWebSocketEnabled: true})
channel.SetOtherSettings(dto.ChannelOtherSettings{AllowServiceTier: true})
require.NoError(t, model.DB.Create(channel).Error)
fixture.channel = channel
require.NoError(t, model.DB.Create(&model.Ability{ChannelId: channel.Id, Model: "ws-billing", Group: "default", Enabled: true}).Error)
t.Cleanup(func() {
require.NoError(t, model.LOG_DB.Where("token_id = ?", token.Id).Delete(&model.Log{}).Error)
require.NoError(t, model.DB.Where("channel_id = ?", channel.Id).Delete(&model.Ability{}).Error)
require.NoError(t, model.DB.Delete(channel).Error)
})
engine := gin.New()
engine.GET("/v1/responses", middleware.TokenAuth(), func(c *gin.Context) {
defer close(fixture.done)
c.Set(common.RequestIdKey, "responses-ws-billing")
ResponsesWebSocket(c)
})
engine.POST("/v1/responses", middleware.TokenAuth(), middleware.ModelRequestRateLimit(), middleware.Distribute(), func(c *gin.Context) {
defer func() { fixture.httpDone <- struct{}{} }()
c.Set(common.RequestIdKey, "responses-http-billing")
Relay(c, types.RelayFormatOpenAIResponses)
})
gateway := httptest.NewServer(engine)
fixture.gatewayURL = gateway.URL
t.Cleanup(gateway.Close)
client, _, err := websocket.DefaultDialer.Dial("ws"+strings.TrimPrefix(gateway.URL, "http")+"/v1/responses", http.Header{"Authorization": []string{"Bearer sk-" + token.Key}})
require.NoError(t, err)
fixture.client = client
t.Cleanup(func() { fixture.closeAndWait(t) })
return fixture
}
func TestResponsesInterruptedStreamHealth(t *testing.T) {
for _, scenario := range []string{"websocket timeout", "sse timeout", "sse client cancel"} {
t.Run(scenario, func(t *testing.T) {
events := []string{`{"type":"response.created","response":{"id":"partial","status":"in_progress"}}`, `{"type":"response.output_text.delta","delta":"hello"}`}
fixture := newResponsesWSBillingTest(t, `tier("request", fixed(0.002))`, func(ws *websocket.Conn, _ *http.Request) {
if _, _, err := ws.ReadMessage(); !assert.NoError(t, err) {
return
}
for _, event := range events {
if !assert.NoError(t, ws.WriteMessage(websocket.TextMessage, []byte(event))) {
return
}
}
_, _, _ = ws.ReadMessage()
})
constant.StreamingTimeout = 1
fixture.httpUpstream = func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "text/event-stream")
for _, event := range events {
_, _ = fmt.Fprintf(w, "data: %s\n\n", event)
}
w.(http.Flusher).Flush()
<-r.Context().Done()
}
if scenario == "websocket timeout" {
require.NoError(t, fixture.client.WriteMessage(websocket.TextMessage, []byte(`{"type":"response.create","model":"ws-billing","input":"hello"}`)))
assert.Equal(t, "response.created", readResponsesWSTestEvent(t, fixture.client)["type"])
assert.Equal(t, "response.output_text.delta", readResponsesWSTestEvent(t, fixture.client)["type"])
_, _, err := fixture.client.ReadMessage()
require.Error(t, err, "idle timeout closes the incomplete stream")
} else {
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
request, err := http.NewRequestWithContext(ctx, http.MethodPost, fixture.gatewayURL+"/v1/responses", strings.NewReader(`{"model":"ws-billing","input":"hello","stream":true}`))
require.NoError(t, err)
request.Header.Set("Authorization", "Bearer sk-"+fixture.token.Key)
request.Header.Set("Content-Type", "application/json")
response, err := (&http.Client{Timeout: 3 * time.Second}).Do(request)
require.NoError(t, err)
if scenario == "sse client cancel" {
reader := bufio.NewReader(response.Body)
for {
line, err := reader.ReadString('\n')
require.NoError(t, err)
if strings.Contains(line, "response.output_text.delta") {
break
}
}
cancel()
} else {
_, err = io.Copy(io.Discard, response.Body)
require.NoError(t, err)
}
require.NoError(t, response.Body.Close())
select {
case <-fixture.httpDone:
case <-time.After(3 * time.Second):
t.Fatal("interrupted SSE request did not finish")
}
}
fixture.closeAndWait(t)
assertResponsesWSAccounting(t, fixture, []int{1000})
if scenario == "sse client cancel" {
keys, err := common.RDB.Keys(context.Background(), "perf:ws-billing:*").Result()
require.NoError(t, err)
assert.Empty(t, keys, "client cancellation must not affect model health")
return
}
requests, successes := waitPerfCounters(t, 1)
assert.Equal(t, int64(1), requests)
assert.Zero(t, successes)
})
}
}
func readResponsesWSTestEvent(t *testing.T, client *websocket.Conn) map[string]any {
t.Helper()
require.NoError(t, client.SetReadDeadline(time.Now().Add(3*time.Second)))
_, data, err := client.ReadMessage()
require.NoError(t, err)
var event map[string]any
require.NoError(t, common.Unmarshal(data, &event))
return event
}
func assertResponsesWSAccounting(t *testing.T, fixture *responsesWSBillingTest, expectedQuotas []int) {
t.Helper()
var logs []model.Log
require.NoError(t, model.LOG_DB.Where("type = ?", model.LogTypeConsume).Order("id").Find(&logs).Error)
require.Len(t, logs, len(expectedQuotas))
var charged int
for index, quota := range expectedQuotas {
assert.Equal(t, quota, logs[index].Quota)
charged += quota
}
require.NoError(t, model.DB.First(fixture.token, fixture.token.Id).Error)
require.NoError(t, model.DB.First(fixture.user, fixture.user.Id).Error)
assert.Equal(t, 3000-charged, fixture.token.RemainQuota)
assert.Equal(t, charged, fixture.token.UsedQuota)
assert.Equal(t, 100000-charged, fixture.user.Quota)
assert.Equal(t, charged, fixture.user.UsedQuota)
}
func TestResponsesWebSocketReusesConnectionAndSettlesEachRequest(t *testing.T) {
type upstreamRequest struct {
Authorization string
Type string `json:"type"`
Model string `json:"model"`
PreviousResponseID string `json:"previous_response_id"`
ServiceTier string `json:"service_tier"`
}
received := make(chan upstreamRequest, 3)
fixture := newResponsesWSBillingTest(t, `param("service_tier") == "priority" ? tier("priority", p * 4) : tier("base", p * 2)`, func(ws *websocket.Conn, r *http.Request) {
for index := 1; ; index++ {
_, data, err := ws.ReadMessage()
if err != nil {
return
}
var event upstreamRequest
if !assert.NoError(t, common.Unmarshal(data, &event)) {
return
}
event.Authorization = r.Header.Get("Authorization")
received <- event
if index == 2 {
duplicate := `{"type":"response.completed","response":{"id":"resp_1","status":"completed","usage":{"input_tokens":1000,"output_tokens":10,"total_tokens":1010}}}`
if !assert.NoError(t, ws.WriteMessage(websocket.TextMessage, []byte(duplicate))) {
return
}
}
terminal := fmt.Sprintf(`{"type":"response.completed","response":{"id":"resp_%d","status":"completed","model":"ws-billing","output":[],"usage":{"input_tokens":1000,"output_tokens":10,"total_tokens":1010}}}`, index)
if !assert.NoError(t, ws.WriteMessage(websocket.TextMessage, []byte(terminal))) {
return
}
}
})
client, user, token := fixture.client, fixture.user, fixture.token
for index, payload := range []string{
`{"type":"response.create","model":"ws-billing","input":"first","service_tier":"default"}`,
`{"type":"response.create","model":"ws-billing","input":"second","previous_response_id":"resp_1","service_tier":"priority"}`,
} {
require.NoError(t, client.WriteMessage(websocket.TextMessage, []byte(payload)))
require.NoError(t, client.SetReadDeadline(time.Now().Add(3*time.Second)))
_, data, err := client.ReadMessage()
require.NoError(t, err)
var terminal struct {
Type string `json:"type"`
Response struct {
ID string `json:"id"`
} `json:"response"`
}
require.NoError(t, common.Unmarshal(data, &terminal))
require.Equal(t, "response.completed", terminal.Type, "unexpected response: %s", data)
assert.Equal(t, fmt.Sprintf("resp_%d", index+1), terminal.Response.ID)
observed := <-received
assert.Equal(t, "Bearer upstream-first", observed.Authorization)
assert.Equal(t, "gpt-4o", observed.Model)
assert.Equal(t, "response.create", observed.Type)
if index == 0 {
assert.Empty(t, observed.PreviousResponseID)
assert.Equal(t, "default", observed.ServiceTier)
} else {
assert.Equal(t, "resp_1", observed.PreviousResponseID)
assert.Equal(t, "priority", observed.ServiceTier)
}
}
require.NoError(t, client.WriteMessage(websocket.TextMessage, []byte(`{"type":"response.create","model":"ws-billing","input":"exhausted"}`)))
require.NoError(t, client.SetReadDeadline(time.Now().Add(3*time.Second)))
_, data, err := client.ReadMessage()
require.NoError(t, err)
var rejection struct {
Type string `json:"type"`
Status int `json:"status"`
}
require.NoError(t, common.Unmarshal(data, &rejection))
assert.Equal(t, "error", rejection.Type)
assert.Equal(t, http.StatusUnauthorized, rejection.Status)
assert.Equal(t, int32(1), fixture.connections.Load())
assert.Empty(t, received, "exhausted token must be rejected before contacting upstream")
var logs []model.Log
require.NoError(t, model.LOG_DB.Where("type = ?", model.LogTypeConsume).Order("id").Find(&logs).Error)
require.Len(t, logs, 2)
for index, expectedQuota := range []int{1000, 2000} {
assert.Equal(t, expectedQuota, logs[index].Quota)
assert.Equal(t, fmt.Sprintf("responses-ws-billing-ws-%d", index), logs[index].RequestId)
assert.Equal(t, 1000, logs[index].PromptTokens)
assert.Equal(t, 10, logs[index].CompletionTokens)
}
// Routing decisions are persisted per request: the first create selects a
// channel, later creates on the same connection reuse it as a pin.
for index, wantDecision := range []string{"attempt:channel_selected", "select:pinned_channel"} {
var other struct {
AdminInfo struct {
RequestPolicy []struct {
ChannelID int `json:"channel_id"`
Decision struct {
Action string `json:"action"`
Reason string `json:"reason"`
} `json:"decision"`
} `json:"request_policy"`
} `json:"admin_info"`
}
require.NoError(t, common.UnmarshalJsonStr(logs[index].Other, &other))
var decisions []string
for _, event := range other.AdminInfo.RequestPolicy {
assert.Equal(t, fixture.channel.Id, event.ChannelID)
decisions = append(decisions, event.Decision.Action+":"+event.Decision.Reason)
}
assert.Contains(t, decisions, wantDecision, "request %d policy events: %v", index, decisions)
}
require.NoError(t, model.DB.First(token, token.Id).Error)
require.NoError(t, model.DB.First(user, user.Id).Error)
assert.Zero(t, token.RemainQuota)
assert.Equal(t, 3000, token.UsedQuota)
assert.Equal(t, 97000, user.Quota)
assert.Equal(t, 3000, user.UsedQuota)
}
// Both transports must reach the same upstream target with the same
// credential placement and settle identically for every supported channel type.
func TestResponsesWebSocketDialsNativeResponsesChannelTypes(t *testing.T) {
type upstreamTarget struct {
Path, Authorization, QueryKey string
}
queryAuthRoute := dto.ChannelOtherSettings{AdvancedCustom: &dto.AdvancedCustomConfig{Routes: []dto.AdvancedCustomRoute{{
IncomingPath: "/v1/responses", UpstreamPath: "/upstream/responses",
Auth: &dto.AdvancedCustomRouteAuth{Type: dto.AdvancedCustomAuthTypeQuery, Name: "api_key", Value: "{api_key}"},
}}}}
bearer := func(key string) upstreamTarget {
return upstreamTarget{Path: "/v1/responses", Authorization: "Bearer " + key}
}
for _, tc := range []struct {
name string
channelType int
otherSettings dto.ChannelOtherSettings
want func(key string) upstreamTarget
}{
{name: "new api", channelType: constant.ChannelTypeNewAPI, want: bearer},
{name: "sub2api", channelType: constant.ChannelTypeSub2API, want: bearer},
{name: "advanced custom query auth", channelType: constant.ChannelTypeAdvancedCustom, otherSettings: queryAuthRoute, want: func(key string) upstreamTarget {
return upstreamTarget{Path: "/upstream/responses", QueryKey: key}
}},
} {
t.Run(tc.name, func(t *testing.T) {
const terminal = `{"type":"response.completed","response":{"id":"first","status":"completed","usage":{"input_tokens":1000,"output_tokens":10,"total_tokens":1010}}}`
targets := make(chan upstreamTarget, 2)
observe := func(r *http.Request) {
targets <- upstreamTarget{Path: r.URL.Path, Authorization: r.Header.Get("Authorization"), QueryKey: r.URL.Query().Get("api_key")}
}
fixture := newResponsesWSBillingTest(t, `tier("request", fixed(0.002))`, func(ws *websocket.Conn, r *http.Request) {
observe(r)
if _, _, err := ws.ReadMessage(); !assert.NoError(t, err) {
return
}
if !assert.NoError(t, ws.WriteMessage(websocket.TextMessage, []byte(terminal))) {
return
}
_, _, _ = ws.ReadMessage()
})
fixture.httpUpstream = func(w http.ResponseWriter, r *http.Request) {
observe(r)
w.Header().Set("Content-Type", "text/event-stream")
_, _ = fmt.Fprintf(w, "data: %s\n\ndata: [DONE]\n\n", terminal)
}
fixture.channel.SetOtherSettings(tc.otherSettings)
require.NoError(t, model.DB.Model(fixture.channel).Updates(map[string]any{"type": tc.channelType, "settings": fixture.channel.OtherSettings}).Error)
require.NoError(t, fixture.client.WriteMessage(websocket.TextMessage, []byte(`{"type":"response.create","model":"ws-billing","input":"hi"}`)))
assert.Equal(t, "response.completed", readResponsesWSTestEvent(t, fixture.client)["type"])
assert.Equal(t, tc.want("upstream-first"), <-targets)
request, err := http.NewRequest(http.MethodPost, fixture.gatewayURL+"/v1/responses", strings.NewReader(`{"model":"ws-billing","input":"hi","stream":true}`))
require.NoError(t, err)
request.Header.Set("Authorization", "Bearer sk-"+fixture.token.Key)
request.Header.Set("Content-Type", "application/json")
response, err := (&http.Client{Timeout: 3 * time.Second}).Do(request)
require.NoError(t, err)
body, err := io.ReadAll(response.Body)
require.NoError(t, err)
require.NoError(t, response.Body.Close())
assert.Equal(t, http.StatusOK, response.StatusCode)
assert.Contains(t, string(body), terminal)
select {
case <-fixture.httpDone:
case <-time.After(3 * time.Second):
t.Fatal("HTTP request did not finish")
}
// The polling multi-key channel rotates to its second key for the HTTP request.
assert.Equal(t, tc.want("upstream-second"), <-targets)
fixture.closeAndWait(t)
assertResponsesWSAccounting(t, fixture, []int{1000, 1000})
})
}
}
func TestResponsesWebSocketDisconnectSettlesDeliveredOutputOnce(t *testing.T) {
fixture := newResponsesWSBillingTest(t, `tier("output", c * 2)`, func(ws *websocket.Conn, _ *http.Request) {
if _, _, err := ws.ReadMessage(); !assert.NoError(t, err) {
return
}
for _, event := range []string{
`{"type":"response.created","response":{"id":"partial","status":"in_progress"}}`,
`{"type":"response.output_text.delta","delta":"hello"}`,
} {
if !assert.NoError(t, ws.WriteMessage(websocket.TextMessage, []byte(event))) {
return
}
}
_, _, _ = ws.ReadMessage()
})
require.NoError(t, fixture.client.WriteMessage(websocket.TextMessage, []byte(`{"type":"response.create","model":"ws-billing","input":"hi","max_output_tokens":1}`)))
assert.Equal(t, "response.created", readResponsesWSTestEvent(t, fixture.client)["type"])
delta := readResponsesWSTestEvent(t, fixture.client)
require.Equal(t, "response.output_text.delta", delta["type"])
assert.Equal(t, "hello", delta["delta"])
fixture.closeAndWait(t)
assertResponsesWSAccounting(t, fixture, []int{1})
keys, err := common.RDB.Keys(context.Background(), "perf:ws-billing:*").Result()
require.NoError(t, err)
assert.Empty(t, keys, "client cancellation must not affect model health")
}
func TestResponsesWebSocketCancelErrorDoesNotFinishActiveRequest(t *testing.T) {
complete := make(chan struct{})
defer close(complete)
allowCreated := make(chan struct{}, 1)
defer close(allowCreated)
receivedCreate := make(chan struct{}, 1)
fixture := newResponsesWSBillingTest(t, `tier("base", p * 2)`, func(ws *websocket.Conn, _ *http.Request) {
_, create, err := ws.ReadMessage()
if !assert.NoError(t, err) || !assert.Contains(t, string(create), `"type":"response.create"`) {
return
}
receivedCreate <- struct{}{}
<-allowCreated
if !assert.NoError(t, ws.WriteMessage(websocket.TextMessage, []byte(`{"type":"response.created","response":{"id":"active","status":"in_progress"}}`))) {
return
}
_, cancel, err := ws.ReadMessage()
if !assert.NoError(t, err) || !assert.JSONEq(t, `{"type":"response.cancel","response_id":"wrong"}`, string(cancel)) {
return
}
if !assert.NoError(t, ws.WriteMessage(websocket.TextMessage, []byte(`{"type":"error","status":400,"error":{"type":"invalid_request_error","code":"response_not_found","message":"No such response"}}`))) {
return
}
<-complete
if !assert.NoError(t, ws.WriteMessage(websocket.TextMessage, []byte(`{"type":"response.completed","response":{"id":"active","status":"completed","usage":{"input_tokens":1000,"output_tokens":10,"total_tokens":1010}}}`))) {
return
}
_, _, _ = ws.ReadMessage()
})
require.NoError(t, fixture.client.WriteMessage(websocket.TextMessage, []byte(`{"type":"response.create","model":"ws-billing","input":"hi"}`)))
select {
case <-receivedCreate:
case <-time.After(3 * time.Second):
t.Fatal("upstream did not receive the initial request")
}
// Queue cancellation before the upstream accepts the response. The following
// conflict is an acknowledgement that the client loop processed both frames.
require.NoError(t, fixture.client.WriteMessage(websocket.TextMessage, []byte(`{"type":"response.cancel","response_id":"wrong"}`)))
require.NoError(t, fixture.client.WriteMessage(websocket.TextMessage, []byte(`{"type":"response.create","model":"ws-billing","input":"before-acceptance"}`)))
assert.Equal(t, float64(http.StatusConflict), readResponsesWSTestEvent(t, fixture.client)["status"])
allowCreated <- struct{}{}
assert.Equal(t, "response.created", readResponsesWSTestEvent(t, fixture.client)["type"])
cancelError := readResponsesWSTestEvent(t, fixture.client)
assert.Equal(t, "error", cancelError["type"])
assert.Equal(t, float64(http.StatusBadRequest), cancelError["status"])
require.NoError(t, fixture.client.WriteMessage(websocket.TextMessage, []byte(`{"type":"response.create","model":"ws-billing","input":"overlapping","stream_id":"other","event_id":"busy"}`)))
conflict := readResponsesWSTestEvent(t, fixture.client)
assert.Equal(t, "error", conflict["type"])
assert.Equal(t, float64(http.StatusConflict), conflict["status"])
assert.Equal(t, "other", conflict["stream_id"])
assert.Equal(t, "busy", conflict["event_id"])
complete <- struct{}{}
assert.Equal(t, "response.completed", readResponsesWSTestEvent(t, fixture.client)["type"])
fixture.closeAndWait(t)
assertResponsesWSAccounting(t, fixture, []int{1000})
}
func TestResponsesWebSocketInitialUpstreamRejectionRefundsReservation(t *testing.T) {
for _, tc := range []struct {
name, upstream, wantType, wantMessage string
status int
}{
{name: "structured error", upstream: `{"type":"error","response_id":"rejected","status":400,"error":{"type":"invalid_request_error","code":"invalid_input","message":"Invalid input"}}`, status: http.StatusBadRequest, wantType: "invalid_request_error", wantMessage: "Invalid input"},
// A frame without an error object is still reported as a request error.
{name: "bare error", upstream: `{"type":"error","status":500,"message":"upstream rejected"}`, status: http.StatusInternalServerError, wantType: "invalid_request_error", wantMessage: "upstream rejected"},
} {
t.Run(tc.name, func(t *testing.T) {
preConsumed := make(chan int, 1)
tokenID := make(chan int, 1)
// Pre-consume no longer estimates completion tokens, so an output-priced expression reserves nothing to refund.
fixture := newResponsesWSBillingTest(t, `tier("request", fixed(0.002))`, func(ws *websocket.Conn, _ *http.Request) {
if _, _, err := ws.ReadMessage(); !assert.NoError(t, err) {
return
}
var token model.Token
if !assert.NoError(t, model.DB.First(&token, <-tokenID).Error) {
return
}
preConsumed <- token.RemainQuota
if !assert.NoError(t, ws.WriteMessage(websocket.TextMessage, []byte(tc.upstream))) {
return
}
_, _, _ = ws.ReadMessage()
})
tokenID <- fixture.token.Id
// Refund is asynchronous. Observe committed token writes instead of waiting
// a fixed delay or returning while its worker still uses the test database.
updates := make(chan struct{}, 4)
require.NoError(t, model.DB.Callback().Update().After("gorm:commit_or_rollback_transaction").Register("responses-ws-refund", func(tx *gorm.DB) {
if tx.Error == nil && tx.Statement.Table == "tokens" {
select {
case updates <- struct{}{}:
default:
}
}
}))
require.NoError(t, fixture.client.WriteMessage(websocket.TextMessage, []byte(`{"type":"response.create","model":"ws-billing","input":"hi","max_output_tokens":10}`)))
rejection := readResponsesWSTestEvent(t, fixture.client)
assert.Equal(t, "error", rejection["type"])
assert.Equal(t, float64(tc.status), rejection["status"])
rejectionError, _ := rejection["error"].(map[string]any)
assert.Equal(t, tc.wantType, rejectionError["type"])
assert.Equal(t, tc.wantMessage, rejectionError["message"])
assert.Equal(t, 2000, <-preConsumed, "the rejected request reserved quota before contacting upstream")
deadline := time.NewTimer(3 * time.Second)
defer deadline.Stop()
for {
require.NoError(t, model.DB.First(fixture.token, fixture.token.Id).Error)
if fixture.token.RemainQuota == 3000 {
break
}
select {
case <-updates:
case <-deadline.C:
t.Fatal("initial rejection did not refund the token reservation")
}
}
fixture.closeAndWait(t)
assertResponsesWSAccounting(t, fixture, nil)
})
}
}
// TEST_RESPONSES_SQL_DSN / TEST_RESPONSES_LOG_SQL_DSN optionally run these
// entry-point regressions against isolated real MySQL/PostgreSQL databases.
func TestResponsesStreamOutcomesPreserveAccounting(t *testing.T) {
for _, transport := range []string{"websocket", "http-sse"} {
for _, tc := range []struct {
name, expression, terminal string
// sseTerminal is the flat error envelope used by the Responses SSE
// protocol; the WebSocket protocol nests the error instead.
sseTerminal string
delta bool
failed bool
ignored bool
}{
{name: "failed-null-fixed", expression: `tier("request", fixed(0.002))`, terminal: `{"type":"response.failed","response":{"id":"first","status":"failed","usage":null,"error":{"code":"server_error","message":"sensitive upstream detail"}}}`, failed: true},
{name: "failed-missing-fixed", expression: `tier("request", fixed(0.002))`, terminal: `{"type":"response.failed","response":{"id":"first","status":"failed"}}`, failed: true},
{name: "failed-actual-usage", expression: `tier("input", p * 2)`, terminal: `{"type":"response.failed","response":{"id":"first","status":"failed","usage":{"input_tokens":1000,"output_tokens":10,"total_tokens":1010}}}`, failed: true},
{name: "failed-estimated-output", expression: `tier("output", c * 2000)`, terminal: `{"type":"response.failed","response":{"id":"first","status":"failed","usage":null}}`, delta: true, failed: true},
{name: "error-after-created", expression: `tier("request", fixed(0.002))`, terminal: `{"type":"error","status":500,"error":{"type":"server_error","code":"server_error","message":"Internal server error"}}`, sseTerminal: `{"type":"error","code":"server_error","message":"Internal server error","param":null,"sequence_number":2}`, failed: true},
{name: "business-error-after-created", expression: `tier("request", fixed(0.002))`, terminal: `{"type":"error","status":400,"error":{"type":"invalid_request_error","code":"context_length_exceeded","message":"Input too long"}}`, sseTerminal: `{"type":"error","code":"context_length_exceeded","message":"Input too long","param":null,"sequence_number":2}`, failed: true, ignored: true},
{name: "completed-at-output-limit", expression: `tier("request", fixed(0.002))`, terminal: `{"type":"response.incomplete","response":{"id":"first","status":"incomplete","incomplete_details":{"reason":"max_output_tokens"},"usage":{"input_tokens":1000,"output_tokens":1,"total_tokens":1001}}}`},
{name: "completed-zero-fixed", expression: `tier("request", fixed(0.002))`, terminal: `{"type":"response.completed","response":{"id":"first","status":"completed","usage":{"input_tokens":0,"output_tokens":0,"total_tokens":0}}}`},
} {
t.Run(transport+"/"+tc.name, func(t *testing.T) {
events := []string{`{"type":"response.created","response":{"id":"first","status":"in_progress"}}`}
if tc.delta {
events = append(events, `{"type":"response.output_text.delta","delta":"hello"}`)
}
terminal := tc.terminal
if transport == "http-sse" && tc.sseTerminal != "" {
terminal = tc.sseTerminal
}
events = append(events, terminal)
fixture := newResponsesWSBillingTest(t, tc.expression, func(ws *websocket.Conn, _ *http.Request) {
for turn := 0; ; turn++ {
_, request, err := ws.ReadMessage()
if err != nil {
return
}
var body map[string]any
if !assert.NoError(t, common.Unmarshal(request, &body)) {
return
}
assert.Equal(t, "planner", body["stream_id"])
assert.Equal(t, "gpt-4o", body["model"])
output := events
if transport == "http-sse" || turn > 0 {
output = []string{`{"type":"response.completed","stream_id":"planner","response":{"id":"second","status":"completed","usage":{"input_tokens":1000,"output_tokens":1,"total_tokens":1001}}}`}
}
for _, event := range output {
if !assert.NoError(t, ws.WriteMessage(websocket.TextMessage, []byte(event))) {
return
}
}
}
}, events...)
setting.ModelRequestRateLimitEnabled = true
setting.ModelRequestRateLimitDurationMinutes = 1
setting.ModelRequestRateLimitSuccessCount = 1
setting.ModelRequestRateLimitCount = 0
if transport == "http-sse" {
request, err := http.NewRequest(http.MethodPost, fixture.gatewayURL+"/v1/responses", strings.NewReader(`{"model":"ws-billing","input":"hi","stream":true,"max_output_tokens":1}`))
require.NoError(t, err)
request.Header.Set("Authorization", "Bearer sk-"+fixture.token.Key)
request.Header.Set("Content-Type", "application/json")
response, err := (&http.Client{Timeout: 3 * time.Second}).Do(request)
require.NoError(t, err)
body, err := io.ReadAll(response.Body)
require.NoError(t, err)
require.NoError(t, response.Body.Close())
assert.Equal(t, http.StatusOK, response.StatusCode)
assert.Contains(t, string(body), terminal)
} else {
require.NoError(t, fixture.client.WriteMessage(websocket.TextMessage, []byte(`{"type":"response.create","stream_id":"planner","model":"ws-billing","input":"hi","max_output_tokens":1}`)))
for _, expected := range events {
actual, err := common.Marshal(readResponsesWSTestEvent(t, fixture.client))
require.NoError(t, err)
assert.JSONEq(t, expected, string(actual), "upstream frames must be delivered once and unchanged")
}
}
quotas := []int{1000}
if tc.failed {
// The first terminal must complete settlement and middleware before
// admitting this immediately following request.
require.NoError(t, fixture.client.WriteMessage(websocket.TextMessage, []byte(`{"type":"response.create","stream_id":"planner","model":"ws-billing","input":"next","max_output_tokens":1}`)))
assert.Equal(t, "response.completed", readResponsesWSTestEvent(t, fixture.client)["type"])
quotas = append(quotas, 1000)
}
require.NoError(t, fixture.client.WriteMessage(websocket.TextMessage, []byte(`{"type":"response.create","event_id":"limited","stream_id":"planner","model":"ws-billing","input":"limited"}`)))
limited := readResponsesWSTestEvent(t, fixture.client)
assert.Equal(t, float64(http.StatusTooManyRequests), limited["status"])
assert.Equal(t, "planner", limited["stream_id"])
assert.Equal(t, "limited", limited["event_id"])
fixture.closeAndWait(t)
assertResponsesWSAccounting(t, fixture, quotas)
var logs []model.Log
require.NoError(t, model.LOG_DB.Where("token_id = ? AND type = ?", fixture.token.Id, model.LogTypeConsume).Order("id").Find(&logs).Error)
require.Len(t, logs, len(quotas))
var other map[string]any
require.NoError(t, common.UnmarshalJsonStr(logs[0].Other, &other))
stream := other["stream_status"].(map[string]any)
if tc.failed {
assert.Equal(t, "error", stream["status"])
assert.Equal(t, "failed", stream["response_status"])
} else {
assert.Equal(t, "ok", stream["status"])
if tc.name == "completed-at-output-limit" {
assert.Equal(t, "incomplete", stream["response_status"])
} else {
assert.Equal(t, "completed", stream["response_status"])
}
}
assert.NotContains(t, logs[0].Other, "sensitive upstream detail")
expectedRequests := int64(len(quotas))
if tc.ignored {
expectedRequests--
}
requests, successes := waitPerfCounters(t, expectedRequests)
assert.Equal(t, expectedRequests, requests)
assert.Equal(t, int64(1), successes)
})
}
}
}
func TestResponsesHTTPHealthCountsFinalResult(t *testing.T) {
for _, tc := range []struct {
name, code string
firstStatus, attempts int
success bool
ignored bool
}{
{"business rejection", "context_length_exceeded", 400, 1, false, true},
{"credentials rejected as 400", "invalid_api_key", 400, 1, false, false},
{"retry succeeds", "server_error", 500, 2, true, false},
} {
t.Run(tc.name, func(t *testing.T) {
fixture := newResponsesWSBillingTest(t, `tier("request", fixed(0.002))`, func(*websocket.Conn, *http.Request) {})
oldRetries := common.RetryTimes
common.RetryTimes = 1
t.Cleanup(func() { common.RetryTimes = oldRetries })
var attempts atomic.Int64
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "application/json")
if attempts.Add(1) == 1 {
w.WriteHeader(tc.firstStatus)
_, _ = fmt.Fprintf(w, `{"error":{"type":"%s","code":"%s","message":"test rejection"}}`, tc.code, tc.code)
return
}
_, _ = fmt.Fprint(w, `{"id":"completed","status":"completed","usage":{"input_tokens":1000,"output_tokens":1,"total_tokens":1001}}`)
}))
t.Cleanup(upstream.Close)
require.NoError(t, model.DB.Model(&model.Channel{}).Where("name = ?", "responses-ws-upstream").Update("base_url", upstream.URL).Error)
request, err := http.NewRequest(http.MethodPost, fixture.gatewayURL+"/v1/responses", strings.NewReader(`{"model":"ws-billing","input":"hello"}`))
require.NoError(t, err)
request.Header.Set("Authorization", "Bearer sk-"+fixture.token.Key)
request.Header.Set("Content-Type", "application/json")
response, err := http.DefaultClient.Do(request)
require.NoError(t, err)
_, err = io.Copy(io.Discard, response.Body)
require.NoError(t, err)
require.NoError(t, response.Body.Close())
select {
case <-fixture.httpDone:
case <-time.After(3 * time.Second):
t.Fatal("HTTP request did not finish")
}
fixture.closeAndWait(t)
assert.Equal(t, int64(tc.attempts), attempts.Load())
if !tc.success {
assert.Equal(t, tc.firstStatus, response.StatusCode)
} else {
assert.Equal(t, http.StatusOK, response.StatusCode)
}
if tc.ignored {
keys, err := common.RDB.Keys(context.Background(), "perf:ws-billing:*").Result()
require.NoError(t, err)
assert.Empty(t, keys)
return
}
requests, successes := waitPerfCounters(t, 1)
assert.Equal(t, int64(1), requests)
if tc.success {
assert.Equal(t, int64(1), successes)
} else {
assert.Zero(t, successes)
}
})
}
}
func TestResponsesWebSocketLocalErrorsKeepStreamIdentity(t *testing.T) {
fixture := newResponsesWSBillingTest(t, `tier("request", fixed(0.002))`, func(*websocket.Conn, *http.Request) {
t.Error("local rejection reached upstream")
})
require.NoError(t, fixture.client.WriteMessage(websocket.TextMessage, []byte(`{"type":"response.create","event_id":"invalid","stream_id":"planner","model":"ws-billing","max_output_tokens":18446744073709551615}`)))
invalid := readResponsesWSTestEvent(t, fixture.client)
assert.Equal(t, float64(http.StatusBadRequest), invalid["status"])
assert.Equal(t, "planner", invalid["stream_id"])
assert.Equal(t, "invalid", invalid["event_id"])
fixture.token.Status = common.TokenStatusDisabled
require.NoError(t, fixture.token.SelectUpdate())
require.NoError(t, fixture.client.WriteMessage(websocket.TextMessage, []byte(`{"type":"response.create","event_id":"revoked","stream_id":"planner","model":"ws-billing","input":"hi"}`)))
revoked := readResponsesWSTestEvent(t, fixture.client)
assert.Equal(t, float64(http.StatusUnauthorized), revoked["status"])
assert.Equal(t, "planner", revoked["stream_id"])
assert.Equal(t, "revoked", revoked["event_id"])
fixture.closeAndWait(t)
assertResponsesWSAccounting(t, fixture, nil)
}
func TestResponsesWebSocketAmbiguousControlErrorClosesAndSettlesOnce(t *testing.T) {
const terminal = `{"type":"error","status":500,"error":{"type":"server_error","code":"server_error","message":"Internal server error"}}`
fixture := newResponsesWSBillingTest(t, `tier("request", fixed(0.002))`, func(ws *websocket.Conn, _ *http.Request) {
if _, _, err := ws.ReadMessage(); !assert.NoError(t, err) {
return
}
if !assert.NoError(t, ws.WriteMessage(websocket.TextMessage, []byte(`{"type":"response.created","response":{"id":"active","status":"in_progress"}}`))) {
return
}
_, cancel, err := ws.ReadMessage()
if !assert.NoError(t, err) {
return
}
assert.Contains(t, string(cancel), `"type":"response.cancel"`)
if !assert.NoError(t, ws.WriteMessage(websocket.TextMessage, []byte(terminal))) {
return
}
_, _, _ = ws.ReadMessage()
})
require.NoError(t, fixture.client.WriteMessage(websocket.TextMessage, []byte(`{"type":"response.create","model":"ws-billing","input":"hi"}`)))
assert.Equal(t, "response.created", readResponsesWSTestEvent(t, fixture.client)["type"])
require.NoError(t, fixture.client.WriteMessage(websocket.TextMessage, []byte(`{"type":"response.cancel","response_id":"active"}`)))
body, err := common.Marshal(readResponsesWSTestEvent(t, fixture.client))
require.NoError(t, err)
assert.JSONEq(t, terminal, string(body))
_, _, err = fixture.client.ReadMessage()
// A dropped socket would read as a synthetic 1006; 1011 is the gateway's frame.
var closeErr *websocket.CloseError
require.ErrorAs(t, err, &closeErr, "an unattributed error must close the connection instead of remaining busy")
assert.Equal(t, websocket.CloseInternalServerErr, closeErr.Code)
fixture.closeAndWait(t)
assertResponsesWSAccounting(t, fixture, []int{1000})
}
// Every way a session ends sends the client a close frame whose code says
// why. A dropped socket would read as a synthetic 1006 close error, so a
// received 1000, 1001, 1008 or 1011 proves the gateway sent the frame.
func TestResponsesWebSocketCloseCodes(t *testing.T) {
require.NoError(t, i18n.Init())
const (
created = `{"type":"response.created","response":{"id":"active","status":"in_progress"}}`
delta = `{"type":"response.output_text.delta","delta":"hello"}`
completed = `{"type":"response.completed","response":{"id":"active","status":"completed","usage":{"input_tokens":1000,"output_tokens":10,"total_tokens":1010}}}`
)
for _, tc := range []struct {
name string
events []string // upstream events after the create, all relayed
// After the events, and after the client read them with waitClient,
// upstream sends upstreamClose or drops the socket. Otherwise it reads
// until the gateway closes the connection.
upstreamClose int
drop, waitClient bool
timeout bool
end func(*testing.T, *responsesWSBillingTest)
code int
reason string
upstreamCode int // the close code upstream receives, when checked
}{
{name: "upstream closes after completion", events: []string{completed}, upstreamClose: websocket.CloseNormalClosure, code: websocket.CloseNormalClosure, reason: "upstream connection closed"},
{name: "idle upstream goes away", events: []string{completed}, waitClient: true, upstreamClose: websocket.CloseGoingAway, code: websocket.CloseGoingAway, reason: "upstream connection closed"},
{name: "idle upstream drops", events: []string{completed}, waitClient: true, drop: true, code: websocket.CloseInternalServerErr, reason: "upstream connection closed"},
// A normal upstream close still truncates the active response.
{name: "upstream closes mid-response", events: []string{created, delta}, upstreamClose: websocket.CloseNormalClosure, code: websocket.CloseInternalServerErr, reason: "upstream connection closed"},
{name: "upstream drops mid-response", events: []string{created, delta}, drop: true, code: websocket.CloseInternalServerErr, reason: "upstream connection closed"},
{name: "upstream response times out", events: []string{created}, timeout: true, code: websocket.CloseInternalServerErr, reason: "upstream response timed out", upstreamCode: websocket.CloseInternalServerErr},
{name: "client closes", events: []string{completed}, end: func(t *testing.T, fixture *responsesWSBillingTest) {
require.NoError(t, fixture.client.WriteControl(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""), time.Now().Add(time.Second)))
}, code: websocket.CloseNormalClosure, upstreamCode: websocket.CloseNormalClosure},
{name: "channel disabled", events: []string{completed}, end: func(t *testing.T, fixture *responsesWSBillingTest) {
require.Equal(t, 1, service.CloseActiveWebSocketsForChannel(fixture.channel.Id, service.ChannelDisabledCloseReason))
}, code: websocket.ClosePolicyViolation, reason: service.ChannelDisabledCloseReason, upstreamCode: websocket.ClosePolicyViolation},
{name: "token revoked", events: []string{completed}, end: func(t *testing.T, fixture *responsesWSBillingTest) {
fixture.token.Status = common.TokenStatusDisabled
require.NoError(t, fixture.token.SelectUpdate())
require.NoError(t, fixture.client.WriteMessage(websocket.TextMessage, []byte(`{"type":"response.create","model":"ws-billing","input":"again"}`)))
assert.Equal(t, float64(http.StatusUnauthorized), readResponsesWSTestEvent(t, fixture.client)["status"])
}, code: websocket.ClosePolicyViolation, upstreamCode: websocket.ClosePolicyViolation},
} {
t.Run(tc.name, func(t *testing.T) {
clientRead := make(chan struct{})
upstreamCodes := make(chan int, 2)
fixture := newResponsesWSBillingTest(t, `tier("request", fixed(0.002))`, func(ws *websocket.Conn, _ *http.Request) {
defaultClose := ws.CloseHandler()
ws.SetCloseHandler(func(code int, text string) error {
upstreamCodes <- code
return defaultClose(code, text)
})
if _, _, err := ws.ReadMessage(); !assert.NoError(t, err) {
return
}
for _, event := range tc.events {
if !assert.NoError(t, ws.WriteMessage(websocket.TextMessage, []byte(event))) {
return
}
}
if tc.waitClient {
<-clientRead
}
if tc.drop {
return
}
if tc.upstreamClose != 0 && !assert.NoError(t, ws.WriteControl(websocket.CloseMessage, websocket.FormatCloseMessage(tc.upstreamClose, "upstream shutdown"), time.Now().Add(time.Second))) {
return
}
for {
if _, _, err := ws.ReadMessage(); err != nil {
return
}
}
})
if tc.timeout {
constant.StreamingTimeout = 1
}
require.NoError(t, fixture.client.WriteMessage(websocket.TextMessage, []byte(`{"type":"response.create","model":"ws-billing","input":"hi"}`)))
for _, event := range tc.events {
actual, err := common.Marshal(readResponsesWSTestEvent(t, fixture.client))
require.NoError(t, err)
assert.JSONEq(t, event, string(actual))
}
close(clientRead)
if tc.end != nil {
tc.end(t, fixture)
}
require.NoError(t, fixture.client.SetReadDeadline(time.Now().Add(5*time.Second)))
_, _, err := fixture.client.ReadMessage()
var closeErr *websocket.CloseError
require.ErrorAs(t, err, &closeErr)
assert.Equal(t, tc.code, closeErr.Code)
assert.Equal(t, tc.reason, closeErr.Text)
fixture.closeAndWait(t)
assertResponsesWSAccounting(t, fixture, []int{1000})
if tc.upstreamCode != 0 {
require.Len(t, upstreamCodes, 1)
assert.Equal(t, tc.upstreamCode, <-upstreamCodes)
}
})
}
}
// A control event queued while another is still outstanding is rejected
// locally instead of reaching upstream twice.
func TestResponsesWebSocketRejectsSecondControlWhileOnePending(t *testing.T) {
receivedCancel := make(chan struct{}, 1)
release := make(chan struct{})
defer close(release)
fixture := newResponsesWSBillingTest(t, `tier("request", fixed(0.002))`, func(ws *websocket.Conn, _ *http.Request) {
if _, _, err := ws.ReadMessage(); !assert.NoError(t, err) {
return
}
if !assert.NoError(t, ws.WriteMessage(websocket.TextMessage, []byte(`{"type":"response.created","response":{"id":"active","status":"in_progress"}}`))) {
return
}
_, cancel, err := ws.ReadMessage()
if !assert.NoError(t, err) || !assert.Contains(t, string(cancel), `"event_id":"first-cancel"`) {
return
}
receivedCancel <- struct{}{}
<-release
if !assert.NoError(t, ws.WriteMessage(websocket.TextMessage, []byte(`{"type":"response.completed","response":{"id":"active","status":"completed","usage":{"input_tokens":1000,"output_tokens":10,"total_tokens":1010}}}`))) {
return
}
_, _, _ = ws.ReadMessage()
})
require.NoError(t, fixture.client.WriteMessage(websocket.TextMessage, []byte(`{"type":"response.create","model":"ws-billing","input":"hi"}`)))
assert.Equal(t, "response.created", readResponsesWSTestEvent(t, fixture.client)["type"])
require.NoError(t, fixture.client.WriteMessage(websocket.TextMessage, []byte(`{"type":"response.cancel","event_id":"first-cancel","response_id":"active"}`)))
select {
case <-receivedCancel:
case <-time.After(3 * time.Second):
t.Fatal("upstream did not receive the first cancel")
}
require.NoError(t, fixture.client.WriteMessage(websocket.TextMessage, []byte(`{"type":"response.cancel","event_id":"second-cancel","response_id":"active"}`)))
rejected := readResponsesWSTestEvent(t, fixture.client)
assert.Equal(t, "error", rejected["type"])
assert.Equal(t, float64(http.StatusBadRequest), rejected["status"])
assert.Equal(t, "second-cancel", rejected["event_id"])
rejectedError, _ := rejected["error"].(map[string]any)
assert.Contains(t, rejectedError["message"], "already pending")
release <- struct{}{}
assert.Equal(t, "response.completed", readResponsesWSTestEvent(t, fixture.client)["type"])
fixture.closeAndWait(t)
assertResponsesWSAccounting(t, fixture, []int{1000})
}
// A control error resolves the outstanding control, so a corrected cancel is
// forwarded instead of being rejected as still pending.
func TestResponsesWebSocketAcceptsCancelAfterControlError(t *testing.T) {
fixture := newResponsesWSBillingTest(t, `tier("request", fixed(0.002))`, func(ws *websocket.Conn, _ *http.Request) {
if _, _, err := ws.ReadMessage(); !assert.NoError(t, err) {
return
}
if !assert.NoError(t, ws.WriteMessage(websocket.TextMessage, []byte(`{"type":"response.created","response":{"id":"active","status":"in_progress"}}`))) {
return
}
_, cancel, err := ws.ReadMessage()
if !assert.NoError(t, err) || !assert.JSONEq(t, `{"type":"response.cancel","response_id":"wrong"}`, string(cancel)) {
return
}
if !assert.NoError(t, ws.WriteMessage(websocket.TextMessage, []byte(`{"type":"error","status":400,"error":{"type":"invalid_request_error","code":"response_not_found","message":"No such response"}}`))) {
return
}
_, cancel, err = ws.ReadMessage()
if !assert.NoError(t, err) || !assert.JSONEq(t, `{"type":"response.cancel","response_id":"active"}`, string(cancel)) {
return
}
if !assert.NoError(t, ws.WriteMessage(websocket.TextMessage, []byte(`{"type":"response.cancelled","response":{"id":"active","status":"cancelled","usage":{"input_tokens":1000,"output_tokens":0,"total_tokens":1000}}}`))) {
return
}
_, _, _ = ws.ReadMessage()
})
require.NoError(t, fixture.client.WriteMessage(websocket.TextMessage, []byte(`{"type":"response.create","model":"ws-billing","input":"hi"}`)))
assert.Equal(t, "response.created", readResponsesWSTestEvent(t, fixture.client)["type"])
require.NoError(t, fixture.client.WriteMessage(websocket.TextMessage, []byte(`{"type":"response.cancel","response_id":"wrong"}`)))
cancelError := readResponsesWSTestEvent(t, fixture.client)
assert.Equal(t, "error", cancelError["type"])
assert.Equal(t, float64(http.StatusBadRequest), cancelError["status"])
require.NoError(t, fixture.client.WriteMessage(websocket.TextMessage, []byte(`{"type":"response.cancel","response_id":"active"}`)))
assert.Equal(t, "response.cancelled", readResponsesWSTestEvent(t, fixture.client)["type"])
fixture.closeAndWait(t)
assertResponsesWSAccounting(t, fixture, []int{1000})
}
// A late error that names the previous response must not end the request
// that follows it on the same connection, whether that response completed or
// failed.
func TestResponsesWebSocketIgnoresLateErrorForPreviousResponse(t *testing.T) {
for _, tc := range []struct {
name, firstTerminal, firstTerminalType string
}{
{name: "after completed", firstTerminalType: "response.completed", firstTerminal: `{"type":"response.completed","response":{"id":"first","status":"completed","usage":{"input_tokens":1000,"output_tokens":10,"total_tokens":1010}}}`},
{name: "after error", firstTerminalType: "error", firstTerminal: `{"type":"error","response_id":"first","status":500,"error":{"type":"server_error","code":"server_error","message":"Internal server error"}}`},
} {
t.Run(tc.name, func(t *testing.T) {
fixture := newResponsesWSBillingTest(t, `tier("request", fixed(0.002))`, func(ws *websocket.Conn, _ *http.Request) {
for _, turn := range [][]string{
{`{"type":"response.created","response":{"id":"first","status":"in_progress"}}`, tc.firstTerminal},
{
`{"type":"error","response_id":"first","status":500,"error":{"type":"server_error","code":"server_error","message":"late failure"}}`,
`{"type":"response.created","response":{"id":"second","status":"in_progress"}}`,
`{"type":"response.completed","response":{"id":"second","status":"completed","usage":{"input_tokens":1000,"output_tokens":10,"total_tokens":1010}}}`,
},
} {
if _, _, err := ws.ReadMessage(); !assert.NoError(t, err) {
return
}
for _, event := range turn {
if !assert.NoError(t, ws.WriteMessage(websocket.TextMessage, []byte(event))) {
return
}
}
}
_, _, _ = ws.ReadMessage()
})
require.NoError(t, fixture.client.WriteMessage(websocket.TextMessage, []byte(`{"type":"response.create","model":"ws-billing","input":"first"}`)))
assert.Equal(t, "response.created", readResponsesWSTestEvent(t, fixture.client)["type"])
assert.Equal(t, tc.firstTerminalType, readResponsesWSTestEvent(t, fixture.client)["type"])
require.NoError(t, fixture.client.WriteMessage(websocket.TextMessage, []byte(`{"type":"response.create","model":"ws-billing","input":"second"}`)))
created := readResponsesWSTestEvent(t, fixture.client)
require.Equal(t, "response.created", created["type"], "the late error must not reach the client or end the request: %v", created)
assert.Equal(t, "response.completed", readResponsesWSTestEvent(t, fixture.client)["type"])
fixture.closeAndWait(t)
assertResponsesWSAccounting(t, fixture, []int{1000, 1000})
})
}
}
// An error that names a different response than the active one is relayed to
// the client and leaves the active generation running.
func TestResponsesWebSocketForwardsErrorForOtherResponseWithoutEndingRequest(t *testing.T) {
const strayError = `{"type":"error","response_id":"other","status":500,"error":{"type":"server_error","code":"server_error","message":"unrelated failure"}}`
fixture := newResponsesWSBillingTest(t, `tier("request", fixed(0.002))`, func(ws *websocket.Conn, _ *http.Request) {
if _, _, err := ws.ReadMessage(); !assert.NoError(t, err) {
return
}
for _, event := range []string{
`{"type":"response.created","response":{"id":"active","status":"in_progress"}}`,
strayError,
`{"type":"response.completed","response":{"id":"active","status":"completed","usage":{"input_tokens":1000,"output_tokens":10,"total_tokens":1010}}}`,
} {
if !assert.NoError(t, ws.WriteMessage(websocket.TextMessage, []byte(event))) {
return
}
}
_, _, _ = ws.ReadMessage()
})
require.NoError(t, fixture.client.WriteMessage(websocket.TextMessage, []byte(`{"type":"response.create","model":"ws-billing","input":"hi"}`)))
assert.Equal(t, "response.created", readResponsesWSTestEvent(t, fixture.client)["type"])
stray, err := common.Marshal(readResponsesWSTestEvent(t, fixture.client))
require.NoError(t, err)
assert.JSONEq(t, strayError, string(stray))
assert.Equal(t, "response.completed", readResponsesWSTestEvent(t, fixture.client)["type"])
fixture.closeAndWait(t)
assertResponsesWSAccounting(t, fixture, []int{1000})
var logs []model.Log
require.NoError(t, model.LOG_DB.Where("type = ?", model.LogTypeConsume).Find(&logs).Error)
require.Len(t, logs, 1)
var other map[string]any
require.NoError(t, common.UnmarshalJsonStr(logs[0].Other, &other))
stream, _ := other["stream_status"].(map[string]any)
assert.Equal(t, "ok", stream["status"])
assert.Equal(t, "completed", stream["response_status"])
}
// Requests rejected before a RelayInfo exists reach the model health sampler
// with the HTTP classification: routing failures count against the model,
// business rejections such as a token model limit do not.
func TestResponsesWebSocketPreRoutingRejectionsFollowHealthClassification(t *testing.T) {
require.NoError(t, i18n.Init())
for _, tc := range []struct {
name string
reject func(*testing.T, *responsesWSBillingTest)
status int
sampled bool
}{
{name: "no eligible channel", status: http.StatusServiceUnavailable, sampled: true, reject: func(t *testing.T, fixture *responsesWSBillingTest) {
require.NoError(t, model.DB.Model(&model.Ability{}).Where("channel_id = ?", fixture.channel.Id).Update("enabled", false).Error)
}},
{name: "token model limit", status: http.StatusForbidden, reject: func(t *testing.T, fixture *responsesWSBillingTest) {
// Update through the model so the cached token is invalidated.
fixture.token.ModelLimitsEnabled = true
fixture.token.ModelLimits = "other-model"
require.NoError(t, fixture.token.Update())
}},
} {
t.Run(tc.name, func(t *testing.T) {
fixture := newResponsesWSBillingTest(t, `tier("request", fixed(0.002))`, func(*websocket.Conn, *http.Request) {
t.Error("rejected request reached upstream")
})
tc.reject(t, fixture)
require.NoError(t, fixture.client.WriteMessage(websocket.TextMessage, []byte(`{"type":"response.create","model":"ws-billing","input":"hi"}`)))
rejected := readResponsesWSTestEvent(t, fixture.client)
assert.Equal(t, "error", rejected["type"])
assert.Equal(t, float64(tc.status), rejected["status"])
fixture.closeAndWait(t)
assertResponsesWSAccounting(t, fixture, nil)
if !tc.sampled {
keys, err := common.RDB.Keys(context.Background(), "perf:ws-billing:*").Result()
require.NoError(t, err)
assert.Empty(t, keys, "a business rejection must not affect model health")
return
}
requests, successes := waitPerfCounters(t, 1)
assert.Equal(t, int64(1), requests)
assert.Zero(t, successes)
})
}
}