2026-02-07 22:14:07 +08:00
package handler
import (
2026-02-28 15:01:20 +08:00
"context"
2026-02-07 22:14:07 +08:00
"encoding/json"
2026-02-28 15:01:20 +08:00
"errors"
2026-07-23 20:17:11 +08:00
"fmt"
2026-07-15 03:43:50 +08:00
"io"
2026-02-07 22:14:07 +08:00
"net/http"
"net/http/httptest"
"strings"
2026-05-28 20:05:24 +08:00
"sync"
2026-07-13 15:32:42 +08:00
"sync/atomic"
2026-02-07 22:14:07 +08:00
"testing"
2026-02-28 15:01:20 +08:00
"time"
2026-02-07 22:14:07 +08:00
2026-05-03 19:33:09 +08:00
"github.com/Wei-Shaw/sub2api/internal/config"
2026-10-01 11:27:35 +08:00
"github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey"
2026-02-28 15:01:20 +08:00
pkghttputil "github.com/Wei-Shaw/sub2api/internal/pkg/httputil"
2026-05-07 09:01:48 +08:00
"github.com/Wei-Shaw/sub2api/internal/pkg/pagination"
2026-08-07 16:45:25 +08:00
"github.com/Wei-Shaw/sub2api/internal/pkg/xai"
2026-02-28 15:01:20 +08:00
"github.com/Wei-Shaw/sub2api/internal/server/middleware"
"github.com/Wei-Shaw/sub2api/internal/service"
coderws "github.com/coder/websocket"
2026-02-07 22:14:07 +08:00
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
2026-02-10 08:59:30 +08:00
"github.com/tidwall/gjson"
"github.com/tidwall/sjson"
2026-02-07 22:14:07 +08:00
)
func TestOpenAIHandleStreamingAwareError_JSONEscaping ( t * testing . T ) {
tests := [] struct {
name string
errType string
message string
}{
{
name : "包含双引号的消息" ,
errType : "server_error" ,
message : `upstream returned "invalid" response` ,
},
{
name : "包含反斜杠的消息" ,
errType : "server_error" ,
message : `path C:\Users\test\file.txt not found` ,
},
{
name : "包含双引号和反斜杠的消息" ,
errType : "upstream_error" ,
message : `error parsing "key\value": unexpected token` ,
},
{
name : "包含换行符的消息" ,
errType : "server_error" ,
message : "line1\nline2\ttab" ,
},
{
name : "普通消息" ,
errType : "upstream_error" ,
message : "Upstream service temporarily unavailable" ,
},
}
for _ , tt := range tests {
t . Run ( tt . name , func ( t * testing . T ) {
gin . SetMode ( gin . TestMode )
w := httptest . NewRecorder ()
c , _ := gin . CreateTestContext ( w )
c . Request = httptest . NewRequest ( http . MethodGet , "/" , nil )
h := & OpenAIGatewayHandler {}
h . handleStreamingAwareError ( c , http . StatusBadGateway , tt . errType , tt . message , true )
body := w . Body . String ()
// 验证 SSE 格式:event: error\ndata: {JSON}\n\n
assert . True ( t , strings . HasPrefix ( body , "event: error\n" ), "应以 'event: error\\n' 开头" )
assert . True ( t , strings . HasSuffix ( body , "\n\n" ), "应以 '\\n\\n' 结尾" )
// 提取 data 部分
lines := strings . Split ( strings . TrimSuffix ( body , "\n\n" ), "\n" )
require . Len ( t , lines , 2 , "应有 event 行和 data 行" )
dataLine := lines [ 1 ]
require . True ( t , strings . HasPrefix ( dataLine , "data: " ), "第二行应以 'data: ' 开头" )
jsonStr := strings . TrimPrefix ( dataLine , "data: " )
// 验证 JSON 合法性
var parsed map [ string ] any
err := json . Unmarshal ([] byte ( jsonStr ), & parsed )
require . NoError ( t , err , "JSON 应能被成功解析,原始 JSON: %s" , jsonStr )
// 验证结构
errorObj , ok := parsed [ "error" ].( map [ string ] any )
require . True ( t , ok , "应包含 error 对象" )
assert . Equal ( t , tt . errType , errorObj [ "type" ])
assert . Equal ( t , tt . message , errorObj [ "message" ])
})
}
}
2026-07-18 03:24:04 +08:00
func TestOpenAIHandleStreamingAwareErrorWithCode_EmitsStableClassification ( t * testing . T ) {
gin . SetMode ( gin . TestMode )
w := httptest . NewRecorder ()
c , _ := gin . CreateTestContext ( w )
c . Request = httptest . NewRequest ( http . MethodPost , "/v1/chat/completions" , nil )
h := & OpenAIGatewayHandler {}
h . handleStreamingAwareErrorWithCode (
c ,
http . StatusBadGateway ,
"upstream_error" ,
service . OpenAIUpstreamHTTP2StreamErrorCode ,
"Upstream HTTP/2 stream failed" ,
true ,
true ,
)
body := w . Body . String ()
require . Contains ( t , body , "event: error\n" )
require . Equal ( t , "upstream_error" , gjson . Get ( body [ strings . Index ( body , "{" ):], "error.type" ). String ())
require . Equal ( t , service . OpenAIUpstreamHTTP2StreamErrorCode , gjson . Get ( body [ strings . Index ( body , "{" ):], "error.code" ). String ())
require . NotContains ( t , body , "stream ID" )
streamErr , ok := service . GetOpsStreamError ( c )
require . True ( t , ok )
require . True ( t , streamErr . CountTowardsSLA )
require . Equal ( t , http . StatusBadGateway , streamErr . IntendedStatus )
}
2026-07-15 21:56:13 +08:00
func TestOpenAIForwardSucceededForScheduling ( t * testing . T ) {
require . True ( t , openAIForwardSucceededForScheduling ( nil ))
require . True ( t , openAIForwardSucceededForScheduling ( & service . OpenAIForwardResult {}))
require . True ( t , openAIForwardSucceededForScheduling ( & service . OpenAIForwardResult {
OpenAIWSMode : true ,
UpstreamTerminalEvent : "response.completed" ,
}))
require . False ( t , openAIForwardSucceededForScheduling ( & service . OpenAIForwardResult {
OpenAIWSMode : true ,
UpstreamTerminalEvent : "response.failed" ,
}))
}
2026-07-19 18:46:20 +08:00
func TestOpenAIResponsesRequiredCapability ( t * testing . T ) {
tests := [] struct {
name string
imageIntent bool
platform string
want service . OpenAIEndpointCapability
}{
{
name : "OpenAI explicit image intent requires Responses" ,
imageIntent : true ,
platform : service . PlatformOpenAI ,
want : service . OpenAIEndpointCapabilityResponses ,
},
{
name : "Grok explicit image intent keeps chat capability" ,
imageIntent : true ,
platform : service . PlatformGrok ,
want : service . OpenAIEndpointCapabilityChatCompletions ,
},
{
name : "non-image intent keeps chat capability" ,
platform : service . PlatformOpenAI ,
want : service . OpenAIEndpointCapabilityChatCompletions ,
},
}
for _ , tt := range tests {
t . Run ( tt . name , func ( t * testing . T ) {
require . Equal ( t , tt . want , openAIResponsesRequiredCapability ( tt . imageIntent , tt . platform ))
})
}
}
2026-05-05 19:32:41 +08:00
func TestResolveOpenAIMessagesMetadataSession_DoesNotDerivePromptCacheKey ( t * testing . T ) {
body := [] byte ( `{"model":"claude-sonnet-4-5","metadata":{"user_id":"claude-code-session"},"messages":[{"role":"user","content":"hello"}]}` )
2026-08-28 14:05:51 +08:00
sessionHash , promptCacheKey := resolveOpenAIMessagesMetadataSession ( nil , "" , "" , "claude-sonnet-4-5" , body )
2026-05-05 19:32:41 +08:00
require . NotEmpty ( t , sessionHash )
require . Empty ( t , promptCacheKey )
}
func TestResolveOpenAIMessagesMetadataSession_PreservesExplicitPromptCacheKey ( t * testing . T ) {
body := [] byte ( `{"metadata":{"user_id":"claude-code-session"}}` )
2026-08-28 14:05:51 +08:00
sessionHash , promptCacheKey := resolveOpenAIMessagesMetadataSession ( nil , "" , "explicit-cache" , "claude-sonnet-4-5" , body )
2026-05-05 19:32:41 +08:00
require . NotEmpty ( t , sessionHash )
require . Equal ( t , "explicit-cache" , promptCacheKey )
}
2026-08-28 14:05:51 +08:00
func TestResolveOpenAIMessagesMetadataSession_ClaudeCodeHeaderOverridesContentFallback ( t * testing . T ) {
gin . SetMode ( gin . TestMode )
recorder := httptest . NewRecorder ()
c , _ := gin . CreateTestContext ( recorder )
c . Request = httptest . NewRequest ( http . MethodPost , "/v1/messages" , nil )
c . Request . Header . Set ( "X-Claude-Code-Session-Id" , "claude-session-001" )
body1 := [] byte ( `{"model":"gpt-5.6-sol","system":"parent","messages":[{"role":"user","content":"parent task"}]}` )
body2 := [] byte ( `{"model":"gpt-5.6-sol","system":"subagent","messages":[{"role":"user","content":"child task"}]}` )
contentHash1 := ( & service . OpenAIGatewayService {}). GenerateSessionHash ( c , body1 )
contentHash2 := ( & service . OpenAIGatewayService {}). GenerateSessionHash ( c , body2 )
require . NotEqual ( t , contentHash1 , contentHash2 , "different bodies should prove the content fallback differs" )
hash1 , cacheKey1 := resolveOpenAIMessagesMetadataSession ( c , contentHash1 , "" , "gpt-5.6-sol" , body1 )
hash2 , cacheKey2 := resolveOpenAIMessagesMetadataSession ( c , contentHash2 , "" , "gpt-5.6-sol" , body2 )
require . Equal ( t , service . DeriveSessionHashFromSeed ( "claude-session-001" ), hash1 )
require . Equal ( t , hash1 , hash2 , "the same Claude Code session must keep one sticky account across changed turn bodies" )
require . Empty ( t , cacheKey1 , "routing-only fix must not create an upstream prompt cache key" )
require . Empty ( t , cacheKey2 , "routing-only fix must not create an upstream prompt cache key" )
}
func TestResolveOpenAIMessagesMetadataSession_OpenAISignalWinsOverClaudeHeader ( t * testing . T ) {
gin . SetMode ( gin . TestMode )
recorder := httptest . NewRecorder ()
c , _ := gin . CreateTestContext ( recorder )
c . Request = httptest . NewRequest ( http . MethodPost , "/v1/messages" , nil )
c . Request . Header . Set ( "X-Claude-Code-Session-Id" , "claude-session-001" )
hash , cacheKey := resolveOpenAIMessagesMetadataSession ( c , "content-hash" , "explicit-openai-session" , "gpt-5.6-sol" , [] byte ( `{"metadata":{"user_id":"opaque"}}` ))
require . Equal ( t , "content-hash" , hash , "existing OpenAI session resolution must remain authoritative" )
require . Equal ( t , "explicit-openai-session" , cacheKey )
}
func TestResolveOpenAIMessagesMetadataSession_BlankClaudeHeaderKeepsContentFallback ( t * testing . T ) {
gin . SetMode ( gin . TestMode )
recorder := httptest . NewRecorder ()
c , _ := gin . CreateTestContext ( recorder )
c . Request = httptest . NewRequest ( http . MethodPost , "/v1/messages" , nil )
c . Request . Header . Set ( "X-Claude-Code-Session-Id" , " " )
hash , cacheKey := resolveOpenAIMessagesMetadataSession ( c , "content-hash" , "" , "gpt-5.6-sol" , [] byte ( `{"metadata":{"user_id":"opaque"}}` ))
require . Equal ( t , "content-hash" , hash )
require . Empty ( t , cacheKey )
}
2026-02-07 22:14:07 +08:00
func TestOpenAIHandleStreamingAwareError_NonStreaming ( t * testing . T ) {
gin . SetMode ( gin . TestMode )
w := httptest . NewRecorder ()
c , _ := gin . CreateTestContext ( w )
c . Request = httptest . NewRequest ( http . MethodGet , "/" , nil )
h := & OpenAIGatewayHandler {}
h . handleStreamingAwareError ( c , http . StatusBadGateway , "upstream_error" , "test error" , false )
// 非流式应返回 JSON 响应
assert . Equal ( t , http . StatusBadGateway , w . Code )
var parsed map [ string ] any
err := json . Unmarshal ( w . Body . Bytes (), & parsed )
require . NoError ( t , err )
errorObj , ok := parsed [ "error" ].( map [ string ] any )
require . True ( t , ok )
assert . Equal ( t , "upstream_error" , errorObj [ "type" ])
assert . Equal ( t , "test error" , errorObj [ "message" ])
}
2026-02-10 08:59:30 +08:00
2026-02-28 15:01:20 +08:00
func TestReadRequestBodyWithPrealloc ( t * testing . T ) {
payload := `{"model":"gpt-5","input":"hello"}`
req := httptest . NewRequest ( http . MethodPost , "/v1/responses" , strings . NewReader ( payload ))
req . ContentLength = int64 ( len ( payload ))
body , err := pkghttputil . ReadRequestBodyWithPrealloc ( req )
require . NoError ( t , err )
require . Equal ( t , payload , string ( body ))
}
func TestReadRequestBodyWithPrealloc_MaxBytesError ( t * testing . T ) {
rec := httptest . NewRecorder ()
req := httptest . NewRequest ( http . MethodPost , "/v1/responses" , strings . NewReader ( strings . Repeat ( "x" , 8 )))
req . Body = http . MaxBytesReader ( rec , req . Body , 4 )
_ , err := pkghttputil . ReadRequestBodyWithPrealloc ( req )
require . Error ( t , err )
var maxErr * http . MaxBytesError
require . ErrorAs ( t , err , & maxErr )
}
2026-02-14 11:23:10 +08:00
func TestOpenAIEnsureForwardErrorResponse_WritesFallbackWhenNotWritten ( t * testing . T ) {
gin . SetMode ( gin . TestMode )
w := httptest . NewRecorder ()
c , _ := gin . CreateTestContext ( w )
c . Request = httptest . NewRequest ( http . MethodGet , "/" , nil )
h := & OpenAIGatewayHandler {}
wrote := h . ensureForwardErrorResponse ( c , false )
require . True ( t , wrote )
require . Equal ( t , http . StatusBadGateway , w . Code )
var parsed map [ string ] any
err := json . Unmarshal ( w . Body . Bytes (), & parsed )
require . NoError ( t , err )
errorObj , ok := parsed [ "error" ].( map [ string ] any )
require . True ( t , ok )
assert . Equal ( t , "upstream_error" , errorObj [ "type" ])
assert . Equal ( t , "Upstream request failed" , errorObj [ "message" ])
}
2026-05-24 22:00:56 +08:00
// Writer 已写后 ensureForwardErrorResponse 必须仍然把错误信息以 SSE
// 形式追加给客户端(streamStarted 强制 true)。
// 这是 case B 修复:旧实现遇到 Writer.Written 直接 return false,
// 客户端只能拿到 silent EOF; Codex CLI 报 "stream closed before response.completed"。
func TestOpenAIEnsureForwardErrorResponse_AppendsSSEAfterWritten ( t * testing . T ) {
2026-02-14 11:23:10 +08:00
gin . SetMode ( gin . TestMode )
w := httptest . NewRecorder ()
c , _ := gin . CreateTestContext ( w )
c . Request = httptest . NewRequest ( http . MethodGet , "/" , nil )
c . String ( http . StatusTeapot , "already written" )
h := & OpenAIGatewayHandler {}
wrote := h . ensureForwardErrorResponse ( c , false )
2026-05-24 22:00:56 +08:00
require . True ( t , wrote , "must attempt to communicate the failure to the client via SSE" )
// 状态码改不了(headers 已 flush),但 body 应该追加 SSE 错误事件。
2026-02-14 11:23:10 +08:00
require . Equal ( t , http . StatusTeapot , w . Code )
2026-05-24 22:00:56 +08:00
assert . Contains ( t , w . Body . String (), "already written" )
// 非 /responses 路径走 legacy event: error 分支。
assert . Contains ( t , w . Body . String (), "event: error\n" )
}
// case B 回归测试:/responses 路径,Writer 已被写过(模拟 ping flushed),
// ensureForwardErrorResponse 必须发 response.failed,让 Codex 收到合规终止事件。
func TestOpenAIEnsureForwardErrorResponse_ResponsesRouteAfterWrittenEmitsResponseFailed ( t * testing . T ) {
gin . SetMode ( gin . TestMode )
w := httptest . NewRecorder ()
c , _ := gin . CreateTestContext ( w )
c . Request = httptest . NewRequest ( http . MethodPost , EndpointResponses , nil )
// 模拟 ping 已 flush 的状态:Writer 已写过 1 个字节
_ , _ = c . Writer . WriteString ( ":\n\n" )
h := & OpenAIGatewayHandler {}
wrote := h . ensureForwardErrorResponse ( c , false )
require . True ( t , wrote )
body := w . Body . String ()
assert . Contains ( t , body , ":\n\n" , "earlier ping bytes preserved" )
assert . Contains ( t , body , "event: response.failed\n" , "appended a Responses terminal event" )
assert . Contains ( t , body , `"type":"response.failed"` )
assert . Contains ( t , body , `"code":"upstream_error"` )
assert . Contains ( t , body , "Upstream request failed" )
2026-02-14 11:23:10 +08:00
}
2026-07-15 20:59:56 +08:00
func TestOpenAIEnsureForwardErrorResponse_AfterDeltaAppendsSingleValidResponseFailed ( t * testing . T ) {
gin . SetMode ( gin . TestMode )
w := httptest . NewRecorder ()
c , _ := gin . CreateTestContext ( w )
c . Request = httptest . NewRequest ( http . MethodPost , EndpointResponses , nil )
delta := `{"type":"response.output_text.delta","delta":"ok","sequence_number":1}`
_ , err := c . Writer . WriteString ( "event: response.output_text.delta\ndata: " + delta + "\n\n" )
require . NoError ( t , err )
h := & OpenAIGatewayHandler {}
require . True ( t , h . ensureForwardErrorResponse ( c , true ))
frames := strings . Split ( strings . TrimSuffix ( w . Body . String (), "\n\n" ), "\n\n" )
require . Len ( t , frames , 2 )
errorEvents := 0
for _ , frame := range frames {
lines := strings . Split ( frame , "\n" )
require . Len ( t , lines , 2 )
require . True ( t , strings . HasPrefix ( lines [ 0 ], "event: " ))
require . True ( t , strings . HasPrefix ( lines [ 1 ], "data: " ))
eventType := strings . TrimPrefix ( lines [ 0 ], "event: " )
data := strings . TrimPrefix ( lines [ 1 ], "data: " )
require . True ( t , json . Valid ([] byte ( data )), "each downstream SSE frame must contain valid JSON" )
var event struct {
Type string `json:"type"`
}
require . NoError ( t , json . Unmarshal ([] byte ( data ), & event ))
require . Equal ( t , eventType , event . Type )
if eventType == "response.failed" {
errorEvents ++
}
}
require . Equal ( t , 1 , errorEvents )
}
2026-09-18 12:38:05 +08:00
func TestOpenAIEnsureForwardErrorResponse_PreservesCommittedStreamError ( t * testing . T ) {
gin . SetMode ( gin . TestMode )
w := httptest . NewRecorder ()
c , _ := gin . CreateTestContext ( w )
c . Request = httptest . NewRequest ( http . MethodPost , EndpointResponses , nil )
body := "event: error\ndata: {\"type\":\"error\",\"code\":\"upstream_stream_read_error\",\"message\":\"Upstream response stream was interrupted\",\"param\":null,\"sequence_number\":0}\n\n"
_ , err := c . Writer . WriteString ( body )
require . NoError ( t , err )
service . MarkResponseCommitted ( c )
h := & OpenAIGatewayHandler {}
require . False ( t , h . ensureForwardErrorResponse ( c , true ))
require . Equal ( t , body , w . Body . String (), "preserve the service error without a second terminal event" )
}
2026-08-05 16:28:16 +08:00
func TestOpenAIEnsureForwardErrorResponse_CompactKeepaliveOnlyWritesResponseFailed ( t * testing . T ) {
gin . SetMode ( gin . TestMode )
w := httptest . NewRecorder ()
c , _ := gin . CreateTestContext ( w )
c . Request = httptest . NewRequest ( http . MethodPost , EndpointResponses , nil )
service . MarkOpenAICompactClientStream ( c )
stop := service . StartOpenAICompactSSEKeepalive ( c , 5 * time . Millisecond )
defer stop ()
before := service . OpenAICompactKeepaliveAdjustedWrittenSize ( c )
require . Eventually ( t , c . Writer . Written , time . Second , time . Millisecond )
require . Equal ( t , before , service . OpenAICompactKeepaliveAdjustedWrittenSize ( c ))
h := & OpenAIGatewayHandler {}
require . True ( t , h . ensureForwardErrorResponse ( c , false ))
require . Equal ( t , http . StatusOK , w . Code )
require . Contains ( t , w . Body . String (), "event: response.failed\n" )
require . NotContains ( t , w . Body . String (), "event: error\n" )
}
2026-07-15 02:54:22 +08:00
func TestOpenAIEnsureForwardErrorResponse_ImageJSONKeepaliveWritesSingleJSONFallback ( t * testing . T ) {
gin . SetMode ( gin . TestMode )
w := httptest . NewRecorder ()
c , _ := gin . CreateTestContext ( w )
c . Request = httptest . NewRequest ( http . MethodPost , "/v1/images/generations" , nil )
stop := service . StartOpenAIImagesJSONKeepalive ( c , 5 * time . Millisecond )
defer stop ()
before := service . OpenAIImagesJSONKeepaliveAdjustedWrittenSize ( c )
require . Eventually ( t , c . Writer . Written , time . Second , time . Millisecond )
require . Equal ( t , before , service . OpenAIImagesJSONKeepaliveAdjustedWrittenSize ( c ))
require . False ( t , openAIForwardErrorAlreadyCommunicated ( c , before , errors . New ( "read upstream response: unexpected EOF" )))
h := & OpenAIGatewayHandler {}
wrote := h . ensureForwardErrorResponse ( c , false )
require . True ( t , wrote )
require . Equal ( t , http . StatusOK , w . Code , "heartbeat already committed the status" )
require . True ( t , json . Valid ( w . Body . Bytes ()), w . Body . String ())
require . NotContains ( t , w . Body . String (), "event:" )
require . NotContains ( t , w . Body . String (), "data:" )
decoder := json . NewDecoder ( strings . NewReader ( w . Body . String ()))
var payload map [ string ] any
require . NoError ( t , decoder . Decode ( & payload ))
require . ErrorIs ( t , decoder . Decode ( & payload ), io . EOF )
require . Equal ( t , "upstream_error" , gjson . Get ( w . Body . String (), "error.type" ). String ())
require . Equal ( t , "Upstream request failed" , gjson . Get ( w . Body . String (), "error.message" ). String ())
}
func TestOpenAIEnsureForwardErrorResponse_ImageJSONKeepalivePreservesCompletedJSON ( t * testing . T ) {
gin . SetMode ( gin . TestMode )
w := httptest . NewRecorder ()
c , _ := gin . CreateTestContext ( w )
c . Request = httptest . NewRequest ( http . MethodPost , "/v1/images/generations" , nil )
stop := service . StartOpenAIImagesJSONKeepalive ( c , 5 * time . Millisecond )
defer stop ()
before := service . OpenAIImagesJSONKeepaliveAdjustedWrittenSize ( c )
require . Eventually ( t , c . Writer . Written , time . Second , time . Millisecond )
c . JSON ( http . StatusOK , gin . H { "data" : [] gin . H {{ "b64_json" : "aW1hZ2U=" }}})
completedBody := w . Body . String ()
require . True ( t , json . Valid ([] byte ( completedBody )), completedBody )
require . Greater ( t , service . OpenAIImagesJSONKeepaliveAdjustedWrittenSize ( c ), before )
require . False ( t , openAIForwardErrorAlreadyCommunicated ( c , before , errors . New ( "read upstream trailer: unexpected EOF" )))
h := & OpenAIGatewayHandler {}
wrote := h . ensureForwardErrorResponse ( c , false )
require . False ( t , wrote , "the completed Images JSON already communicated the response" )
require . Equal ( t , completedBody , w . Body . String ())
require . NotContains ( t , w . Body . String (), "event:" )
require . NotContains ( t , w . Body . String (), "data:" )
decoder := json . NewDecoder ( strings . NewReader ( w . Body . String ()))
var payload map [ string ] any
require . NoError ( t , decoder . Decode ( & payload ))
require . ErrorIs ( t , decoder . Decode ( & payload ), io . EOF )
require . Equal ( t , "aW1hZ2U=" , gjson . Get ( w . Body . String (), "data.0.b64_json" ). String ())
}
func TestOpenAIEnsureForwardErrorResponse_FastImageJSONKeepalivePreservesCompletedJSON ( t * testing . T ) {
gin . SetMode ( gin . TestMode )
w := httptest . NewRecorder ()
c , _ := gin . CreateTestContext ( w )
c . Request = httptest . NewRequest ( http . MethodPost , "/v1/images/generations" , nil )
stop := service . StartOpenAIImagesJSONKeepalive ( c , time . Hour )
defer stop ()
before := service . OpenAIImagesJSONKeepaliveAdjustedWrittenSize ( c )
c . JSON ( http . StatusOK , gin . H { "data" : [] gin . H {{ "b64_json" : "ZmFzdC1pbWFnZQ==" }}})
completedBody := w . Body . String ()
require . True ( t , json . Valid ([] byte ( completedBody )), completedBody )
require . Greater ( t , service . OpenAIImagesJSONKeepaliveAdjustedWrittenSize ( c ), before )
require . False ( t , openAIForwardErrorAlreadyCommunicated ( c , before , errors . New ( "read upstream trailer: unexpected EOF" )))
h := & OpenAIGatewayHandler {}
wrote := h . ensureForwardErrorResponse ( c , false )
require . False ( t , wrote , "fast completed Images JSON already communicated the response" )
require . Equal ( t , completedBody , w . Body . String ())
require . NotContains ( t , w . Body . String (), "event:" )
require . NotContains ( t , w . Body . String (), "data:" )
decoder := json . NewDecoder ( strings . NewReader ( w . Body . String ()))
var payload map [ string ] any
require . NoError ( t , decoder . Decode ( & payload ))
require . ErrorIs ( t , decoder . Decode ( & payload ), io . EOF )
require . Equal ( t , "ZmFzdC1pbWFnZQ==" , gjson . Get ( w . Body . String (), "data.0.b64_json" ). String ())
}
2026-02-28 15:01:20 +08:00
func TestShouldLogOpenAIForwardFailureAsWarn ( t * testing . T ) {
gin . SetMode ( gin . TestMode )
t . Run ( "fallback_written_should_not_downgrade" , func ( t * testing . T ) {
w := httptest . NewRecorder ()
c , _ := gin . CreateTestContext ( w )
c . Request = httptest . NewRequest ( http . MethodGet , "/" , nil )
require . False ( t , shouldLogOpenAIForwardFailureAsWarn ( c , true ))
})
t . Run ( "context_nil_should_not_downgrade" , func ( t * testing . T ) {
require . False ( t , shouldLogOpenAIForwardFailureAsWarn ( nil , false ))
})
t . Run ( "response_not_written_should_not_downgrade" , func ( t * testing . T ) {
w := httptest . NewRecorder ()
c , _ := gin . CreateTestContext ( w )
c . Request = httptest . NewRequest ( http . MethodGet , "/" , nil )
require . False ( t , shouldLogOpenAIForwardFailureAsWarn ( c , false ))
})
t . Run ( "response_already_written_should_downgrade" , func ( t * testing . T ) {
w := httptest . NewRecorder ()
c , _ := gin . CreateTestContext ( w )
c . Request = httptest . NewRequest ( http . MethodGet , "/" , nil )
c . String ( http . StatusForbidden , "already written" )
require . True ( t , shouldLogOpenAIForwardFailureAsWarn ( c , false ))
})
}
func TestOpenAIRecoverResponsesPanic_WritesFallbackResponse ( t * testing . T ) {
gin . SetMode ( gin . TestMode )
w := httptest . NewRecorder ()
c , _ := gin . CreateTestContext ( w )
c . Request = httptest . NewRequest ( http . MethodPost , "/v1/responses" , nil )
h := & OpenAIGatewayHandler {}
streamStarted := false
require . NotPanics ( t , func () {
func () {
defer h . recoverResponsesPanic ( c , & streamStarted )
panic ( "test panic" )
}()
})
require . Equal ( t , http . StatusBadGateway , w . Code )
var parsed map [ string ] any
err := json . Unmarshal ( w . Body . Bytes (), & parsed )
require . NoError ( t , err )
errorObj , ok := parsed [ "error" ].( map [ string ] any )
require . True ( t , ok )
assert . Equal ( t , "upstream_error" , errorObj [ "type" ])
assert . Equal ( t , "Upstream request failed" , errorObj [ "message" ])
}
func TestOpenAIRecoverResponsesPanic_NoPanicNoWrite ( t * testing . T ) {
gin . SetMode ( gin . TestMode )
w := httptest . NewRecorder ()
c , _ := gin . CreateTestContext ( w )
c . Request = httptest . NewRequest ( http . MethodPost , "/v1/responses" , nil )
h := & OpenAIGatewayHandler {}
streamStarted := false
require . NotPanics ( t , func () {
func () {
defer h . recoverResponsesPanic ( c , & streamStarted )
}()
})
require . False ( t , c . Writer . Written ())
assert . Equal ( t , "" , w . Body . String ())
}
2026-05-24 22:00:56 +08:00
// Panic 在已 flush 的 /v1/responses 流中:状态码无法改(已 written),
// 但 body 应追加 response.failed 让客户端识别为合规截断而不是 silent EOF。
func TestOpenAIRecoverResponsesPanic_AppendsResponseFailedAfterWritten ( t * testing . T ) {
2026-02-28 15:01:20 +08:00
gin . SetMode ( gin . TestMode )
w := httptest . NewRecorder ()
c , _ := gin . CreateTestContext ( w )
c . Request = httptest . NewRequest ( http . MethodPost , "/v1/responses" , nil )
c . String ( http . StatusTeapot , "already written" )
h := & OpenAIGatewayHandler {}
streamStarted := false
require . NotPanics ( t , func () {
func () {
defer h . recoverResponsesPanic ( c , & streamStarted )
panic ( "test panic" )
}()
})
require . Equal ( t , http . StatusTeapot , w . Code )
2026-05-24 22:00:56 +08:00
body := w . Body . String ()
assert . Contains ( t , body , "already written" )
assert . Contains ( t , body , "event: response.failed\n" )
2026-02-28 15:01:20 +08:00
}
func TestOpenAIMissingResponsesDependencies ( t * testing . T ) {
t . Run ( "nil_handler" , func ( t * testing . T ) {
var h * OpenAIGatewayHandler
require . Equal ( t , [] string { "handler" }, h . missingResponsesDependencies ())
})
t . Run ( "all_dependencies_missing" , func ( t * testing . T ) {
h := & OpenAIGatewayHandler {}
require . Equal ( t ,
[] string { "gatewayService" , "billingCacheService" , "apiKeyService" , "concurrencyHelper" },
h . missingResponsesDependencies (),
)
})
t . Run ( "all_dependencies_present" , func ( t * testing . T ) {
h := & OpenAIGatewayHandler {
gatewayService : & service . OpenAIGatewayService {},
billingCacheService : & service . BillingCacheService {},
apiKeyService : & service . APIKeyService {},
concurrencyHelper : & ConcurrencyHelper {
concurrencyService : & service . ConcurrencyService {},
},
}
require . Empty ( t , h . missingResponsesDependencies ())
})
}
func TestOpenAIEnsureResponsesDependencies ( t * testing . T ) {
t . Run ( "missing_dependencies_returns_503" , func ( t * testing . T ) {
gin . SetMode ( gin . TestMode )
w := httptest . NewRecorder ()
c , _ := gin . CreateTestContext ( w )
c . Request = httptest . NewRequest ( http . MethodPost , "/v1/responses" , nil )
h := & OpenAIGatewayHandler {}
ok := h . ensureResponsesDependencies ( c , nil )
require . False ( t , ok )
require . Equal ( t , http . StatusServiceUnavailable , w . Code )
var parsed map [ string ] any
err := json . Unmarshal ( w . Body . Bytes (), & parsed )
require . NoError ( t , err )
errorObj , exists := parsed [ "error" ].( map [ string ] any )
require . True ( t , exists )
assert . Equal ( t , "api_error" , errorObj [ "type" ])
assert . Equal ( t , "Service temporarily unavailable" , errorObj [ "message" ])
})
t . Run ( "already_written_response_not_overridden" , func ( t * testing . T ) {
gin . SetMode ( gin . TestMode )
w := httptest . NewRecorder ()
c , _ := gin . CreateTestContext ( w )
c . Request = httptest . NewRequest ( http . MethodPost , "/v1/responses" , nil )
c . String ( http . StatusTeapot , "already written" )
h := & OpenAIGatewayHandler {}
ok := h . ensureResponsesDependencies ( c , nil )
require . False ( t , ok )
require . Equal ( t , http . StatusTeapot , w . Code )
assert . Equal ( t , "already written" , w . Body . String ())
})
t . Run ( "dependencies_ready_returns_true_and_no_write" , func ( t * testing . T ) {
gin . SetMode ( gin . TestMode )
w := httptest . NewRecorder ()
c , _ := gin . CreateTestContext ( w )
c . Request = httptest . NewRequest ( http . MethodPost , "/v1/responses" , nil )
h := & OpenAIGatewayHandler {
gatewayService : & service . OpenAIGatewayService {},
billingCacheService : & service . BillingCacheService {},
apiKeyService : & service . APIKeyService {},
concurrencyHelper : & ConcurrencyHelper {
concurrencyService : & service . ConcurrencyService {},
},
}
ok := h . ensureResponsesDependencies ( c , nil )
require . True ( t , ok )
require . False ( t , c . Writer . Written ())
assert . Equal ( t , "" , w . Body . String ())
})
}
2026-04-09 12:29:49 +08:00
func TestResolveOpenAIMessagesDispatchMappedModel ( t * testing . T ) {
t . Run ( "exact_claude_model_override_wins" , func ( t * testing . T ) {
apiKey := & service . APIKey {
Group : & service . Group {
MessagesDispatchModelConfig : service . OpenAIMessagesDispatchModelConfig {
SonnetMappedModel : "gpt-5.2" ,
ExactModelMappings : map [ string ] string {
"claude-sonnet-4-5-20250929" : "gpt-5.4-mini-high" ,
2026-07-13 17:57:08 +08:00
"claude-fable-5" : "gpt-5.6-sol" ,
2026-04-09 12:29:49 +08:00
},
},
},
}
2026-08-19 15:24:26 +08:00
require . Equal ( t , "gpt-5.4-mini" , resolveOpenAIMessagesDispatchMappedModel ( nil , apiKey , "claude-sonnet-4-5-20250929" ))
require . Equal ( t , "gpt-5.6-sol" , resolveOpenAIMessagesDispatchMappedModel ( nil , apiKey , "claude-fable-5" ))
2026-04-09 12:29:49 +08:00
})
t . Run ( "uses_family_default_when_no_override" , func ( t * testing . T ) {
apiKey := & service . APIKey { Group : & service . Group {}}
2026-08-19 15:24:26 +08:00
require . Equal ( t , "gpt-5.4" , resolveOpenAIMessagesDispatchMappedModel ( nil , apiKey , "claude-opus-4-6" ))
require . Equal ( t , "gpt-5.3-codex" , resolveOpenAIMessagesDispatchMappedModel ( nil , apiKey , "claude-sonnet-4-5-20250929" ))
require . Equal ( t , "gpt-5.4-mini" , resolveOpenAIMessagesDispatchMappedModel ( nil , apiKey , "claude-haiku-4-5-20251001" ))
2026-04-09 12:29:49 +08:00
})
t . Run ( "returns_empty_for_non_claude_or_missing_group" , func ( t * testing . T ) {
2026-08-19 15:24:26 +08:00
require . Empty ( t , resolveOpenAIMessagesDispatchMappedModel ( nil , nil , "claude-sonnet-4-5-20250929" ))
require . Empty ( t , resolveOpenAIMessagesDispatchMappedModel ( nil , & service . APIKey {}, "claude-sonnet-4-5-20250929" ))
require . Empty ( t , resolveOpenAIMessagesDispatchMappedModel ( nil , & service . APIKey { Group : & service . Group {}}, "gpt-5.4" ))
2026-04-09 12:29:49 +08:00
})
2026-06-29 18:33:37 +08:00
t . Run ( "grok_group_maps_claude_cli_model_to_grok_default" , func ( t * testing . T ) {
2026-08-07 16:45:25 +08:00
original := xai . RuntimeModelMappingOptions ()
t . Cleanup ( func () { xai . SetRuntimeModelMappingOptions ( original ) })
xai . SetRuntimeModelMappingOptions ( xai . ModelMappingOptions { EnableCrossClientMap : true })
2026-06-29 18:33:37 +08:00
apiKey := & service . APIKey {
Group : & service . Group {
Platform : service . PlatformGrok ,
},
}
2026-08-20 04:26:24 -07:00
require . Equal ( t , "grok-4.6" , resolveOpenAIMessagesDispatchMappedModel ( nil , apiKey , "claude-sonnet-4-5" ))
2026-08-19 15:24:26 +08:00
require . Empty ( t , resolveOpenAIMessagesDispatchMappedModel ( nil , apiKey , "grok" ))
2026-06-29 18:33:37 +08:00
})
2026-04-09 12:29:49 +08:00
t . Run ( "does_not_fall_back_to_group_default_mapped_model" , func ( t * testing . T ) {
apiKey := & service . APIKey {
Group : & service . Group {
DefaultMappedModel : "gpt-5.4" ,
},
}
2026-08-19 15:24:26 +08:00
require . Empty ( t , resolveOpenAIMessagesDispatchMappedModel ( nil , apiKey , "gpt-5.4" ))
require . Equal ( t , "gpt-5.3-codex" , resolveOpenAIMessagesDispatchMappedModel ( nil , apiKey , "claude-sonnet-4-5-20250929" ))
2026-04-09 12:29:49 +08:00
})
}
2026-06-29 18:33:37 +08:00
func TestOpenAIGatewayMessagesDispatchGateAllowsGrokGroups ( t * testing . T ) {
gin . SetMode ( gin . TestMode )
t . Run ( "openai_group_without_dispatch_flag_is_rejected" , func ( t * testing . T ) {
rec := httptest . NewRecorder ()
c , _ := gin . CreateTestContext ( rec )
c . Request = httptest . NewRequest ( http . MethodPost , "/v1/messages" , strings . NewReader ( `{"model":"claude-sonnet-4-5","messages":[{"role":"user","content":"hi"}]}` ))
groupID := int64 ( 4101 )
c . Set ( string ( middleware . ContextKeyAPIKey ), & service . APIKey {
ID : 5101 ,
GroupID : & groupID ,
User : & service . User { ID : 6101 },
Group : & service . Group {
ID : groupID ,
Platform : service . PlatformOpenAI ,
AllowMessagesDispatch : false ,
},
})
c . Set ( string ( middleware . ContextKeyUser ), middleware . AuthSubject { UserID : 6101 , Concurrency : 1 })
h := & OpenAIGatewayHandler {}
h . Messages ( c )
require . Equal ( t , http . StatusForbidden , rec . Code )
require . Equal ( t , "permission_error" , gjson . GetBytes ( rec . Body . Bytes (), "error.type" ). String ())
require . Contains ( t , rec . Body . String (), "This group does not allow /v1/messages dispatch" )
})
t . Run ( "grok_group_without_dispatch_flag_reaches_gateway_dependencies" , func ( t * testing . T ) {
rec := httptest . NewRecorder ()
c , _ := gin . CreateTestContext ( rec )
c . Request = httptest . NewRequest ( http . MethodPost , "/v1/messages" , strings . NewReader ( `{"model":"grok-4.3","messages":[{"role":"user","content":"hi"}]}` ))
groupID := int64 ( 4102 )
c . Set ( string ( middleware . ContextKeyAPIKey ), & service . APIKey {
ID : 5102 ,
GroupID : & groupID ,
User : & service . User { ID : 6102 },
Group : & service . Group {
ID : groupID ,
Platform : service . PlatformGrok ,
AllowMessagesDispatch : false ,
},
})
c . Set ( string ( middleware . ContextKeyUser ), middleware . AuthSubject { UserID : 6102 , Concurrency : 1 })
h := & OpenAIGatewayHandler {}
h . Messages ( c )
require . Equal ( t , http . StatusServiceUnavailable , rec . Code )
require . Equal ( t , "api_error" , gjson . GetBytes ( rec . Body . Bytes (), "error.type" ). String ())
require . NotContains ( t , rec . Body . String (), "This group does not allow /v1/messages dispatch" )
})
}
2026-06-09 21:52:04 +08:00
func TestOpenAIModelMappedBody ( t * testing . T ) {
body := [] byte ( `{"model":"alias","input":"hello"}` )
calls := 0
forwardBody := openAIModelMappedBody ( body , true , "gpt-5.4" , func ( body [] byte , newModel string ) [] byte {
calls ++
return service . ReplaceModelInBody ( body , newModel )
})
require . Equal ( t , 1 , calls )
require . Equal ( t , "gpt-5.4" , gjson . GetBytes ( forwardBody , "model" ). String ())
require . Equal ( t , "alias" , gjson . GetBytes ( body , "model" ). String ())
}
func TestOpenAIModelMappedBodyCache ( t * testing . T ) {
body := [] byte ( `{"model":"alias","input":"hello"}` )
calls := 0
mappedBody := newOpenAIModelMappedBodyCache ( body , func ( body [] byte , newModel string ) [] byte {
calls ++
return service . ReplaceModelInBody ( body , newModel )
})
first := mappedBody ( true , "gpt-5.4" )
second := mappedBody ( true , "gpt-5.4" )
third := mappedBody ( true , "gpt-5.3-codex" )
unmapped := mappedBody ( false , "ignored" )
require . Equal ( t , 2 , calls )
require . Equal ( t , "gpt-5.4" , gjson . GetBytes ( first , "model" ). String ())
require . Equal ( t , "gpt-5.4" , gjson . GetBytes ( second , "model" ). String ())
require . Equal ( t , "gpt-5.3-codex" , gjson . GetBytes ( third , "model" ). String ())
require . Equal ( t , body , unmapped )
require . Same ( t , & first [ 0 ], & second [ 0 ])
}
2026-02-28 15:01:20 +08:00
func TestOpenAIResponses_MissingDependencies_ReturnsServiceUnavailable ( t * testing . T ) {
gin . SetMode ( gin . TestMode )
w := httptest . NewRecorder ()
c , _ := gin . CreateTestContext ( w )
c . Request = httptest . NewRequest ( http . MethodPost , "/v1/responses" , strings . NewReader ( `{"model":"gpt-5","stream":false}` ))
c . Request . Header . Set ( "Content-Type" , "application/json" )
groupID := int64 ( 2 )
c . Set ( string ( middleware . ContextKeyAPIKey ), & service . APIKey {
ID : 10 ,
GroupID : & groupID ,
})
c . Set ( string ( middleware . ContextKeyUser ), middleware . AuthSubject {
UserID : 1 ,
Concurrency : 1 ,
})
// 故意使用未初始化依赖,验证快速失败而不是崩溃。
h := & OpenAIGatewayHandler {}
require . NotPanics ( t , func () {
h . Responses ( c )
})
require . Equal ( t , http . StatusServiceUnavailable , w . Code )
var parsed map [ string ] any
err := json . Unmarshal ( w . Body . Bytes (), & parsed )
require . NoError ( t , err )
errorObj , ok := parsed [ "error" ].( map [ string ] any )
require . True ( t , ok )
assert . Equal ( t , "api_error" , errorObj [ "type" ])
assert . Equal ( t , "Service temporarily unavailable" , errorObj [ "message" ])
}
func TestOpenAIResponses_SetsClientTransportHTTP ( t * testing . T ) {
gin . SetMode ( gin . TestMode )
w := httptest . NewRecorder ()
c , _ := gin . CreateTestContext ( w )
c . Request = httptest . NewRequest ( http . MethodPost , "/openai/v1/responses" , strings . NewReader ( `{"model":"gpt-5"}` ))
c . Request . Header . Set ( "Content-Type" , "application/json" )
h := & OpenAIGatewayHandler {}
h . Responses ( c )
require . Equal ( t , http . StatusUnauthorized , w . Code )
require . Equal ( t , service . OpenAIClientTransportHTTP , service . GetOpenAIClientTransport ( c ))
}
func TestOpenAIResponses_RejectsMessageIDAsPreviousResponseID ( t * testing . T ) {
gin . SetMode ( gin . TestMode )
w := httptest . NewRecorder ()
c , _ := gin . CreateTestContext ( w )
c . Request = httptest . NewRequest ( http . MethodPost , "/openai/v1/responses" , strings . NewReader (
`{"model":"gpt-5.1","stream":false,"previous_response_id":"msg_123456","input":[{"type":"input_text","text":"hello"}]}` ,
))
c . Request . Header . Set ( "Content-Type" , "application/json" )
groupID := int64 ( 2 )
c . Set ( string ( middleware . ContextKeyAPIKey ), & service . APIKey {
ID : 101 ,
GroupID : & groupID ,
User : & service . User { ID : 1 },
})
c . Set ( string ( middleware . ContextKeyUser ), middleware . AuthSubject {
UserID : 1 ,
Concurrency : 1 ,
})
h := newOpenAIHandlerForPreviousResponseIDValidation ( t , nil )
h . Responses ( c )
require . Equal ( t , http . StatusBadRequest , w . Code )
require . Contains ( t , w . Body . String (), "previous_response_id must be a response.id" )
}
2026-08-19 23:13:40 -07:00
func TestOpenAIResponses_AcceptsHTTPContinuationPreviousResponseIDBeforeRouting ( t * testing . T ) {
2026-04-21 13:53:12 +08:00
gin . SetMode ( gin . TestMode )
w := httptest . NewRecorder ()
c , _ := gin . CreateTestContext ( w )
c . Request = httptest . NewRequest ( http . MethodPost , "/openai/v1/responses" , strings . NewReader (
`{"model":"gpt-5.1","stream":false,"previous_response_id":"resp_123456","input":[{"type":"input_text","text":"hello"}]}` ,
))
c . Request . Header . Set ( "Content-Type" , "application/json" )
groupID := int64 ( 2 )
c . Set ( string ( middleware . ContextKeyAPIKey ), & service . APIKey {
ID : 101 ,
GroupID : & groupID ,
User : & service . User { ID : 1 },
})
c . Set ( string ( middleware . ContextKeyUser ), middleware . AuthSubject {
UserID : 1 ,
Concurrency : 1 ,
})
h := newOpenAIHandlerForPreviousResponseIDValidation ( t , nil )
2026-08-20 21:06:48 -07:00
require . NoError ( t , h . gatewayService . BindOpenAIHTTPResponseOwner ( context . Background (), groupID , "resp_123456" , 1 , 101 ))
2026-04-21 13:53:12 +08:00
h . Responses ( c )
2026-08-19 23:13:40 -07:00
require . NotEqual ( t , http . StatusBadRequest , w . Code )
require . NotContains ( t , w . Body . String (), "Responses WebSocket v2" )
2026-04-21 13:53:12 +08:00
}
2026-08-20 21:06:48 -07:00
func TestOpenAIResponses_RejectsHTTPContinuationOwnedByAnotherUser ( t * testing . T ) {
gin . SetMode ( gin . TestMode )
w := httptest . NewRecorder ()
c , _ := gin . CreateTestContext ( w )
c . Request = httptest . NewRequest ( http . MethodPost , "/openai/v1/responses" , strings . NewReader (
`{"model":"gpt-5.1","stream":false,"previous_response_id":"resp_other_tenant","input":"hello"}` ,
))
c . Request . Header . Set ( "Content-Type" , "application/json" )
groupID := int64 ( 2 )
c . Set ( string ( middleware . ContextKeyAPIKey ), & service . APIKey {
ID : 202 ,
UserID : 2 ,
GroupID : & groupID ,
User : & service . User { ID : 2 },
})
c . Set ( string ( middleware . ContextKeyUser ), middleware . AuthSubject { UserID : 2 , Concurrency : 1 })
h := newOpenAIHandlerForPreviousResponseIDValidation ( t , nil )
require . NoError ( t , h . gatewayService . BindOpenAIHTTPResponseOwner ( context . Background (), groupID , "resp_other_tenant" , 1 , 101 ))
h . Responses ( c )
require . Equal ( t , http . StatusBadRequest , w . Code )
require . Contains ( t , w . Body . String (), "previous_response_id is not available for this user" )
}
func TestOpenAIResponses_RejectsUnownedHTTPContinuation ( t * testing . T ) {
gin . SetMode ( gin . TestMode )
w := httptest . NewRecorder ()
c , _ := gin . CreateTestContext ( w )
c . Request = httptest . NewRequest ( http . MethodPost , "/openai/v1/responses" , strings . NewReader (
`{"model":"gpt-5.1","stream":false,"previous_response_id":"resp_unknown","input":"hello"}` ,
))
c . Request . Header . Set ( "Content-Type" , "application/json" )
groupID := int64 ( 2 )
c . Set ( string ( middleware . ContextKeyAPIKey ), & service . APIKey { ID : 101 , UserID : 1 , GroupID : & groupID , User : & service . User { ID : 1 }})
c . Set ( string ( middleware . ContextKeyUser ), middleware . AuthSubject { UserID : 1 , Concurrency : 1 })
h := newOpenAIHandlerForPreviousResponseIDValidation ( t , nil )
h . Responses ( c )
require . Equal ( t , http . StatusBadRequest , w . Code )
require . Contains ( t , w . Body . String (), "previous_response_id is not available for this user" )
}
2026-04-21 13:53:12 +08:00
func TestOpenAIResponses_FunctionCallOutputHTTPGuidanceDoesNotSuggestPreviousResponseReuse ( t * testing . T ) {
gin . SetMode ( gin . TestMode )
w := httptest . NewRecorder ()
c , _ := gin . CreateTestContext ( w )
c . Request = httptest . NewRequest ( http . MethodPost , "/openai/v1/responses" , strings . NewReader (
`{"model":"gpt-5.1","stream":false,"input":[{"type":"function_call_output","output":"{}"}]}` ,
))
c . Request . Header . Set ( "Content-Type" , "application/json" )
groupID := int64 ( 2 )
c . Set ( string ( middleware . ContextKeyAPIKey ), & service . APIKey {
ID : 101 ,
GroupID : & groupID ,
User : & service . User { ID : 1 },
})
c . Set ( string ( middleware . ContextKeyUser ), middleware . AuthSubject {
UserID : 1 ,
Concurrency : 1 ,
})
h := newOpenAIHandlerForPreviousResponseIDValidation ( t , nil )
h . Responses ( c )
require . Equal ( t , http . StatusBadRequest , w . Code )
require . Contains ( t , w . Body . String (), "Responses WebSocket v2" )
require . NotContains ( t , w . Body . String (), "reuse previous_response_id" )
}
2026-02-28 15:01:20 +08:00
func TestOpenAIResponsesWebSocket_SetsClientTransportWSWhenUpgradeValid ( t * testing . T ) {
gin . SetMode ( gin . TestMode )
w := httptest . NewRecorder ()
c , _ := gin . CreateTestContext ( w )
c . Request = httptest . NewRequest ( http . MethodGet , "/openai/v1/responses" , nil )
c . Request . Header . Set ( "Upgrade" , "websocket" )
c . Request . Header . Set ( "Connection" , "Upgrade" )
h := & OpenAIGatewayHandler {}
h . ResponsesWebSocket ( c )
require . Equal ( t , http . StatusUnauthorized , w . Code )
require . Equal ( t , service . OpenAIClientTransportWS , service . GetOpenAIClientTransport ( c ))
}
func TestOpenAIResponsesWebSocket_InvalidUpgradeDoesNotSetTransport ( t * testing . T ) {
gin . SetMode ( gin . TestMode )
w := httptest . NewRecorder ()
c , _ := gin . CreateTestContext ( w )
c . Request = httptest . NewRequest ( http . MethodGet , "/openai/v1/responses" , nil )
h := & OpenAIGatewayHandler {}
h . ResponsesWebSocket ( c )
require . Equal ( t , http . StatusUpgradeRequired , w . Code )
require . Equal ( t , service . OpenAIClientTransportUnknown , service . GetOpenAIClientTransport ( c ))
}
2026-07-13 15:32:42 +08:00
func TestOpenAIResponsesWebSocket_IngressCapacityRejected ( t * testing . T ) {
gin . SetMode ( gin . TestMode )
cache := & concurrencyCacheMock {
acquireIngressLeaseFn : func ( context . Context , int64 , int , string ) ( bool , error ) {
return false , nil
},
}
h := newOpenAIHandlerForPreviousResponseIDValidation ( t , cache )
h . cfg = & config . Config {}
h . cfg . Gateway . OpenAIWS . MaxIngressConnectionsPerAPIKey = 1
wsServer := newOpenAIWSHandlerTestServer ( t , h , middleware . AuthSubject { UserID : 1 , Concurrency : 1 })
defer wsServer . Close ()
dialCtx , cancelDial := context . WithTimeout ( context . Background (), 3 * time . Second )
2026-07-15 22:05:53 +08:00
clientConn , response , err := coderws . Dial ( dialCtx , "ws" + strings . TrimPrefix ( wsServer . URL , "http" ) + "/openai/v1/responses" , nil )
2026-07-13 15:32:42 +08:00
cancelDial ()
2026-07-15 22:05:53 +08:00
require . Error ( t , err )
require . Nil ( t , clientConn )
require . NotNil ( t , response )
require . Equal ( t , http . StatusTooManyRequests , response . StatusCode )
_ = response . Body . Close ()
}
2026-07-13 15:32:42 +08:00
2026-07-15 22:05:53 +08:00
func TestOpenAIResponsesWebSocket_IngressLeaseBackendUnavailableBeforeUpgrade ( t * testing . T ) {
gin . SetMode ( gin . TestMode )
cache := & concurrencyCacheMock {
acquireIngressLeaseFn : func ( context . Context , int64 , int , string ) ( bool , error ) {
return false , errors . New ( "redis unavailable" )
},
}
h := newOpenAIHandlerForPreviousResponseIDValidation ( t , cache )
h . cfg = & config . Config {}
h . cfg . Gateway . OpenAIWS . MaxIngressConnectionsPerAPIKey = 1
wsServer := newOpenAIWSHandlerTestServer ( t , h , middleware . AuthSubject { UserID : 1 , Concurrency : 1 })
defer wsServer . Close ()
dialCtx , cancelDial := context . WithTimeout ( context . Background (), 3 * time . Second )
clientConn , response , err := coderws . Dial ( dialCtx , "ws" + strings . TrimPrefix ( wsServer . URL , "http" ) + "/openai/v1/responses" , nil )
cancelDial ()
require . Error ( t , err )
require . Nil ( t , clientConn )
require . NotNil ( t , response )
require . Equal ( t , http . StatusServiceUnavailable , response . StatusCode )
_ = response . Body . Close ()
2026-07-13 15:32:42 +08:00
}
2026-07-14 22:40:05 +08:00
func TestOpenAIResponsesWebSocket_FirstMessageTimeoutUsesConfig ( t * testing . T ) {
gin . SetMode ( gin . TestMode )
logSink , restore := captureHandlerStructuredLog ( t )
defer restore ()
h := newOpenAIHandlerForPreviousResponseIDValidation ( t , nil )
h . cfg = & config . Config {}
h . cfg . Gateway . OpenAIWS . ClientFirstMessageTimeoutSeconds = 1
wsServer := newOpenAIWSHandlerTestServer ( t , h , middleware . AuthSubject { UserID : 1 , Concurrency : 1 })
defer wsServer . Close ()
dialCtx , cancelDial := context . WithTimeout ( context . Background (), 3 * time . Second )
clientConn , _ , err := coderws . Dial ( dialCtx , "ws" + strings . TrimPrefix ( wsServer . URL , "http" ) + "/openai/v1/responses" , nil )
cancelDial ()
require . NoError ( t , err )
defer func () { _ = clientConn . CloseNow () }()
started := time . Now ()
readCtx , cancelRead := context . WithTimeout ( context . Background (), 5 * time . Second )
_ , _ , err = clientConn . Read ( readCtx )
cancelRead ()
elapsed := time . Since ( started )
require . Error ( t , err )
require . NotErrorIs ( t , err , context . DeadlineExceeded )
require . GreaterOrEqual ( t , elapsed , 500 * time . Millisecond )
require . Less ( t , elapsed , 4 * time . Second )
require . Eventually ( t , func () bool {
readTimeout , ok := logSink . FieldValueForMessage ( "openai.websocket_read_first_message_failed" , "read_timeout" )
return ok && readTimeout == time . Second &&
logSink . ContainsMessageAtLevel ( "openai.websocket_read_first_message_failed" , "warn" )
}, time . Second , 10 * time . Millisecond )
}
2026-07-13 15:32:42 +08:00
func TestOpenAIResponsesWebSocket_IngressLeaseReleasedOnEarlyReturn ( t * testing . T ) {
gin . SetMode ( gin . TestMode )
cache := & concurrencyCacheMock {
acquireIngressLeaseFn : func ( context . Context , int64 , int , string ) ( bool , error ) {
return true , nil
},
}
h := newOpenAIHandlerForPreviousResponseIDValidation ( t , cache )
h . cfg = & config . Config {}
h . cfg . Gateway . OpenAIWS . MaxIngressConnectionsPerAPIKey = 1
wsServer := newOpenAIWSHandlerTestServer ( t , h , middleware . AuthSubject { UserID : 1 , Concurrency : 1 })
defer wsServer . Close ()
dialCtx , cancelDial := context . WithTimeout ( context . Background (), 3 * time . Second )
clientConn , _ , err := coderws . Dial ( dialCtx , "ws" + strings . TrimPrefix ( wsServer . URL , "http" ) + "/openai/v1/responses" , nil )
cancelDial ()
require . NoError ( t , err )
defer func () { _ = clientConn . CloseNow () }()
writeCtx , cancelWrite := context . WithTimeout ( context . Background (), 3 * time . Second )
err = clientConn . Write ( writeCtx , coderws . MessageBinary , [] byte ( "not a response.create frame" ))
cancelWrite ()
require . NoError ( t , err )
readCtx , cancelRead := context . WithTimeout ( context . Background (), 3 * time . Second )
_ , _ , err = clientConn . Read ( readCtx )
cancelRead ()
var closeErr coderws . CloseError
require . ErrorAs ( t , err , & closeErr )
require . Equal ( t , coderws . StatusPolicyViolation , closeErr . Code )
require . Eventually ( t , func () bool {
return atomic . LoadInt32 ( & cache . releaseIngressCalled ) == 1
}, time . Second , 10 * time . Millisecond )
}
2026-07-15 22:05:53 +08:00
func TestOpenAIResponsesWebSocket_IngressLeaseReleasedWhenUpgradeFails ( t * testing . T ) {
gin . SetMode ( gin . TestMode )
cache := & concurrencyCacheMock {
acquireIngressLeaseFn : func ( context . Context , int64 , int , string ) ( bool , error ) {
return true , nil
},
}
h := newOpenAIHandlerForPreviousResponseIDValidation ( t , cache )
h . cfg = & config . Config {}
h . cfg . Gateway . OpenAIWS . MaxIngressConnectionsPerAPIKey = 1
wsServer := newOpenAIWSHandlerTestServer ( t , h , middleware . AuthSubject { UserID : 1 , Concurrency : 1 })
defer wsServer . Close ()
req , err := http . NewRequest ( http . MethodGet , wsServer . URL + "/openai/v1/responses" , nil )
require . NoError ( t , err )
req . Header . Set ( "Upgrade" , "websocket" )
req . Header . Set ( "Connection" , "Upgrade" )
req . Header . Set ( "Sec-WebSocket-Version" , "13" )
resp , err := http . DefaultClient . Do ( req )
require . NoError ( t , err )
_ = resp . Body . Close ()
require . NotEqual ( t , http . StatusSwitchingProtocols , resp . StatusCode )
require . Eventually ( t , func () bool {
return atomic . LoadInt32 ( & cache . releaseIngressCalled ) == 1
}, time . Second , 10 * time . Millisecond )
}
2026-02-28 15:01:20 +08:00
func TestOpenAIResponsesWebSocket_RejectsMessageIDAsPreviousResponseID ( t * testing . T ) {
gin . SetMode ( gin . TestMode )
h := newOpenAIHandlerForPreviousResponseIDValidation ( t , nil )
wsServer := newOpenAIWSHandlerTestServer ( t , h , middleware . AuthSubject { UserID : 1 , Concurrency : 1 })
defer wsServer . Close ()
dialCtx , cancelDial := context . WithTimeout ( context . Background (), 3 * time . Second )
clientConn , _ , err := coderws . Dial ( dialCtx , "ws" + strings . TrimPrefix ( wsServer . URL , "http" ) + "/openai/v1/responses" , nil )
cancelDial ()
require . NoError ( t , err )
defer func () {
_ = clientConn . CloseNow ()
}()
writeCtx , cancelWrite := context . WithTimeout ( context . Background (), 3 * time . Second )
err = clientConn . Write ( writeCtx , coderws . MessageText , [] byte (
`{"type":"response.create","model":"gpt-5.1","stream":false,"previous_response_id":"msg_abc123"}` ,
))
cancelWrite ()
require . NoError ( t , err )
readCtx , cancelRead := context . WithTimeout ( context . Background (), 3 * time . Second )
_ , _ , err = clientConn . Read ( readCtx )
cancelRead ()
require . Error ( t , err )
var closeErr coderws . CloseError
require . ErrorAs ( t , err , & closeErr )
require . Equal ( t , coderws . StatusPolicyViolation , closeErr . Code )
require . Contains ( t , strings . ToLower ( closeErr . Reason ), "previous_response_id" )
}
func TestOpenAIResponsesWebSocket_PreviousResponseIDKindLoggedBeforeAcquireFailure ( t * testing . T ) {
gin . SetMode ( gin . TestMode )
cache := & concurrencyCacheMock {
acquireUserSlotFn : func ( ctx context . Context , userID int64 , maxConcurrency int , requestID string ) ( bool , error ) {
return false , errors . New ( "user slot unavailable" )
},
}
h := newOpenAIHandlerForPreviousResponseIDValidation ( t , cache )
wsServer := newOpenAIWSHandlerTestServer ( t , h , middleware . AuthSubject { UserID : 1 , Concurrency : 1 })
defer wsServer . Close ()
dialCtx , cancelDial := context . WithTimeout ( context . Background (), 3 * time . Second )
clientConn , _ , err := coderws . Dial ( dialCtx , "ws" + strings . TrimPrefix ( wsServer . URL , "http" ) + "/openai/v1/responses" , nil )
cancelDial ()
require . NoError ( t , err )
defer func () {
_ = clientConn . CloseNow ()
}()
writeCtx , cancelWrite := context . WithTimeout ( context . Background (), 3 * time . Second )
err = clientConn . Write ( writeCtx , coderws . MessageText , [] byte (
`{"type":"response.create","model":"gpt-5.1","stream":false,"previous_response_id":"resp_prev_123"}` ,
))
cancelWrite ()
require . NoError ( t , err )
readCtx , cancelRead := context . WithTimeout ( context . Background (), 3 * time . Second )
_ , _ , err = clientConn . Read ( readCtx )
cancelRead ()
require . Error ( t , err )
var closeErr coderws . CloseError
require . ErrorAs ( t , err , & closeErr )
require . Equal ( t , coderws . StatusInternalError , closeErr . Code )
require . Contains ( t , strings . ToLower ( closeErr . Reason ), "failed to acquire user concurrency slot" )
}
2026-05-07 09:01:48 +08:00
type contentModerationHandlerSettingRepo struct {
values map [ string ] string
}
func ( r * contentModerationHandlerSettingRepo ) Get ( ctx context . Context , key string ) ( * service . Setting , error ) {
if value , ok := r . values [ key ]; ok {
return & service . Setting { Key : key , Value : value }, nil
}
return nil , service . ErrSettingNotFound
}
func ( r * contentModerationHandlerSettingRepo ) GetValue ( ctx context . Context , key string ) ( string , error ) {
if value , ok := r . values [ key ]; ok {
return value , nil
}
return "" , service . ErrSettingNotFound
}
func ( r * contentModerationHandlerSettingRepo ) Set ( ctx context . Context , key , value string ) error {
if r . values == nil {
r . values = map [ string ] string {}
}
r . values [ key ] = value
return nil
}
func ( r * contentModerationHandlerSettingRepo ) GetMultiple ( ctx context . Context , keys [] string ) ( map [ string ] string , error ) {
out := map [ string ] string {}
for _ , key := range keys {
if value , ok := r . values [ key ]; ok {
out [ key ] = value
}
}
return out , nil
}
func ( r * contentModerationHandlerSettingRepo ) SetMultiple ( ctx context . Context , settings map [ string ] string ) error {
if r . values == nil {
r . values = map [ string ] string {}
}
for key , value := range settings {
r . values [ key ] = value
}
return nil
}
func ( r * contentModerationHandlerSettingRepo ) GetAll ( ctx context . Context ) ( map [ string ] string , error ) {
out := make ( map [ string ] string , len ( r . values ))
for key , value := range r . values {
out [ key ] = value
}
return out , nil
}
func ( r * contentModerationHandlerSettingRepo ) Delete ( ctx context . Context , key string ) error {
delete ( r . values , key )
return nil
}
type contentModerationHandlerTestRepo struct {
2026-05-28 20:05:24 +08:00
mu sync . Mutex
2026-05-07 09:01:48 +08:00
logs [] service . ContentModerationLog
}
func ( r * contentModerationHandlerTestRepo ) CreateLog ( ctx context . Context , log * service . ContentModerationLog ) error {
if log != nil {
2026-05-28 20:05:24 +08:00
r . mu . Lock ()
defer r . mu . Unlock ()
2026-05-07 09:01:48 +08:00
r . logs = append ( r . logs , * log )
}
return nil
}
2026-05-28 20:05:24 +08:00
func ( r * contentModerationHandlerTestRepo ) resetLogs () {
r . mu . Lock ()
defer r . mu . Unlock ()
r . logs = nil
}
func ( r * contentModerationHandlerTestRepo ) logSnapshot () [] service . ContentModerationLog {
r . mu . Lock ()
defer r . mu . Unlock ()
return append ([] service . ContentModerationLog ( nil ), r . logs ... )
}
2026-05-07 09:01:48 +08:00
func ( r * contentModerationHandlerTestRepo ) ListLogs ( ctx context . Context , filter service . ContentModerationLogFilter ) ([] service . ContentModerationLog , * pagination . PaginationResult , error ) {
return nil , nil , nil
}
2026-06-11 17:26:14 +08:00
func ( r * contentModerationHandlerTestRepo ) CountFlaggedByUserSince ( ctx context . Context , userID int64 , since time . Time , excludeCyberPolicy bool ) ( int , error ) {
2026-05-07 09:01:48 +08:00
return 0 , nil
}
func ( r * contentModerationHandlerTestRepo ) CleanupExpiredLogs ( ctx context . Context , hitBefore time . Time , nonHitBefore time . Time ) ( * service . ContentModerationCleanupResult , error ) {
return & service . ContentModerationCleanupResult {}, nil
}
2026-06-11 17:26:14 +08:00
func ( r * contentModerationHandlerTestRepo ) UpdateLogEmailSent ( ctx context . Context , id int64 , sent bool ) error {
return nil
}
2026-05-07 09:01:48 +08:00
func TestOpenAIResponsesWebSocket_ContentModerationBlocksFirstFrame ( t * testing . T ) {
gin . SetMode ( gin . TestMode )
moderationServer := httptest . NewServer ( http . HandlerFunc ( func ( w http . ResponseWriter , r * http . Request ) {
require . Equal ( t , "/v1/moderations" , r . URL . Path )
_ , _ = w . Write ([] byte ( `{"results":[{"category_scores":{"sexual":0.9}}]}` ))
}))
defer moderationServer . Close ()
cfg := & service . ContentModerationConfig {
Enabled : true ,
Mode : service . ContentModerationModePreBlock ,
BaseURL : moderationServer . URL ,
Model : "omni-moderation-latest" ,
APIKeys : [] string { "sk-test" },
SampleRate : 100 ,
AllGroups : true ,
BlockMessage : "内容审计测试阻断" ,
}
rawCfg , err := json . Marshal ( cfg )
require . NoError ( t , err )
repo := & contentModerationHandlerTestRepo {}
settingRepo := & contentModerationHandlerSettingRepo { values : map [ string ] string {
service . SettingKeyRiskControlEnabled : "true" ,
service . SettingKeyContentModerationConfig : string ( rawCfg ),
}}
moderationSvc := service . NewContentModerationService (
settingRepo ,
repo ,
nil ,
nil ,
nil ,
nil ,
nil ,
2026-07-31 23:12:22 +08:00
nil ,
2026-05-07 09:01:48 +08:00
)
decision , err := moderationSvc . Check ( context . Background (), service . ContentModerationCheckInput {
UserID : 1 ,
Endpoint : "/v1/responses" ,
Provider : "openai" ,
Model : "gpt-5.5" ,
Protocol : service . ContentModerationProtocolOpenAIResponses ,
Body : [] byte ( `{"model":"gpt-5.5","input":[{"type":"message","role":"user","content":[{"type":"input_text","text":"bad prompt"}]}]}` ),
})
require . NoError ( t , err )
require . True ( t , decision . Blocked )
2026-05-28 20:05:24 +08:00
require . Eventually ( t , func () bool {
return len ( repo . logSnapshot ()) == 1
}, time . Second , 10 * time . Millisecond )
repo . resetLogs ()
2026-05-07 09:01:48 +08:00
h := & OpenAIGatewayHandler {
gatewayService : & service . OpenAIGatewayService {},
billingCacheService : & service . BillingCacheService {},
apiKeyService : & service . APIKeyService {},
contentModerationService : moderationSvc ,
concurrencyHelper : NewConcurrencyHelper ( service . NewConcurrencyService ( & concurrencyCacheMock {}), SSEPingFormatNone , time . Second ),
}
wsServer := newOpenAIWSHandlerTestServer ( t , h , middleware . AuthSubject { UserID : 1 , Concurrency : 1 })
defer wsServer . Close ()
dialCtx , cancelDial := context . WithTimeout ( context . Background (), 3 * time . Second )
clientConn , _ , err := coderws . Dial ( dialCtx , "ws" + strings . TrimPrefix ( wsServer . URL , "http" ) + "/openai/v1/responses" , nil )
cancelDial ()
require . NoError ( t , err )
defer func () {
_ = clientConn . CloseNow ()
}()
writeCtx , cancelWrite := context . WithTimeout ( context . Background (), 3 * time . Second )
err = clientConn . Write ( writeCtx , coderws . MessageText , [] byte ( `{
"type":"response.create",
"model":"gpt-5.5",
"input":[{"type":"message","role":"user","content":[{"type":"input_text","text":"bad prompt"}]}]
}` ))
cancelWrite ()
require . NoError ( t , err )
readCtx , cancelRead := context . WithTimeout ( context . Background (), 3 * time . Second )
_ , payload , readErr := clientConn . Read ( readCtx )
cancelRead ()
if readErr == nil {
require . Contains ( t , string ( payload ), "content_policy_violation" )
require . Contains ( t , string ( payload ), "内容审计测试阻断" )
} else {
var closeErr coderws . CloseError
require . ErrorAs ( t , readErr , & closeErr )
require . Equal ( t , coderws . StatusPolicyViolation , closeErr . Code )
require . Contains ( t , closeErr . Reason , "内容审计测试阻断" )
}
2026-05-29 09:57:02 +08:00
var logs [] service . ContentModerationLog
require . Eventually ( t , func () bool {
logs = repo . logSnapshot ()
return len ( logs ) == 1
}, time . Second , 10 * time . Millisecond )
2026-05-28 20:05:24 +08:00
require . True ( t , logs [ 0 ]. Flagged )
require . Equal ( t , service . ContentModerationActionBlock , logs [ 0 ]. Action )
require . Equal ( t , "bad prompt" , logs [ 0 ]. InputExcerpt )
2026-05-07 09:01:48 +08:00
}
2026-05-03 19:33:09 +08:00
func TestOpenAIResponsesWebSocket_PassthroughUsageLogPersistsUserAgentAndReasoningEffort ( t * testing . T ) {
got := runOpenAIResponsesWebSocketUsageLogCase ( t , openAIResponsesWSUsageLogCase {
firstPayload : `{"type":"response.create","model":"gpt-5.4","stream":false,"reasoning":{"effort":"HIGH"}}` ,
userAgent : testStringPtr ( "codex_cli_rs/0.125.0 test" ),
})
require . NotNil ( t , got . log . UserAgent )
require . Equal ( t , "codex_cli_rs/0.125.0 test" , * got . log . UserAgent )
require . NotNil ( t , got . log . ReasoningEffort )
require . Equal ( t , "high" , * got . log . ReasoningEffort )
require . True ( t , got . log . OpenAIWSMode )
}
func TestOpenAIResponsesWebSocket_PassthroughUsageLogInfersReasoningFromInitialRequestModel ( t * testing . T ) {
got := runOpenAIResponsesWebSocketUsageLogCase ( t , openAIResponsesWSUsageLogCase {
firstPayload : `{"type":"response.create","model":"gpt-5.4-xhigh","stream":false}` ,
userAgent : testStringPtr ( "codex_cli_rs/0.125.0 mapped" ),
channelMapping : map [ string ] string {
"gpt-5.4-xhigh" : "gpt-5.4" ,
},
})
require . Equal ( t , "gpt-5.4" , gjson . GetBytes ( got . upstreamFirstPayload , "model" ). String (),
"上游首帧应使用渠道映射后的模型" )
require . NotNil ( t , got . log . ReasoningEffort )
require . Equal ( t , "xhigh" , * got . log . ReasoningEffort ,
"usage log reasoning effort 必须使用渠道映射前首帧模型后缀推导" )
}
func TestOpenAIResponsesWebSocket_PassthroughUsageLogLeavesUserAgentNilWhenMissing ( t * testing . T ) {
got := runOpenAIResponsesWebSocketUsageLogCase ( t , openAIResponsesWSUsageLogCase {
firstPayload : `{"type":"response.create","model":"gpt-5.4","stream":false,"reasoning":{"effort":"medium"}}` ,
userAgent : testStringPtr ( "" ),
})
require . Nil ( t , got . log . UserAgent , "空入站 User-Agent 不应由上游握手 UA 或默认 UA 兜底" )
require . NotNil ( t , got . log . ReasoningEffort )
require . Equal ( t , "medium" , * got . log . ReasoningEffort )
}
2026-07-23 20:17:11 +08:00
func TestOpenAIResponsesWebSocket_PassthroughTracksModelPerTurn ( t * testing . T ) {
got := runOpenAIResponsesWebSocketUsageLogCase ( t , openAIResponsesWSUsageLogCase {
firstPayload : `{"type":"response.create","model":"sol","stream":false}` ,
secondPayload : `{"type":"response.create","model":"terra","stream":false}` ,
channelMapping : map [ string ] string {
"sol" : "sol-channel" ,
"terra" : "terra-channel" ,
},
accountModelMapping : map [ string ] any {
"sol" : "gpt-5.6-sol" ,
"terra" : "gpt-5.6-terra" ,
"sol-channel" : "gpt-5.6-sol" ,
"terra-channel" : "gpt-5.6-terra" ,
},
})
require . Len ( t , got . upstreamPayloads , 2 )
2026-08-19 23:13:40 -07:00
require . Equal ( t , "sol-channel" , gjson . GetBytes ( got . upstreamPayloads [ 0 ], "model" ). String ())
require . Equal ( t , "terra-channel" , gjson . GetBytes ( got . upstreamPayloads [ 1 ], "model" ). String ())
2026-07-23 20:17:11 +08:00
require . Len ( t , got . clientEvents , 2 )
require . Equal ( t , "sol" , gjson . GetBytes ( got . clientEvents [ 0 ], "response.model" ). String ())
require . Equal ( t , "terra" , gjson . GetBytes ( got . clientEvents [ 1 ], "response.model" ). String ())
require . Len ( t , got . logs , 2 )
require . Equal ( t , "sol" , got . logs [ 0 ]. Model )
require . Equal ( t , "sol" , got . logs [ 0 ]. RequestedModel )
require . NotNil ( t , got . logs [ 0 ]. UpstreamModel )
2026-08-19 23:13:40 -07:00
require . Equal ( t , "sol-channel" , * got . logs [ 0 ]. UpstreamModel )
2026-07-23 20:17:11 +08:00
require . NotNil ( t , got . logs [ 0 ]. ModelMappingChain )
2026-08-19 23:13:40 -07:00
require . Equal ( t , "sol→sol-channel" , * got . logs [ 0 ]. ModelMappingChain )
2026-07-23 20:17:11 +08:00
require . Equal ( t , "terra" , got . logs [ 1 ]. Model )
require . Equal ( t , "terra" , got . logs [ 1 ]. RequestedModel )
require . NotNil ( t , got . logs [ 1 ]. UpstreamModel )
2026-08-19 23:13:40 -07:00
require . Equal ( t , "terra-channel" , * got . logs [ 1 ]. UpstreamModel )
2026-07-23 20:17:11 +08:00
require . NotNil ( t , got . logs [ 1 ]. ModelMappingChain )
2026-08-19 23:13:40 -07:00
require . Equal ( t , "terra→terra-channel" , * got . logs [ 1 ]. ModelMappingChain )
2026-07-31 09:55:16 +08:00
require . InDelta ( t , got . logs [ 1 ]. TotalCost * 2.5 , got . logs [ 0 ]. TotalCost , 1e-12 ,
2026-07-23 20:17:11 +08:00
"each turn must be billed with its own channel-mapped model" )
}
2026-09-03 10:51:12 +08:00
func TestOpenAIResponsesWebSocket_ChannelMappedTargetSelectsAccountWithoutRequestedAlias ( t * testing . T ) {
2026-07-23 20:17:11 +08:00
got := runOpenAIResponsesWebSocketUsageLogCase ( t , openAIResponsesWSUsageLogCase {
firstPayload : `{"type":"response.create","model":"public-alias","stream":false}` ,
secondPayload : `{"type":"response.create","stream":false}` ,
channelMapping : map [ string ] string {
"public-alias" : "gpt-5.6-sol" ,
},
accountModelMapping : map [ string ] any {
2026-09-03 10:51:12 +08:00
"gpt-5.6-sol" : "gpt-5.6-sol" ,
2026-07-23 20:17:11 +08:00
},
})
require . Len ( t , got . upstreamPayloads , 2 )
require . Equal ( t , "gpt-5.6-sol" , gjson . GetBytes ( got . upstreamPayloads [ 0 ], "model" ). String ())
require . Equal ( t , "gpt-5.6-sol" , gjson . GetBytes ( got . upstreamPayloads [ 1 ], "model" ). String ())
require . Len ( t , got . clientEvents , 2 )
require . Equal ( t , "public-alias" , gjson . GetBytes ( got . clientEvents [ 0 ], "response.model" ). String ())
require . Equal ( t , "public-alias" , gjson . GetBytes ( got . clientEvents [ 1 ], "response.model" ). String ())
require . Len ( t , got . logs , 2 )
for _ , usageLog := range got . logs {
require . Equal ( t , "public-alias" , usageLog . RequestedModel )
require . NotNil ( t , usageLog . UpstreamModel )
require . Equal ( t , "gpt-5.6-sol" , * usageLog . UpstreamModel )
require . NotNil ( t , usageLog . ModelMappingChain )
require . Equal ( t , "public-alias→gpt-5.6-sol" , * usageLog . ModelMappingChain )
}
}
func TestOpenAIResponsesWebSocket_PassthroughKeepsTurnMappingSnapshot ( t * testing . T ) {
got := runOpenAIResponsesWebSocketUsageLogCase ( t , openAIResponsesWSUsageLogCase {
firstPayload : `{"type":"response.create","model":"sol","stream":false}` ,
secondPayload : `{"type":"response.create","model":"sol","stream":false}` ,
channelMapping : map [ string ] string {
"sol" : "gpt-5.6-sol" ,
},
afterFirstUpstreamRequest : func ( channelSvc * service . ChannelService ) error {
if channelSvc == nil {
return errors . New ( "channel service is nil" )
}
_ , err := channelSvc . Update ( context . Background (), 7701 , & service . UpdateChannelInput {
ModelMapping : map [ string ] map [ string ] string {
service . PlatformOpenAI : { "sol" : "gpt-5.6-terra" },
},
})
return err
},
})
require . Len ( t , got . upstreamPayloads , 2 )
require . Equal ( t , "gpt-5.6-sol" , gjson . GetBytes ( got . upstreamPayloads [ 0 ], "model" ). String ())
require . Equal ( t , "gpt-5.6-terra" , gjson . GetBytes ( got . upstreamPayloads [ 1 ], "model" ). String ())
require . Len ( t , got . logs , 2 )
require . Equal ( t , "sol" , got . logs [ 0 ]. Model )
require . NotNil ( t , got . logs [ 0 ]. UpstreamModel )
require . Equal ( t , "gpt-5.6-sol" , * got . logs [ 0 ]. UpstreamModel )
require . NotNil ( t , got . logs [ 0 ]. ModelMappingChain )
require . Equal ( t , "sol→gpt-5.6-sol" , * got . logs [ 0 ]. ModelMappingChain )
require . InDelta ( t , 40e-6 , got . logs [ 0 ]. TotalCost , 1e-12 ,
"the in-flight turn must retain the channel-mapped billing model used when it was sent" )
require . Equal ( t , "sol" , got . logs [ 1 ]. Model )
require . NotNil ( t , got . logs [ 1 ]. UpstreamModel )
require . Equal ( t , "gpt-5.6-terra" , * got . logs [ 1 ]. UpstreamModel )
require . NotNil ( t , got . logs [ 1 ]. ModelMappingChain )
require . Equal ( t , "sol→gpt-5.6-terra" , * got . logs [ 1 ]. ModelMappingChain )
2026-07-31 09:55:16 +08:00
require . InDelta ( t , got . logs [ 1 ]. TotalCost * 2.5 , got . logs [ 0 ]. TotalCost , 1e-12 ,
2026-07-23 20:17:11 +08:00
"the next turn must use the updated channel mapping" )
}
func TestOpenAIResponsesWebSocket_CtxPoolAppliesPerTurnMappingAndPreservesRequestedModel ( t * testing . T ) {
got := runOpenAIResponsesWebSocketUsageLogCase ( t , openAIResponsesWSUsageLogCase {
firstPayload : `{"type":"response.create","model":"gpt-5.6-sol","stream":false}` ,
secondPayload : `{"type":"response.create","model":"gpt-5.6-terra","stream":false}` ,
ingressMode : service . OpenAIWSIngressModeCtxPool ,
billingModelSource : service . BillingModelSourceRequested ,
channelMapping : map [ string ] string {
"gpt-5.6-terra" : "gpt-5.6-sol" ,
},
})
require . Len ( t , got . upstreamPayloads , 2 )
require . Equal ( t , "gpt-5.6-sol" , gjson . GetBytes ( got . upstreamPayloads [ 0 ], "model" ). String ())
require . Equal ( t , "gpt-5.6-sol" , gjson . GetBytes ( got . upstreamPayloads [ 1 ], "model" ). String ())
require . Len ( t , got . clientEvents , 2 )
require . Equal ( t , "gpt-5.6-sol" , gjson . GetBytes ( got . clientEvents [ 0 ], "response.model" ). String ())
require . Equal ( t , "gpt-5.6-terra" , gjson . GetBytes ( got . clientEvents [ 1 ], "response.model" ). String ())
require . Len ( t , got . logs , 2 )
require . Equal ( t , "gpt-5.6-sol" , got . logs [ 0 ]. RequestedModel )
require . Nil ( t , got . logs [ 0 ]. ModelMappingChain )
require . InDelta ( t , 40e-6 , got . logs [ 0 ]. TotalCost , 1e-12 )
require . Equal ( t , "gpt-5.6-terra" , got . logs [ 1 ]. RequestedModel )
require . NotNil ( t , got . logs [ 1 ]. ModelMappingChain )
require . Equal ( t , "gpt-5.6-terra→gpt-5.6-sol" , * got . logs [ 1 ]. ModelMappingChain )
2026-07-31 10:02:04 +08:00
require . InDelta ( t , 16e-6 , got . logs [ 1 ]. TotalCost , 1e-12 ,
2026-07-23 20:17:11 +08:00
"BillingModelSourceRequested must use the client model before channel mapping" )
}
func TestOpenAIWSTurnBillingModelPreservesImagePricingModel ( t * testing . T ) {
tests := [] struct {
name string
resultModel string
mapping service . ChannelMappingResult
requestedModel string
upstreamModel string
wantBillingModel string
}{
{
name : "upstream billing preserves image model" ,
resultModel : "gpt-image-2" ,
mapping : service . ChannelMappingResult { BillingModelSource : service . BillingModelSourceUpstream },
requestedModel : "gpt-5.6-sol" ,
upstreamModel : "gpt-5.6-sol" ,
wantBillingModel : "gpt-image-2" ,
},
{
name : "unmapped channel preserves image model" ,
resultModel : "gpt-image-2" ,
mapping : service . ChannelMappingResult { MappedModel : "gpt-5.6-sol" , BillingModelSource : service . BillingModelSourceChannelMapped },
requestedModel : "gpt-5.6-sol" ,
upstreamModel : "gpt-5.6-sol" ,
wantBillingModel : "gpt-image-2" ,
},
{
name : "requested source overrides image model" ,
resultModel : "gpt-image-2" ,
mapping : service . ChannelMappingResult { BillingModelSource : service . BillingModelSourceRequested },
requestedModel : "public-image-alias" ,
upstreamModel : "gpt-5.6-sol" ,
wantBillingModel : "public-image-alias" ,
},
{
name : "mapped channel source overrides image model" ,
resultModel : "gpt-image-2" ,
mapping : service . ChannelMappingResult { MappedModel : "priced-channel-model" , BillingModelSource : service . BillingModelSourceChannelMapped },
requestedModel : "public-image-alias" ,
upstreamModel : "gpt-5.6-sol" ,
wantBillingModel : "priced-channel-model" ,
},
{
name : "text turn falls back to upstream model" ,
mapping : service . ChannelMappingResult { BillingModelSource : service . BillingModelSourceUpstream },
requestedModel : "public-alias" ,
upstreamModel : "gpt-5.6-sol" ,
wantBillingModel : "gpt-5.6-sol" ,
},
}
for _ , tt := range tests {
t . Run ( tt . name , func ( t * testing . T ) {
result := & service . OpenAIForwardResult { BillingModel : tt . resultModel }
require . Equal ( t , tt . wantBillingModel , openAIWSTurnBillingModel ( result , tt . mapping , tt . requestedModel , tt . upstreamModel ))
})
}
}
2026-08-19 23:13:40 -07:00
func TestOpenAIAccountScheduleModelUsesActualOrSharedResolver ( t * testing . T ) {
account := & service . Account {
Platform : service . PlatformOpenAI ,
Type : service . AccountTypeOAuth ,
Credentials : map [ string ] any {
"model_mapping" : map [ string ] any { "public" : "billing" },
"compact_model_mapping" : map [ string ] any { "public" : "compact-actual" },
},
}
reported := & service . OpenAIForwardResult { UpstreamModel : "observed-actual" }
require . Equal ( t , "observed-actual" , openAIAccountScheduleModel ( nil , account , "public" , true , reported ))
require . Equal ( t , "compact-actual" , openAIAccountScheduleModel ( nil , account , "public" , true , nil ))
require . Equal ( t , "billing" , openAIAccountScheduleModel ( nil , account , "public" , false , nil ))
c , _ := gin . CreateTestContext ( nil )
service . SetOpsUpstreamModel ( c , "attempt-actual" )
require . Equal ( t , "attempt-actual" , openAIAccountScheduleModel ( c , account , "public" , true , nil ))
setOpsSelectedAccount ( c , account . ID , account . Platform )
require . Equal ( t , "attempt-actual" , openAIAccountScheduleModel ( c , account , "public" , true , nil ))
}
2026-09-03 10:51:12 +08:00
func TestOpenAIChannelForwardModelForScheduler ( t * testing . T ) {
tests := [] struct {
name string
mapping service . ChannelMappingResult
requested string
want string
}{
{
name : "mapped model is used for capability filtering" ,
mapping : service . ChannelMappingResult { Mapped : true , MappedModel : " gpt-forward " },
requested : "client-alias" ,
want : "gpt-forward" ,
},
{
name : "unmapped request preserves client model" ,
mapping : service . ChannelMappingResult { MappedModel : "client-model" },
requested : "client-model" ,
want : "client-model" ,
},
{
name : "empty mapped model safely preserves client model" ,
mapping : service . ChannelMappingResult { Mapped : true , MappedModel : " " },
requested : "client-model" ,
want : "client-model" ,
},
}
for _ , tt := range tests {
t . Run ( tt . name , func ( t * testing . T ) {
require . Equal ( t , tt . want , openAIChannelForwardModel ( tt . mapping , tt . requested ))
})
}
}
2026-07-23 20:17:11 +08:00
func TestShouldReportOpenAIWSProxyAccountFailure ( t * testing . T ) {
t . Run ( "unsupported client model switch does not penalize account" , func ( t * testing . T ) {
err := fmt . Errorf ( "wrapped ingress turn: %w" , newOpenAIWSUnsupportedModelSwitchError ( "gpt-unsupported" ))
require . False ( t , shouldReportOpenAIWSProxyAccountFailure ( err ))
var closeErr * service . OpenAIWSClientCloseError
require . ErrorAs ( t , err , & closeErr )
require . Equal ( t , coderws . StatusPolicyViolation , closeErr . StatusCode ())
require . Equal ( t , "model switch requires reconnect" , closeErr . Reason ())
})
t . Run ( "upstream policy violation still penalizes account" , func ( t * testing . T ) {
err := service . NewOpenAIWSClientCloseError (
coderws . StatusPolicyViolation ,
"upstream websocket authentication failed" ,
errors . New ( "upstream rejected credentials" ),
)
require . True ( t , shouldReportOpenAIWSProxyAccountFailure ( err ))
})
t . Run ( "generic proxy failure still penalizes account" , func ( t * testing . T ) {
require . True ( t , shouldReportOpenAIWSProxyAccountFailure ( errors . New ( "upstream websocket read failed" )))
})
}
2026-02-28 15:01:20 +08:00
func TestSetOpenAIClientTransportHTTP ( t * testing . T ) {
gin . SetMode ( gin . TestMode )
w := httptest . NewRecorder ()
c , _ := gin . CreateTestContext ( w )
setOpenAIClientTransportHTTP ( c )
require . Equal ( t , service . OpenAIClientTransportHTTP , service . GetOpenAIClientTransport ( c ))
}
func TestSetOpenAIClientTransportWS ( t * testing . T ) {
gin . SetMode ( gin . TestMode )
w := httptest . NewRecorder ()
c , _ := gin . CreateTestContext ( w )
setOpenAIClientTransportWS ( c )
require . Equal ( t , service . OpenAIClientTransportWS , service . GetOpenAIClientTransport ( c ))
}
2026-02-10 08:59:30 +08:00
// TestOpenAIHandler_GjsonExtraction 验证 gjson 从请求体中提取 model/stream 的正确性
func TestOpenAIHandler_GjsonExtraction ( t * testing . T ) {
tests := [] struct {
name string
body string
wantModel string
wantStream bool
}{
{ "正常提取" , `{"model":"gpt-4","stream":true,"input":"hello"}` , "gpt-4" , true },
{ "stream false" , `{"model":"gpt-4","stream":false}` , "gpt-4" , false },
{ "无 stream 字段" , `{"model":"gpt-4"}` , "gpt-4" , false },
{ "model 缺失" , `{"stream":true}` , "" , true },
}
for _ , tt := range tests {
t . Run ( tt . name , func ( t * testing . T ) {
body := [] byte ( tt . body )
2026-02-10 09:13:20 +08:00
modelResult := gjson . GetBytes ( body , "model" )
model := ""
if modelResult . Type == gjson . String {
model = modelResult . String ()
}
2026-02-10 08:59:30 +08:00
stream := gjson . GetBytes ( body , "stream" ). Bool ()
require . Equal ( t , tt . wantModel , model )
require . Equal ( t , tt . wantStream , stream )
})
}
}
2026-02-10 09:13:20 +08:00
// TestOpenAIHandler_GjsonValidation 验证修复后的 JSON 合法性和类型校验
func TestOpenAIHandler_GjsonValidation ( t * testing . T ) {
// 非法 JSON 被 gjson.ValidBytes 拦截
require . False ( t , gjson . ValidBytes ([] byte ( `{invalid json` )))
// model 为数字 → 类型不是 gjson.String,应被拒绝
body := [] byte ( `{"model":123}` )
modelResult := gjson . GetBytes ( body , "model" )
require . True ( t , modelResult . Exists ())
require . NotEqual ( t , gjson . String , modelResult . Type )
// model 为 null → 类型不是 gjson.String,应被拒绝
body2 := [] byte ( `{"model":null}` )
modelResult2 := gjson . GetBytes ( body2 , "model" )
require . True ( t , modelResult2 . Exists ())
require . NotEqual ( t , gjson . String , modelResult2 . Type )
// stream 为 string → 类型既不是 True 也不是 False,应被拒绝
body3 := [] byte ( `{"model":"gpt-4","stream":"true"}` )
streamResult := gjson . GetBytes ( body3 , "stream" )
require . True ( t , streamResult . Exists ())
require . NotEqual ( t , gjson . True , streamResult . Type )
require . NotEqual ( t , gjson . False , streamResult . Type )
// stream 为 int → 同上
body4 := [] byte ( `{"model":"gpt-4","stream":1}` )
streamResult2 := gjson . GetBytes ( body4 , "stream" )
require . True ( t , streamResult2 . Exists ())
require . NotEqual ( t , gjson . True , streamResult2 . Type )
require . NotEqual ( t , gjson . False , streamResult2 . Type )
}
2026-02-10 08:59:30 +08:00
// TestOpenAIHandler_InstructionsInjection 验证 instructions 的 gjson/sjson 注入逻辑
func TestOpenAIHandler_InstructionsInjection ( t * testing . T ) {
// 测试 1:无 instructions → 注入
body := [] byte ( `{"model":"gpt-4"}` )
existing := gjson . GetBytes ( body , "instructions" ). String ()
require . Empty ( t , existing )
newBody , err := sjson . SetBytes ( body , "instructions" , "test instruction" )
require . NoError ( t , err )
require . Equal ( t , "test instruction" , gjson . GetBytes ( newBody , "instructions" ). String ())
// 测试 2:已有 instructions → 不覆盖
body2 := [] byte ( `{"model":"gpt-4","instructions":"existing"}` )
existing2 := gjson . GetBytes ( body2 , "instructions" ). String ()
require . Equal ( t , "existing" , existing2 )
// 测试 3:空白 instructions → 注入
body3 := [] byte ( `{"model":"gpt-4","instructions":" "}` )
existing3 := strings . TrimSpace ( gjson . GetBytes ( body3 , "instructions" ). String ())
require . Empty ( t , existing3 )
2026-02-10 09:13:20 +08:00
// 测试 4: sjson.SetBytes 返回错误时不应 panic
// 正常 JSON 不会产生 sjson 错误,验证返回值被正确处理
validBody := [] byte ( `{"model":"gpt-4"}` )
result , setErr := sjson . SetBytes ( validBody , "instructions" , "hello" )
require . NoError ( t , setErr )
require . True ( t , gjson . ValidBytes ( result ))
2026-02-10 08:59:30 +08:00
}
2026-02-28 15:01:20 +08:00
func newOpenAIHandlerForPreviousResponseIDValidation ( t * testing . T , cache * concurrencyCacheMock ) * OpenAIGatewayHandler {
t . Helper ()
if cache == nil {
cache = & concurrencyCacheMock {
acquireUserSlotFn : func ( ctx context . Context , userID int64 , maxConcurrency int , requestID string ) ( bool , error ) {
return true , nil
},
acquireAccountSlotFn : func ( ctx context . Context , accountID int64 , maxConcurrency int , requestID string ) ( bool , error ) {
return true , nil
},
}
}
return & OpenAIGatewayHandler {
gatewayService : & service . OpenAIGatewayService {},
billingCacheService : & service . BillingCacheService {},
apiKeyService : & service . APIKeyService {},
concurrencyHelper : NewConcurrencyHelper ( service . NewConcurrencyService ( cache ), SSEPingFormatNone , time . Second ),
}
}
func newOpenAIWSHandlerTestServer ( t * testing . T , h * OpenAIGatewayHandler , subject middleware . AuthSubject ) * httptest . Server {
t . Helper ()
groupID := int64 ( 2 )
apiKey := & service . APIKey {
ID : 101 ,
GroupID : & groupID ,
User : & service . User { ID : subject . UserID },
}
router := gin . New ()
router . Use ( func ( c * gin . Context ) {
c . Set ( string ( middleware . ContextKeyAPIKey ), apiKey )
c . Set ( string ( middleware . ContextKeyUser ), subject )
c . Next ()
})
router . GET ( "/openai/v1/responses" , h . ResponsesWebSocket )
return httptest . NewServer ( router )
}
2026-05-03 19:33:09 +08:00
type openAIResponsesWSUsageLogCase struct {
2026-09-22 15:23:09 +08:00
simpleModeRejectAtRead int64
2026-09-19 20:50:28 +08:00
compositeResolver * service . CompositeRouteResolver
accountPlatform string
2026-09-22 15:23:09 +08:00
closeReason string
2026-09-19 20:50:28 +08:00
closeStatus coderws . StatusCode
2026-09-22 15:23:09 +08:00
firstPayload string
2026-09-05 15:58:26 +08:00
// midPayload 在首个 turn 完成后发送(如 session.update),上游桩会为它
// 回一个 response.completed,客户端按普通事件读取。
midPayload string
2026-07-23 20:17:11 +08:00
secondPayload string
userAgent * string
ingressMode string
channelMapping map [ string ] string
billingModelSource string
accountModelMapping map [ string ] any
afterFirstUpstreamRequest func ( channelSvc * service . ChannelService ) error
2026-09-05 15:58:26 +08:00
// group 覆盖 apiKey.Group(分组级模型白名单测试用);nil 保持原有无分组行为。
group * service . Group
2026-10-01 11:27:35 +08:00
// apiKeyService 非 nil 时模拟 API Key 认证中间件:连接认证快照经它按
// apiKeyCredential 取得,并把其分组放入请求 ctx;handler 也使用它。
apiKeyService * service . APIKeyService
apiKeyCredential string
// accountRateMultiplier 覆盖账号倍率(利润门测试用)。
accountRateMultiplier * float64
2026-09-05 15:58:26 +08:00
// firstFrameCloseExpected:首帧即被拒(连接被 1008 关闭),不期待任何响应帧。
firstFrameCloseExpected bool
// secondTurnCloseExpected:第二个 turn 被拒(连接被 1008 关闭)。
secondTurnCloseExpected bool
2026-05-03 19:33:09 +08:00
}
type openAIResponsesWSUsageLogResult struct {
log * service . UsageLog
2026-07-23 20:17:11 +08:00
logs [] * service . UsageLog
2026-05-03 19:33:09 +08:00
upstreamFirstPayload [] byte
2026-07-23 20:17:11 +08:00
upstreamPayloads [][] byte
clientEvents [][] byte
2026-05-03 19:33:09 +08:00
}
type openAIWSUsageHandlerAccountRepoStub struct {
service . AccountRepository
account service . Account
}
func ( s * openAIWSUsageHandlerAccountRepoStub ) ListSchedulableByPlatform ( ctx context . Context , platform string ) ([] service . Account , error ) {
if s . account . Platform != platform {
return nil , nil
}
return [] service . Account { s . account }, nil
}
func ( s * openAIWSUsageHandlerAccountRepoStub ) ListSchedulableByGroupIDAndPlatform ( ctx context . Context , groupID int64 , platform string ) ([] service . Account , error ) {
return s . ListSchedulableByPlatform ( ctx , platform )
}
func ( s * openAIWSUsageHandlerAccountRepoStub ) GetByID ( ctx context . Context , id int64 ) ( * service . Account , error ) {
if s . account . ID != id {
return nil , nil
}
account := s . account
return & account , nil
}
2026-05-26 20:07:00 +08:00
type openAIWSFailoverHandlerAccountRepoStub struct {
service . AccountRepository
accounts [] service . Account
rateLimitedIDs [] int64
}
2026-07-15 03:43:50 +08:00
type openAIHTTPPassthroughFailoverUpstream struct {
service . HTTPUpstream
mu sync . Mutex
accountIDs [] int64
}
func ( u * openAIHTTPPassthroughFailoverUpstream ) Do ( _ * http . Request , _ string , accountID int64 , _ int ) ( * http . Response , error ) {
u . mu . Lock ()
u . accountIDs = append ( u . accountIDs , accountID )
u . mu . Unlock ()
return & http . Response {
StatusCode : http . StatusBadGateway ,
Header : http . Header { "Content-Type" : [] string { "application/json" }},
Body : io . NopCloser ( strings . NewReader ( `{"error":{"message":"temporary upstream failure"}}` )),
}, nil
}
func ( u * openAIHTTPPassthroughFailoverUpstream ) calls () [] int64 {
u . mu . Lock ()
defer u . mu . Unlock ()
return append ([] int64 ( nil ), u . accountIDs ... )
}
2026-07-29 18:47:07 +08:00
type openAIHTTPPassthroughAuthFailoverUpstream struct {
service . HTTPUpstream
mu sync . Mutex
accountIDs [] int64
statusCode int
}
func ( u * openAIHTTPPassthroughAuthFailoverUpstream ) Do ( _ * http . Request , _ string , accountID int64 , _ int ) ( * http . Response , error ) {
u . mu . Lock ()
u . accountIDs = append ( u . accountIDs , accountID )
u . mu . Unlock ()
if accountID == 9911 {
return & http . Response {
StatusCode : http . StatusOK ,
Header : http . Header { "Content-Type" : [] string { "application/json" }},
Body : io . NopCloser ( strings . NewReader ( `{"id":"resp_healthy","object":"response","model":"gpt-5.2","status":"completed","usage":{"input_tokens":1,"output_tokens":1,"total_tokens":2}}` )),
}, nil
}
return & http . Response {
StatusCode : u . statusCode ,
Header : http . Header { "Content-Type" : [] string { "application/json" }},
Body : io . NopCloser ( strings . NewReader ( `{"error":{"message":"upstream credential rejected"}}` )),
}, nil
}
func ( u * openAIHTTPPassthroughAuthFailoverUpstream ) calls () [] int64 {
u . mu . Lock ()
defer u . mu . Unlock ()
return append ([] int64 ( nil ), u . accountIDs ... )
}
2026-07-30 15:22:16 +08:00
type openAIHTTPPassthroughSSERateLimitUpstream struct {
service . HTTPUpstream
mu sync . Mutex
accountIDs [] int64
}
func ( u * openAIHTTPPassthroughSSERateLimitUpstream ) Do ( _ * http . Request , _ string , accountID int64 , _ int ) ( * http . Response , error ) {
u . mu . Lock ()
u . accountIDs = append ( u . accountIDs , accountID )
u . mu . Unlock ()
body := strings . Join ([] string {
"event: response.created" ,
`data: {"type":"response.created","response":{"id":"resp_rate_limited"}}` ,
"" ,
"event: response.failed" ,
`data: {"type":"response.failed","response":{"id":"resp_rate_limited","status":"failed","error":{"code":"rate_limit_exceeded","message":"Concurrency limit exceeded for account, please retry later"}}}` ,
"" ,
}, "\n" )
return & http . Response {
StatusCode : http . StatusOK ,
Header : http . Header {
"Content-Type" : [] string { "text/event-stream" },
"Retry-After" : [] string { "1" },
},
Body : io . NopCloser ( strings . NewReader ( body )),
}, nil
}
func ( u * openAIHTTPPassthroughSSERateLimitUpstream ) calls () [] int64 {
u . mu . Lock ()
defer u . mu . Unlock ()
return append ([] int64 ( nil ), u . accountIDs ... )
}
2026-05-26 20:07:00 +08:00
func ( s * openAIWSFailoverHandlerAccountRepoStub ) ListSchedulableByPlatform ( ctx context . Context , platform string ) ([] service . Account , error ) {
out := make ([] service . Account , 0 , len ( s . accounts ))
for _ , account := range s . accounts {
if account . Platform == platform && account . IsSchedulable () {
out = append ( out , account )
}
}
return out , nil
}
func ( s * openAIWSFailoverHandlerAccountRepoStub ) ListSchedulableByGroupIDAndPlatform ( ctx context . Context , groupID int64 , platform string ) ([] service . Account , error ) {
return s . ListSchedulableByPlatform ( ctx , platform )
}
func ( s * openAIWSFailoverHandlerAccountRepoStub ) ListSchedulableUngroupedByPlatform ( ctx context . Context , platform string ) ([] service . Account , error ) {
return s . ListSchedulableByPlatform ( ctx , platform )
}
func ( s * openAIWSFailoverHandlerAccountRepoStub ) GetByID ( ctx context . Context , id int64 ) ( * service . Account , error ) {
for _ , account := range s . accounts {
if account . ID == id {
acc := account
return & acc , nil
}
}
return nil , nil
}
func ( s * openAIWSFailoverHandlerAccountRepoStub ) SetRateLimited ( ctx context . Context , id int64 , resetAt time . Time ) error {
s . rateLimitedIDs = append ( s . rateLimitedIDs , id )
for i := range s . accounts {
if s . accounts [ i ]. ID == id {
reset := resetAt
s . accounts [ i ]. RateLimitResetAt = & reset
break
}
}
return nil
}
2026-05-03 19:33:09 +08:00
type openAIWSUsageHandlerUsageLogRepoStub struct {
service . UsageLogRepository
created chan * service . UsageLog
}
func ( s * openAIWSUsageHandlerUsageLogRepoStub ) Create ( ctx context . Context , log * service . UsageLog ) ( bool , error ) {
if s . created != nil {
s . created <- log
}
return true , nil
}
type openAIWSUsageHandlerChannelRepoStub struct {
service . ChannelRepository
2026-07-23 20:17:11 +08:00
mu sync . Mutex
2026-05-03 19:33:09 +08:00
channels [] service . Channel
groupPlatforms map [ int64 ] string
}
func ( s * openAIWSUsageHandlerChannelRepoStub ) ListAll ( ctx context . Context ) ([] service . Channel , error ) {
2026-07-23 20:17:11 +08:00
s . mu . Lock ()
defer s . mu . Unlock ()
out := make ([] service . Channel , 0 , len ( s . channels ))
for i := range s . channels {
out = append ( out , * s . channels [ i ]. Clone ())
}
return out , nil
}
func ( s * openAIWSUsageHandlerChannelRepoStub ) GetByID ( ctx context . Context , id int64 ) ( * service . Channel , error ) {
s . mu . Lock ()
defer s . mu . Unlock ()
for i := range s . channels {
if s . channels [ i ]. ID == id {
return s . channels [ i ]. Clone (), nil
}
}
return nil , service . ErrChannelNotFound
}
func ( s * openAIWSUsageHandlerChannelRepoStub ) Update ( ctx context . Context , channel * service . Channel ) error {
s . mu . Lock ()
defer s . mu . Unlock ()
for i := range s . channels {
if s . channels [ i ]. ID == channel . ID {
s . channels [ i ] = * channel . Clone ()
return nil
}
}
return service . ErrChannelNotFound
}
func ( s * openAIWSUsageHandlerChannelRepoStub ) GetGroupIDs ( ctx context . Context , channelID int64 ) ([] int64 , error ) {
channel , err := s . GetByID ( ctx , channelID )
if err != nil {
return nil , err
}
return append ([] int64 ( nil ), channel . GroupIDs ... ), nil
2026-05-03 19:33:09 +08:00
}
func ( s * openAIWSUsageHandlerChannelRepoStub ) GetGroupPlatforms ( ctx context . Context , groupIDs [] int64 ) ( map [ int64 ] string , error ) {
out := make ( map [ int64 ] string , len ( groupIDs ))
for _ , groupID := range groupIDs {
if platform := strings . TrimSpace ( s . groupPlatforms [ groupID ]); platform != "" {
out [ groupID ] = platform
}
}
return out , nil
}
2026-07-15 03:43:50 +08:00
func TestOpenAIResponses_APIKeyPassthroughPool5xxRetriesThenExhaustsMaxSwitches ( t * testing . T ) {
gin . SetMode ( gin . TestMode )
groupID := int64 ( 4203 )
accounts := [] service . Account {
{
ID : 9910 , Name : "pool-api-key" , Platform : service . PlatformOpenAI ,
Type : service . AccountTypeAPIKey , Status : service . StatusActive , Schedulable : true , Priority : 1 ,
Credentials : map [ string ] any {
"api_key" : "sk-pool" ,
"base_url" : "https://api.example.test" ,
"pool_mode" : true ,
"pool_mode_retry_count" : float64 ( 1 ),
"pool_mode_retry_status_codes" : [] any { float64 ( http . StatusBadGateway )},
},
Extra : map [ string ] any { "openai_passthrough" : true },
},
{
ID : 9911 , Name : "fallback-api-key" , Platform : service . PlatformOpenAI ,
Type : service . AccountTypeAPIKey , Status : service . StatusActive , Schedulable : true , Priority : 2 ,
Credentials : map [ string ] any {
"api_key" : "sk-fallback" ,
"base_url" : "https://api.example.test" ,
},
Extra : map [ string ] any { "openai_passthrough" : true },
},
}
cfg := & config . Config { RunMode : config . RunModeSimple }
cfg . Default . RateMultiplier = 1
cfg . Security . URLAllowlist . Enabled = false
cfg . Gateway . MaxAccountSwitches = 1
accountRepo := & openAIWSFailoverHandlerAccountRepoStub { accounts : accounts }
upstream := & openAIHTTPPassthroughFailoverUpstream {}
billingCacheSvc := service . NewBillingCacheService ( nil , nil , nil , nil , nil , nil , cfg , nil )
t . Cleanup ( billingCacheSvc . Stop )
gatewaySvc := service . NewOpenAIGatewayService (
accountRepo ,
nil ,
nil ,
nil ,
nil ,
nil ,
nil ,
cfg ,
nil ,
nil ,
service . NewBillingService ( cfg , nil ),
nil ,
billingCacheSvc ,
upstream ,
& service . DeferredService {},
nil ,
nil ,
nil ,
nil ,
nil ,
nil ,
nil ,
)
h := NewOpenAIGatewayHandler (
gatewaySvc ,
service . NewConcurrencyService ( nil ),
billingCacheSvc ,
service . NewAPIKeyService ( nil , nil , nil , nil , nil , nil , cfg ),
nil ,
nil ,
nil ,
nil ,
cfg ,
)
rec := httptest . NewRecorder ()
c , _ := gin . CreateTestContext ( rec )
c . Request = httptest . NewRequest ( http . MethodPost , "/openai/v1/responses" , strings . NewReader ( `{"model":"gpt-5.2","input":"hello","stream":false}` ))
c . Request . Header . Set ( "Content-Type" , "application/json" )
c . Set ( string ( middleware . ContextKeyAPIKey ), & service . APIKey {
ID : 1803 , GroupID : & groupID ,
User : & service . User { ID : 1703 , Status : service . StatusActive },
Group : & service . Group { ID : groupID , Platform : service . PlatformOpenAI , Status : service . StatusActive },
})
c . Set ( string ( middleware . ContextKeyUser ), middleware . AuthSubject { UserID : 1703 , Concurrency : 0 })
h . Responses ( c )
require . Equal ( t , [] int64 { 9910 , 9910 , 9911 }, upstream . calls ())
require . Equal ( t , http . StatusBadGateway , rec . Code )
require . Equal ( t , "upstream_error" , gjson . GetBytes ( rec . Body . Bytes (), "error.type" ). String ())
require . Equal ( t , "Upstream service temporarily unavailable" , gjson . GetBytes ( rec . Body . Bytes (), "error.message" ). String ())
}
2026-07-29 18:47:07 +08:00
func TestOpenAIResponses_APIKeyPassthroughPoolAuthFailureRetriesThenSwitchesToHealthyAccount ( t * testing . T ) {
tests := [] struct {
name string
statusCode int
}{
{ name : "401" , statusCode : http . StatusUnauthorized },
{ name : "403" , statusCode : http . StatusForbidden },
}
for _ , tt := range tests {
t . Run ( tt . name , func ( t * testing . T ) {
gin . SetMode ( gin . TestMode )
groupID := int64 ( 4203 )
accounts := [] service . Account {
{
ID : 9910 , Name : "pool-api-key" , Platform : service . PlatformOpenAI ,
Type : service . AccountTypeAPIKey , Status : service . StatusActive , Schedulable : true , Priority : 1 ,
Credentials : map [ string ] any {
"api_key" : "sk-pool" ,
"base_url" : "https://api.example.test" ,
"pool_mode" : true ,
"pool_mode_retry_count" : float64 ( 1 ),
"pool_mode_retry_status_codes" : [] any { float64 ( tt . statusCode )},
},
Extra : map [ string ] any { "openai_passthrough" : true },
},
{
ID : 9911 , Name : "fallback-api-key" , Platform : service . PlatformOpenAI ,
Type : service . AccountTypeAPIKey , Status : service . StatusActive , Schedulable : true , Priority : 2 ,
Credentials : map [ string ] any {
"api_key" : "sk-fallback" ,
"base_url" : "https://api.example.test" ,
},
Extra : map [ string ] any { "openai_passthrough" : true },
},
}
cfg := & config . Config { RunMode : config . RunModeSimple }
cfg . Default . RateMultiplier = 1
cfg . Security . URLAllowlist . Enabled = false
cfg . Gateway . MaxAccountSwitches = 1
accountRepo := & openAIWSFailoverHandlerAccountRepoStub { accounts : accounts }
upstream := & openAIHTTPPassthroughAuthFailoverUpstream { statusCode : tt . statusCode }
rateLimitSvc := service . NewRateLimitService ( accountRepo , nil , cfg , nil , nil )
billingCacheSvc := service . NewBillingCacheService ( nil , nil , nil , nil , nil , nil , cfg , nil )
t . Cleanup ( billingCacheSvc . Stop )
gatewaySvc := service . NewOpenAIGatewayService (
accountRepo ,
nil ,
nil ,
nil ,
nil ,
nil ,
nil ,
cfg ,
nil ,
nil ,
service . NewBillingService ( cfg , nil ),
rateLimitSvc ,
billingCacheSvc ,
upstream ,
& service . DeferredService {},
nil ,
nil ,
nil ,
nil ,
nil ,
nil ,
nil ,
)
h := NewOpenAIGatewayHandler (
gatewaySvc ,
service . NewConcurrencyService ( nil ),
billingCacheSvc ,
service . NewAPIKeyService ( nil , nil , nil , nil , nil , nil , cfg ),
nil ,
nil ,
nil ,
nil ,
cfg ,
)
rec := httptest . NewRecorder ()
c , _ := gin . CreateTestContext ( rec )
c . Request = httptest . NewRequest ( http . MethodPost , "/openai/v1/responses" , strings . NewReader ( `{"model":"gpt-5.2","input":"hello","stream":false}` ))
c . Request . Header . Set ( "Content-Type" , "application/json" )
c . Set ( string ( middleware . ContextKeyAPIKey ), & service . APIKey {
ID : 1803 , GroupID : & groupID ,
User : & service . User { ID : 1703 , Status : service . StatusActive },
Group : & service . Group { ID : groupID , Platform : service . PlatformOpenAI , Status : service . StatusActive },
})
c . Set ( string ( middleware . ContextKeyUser ), middleware . AuthSubject { UserID : 1703 , Concurrency : 0 })
h . Responses ( c )
require . Equal ( t , [] int64 { 9910 , 9910 , 9911 }, upstream . calls ())
require . Equal ( t , http . StatusOK , rec . Code )
require . Equal ( t , "resp_healthy" , gjson . GetBytes ( rec . Body . Bytes (), "id" ). String ())
})
}
}
2026-07-30 15:22:16 +08:00
func TestOpenAIResponses_APIKeyPassthroughSSERateLimitUsesConfiguredPoolRetry ( t * testing . T ) {
gin . SetMode ( gin . TestMode )
groupID := int64 ( 4204 )
accounts := [] service . Account {
{
ID : 9912 , Name : "pool-sse-rate-limit" , Platform : service . PlatformOpenAI ,
Type : service . AccountTypeAPIKey , Status : service . StatusActive , Schedulable : true , Priority : 1 ,
Credentials : map [ string ] any {
"api_key" : "sk-pool" ,
"base_url" : "https://api.example.test" ,
"pool_mode" : true ,
"pool_mode_retry_count" : float64 ( 1 ),
"pool_mode_retry_status_codes" : [] any { float64 ( http . StatusTooManyRequests )},
},
Extra : map [ string ] any { "openai_passthrough" : true },
},
}
cfg := & config . Config { RunMode : config . RunModeSimple }
cfg . Default . RateMultiplier = 1
cfg . Security . URLAllowlist . Enabled = false
cfg . Gateway . MaxAccountSwitches = 1
accountRepo := & openAIWSFailoverHandlerAccountRepoStub { accounts : accounts }
upstream := & openAIHTTPPassthroughSSERateLimitUpstream {}
billingCacheSvc := service . NewBillingCacheService ( nil , nil , nil , nil , nil , nil , cfg , nil )
t . Cleanup ( billingCacheSvc . Stop )
gatewaySvc := service . NewOpenAIGatewayService (
accountRepo ,
nil ,
nil ,
nil ,
nil ,
nil ,
nil ,
cfg ,
nil ,
nil ,
service . NewBillingService ( cfg , nil ),
nil ,
billingCacheSvc ,
upstream ,
& service . DeferredService {},
nil ,
nil ,
nil ,
nil ,
nil ,
nil ,
nil ,
)
h := NewOpenAIGatewayHandler (
gatewaySvc ,
service . NewConcurrencyService ( nil ),
billingCacheSvc ,
service . NewAPIKeyService ( nil , nil , nil , nil , nil , nil , cfg ),
nil ,
nil ,
nil ,
nil ,
cfg ,
)
rec := httptest . NewRecorder ()
c , _ := gin . CreateTestContext ( rec )
c . Request = httptest . NewRequest ( http . MethodPost , "/openai/v1/responses" , strings . NewReader ( `{"model":"gpt-5.6-sol","input":"hello","stream":true}` ))
c . Request . Header . Set ( "Content-Type" , "application/json" )
c . Set ( string ( middleware . ContextKeyAPIKey ), & service . APIKey {
ID : 1804 , GroupID : & groupID ,
User : & service . User { ID : 1704 , Status : service . StatusActive },
Group : & service . Group { ID : groupID , Platform : service . PlatformOpenAI , Status : service . StatusActive },
})
c . Set ( string ( middleware . ContextKeyUser ), middleware . AuthSubject { UserID : 1704 , Concurrency : 0 })
h . Responses ( c )
require . Equal ( t , [] int64 { 9912 , 9912 }, upstream . calls ())
require . Equal ( t , http . StatusTooManyRequests , rec . Code )
require . Equal ( t , "1" , rec . Header (). Get ( "Retry-After" ))
require . Equal ( t , "rate_limit_error" , gjson . GetBytes ( rec . Body . Bytes (), "error.type" ). String ())
require . Equal ( t , "Upstream rate limit exceeded, please retry later" , gjson . GetBytes ( rec . Body . Bytes (), "error.message" ). String ())
}
2026-05-26 20:07:00 +08:00
func TestOpenAIResponsesWebSocket_FailoverOnUpstreamUsageLimitEvent ( t * testing . T ) {
gin . SetMode ( gin . TestMode )
firstHitCh := make ( chan [] byte , 1 )
secondHitCh := make ( chan [] byte , 1 )
firstUpstream := httptest . NewServer ( http . HandlerFunc ( func ( w http . ResponseWriter , r * http . Request ) {
conn , err := coderws . Accept ( w , r , & coderws . AcceptOptions { CompressionMode : coderws . CompressionContextTakeover })
if err != nil {
return
}
defer func () { _ = conn . CloseNow () }()
readCtx , cancelRead := context . WithTimeout ( r . Context (), 3 * time . Second )
_ , payload , readErr := conn . Read ( readCtx )
cancelRead ()
if readErr == nil {
firstHitCh <- payload
}
writeCtx , cancelWrite := context . WithTimeout ( r . Context (), 3 * time . Second )
_ = conn . Write ( writeCtx , coderws . MessageText , [] byte ( `{"type":"error","error":{"code":"rate_limit_exceeded","type":"usage_limit_reached","message":"The usage limit has been reached"}}` ))
cancelWrite ()
}))
defer firstUpstream . Close ()
secondUpstream := httptest . NewServer ( http . HandlerFunc ( func ( w http . ResponseWriter , r * http . Request ) {
conn , err := coderws . Accept ( w , r , & coderws . AcceptOptions { CompressionMode : coderws . CompressionContextTakeover })
if err != nil {
return
}
defer func () { _ = conn . CloseNow () }()
readCtx , cancelRead := context . WithTimeout ( r . Context (), 3 * time . Second )
_ , payload , readErr := conn . Read ( readCtx )
cancelRead ()
if readErr == nil {
secondHitCh <- payload
}
writeCtx , cancelWrite := context . WithTimeout ( r . Context (), 3 * time . Second )
_ = conn . Write ( writeCtx , coderws . MessageText , [] byte ( `{"type":"response.completed","response":{"id":"resp_ws_failover_ok","model":"gpt-5.1","usage":{"input_tokens":1,"output_tokens":1}}}` ))
cancelWrite ()
_ = conn . Close ( coderws . StatusNormalClosure , "done" )
}))
defer secondUpstream . Close ()
groupID := int64 ( 4202 )
accounts := [] service . Account {
{
ID : 9902 ,
Name : "openai-ws-rate-limited" ,
Platform : service . PlatformOpenAI ,
Type : service . AccountTypeAPIKey ,
Status : service . StatusActive ,
Schedulable : true ,
Concurrency : 1 ,
Priority : 1 ,
Credentials : map [ string ] any {
"api_key" : "sk-first" ,
"base_url" : firstUpstream . URL ,
},
Extra : map [ string ] any {
"openai_apikey_responses_websockets_v2_enabled" : true ,
"openai_apikey_responses_websockets_v2_mode" : service . OpenAIWSIngressModePassthrough ,
},
},
{
ID : 9903 ,
Name : "openai-ws-healthy" ,
Platform : service . PlatformOpenAI ,
Type : service . AccountTypeAPIKey ,
Status : service . StatusActive ,
Schedulable : true ,
Concurrency : 1 ,
Priority : 2 ,
Credentials : map [ string ] any {
"api_key" : "sk-second" ,
"base_url" : secondUpstream . URL ,
},
Extra : map [ string ] any {
"openai_apikey_responses_websockets_v2_enabled" : true ,
"openai_apikey_responses_websockets_v2_mode" : service . OpenAIWSIngressModePassthrough ,
},
},
}
cfg := & config . Config {}
cfg . RunMode = config . RunModeSimple
cfg . Default . RateMultiplier = 1
cfg . Security . URLAllowlist . Enabled = false
cfg . Security . URLAllowlist . AllowInsecureHTTP = true
cfg . Gateway . OpenAIWS . Enabled = true
cfg . Gateway . OpenAIWS . APIKeyEnabled = true
cfg . Gateway . OpenAIWS . ResponsesWebsocketsV2 = true
cfg . Gateway . OpenAIWS . ModeRouterV2Enabled = true
cfg . Gateway . OpenAIWS . DialTimeoutSeconds = 3
cfg . Gateway . OpenAIWS . ReadTimeoutSeconds = 3
cfg . Gateway . OpenAIWS . WriteTimeoutSeconds = 3
cfg . Gateway . MaxAccountSwitches = 3
accountRepo := & openAIWSFailoverHandlerAccountRepoStub { accounts : accounts }
rateLimitSvc := service . NewRateLimitService ( accountRepo , nil , cfg , nil , nil )
billingCacheSvc := service . NewBillingCacheService ( nil , nil , nil , nil , nil , nil , cfg , nil )
gatewaySvc := service . NewOpenAIGatewayService (
accountRepo ,
nil ,
nil ,
nil ,
nil ,
nil ,
nil ,
cfg ,
nil ,
nil ,
service . NewBillingService ( cfg , nil ),
rateLimitSvc ,
billingCacheSvc ,
nil ,
& service . DeferredService {},
nil ,
nil ,
nil ,
nil ,
nil ,
nil ,
2026-06-16 18:38:07 +08:00
nil ,
2026-05-26 20:07:00 +08:00
)
cache := & concurrencyCacheMock {
acquireUserSlotFn : func ( ctx context . Context , userID int64 , maxConcurrency int , requestID string ) ( bool , error ) {
return true , nil
},
acquireAccountSlotFn : func ( ctx context . Context , accountID int64 , maxConcurrency int , requestID string ) ( bool , error ) {
return true , nil
},
}
h := & OpenAIGatewayHandler {
gatewayService : gatewaySvc ,
billingCacheService : billingCacheSvc ,
apiKeyService : & service . APIKeyService {},
concurrencyHelper : NewConcurrencyHelper ( service . NewConcurrencyService ( cache ), SSEPingFormatNone , time . Second ),
maxAccountSwitches : 3 ,
}
apiKey := & service . APIKey {
ID : 1802 ,
GroupID : & groupID ,
User : & service . User { ID : 1702 , Status : service . StatusActive },
Group : & service . Group { ID : groupID , Platform : service . PlatformOpenAI , Status : service . StatusActive },
}
router := gin . New ()
router . Use ( func ( c * gin . Context ) {
c . Set ( string ( middleware . ContextKeyAPIKey ), apiKey )
c . Set ( string ( middleware . ContextKeyUser ), middleware . AuthSubject { UserID : apiKey . User . ID , Concurrency : 1 })
c . Next ()
})
router . GET ( "/openai/v1/responses" , h . ResponsesWebSocket )
handlerServer := httptest . NewServer ( router )
defer handlerServer . Close ()
dialCtx , cancelDial := context . WithTimeout ( context . Background (), 3 * time . Second )
clientConn , _ , err := coderws . Dial (
dialCtx ,
"ws" + strings . TrimPrefix ( handlerServer . URL , "http" ) + "/openai/v1/responses" ,
& coderws . DialOptions { CompressionMode : coderws . CompressionContextTakeover },
)
cancelDial ()
require . NoError ( t , err )
defer func () { _ = clientConn . CloseNow () }()
writeCtx , cancelWrite := context . WithTimeout ( context . Background (), 3 * time . Second )
err = clientConn . Write ( writeCtx , coderws . MessageText , [] byte ( `{"type":"response.create","model":"gpt-5.1","stream":false}` ))
cancelWrite ()
require . NoError ( t , err )
readCtx , cancelRead := context . WithTimeout ( context . Background (), 5 * time . Second )
_ , event , err := clientConn . Read ( readCtx )
cancelRead ()
require . NoError ( t , err )
require . Equal ( t , "response.completed" , gjson . GetBytes ( event , "type" ). String ())
require . Equal ( t , "resp_ws_failover_ok" , gjson . GetBytes ( event , "response.id" ). String ())
select {
case <- firstHitCh :
case <- time . After ( 3 * time . Second ):
t . Fatal ( "等待第一个上游收到首帧超时" )
}
select {
case <- secondHitCh :
case <- time . After ( 3 * time . Second ):
t . Fatal ( "等待第二个上游收到重放首帧超时" )
}
require . Equal ( t , [] int64 { int64 ( 9902 )}, accountRepo . rateLimitedIDs )
}
2026-07-15 22:44:19 +08:00
func TestOpenAIResponsesWebSocket_FirstOutputTimeoutWithoutDownstreamReusesClientForOneFailover ( t * testing . T ) {
gin . SetMode ( gin . TestMode )
firstHitCh := make ( chan [] byte , 1 )
secondHitCh := make ( chan [] byte , 1 )
var firstConnections atomic . Int32
var secondConnections atomic . Int32
firstUpstream := httptest . NewServer ( http . HandlerFunc ( func ( w http . ResponseWriter , r * http . Request ) {
firstConnections . Add ( 1 )
conn , err := coderws . Accept ( w , r , & coderws . AcceptOptions { CompressionMode : coderws . CompressionContextTakeover })
if err != nil {
return
}
defer func () { _ = conn . CloseNow () }()
readCtx , cancelRead := context . WithTimeout ( r . Context (), 3 * time . Second )
_ , payload , readErr := conn . Read ( readCtx )
cancelRead ()
if readErr == nil {
firstHitCh <- payload
}
select {
case <- r . Context (). Done ():
case <- time . After ( 3 * time . Second ):
}
}))
defer firstUpstream . Close ()
secondUpstream := httptest . NewServer ( http . HandlerFunc ( func ( w http . ResponseWriter , r * http . Request ) {
secondConnections . Add ( 1 )
conn , err := coderws . Accept ( w , r , & coderws . AcceptOptions { CompressionMode : coderws . CompressionContextTakeover })
if err != nil {
return
}
defer func () { _ = conn . CloseNow () }()
readCtx , cancelRead := context . WithTimeout ( r . Context (), 3 * time . Second )
_ , payload , readErr := conn . Read ( readCtx )
cancelRead ()
if readErr == nil {
secondHitCh <- payload
}
for _ , event := range [] string {
`{"type":"response.created","response":{"id":"resp_ws_timeout_b","model":"gpt-5.1"}}` ,
`{"type":"response.output_text.delta","response_id":"resp_ws_timeout_b","delta":"recovered"}` ,
`{"type":"response.completed","response":{"id":"resp_ws_timeout_b","model":"gpt-5.1","usage":{"input_tokens":1,"output_tokens":1}}}` ,
} {
writeCtx , cancelWrite := context . WithTimeout ( r . Context (), 3 * time . Second )
writeErr := conn . Write ( writeCtx , coderws . MessageText , [] byte ( event ))
cancelWrite ()
if writeErr != nil {
return
}
}
readCtx , cancelRead = context . WithTimeout ( r . Context (), 3 * time . Second )
_ , _ , _ = conn . Read ( readCtx )
cancelRead ()
}))
defer secondUpstream . Close ()
groupID := int64 ( 4212 )
accounts := [] service . Account {
{
ID : 9912 ,
Name : "openai-ws-first-semantic-timeout" ,
Platform : service . PlatformOpenAI ,
Type : service . AccountTypeAPIKey ,
Status : service . StatusActive ,
Schedulable : true ,
Concurrency : 1 ,
Priority : 1 ,
Credentials : map [ string ] any { "api_key" : "sk-first" , "base_url" : firstUpstream . URL },
Extra : map [ string ] any {
"openai_apikey_responses_websockets_v2_enabled" : true ,
"openai_apikey_responses_websockets_v2_mode" : service . OpenAIWSIngressModePassthrough ,
},
},
{
ID : 9913 ,
Name : "openai-ws-failover-healthy" ,
Platform : service . PlatformOpenAI ,
Type : service . AccountTypeAPIKey ,
Status : service . StatusActive ,
Schedulable : true ,
Concurrency : 1 ,
Priority : 2 ,
Credentials : map [ string ] any { "api_key" : "sk-second" , "base_url" : secondUpstream . URL },
Extra : map [ string ] any {
"openai_apikey_responses_websockets_v2_enabled" : true ,
"openai_apikey_responses_websockets_v2_mode" : service . OpenAIWSIngressModePassthrough ,
},
},
}
cfg := & config . Config {}
cfg . RunMode = config . RunModeSimple
cfg . Default . RateMultiplier = 1
cfg . Security . URLAllowlist . Enabled = false
cfg . Security . URLAllowlist . AllowInsecureHTTP = true
cfg . Gateway . OpenAIFirstOutputTimeoutSeconds = 1
cfg . Gateway . OpenAIWS . Enabled = true
cfg . Gateway . OpenAIWS . APIKeyEnabled = true
cfg . Gateway . OpenAIWS . ResponsesWebsocketsV2 = true
cfg . Gateway . OpenAIWS . ModeRouterV2Enabled = true
cfg . Gateway . OpenAIWS . DialTimeoutSeconds = 3
cfg . Gateway . OpenAIWS . ReadTimeoutSeconds = 3
cfg . Gateway . OpenAIWS . WriteTimeoutSeconds = 3
cfg . Gateway . OpenAIWS . IngressInterTurnIdleTimeoutSeconds = 3
cfg . Gateway . MaxAccountSwitches = 3
accountRepo := & openAIWSFailoverHandlerAccountRepoStub { accounts : accounts }
rateLimitSvc := service . NewRateLimitService ( accountRepo , nil , cfg , nil , nil )
billingCacheSvc := service . NewBillingCacheService ( nil , nil , nil , nil , nil , nil , cfg , nil )
gatewaySvc := service . NewOpenAIGatewayService (
accountRepo , nil , nil , nil , nil , nil , nil , cfg , nil , nil ,
service . NewBillingService ( cfg , nil ), rateLimitSvc , billingCacheSvc ,
nil , & service . DeferredService {}, nil , nil , nil , nil , nil , nil , nil ,
)
cache := & concurrencyCacheMock {
acquireUserSlotFn : func ( context . Context , int64 , int , string ) ( bool , error ) { return true , nil },
acquireAccountSlotFn : func ( context . Context , int64 , int , string ) ( bool , error ) {
return true , nil
},
}
h := & OpenAIGatewayHandler {
gatewayService : gatewaySvc ,
billingCacheService : billingCacheSvc ,
apiKeyService : & service . APIKeyService {},
concurrencyHelper : NewConcurrencyHelper ( service . NewConcurrencyService ( cache ), SSEPingFormatNone , time . Second ),
maxAccountSwitches : 3 ,
}
apiKey := & service . APIKey {
ID : 1812 ,
GroupID : & groupID ,
User : & service . User { ID : 1712 , Status : service . StatusActive },
Group : & service . Group { ID : groupID , Platform : service . PlatformOpenAI , Status : service . StatusActive },
}
handlerDone := make ( chan struct {})
router := gin . New ()
router . Use ( func ( c * gin . Context ) {
c . Set ( string ( middleware . ContextKeyAPIKey ), apiKey )
c . Set ( string ( middleware . ContextKeyUser ), middleware . AuthSubject { UserID : apiKey . User . ID , Concurrency : 1 })
c . Next ()
})
router . GET ( "/openai/v1/responses" , func ( c * gin . Context ) {
h . ResponsesWebSocket ( c )
close ( handlerDone )
})
handlerServer := httptest . NewServer ( router )
defer handlerServer . Close ()
dialCtx , cancelDial := context . WithTimeout ( context . Background (), 3 * time . Second )
clientConn , _ , err := coderws . Dial (
dialCtx ,
"ws" + strings . TrimPrefix ( handlerServer . URL , "http" ) + "/openai/v1/responses" ,
& coderws . DialOptions { CompressionMode : coderws . CompressionContextTakeover },
)
cancelDial ()
require . NoError ( t , err )
defer func () { _ = clientConn . CloseNow () }()
writeCtx , cancelWrite := context . WithTimeout ( context . Background (), 3 * time . Second )
err = clientConn . Write ( writeCtx , coderws . MessageText , [] byte ( `{"type":"response.create","model":"gpt-5.1","stream":false}` ))
cancelWrite ()
require . NoError ( t , err )
var eventTypes [] string
readCtx , cancelRead := context . WithTimeout ( context . Background (), 6 * time . Second )
for {
_ , event , readErr := clientConn . Read ( readCtx )
require . NoError ( t , readErr )
eventType := gjson . GetBytes ( event , "type" ). String ()
eventTypes = append ( eventTypes , eventType )
if eventType == "response.completed" {
require . Equal ( t , "resp_ws_timeout_b" , gjson . GetBytes ( event , "response.id" ). String ())
break
}
}
cancelRead ()
require . Contains ( t , eventTypes , "response.output_text.delta" )
require . NoError ( t , clientConn . Close ( coderws . StatusNormalClosure , "done" ))
select {
case <- handlerDone :
case <- time . After ( 3 * time . Second ):
t . Fatal ( "websocket handler did not finish after healthy failover turn" )
}
select {
case <- firstHitCh :
case <- time . After ( 3 * time . Second ):
t . Fatal ( "first upstream did not receive replayable request" )
}
select {
case <- secondHitCh :
case <- time . After ( 3 * time . Second ):
t . Fatal ( "second upstream did not receive replayed request" )
}
require . Equal ( t , int32 ( 1 ), firstConnections . Load ())
require . Equal ( t , int32 ( 1 ), secondConnections . Load ())
require . NotContains ( t , accountRepo . rateLimitedIDs , int64 ( 9913 ), "healthy failover account must not be penalized" )
}
2026-05-03 19:33:09 +08:00
func runOpenAIResponsesWebSocketUsageLogCase ( t * testing . T , tc openAIResponsesWSUsageLogCase ) openAIResponsesWSUsageLogResult {
t . Helper ()
gin . SetMode ( gin . TestMode )
2026-07-23 20:17:11 +08:00
turnCount := 1
2026-09-05 15:58:26 +08:00
if strings . TrimSpace ( tc . midPayload ) != "" {
turnCount ++
}
2026-07-23 20:17:11 +08:00
if strings . TrimSpace ( tc . secondPayload ) != "" {
2026-09-05 15:58:26 +08:00
turnCount ++
2026-07-23 20:17:11 +08:00
}
upstreamPayloadCh := make ( chan [] byte , turnCount )
2026-05-03 19:33:09 +08:00
upstreamErrCh := make ( chan error , 1 )
2026-07-23 20:17:11 +08:00
var channelSvc * service . ChannelService
2026-05-03 19:33:09 +08:00
upstreamServer := httptest . NewServer ( http . HandlerFunc ( func ( w http . ResponseWriter , r * http . Request ) {
2026-09-19 20:50:28 +08:00
if tc . accountPlatform == service . PlatformGrok {
payload , err := io . ReadAll ( r . Body )
if err != nil {
upstreamErrCh <- err
return
}
upstreamPayloadCh <- payload
w . Header (). Set ( "Content-Type" , "text/event-stream" )
_ , _ = fmt . Fprintf ( w , "data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_grok_test\",\"model\":%q,\"usage\":{\"input_tokens\":2,\"output_tokens\":1}}}\n\n" , gjson . GetBytes ( payload , "model" ). String ())
return
}
2026-05-03 19:33:09 +08:00
conn , err := coderws . Accept ( w , r , & coderws . AcceptOptions {
CompressionMode : coderws . CompressionContextTakeover ,
})
if err != nil {
upstreamErrCh <- err
return
}
defer func () {
_ = conn . CloseNow ()
}()
2026-07-23 20:17:11 +08:00
for turn := 1 ; turn <= turnCount ; turn ++ {
readCtx , cancelRead := context . WithTimeout ( r . Context (), 3 * time . Second )
msgType , payload , readErr := conn . Read ( readCtx )
cancelRead ()
if readErr != nil {
upstreamErrCh <- readErr
return
}
if msgType != coderws . MessageText && msgType != coderws . MessageBinary {
upstreamErrCh <- errors . New ( "unexpected upstream websocket message type" )
return
}
upstreamPayloadCh <- payload
if turn == 1 && tc . afterFirstUpstreamRequest != nil {
if callbackErr := tc . afterFirstUpstreamRequest ( channelSvc ); callbackErr != nil {
upstreamErrCh <- callbackErr
return
}
}
2026-05-03 19:33:09 +08:00
2026-07-23 20:17:11 +08:00
response := fmt . Sprintf (
`{"type":"response.completed","response":{"id":"resp_usage_e2e_%d","model":%q,"usage":{"input_tokens":2,"output_tokens":1}}}` ,
turn ,
gjson . GetBytes ( payload , "model" ). String (),
)
writeCtx , cancelWrite := context . WithTimeout ( r . Context (), 3 * time . Second )
writeErr := conn . Write ( writeCtx , coderws . MessageText , [] byte ( response ))
cancelWrite ()
if writeErr != nil {
upstreamErrCh <- writeErr
return
}
2026-05-03 19:33:09 +08:00
}
upstreamErrCh <- nil
}))
defer upstreamServer . Close ()
groupID := int64 ( 4201 )
account := service . Account {
ID : 9901 ,
Name : "openai-ws-passthrough-usage-e2e" ,
Platform : service . PlatformOpenAI ,
Type : service . AccountTypeAPIKey ,
Status : service . StatusActive ,
Schedulable : true ,
Concurrency : 1 ,
Credentials : map [ string ] any {
2026-07-23 20:17:11 +08:00
"api_key" : "sk-test" ,
"base_url" : upstreamServer . URL ,
"model_mapping" : tc . accountModelMapping ,
2026-05-03 19:33:09 +08:00
},
Extra : map [ string ] any {
"openai_apikey_responses_websockets_v2_enabled" : true ,
"openai_apikey_responses_websockets_v2_mode" : service . OpenAIWSIngressModePassthrough ,
},
}
2026-09-19 20:50:28 +08:00
if tc . accountPlatform != "" {
account . Platform = tc . accountPlatform
}
2026-10-01 11:27:35 +08:00
account . RateMultiplier = tc . accountRateMultiplier
2026-07-23 20:17:11 +08:00
if strings . TrimSpace ( tc . ingressMode ) != "" {
account . Extra [ "openai_apikey_responses_websockets_v2_mode" ] = tc . ingressMode
}
2026-05-03 19:33:09 +08:00
cfg := & config . Config {}
cfg . RunMode = config . RunModeSimple
cfg . Default . RateMultiplier = 1
cfg . Security . URLAllowlist . Enabled = false
cfg . Security . URLAllowlist . AllowInsecureHTTP = true
cfg . Gateway . OpenAIWS . Enabled = true
cfg . Gateway . OpenAIWS . APIKeyEnabled = true
cfg . Gateway . OpenAIWS . ResponsesWebsocketsV2 = true
cfg . Gateway . OpenAIWS . ModeRouterV2Enabled = true
cfg . Gateway . OpenAIWS . DialTimeoutSeconds = 3
cfg . Gateway . OpenAIWS . ReadTimeoutSeconds = 3
cfg . Gateway . OpenAIWS . WriteTimeoutSeconds = 3
accountRepo := & openAIWSUsageHandlerAccountRepoStub { account : account }
2026-07-23 20:17:11 +08:00
usageRepo := & openAIWSUsageHandlerUsageLogRepoStub { created : make ( chan * service . UsageLog , turnCount )}
2026-05-03 19:33:09 +08:00
if len ( tc . channelMapping ) > 0 {
channelSvc = service . NewChannelService ( & openAIWSUsageHandlerChannelRepoStub {
channels : [] service . Channel {{
2026-07-23 20:17:11 +08:00
ID : 7701 ,
Name : "openai-ws-e2e-channel" ,
Status : service . StatusActive ,
GroupIDs : [] int64 { groupID },
ModelMapping : map [ string ] map [ string ] string { service . PlatformOpenAI : tc . channelMapping },
BillingModelSource : tc . billingModelSource ,
2026-05-03 19:33:09 +08:00
}},
groupPlatforms : map [ int64 ] string { groupID : service . PlatformOpenAI },
2026-08-26 11:16:12 +08:00
}, nil , nil , nil , nil )
2026-05-03 19:33:09 +08:00
}
2026-09-22 15:23:09 +08:00
var keyRepo service . APIKeyRepository
if tc . simpleModeRejectAtRead > 0 {
cfg . SimpleModeKeyRateLimitEnabled = true
keyRepo = & simpleModeWSRateLimitRepo { rejectAt : tc . simpleModeRejectAtRead }
}
billingCacheSvc := service . NewBillingCacheService ( nil , nil , nil , keyRepo , nil , nil , cfg , nil )
t . Cleanup ( billingCacheSvc . Stop )
2026-05-03 19:33:09 +08:00
gatewaySvc := service . NewOpenAIGatewayService (
accountRepo ,
usageRepo ,
nil ,
nil ,
nil ,
nil ,
nil ,
cfg ,
nil ,
nil ,
service . NewBillingService ( cfg , nil ),
nil ,
billingCacheSvc ,
2026-09-19 20:50:28 +08:00
& compositeWSHTTPUpstream {},
2026-05-03 19:33:09 +08:00
& service . DeferredService {},
nil ,
nil ,
2026-06-16 18:38:07 +08:00
nil ,
2026-05-03 19:33:09 +08:00
channelSvc ,
nil ,
nil ,
2026-05-25 14:26:11 +08:00
nil , // userPlatformQuotaRepo
2026-05-03 19:33:09 +08:00
)
cache := & concurrencyCacheMock {
acquireUserSlotFn : func ( ctx context . Context , userID int64 , maxConcurrency int , requestID string ) ( bool , error ) {
return true , nil
},
acquireAccountSlotFn : func ( ctx context . Context , accountID int64 , maxConcurrency int , requestID string ) ( bool , error ) {
return true , nil
},
}
h := & OpenAIGatewayHandler {
2026-09-22 15:23:09 +08:00
cfg : cfg ,
2026-09-19 20:50:28 +08:00
compositeResolver : tc . compositeResolver ,
2026-05-03 19:33:09 +08:00
gatewayService : gatewaySvc ,
billingCacheService : billingCacheSvc ,
apiKeyService : & service . APIKeyService {},
concurrencyHelper : NewConcurrencyHelper ( service . NewConcurrencyService ( cache ), SSEPingFormatNone , time . Second ),
}
apiKey := & service . APIKey {
ID : 1801 ,
GroupID : & groupID ,
User : & service . User { ID : 1701 , Status : service . StatusActive },
}
2026-10-01 11:27:35 +08:00
if tc . apiKeyService != nil {
h . apiKeyService = tc . apiKeyService
authKey , err := tc . apiKeyService . GetByKey ( context . Background (), tc . apiKeyCredential )
require . NoError ( t , err )
apiKey = authKey
}
2026-09-22 15:23:09 +08:00
if tc . simpleModeRejectAtRead > 0 {
apiKey . RateLimit5h = 1
}
2026-09-05 15:58:26 +08:00
if tc . group != nil {
apiKey . Group = tc . group
}
2026-05-03 19:33:09 +08:00
router := gin . New ()
router . Use ( func ( c * gin . Context ) {
c . Set ( string ( middleware . ContextKeyAPIKey ), apiKey )
c . Set ( string ( middleware . ContextKeyUser ), middleware . AuthSubject { UserID : apiKey . User . ID , Concurrency : 1 })
2026-10-01 11:27:35 +08:00
if tc . apiKeyService != nil && apiKey . Group != nil {
c . Request = c . Request . WithContext ( context . WithValue ( c . Request . Context (), ctxkey . Group , apiKey . Group ))
}
2026-05-03 19:33:09 +08:00
c . Next ()
})
router . GET ( "/openai/v1/responses" , h . ResponsesWebSocket )
handlerServer := httptest . NewServer ( router )
defer handlerServer . Close ()
headers := http . Header {}
if tc . userAgent != nil {
headers . Set ( "User-Agent" , * tc . userAgent )
}
dialCtx , cancelDial := context . WithTimeout ( context . Background (), 3 * time . Second )
clientConn , _ , err := coderws . Dial (
dialCtx ,
"ws" + strings . TrimPrefix ( handlerServer . URL , "http" ) + "/openai/v1/responses" ,
& coderws . DialOptions { HTTPHeader : headers , CompressionMode : coderws . CompressionContextTakeover },
)
cancelDial ()
require . NoError ( t , err )
defer func () {
_ = clientConn . CloseNow ()
}()
writeCtx , cancelWrite := context . WithTimeout ( context . Background (), 3 * time . Second )
err = clientConn . Write ( writeCtx , coderws . MessageText , [] byte ( tc . firstPayload ))
cancelWrite ()
require . NoError ( t , err )
2026-09-19 20:50:28 +08:00
if tc . closeReason == "" {
tc . closeReason = "not available for this group"
}
if tc . closeStatus == 0 {
tc . closeStatus = coderws . StatusPolicyViolation
}
2026-09-05 15:58:26 +08:00
if tc . firstFrameCloseExpected {
readCtx , cancelRead := context . WithTimeout ( context . Background (), 3 * time . Second )
_ , _ , readErr := clientConn . Read ( readCtx )
cancelRead ()
require . Error ( t , readErr , "first frame should have been rejected with a close" )
var closeErr coderws . CloseError
require . ErrorAs ( t , readErr , & closeErr )
2026-09-19 20:50:28 +08:00
status := tc . closeStatus
if status == 0 {
status = coderws . StatusPolicyViolation
}
require . Equal ( t , status , closeErr . Code )
2026-09-22 15:23:09 +08:00
reason := tc . closeReason
if reason == "" {
reason = "not available for this group"
}
require . Contains ( t , closeErr . Reason , reason )
2026-09-19 20:50:28 +08:00
require . Empty ( t , upstreamPayloadCh , "rejected first frame must not reach upstream" )
2026-09-05 15:58:26 +08:00
_ = clientConn . CloseNow ()
return openAIResponsesWSUsageLogResult {}
}
2026-07-23 20:17:11 +08:00
clientEvents := make ([][] byte , 0 , turnCount )
readCompleted := func () {
readCtx , cancelRead := context . WithTimeout ( context . Background (), 3 * time . Second )
_ , event , readErr := clientConn . Read ( readCtx )
cancelRead ()
require . NoError ( t , readErr )
require . Equal ( t , "response.completed" , gjson . GetBytes ( event , "type" ). String ())
clientEvents = append ( clientEvents , append ([] byte ( nil ), event ... ))
}
readCompleted ()
2026-09-05 15:58:26 +08:00
if tc . midPayload != "" {
writeCtx , cancelWrite = context . WithTimeout ( context . Background (), 3 * time . Second )
err = clientConn . Write ( writeCtx , coderws . MessageText , [] byte ( tc . midPayload ))
cancelWrite ()
require . NoError ( t , err )
readCompleted ()
}
if strings . TrimSpace ( tc . secondPayload ) != "" && ( turnCount >= 2 ) {
2026-07-23 20:17:11 +08:00
writeCtx , cancelWrite = context . WithTimeout ( context . Background (), 3 * time . Second )
err = clientConn . Write ( writeCtx , coderws . MessageText , [] byte ( tc . secondPayload ))
cancelWrite ()
require . NoError ( t , err )
2026-09-05 15:58:26 +08:00
if tc . secondTurnCloseExpected {
readCtx , cancelRead := context . WithTimeout ( context . Background (), 3 * time . Second )
_ , _ , readErr := clientConn . Read ( readCtx )
cancelRead ()
require . Error ( t , readErr , "second turn should have been rejected with a close" )
var closeErr coderws . CloseError
require . ErrorAs ( t , readErr , & closeErr )
2026-09-19 20:50:28 +08:00
status := tc . closeStatus
if status == 0 {
status = coderws . StatusPolicyViolation
}
require . Equal ( t , status , closeErr . Code )
2026-09-22 15:23:09 +08:00
reason := tc . closeReason
if reason == "" {
reason = "not available for this group"
}
require . Contains ( t , closeErr . Reason , reason )
2026-09-19 20:50:28 +08:00
require . Len ( t , upstreamPayloadCh , turnCount - 1 , "rejected turn must not reach upstream" )
2026-09-05 15:58:26 +08:00
_ = clientConn . CloseNow ()
return openAIResponsesWSUsageLogResult {}
}
2026-07-23 20:17:11 +08:00
readCompleted ()
}
2026-05-03 19:33:09 +08:00
_ = clientConn . Close ( coderws . StatusNormalClosure , "done" )
2026-07-23 20:17:11 +08:00
usageLogs := make ([] * service . UsageLog , 0 , turnCount )
for len ( usageLogs ) < turnCount {
select {
case usageLog := <- usageRepo . created :
require . NotNil ( t , usageLog )
usageLogs = append ( usageLogs , usageLog )
case <- time . After ( 3 * time . Second ):
t . Fatal ( "等待 WebSocket usage log 写入超时" )
}
2026-05-03 19:33:09 +08:00
}
2026-07-23 20:17:11 +08:00
upstreamPayloads := make ([][] byte , 0 , turnCount )
for len ( upstreamPayloads ) < turnCount {
select {
case payload := <- upstreamPayloadCh :
upstreamPayloads = append ( upstreamPayloads , payload )
case <- time . After ( 3 * time . Second ):
t . Fatal ( "等待上游 WebSocket 请求帧超时" )
}
2026-05-03 19:33:09 +08:00
}
2026-09-19 20:50:28 +08:00
if tc . accountPlatform == service . PlatformGrok {
upstreamErrCh <- nil
}
2026-05-03 19:33:09 +08:00
select {
case upstreamErr := <- upstreamErrCh :
require . NoError ( t , upstreamErr )
case <- time . After ( 3 * time . Second ):
t . Fatal ( "等待上游 WebSocket 结束超时" )
}
return openAIResponsesWSUsageLogResult {
2026-07-23 20:17:11 +08:00
log : usageLogs [ 0 ],
logs : usageLogs ,
upstreamFirstPayload : upstreamPayloads [ 0 ],
upstreamPayloads : upstreamPayloads ,
clientEvents : clientEvents ,
2026-05-03 19:33:09 +08:00
}
}
func testStringPtr ( v string ) * string {
return & v
}
2026-06-04 11:17:22 +00:00
func TestOpenAIForwardErrorAlreadyCommunicated ( t * testing . T ) {
gin . SetMode ( gin . TestMode )
t . Run ( "upstream response failed after write" , func ( t * testing . T ) {
w := httptest . NewRecorder ()
c , _ := gin . CreateTestContext ( w )
c . Request = httptest . NewRequest ( http . MethodPost , EndpointResponses , nil )
before := c . Writer . Size ()
_ , _ = c . Writer . WriteString ( `event: response.failed
data: {"type":"response.failed","error":{"message":"This content was flagged"}}
` )
reported := openAIForwardErrorAlreadyCommunicated ( c , before , errors . New ( "upstream response failed: This content was flagged" ))
require . True ( t , reported )
})
t . Run ( "no write still needs fallback" , func ( t * testing . T ) {
w := httptest . NewRecorder ()
c , _ := gin . CreateTestContext ( w )
c . Request = httptest . NewRequest ( http . MethodPost , EndpointResponses , nil )
reported := openAIForwardErrorAlreadyCommunicated ( c , c . Writer . Size (), errors . New ( "upstream response failed: This content was flagged" ))
require . False ( t , reported )
})
t . Run ( "generic error after write still needs fallback" , func ( t * testing . T ) {
w := httptest . NewRecorder ()
c , _ := gin . CreateTestContext ( w )
c . Request = httptest . NewRequest ( http . MethodPost , EndpointResponses , nil )
before := c . Writer . Size ()
_ , _ = c . Writer . WriteString ( ":\n\n" )
reported := openAIForwardErrorAlreadyCommunicated ( c , before , errors . New ( "stream read error: unexpected EOF" ))
require . False ( t , reported )
})
2026-06-11 17:26:14 +08:00
// H-2: cyber_policy 命中且响应已写出时,即便 err 前缀不在白名单(非流式 400 cyber
// 返回 "openai cyber_policy:"、透传账号返回 "upstream error:"),也须判定已透传,避免
// ensureForwardErrorResponse 在已写出的完整响应尾部追加 SSE 污染响应体。
t . Run ( "cyber policy hit after write is already communicated" , func ( t * testing . T ) {
w := httptest . NewRecorder ()
c , _ := gin . CreateTestContext ( w )
c . Request = httptest . NewRequest ( http . MethodPost , EndpointResponses , nil )
service . MarkOpsCyberPolicy ( c , service . CyberPolicyMark { Message : "blocked" , UpstreamStatus : 400 })
before := c . Writer . Size ()
_ , _ = c . Writer . WriteString ( `{"error":{"code":"cyber_policy","message":"blocked"}}` )
require . True ( t , openAIForwardErrorAlreadyCommunicated ( c , before , errors . New ( "openai cyber_policy: blocked" )))
})
// Size 守卫优先于 cyber 短路:cyber 命中但未写出任何响应时仍需补写错误。
t . Run ( "cyber policy without write still needs fallback" , func ( t * testing . T ) {
w := httptest . NewRecorder ()
c , _ := gin . CreateTestContext ( w )
c . Request = httptest . NewRequest ( http . MethodPost , EndpointResponses , nil )
service . MarkOpsCyberPolicy ( c , service . CyberPolicyMark { Message : "blocked" , UpstreamStatus : 400 })
require . False ( t , openAIForwardErrorAlreadyCommunicated ( c , c . Writer . Size (), errors . New ( "openai cyber_policy: blocked" )))
})
2026-06-04 11:17:22 +00:00
}
2026-09-22 15:23:09 +08:00
// The loader simulates window usage reaching its limit while a connection is
// waiting for its first account or between completed WebSocket turns.
type simpleModeWSRateLimitRepo struct {
service . APIKeyRepository
reads atomic . Int64
rejectAt int64
}
func ( r * simpleModeWSRateLimitRepo ) GetRateLimitData ( context . Context , int64 ) ( * service . APIKeyRateLimitData , error ) {
now := time . Now ()
usage := 0.0
if r . reads . Add ( 1 ) >= r . rejectAt {
usage = 1
}
return & service . APIKeyRateLimitData { Usage5h : usage , Window5hStart : & now }, nil
}
func TestOpenAIResponsesWebSocketSimpleModeRechecksKeyWindows ( t * testing . T ) {
for _ , mode := range [] string { service . OpenAIWSIngressModePassthrough , service . OpenAIWSIngressModeCtxPool } {
t . Run ( mode + "/after-account-selection" , func ( t * testing . T ) {
runOpenAIResponsesWebSocketUsageLogCase ( t , openAIResponsesWSUsageLogCase {
firstPayload : `{"type":"response.create","model":"gpt-5.1"}` ,
ingressMode : mode , simpleModeRejectAtRead : 2 , firstFrameCloseExpected : true , closeReason : "billing check failed" ,
})
})
t . Run ( mode + "/followup-turn" , func ( t * testing . T ) {
runOpenAIResponsesWebSocketUsageLogCase ( t , openAIResponsesWSUsageLogCase {
firstPayload : `{"type":"response.create","model":"gpt-5.1"}` ,
secondPayload : `{"type":"response.create","model":"gpt-5.1"}` ,
ingressMode : mode , simpleModeRejectAtRead : 3 , secondTurnCloseExpected : true , closeReason : "billing check failed" ,
})
})
}
}