Преглед изворни кода

feat: 端点格式限制功能 + 完善测试

当开启端点格式限制后,不同协议的渠道只能通过对应的端点访问:
- Anthropic 渠道只能用 /v1/messages
- Gemini 渠道只能用 /v1beta/models/*
- OpenAI 系渠道只能用 /v1/chat/completions 等
- VertexAI/阿里云等支持多协议的渠道不受限

核心实现:EndpointFormatGuard 中间件从请求路径推导协议家族,
不依赖 gin context 中的 relay_format(因 group middleware 执行
时 context 中尚无此值)。

修复了 /v1/models/*path 和 /v1/engines/ 路径被错误归类为
openai 家族的 bug(实际是 Gemini 格式)。

单元测试 58 个,E2E 测试 15 个,覆盖:
- 路径分类边界(含 Gemini relay 路径修复验证)
- 渠道类型覆盖(含非 LLM 渠道透传)
- 并发安全、运行时开关切换
- 错误消息 JSON 结构和内容验证
master
fengsilin пре 3 недеља
родитељ
комит
84eb6d6a88
11 измењених фајлова са 1459 додато и 59 уклоњено
  1. +2
    -0
      constant/context_key.go
  2. +1
    -0
      main.go
  3. +142
    -0
      middleware/endpoint_format_guard.go
  4. +591
    -0
      middleware/endpoint_format_guard_test.go
  5. +4
    -4
      model/option.go
  6. +28
    -54
      router/relay-router.go
  7. +7
    -1
      setting/model_setting/global.go
  8. +245
    -0
      test/e2e_endpoint_format.go
  9. +420
    -0
      test/e2e_endpoint_format.py
  10. +1
    -0
      web/src/components/settings/ModelSetting.jsx
  11. +18
    -0
      web/src/pages/Setting/Model/SettingGlobalModel.jsx

+ 2
- 0
constant/context_key.go Прегледај датотеку

@@ -37,6 +37,8 @@ const (
ContextKeyChannelMultiKeyIndex ContextKey = "channel_multi_key_index"
ContextKeyChannelKey ContextKey = "channel_key"

ContextKeyRelayFormat ContextKey = "relay_format"

ContextKeyAutoGroup ContextKey = "auto_group"
ContextKeyAutoGroupIndex ContextKey = "auto_group_index"
ContextKeyAutoGroupRetryIndex ContextKey = "auto_group_retry_index"


+ 1
- 0
main.go Прегледај датотеку

@@ -286,6 +286,7 @@ func InitResources() error {

// Initialize options, should after model.InitDB()
model.InitOptionMap()
model.LoadOptionsFromDatabase()

// 清理旧的磁盘缓存文件
common.CleanupOldCacheFiles()


+ 142
- 0
middleware/endpoint_format_guard.go Прегледај датотеку

@@ -0,0 +1,142 @@
package middleware

import (
"fmt"
"net/http"
"slices"
"strings"

"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/setting/model_setting"

"github.com/gin-gonic/gin"
)

// pathToFamily maps URL path prefixes to format families.
// The middleware uses this to determine the request's protocol family
// directly from the request path, without relying on context values.
var pathToFamily = []struct {
prefix string
family string
}{
{"/v1/messages", "claude"},
{"/v1beta/", "gemini"},
{"/v1/models/", "gemini"}, // Gemini relay via /v1/models/*path
{"/v1/engines/", "gemini"}, // Gemini relay via /v1/engines/:model/embeddings
// All other /v1/ paths default to "openai"
}

// getFamilyFromPath returns the format family for a given request path.
func getFamilyFromPath(path string) string {
for _, rule := range pathToFamily {
if strings.HasPrefix(path, rule.prefix) {
return rule.family
}
}
if strings.HasPrefix(path, "/v1/") {
return "openai"
}
return ""
}

var channelNativeFamilies = map[int][]string{
constant.ChannelTypeOpenAI: {"openai"},
constant.ChannelTypeAzure: {"openai"},
constant.ChannelTypeOllama: {"openai"},
constant.ChannelTypeDeepSeek: {"openai"},
constant.ChannelTypeOpenRouter: {"openai"},
constant.ChannelTypeXinference: {"openai"},
constant.ChannelTypeXai: {"openai"},
constant.ChannelTypeMistral: {"openai"},
constant.ChannelTypeSiliconFlow: {"openai"},
constant.ChannelTypePerplexity: {"openai"},
constant.ChannelTypeCodex: {"openai"},
constant.ChannelTypeMiniMax: {"openai"},
constant.ChannelTypeMoonshot: {"openai"},
constant.ChannelTypeBaidu: {"openai"},
constant.ChannelTypeBaiduV2: {"openai"},
constant.ChannelTypeZhipu: {"openai"},
constant.ChannelTypeZhipu_v4: {"openai"},
constant.ChannelTypeTencent: {"openai"},
constant.ChannelTypeXunfei: {"openai"},
constant.ChannelTypeCohere: {"openai"},
constant.ChannelTypeDify: {"openai"},
constant.ChannelTypeCoze: {"openai"},
constant.ChannelTypeVolcEngine: {"openai"},
constant.ChannelTypeReplicate: {"openai"},
constant.ChannelTypeSubmodel: {"openai"},
constant.ChannelTypeJina: {"openai"},
constant.ChannelTypePaLM: {"openai"},
constant.ChannelTypeAnthropic: {"claude"},
constant.ChannelTypeAws: {"claude"},
constant.ChannelTypeGemini: {"gemini"},
constant.ChannelTypeVertexAi: {"openai", "claude", "gemini"},
constant.ChannelTypeAli: {"openai", "claude"},
}

var nativeEndpoints = map[string]string{
"openai": "/v1/chat/completions",
"claude": "/v1/messages",
"gemini": "/v1beta/models/{model}:generateContent",
}

func isChannelAllowed(channelType int, family string) bool {
families, ok := channelNativeFamilies[channelType]
if !ok {
return true
}
return slices.Contains(families, family)
}

func EndpointFormatGuard() gin.HandlerFunc {
return func(c *gin.Context) {
if !model_setting.IsEndpointFormatRestrictionEnabled() {
c.Next()
return
}

family := getFamilyFromPath(c.Request.URL.Path)
if family == "" {
c.Next()
return
}

channelType := common.GetContextKeyInt(c, constant.ContextKeyChannelType)
if channelType == 0 {
c.Next()
return
}

if isChannelAllowed(channelType, family) {
c.Next()
return
}

channelTypeName := constant.GetChannelTypeName(channelType)

var nativeEndpoint string
if families, ok := channelNativeFamilies[channelType]; ok && len(families) > 0 {
nativeEndpoint = nativeEndpoints[families[0]]
}

c.JSON(http.StatusBadRequest, gin.H{
"error": gin.H{
"message": fmt.Sprintf(
"Model is on a %s channel and cannot be accessed via %s endpoints. Please use the native %s endpoint (%s) instead.",
channelTypeName, family, channelTypeName, nativeEndpoint,
),
"type": "new_api_error",
"code": "endpoint_format_restricted",
},
})
c.Abort()
}
}

// setTestRestrictionEnabled is a test helper.
func setTestRestrictionEnabled(enabled bool) bool {
prev := model_setting.IsEndpointFormatRestrictionEnabled()
model_setting.GetGlobalSettings().EndpointFormatRestrictionEnabled = enabled
return prev
}

+ 591
- 0
middleware/endpoint_format_guard_test.go Прегледај датотеку

@@ -0,0 +1,591 @@
package middleware

import (
"encoding/json"
"net/http/httptest"
"sync"
"sync/atomic"
"testing"

"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)

// ============================================================================
// Group A: getFamilyFromPath 纯函数测试
// ============================================================================

func TestGetFamilyFromPath_Claude(t *testing.T) {
assert.Equal(t, "claude", getFamilyFromPath("/v1/messages"))
}

func TestGetFamilyFromPath_ClaudeSubpath(t *testing.T) {
assert.Equal(t, "claude", getFamilyFromPath("/v1/messages/"))
}

func TestGetFamilyFromPath_ClaudeBatchPath(t *testing.T) {
assert.Equal(t, "claude", getFamilyFromPath("/v1/messages/batches"))
}

func TestGetFamilyFromPath_Gemini(t *testing.T) {
assert.Equal(t, "gemini", getFamilyFromPath("/v1beta/models/gemini-pro:generateContent"))
assert.Equal(t, "gemini", getFamilyFromPath("/v1beta/"))
}

func TestGetFamilyFromPath_GeminiCompatPath(t *testing.T) {
assert.Equal(t, "gemini", getFamilyFromPath("/v1beta/openai/models"))
}

func TestGetFamilyFromPath_GeminiStreamPath(t *testing.T) {
assert.Equal(t, "gemini", getFamilyFromPath("/v1beta/models/gemini-2.0-flash:streamGenerateContent"))
}

func TestGetFamilyFromPath_V1ModelsPath(t *testing.T) {
// /v1/models/*path 用于 Gemini relay,应归类为 gemini 家族
assert.Equal(t, "gemini", getFamilyFromPath("/v1/models/gemini-pro:generateContent"))
assert.Equal(t, "gemini", getFamilyFromPath("/v1/models/gemini-2.0-flash:streamGenerateContent"))
}

func TestGetFamilyFromPath_V1EnginesPath(t *testing.T) {
// /v1/engines/:model/embeddings 用于 Gemini relay,应归类为 gemini 家族
assert.Equal(t, "gemini", getFamilyFromPath("/v1/engines/text-embedding-004/embeddings"))
}

func TestGetFamilyFromPath_OpenAI(t *testing.T) {
assert.Equal(t, "openai", getFamilyFromPath("/v1/chat/completions"))
assert.Equal(t, "openai", getFamilyFromPath("/v1/completions"))
assert.Equal(t, "openai", getFamilyFromPath("/v1/embeddings"))
assert.Equal(t, "openai", getFamilyFromPath("/v1/audio/speech"))
assert.Equal(t, "openai", getFamilyFromPath("/v1/images/generations"))
assert.Equal(t, "openai", getFamilyFromPath("/v1/rerank"))
assert.Equal(t, "openai", getFamilyFromPath("/v1/moderations"))
}

func TestGetFamilyFromPath_Realtime(t *testing.T) {
assert.Equal(t, "openai", getFamilyFromPath("/v1/realtime"))
}

func TestGetFamilyFromPath_Responses(t *testing.T) {
assert.Equal(t, "openai", getFamilyFromPath("/v1/responses"))
assert.Equal(t, "openai", getFamilyFromPath("/v1/responses/compact"))
}

func TestGetFamilyFromPath_Unknown(t *testing.T) {
assert.Equal(t, "", getFamilyFromPath("/api/status"))
assert.Equal(t, "", getFamilyFromPath("/pg/chat/completions"))
assert.Equal(t, "", getFamilyFromPath("/"))
assert.Equal(t, "", getFamilyFromPath("/v1"))
}

// ============================================================================
// Group B: isChannelAllowed 纯函数测试
// ============================================================================

func TestIsChannelAllowed_OpenAIChannel_OpenAIFormat(t *testing.T) {
assert.True(t, isChannelAllowed(constant.ChannelTypeOpenAI, "openai"))
}

func TestIsChannelAllowed_OpenAIChannel_ClaudeFormat(t *testing.T) {
assert.False(t, isChannelAllowed(constant.ChannelTypeOpenAI, "claude"))
}

func TestIsChannelAllowed_OpenAIChannel_GeminiFormat(t *testing.T) {
assert.False(t, isChannelAllowed(constant.ChannelTypeOpenAI, "gemini"))
}

func TestIsChannelAllowed_ClaudeChannel_ClaudeFormat(t *testing.T) {
assert.True(t, isChannelAllowed(constant.ChannelTypeAnthropic, "claude"))
}

func TestIsChannelAllowed_ClaudeChannel_OpenAIFormat(t *testing.T) {
assert.False(t, isChannelAllowed(constant.ChannelTypeAnthropic, "openai"))
}

func TestIsChannelAllowed_GeminiChannel(t *testing.T) {
assert.True(t, isChannelAllowed(constant.ChannelTypeGemini, "gemini"))
assert.False(t, isChannelAllowed(constant.ChannelTypeGemini, "openai"))
assert.False(t, isChannelAllowed(constant.ChannelTypeGemini, "claude"))
}

func TestIsChannelAllowed_AzureChannel(t *testing.T) {
assert.True(t, isChannelAllowed(constant.ChannelTypeAzure, "openai"))
assert.False(t, isChannelAllowed(constant.ChannelTypeAzure, "claude"))
assert.False(t, isChannelAllowed(constant.ChannelTypeAzure, "gemini"))
}

func TestIsChannelAllowed_DeepSeekChannel(t *testing.T) {
assert.True(t, isChannelAllowed(constant.ChannelTypeDeepSeek, "openai"))
assert.False(t, isChannelAllowed(constant.ChannelTypeDeepSeek, "claude"))
assert.False(t, isChannelAllowed(constant.ChannelTypeDeepSeek, "gemini"))
}

func TestIsChannelAllowed_VertexAI_AllFormats(t *testing.T) {
assert.True(t, isChannelAllowed(constant.ChannelTypeVertexAi, "openai"))
assert.True(t, isChannelAllowed(constant.ChannelTypeVertexAi, "claude"))
assert.True(t, isChannelAllowed(constant.ChannelTypeVertexAi, "gemini"))
}

func TestIsChannelAllowed_Ali(t *testing.T) {
assert.True(t, isChannelAllowed(constant.ChannelTypeAli, "openai"))
assert.True(t, isChannelAllowed(constant.ChannelTypeAli, "claude"))
assert.False(t, isChannelAllowed(constant.ChannelTypeAli, "gemini"))
}

func TestIsChannelAllowed_Midjourney_PassThrough(t *testing.T) {
// Midjourney 不在 channelNativeFamilies 中,所有格式都应放行
assert.True(t, isChannelAllowed(constant.ChannelTypeMidjourney, "openai"))
assert.True(t, isChannelAllowed(constant.ChannelTypeMidjourney, "claude"))
assert.True(t, isChannelAllowed(constant.ChannelTypeMidjourney, "gemini"))
}

func TestIsChannelAllowed_Kling_PassThrough(t *testing.T) {
// Kling 不在 channelNativeFamilies 中,所有格式都应放行
assert.True(t, isChannelAllowed(constant.ChannelTypeKling, "openai"))
assert.True(t, isChannelAllowed(constant.ChannelTypeKling, "claude"))
}

func TestIsChannelAllowed_Suno_PassThrough(t *testing.T) {
// Suno 不在 channelNativeFamilies 中
assert.True(t, isChannelAllowed(constant.ChannelTypeSunoAPI, "openai"))
}

func TestIsChannelAllowed_UnknownChannel(t *testing.T) {
assert.True(t, isChannelAllowed(999, "openai"))
assert.True(t, isChannelAllowed(999, "claude"))
assert.True(t, isChannelAllowed(999, "gemini"))
}

func TestIsChannelAllowed_EmptyFamily(t *testing.T) {
// 空 family 表示路径不匹配任何已知格式,guard 会直接放行
// isChannelAllowed 不会被空 family 调用,但验证行为:openai 在 families 中不含 ""
assert.False(t, isChannelAllowed(constant.ChannelTypeOpenAI, ""))
}

// ============================================================================
// Helper: runGuardTest
// ============================================================================

func runGuardTestWithMethod(method, path string, channelType int) (bool, *httptest.ResponseRecorder) {
called := false
gin.SetMode(gin.TestMode)

r := gin.New()
r.Use(func(c *gin.Context) {
if channelType >= 0 {
c.Set(string(constant.ContextKeyChannelType), channelType)
}
c.Next()
})
r.Use(EndpointFormatGuard())

// 注册所有相关路由
r.POST("/v1/chat/completions", func(ctx *gin.Context) { called = true })
r.POST("/v1/completions", func(ctx *gin.Context) { called = true })
r.POST("/v1/messages", func(ctx *gin.Context) { called = true })
r.POST("/v1/embeddings", func(ctx *gin.Context) { called = true })
r.POST("/v1/rerank", func(ctx *gin.Context) { called = true })
r.POST("/v1/responses", func(ctx *gin.Context) { called = true })
r.POST("/v1/responses/compact", func(ctx *gin.Context) { called = true })
r.POST("/v1/moderations", func(ctx *gin.Context) { called = true })
r.POST("/v1/audio/speech", func(ctx *gin.Context) { called = true })
r.POST("/v1/images/generations", func(ctx *gin.Context) { called = true })
r.POST("/v1/models/*path", func(ctx *gin.Context) { called = true })
r.POST("/v1/engines/:model/embeddings", func(ctx *gin.Context) { called = true })
r.GET("/v1/realtime", func(ctx *gin.Context) { called = true })
r.Any("/v1beta/models/*path", func(ctx *gin.Context) { called = true })

w := httptest.NewRecorder()
req := httptest.NewRequest(method, path, nil)
r.ServeHTTP(w, req)

return called, w
}

func runGuardTest(path string, channelType int) (bool, *httptest.ResponseRecorder) {
return runGuardTestWithMethod("POST", path, channelType)
}

// ============================================================================
// Group C: 中间件集成测试
// ============================================================================

func TestEndpointFormatGuard_Disabled(t *testing.T) {
orig := setTestRestrictionEnabled(false)
defer setTestRestrictionEnabled(orig)

called, _ := runGuardTest("/v1/chat/completions", constant.ChannelTypeAnthropic)
assert.True(t, called, "handler should be called when restriction is disabled")
}

func TestEndpointFormatGuard_Disabled_MultipleChannelTypes(t *testing.T) {
orig := setTestRestrictionEnabled(false)
defer setTestRestrictionEnabled(orig)

for _, ch := range []int{constant.ChannelTypeOpenAI, constant.ChannelTypeAnthropic, constant.ChannelTypeGemini, constant.ChannelTypeMidjourney} {
called, _ := runGuardTest("/v1/chat/completions", ch)
assert.True(t, called, "all channel types should pass when restriction is disabled")
}
}

func TestEndpointFormatGuard_Enabled_Matched(t *testing.T) {
orig := setTestRestrictionEnabled(true)
defer setTestRestrictionEnabled(orig)

called, _ := runGuardTest("/v1/chat/completions", constant.ChannelTypeOpenAI)
assert.True(t, called, "handler should be called when format matches channel type")
}

func TestEndpointFormatGuard_Enabled_Mismatched(t *testing.T) {
orig := setTestRestrictionEnabled(true)
defer setTestRestrictionEnabled(orig)

called, w := runGuardTest("/v1/chat/completions", constant.ChannelTypeAnthropic)
assert.False(t, called, "handler should NOT be called when format mismatches")
assert.Equal(t, 400, w.Code)
}

func TestEndpointFormatGuard_Enabled_ClaudeEndpoint_ClaudeChannel(t *testing.T) {
orig := setTestRestrictionEnabled(true)
defer setTestRestrictionEnabled(orig)

called, _ := runGuardTest("/v1/messages", constant.ChannelTypeAnthropic)
assert.True(t, called, "claude endpoint + claude channel should be allowed")
}

func TestEndpointFormatGuard_Enabled_ClaudeEndpoint_OpenAIChannel(t *testing.T) {
orig := setTestRestrictionEnabled(true)
defer setTestRestrictionEnabled(orig)

called, w := runGuardTest("/v1/messages", constant.ChannelTypeOpenAI)
assert.False(t, called, "claude endpoint + openai channel should be blocked")
assert.Equal(t, 400, w.Code)
}

func TestEndpointFormatGuard_Enabled_GeminiEndpoint_GeminiChannel(t *testing.T) {
orig := setTestRestrictionEnabled(true)
defer setTestRestrictionEnabled(orig)

called, _ := runGuardTest("/v1beta/models/gemini-pro:generateContent", constant.ChannelTypeGemini)
assert.True(t, called, "gemini endpoint + gemini channel should be allowed")
}

func TestEndpointFormatGuard_Enabled_GeminiEndpoint_OpenAIChannel(t *testing.T) {
orig := setTestRestrictionEnabled(true)
defer setTestRestrictionEnabled(orig)

called, w := runGuardTest("/v1beta/models/gemini-pro:generateContent", constant.ChannelTypeOpenAI)
assert.False(t, called, "gemini endpoint + openai channel should be blocked")
assert.Equal(t, 400, w.Code)
}

func TestEndpointFormatGuard_Enabled_GeminiChannel_OpenAIEndpoint(t *testing.T) {
orig := setTestRestrictionEnabled(true)
defer setTestRestrictionEnabled(orig)

called, w := runGuardTest("/v1/chat/completions", constant.ChannelTypeGemini)
assert.False(t, called, "gemini channel should be blocked on openai endpoint")
assert.Equal(t, 400, w.Code)
}

func TestEndpointFormatGuard_Enabled_GeminiChannel_V1ModelsPath(t *testing.T) {
orig := setTestRestrictionEnabled(true)
defer setTestRestrictionEnabled(orig)

called, _ := runGuardTest("/v1/models/gemini-pro:generateContent", constant.ChannelTypeGemini)
assert.True(t, called, "gemini channel should be allowed on /v1/models/* path (gemini family)")
}

func TestEndpointFormatGuard_Enabled_DeepSeekAllowedOnOpenAI(t *testing.T) {
orig := setTestRestrictionEnabled(true)
defer setTestRestrictionEnabled(orig)

called, _ := runGuardTest("/v1/chat/completions", constant.ChannelTypeDeepSeek)
assert.True(t, called, "deepseek channel should be allowed on openai endpoint")
}

func TestEndpointFormatGuard_Enabled_DeepSeekBlockedOnClaude(t *testing.T) {
orig := setTestRestrictionEnabled(true)
defer setTestRestrictionEnabled(orig)

called, w := runGuardTest("/v1/messages", constant.ChannelTypeDeepSeek)
assert.False(t, called, "deepseek channel should be blocked on claude endpoint")
assert.Equal(t, 400, w.Code)
}

func TestEndpointFormatGuard_Enabled_MidjourneyPassThrough(t *testing.T) {
orig := setTestRestrictionEnabled(true)
defer setTestRestrictionEnabled(orig)

called, _ := runGuardTest("/v1/chat/completions", constant.ChannelTypeMidjourney)
assert.True(t, called, "midjourney channel should bypass restriction (not LLM)")
}

func TestEndpointFormatGuard_Enabled_KlingPassThrough(t *testing.T) {
orig := setTestRestrictionEnabled(true)
defer setTestRestrictionEnabled(orig)

called, _ := runGuardTest("/v1/chat/completions", constant.ChannelTypeKling)
assert.True(t, called, "kling channel should bypass restriction (not LLM)")
}

func TestEndpointFormatGuard_Enabled_RealtimeWebSocket(t *testing.T) {
orig := setTestRestrictionEnabled(true)
defer setTestRestrictionEnabled(orig)

called, _ := runGuardTestWithMethod("GET", "/v1/realtime", constant.ChannelTypeOpenAI)
assert.True(t, called, "openai channel should be allowed on realtime websocket")
}

func TestEndpointFormatGuard_Enabled_RealtimeWebSocket_Blocked(t *testing.T) {
orig := setTestRestrictionEnabled(true)
defer setTestRestrictionEnabled(orig)

called, w := runGuardTestWithMethod("GET", "/v1/realtime", constant.ChannelTypeAnthropic)
assert.False(t, called, "anthropic channel should be blocked on realtime websocket")
assert.Equal(t, 400, w.Code)
}

func TestEndpointFormatGuard_Enabled_ResponsesEndpoint(t *testing.T) {
orig := setTestRestrictionEnabled(true)
defer setTestRestrictionEnabled(orig)

called, _ := runGuardTest("/v1/responses", constant.ChannelTypeOpenAI)
assert.True(t, called, "openai channel should be allowed on /v1/responses")

called, w := runGuardTest("/v1/responses", constant.ChannelTypeAnthropic)
assert.False(t, called, "anthropic channel should be blocked on /v1/responses")
assert.Equal(t, 400, w.Code)
}

func TestEndpointFormatGuard_Enabled_AWS_ClaudeAllowed(t *testing.T) {
orig := setTestRestrictionEnabled(true)
defer setTestRestrictionEnabled(orig)

called, _ := runGuardTest("/v1/messages", constant.ChannelTypeAws)
assert.True(t, called, "AWS channel should be allowed for claude endpoint")
}

func TestEndpointFormatGuard_Enabled_AWS_OpenAIBlocked(t *testing.T) {
orig := setTestRestrictionEnabled(true)
defer setTestRestrictionEnabled(orig)

called, w := runGuardTest("/v1/chat/completions", constant.ChannelTypeAws)
assert.False(t, called, "AWS channel should be blocked for openai endpoint")
assert.Equal(t, 400, w.Code)
}

func TestEndpointFormatGuard_Enabled_VertexAI_AllEndpoints(t *testing.T) {
orig := setTestRestrictionEnabled(true)
defer setTestRestrictionEnabled(orig)

called, _ := runGuardTest("/v1/chat/completions", constant.ChannelTypeVertexAi)
assert.True(t, called, "vertexai should be allowed on openai endpoint")

called, _ = runGuardTest("/v1/messages", constant.ChannelTypeVertexAi)
assert.True(t, called, "vertexai should be allowed on claude endpoint")

called, _ = runGuardTest("/v1beta/models/gemini-pro:generateContent", constant.ChannelTypeVertexAi)
assert.True(t, called, "vertexai should be allowed on gemini endpoint")
}

func TestEndpointFormatGuard_Enabled_Ali_GeminiBlocked(t *testing.T) {
orig := setTestRestrictionEnabled(true)
defer setTestRestrictionEnabled(orig)

called, _ := runGuardTest("/v1/chat/completions", constant.ChannelTypeAli)
assert.True(t, called, "ali should be allowed on openai endpoint")

called, _ = runGuardTest("/v1/messages", constant.ChannelTypeAli)
assert.True(t, called, "ali should be allowed on claude endpoint")

called, w := runGuardTest("/v1beta/models/gemini-pro:generateContent", constant.ChannelTypeAli)
assert.False(t, called, "ali should be blocked on gemini endpoint")
assert.Equal(t, 400, w.Code)
}

func TestEndpointFormatGuard_Enabled_NoChannelType(t *testing.T) {
orig := setTestRestrictionEnabled(true)
defer setTestRestrictionEnabled(orig)

called, _ := runGuardTest("/v1/chat/completions", -1)
assert.True(t, called, "handler should be called when no channelType in context")
}

func TestEndpointFormatGuard_Enabled_ChannelTypeZero(t *testing.T) {
orig := setTestRestrictionEnabled(true)
defer setTestRestrictionEnabled(orig)

called, _ := runGuardTest("/v1/chat/completions", 0)
assert.True(t, called, "handler should be called when channelType is 0")
}

func TestEndpointFormatGuard_NonV1Path(t *testing.T) {
orig := setTestRestrictionEnabled(true)
defer setTestRestrictionEnabled(orig)

// /api/status 不在 guard 路由组中,无法通过 runGuardTest 测试
// 改用直接调用 getFamilyFromPath 验证返回空
family := getFamilyFromPath("/api/status")
assert.Equal(t, "", family, "non-v1 paths should have empty family")
}

// ============================================================================
// Group D: 错误消息验证
// ============================================================================

func TestEndpointFormatGuard_ErrorMessage_ClaudeChannel(t *testing.T) {
orig := setTestRestrictionEnabled(true)
defer setTestRestrictionEnabled(orig)

_, w := runGuardTest("/v1/chat/completions", constant.ChannelTypeAnthropic)
require.Equal(t, 400, w.Code)

var body map[string]interface{}
require.NoError(t, json.Unmarshal(w.Body.Bytes(), &body))

errObj := body["error"].(map[string]interface{})
assert.Equal(t, "endpoint_format_restricted", errObj["code"])
assert.Equal(t, "new_api_error", errObj["type"])

msg := errObj["message"].(string)
assert.Contains(t, msg, "Anthropic")
assert.Contains(t, msg, "/v1/messages")
assert.Contains(t, msg, "openai endpoints")
}

func TestEndpointFormatGuard_ErrorMessage_OpenAIChannel(t *testing.T) {
orig := setTestRestrictionEnabled(true)
defer setTestRestrictionEnabled(orig)

_, w := runGuardTest("/v1/messages", constant.ChannelTypeOpenAI)
require.Equal(t, 400, w.Code)

var body map[string]interface{}
require.NoError(t, json.Unmarshal(w.Body.Bytes(), &body))

errObj := body["error"].(map[string]interface{})
assert.Equal(t, "endpoint_format_restricted", errObj["code"])

msg := errObj["message"].(string)
assert.Contains(t, msg, "OpenAI")
assert.Contains(t, msg, "/v1/chat/completions")
assert.Contains(t, msg, "claude endpoints")
}

func TestEndpointFormatGuard_ErrorMessage_GeminiChannel(t *testing.T) {
orig := setTestRestrictionEnabled(true)
defer setTestRestrictionEnabled(orig)

_, w := runGuardTest("/v1/chat/completions", constant.ChannelTypeGemini)
require.Equal(t, 400, w.Code)

var body map[string]interface{}
require.NoError(t, json.Unmarshal(w.Body.Bytes(), &body))

errObj := body["error"].(map[string]interface{})
assert.Equal(t, "endpoint_format_restricted", errObj["code"])

msg := errObj["message"].(string)
assert.Contains(t, msg, "Gemini")
assert.Contains(t, msg, "/v1beta/models/{model}:generateContent")
assert.Contains(t, msg, "openai endpoints")
}

func TestEndpointFormatGuard_ErrorMessage_DeepSeekChannel(t *testing.T) {
orig := setTestRestrictionEnabled(true)
defer setTestRestrictionEnabled(orig)

_, w := runGuardTest("/v1/messages", constant.ChannelTypeDeepSeek)
require.Equal(t, 400, w.Code)

var body map[string]interface{}
require.NoError(t, json.Unmarshal(w.Body.Bytes(), &body))

errObj := body["error"].(map[string]interface{})
assert.Equal(t, "endpoint_format_restricted", errObj["code"])

msg := errObj["message"].(string)
assert.Contains(t, msg, "DeepSeek")
assert.Contains(t, msg, "/v1/chat/completions")
}

// ============================================================================
// Group E: 并发测试
// ============================================================================

func TestEndpointFormatGuard_ConcurrentRequests(t *testing.T) {
orig := setTestRestrictionEnabled(true)
defer setTestRestrictionEnabled(orig)

var allowedCount, blockedCount int64
var wg sync.WaitGroup

for i := 0; i < 20; i++ {
wg.Add(1)
go func(idx int) {
defer wg.Done()
var ch int
if idx%2 == 0 {
ch = constant.ChannelTypeOpenAI
} else {
ch = constant.ChannelTypeAnthropic
}
called, w := runGuardTest("/v1/chat/completions", ch)
if idx%2 == 0 {
if called && w.Code == 200 {
atomic.AddInt64(&allowedCount, 1)
}
} else {
if !called && w.Code == 400 {
atomic.AddInt64(&blockedCount, 1)
}
}
}(i)
}
wg.Wait()

assert.Equal(t, int64(10), atomic.LoadInt64(&allowedCount), "10 个 OpenAI 请求应放行")
assert.Equal(t, int64(10), atomic.LoadInt64(&blockedCount), "10 个 Anthropic 请求应被拦截")
}

// ============================================================================
// Group F: 功能开关测试
// ============================================================================

func TestEndpointFormatGuard_RuntimeToggle(t *testing.T) {
orig := setTestRestrictionEnabled(false)
defer setTestRestrictionEnabled(orig)

// 阶段 1: 关闭 → 应放行
called, _ := runGuardTest("/v1/chat/completions", constant.ChannelTypeAnthropic)
assert.True(t, called, "should pass when restriction is OFF")

// 阶段 2: 开启 → 应拦截
setTestRestrictionEnabled(true)
called, w := runGuardTest("/v1/chat/completions", constant.ChannelTypeAnthropic)
assert.False(t, called, "should be blocked when restriction is ON")
assert.Equal(t, 400, w.Code)

// 阶段 3: 再次关闭 → 应放行
setTestRestrictionEnabled(false)
called, _ = runGuardTest("/v1/chat/completions", constant.ChannelTypeAnthropic)
assert.True(t, called, "should pass again when restriction is turned OFF")
}

// ============================================================================
// 辅助测试
// ============================================================================

func TestEndpointFormatGuard_ContextKeyInt(t *testing.T) {
gin.SetMode(gin.TestMode)
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
c.Request = httptest.NewRequest("POST", "/v1/chat/completions", nil)
c.Set(string(constant.ContextKeyChannelType), 14)

val := common.GetContextKeyInt(c, constant.ContextKeyChannelType)
assert.Equal(t, 14, val)
}

+ 4
- 4
model/option.go Прегледај датотеку

@@ -175,7 +175,7 @@ func InitOptionMap() {
}

common.OptionMapRWMutex.Unlock()
loadOptionsFromDatabase()
LoadOptionsFromDatabase()
}

// triggerWechatPayReset 安全触发微信支付客户端重置
@@ -192,7 +192,7 @@ func triggerAlipayReset() {
}
}

func loadOptionsFromDatabase() {
func LoadOptionsFromDatabase() {
options, _ := AllOption()
for _, option := range options {
err := updateOptionMap(option.Key, option.Value)
@@ -206,7 +206,7 @@ func SyncOptions(frequency int) {
for {
time.Sleep(time.Duration(frequency) * time.Second)
common.SysLog("syncing options from database")
loadOptionsFromDatabase()
LoadOptionsFromDatabase()
}
}

@@ -567,7 +567,7 @@ func handleConfigUpdate(key, value string) bool {
configMap := map[string]string{
configKey: value,
}
config.UpdateConfigFromMap(cfg, configMap)
_ = config.UpdateConfigFromMap(cfg, configMap)

// 特定配置的后处理
if configName == "performance_setting" {


+ 28
- 54
router/relay-router.go Прегледај датотеку

@@ -10,6 +10,13 @@ import (
"github.com/gin-gonic/gin"
)

// relayHandler returns a handler that calls controller.Relay with the given format.
func relayHandler(format types.RelayFormat) gin.HandlerFunc {
return func(c *gin.Context) {
controller.Relay(c, format)
}
}

func SetRelayRouter(router *gin.Engine) {
router.Use(middleware.CORS())
router.Use(middleware.DecompressRequestMiddleware())
@@ -71,81 +78,49 @@ func SetRelayRouter(router *gin.Engine) {
// WebSocket 路由(统一到 Relay)
wsRouter := relayV1Router.Group("")
wsRouter.Use(middleware.Distribute())
wsRouter.GET("/realtime", func(c *gin.Context) {
controller.Relay(c, types.RelayFormatOpenAIRealtime)
})
wsRouter.Use(middleware.EndpointFormatGuard())
wsRouter.GET("/realtime", relayHandler(types.RelayFormatOpenAIRealtime))
}
{
//http router
httpRouter := relayV1Router.Group("")
httpRouter.Use(middleware.Distribute())
httpRouter.Use(middleware.EndpointFormatGuard())
httpRouter.Use(middleware.UserModelRateLimit())

// claude related routes
httpRouter.POST("/messages", func(c *gin.Context) {
controller.Relay(c, types.RelayFormatClaude)
})
httpRouter.POST("/messages", relayHandler(types.RelayFormatClaude))

// chat related routes
httpRouter.POST("/completions", func(c *gin.Context) {
controller.Relay(c, types.RelayFormatOpenAI)
})
httpRouter.POST("/chat/completions", func(c *gin.Context) {
controller.Relay(c, types.RelayFormatOpenAI)
})
httpRouter.POST("/completions", relayHandler(types.RelayFormatOpenAI))
httpRouter.POST("/chat/completions", relayHandler(types.RelayFormatOpenAI))

// response related routes
httpRouter.POST("/responses", func(c *gin.Context) {
controller.Relay(c, types.RelayFormatOpenAIResponses)
})
httpRouter.POST("/responses/compact", func(c *gin.Context) {
controller.Relay(c, types.RelayFormatOpenAIResponsesCompaction)
})
httpRouter.POST("/responses", relayHandler(types.RelayFormatOpenAIResponses))
httpRouter.POST("/responses/compact", relayHandler(types.RelayFormatOpenAIResponsesCompaction))

// image related routes
httpRouter.POST("/edits", func(c *gin.Context) {
controller.Relay(c, types.RelayFormatOpenAIImage)
})
httpRouter.POST("/images/generations", func(c *gin.Context) {
controller.Relay(c, types.RelayFormatOpenAIImage)
})
httpRouter.POST("/images/edits", func(c *gin.Context) {
controller.Relay(c, types.RelayFormatOpenAIImage)
})
httpRouter.POST("/edits", relayHandler(types.RelayFormatOpenAIImage))
httpRouter.POST("/images/generations", relayHandler(types.RelayFormatOpenAIImage))
httpRouter.POST("/images/edits", relayHandler(types.RelayFormatOpenAIImage))

// embedding related routes
httpRouter.POST("/embeddings", func(c *gin.Context) {
controller.Relay(c, types.RelayFormatEmbedding)
})
httpRouter.POST("/embeddings", relayHandler(types.RelayFormatEmbedding))

// audio related routes
httpRouter.POST("/audio/transcriptions", func(c *gin.Context) {
controller.Relay(c, types.RelayFormatOpenAIAudio)
})
httpRouter.POST("/audio/translations", func(c *gin.Context) {
controller.Relay(c, types.RelayFormatOpenAIAudio)
})
httpRouter.POST("/audio/speech", func(c *gin.Context) {
controller.Relay(c, types.RelayFormatOpenAIAudio)
})
httpRouter.POST("/audio/transcriptions", relayHandler(types.RelayFormatOpenAIAudio))
httpRouter.POST("/audio/translations", relayHandler(types.RelayFormatOpenAIAudio))
httpRouter.POST("/audio/speech", relayHandler(types.RelayFormatOpenAIAudio))

// rerank related routes
httpRouter.POST("/rerank", func(c *gin.Context) {
controller.Relay(c, types.RelayFormatRerank)
})
httpRouter.POST("/rerank", relayHandler(types.RelayFormatRerank))

// gemini relay routes
httpRouter.POST("/engines/:model/embeddings", func(c *gin.Context) {
controller.Relay(c, types.RelayFormatGemini)
})
httpRouter.POST("/models/*path", func(c *gin.Context) {
controller.Relay(c, types.RelayFormatGemini)
})
httpRouter.POST("/engines/:model/embeddings", relayHandler(types.RelayFormatGemini))
httpRouter.POST("/models/*path", relayHandler(types.RelayFormatGemini))

// other relay routes
httpRouter.POST("/moderations", func(c *gin.Context) {
controller.Relay(c, types.RelayFormatOpenAI)
})
httpRouter.POST("/moderations", relayHandler(types.RelayFormatOpenAI))

// not implemented
httpRouter.POST("/images/variations", controller.RelayNotImplemented)
@@ -186,12 +161,11 @@ func SetRelayRouter(router *gin.Engine) {
relayGeminiRouter.Use(middleware.ModelRequestRateLimit())
relayGeminiRouter.Use(middleware.RelayCaptureMiddleware())
relayGeminiRouter.Use(middleware.Distribute())
relayGeminiRouter.Use(middleware.EndpointFormatGuard())
relayGeminiRouter.Use(middleware.UserModelRateLimit())
{
// Gemini API 路径格式: /v1beta/models/{model_name}:{action}
relayGeminiRouter.POST("/models/*path", func(c *gin.Context) {
controller.Relay(c, types.RelayFormatGemini)
})
relayGeminiRouter.POST("/models/*path", relayHandler(types.RelayFormatGemini))
}
}



+ 7
- 1
setting/model_setting/global.go Прегледај датотеку

@@ -35,7 +35,8 @@ func (p ChatCompletionsToResponsesPolicy) IsChannelEnabled(channelID int, channe
type GlobalSettings struct {
PassThroughRequestEnabled bool `json:"pass_through_request_enabled"`
ThinkingModelBlacklist []string `json:"thinking_model_blacklist"`
ChatCompletionsToResponsesPolicy ChatCompletionsToResponsesPolicy `json:"chat_completions_to_responses_policy"`
ChatCompletionsToResponsesPolicy ChatCompletionsToResponsesPolicy `json:"chat_completions_to_responses_policy"`
EndpointFormatRestrictionEnabled bool `json:"endpoint_format_restriction_enabled"`
}

// 默认配置
@@ -49,6 +50,7 @@ var defaultOpenaiSettings = GlobalSettings{
Enabled: false,
AllChannels: true,
},
EndpointFormatRestrictionEnabled: false,
}

// 全局实例
@@ -63,6 +65,10 @@ func GetGlobalSettings() *GlobalSettings {
return &globalSettings
}

func IsEndpointFormatRestrictionEnabled() bool {
return globalSettings.EndpointFormatRestrictionEnabled
}

// ShouldPreserveThinkingSuffix 判断模型是否配置为保留 thinking/-nothinking/-low/-high/-medium 后缀
func ShouldPreserveThinkingSuffix(modelName string) bool {
target := strings.TrimSpace(modelName)


+ 245
- 0
test/e2e_endpoint_format.go Прегледај датотеку

@@ -0,0 +1,245 @@
//go:build ignore

package main

import (
"crypto/tls"
"encoding/json"
"fmt"
"io"
"net/http"
"os"
"strings"
"time"
)

var (
baseURL = "http://localhost:3001"
apiKey string
client = &http.Client{Timeout: 30 * time.Second}
)

func init() {
// Skip TLS verification if needed
client.Transport = &http.Transport{
TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
}
}

func main() {
// Use test token created earlier
apiKey = "sk-3bJC6cCSSFuTqOxaZQ8lL6kOfYAl3pt0uuwWMiU40F6KkxOS"
if len(os.Args) > 1 {
apiKey = os.Args[1]
}
if len(os.Args) > 2 {
baseURL = os.Args[2]
}

pass, fail := 0, 0

fmt.Println("==========================================")
fmt.Println("E2E: Endpoint Format Restriction")
fmt.Println("==========================================")

// Test 1: Feature OFF - Claude via /v1/chat/completions should pass
fmt.Println("\n--- Test 1: Feature OFF, Claude via /v1/chat/completions ---")
code, body := doRequest("POST", "/v1/chat/completions", `{"model":"anthropic/claude-sonnet-4-20250514","messages":[{"role":"user","content":"hi"}],"max_tokens":5}`)
fmt.Printf(" HTTP: %d\n", code)
if code != 400 || !strings.Contains(body, "endpoint_format_restricted") {
fmt.Println(" PASS (not blocked)")
pass++
} else {
fmt.Println(" FAIL: should not be blocked when feature is off")
fail++
}

// Enable feature via admin API
fmt.Println("\n--- Enabling feature via admin login ---")
session := adminLogin()
if session == "" {
fmt.Println(" WARNING: Could not login as admin, trying direct DB approach")
} else {
enableFeature(session)
}

// Test 2: Feature ON - Claude via /v1/chat/completions should be BLOCKED
fmt.Println("\n--- Test 2: Feature ON, Claude via /v1/chat/completions (should BLOCK) ---")
code, body = doRequest("POST", "/v1/chat/completions", `{"model":"anthropic/claude-sonnet-4-20250514","messages":[{"role":"user","content":"hi"}],"max_tokens":5}`)
fmt.Printf(" HTTP: %d\n", code)
fmt.Printf(" Body: %s\n", truncate(body, 300))
if code == 400 && strings.Contains(body, "endpoint_format_restricted") {
fmt.Println(" PASS (correctly blocked)")
pass++
} else {
fmt.Println(" FAIL: expected 400 with endpoint_format_restricted")
fail++
}

// Test 3: Feature ON - Claude via /v1/messages should PASS
fmt.Println("\n--- Test 3: Feature ON, Claude via /v1/messages (should PASS) ---")
code, body = doRequestWithHeaders("POST", "/v1/messages",
`{"model":"anthropic/claude-sonnet-4-20250514","messages":[{"role":"user","content":"hi"}],"max_tokens":5}`,
map[string]string{"anthropic-version": "2023-06-01"})
fmt.Printf(" HTTP: %d\n", code)
if code != 400 || !strings.Contains(body, "endpoint_format_restricted") {
fmt.Println(" PASS (correctly allowed)")
pass++
} else {
fmt.Println(" FAIL: should not be blocked for matching format")
fail++
}

// Test 4: Feature ON - DeepSeek via /v1/chat/completions should PASS
fmt.Println("\n--- Test 4: Feature ON, DeepSeek via /v1/chat/completions (should PASS) ---")
code, body = doRequest("POST", "/v1/chat/completions", `{"model":"deepseek-ai/DeepSeek-V3-0324","messages":[{"role":"user","content":"hi"}],"max_tokens":5}`)
fmt.Printf(" HTTP: %d\n", code)
if code != 400 || !strings.Contains(body, "endpoint_format_restricted") {
fmt.Println(" PASS (correctly allowed)")
pass++
} else {
fmt.Println(" FAIL: should not be blocked for openai channel")
fail++
}

// Test 5: Check error message content
fmt.Println("\n--- Test 5: Error message contains native endpoint recommendation ---")
if strings.Contains(body, "/v1/messages") || strings.Contains(body, "Anthropic") || strings.Contains(body, "endpoint_format_restricted") {
// Re-test to get the error body
code2, body2 := doRequest("POST", "/v1/chat/completions", `{"model":"anthropic/claude-sonnet-4-20250514","messages":[{"role":"user","content":"hi"}],"max_tokens":5}`)
_ = code2
if strings.Contains(body2, "/v1/messages") && strings.Contains(body2, "endpoint_format_restricted") {
fmt.Println(" PASS (error recommends /v1/messages)")
pass++
} else {
fmt.Printf(" Body: %s\n", truncate(body2, 300))
fmt.Println(" FAIL: error message missing native endpoint")
fail++
}
} else {
fmt.Println(" SKIP (previous test didn't produce blocking error)")
}

// Disable feature
if session != "" {
fmt.Println("\n--- Disabling feature ---")
disableFeature(session)
}

fmt.Printf("\n==========================================\n")
fmt.Printf("Results: %d passed, %d failed\n", pass, fail)
fmt.Printf("==========================================\n")

if fail > 0 {
os.Exit(1)
}
}

func doRequest(method, path, body string) (int, string) {
return doRequestWithHeaders(method, path, body, nil)
}

func doRequestWithHeaders(method, path, body string, headers map[string]string) (int, string) {
var req *http.Request
var err error
if body != "" {
req, err = http.NewRequest(method, baseURL+path, strings.NewReader(body))
} else {
req, err = http.NewRequest(method, baseURL+path, nil)
}
if err != nil {
return 0, err.Error()
}
req.Header.Set("Authorization", "Bearer "+apiKey)
req.Header.Set("Content-Type", "application/json")
for k, v := range headers {
req.Header.Set(k, v)
}

resp, err := client.Do(req)
if err != nil {
return 0, err.Error()
}
defer resp.Body.Close()
respBody, _ := io.ReadAll(resp.Body)
return resp.StatusCode, string(respBody)
}

func adminLogin() string {
// Try common passwords
passwords := []string{"123456", "password", "admin", "root", "root123", "admin123"}
for _, pw := range passwords {
code, body := doRawRequest("POST", baseURL+"/api/user/login",
fmt.Sprintf(`{"username":"root","password":"%s"}`, pw))
if code == 200 {
var result struct {
Success bool `json:"success"`
Data string `json:"data"`
}
json.Unmarshal([]byte(body), &result)
if result.Success {
fmt.Printf(" Logged in with password: %s\n", pw)
return result.Data
}
}
}
fmt.Println(" Could not login with any common password")
return ""
}

func doRawRequest(method, url, body string) (int, string) {
req, err := http.NewRequest(method, url, strings.NewReader(body))
if err != nil {
return 0, err.Error()
}
req.Header.Set("Content-Type", "application/json")
resp, err := client.Do(req)
if err != nil {
return 0, err.Error()
}
defer resp.Body.Close()
respBody, _ := io.ReadAll(resp.Body)
return resp.StatusCode, string(respBody)
}

func enableFeature(session string) {
code, body := doRawRequest("PUT", baseURL+"/api/option/",
`{"key":"global.endpoint_format_restriction_enabled","value":"true"}`)
req, _ := http.NewRequest("PUT", baseURL+"/api/option/", strings.NewReader(
`{"key":"global.endpoint_format_restriction_enabled","value":"true"}`))
req.Header.Set("Content-Type", "application/json")
req.Header.Set("Authorization", "Bearer "+session)
resp, err := client.Do(req)
if err != nil {
fmt.Printf(" Error: %v\n", err)
return
}
defer resp.Body.Close()
respBody, _ := io.ReadAll(resp.Body)
fmt.Printf(" Enable response: %d %s\n", resp.StatusCode, truncate(string(respBody), 100))
_ = code
_ = body
time.Sleep(1 * time.Second)
}

func disableFeature(session string) {
req, _ := http.NewRequest("PUT", baseURL+"/api/option/", strings.NewReader(
`{"key":"global.endpoint_format_restriction_enabled","value":"false"}`))
req.Header.Set("Content-Type", "application/json")
req.Header.Set("Authorization", "Bearer "+session)
resp, err := client.Do(req)
if err != nil {
fmt.Printf(" Error: %v\n", err)
return
}
defer resp.Body.Close()
respBody, _ := io.ReadAll(resp.Body)
fmt.Printf(" Disable response: %d %s\n", resp.StatusCode, truncate(string(respBody), 100))
}

func truncate(s string, n int) string {
if len(s) <= n {
return s
}
return s[:n] + "..."
}

+ 420
- 0
test/e2e_endpoint_format.py Прегледај датотеку

@@ -0,0 +1,420 @@
"""E2E test for Endpoint Format Restriction feature.

完整测试 6 个阶段:
0. 环境发现(自动查找渠道和模型)
1. 功能关闭 - 所有跨格式请求不被拦截
2. 功能开启 - 匹配的渠道+端点放行
3. 功能开启 - 不匹配的渠道+端点被拦截
4. 错误消息验证
5. 清理
"""

import argparse
import json
import sys
import time

import mysql.connector
import requests

DEFAULT_BASE_URL = "http://localhost:3000"
DEFAULT_DB_HOST = "localhost"
DEFAULT_DB_PORT = 33306
DEFAULT_DB_USER = "root"
DEFAULT_DB_PASS = "Lanqi123456"
DEFAULT_DB_NAME = "new-api"

OPTION_KEY = "global.endpoint_format_restriction_enabled"

# 渠道类型常量 (与 Go constant 对应)
CH_OPENAI = 1
CH_MIDJOURNEY = 2
CH_AZURE = 3
CH_OLLAMA = 4
CH_ANTHROPIC = 14
CH_GEMINI = 24
CH_DEEPSEEK = 43
CH_VERTEX_AI = 41
CH_ALI = 17
CH_AWS = 33
CH_KLING = 50
CH_SUNO = 36

# 渠道类型 → 默认格式家族
CHANNEL_FAMILIES = {
CH_OPENAI: ["openai"],
CH_AZURE: ["openai"],
CH_DEEPSEEK: ["openai"],
CH_ANTHROPIC: ["claude"],
CH_AWS: ["claude"],
CH_GEMINI: ["gemini"],
CH_VERTEX_AI: ["openai", "claude", "gemini"],
CH_ALI: ["openai", "claude"],
}

# 非LLM渠道(不在 channelNativeFamilies 中)
NON_LLM_CHANNELS = [CH_MIDJOURNEY, CH_KLING, CH_SUNO]


def get_db(args):
return mysql.connector.connect(
host=args.db_host, port=args.db_port,
user=args.db_user, password=args.db_pass, database=args.db_name,
)


def get_token_key(args) -> str:
db = get_db(args)
cur = db.cursor()
cur.execute("SELECT `key` FROM tokens WHERE id = %s", (args.token_id,))
row = cur.fetchone()
cur.close()
db.close()
if not row:
raise ValueError(f"Token id={args.token_id} not found")
return row[0]


def set_restriction(args, enabled: bool):
db = get_db(args)
cur = db.cursor()
val = "true" if enabled else "false"
cur.execute(
"INSERT INTO options (`key`, value) VALUES (%s, %s) ON DUPLICATE KEY UPDATE value = %s",
(OPTION_KEY, val, val),
)
db.commit()
cur.close()
db.close()


def discover_channels(args) -> dict:
"""从 DB 查询可用渠道,按 type 分组。返回 {type: [(id, model_name), ...]}"""
db = get_db(args)
cur = db.cursor()
cur.execute("SELECT id, type, models FROM channels WHERE status = 1")
rows = cur.fetchall()
cur.close()
db.close()

result = {}
for ch_id, ch_type, models_str in rows:
if not models_str:
continue
models = [m.strip() for m in models_str.split(",") if m.strip()]
if models:
result.setdefault(ch_type, []).append((ch_id, models[0]))
return result


def find_channel(channels: dict, ch_type: int) -> tuple:
"""查找指定类型的渠道,返回 (channel_id, model_name) 或 None"""
entries = channels.get(ch_type)
if entries:
return entries[0]
return None


# ============================================================================
# Request helpers
# ============================================================================

def chat_req(key: str, model: str) -> requests.Response:
return requests.post(
f"{BASE_URL}/v1/chat/completions",
headers={"Authorization": f"Bearer sk-{key}", "Content-Type": "application/json"},
json={"model": model, "messages": [{"role": "user", "content": "hi"}], "max_tokens": 5},
timeout=30,
)


def messages_req(key: str, model: str) -> requests.Response:
return requests.post(
f"{BASE_URL}/v1/messages",
headers={"Authorization": f"Bearer sk-{key}", "Content-Type": "application/json"},
json={"model": model, "messages": [{"role": "user", "content": "hi"}], "max_tokens": 5},
timeout=30,
)


def gemini_req(key: str, model: str) -> requests.Response:
return requests.post(
f"{BASE_URL}/v1beta/models/{model}:generateContent",
headers={"Authorization": f"Bearer sk-{key}", "Content-Type": "application/json"},
json={"contents": [{"parts": [{"text": "hi"}]}]},
timeout=30,
)


def is_blocked(resp: requests.Response) -> bool:
if resp.status_code != 400:
return False
try:
return resp.json().get("error", {}).get("code") == "endpoint_format_restricted"
except Exception:
return False


def wait_for_restriction(args, token_key: str, model: str, enabled: bool) -> bool:
"""轮询等待限制设置生效(最长 70s,每 3s 检查一次)"""
for i in range(24):
r = chat_req(token_key, model)
blocked = is_blocked(r)
if enabled and blocked:
print(f" Restriction ACTIVE (after ~{i*3}s)")
return True
if not enabled and not blocked:
print(f" Restriction INACTIVE (after ~{i*3}s)")
return True
time.sleep(3)
print(f" WARNING: Restriction state not confirmed after 70s")
return False


# ============================================================================
# Test runners
# ============================================================================

BASE_URL = DEFAULT_BASE_URL
passed = 0
failed = 0


def run_test(name: str, fn) -> bool:
global passed, failed
try:
fn()
print(f" [PASS] {name}")
passed += 1
return True
except AssertionError as e:
print(f" [FAIL] {name}: {e}")
failed += 1
return False


def phase1_off(token_key, openai_model, claude_model, gemini_info):
"""阶段 1: 功能关闭 - 所有跨格式请求不被拦截"""
print("\n--- 阶段 1: 功能关闭 ---")

def test_cross_openai_claude():
r = chat_req(token_key, claude_model)
assert not is_blocked(r), f"Claude模型 + OpenAI端点 不应被拦截, got {r.status_code}: {r.text[:200]}"

def test_cross_claude_openai():
r = messages_req(token_key, openai_model)
assert not is_blocked(r), f"OpenAI模型 + Claude端点 不应被拦截, got {r.status_code}: {r.text[:200]}"

def test_cross_gemini_openai():
if gemini_info:
r = chat_req(token_key, gemini_info[1])
assert not is_blocked(r), f"Gemini模型 + OpenAI端点 不应被拦截"

run_test("Claude模型 + OpenAI端点 不拦截", test_cross_openai_claude)
run_test("OpenAI模型 + Claude端点 不拦截", test_cross_claude_openai)
run_test("Gemini模型 + OpenAI端点 不拦截", test_cross_gemini_openai)


def phase2_allowed(token_key, openai_model, claude_model, gemini_info, deepseek_info):
"""阶段 2: 功能开启 - 匹配的渠道+端点放行"""
print("\n--- 阶段 2: 功能开启 - 匹配放行 ---")

def test_openai_openai():
r = chat_req(token_key, openai_model)
assert not is_blocked(r), f"OpenAI渠道 + OpenAI端点 应放行, got {r.status_code}"

def test_claude_claude():
r = messages_req(token_key, claude_model)
assert not is_blocked(r), f"Claude渠道 + Claude端点 应放行, got {r.status_code}"

def test_gemini_gemini():
if gemini_info:
r = gemini_req(token_key, gemini_info[1])
assert not is_blocked(r), f"Gemini渠道 + Gemini端点 应放行"

def test_deepseek_openai():
if deepseek_info:
r = chat_req(token_key, deepseek_info[1])
assert not is_blocked(r), f"DeepSeek渠道 + OpenAI端点 应放行"

run_test("OpenAI渠道 + /v1/chat/completions 放行", test_openai_openai)
run_test("Claude渠道 + /v1/messages 放行", test_claude_claude)
run_test("Gemini渠道 + /v1beta 放行", test_gemini_gemini)
run_test("DeepSeek渠道 + /v1/chat/completions 放行", test_deepseek_openai)


def phase3_blocked(token_key, openai_model, claude_model, gemini_info, deepseek_info):
"""阶段 3: 功能开启 - 不匹配的渠道+端点被拦截"""
print("\n--- 阶段 3: 功能开启 - 不匹配拦截 ---")

def test_claude_openai_blocked():
r = chat_req(token_key, claude_model)
assert is_blocked(r), f"Claude渠道 + OpenAI端点 应被拦截, got {r.status_code}: {r.text[:200]}"

def test_openai_claude_blocked():
r = messages_req(token_key, openai_model)
assert is_blocked(r), f"OpenAI渠道 + Claude端点 应被拦截, got {r.status_code}: {r.text[:200]}"

def test_gemini_openai_blocked():
if gemini_info:
r = chat_req(token_key, gemini_info[1])
assert is_blocked(r), f"Gemini渠道 + OpenAI端点 应被拦截"

def test_openai_gemini_blocked():
if gemini_info:
r = gemini_req(token_key, openai_model)
assert is_blocked(r), f"OpenAI渠道 + Gemini端点 应被拦截"

def test_deepseek_claude_blocked():
if deepseek_info:
r = messages_req(token_key, deepseek_info[1])
assert is_blocked(r), f"DeepSeek渠道 + Claude端点 应被拦截"

run_test("Claude渠道 + /v1/chat/completions 拦截", test_claude_openai_blocked)
run_test("OpenAI渠道 + /v1/messages 拦截", test_openai_claude_blocked)
run_test("Gemini渠道 + /v1/chat/completions 拦截", test_gemini_openai_blocked)
run_test("OpenAI渠道 + /v1beta 拦截", test_openai_gemini_blocked)
run_test("DeepSeek渠道 + /v1/messages 拦截", test_deepseek_claude_blocked)


def phase4_error_messages(token_key, openai_model, claude_model, gemini_info):
"""阶段 4: 错误消息验证"""
print("\n--- 阶段 4: 错误消息验证 ---")

def test_claude_error_msg():
r = chat_req(token_key, claude_model)
assert is_blocked(r)
body = r.json()
err = body["error"]
assert err["code"] == "endpoint_format_restricted"
assert err["type"] == "new_api_error"
assert "Anthropic" in err["message"]
assert "/v1/messages" in err["message"]

def test_openai_error_msg():
r = messages_req(token_key, openai_model)
assert is_blocked(r)
body = r.json()
err = body["error"]
assert err["code"] == "endpoint_format_restricted"
assert "OpenAI" in err["message"]
assert "/v1/chat/completions" in err["message"]

def test_gemini_error_msg():
if gemini_info:
r = chat_req(token_key, gemini_info[1])
assert is_blocked(r)
body = r.json()
err = body["error"]
assert err["code"] == "endpoint_format_restricted"
assert "Gemini" in err["message"]
assert "/v1beta/models/{model}:generateContent" in err["message"]

run_test("Claude渠道错误消息 (Anthropic + /v1/messages)", test_claude_error_msg)
run_test("OpenAI渠道错误消息 (OpenAI + /v1/chat/completions)", test_openai_error_msg)
run_test("Gemini渠道错误消息 (Gemini + /v1beta)", test_gemini_error_msg)


def main():
global BASE_URL, passed, failed

parser = argparse.ArgumentParser(description="E2E test: Endpoint Format Restriction")
parser.add_argument("--base-url", default=DEFAULT_BASE_URL)
parser.add_argument("--db-host", default=DEFAULT_DB_HOST)
parser.add_argument("--db-port", type=int, default=DEFAULT_DB_PORT)
parser.add_argument("--db-user", default=DEFAULT_DB_USER)
parser.add_argument("--db-pass", default=DEFAULT_DB_PASS)
parser.add_argument("--db-name", default=DEFAULT_DB_NAME)
parser.add_argument("--token-id", type=int, default=29)
args = parser.parse_args()

BASE_URL = args.base_url

print("=" * 60)
print("E2E Test: Endpoint Format Restriction")
print("=" * 60)

# ---- 阶段 0: 环境发现 ----
print("\n--- 阶段 0: 环境发现 ---")

# 健康检查
try:
r = requests.get(f"{BASE_URL}/api/status", timeout=5)
assert r.status_code == 200
print(f" Server: {BASE_URL} OK")
except Exception as e:
print(f" ERROR: Server not reachable: {e}")
sys.exit(1)

# 获取 token
try:
token_key = get_token_key(args)
print(f" Token: id={args.token_id} key=sk-{token_key[:8]}...")
except ValueError:
print(f" ERROR: Token id={args.token_id} not found")
sys.exit(1)

# 发现渠道
channels = discover_channels(args)
openai_ch = find_channel(channels, CH_OPENAI)
claude_ch = find_channel(channels, CH_ANTHROPIC)
gemini_ch = find_channel(channels, CH_GEMINI)
deepseek_ch = find_channel(channels, CH_DEEPSEEK)

if not openai_ch:
print(" ERROR: No OpenAI channel found")
sys.exit(1)
if not claude_ch:
print(" ERROR: No Anthropic channel found")
sys.exit(1)

openai_model = openai_ch[1]
claude_model = claude_ch[1]
gemini_info = gemini_ch # 可能为 None
deepseek_info = deepseek_ch # 可能为 None

print(f" OpenAI model: {openai_model} (ch={openai_ch[0]})")
print(f" Claude model: {claude_model} (ch={claude_ch[0]})")
if gemini_info:
print(f" Gemini model: {gemini_info[1]} (ch={gemini_info[0]})")
else:
print(f" Gemini model: (no channel, skip gemini tests)")
if deepseek_info:
print(f" DeepSeek model: {deepseek_info[1]} (ch={deepseek_info[0]})")

# ---- 阶段 1: 功能关闭 ----
set_restriction(args, False)
print("\n Set restriction=false, waiting for sync...")
wait_for_restriction(args, token_key, claude_model, enabled=False)
phase1_off(token_key, openai_model, claude_model, gemini_info)

# ---- 阶段 2: 功能开启 - 放行 ----
set_restriction(args, True)
print("\n Set restriction=true, waiting for sync...")
if not wait_for_restriction(args, token_key, claude_model, enabled=True):
print(" ERROR: Restriction did not activate. Aborting.")
sys.exit(1)
phase2_allowed(token_key, openai_model, claude_model, gemini_info, deepseek_info)

# ---- 阶段 3: 功能开启 - 拦截 ----
phase3_blocked(token_key, openai_model, claude_model, gemini_info, deepseek_info)

# ---- 阶段 4: 错误消息 ----
phase4_error_messages(token_key, openai_model, claude_model, gemini_info)

# ---- 阶段 5: 清理 ----
set_restriction(args, False)
print(f"\n--- 阶段 5: 清理 ---")
print(f" Restriction reset to: false")

# ---- 结果 ----
print("\n" + "=" * 60)
print(f"Results: {passed} passed, {failed} failed")
if failed > 0:
sys.exit(1)
else:
print("ALL TESTS PASSED!")


if __name__ == "__main__":
main()

+ 1
- 0
web/src/components/settings/ModelSetting.jsx Прегледај датотеку

@@ -40,6 +40,7 @@ const ModelSetting = () => {
'claude.default_max_tokens': '',
'claude.thinking_adapter_budget_tokens_percentage': 0.8,
'global.pass_through_request_enabled': false,
'global.endpoint_format_restriction_enabled': false,
'global.thinking_model_blacklist': '[]',
'global.chat_completions_to_responses_policy': '{}',
'general_setting.ping_interval_enabled': false,


+ 18
- 0
web/src/pages/Setting/Model/SettingGlobalModel.jsx Прегледај датотеку

@@ -68,6 +68,7 @@ const chatCompletionsToResponsesPolicyAllChannelsExample = JSON.stringify(

const defaultGlobalSettingInputs = {
'global.pass_through_request_enabled': false,
'global.endpoint_format_restriction_enabled': false,
'global.thinking_model_blacklist': '[]',
'global.chat_completions_to_responses_policy': '{}',
'general_setting.ping_interval_enabled': false,
@@ -205,6 +206,23 @@ export default function SettingGlobalModel(props) {
/>
</Col>
</Row>
<Row>
<Col xs={24} sm={12} md={8} lg={8} xl={8}>
<Form.Switch
label={t('端点格式限制')}
field={'global.endpoint_format_restriction_enabled'}
onChange={(value) =>
setInputs({
...inputs,
'global.endpoint_format_restriction_enabled': value,
})
}
extraText={t(
'开启后,渠道模型只能通过原生协议端点访问(如 Claude 模型只能走 /v1/messages),拒绝跨协议转换请求',
)}
/>
</Col>
</Row>
<Row>
<Col span={24}>
<Form.TextArea


Loading…
Откажи
Сачувај