Просмотр исходного кода

chore: commit relay SSE error handling fixes and E2E mock server

- relay: SSE stream error exposure improvements (Claude/Gemini/OpenAI)
- test: add mock capture server for E2E relay capture testing

Co-Authored-By: Claude <noreply@anthropic.com>
master
fengsilin 2 недель назад
Родитель
Сommit
0e3b98efe0
10 измененных файлов: 506 добавлений и 7 удалений
  1. +7
    -0
      relay/channel/claude/relay-claude.go
  2. +70
    -0
      relay/channel/claude/relay_claude_test.go
  3. +9
    -2
      relay/channel/gemini/relay-gemini.go
  4. +59
    -0
      relay/channel/gemini/relay_gemini_usage_test.go
  5. +1
    -1
      relay/channel/openai/helper.go
  6. +4
    -2
      relay/channel/openai/relay-openai.go
  7. +8
    -1
      relay/channel/openai/relay_responses.go
  8. +100
    -0
      relay/channel/openai/upstream_body_test.go
  9. +40
    -1
      relay/helper/common.go
  10. +208
    -0
      test/e2e/mock_capture_server.go

+ 7
- 0
relay/channel/claude/relay-claude.go Просмотреть файл

@@ -778,14 +778,21 @@ func ClaudeStreamHandler(c *gin.Context, resp *http.Response, info *relaycommon.
Usage: &dto.Usage{}, Usage: &dto.Usage{},
} }
var err *types.NewAPIError var err *types.NewAPIError
var normalStreamStarted bool
helper.StreamScannerHandler(c, resp, info, func(data string) bool { helper.StreamScannerHandler(c, resp, info, func(data string) bool {
err = HandleStreamResponseData(c, info, claudeInfo, data) err = HandleStreamResponseData(c, info, claudeInfo, data)
if err != nil { if err != nil {
if normalStreamStarted || helper.StreamDataWritten(c) {
logger.LogError(c, "upstream claude stream error after normal stream started: "+err.Error())
err = nil
}
return false return false
} }
normalStreamStarted = true
return true return true
}) })
if err != nil { if err != nil {
helper.ClearEventStreamHeadersIfNotWritten(c)
return nil, err return nil, err
} }




+ 70
- 0
relay/channel/claude/relay_claude_test.go Просмотреть файл

@@ -1,10 +1,18 @@
package claude package claude


import ( import (
"io"
"net/http"
"net/http/httptest"
"strings" "strings"
"testing" "testing"


"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/dto" "github.com/QuantumNous/new-api/dto"
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/types"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/require"
) )


func TestFormatClaudeResponseInfo_MessageStart(t *testing.T) { func TestFormatClaudeResponseInfo_MessageStart(t *testing.T) {
@@ -173,3 +181,65 @@ func TestFormatClaudeResponseInfo_ContentBlockDelta(t *testing.T) {
t.Errorf("ResponseText = %q, want %q", claudeInfo.ResponseText.String(), "hello") t.Errorf("ResponseText = %q, want %q", claudeInfo.ResponseText.String(), "hello")
} }
} }

func TestClaudeStreamHandlerReturnsErrorOnInitialErrorEvent(t *testing.T) {
gin.SetMode(gin.TestMode)
oldTimeout := constant.StreamingTimeout
constant.StreamingTimeout = 30
t.Cleanup(func() {
constant.StreamingTimeout = oldTimeout
})

w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", nil)

errorEvent := `{"type":"error","error":{"type":"overloaded_error","message":"Overloaded"}}`
resp := &http.Response{
StatusCode: http.StatusOK,
Body: io.NopCloser(strings.NewReader("data: " + errorEvent + "\n\n")),
}
info := &relaycommon.RelayInfo{
RelayFormat: types.RelayFormatClaude,
ChannelMeta: &relaycommon.ChannelMeta{
UpstreamModelName: "claude-test",
},
}

usage, err := ClaudeStreamHandler(c, resp, info)
require.Nil(t, usage)
require.Error(t, err)
require.Empty(t, w.Body.String())
}

func TestClaudeStreamHandlerDoesNotRetryAfterNormalStreamData(t *testing.T) {
gin.SetMode(gin.TestMode)
oldTimeout := constant.StreamingTimeout
constant.StreamingTimeout = 30
t.Cleanup(func() {
constant.StreamingTimeout = oldTimeout
})

w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", nil)

normalEvent := `{"type":"message_start","message":{"id":"msg_test","model":"claude-test","usage":{"input_tokens":1,"output_tokens":1}}}`
errorEvent := `{"type":"error","error":{"type":"overloaded_error","message":"Overloaded"}}`
resp := &http.Response{
StatusCode: http.StatusOK,
Body: io.NopCloser(strings.NewReader("data: " + normalEvent + "\n\ndata: " + errorEvent + "\n\n")),
}
info := &relaycommon.RelayInfo{
RelayFormat: types.RelayFormatClaude,
ChannelMeta: &relaycommon.ChannelMeta{
UpstreamModelName: "claude-test",
},
}

usage, err := ClaudeStreamHandler(c, resp, info)
require.Nil(t, err)
require.NotNil(t, usage)
require.Contains(t, w.Body.String(), "message_start")
require.NotContains(t, w.Body.String(), "overloaded_error")
}

+ 9
- 2
relay/channel/gemini/relay-gemini.go Просмотреть файл

@@ -1286,6 +1286,7 @@ func geminiStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http
var imageCount int var imageCount int
responseText := strings.Builder{} responseText := strings.Builder{}
var streamErr *types.NewAPIError var streamErr *types.NewAPIError
var normalStreamStarted bool


helper.StreamScannerHandler(c, resp, info, func(data string) bool { helper.StreamScannerHandler(c, resp, info, func(data string) bool {
var geminiResponse dto.GeminiChatResponse var geminiResponse dto.GeminiChatResponse
@@ -1297,15 +1298,21 @@ func geminiStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http


if len(geminiResponse.Candidates) == 0 && geminiResponse.PromptFeedback != nil && geminiResponse.PromptFeedback.BlockReason != nil { if len(geminiResponse.Candidates) == 0 && geminiResponse.PromptFeedback != nil && geminiResponse.PromptFeedback.BlockReason != nil {
common.SetContextKey(c, constant.ContextKeyAdminRejectReason, fmt.Sprintf("gemini_block_reason=%s", *geminiResponse.PromptFeedback.BlockReason)) common.SetContextKey(c, constant.ContextKeyAdminRejectReason, fmt.Sprintf("gemini_block_reason=%s", *geminiResponse.PromptFeedback.BlockReason))
streamErr = types.NewOpenAIError(
apiErr := types.NewOpenAIError(
errors.New("request blocked by Gemini API: "+*geminiResponse.PromptFeedback.BlockReason), errors.New("request blocked by Gemini API: "+*geminiResponse.PromptFeedback.BlockReason),
types.ErrorCodePromptBlocked, types.ErrorCodePromptBlocked,
http.StatusBadRequest, http.StatusBadRequest,
) )
streamErr.UpstreamBody = service.TruncateBody(data)
apiErr.UpstreamBody = service.TruncateBody(data)
if !normalStreamStarted && !helper.StreamDataWritten(c) {
streamErr = apiErr
} else {
logger.LogError(c, "upstream gemini stream blocked after normal stream started: "+apiErr.Error())
}
return false return false
} }


normalStreamStarted = true
// 统计图片数量 // 统计图片数量
for _, candidate := range geminiResponse.Candidates { for _, candidate := range geminiResponse.Candidates {
for _, part := range candidate.Content.Parts { for _, part := range candidate.Content.Parts {


+ 59
- 0
relay/channel/gemini/relay_gemini_usage_test.go Просмотреть файл

@@ -328,6 +328,65 @@ func TestGeminiStreamHandlerReturnsErrorOnPromptBlock(t *testing.T) {
require.Contains(t, newAPIError.UpstreamBody, "SAFETY") require.Contains(t, newAPIError.UpstreamBody, "SAFETY")
} }


func TestGeminiStreamHandlerDoesNotRetryAfterNormalStreamData(t *testing.T) {
gin.SetMode(gin.TestMode)
oldStreamingTimeout := constant.StreamingTimeout
constant.StreamingTimeout = 300
t.Cleanup(func() {
constant.StreamingTimeout = oldStreamingTimeout
})

w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil)

normalChunk := dto.GeminiChatResponse{
Candidates: []dto.GeminiChatCandidate{
{
Content: dto.GeminiChatContent{
Role: "model",
Parts: []dto.GeminiPart{
{Text: "partial"},
},
},
},
},
}
blockReason := "SAFETY"
blockChunk := dto.GeminiChatResponse{
PromptFeedback: &dto.GeminiChatPromptFeedback{
BlockReason: &blockReason,
},
}

normalData, err := common.Marshal(normalChunk)
require.NoError(t, err)
blockData, err := common.Marshal(blockChunk)
require.NoError(t, err)

resp := &http.Response{
StatusCode: http.StatusOK,
Body: io.NopCloser(bytes.NewReader([]byte("data: " + string(normalData) + "\n\ndata: " + string(blockData) + "\n\n"))),
}
info := &relaycommon.RelayInfo{
RelayFormat: types.RelayFormatOpenAI,
ChannelMeta: &relaycommon.ChannelMeta{
UpstreamModelName: "gemini-test",
},
}

called := false
usage, newAPIError := geminiStreamHandler(c, info, resp, func(_ string, _ *dto.GeminiChatResponse) bool {
called = true
return true
})

require.Nil(t, newAPIError)
require.NotNil(t, usage)
require.True(t, called)
require.Empty(t, w.Body.String())
}

func TestGeminiTextGenerationHandlerUsesEstimatedPromptTokensWhenUsagePromptMissing(t *testing.T) { func TestGeminiTextGenerationHandlerUsesEstimatedPromptTokensWhenUsagePromptMissing(t *testing.T) {
t.Parallel() t.Parallel()




+ 1
- 1
relay/channel/openai/helper.go Просмотреть файл

@@ -29,7 +29,7 @@ func streamOpenAIErrorFromData(data string, statusCode int) *types.NewAPIError {
if openaiErr == nil { if openaiErr == nil {
return nil return nil
} }
if statusCode < 100 || statusCode > 599 {
if statusCode < 100 || statusCode > 599 || (statusCode >= 200 && statusCode < 300) {
statusCode = http.StatusInternalServerError statusCode = http.StatusInternalServerError
} }
apiErr := types.WithOpenAIError(*openaiErr, statusCode) apiErr := types.WithOpenAIError(*openaiErr, statusCode)


+ 4
- 2
relay/channel/openai/relay-openai.go Просмотреть файл

@@ -122,6 +122,7 @@ func OaiStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Re
var streamItems []string // store stream items var streamItems []string // store stream items
var lastStreamData string var lastStreamData string
var streamErr *types.NewAPIError var streamErr *types.NewAPIError
var normalStreamStarted bool
var secondLastStreamData string // 存储倒数第二个stream data,用于音频模型 var secondLastStreamData string // 存储倒数第二个stream data,用于音频模型


// 检查是否为音频模型 // 检查是否为音频模型
@@ -129,14 +130,15 @@ func OaiStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Re


helper.StreamScannerHandler(c, resp, info, func(data string) bool { helper.StreamScannerHandler(c, resp, info, func(data string) bool {
if apiErr := streamOpenAIErrorFromData(data, resp.StatusCode); apiErr != nil { if apiErr := streamOpenAIErrorFromData(data, resp.StatusCode); apiErr != nil {
if !c.Writer.Written() {
if !normalStreamStarted && !helper.StreamDataWritten(c) {
streamErr = apiErr streamErr = apiErr
} else { } else {
logger.LogError(c, "upstream stream error after downstream response was written: "+apiErr.Error())
logger.LogError(c, "upstream stream error after normal stream started: "+apiErr.Error())
} }
return false return false
} }


normalStreamStarted = true
if lastStreamData != "" { if lastStreamData != "" {
err := HandleStreamFormat(c, info, lastStreamData, info.ChannelSetting.ForceFormat, info.ChannelSetting.ThinkingToContent) err := HandleStreamFormat(c, info, lastStreamData, info.ChannelSetting.ForceFormat, info.ChannelSetting.ThinkingToContent)
if err != nil { if err != nil {


+ 8
- 1
relay/channel/openai/relay_responses.go Просмотреть файл

@@ -85,6 +85,7 @@ func OaiResponsesStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp
var usage = &dto.Usage{} var usage = &dto.Usage{}
var responseTextBuilder strings.Builder var responseTextBuilder strings.Builder
var streamErr *types.NewAPIError var streamErr *types.NewAPIError
var normalStreamStarted bool


helper.StreamScannerHandler(c, resp, info, func(data string) bool { helper.StreamScannerHandler(c, resp, info, func(data string) bool {


@@ -92,10 +93,16 @@ func OaiResponsesStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp
if err := common.UnmarshalJsonStr(data, &streamResponse); err == nil { if err := common.UnmarshalJsonStr(data, &streamResponse); err == nil {
switch streamResponse.Type { switch streamResponse.Type {
case "response.error", "response.failed", "error": case "response.error", "response.failed", "error":
streamErr = handleResponsesStreamError(streamResponse, data)
apiErr := handleResponsesStreamError(streamResponse, data)
if !normalStreamStarted && !helper.StreamDataWritten(c) {
streamErr = apiErr
} else if apiErr != nil {
logger.LogError(c, "upstream responses stream error after normal stream started: "+apiErr.Error())
}
return false return false
} }


normalStreamStarted = true
sendResponsesStreamData(c, streamResponse, data) sendResponsesStreamData(c, streamResponse, data)
switch streamResponse.Type { switch streamResponse.Type {
case "response.completed": case "response.completed":


+ 100
- 0
relay/channel/openai/upstream_body_test.go Просмотреть файл

@@ -11,6 +11,7 @@ import (
"github.com/QuantumNous/new-api/constant" "github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/dto" "github.com/QuantumNous/new-api/dto"
relaycommon "github.com/QuantumNous/new-api/relay/common" relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/relay/helper"
"github.com/QuantumNous/new-api/types" "github.com/QuantumNous/new-api/types"
"github.com/gin-gonic/gin" "github.com/gin-gonic/gin"
"github.com/stretchr/testify/assert" "github.com/stretchr/testify/assert"
@@ -109,6 +110,7 @@ func TestOaiResponsesStreamHandlerDoesNotForwardErrorEvent(t *testing.T) {
usage, err := OaiResponsesStreamHandler(c, &relaycommon.RelayInfo{}, resp) usage, err := OaiResponsesStreamHandler(c, &relaycommon.RelayInfo{}, resp)
require.Nil(t, usage) require.Nil(t, usage)
require.Error(t, err) require.Error(t, err)
assert.Equal(t, http.StatusInternalServerError, err.StatusCode)
assert.Contains(t, err.UpstreamBody, "rate_limit_exceeded") assert.Contains(t, err.UpstreamBody, "rate_limit_exceeded")
assert.NotContains(t, w.Body.String(), "rate_limit_exceeded") assert.NotContains(t, w.Body.String(), "rate_limit_exceeded")
assert.Empty(t, w.Body.String()) assert.Empty(t, w.Body.String())
@@ -148,6 +150,104 @@ func TestOaiStreamHandlerDoesNotForwardInitialErrorEvent(t *testing.T) {
assert.Empty(t, w.Header().Get("Content-Type")) assert.Empty(t, w.Header().Get("Content-Type"))
} }


func TestOaiStreamHandlerReturnsInitialErrorAfterPing(t *testing.T) {
gin.SetMode(gin.TestMode)
oldTimeout := constant.StreamingTimeout
constant.StreamingTimeout = 30
t.Cleanup(func() {
constant.StreamingTimeout = oldTimeout
})

w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil)

require.NoError(t, helper.PingData(c))

errorEvent := `{"error":{"type":"tokens","code":"rate_limit_exceeded","message":"Request too large","param":null}}`
resp := &http.Response{
StatusCode: http.StatusOK,
Body: io.NopCloser(strings.NewReader("data: " + errorEvent + "\n\n")),
}
info := &relaycommon.RelayInfo{
RelayFormat: types.RelayFormatOpenAI,
ChannelMeta: &relaycommon.ChannelMeta{
UpstreamModelName: "gpt-test",
},
}

usage, err := OaiStreamHandler(c, info, resp)
require.Nil(t, usage)
require.Error(t, err)
assert.Equal(t, http.StatusInternalServerError, err.StatusCode)
assert.Contains(t, err.UpstreamBody, "rate_limit_exceeded")
assert.Contains(t, w.Body.String(), ": PING")
assert.NotContains(t, w.Body.String(), "rate_limit_exceeded")
}

func TestOaiStreamHandlerDoesNotRetryAfterNormalStreamData(t *testing.T) {
gin.SetMode(gin.TestMode)
oldTimeout := constant.StreamingTimeout
constant.StreamingTimeout = 30
t.Cleanup(func() {
constant.StreamingTimeout = oldTimeout
})

w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil)

normalEvent := `{"id":"chatcmpl-test","object":"chat.completion.chunk","created":1,"model":"gpt-test","choices":[{"index":0,"delta":{"content":"hi"}}]}`
errorEvent := `{"error":{"type":"tokens","code":"rate_limit_exceeded","message":"Request too large","param":null}}`
resp := &http.Response{
StatusCode: http.StatusOK,
Body: io.NopCloser(strings.NewReader("data: " + normalEvent + "\n\ndata: " + errorEvent + "\n\n")),
}
info := &relaycommon.RelayInfo{
RelayFormat: types.RelayFormatOpenAI,
ChannelMeta: &relaycommon.ChannelMeta{
UpstreamModelName: "gpt-test",
},
}

usage, err := OaiStreamHandler(c, info, resp)
require.Nil(t, err)
require.NotNil(t, usage)
assert.Contains(t, w.Body.String(), "hi")
assert.NotContains(t, w.Body.String(), "rate_limit_exceeded")
}

func TestOaiResponsesStreamHandlerDoesNotRetryAfterNormalStreamData(t *testing.T) {
gin.SetMode(gin.TestMode)
oldTimeout := constant.StreamingTimeout
constant.StreamingTimeout = 30
t.Cleanup(func() {
constant.StreamingTimeout = oldTimeout
})

w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil)

normalEvent := `{"type":"response.completed","response":{"id":"resp_test","usage":{"input_tokens":1,"output_tokens":1,"total_tokens":2}}}`
errorEvent := `{"type":"error","error":{"type":"tokens","code":"rate_limit_exceeded","message":"Request too large","param":null},"sequence_number":2}`
resp := &http.Response{
StatusCode: http.StatusOK,
Body: io.NopCloser(strings.NewReader("data: " + normalEvent + "\n\ndata: " + errorEvent + "\n\n")),
}
info := &relaycommon.RelayInfo{
ChannelMeta: &relaycommon.ChannelMeta{
UpstreamModelName: "gpt-test",
},
}

usage, err := OaiResponsesStreamHandler(c, info, resp)
require.Nil(t, err)
require.NotNil(t, usage)
assert.Contains(t, w.Body.String(), "response.completed")
assert.NotContains(t, w.Body.String(), "rate_limit_exceeded")
}

func TestOaiResponsesToChatHandlerAttachesUpstreamBodyOnInvalidJSON(t *testing.T) { func TestOaiResponsesToChatHandlerAttachesUpstreamBodyOnInvalidJSON(t *testing.T) {
gin.SetMode(gin.TestMode) gin.SetMode(gin.TestMode)
w := httptest.NewRecorder() w := httptest.NewRecorder()


+ 40
- 1
relay/helper/common.go Просмотреть файл

@@ -14,6 +14,40 @@ import (
"github.com/gorilla/websocket" "github.com/gorilla/websocket"
) )


const (
streamBodyWrittenKey = "stream_body_written"
streamDataWrittenKey = "stream_data_written"
)

func MarkStreamBodyWritten(c *gin.Context) {
if c == nil {
return
}
c.Set(streamBodyWrittenKey, true)
}

func StreamBodyWritten(c *gin.Context) bool {
if c == nil {
return false
}
return c.GetBool(streamBodyWrittenKey)
}

func MarkStreamDataWritten(c *gin.Context) {
if c == nil {
return
}
MarkStreamBodyWritten(c)
c.Set(streamDataWrittenKey, true)
}

func StreamDataWritten(c *gin.Context) bool {
if c == nil {
return false
}
return c.GetBool(streamDataWrittenKey)
}

func FlushWriter(c *gin.Context) (err error) { func FlushWriter(c *gin.Context) (err error) {
defer func() { defer func() {
if r := recover(); r != nil { if r := recover(); r != nil {
@@ -55,7 +89,7 @@ func SetEventStreamHeaders(c *gin.Context) {
} }


func ClearEventStreamHeadersIfNotWritten(c *gin.Context) { func ClearEventStreamHeadersIfNotWritten(c *gin.Context) {
if c == nil || c.Writer == nil || c.Writer.Written() {
if c == nil || c.Writer == nil || StreamBodyWritten(c) {
return return
} }
c.Set("event_stream_headers_set", false) c.Set("event_stream_headers_set", false)
@@ -74,6 +108,7 @@ func ClaudeData(c *gin.Context, resp dto.ClaudeResponse) error {
} else { } else {
c.Render(-1, common.CustomEvent{Data: fmt.Sprintf("event: %s\n", resp.Type)}) c.Render(-1, common.CustomEvent{Data: fmt.Sprintf("event: %s\n", resp.Type)})
c.Render(-1, common.CustomEvent{Data: "data: " + string(jsonData)}) c.Render(-1, common.CustomEvent{Data: "data: " + string(jsonData)})
MarkStreamDataWritten(c)
} }
_ = FlushWriter(c) _ = FlushWriter(c)
return nil return nil
@@ -82,12 +117,14 @@ func ClaudeData(c *gin.Context, resp dto.ClaudeResponse) error {
func ClaudeChunkData(c *gin.Context, resp dto.ClaudeResponse, data string) { func ClaudeChunkData(c *gin.Context, resp dto.ClaudeResponse, data string) {
c.Render(-1, common.CustomEvent{Data: fmt.Sprintf("event: %s\n", resp.Type)}) c.Render(-1, common.CustomEvent{Data: fmt.Sprintf("event: %s\n", resp.Type)})
c.Render(-1, common.CustomEvent{Data: fmt.Sprintf("data: %s\n", data)}) c.Render(-1, common.CustomEvent{Data: fmt.Sprintf("data: %s\n", data)})
MarkStreamDataWritten(c)
_ = FlushWriter(c) _ = FlushWriter(c)
} }


func ResponseChunkData(c *gin.Context, resp dto.ResponsesStreamResponse, data string) { func ResponseChunkData(c *gin.Context, resp dto.ResponsesStreamResponse, data string) {
c.Render(-1, common.CustomEvent{Data: fmt.Sprintf("event: %s\n", resp.Type)}) c.Render(-1, common.CustomEvent{Data: fmt.Sprintf("event: %s\n", resp.Type)})
c.Render(-1, common.CustomEvent{Data: fmt.Sprintf("data: %s", data)}) c.Render(-1, common.CustomEvent{Data: fmt.Sprintf("data: %s", data)})
MarkStreamDataWritten(c)
_ = FlushWriter(c) _ = FlushWriter(c)
} }


@@ -101,6 +138,7 @@ func StringData(c *gin.Context, str string) error {
} }


c.Render(-1, common.CustomEvent{Data: "data: " + str}) c.Render(-1, common.CustomEvent{Data: "data: " + str})
MarkStreamDataWritten(c)
return FlushWriter(c) return FlushWriter(c)
} }


@@ -116,6 +154,7 @@ func PingData(c *gin.Context) error {
if _, err := c.Writer.Write([]byte(": PING\n\n")); err != nil { if _, err := c.Writer.Write([]byte(": PING\n\n")); err != nil {
return fmt.Errorf("write ping data failed: %w", err) return fmt.Errorf("write ping data failed: %w", err)
} }
MarkStreamBodyWritten(c)
return FlushWriter(c) return FlushWriter(c)
} }




+ 208
- 0
test/e2e/mock_capture_server.go Просмотреть файл

@@ -0,0 +1,208 @@
// +build ignore

package main

import (
"encoding/json"
"fmt"
"log"
"net/http"
"os"
"time"

"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/middleware"
"github.com/gin-gonic/gin"
)

func main() {
gin.SetMode(gin.ReleaseMode)
middleware.SetCaptureEnabled(1, true)

router := gin.New()
router.Use(gin.Recovery())

// 模拟 TokenAuth
router.Use(func(c *gin.Context) {
auth := c.GetHeader("Authorization")
if auth == "" || auth == "Bearer disabled" {
c.JSON(401, gin.H{"error": gin.H{"message": "Unauthorized"}})
c.Abort()
return
}
c.Set("id", 1)
c.Next()
})

// 模拟 RequestId
router.Use(func(c *gin.Context) {
reqID := c.GetHeader(common.RequestIdKey)
if reqID == "" {
reqID = fmt.Sprintf("req-%d", time.Now().UnixNano())
}
c.Set(common.RequestIdKey, reqID)
c.Next()
})

router.Use(middleware.RelayCaptureMiddleware())

// ═══════════════════════════════════════════════════════
// OpenAI Chat Completions API
// ═══════════════════════════════════════════════════════
router.POST("/v1/chat/completions", func(c *gin.Context) {
stream := c.Query("stream") == "true"
if stream {
c.Header("Content-Type", "text/event-stream")
f := c.Writer.(http.Flusher)
for _, s := range []string{
"data: {\"id\":\"chatcmpl-001\",\"choices\":[{\"delta\":{\"content\":\"Hello \"}}]}\n\n",
"data: {\"id\":\"chatcmpl-001\",\"choices\":[{\"delta\":{\"content\":\"World!\"}}]}\n\n",
"data: [DONE]\n\n",
} {
c.Writer.Write([]byte(s))
f.Flush()
}
return
}
c.JSON(200, gin.H{
"id": "chatcmpl-001",
"object": "chat.completion",
"model": "gpt-4",
"choices": []gin.H{{"index": 0, "message": gin.H{"role": "assistant", "content": "Hello from Chat Completions!"}}},
"usage": gin.H{"prompt_tokens": 5, "completion_tokens": 5, "total_tokens": 10},
})
})

// ═══════════════════════════════════════════════════════
// OpenAI Responses API (new)
// ═══════════════════════════════════════════════════════
router.POST("/v1/responses", func(c *gin.Context) {
c.JSON(200, gin.H{
"id": "resp_abc123",
"object": "response",
"status": "completed",
"output": []gin.H{{
"type": "message",
"role": "assistant",
"content": []gin.H{{
"type": "output_text",
"text": "Hello from Responses API!",
}},
}},
"usage": gin.H{"input_tokens": 8, "output_tokens": 4, "total_tokens": 12},
})
})

// ═══════════════════════════════════════════════════════
// Claude Messages API (Anthropic)
// ═══════════════════════════════════════════════════════
router.POST("/v1/messages", func(c *gin.Context) {
c.JSON(200, gin.H{
"id": "msg_01AbCdEfGhIjKlMnOpQrStUv",
"type": "message",
"role": "assistant",
"model": "claude-sonnet-4-6",
"content": []gin.H{{"type": "text", "text": "Hello! How can I assist you today?"}},
"stop_reason": "end_turn",
"usage": gin.H{"input_tokens": 8, "output_tokens": 12},
})
})

// ═══════════════════════════════════════════════════════
// Gemini Native API (/v1beta/models/*)
// ═══════════════════════════════════════════════════════
router.POST("/v1beta/models/*path", func(c *gin.Context) {
path := c.Param("path")
c.JSON(200, gin.H{
"candidates": []gin.H{{
"content": gin.H{
"role": "model",
"parts": []gin.H{{"text": fmt.Sprintf("Gemini response for: /%s", path)}},
},
"finishReason": "STOP",
}},
"usageMetadata": gin.H{"promptTokenCount": 5, "candidatesTokenCount": 10, "totalTokenCount": 15},
})
})

// ═══════════════════════════════════════════════════════
// Gemini via /v1/models/* (main route)
// ═══════════════════════════════════════════════════════
router.POST("/v1/models/*path", func(c *gin.Context) {
path := c.Param("path")
c.JSON(200, gin.H{
"candidates": []gin.H{{
"content": gin.H{
"role": "model",
"parts": []gin.H{{"text": fmt.Sprintf("Gemini(v1) response for: /%s", path)}},
},
"finishReason": "STOP",
}},
})
})

// ═══════════════════════════════════════════════════════
// Models list
// ═══════════════════════════════════════════════════════
router.GET("/v1/models", func(c *gin.Context) {
c.JSON(200, gin.H{"data": []gin.H{{"id": "gpt-4"}, {"id": "claude-3"}, {"id": "gemini-2.0-flash"}}})
})

// ═══════════════════════════════════════════════════════
// Error endpoint for testing error status codes
// ═══════════════════════════════════════════════════════
router.POST("/v1/error", func(c *gin.Context) {
code := 500
if c := c.Query("code"); c != "" {
fmt.Sscanf(c, "%d", &code)
}
c.JSON(code, gin.H{"error": gin.H{"message": "test error", "type": "test_error"}})
})

// ═══════════════════════════════════════════════════════
// Capture file management endpoints
// ═══════════════════════════════════════════════════════
router.GET("/health", func(c *gin.Context) {
c.JSON(200, gin.H{"status": "ok"})
})

router.GET("/capture-files", func(c *gin.Context) {
today := time.Now().Format("2006-01-02")
dir := fmt.Sprintf("./data/relay-capture/%s", today)
entries, err := os.ReadDir(dir)
if err != nil {
c.JSON(200, gin.H{"files": []string{}, "error": err.Error()})
return
}
var files []string
for _, e := range entries {
files = append(files, e.Name())
}
c.JSON(200, gin.H{"dir": dir, "files": files})
})

router.GET("/capture-files/:name", func(c *gin.Context) {
today := time.Now().Format("2006-01-02")
path := fmt.Sprintf("./data/relay-capture/%s/%s", today, c.Param("name"))
data, err := os.ReadFile(path)
if err != nil {
c.JSON(404, gin.H{"error": err.Error()})
return
}
var record map[string]interface{}
json.Unmarshal(data, &record)
c.JSON(200, record)
})

router.DELETE("/capture-files", func(c *gin.Context) {
os.RemoveAll("./data/relay-capture")
c.JSON(200, gin.H{"status": "cleaned"})
})

port := "18080"
if p := os.Getenv("PORT"); p != "" {
port = p
}
log.Printf("Mock Relay Server on :%s (chat|responses|messages|gemini)", port)
http.ListenAndServe(":"+port, router)
}

Загрузка…
Отмена
Сохранить