Files

3004 lines
106 KiB
Go
Raw Permalink Normal View History

package handler
import (
"context"
"encoding/json"
"errors"
"fmt"
"io"
"net/http"
"net/http/httptest"
"strings"
"sync"
"sync/atomic"
"testing"
"time"
"github.com/Wei-Shaw/sub2api/internal/config"
pkghttputil "github.com/Wei-Shaw/sub2api/internal/pkg/httputil"
"github.com/Wei-Shaw/sub2api/internal/pkg/pagination"
"github.com/Wei-Shaw/sub2api/internal/pkg/xai"
"github.com/Wei-Shaw/sub2api/internal/server/middleware"
"github.com/Wei-Shaw/sub2api/internal/service"
coderws "github.com/coder/websocket"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/tidwall/gjson"
"github.com/tidwall/sjson"
)
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"])
})
}
}
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)
}
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",
}))
}
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))
})
}
}
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)
}
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"])
}
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)
}
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"])
}
// Writer 已写后 ensureForwardErrorResponse 必须仍然把错误信息以 SSE
// 形式追加给客户端(streamStarted 强制 true)。
// 这是 case B 修复:旧实现遇到 Writer.Written 直接 return false
// 客户端只能拿到 silent EOFCodex CLI 报 "stream closed before response.completed"。
func TestOpenAIEnsureForwardErrorResponse_AppendsSSEAfterWritten(t *testing.T) {
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)
require.True(t, wrote, "must attempt to communicate the failure to the client via SSE")
// 状态码改不了(headers 已 flush),但 body 应该追加 SSE 错误事件。
require.Equal(t, http.StatusTeapot, w.Code)
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")
}
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)
}
func TestOpenAIEnsureForwardErrorResponse_CompactKeepaliveOnlyWritesResponseFailed(t *testing.T) {
gin.SetMode(gin.TestMode)
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
c.Request = httptest.NewRequest(http.MethodPost, EndpointResponses, nil)
service.MarkOpenAICompactClientStream(c)
stop := service.StartOpenAICompactSSEKeepalive(c, 5*time.Millisecond)
defer stop()
before := service.OpenAICompactKeepaliveAdjustedWrittenSize(c)
require.Eventually(t, c.Writer.Written, time.Second, time.Millisecond)
require.Equal(t, before, service.OpenAICompactKeepaliveAdjustedWrittenSize(c))
h := &OpenAIGatewayHandler{}
require.True(t, h.ensureForwardErrorResponse(c, false))
require.Equal(t, http.StatusOK, w.Code)
require.Contains(t, w.Body.String(), "event: response.failed\n")
require.NotContains(t, w.Body.String(), "event: error\n")
}
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())
}
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())
}
// Panic 在已 flush 的 /v1/responses 流中:状态码无法改(已 written),
// 但 body 应追加 response.failed 让客户端识别为合规截断而不是 silent EOF。
func TestOpenAIRecoverResponsesPanic_AppendsResponseFailedAfterWritten(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{}
streamStarted := false
require.NotPanics(t, func() {
func() {
defer h.recoverResponsesPanic(c, &streamStarted)
panic("test panic")
}()
})
require.Equal(t, http.StatusTeapot, w.Code)
body := w.Body.String()
assert.Contains(t, body, "already written")
assert.Contains(t, body, "event: response.failed\n")
}
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())
})
}
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",
"claude-fable-5": "gpt-5.6-sol",
},
},
},
}
require.Equal(t, "gpt-5.4-mini", resolveOpenAIMessagesDispatchMappedModel(nil, apiKey, "claude-sonnet-4-5-20250929"))
require.Equal(t, "gpt-5.6-sol", resolveOpenAIMessagesDispatchMappedModel(nil, apiKey, "claude-fable-5"))
})
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(nil, apiKey, "claude-opus-4-6"))
require.Equal(t, "gpt-5.3-codex", resolveOpenAIMessagesDispatchMappedModel(nil, apiKey, "claude-sonnet-4-5-20250929"))
require.Equal(t, "gpt-5.4-mini", resolveOpenAIMessagesDispatchMappedModel(nil, apiKey, "claude-haiku-4-5-20251001"))
})
t.Run("returns_empty_for_non_claude_or_missing_group", func(t *testing.T) {
require.Empty(t, resolveOpenAIMessagesDispatchMappedModel(nil, nil, "claude-sonnet-4-5-20250929"))
require.Empty(t, resolveOpenAIMessagesDispatchMappedModel(nil, &service.APIKey{}, "claude-sonnet-4-5-20250929"))
require.Empty(t, resolveOpenAIMessagesDispatchMappedModel(nil, &service.APIKey{Group: &service.Group{}}, "gpt-5.4"))
})
t.Run("grok_group_maps_claude_cli_model_to_grok_default", func(t *testing.T) {
original := xai.RuntimeModelMappingOptions()
t.Cleanup(func() { xai.SetRuntimeModelMappingOptions(original) })
xai.SetRuntimeModelMappingOptions(xai.ModelMappingOptions{EnableCrossClientMap: true})
apiKey := &service.APIKey{
Group: &service.Group{
Platform: service.PlatformGrok,
},
}
require.Equal(t, "grok-4.5", resolveOpenAIMessagesDispatchMappedModel(nil, apiKey, "claude-sonnet-4-5"))
require.Empty(t, resolveOpenAIMessagesDispatchMappedModel(nil, apiKey, "grok"))
})
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(nil, apiKey, "gpt-5.4"))
require.Equal(t, "gpt-5.3-codex", resolveOpenAIMessagesDispatchMappedModel(nil, apiKey, "claude-sonnet-4-5-20250929"))
})
}
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")
})
}
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])
}
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")
}
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")
}
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))
}
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)
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.StatusTooManyRequests, response.StatusCode)
_ = response.Body.Close()
}
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()
}
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)
}
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)
}
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)
}
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")
}
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 {
mu sync.Mutex
logs []service.ContentModerationLog
}
func (r *contentModerationHandlerTestRepo) CreateLog(ctx context.Context, log *service.ContentModerationLog) error {
if log != nil {
r.mu.Lock()
defer r.mu.Unlock()
r.logs = append(r.logs, *log)
}
return nil
}
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...)
}
func (r *contentModerationHandlerTestRepo) ListLogs(ctx context.Context, filter service.ContentModerationLogFilter) ([]service.ContentModerationLog, *pagination.PaginationResult, error) {
return nil, nil, nil
}
func (r *contentModerationHandlerTestRepo) CountFlaggedByUserSince(ctx context.Context, userID int64, since time.Time, excludeCyberPolicy bool) (int, error) {
return 0, nil
}
func (r *contentModerationHandlerTestRepo) CleanupExpiredLogs(ctx context.Context, hitBefore time.Time, nonHitBefore time.Time) (*service.ContentModerationCleanupResult, error) {
return &service.ContentModerationCleanupResult{}, nil
}
func (r *contentModerationHandlerTestRepo) UpdateLogEmailSent(ctx context.Context, id int64, sent bool) error {
return nil
}
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,
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)
require.Eventually(t, func() bool {
return len(repo.logSnapshot()) == 1
}, time.Second, 10*time.Millisecond)
repo.resetLogs()
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, "内容审计测试阻断")
}
var logs []service.ContentModerationLog
require.Eventually(t, func() bool {
logs = repo.logSnapshot()
return len(logs) == 1
}, time.Second, 10*time.Millisecond)
require.True(t, logs[0].Flagged)
require.Equal(t, service.ContentModerationActionBlock, logs[0].Action)
require.Equal(t, "bad prompt", logs[0].InputExcerpt)
}
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)
}
func TestOpenAIResponsesWebSocket_PassthroughTracksModelPerTurn(t *testing.T) {
got := runOpenAIResponsesWebSocketUsageLogCase(t, openAIResponsesWSUsageLogCase{
firstPayload: `{"type":"response.create","model":"sol","stream":false}`,
secondPayload: `{"type":"response.create","model":"terra","stream":false}`,
channelMapping: map[string]string{
"sol": "sol-channel",
"terra": "terra-channel",
},
accountModelMapping: map[string]any{
"sol": "gpt-5.6-sol",
"terra": "gpt-5.6-terra",
"sol-channel": "gpt-5.6-sol",
"terra-channel": "gpt-5.6-terra",
},
})
require.Len(t, got.upstreamPayloads, 2)
require.Equal(t, "gpt-5.6-sol", gjson.GetBytes(got.upstreamPayloads[0], "model").String())
require.Equal(t, "gpt-5.6-terra", gjson.GetBytes(got.upstreamPayloads[1], "model").String())
require.Len(t, got.clientEvents, 2)
require.Equal(t, "sol", gjson.GetBytes(got.clientEvents[0], "response.model").String())
require.Equal(t, "terra", gjson.GetBytes(got.clientEvents[1], "response.model").String())
require.Len(t, got.logs, 2)
require.Equal(t, "sol", got.logs[0].Model)
require.Equal(t, "sol", got.logs[0].RequestedModel)
require.NotNil(t, got.logs[0].UpstreamModel)
require.Equal(t, "gpt-5.6-sol", *got.logs[0].UpstreamModel)
require.NotNil(t, got.logs[0].ModelMappingChain)
require.Equal(t, "sol→sol-channel→gpt-5.6-sol", *got.logs[0].ModelMappingChain)
require.Equal(t, "terra", got.logs[1].Model)
require.Equal(t, "terra", got.logs[1].RequestedModel)
require.NotNil(t, got.logs[1].UpstreamModel)
require.Equal(t, "gpt-5.6-terra", *got.logs[1].UpstreamModel)
require.NotNil(t, got.logs[1].ModelMappingChain)
require.Equal(t, "terra→terra-channel→gpt-5.6-terra", *got.logs[1].ModelMappingChain)
require.InDelta(t, got.logs[1].TotalCost*2.5, got.logs[0].TotalCost, 1e-12,
"each turn must be billed with its own channel-mapped model")
}
func TestOpenAIResponsesWebSocket_UnchangedChannelTargetOutsideAccountMappingKeysRemainsValid(t *testing.T) {
got := runOpenAIResponsesWebSocketUsageLogCase(t, openAIResponsesWSUsageLogCase{
firstPayload: `{"type":"response.create","model":"public-alias","stream":false}`,
secondPayload: `{"type":"response.create","stream":false}`,
channelMapping: map[string]string{
"public-alias": "gpt-5.6-sol",
},
accountModelMapping: map[string]any{
"public-alias": "gpt-5.6-terra",
},
})
require.Len(t, got.upstreamPayloads, 2)
require.Equal(t, "gpt-5.6-sol", gjson.GetBytes(got.upstreamPayloads[0], "model").String())
require.Equal(t, "gpt-5.6-sol", gjson.GetBytes(got.upstreamPayloads[1], "model").String())
require.Len(t, got.clientEvents, 2)
require.Equal(t, "public-alias", gjson.GetBytes(got.clientEvents[0], "response.model").String())
require.Equal(t, "public-alias", gjson.GetBytes(got.clientEvents[1], "response.model").String())
require.Len(t, got.logs, 2)
for _, usageLog := range got.logs {
require.Equal(t, "public-alias", usageLog.RequestedModel)
require.NotNil(t, usageLog.UpstreamModel)
require.Equal(t, "gpt-5.6-sol", *usageLog.UpstreamModel)
require.NotNil(t, usageLog.ModelMappingChain)
require.Equal(t, "public-alias→gpt-5.6-sol", *usageLog.ModelMappingChain)
}
}
func TestOpenAIResponsesWebSocket_PassthroughKeepsTurnMappingSnapshot(t *testing.T) {
got := runOpenAIResponsesWebSocketUsageLogCase(t, openAIResponsesWSUsageLogCase{
firstPayload: `{"type":"response.create","model":"sol","stream":false}`,
secondPayload: `{"type":"response.create","model":"sol","stream":false}`,
channelMapping: map[string]string{
"sol": "gpt-5.6-sol",
},
afterFirstUpstreamRequest: func(channelSvc *service.ChannelService) error {
if channelSvc == nil {
return errors.New("channel service is nil")
}
_, err := channelSvc.Update(context.Background(), 7701, &service.UpdateChannelInput{
ModelMapping: map[string]map[string]string{
service.PlatformOpenAI: {"sol": "gpt-5.6-terra"},
},
})
return err
},
})
require.Len(t, got.upstreamPayloads, 2)
require.Equal(t, "gpt-5.6-sol", gjson.GetBytes(got.upstreamPayloads[0], "model").String())
require.Equal(t, "gpt-5.6-terra", gjson.GetBytes(got.upstreamPayloads[1], "model").String())
require.Len(t, got.logs, 2)
require.Equal(t, "sol", got.logs[0].Model)
require.NotNil(t, got.logs[0].UpstreamModel)
require.Equal(t, "gpt-5.6-sol", *got.logs[0].UpstreamModel)
require.NotNil(t, got.logs[0].ModelMappingChain)
require.Equal(t, "sol→gpt-5.6-sol", *got.logs[0].ModelMappingChain)
require.InDelta(t, 40e-6, got.logs[0].TotalCost, 1e-12,
"the in-flight turn must retain the channel-mapped billing model used when it was sent")
require.Equal(t, "sol", got.logs[1].Model)
require.NotNil(t, got.logs[1].UpstreamModel)
require.Equal(t, "gpt-5.6-terra", *got.logs[1].UpstreamModel)
require.NotNil(t, got.logs[1].ModelMappingChain)
require.Equal(t, "sol→gpt-5.6-terra", *got.logs[1].ModelMappingChain)
require.InDelta(t, got.logs[1].TotalCost*2.5, got.logs[0].TotalCost, 1e-12,
"the next turn must use the updated channel mapping")
}
func TestOpenAIResponsesWebSocket_CtxPoolAppliesPerTurnMappingAndPreservesRequestedModel(t *testing.T) {
got := runOpenAIResponsesWebSocketUsageLogCase(t, openAIResponsesWSUsageLogCase{
firstPayload: `{"type":"response.create","model":"gpt-5.6-sol","stream":false}`,
secondPayload: `{"type":"response.create","model":"gpt-5.6-terra","stream":false}`,
ingressMode: service.OpenAIWSIngressModeCtxPool,
billingModelSource: service.BillingModelSourceRequested,
channelMapping: map[string]string{
"gpt-5.6-terra": "gpt-5.6-sol",
},
})
require.Len(t, got.upstreamPayloads, 2)
require.Equal(t, "gpt-5.6-sol", gjson.GetBytes(got.upstreamPayloads[0], "model").String())
require.Equal(t, "gpt-5.6-sol", gjson.GetBytes(got.upstreamPayloads[1], "model").String())
require.Len(t, got.clientEvents, 2)
require.Equal(t, "gpt-5.6-sol", gjson.GetBytes(got.clientEvents[0], "response.model").String())
require.Equal(t, "gpt-5.6-terra", gjson.GetBytes(got.clientEvents[1], "response.model").String())
require.Len(t, got.logs, 2)
require.Equal(t, "gpt-5.6-sol", got.logs[0].RequestedModel)
require.Nil(t, got.logs[0].ModelMappingChain)
require.InDelta(t, 40e-6, got.logs[0].TotalCost, 1e-12)
require.Equal(t, "gpt-5.6-terra", got.logs[1].RequestedModel)
require.NotNil(t, got.logs[1].ModelMappingChain)
require.Equal(t, "gpt-5.6-terra→gpt-5.6-sol", *got.logs[1].ModelMappingChain)
require.InDelta(t, 16e-6, got.logs[1].TotalCost, 1e-12,
"BillingModelSourceRequested must use the client model before channel mapping")
}
func TestOpenAIWSTurnBillingModelPreservesImagePricingModel(t *testing.T) {
tests := []struct {
name string
resultModel string
mapping service.ChannelMappingResult
requestedModel string
upstreamModel string
wantBillingModel string
}{
{
name: "upstream billing preserves image model",
resultModel: "gpt-image-2",
mapping: service.ChannelMappingResult{BillingModelSource: service.BillingModelSourceUpstream},
requestedModel: "gpt-5.6-sol",
upstreamModel: "gpt-5.6-sol",
wantBillingModel: "gpt-image-2",
},
{
name: "unmapped channel preserves image model",
resultModel: "gpt-image-2",
mapping: service.ChannelMappingResult{MappedModel: "gpt-5.6-sol", BillingModelSource: service.BillingModelSourceChannelMapped},
requestedModel: "gpt-5.6-sol",
upstreamModel: "gpt-5.6-sol",
wantBillingModel: "gpt-image-2",
},
{
name: "requested source overrides image model",
resultModel: "gpt-image-2",
mapping: service.ChannelMappingResult{BillingModelSource: service.BillingModelSourceRequested},
requestedModel: "public-image-alias",
upstreamModel: "gpt-5.6-sol",
wantBillingModel: "public-image-alias",
},
{
name: "mapped channel source overrides image model",
resultModel: "gpt-image-2",
mapping: service.ChannelMappingResult{MappedModel: "priced-channel-model", BillingModelSource: service.BillingModelSourceChannelMapped},
requestedModel: "public-image-alias",
upstreamModel: "gpt-5.6-sol",
wantBillingModel: "priced-channel-model",
},
{
name: "text turn falls back to upstream model",
mapping: service.ChannelMappingResult{BillingModelSource: service.BillingModelSourceUpstream},
requestedModel: "public-alias",
upstreamModel: "gpt-5.6-sol",
wantBillingModel: "gpt-5.6-sol",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
result := &service.OpenAIForwardResult{BillingModel: tt.resultModel}
require.Equal(t, tt.wantBillingModel, openAIWSTurnBillingModel(result, tt.mapping, tt.requestedModel, tt.upstreamModel))
})
}
}
func TestShouldReportOpenAIWSProxyAccountFailure(t *testing.T) {
t.Run("unsupported client model switch does not penalize account", func(t *testing.T) {
err := fmt.Errorf("wrapped ingress turn: %w", newOpenAIWSUnsupportedModelSwitchError("gpt-unsupported"))
require.False(t, shouldReportOpenAIWSProxyAccountFailure(err))
var closeErr *service.OpenAIWSClientCloseError
require.ErrorAs(t, err, &closeErr)
require.Equal(t, coderws.StatusPolicyViolation, closeErr.StatusCode())
require.Equal(t, "model switch requires reconnect", closeErr.Reason())
})
t.Run("upstream policy violation still penalizes account", func(t *testing.T) {
err := service.NewOpenAIWSClientCloseError(
coderws.StatusPolicyViolation,
"upstream websocket authentication failed",
errors.New("upstream rejected credentials"),
)
require.True(t, shouldReportOpenAIWSProxyAccountFailure(err))
})
t.Run("generic proxy failure still penalizes account", func(t *testing.T) {
require.True(t, shouldReportOpenAIWSProxyAccountFailure(errors.New("upstream websocket read failed")))
})
}
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))
}
// 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)
modelResult := gjson.GetBytes(body, "model")
model := ""
if modelResult.Type == gjson.String {
model = modelResult.String()
}
stream := gjson.GetBytes(body, "stream").Bool()
require.Equal(t, tt.wantModel, model)
require.Equal(t, tt.wantStream, stream)
})
}
}
// 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)
}
// 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)
// 测试 4sjson.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))
}
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)
}
type openAIResponsesWSUsageLogCase struct {
firstPayload string
secondPayload string
userAgent *string
ingressMode string
channelMapping map[string]string
billingModelSource string
accountModelMapping map[string]any
afterFirstUpstreamRequest func(channelSvc *service.ChannelService) error
}
type openAIResponsesWSUsageLogResult struct {
log *service.UsageLog
logs []*service.UsageLog
upstreamFirstPayload []byte
upstreamPayloads [][]byte
clientEvents [][]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
}
type openAIWSFailoverHandlerAccountRepoStub struct {
service.AccountRepository
accounts []service.Account
rateLimitedIDs []int64
}
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...)
}
type openAIHTTPPassthroughAuthFailoverUpstream struct {
service.HTTPUpstream
mu sync.Mutex
accountIDs []int64
statusCode int
}
func (u *openAIHTTPPassthroughAuthFailoverUpstream) Do(_ *http.Request, _ string, accountID int64, _ int) (*http.Response, error) {
u.mu.Lock()
u.accountIDs = append(u.accountIDs, accountID)
u.mu.Unlock()
if accountID == 9911 {
return &http.Response{
StatusCode: http.StatusOK,
Header: http.Header{"Content-Type": []string{"application/json"}},
Body: io.NopCloser(strings.NewReader(`{"id":"resp_healthy","object":"response","model":"gpt-5.2","status":"completed","usage":{"input_tokens":1,"output_tokens":1,"total_tokens":2}}`)),
}, nil
}
return &http.Response{
StatusCode: u.statusCode,
Header: http.Header{"Content-Type": []string{"application/json"}},
Body: io.NopCloser(strings.NewReader(`{"error":{"message":"upstream credential rejected"}}`)),
}, nil
}
func (u *openAIHTTPPassthroughAuthFailoverUpstream) calls() []int64 {
u.mu.Lock()
defer u.mu.Unlock()
return append([]int64(nil), u.accountIDs...)
}
type openAIHTTPPassthroughSSERateLimitUpstream struct {
service.HTTPUpstream
mu sync.Mutex
accountIDs []int64
}
func (u *openAIHTTPPassthroughSSERateLimitUpstream) Do(_ *http.Request, _ string, accountID int64, _ int) (*http.Response, error) {
u.mu.Lock()
u.accountIDs = append(u.accountIDs, accountID)
u.mu.Unlock()
body := strings.Join([]string{
"event: response.created",
`data: {"type":"response.created","response":{"id":"resp_rate_limited"}}`,
"",
"event: response.failed",
`data: {"type":"response.failed","response":{"id":"resp_rate_limited","status":"failed","error":{"code":"rate_limit_exceeded","message":"Concurrency limit exceeded for account, please retry later"}}}`,
"",
}, "\n")
return &http.Response{
StatusCode: http.StatusOK,
Header: http.Header{
"Content-Type": []string{"text/event-stream"},
"Retry-After": []string{"1"},
},
Body: io.NopCloser(strings.NewReader(body)),
}, nil
}
func (u *openAIHTTPPassthroughSSERateLimitUpstream) calls() []int64 {
u.mu.Lock()
defer u.mu.Unlock()
return append([]int64(nil), u.accountIDs...)
}
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
}
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
mu sync.Mutex
channels []service.Channel
groupPlatforms map[int64]string
}
func (s *openAIWSUsageHandlerChannelRepoStub) ListAll(ctx context.Context) ([]service.Channel, error) {
s.mu.Lock()
defer s.mu.Unlock()
out := make([]service.Channel, 0, len(s.channels))
for i := range s.channels {
out = append(out, *s.channels[i].Clone())
}
return out, nil
}
func (s *openAIWSUsageHandlerChannelRepoStub) GetByID(ctx context.Context, id int64) (*service.Channel, error) {
s.mu.Lock()
defer s.mu.Unlock()
for i := range s.channels {
if s.channels[i].ID == id {
return s.channels[i].Clone(), nil
}
}
return nil, service.ErrChannelNotFound
}
func (s *openAIWSUsageHandlerChannelRepoStub) Update(ctx context.Context, channel *service.Channel) error {
s.mu.Lock()
defer s.mu.Unlock()
for i := range s.channels {
if s.channels[i].ID == channel.ID {
s.channels[i] = *channel.Clone()
return nil
}
}
return service.ErrChannelNotFound
}
func (s *openAIWSUsageHandlerChannelRepoStub) GetGroupIDs(ctx context.Context, channelID int64) ([]int64, error) {
channel, err := s.GetByID(ctx, channelID)
if err != nil {
return nil, err
}
return append([]int64(nil), channel.GroupIDs...), nil
}
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
}
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())
}
func TestOpenAIResponses_APIKeyPassthroughPoolAuthFailureRetriesThenSwitchesToHealthyAccount(t *testing.T) {
tests := []struct {
name string
statusCode int
}{
{name: "401", statusCode: http.StatusUnauthorized},
{name: "403", statusCode: http.StatusForbidden},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
gin.SetMode(gin.TestMode)
groupID := int64(4203)
accounts := []service.Account{
{
ID: 9910, Name: "pool-api-key", Platform: service.PlatformOpenAI,
Type: service.AccountTypeAPIKey, Status: service.StatusActive, Schedulable: true, Priority: 1,
Credentials: map[string]any{
"api_key": "sk-pool",
"base_url": "https://api.example.test",
"pool_mode": true,
"pool_mode_retry_count": float64(1),
"pool_mode_retry_status_codes": []any{float64(tt.statusCode)},
},
Extra: map[string]any{"openai_passthrough": true},
},
{
ID: 9911, Name: "fallback-api-key", Platform: service.PlatformOpenAI,
Type: service.AccountTypeAPIKey, Status: service.StatusActive, Schedulable: true, Priority: 2,
Credentials: map[string]any{
"api_key": "sk-fallback",
"base_url": "https://api.example.test",
},
Extra: map[string]any{"openai_passthrough": true},
},
}
cfg := &config.Config{RunMode: config.RunModeSimple}
cfg.Default.RateMultiplier = 1
cfg.Security.URLAllowlist.Enabled = false
cfg.Gateway.MaxAccountSwitches = 1
accountRepo := &openAIWSFailoverHandlerAccountRepoStub{accounts: accounts}
upstream := &openAIHTTPPassthroughAuthFailoverUpstream{statusCode: tt.statusCode}
rateLimitSvc := service.NewRateLimitService(accountRepo, nil, cfg, nil, nil)
billingCacheSvc := service.NewBillingCacheService(nil, nil, nil, nil, nil, nil, cfg, nil)
t.Cleanup(billingCacheSvc.Stop)
gatewaySvc := service.NewOpenAIGatewayService(
accountRepo,
nil,
nil,
nil,
nil,
nil,
nil,
cfg,
nil,
nil,
service.NewBillingService(cfg, nil),
rateLimitSvc,
billingCacheSvc,
upstream,
&service.DeferredService{},
nil,
nil,
nil,
nil,
nil,
nil,
nil,
)
h := NewOpenAIGatewayHandler(
gatewaySvc,
service.NewConcurrencyService(nil),
billingCacheSvc,
service.NewAPIKeyService(nil, nil, nil, nil, nil, nil, cfg),
nil,
nil,
nil,
nil,
cfg,
)
rec := httptest.NewRecorder()
c, _ := gin.CreateTestContext(rec)
c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", strings.NewReader(`{"model":"gpt-5.2","input":"hello","stream":false}`))
c.Request.Header.Set("Content-Type", "application/json")
c.Set(string(middleware.ContextKeyAPIKey), &service.APIKey{
ID: 1803, GroupID: &groupID,
User: &service.User{ID: 1703, Status: service.StatusActive},
Group: &service.Group{ID: groupID, Platform: service.PlatformOpenAI, Status: service.StatusActive},
})
c.Set(string(middleware.ContextKeyUser), middleware.AuthSubject{UserID: 1703, Concurrency: 0})
h.Responses(c)
require.Equal(t, []int64{9910, 9910, 9911}, upstream.calls())
require.Equal(t, http.StatusOK, rec.Code)
require.Equal(t, "resp_healthy", gjson.GetBytes(rec.Body.Bytes(), "id").String())
})
}
}
func TestOpenAIResponses_APIKeyPassthroughSSERateLimitUsesConfiguredPoolRetry(t *testing.T) {
gin.SetMode(gin.TestMode)
groupID := int64(4204)
accounts := []service.Account{
{
ID: 9912, Name: "pool-sse-rate-limit", Platform: service.PlatformOpenAI,
Type: service.AccountTypeAPIKey, Status: service.StatusActive, Schedulable: true, Priority: 1,
Credentials: map[string]any{
"api_key": "sk-pool",
"base_url": "https://api.example.test",
"pool_mode": true,
"pool_mode_retry_count": float64(1),
"pool_mode_retry_status_codes": []any{float64(http.StatusTooManyRequests)},
},
Extra: map[string]any{"openai_passthrough": true},
},
}
cfg := &config.Config{RunMode: config.RunModeSimple}
cfg.Default.RateMultiplier = 1
cfg.Security.URLAllowlist.Enabled = false
cfg.Gateway.MaxAccountSwitches = 1
accountRepo := &openAIWSFailoverHandlerAccountRepoStub{accounts: accounts}
upstream := &openAIHTTPPassthroughSSERateLimitUpstream{}
billingCacheSvc := service.NewBillingCacheService(nil, nil, nil, nil, nil, nil, cfg, nil)
t.Cleanup(billingCacheSvc.Stop)
gatewaySvc := service.NewOpenAIGatewayService(
accountRepo,
nil,
nil,
nil,
nil,
nil,
nil,
cfg,
nil,
nil,
service.NewBillingService(cfg, nil),
nil,
billingCacheSvc,
upstream,
&service.DeferredService{},
nil,
nil,
nil,
nil,
nil,
nil,
nil,
)
h := NewOpenAIGatewayHandler(
gatewaySvc,
service.NewConcurrencyService(nil),
billingCacheSvc,
service.NewAPIKeyService(nil, nil, nil, nil, nil, nil, cfg),
nil,
nil,
nil,
nil,
cfg,
)
rec := httptest.NewRecorder()
c, _ := gin.CreateTestContext(rec)
c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", strings.NewReader(`{"model":"gpt-5.6-sol","input":"hello","stream":true}`))
c.Request.Header.Set("Content-Type", "application/json")
c.Set(string(middleware.ContextKeyAPIKey), &service.APIKey{
ID: 1804, GroupID: &groupID,
User: &service.User{ID: 1704, Status: service.StatusActive},
Group: &service.Group{ID: groupID, Platform: service.PlatformOpenAI, Status: service.StatusActive},
})
c.Set(string(middleware.ContextKeyUser), middleware.AuthSubject{UserID: 1704, Concurrency: 0})
h.Responses(c)
require.Equal(t, []int64{9912, 9912}, upstream.calls())
require.Equal(t, http.StatusTooManyRequests, rec.Code)
require.Equal(t, "1", rec.Header().Get("Retry-After"))
require.Equal(t, "rate_limit_error", gjson.GetBytes(rec.Body.Bytes(), "error.type").String())
require.Equal(t, "Upstream rate limit exceeded, please retry later", gjson.GetBytes(rec.Body.Bytes(), "error.message").String())
}
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,
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
},
}
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)
}
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")
}
func runOpenAIResponsesWebSocketUsageLogCase(t *testing.T, tc openAIResponsesWSUsageLogCase) openAIResponsesWSUsageLogResult {
t.Helper()
gin.SetMode(gin.TestMode)
turnCount := 1
if strings.TrimSpace(tc.secondPayload) != "" {
turnCount = 2
}
upstreamPayloadCh := make(chan []byte, turnCount)
upstreamErrCh := make(chan error, 1)
var channelSvc *service.ChannelService
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()
}()
for turn := 1; turn <= turnCount; turn++ {
readCtx, cancelRead := context.WithTimeout(r.Context(), 3*time.Second)
msgType, payload, readErr := conn.Read(readCtx)
cancelRead()
if readErr != nil {
upstreamErrCh <- readErr
return
}
if msgType != coderws.MessageText && msgType != coderws.MessageBinary {
upstreamErrCh <- errors.New("unexpected upstream websocket message type")
return
}
upstreamPayloadCh <- payload
if turn == 1 && tc.afterFirstUpstreamRequest != nil {
if callbackErr := tc.afterFirstUpstreamRequest(channelSvc); callbackErr != nil {
upstreamErrCh <- callbackErr
return
}
}
response := fmt.Sprintf(
`{"type":"response.completed","response":{"id":"resp_usage_e2e_%d","model":%q,"usage":{"input_tokens":2,"output_tokens":1}}}`,
turn,
gjson.GetBytes(payload, "model").String(),
)
writeCtx, cancelWrite := context.WithTimeout(r.Context(), 3*time.Second)
writeErr := conn.Write(writeCtx, coderws.MessageText, []byte(response))
cancelWrite()
if writeErr != nil {
upstreamErrCh <- writeErr
return
}
}
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,
"model_mapping": tc.accountModelMapping,
},
Extra: map[string]any{
"openai_apikey_responses_websockets_v2_enabled": true,
"openai_apikey_responses_websockets_v2_mode": service.OpenAIWSIngressModePassthrough,
},
}
if strings.TrimSpace(tc.ingressMode) != "" {
account.Extra["openai_apikey_responses_websockets_v2_mode"] = tc.ingressMode
}
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, turnCount)}
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},
BillingModelSource: tc.billingModelSource,
}},
groupPlatforms: map[int64]string{groupID: service.PlatformOpenAI},
}, nil, nil, nil)
}
billingCacheSvc := service.NewBillingCacheService(nil, nil, nil, nil, nil, nil, cfg, nil)
gatewaySvc := service.NewOpenAIGatewayService(
accountRepo,
usageRepo,
nil,
nil,
nil,
nil,
nil,
cfg,
nil,
nil,
service.NewBillingService(cfg, nil),
nil,
billingCacheSvc,
nil,
&service.DeferredService{},
nil,
nil,
nil,
channelSvc,
nil,
nil,
nil, // userPlatformQuotaRepo
)
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)
clientEvents := make([][]byte, 0, turnCount)
readCompleted := func() {
readCtx, cancelRead := context.WithTimeout(context.Background(), 3*time.Second)
_, event, readErr := clientConn.Read(readCtx)
cancelRead()
require.NoError(t, readErr)
require.Equal(t, "response.completed", gjson.GetBytes(event, "type").String())
clientEvents = append(clientEvents, append([]byte(nil), event...))
}
readCompleted()
if turnCount == 2 {
writeCtx, cancelWrite = context.WithTimeout(context.Background(), 3*time.Second)
err = clientConn.Write(writeCtx, coderws.MessageText, []byte(tc.secondPayload))
cancelWrite()
require.NoError(t, err)
readCompleted()
}
_ = clientConn.Close(coderws.StatusNormalClosure, "done")
usageLogs := make([]*service.UsageLog, 0, turnCount)
for len(usageLogs) < turnCount {
select {
case usageLog := <-usageRepo.created:
require.NotNil(t, usageLog)
usageLogs = append(usageLogs, usageLog)
case <-time.After(3 * time.Second):
t.Fatal("等待 WebSocket usage log 写入超时")
}
}
upstreamPayloads := make([][]byte, 0, turnCount)
for len(upstreamPayloads) < turnCount {
select {
case payload := <-upstreamPayloadCh:
upstreamPayloads = append(upstreamPayloads, payload)
case <-time.After(3 * time.Second):
t.Fatal("等待上游 WebSocket 请求帧超时")
}
}
select {
case upstreamErr := <-upstreamErrCh:
require.NoError(t, upstreamErr)
case <-time.After(3 * time.Second):
t.Fatal("等待上游 WebSocket 结束超时")
}
return openAIResponsesWSUsageLogResult{
log: usageLogs[0],
logs: usageLogs,
upstreamFirstPayload: upstreamPayloads[0],
upstreamPayloads: upstreamPayloads,
clientEvents: clientEvents,
}
}
func testStringPtr(v string) *string {
return &v
}
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)
})
// 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")))
})
}