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-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-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-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"}]}` )
sessionHash , promptCacheKey := resolveOpenAIMessagesMetadataSession ( "" , "" , "claude-sonnet-4-5" , body )
require . NotEmpty ( t , sessionHash )
require . Empty ( t , promptCacheKey )
}
func TestResolveOpenAIMessagesMetadataSession_PreservesExplicitPromptCacheKey ( t * testing . T ) {
body := [] byte ( `{"metadata":{"user_id":"claude-code-session"}}` )
sessionHash , promptCacheKey := resolveOpenAIMessagesMetadataSession ( "" , "explicit-cache" , "claude-sonnet-4-5" , body )
require . NotEmpty ( t , sessionHash )
require . Equal ( t , "explicit-cache" , promptCacheKey )
}
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-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
},
},
},
}
require . Equal ( t , "gpt-5.4-mini" , resolveOpenAIMessagesDispatchMappedModel ( apiKey , "claude-sonnet-4-5-20250929" ))
2026-07-13 17:57:08 +08:00
require . Equal ( t , "gpt-5.6-sol" , resolveOpenAIMessagesDispatchMappedModel ( 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 {}}
require . Equal ( t , "gpt-5.4" , resolveOpenAIMessagesDispatchMappedModel ( apiKey , "claude-opus-4-6" ))
require . Equal ( t , "gpt-5.3-codex" , resolveOpenAIMessagesDispatchMappedModel ( apiKey , "claude-sonnet-4-5-20250929" ))
require . Equal ( t , "gpt-5.4-mini" , resolveOpenAIMessagesDispatchMappedModel ( apiKey , "claude-haiku-4-5-20251001" ))
})
t . Run ( "returns_empty_for_non_claude_or_missing_group" , func ( t * testing . T ) {
require . Empty ( t , resolveOpenAIMessagesDispatchMappedModel ( nil , "claude-sonnet-4-5-20250929" ))
require . Empty ( t , resolveOpenAIMessagesDispatchMappedModel ( & service . APIKey {}, "claude-sonnet-4-5-20250929" ))
require . Empty ( t , resolveOpenAIMessagesDispatchMappedModel ( & service . APIKey { Group : & service . Group {}}, "gpt-5.4" ))
})
2026-06-29 18:33:37 +08:00
t . Run ( "grok_group_maps_claude_cli_model_to_grok_default" , func ( t * testing . T ) {
apiKey := & service . APIKey {
Group : & service . Group {
Platform : service . PlatformGrok ,
},
}
2026-07-09 14:26:10 +08:00
require . Equal ( t , "grok-4.5" , resolveOpenAIMessagesDispatchMappedModel ( apiKey , "claude-sonnet-4-5" ))
2026-06-29 18:33:37 +08:00
require . Empty ( t , resolveOpenAIMessagesDispatchMappedModel ( apiKey , "grok" ))
})
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" ,
},
}
require . Empty ( t , resolveOpenAIMessagesDispatchMappedModel ( apiKey , "gpt-5.4" ))
require . Equal ( t , "gpt-5.3-codex" , resolveOpenAIMessagesDispatchMappedModel ( apiKey , "claude-sonnet-4-5-20250929" ))
})
}
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-04-21 13:53:12 +08:00
func TestOpenAIResponses_RejectsHTTPContinuationPreviousResponseID ( 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_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 (), "Responses WebSocket v2" )
require . Contains ( t , w . Body . String (), "previous_response_id" )
}
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 ,
)
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-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 {
firstPayload string
userAgent * string
channelMapping map [ string ] string
}
type openAIResponsesWSUsageLogResult struct {
log * service . UsageLog
upstreamFirstPayload [] byte
}
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-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
channels [] service . Channel
groupPlatforms map [ int64 ] string
}
func ( s * openAIWSUsageHandlerChannelRepoStub ) ListAll ( ctx context . Context ) ([] service . Channel , error ) {
return s . channels , nil
}
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-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 )
upstreamPayloadCh := make ( chan [] byte , 1 )
upstreamErrCh := make ( chan error , 1 )
upstreamServer := 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 {
upstreamErrCh <- err
return
}
defer func () {
_ = conn . CloseNow ()
}()
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
writeCtx , cancelWrite := context . WithTimeout ( r . Context (), 3 * time . Second )
writeErr := conn . Write ( writeCtx , coderws . MessageText , [] byte (
`{"type":"response.completed","response":{"id":"resp_usage_e2e","model":"gpt-5.4","usage":{"input_tokens":2,"output_tokens":1}}}` ,
))
cancelWrite ()
if writeErr != nil {
upstreamErrCh <- writeErr
return
}
_ = conn . Close ( coderws . StatusNormalClosure , "done" )
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 {
"api_key" : "sk-test" ,
"base_url" : upstreamServer . 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
accountRepo := & openAIWSUsageHandlerAccountRepoStub { account : account }
usageRepo := & openAIWSUsageHandlerUsageLogRepoStub { created : make ( chan * service . UsageLog , 1 )}
var channelSvc * service . ChannelService
if len ( tc . channelMapping ) > 0 {
channelSvc = service . NewChannelService ( & openAIWSUsageHandlerChannelRepoStub {
channels : [] service . Channel {{
ID : 7701 ,
Name : "openai-ws-e2e-channel" ,
Status : service . StatusActive ,
GroupIDs : [] int64 { groupID },
ModelMapping : map [ string ] map [ string ] string { service . PlatformOpenAI : tc . channelMapping },
}},
groupPlatforms : map [ int64 ] string { groupID : service . PlatformOpenAI },
}, nil , nil , nil )
}
2026-05-25 14:26:11 +08:00
billingCacheSvc := service . NewBillingCacheService ( nil , nil , nil , nil , nil , nil , cfg , nil )
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 ,
nil ,
& 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 {
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 },
}
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 ()
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 )
readCtx , cancelRead := context . WithTimeout ( context . Background (), 3 * time . Second )
_ , event , err := clientConn . Read ( readCtx )
cancelRead ()
require . NoError ( t , err )
require . Equal ( t , "response.completed" , gjson . GetBytes ( event , "type" ). String ())
_ = clientConn . Close ( coderws . StatusNormalClosure , "done" )
var usageLog * service . UsageLog
select {
case usageLog = <- usageRepo . created :
require . NotNil ( t , usageLog )
case <- time . After ( 3 * time . Second ):
t . Fatal ( "等待 WebSocket usage log 写入超时" )
}
var upstreamFirstPayload [] byte
select {
case upstreamFirstPayload = <- upstreamPayloadCh :
case <- time . After ( 3 * time . Second ):
t . Fatal ( "等待上游 WebSocket 首帧超时" )
}
select {
case upstreamErr := <- upstreamErrCh :
require . NoError ( t , upstreamErr )
case <- time . After ( 3 * time . Second ):
t . Fatal ( "等待上游 WebSocket 结束超时" )
}
return openAIResponsesWSUsageLogResult {
log : usageLog ,
upstreamFirstPayload : upstreamFirstPayload ,
}
}
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
}