当开启端点格式限制后,不同协议的渠道只能通过对应的端点访问: - 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
| @@ -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" | |||
| @@ -286,6 +286,7 @@ func InitResources() error { | |||
| // Initialize options, should after model.InitDB() | |||
| model.InitOptionMap() | |||
| model.LoadOptionsFromDatabase() | |||
| // 清理旧的磁盘缓存文件 | |||
| common.CleanupOldCacheFiles() | |||
| @@ -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 | |||
| } | |||
| @@ -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) | |||
| } | |||
| @@ -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" { | |||
| @@ -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)) | |||
| } | |||
| } | |||
| @@ -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) | |||
| @@ -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] + "..." | |||
| } | |||
| @@ -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() | |||
| @@ -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, | |||
| @@ -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 | |||