Files
sub2api/backend/internal/handler/openai_gateway_handler_test.go
deqiying 23555be380 fix(openai): 修复 WS passthrough 使用记录缺失推理强度和 User-Agent
- 为 OpenAI Responses WebSocket v2 passthrough 补齐每轮 reasoning_effort 元数据
- 传递首帧渠道映射前模型,保留模型后缀推理强度推导能力
- 增加 usage log 端到端回归,覆盖入站 User-Agent、显式 effort 和渠道映射场景
2026-05-03 19:33:09 +08:00

1115 lines
37 KiB
Go
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
package handler
import (
"context"
"encoding/json"
"errors"
"net/http"
"net/http/httptest"
"strings"
"testing"
"time"
"github.com/Wei-Shaw/sub2api/internal/config"
pkghttputil "github.com/Wei-Shaw/sub2api/internal/pkg/httputil"
"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 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"])
}
func TestOpenAIEnsureForwardErrorResponse_DoesNotOverrideWrittenResponse(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.False(t, wrote)
require.Equal(t, http.StatusTeapot, w.Code)
assert.Equal(t, "already written", w.Body.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())
}
func TestOpenAIRecoverResponsesPanic_DoesNotOverrideWrittenResponse(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)
assert.Equal(t, "already written", w.Body.String())
}
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 TestResolveOpenAIForwardDefaultMappedModel(t *testing.T) {
t.Run("prefers_explicit_fallback_model", func(t *testing.T) {
apiKey := &service.APIKey{
Group: &service.Group{DefaultMappedModel: "gpt-5.4"},
}
require.Equal(t, "gpt-5.2", resolveOpenAIForwardDefaultMappedModel(apiKey, " gpt-5.2 "))
})
t.Run("uses_group_default_when_explicit_fallback_absent", func(t *testing.T) {
apiKey := &service.APIKey{
Group: &service.Group{DefaultMappedModel: "gpt-5.4"},
}
require.Equal(t, "gpt-5.4", resolveOpenAIForwardDefaultMappedModel(apiKey, ""))
})
t.Run("returns_empty_without_group_default", func(t *testing.T) {
require.Empty(t, resolveOpenAIForwardDefaultMappedModel(nil, ""))
require.Empty(t, resolveOpenAIForwardDefaultMappedModel(&service.APIKey{}, ""))
require.Empty(t, resolveOpenAIForwardDefaultMappedModel(&service.APIKey{
Group: &service.Group{},
}, ""))
})
}
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",
},
},
},
}
require.Equal(t, "gpt-5.4-mini", resolveOpenAIMessagesDispatchMappedModel(apiKey, "claude-sonnet-4-5-20250929"))
})
t.Run("uses_family_default_when_no_override", func(t *testing.T) {
apiKey := &service.APIKey{Group: &service.Group{}}
require.Equal(t, "gpt-5.4", resolveOpenAIMessagesDispatchMappedModel(apiKey, "claude-opus-4-6"))
require.Equal(t, "gpt-5.3-codex", resolveOpenAIMessagesDispatchMappedModel(apiKey, "claude-sonnet-4-5-20250929"))
require.Equal(t, "gpt-5.4-mini", resolveOpenAIMessagesDispatchMappedModel(apiKey, "claude-haiku-4-5-20251001"))
})
t.Run("returns_empty_for_non_claude_or_missing_group", func(t *testing.T) {
require.Empty(t, resolveOpenAIMessagesDispatchMappedModel(nil, "claude-sonnet-4-5-20250929"))
require.Empty(t, resolveOpenAIMessagesDispatchMappedModel(&service.APIKey{}, "claude-sonnet-4-5-20250929"))
require.Empty(t, resolveOpenAIMessagesDispatchMappedModel(&service.APIKey{Group: &service.Group{}}, "gpt-5.4"))
})
t.Run("does_not_fall_back_to_group_default_mapped_model", func(t *testing.T) {
apiKey := &service.APIKey{
Group: &service.Group{
DefaultMappedModel: "gpt-5.4",
},
}
require.Empty(t, resolveOpenAIMessagesDispatchMappedModel(apiKey, "gpt-5.4"))
require.Equal(t, "gpt-5.3-codex", resolveOpenAIMessagesDispatchMappedModel(apiKey, "claude-sonnet-4-5-20250929"))
})
}
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_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")
}
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 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
userAgent *string
channelMapping map[string]string
}
type openAIResponsesWSUsageLogResult struct {
log *service.UsageLog
upstreamFirstPayload []byte
}
type openAIWSUsageHandlerAccountRepoStub struct {
service.AccountRepository
account service.Account
}
func (s *openAIWSUsageHandlerAccountRepoStub) ListSchedulableByPlatform(ctx context.Context, platform string) ([]service.Account, error) {
if s.account.Platform != platform {
return nil, nil
}
return []service.Account{s.account}, nil
}
func (s *openAIWSUsageHandlerAccountRepoStub) ListSchedulableByGroupIDAndPlatform(ctx context.Context, groupID int64, platform string) ([]service.Account, error) {
return s.ListSchedulableByPlatform(ctx, platform)
}
func (s *openAIWSUsageHandlerAccountRepoStub) GetByID(ctx context.Context, id int64) (*service.Account, error) {
if s.account.ID != id {
return nil, nil
}
account := s.account
return &account, nil
}
type openAIWSUsageHandlerUsageLogRepoStub struct {
service.UsageLogRepository
created chan *service.UsageLog
}
func (s *openAIWSUsageHandlerUsageLogRepoStub) Create(ctx context.Context, log *service.UsageLog) (bool, error) {
if s.created != nil {
s.created <- log
}
return true, nil
}
type openAIWSUsageHandlerChannelRepoStub struct {
service.ChannelRepository
channels []service.Channel
groupPlatforms map[int64]string
}
func (s *openAIWSUsageHandlerChannelRepoStub) ListAll(ctx context.Context) ([]service.Channel, error) {
return s.channels, nil
}
func (s *openAIWSUsageHandlerChannelRepoStub) GetGroupPlatforms(ctx context.Context, groupIDs []int64) (map[int64]string, error) {
out := make(map[int64]string, len(groupIDs))
for _, groupID := range groupIDs {
if platform := strings.TrimSpace(s.groupPlatforms[groupID]); platform != "" {
out[groupID] = platform
}
}
return out, nil
}
func runOpenAIResponsesWebSocketUsageLogCase(t *testing.T, tc openAIResponsesWSUsageLogCase) openAIResponsesWSUsageLogResult {
t.Helper()
gin.SetMode(gin.TestMode)
upstreamPayloadCh := make(chan []byte, 1)
upstreamErrCh := make(chan error, 1)
upstreamServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
conn, err := coderws.Accept(w, r, &coderws.AcceptOptions{
CompressionMode: coderws.CompressionContextTakeover,
})
if err != nil {
upstreamErrCh <- err
return
}
defer func() {
_ = conn.CloseNow()
}()
readCtx, cancelRead := context.WithTimeout(r.Context(), 3*time.Second)
msgType, payload, readErr := conn.Read(readCtx)
cancelRead()
if readErr != nil {
upstreamErrCh <- readErr
return
}
if msgType != coderws.MessageText && msgType != coderws.MessageBinary {
upstreamErrCh <- errors.New("unexpected upstream websocket message type")
return
}
upstreamPayloadCh <- payload
writeCtx, cancelWrite := context.WithTimeout(r.Context(), 3*time.Second)
writeErr := conn.Write(writeCtx, coderws.MessageText, []byte(
`{"type":"response.completed","response":{"id":"resp_usage_e2e","model":"gpt-5.4","usage":{"input_tokens":2,"output_tokens":1}}}`,
))
cancelWrite()
if writeErr != nil {
upstreamErrCh <- writeErr
return
}
_ = conn.Close(coderws.StatusNormalClosure, "done")
upstreamErrCh <- nil
}))
defer upstreamServer.Close()
groupID := int64(4201)
account := service.Account{
ID: 9901,
Name: "openai-ws-passthrough-usage-e2e",
Platform: service.PlatformOpenAI,
Type: service.AccountTypeAPIKey,
Status: service.StatusActive,
Schedulable: true,
Concurrency: 1,
Credentials: map[string]any{
"api_key": "sk-test",
"base_url": upstreamServer.URL,
},
Extra: map[string]any{
"openai_apikey_responses_websockets_v2_enabled": true,
"openai_apikey_responses_websockets_v2_mode": service.OpenAIWSIngressModePassthrough,
},
}
cfg := &config.Config{}
cfg.RunMode = config.RunModeSimple
cfg.Default.RateMultiplier = 1
cfg.Security.URLAllowlist.Enabled = false
cfg.Security.URLAllowlist.AllowInsecureHTTP = true
cfg.Gateway.OpenAIWS.Enabled = true
cfg.Gateway.OpenAIWS.APIKeyEnabled = true
cfg.Gateway.OpenAIWS.ResponsesWebsocketsV2 = true
cfg.Gateway.OpenAIWS.ModeRouterV2Enabled = true
cfg.Gateway.OpenAIWS.DialTimeoutSeconds = 3
cfg.Gateway.OpenAIWS.ReadTimeoutSeconds = 3
cfg.Gateway.OpenAIWS.WriteTimeoutSeconds = 3
accountRepo := &openAIWSUsageHandlerAccountRepoStub{account: account}
usageRepo := &openAIWSUsageHandlerUsageLogRepoStub{created: make(chan *service.UsageLog, 1)}
var channelSvc *service.ChannelService
if len(tc.channelMapping) > 0 {
channelSvc = service.NewChannelService(&openAIWSUsageHandlerChannelRepoStub{
channels: []service.Channel{{
ID: 7701,
Name: "openai-ws-e2e-channel",
Status: service.StatusActive,
GroupIDs: []int64{groupID},
ModelMapping: map[string]map[string]string{service.PlatformOpenAI: tc.channelMapping},
}},
groupPlatforms: map[int64]string{groupID: service.PlatformOpenAI},
}, nil, nil, nil)
}
billingCacheSvc := service.NewBillingCacheService(nil, nil, nil, nil, nil, nil, cfg)
gatewaySvc := service.NewOpenAIGatewayService(
accountRepo,
usageRepo,
nil,
nil,
nil,
nil,
nil,
cfg,
nil,
nil,
service.NewBillingService(cfg, nil),
nil,
billingCacheSvc,
nil,
&service.DeferredService{},
nil,
nil,
channelSvc,
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),
}
apiKey := &service.APIKey{
ID: 1801,
GroupID: &groupID,
User: &service.User{ID: 1701, Status: service.StatusActive},
}
router := gin.New()
router.Use(func(c *gin.Context) {
c.Set(string(middleware.ContextKeyAPIKey), apiKey)
c.Set(string(middleware.ContextKeyUser), middleware.AuthSubject{UserID: apiKey.User.ID, Concurrency: 1})
c.Next()
})
router.GET("/openai/v1/responses", h.ResponsesWebSocket)
handlerServer := httptest.NewServer(router)
defer handlerServer.Close()
headers := http.Header{}
if tc.userAgent != nil {
headers.Set("User-Agent", *tc.userAgent)
}
dialCtx, cancelDial := context.WithTimeout(context.Background(), 3*time.Second)
clientConn, _, err := coderws.Dial(
dialCtx,
"ws"+strings.TrimPrefix(handlerServer.URL, "http")+"/openai/v1/responses",
&coderws.DialOptions{HTTPHeader: headers, CompressionMode: coderws.CompressionContextTakeover},
)
cancelDial()
require.NoError(t, err)
defer func() {
_ = clientConn.CloseNow()
}()
writeCtx, cancelWrite := context.WithTimeout(context.Background(), 3*time.Second)
err = clientConn.Write(writeCtx, coderws.MessageText, []byte(tc.firstPayload))
cancelWrite()
require.NoError(t, err)
readCtx, cancelRead := context.WithTimeout(context.Background(), 3*time.Second)
_, event, err := clientConn.Read(readCtx)
cancelRead()
require.NoError(t, err)
require.Equal(t, "response.completed", gjson.GetBytes(event, "type").String())
_ = clientConn.Close(coderws.StatusNormalClosure, "done")
var usageLog *service.UsageLog
select {
case usageLog = <-usageRepo.created:
require.NotNil(t, usageLog)
case <-time.After(3 * time.Second):
t.Fatal("等待 WebSocket usage log 写入超时")
}
var upstreamFirstPayload []byte
select {
case upstreamFirstPayload = <-upstreamPayloadCh:
case <-time.After(3 * time.Second):
t.Fatal("等待上游 WebSocket 首帧超时")
}
select {
case upstreamErr := <-upstreamErrCh:
require.NoError(t, upstreamErr)
case <-time.After(3 * time.Second):
t.Fatal("等待上游 WebSocket 结束超时")
}
return openAIResponsesWSUsageLogResult{
log: usageLog,
upstreamFirstPayload: upstreamFirstPayload,
}
}
func testStringPtr(v string) *string {
return &v
}