mirror of
https://github.com/Wei-Shaw/sub2api.git
synced 2026-10-14 12:02:27 +08:00
Merge pull request #7378 from wucm667/fix/issue-7368-composite-ws-routing
fix(gateway): resolve composite routes for websocket aliases
This commit is contained in:
@@ -295,7 +295,7 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) {
|
||||
legacyEngine := securityaudit.NewLegacyModerationAdapter(contentModerationService)
|
||||
coordinator := securityaudit.NewCoordinator(legacyEngine, promptService)
|
||||
gatewayHandler := handler.ProvideGatewayHandler(gatewayService, openAIGatewayService, geminiMessagesCompatService, antigravityGatewayService, userService, concurrencyService, billingCacheService, usageService, apiKeyService, usageRecordWorkerPool, errorPassthroughService, contentModerationService, userMessageQueueService, configConfig, settingService, coordinator)
|
||||
openAIGatewayHandler := handler.ProvideOpenAIGatewayHandler(openAIGatewayService, pluginManager, concurrencyService, billingCacheService, apiKeyService, usageRecordWorkerPool, errorPassthroughService, contentModerationService, opsService, grokQuotaService, configConfig, coordinator)
|
||||
openAIGatewayHandler := handler.ProvideOpenAIGatewayHandler(openAIGatewayService, pluginManager, concurrencyService, billingCacheService, apiKeyService, usageRecordWorkerPool, errorPassthroughService, contentModerationService, opsService, grokQuotaService, configConfig, coordinator, compositeRouteResolver)
|
||||
handlerSettingHandler := handler.ProvideSettingHandler(settingService, buildInfo, notificationEmailService)
|
||||
totpHandler := handler.NewTotpHandler(totpService)
|
||||
passkeyRepository := repository.NewPasskeyRepository(db)
|
||||
|
||||
@@ -34,6 +34,7 @@ import (
|
||||
|
||||
// OpenAIGatewayHandler handles OpenAI API gateway requests
|
||||
type OpenAIGatewayHandler struct {
|
||||
compositeResolver *service.CompositeRouteResolver
|
||||
gatewayService *service.OpenAIGatewayService
|
||||
billingCacheService *service.BillingCacheService
|
||||
apiKeyService *service.APIKeyService
|
||||
@@ -2391,7 +2392,7 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
// 分组级模型白名单:首帧校验客户端模型,不通过则关闭连接并标记运维原因。
|
||||
// 必须在 ensureCompositeTargetPlatform(合成路由改写)之前执行。
|
||||
// 必须在合成路由解析和上游模型映射之前执行。
|
||||
// 与 HTTP 准入一致:帧内重复 model 键/大小写变体可能被上游按末值绑定,
|
||||
// 全部候选值逐一校验,任一未命中即拒绝。
|
||||
if blocked := blockedModelAllowlistCandidate(apiKey.Group, requestmodel.FromBodyCandidates("", "application/json", firstMessage)); blocked != "" {
|
||||
@@ -2400,7 +2401,21 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) {
|
||||
closeOpenAIClientWS(wsConn, coderws.StatusPolicyViolation, fmt.Sprintf("Model %q is not available for this group", blocked))
|
||||
return
|
||||
}
|
||||
ensureCompositeTargetPlatform(c, apiKey, reqModel)
|
||||
// Keep the client model in the frame for admission, response identity and usage.
|
||||
// Apply the resolved upstream model through MapRequestModel on every turn.
|
||||
wsRouteModel := reqModel
|
||||
if apiKey.Group != nil && apiKey.Group.Platform == service.PlatformComposite {
|
||||
decision, resolveErr := h.compositeResolver.Resolve(c.Request.Context(), apiKey.Group.ID, reqModel, service.CompositeRouteEndpointResponses)
|
||||
if resolveErr != nil {
|
||||
reqLog.Error("openai.websocket_composite_route_failed", zap.Error(resolveErr))
|
||||
closeOpenAIClientWS(wsConn, coderws.StatusInternalError, "Failed to resolve composite model route")
|
||||
return
|
||||
}
|
||||
if decision.Matched {
|
||||
c.Request = c.Request.WithContext(service.WithCompositeRouteDecision(c.Request.Context(), decision))
|
||||
wsRouteModel = decision.UpstreamModel
|
||||
}
|
||||
}
|
||||
ctx = c.Request.Context()
|
||||
if apiKey.Group != nil && apiKey.Group.Platform == service.PlatformComposite {
|
||||
platform, ok := service.ResolvedTargetPlatformFromContext(ctx)
|
||||
@@ -2464,8 +2479,8 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) {
|
||||
}
|
||||
|
||||
// 解析渠道级模型映射
|
||||
channelMappingWS, _ := h.gatewayService.ResolveChannelMappingAndRestrict(ctx, apiKey.GroupID, reqModel)
|
||||
wsForwardModel := openAIChannelForwardModel(channelMappingWS, reqModel)
|
||||
channelMappingWS, _ := h.gatewayService.ResolveChannelMappingAndRestrict(ctx, apiKey.GroupID, wsRouteModel)
|
||||
wsForwardModel := openAIChannelForwardModel(channelMappingWS, wsRouteModel)
|
||||
|
||||
var currentUserRelease func()
|
||||
var currentAccountRelease func()
|
||||
@@ -2856,7 +2871,16 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) {
|
||||
model = reqModel
|
||||
}
|
||||
setOpsRequestContext(c, model, true)
|
||||
mapping, _ := h.gatewayService.ResolveChannelMappingAndRestrict(ctx, apiKey.GroupID, model)
|
||||
routeModel := model
|
||||
if apiKey.Group != nil && apiKey.Group.Platform == service.PlatformComposite {
|
||||
// The account, target platform and route context are connection-scoped.
|
||||
// A different public model needs a new connection and fresh resolution.
|
||||
if model != reqModel {
|
||||
return "", newOpenAIWSUnsupportedModelSwitchError(model)
|
||||
}
|
||||
routeModel = wsRouteModel
|
||||
}
|
||||
mapping, _ := h.gatewayService.ResolveChannelMappingAndRestrict(ctx, apiKey.GroupID, routeModel)
|
||||
mappedModelUnchanged := false
|
||||
if previous := turnChannelMapping.Load(); previous != nil && previous.turn < turn {
|
||||
mappedModelUnchanged = strings.TrimSpace(previous.mapping.MappedModel) == strings.TrimSpace(mapping.MappedModel)
|
||||
|
||||
@@ -1931,7 +1931,10 @@ func newOpenAIWSHandlerTestServer(t *testing.T, h *OpenAIGatewayHandler, subject
|
||||
|
||||
type openAIResponsesWSUsageLogCase struct {
|
||||
simpleModeRejectAtRead int64
|
||||
compositeResolver *service.CompositeRouteResolver
|
||||
accountPlatform string
|
||||
closeReason string
|
||||
closeStatus coderws.StatusCode
|
||||
firstPayload string
|
||||
// midPayload 在首个 turn 完成后发送(如 session.update),上游桩会为它
|
||||
// 回一个 response.completed,客户端按普通事件读取。
|
||||
@@ -2878,6 +2881,17 @@ func runOpenAIResponsesWebSocketUsageLogCase(t *testing.T, tc openAIResponsesWSU
|
||||
upstreamErrCh := make(chan error, 1)
|
||||
var channelSvc *service.ChannelService
|
||||
upstreamServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if tc.accountPlatform == service.PlatformGrok {
|
||||
payload, err := io.ReadAll(r.Body)
|
||||
if err != nil {
|
||||
upstreamErrCh <- err
|
||||
return
|
||||
}
|
||||
upstreamPayloadCh <- payload
|
||||
w.Header().Set("Content-Type", "text/event-stream")
|
||||
_, _ = fmt.Fprintf(w, "data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_grok_test\",\"model\":%q,\"usage\":{\"input_tokens\":2,\"output_tokens\":1}}}\n\n", gjson.GetBytes(payload, "model").String())
|
||||
return
|
||||
}
|
||||
conn, err := coderws.Accept(w, r, &coderws.AcceptOptions{
|
||||
CompressionMode: coderws.CompressionContextTakeover,
|
||||
})
|
||||
@@ -2945,6 +2959,9 @@ func runOpenAIResponsesWebSocketUsageLogCase(t *testing.T, tc openAIResponsesWSU
|
||||
"openai_apikey_responses_websockets_v2_mode": service.OpenAIWSIngressModePassthrough,
|
||||
},
|
||||
}
|
||||
if tc.accountPlatform != "" {
|
||||
account.Platform = tc.accountPlatform
|
||||
}
|
||||
if strings.TrimSpace(tc.ingressMode) != "" {
|
||||
account.Extra["openai_apikey_responses_websockets_v2_mode"] = tc.ingressMode
|
||||
}
|
||||
@@ -3000,7 +3017,7 @@ func runOpenAIResponsesWebSocketUsageLogCase(t *testing.T, tc openAIResponsesWSU
|
||||
service.NewBillingService(cfg, nil),
|
||||
nil,
|
||||
billingCacheSvc,
|
||||
nil,
|
||||
&compositeWSHTTPUpstream{},
|
||||
&service.DeferredService{},
|
||||
nil,
|
||||
nil,
|
||||
@@ -3021,6 +3038,7 @@ func runOpenAIResponsesWebSocketUsageLogCase(t *testing.T, tc openAIResponsesWSU
|
||||
}
|
||||
h := &OpenAIGatewayHandler{
|
||||
cfg: cfg,
|
||||
compositeResolver: tc.compositeResolver,
|
||||
gatewayService: gatewaySvc,
|
||||
billingCacheService: billingCacheSvc,
|
||||
apiKeyService: &service.APIKeyService{},
|
||||
@@ -3069,6 +3087,12 @@ func runOpenAIResponsesWebSocketUsageLogCase(t *testing.T, tc openAIResponsesWSU
|
||||
cancelWrite()
|
||||
require.NoError(t, err)
|
||||
|
||||
if tc.closeReason == "" {
|
||||
tc.closeReason = "not available for this group"
|
||||
}
|
||||
if tc.closeStatus == 0 {
|
||||
tc.closeStatus = coderws.StatusPolicyViolation
|
||||
}
|
||||
if tc.firstFrameCloseExpected {
|
||||
readCtx, cancelRead := context.WithTimeout(context.Background(), 3*time.Second)
|
||||
_, _, readErr := clientConn.Read(readCtx)
|
||||
@@ -3076,12 +3100,17 @@ func runOpenAIResponsesWebSocketUsageLogCase(t *testing.T, tc openAIResponsesWSU
|
||||
require.Error(t, readErr, "first frame should have been rejected with a close")
|
||||
var closeErr coderws.CloseError
|
||||
require.ErrorAs(t, readErr, &closeErr)
|
||||
require.Equal(t, coderws.StatusPolicyViolation, closeErr.Code)
|
||||
status := tc.closeStatus
|
||||
if status == 0 {
|
||||
status = coderws.StatusPolicyViolation
|
||||
}
|
||||
require.Equal(t, status, closeErr.Code)
|
||||
reason := tc.closeReason
|
||||
if reason == "" {
|
||||
reason = "not available for this group"
|
||||
}
|
||||
require.Contains(t, closeErr.Reason, reason)
|
||||
require.Empty(t, upstreamPayloadCh, "rejected first frame must not reach upstream")
|
||||
_ = clientConn.CloseNow()
|
||||
return openAIResponsesWSUsageLogResult{}
|
||||
}
|
||||
@@ -3115,12 +3144,17 @@ func runOpenAIResponsesWebSocketUsageLogCase(t *testing.T, tc openAIResponsesWSU
|
||||
require.Error(t, readErr, "second turn should have been rejected with a close")
|
||||
var closeErr coderws.CloseError
|
||||
require.ErrorAs(t, readErr, &closeErr)
|
||||
require.Equal(t, coderws.StatusPolicyViolation, closeErr.Code)
|
||||
status := tc.closeStatus
|
||||
if status == 0 {
|
||||
status = coderws.StatusPolicyViolation
|
||||
}
|
||||
require.Equal(t, status, closeErr.Code)
|
||||
reason := tc.closeReason
|
||||
if reason == "" {
|
||||
reason = "not available for this group"
|
||||
}
|
||||
require.Contains(t, closeErr.Reason, reason)
|
||||
require.Len(t, upstreamPayloadCh, turnCount-1, "rejected turn must not reach upstream")
|
||||
_ = clientConn.CloseNow()
|
||||
return openAIResponsesWSUsageLogResult{}
|
||||
}
|
||||
@@ -3149,6 +3183,9 @@ func runOpenAIResponsesWebSocketUsageLogCase(t *testing.T, tc openAIResponsesWSU
|
||||
}
|
||||
}
|
||||
|
||||
if tc.accountPlatform == service.PlatformGrok {
|
||||
upstreamErrCh <- nil
|
||||
}
|
||||
select {
|
||||
case upstreamErr := <-upstreamErrCh:
|
||||
require.NoError(t, upstreamErr)
|
||||
|
||||
@@ -0,0 +1,159 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"net/http"
|
||||
"testing"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
coderws "github.com/coder/websocket"
|
||||
"github.com/stretchr/testify/require"
|
||||
"github.com/tidwall/gjson"
|
||||
)
|
||||
|
||||
type compositeWSRouteRepo struct {
|
||||
service.CompositeModelRouteRepository
|
||||
routes []service.CompositeModelRoute
|
||||
err error
|
||||
}
|
||||
|
||||
func (r *compositeWSRouteRepo) ListByGroup(context.Context, int64, bool) ([]service.CompositeModelRoute, error) {
|
||||
return r.routes, r.err
|
||||
}
|
||||
|
||||
type compositeWSHTTPUpstream struct{ service.HTTPUpstream }
|
||||
|
||||
func (*compositeWSHTTPUpstream) Do(req *http.Request, _ string, _ int64, _ int) (*http.Response, error) {
|
||||
return http.DefaultClient.Do(req)
|
||||
}
|
||||
|
||||
func compositeWSResolver(platform, endpoint, upstream string) *service.CompositeRouteResolver {
|
||||
return service.NewCompositeRouteResolver(&compositeWSRouteRepo{routes: []service.CompositeModelRoute{{
|
||||
GroupID: 4201, PublicModel: "public-alias", MatchType: service.CompositeRouteMatchExact,
|
||||
TargetPlatform: platform, Endpoint: endpoint, UpstreamModel: upstream, Enabled: true,
|
||||
}}})
|
||||
}
|
||||
func compositeWSGroup(models ...string) *service.Group {
|
||||
g := wsAllowlistGroup(len(models) > 0, models...)
|
||||
g.Platform = service.PlatformComposite
|
||||
return g
|
||||
}
|
||||
|
||||
func TestOpenAIResponsesWebSocket_CompositeAlias(t *testing.T) {
|
||||
for _, platform := range []string{service.PlatformOpenAI, service.PlatformGrok} {
|
||||
for _, endpoint := range []string{service.CompositeRouteEndpointResponses, service.CompositeRouteEndpointAny} {
|
||||
for _, mode := range []string{service.OpenAIWSIngressModePassthrough, service.OpenAIWSIngressModeDedicated} {
|
||||
t.Run(platform+"/"+endpoint+"/"+mode, func(t *testing.T) {
|
||||
upstream := "gpt-5.4"
|
||||
if platform == service.PlatformGrok {
|
||||
upstream = "grok-4.3"
|
||||
}
|
||||
got := runOpenAIResponsesWebSocketUsageLogCase(t, openAIResponsesWSUsageLogCase{
|
||||
firstPayload: `{"type":"response.create","model":"public-alias","input":"hi"}`,
|
||||
secondPayload: `{"type":"response.create","input":"again"}`,
|
||||
group: compositeWSGroup("public-alias"), accountPlatform: platform, ingressMode: mode,
|
||||
compositeResolver: compositeWSResolver(platform, endpoint, upstream),
|
||||
})
|
||||
require.Len(t, got.upstreamPayloads, 2)
|
||||
for i, payload := range got.upstreamPayloads {
|
||||
require.Equal(t, upstream, gjson.GetBytes(payload, "model").String())
|
||||
require.Equal(t, "public-alias", gjson.GetBytes(got.clientEvents[i], "response.model").String())
|
||||
require.Equal(t, "public-alias", got.logs[i].RequestedModel)
|
||||
require.NotNil(t, got.logs[i].UpstreamModel)
|
||||
require.Equal(t, upstream, *got.logs[i].UpstreamModel)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenAIResponsesWebSocket_CompositeRouteRejections(t *testing.T) {
|
||||
for _, tc := range []struct {
|
||||
name, platform, endpoint, model, reason string
|
||||
repoErr error
|
||||
status coderws.StatusCode
|
||||
}{
|
||||
{name: "disallowed platform", platform: service.PlatformAnthropic, endpoint: "responses", model: "public-alias", reason: "only supports OpenAI-compatible"},
|
||||
{name: "wrong endpoint", platform: service.PlatformOpenAI, endpoint: "messages", model: "public-alias", reason: "only supports OpenAI-compatible"},
|
||||
{name: "unknown alias", platform: service.PlatformOpenAI, endpoint: "responses", model: "unknown-alias", reason: "only supports OpenAI-compatible"},
|
||||
{name: "resolver error", model: "gpt-5.4", repoErr: errors.New("database unavailable"), reason: "Failed to resolve composite model route", status: coderws.StatusInternalError},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
resolver := compositeWSResolver(tc.platform, tc.endpoint, "gpt-5.4")
|
||||
if tc.repoErr != nil {
|
||||
resolver = service.NewCompositeRouteResolver(&compositeWSRouteRepo{err: tc.repoErr})
|
||||
}
|
||||
runOpenAIResponsesWebSocketUsageLogCase(t, openAIResponsesWSUsageLogCase{
|
||||
firstPayload: `{"type":"response.create","model":"` + tc.model + `"}`, group: compositeWSGroup(), compositeResolver: resolver,
|
||||
firstFrameCloseExpected: true, closeReason: tc.reason, closeStatus: tc.status,
|
||||
})
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenAIResponsesWebSocket_CompositeAdmissionUsesPublicModel(t *testing.T) {
|
||||
runOpenAIResponsesWebSocketUsageLogCase(t, openAIResponsesWSUsageLogCase{
|
||||
firstPayload: `{"type":"response.create","model":"public-alias"}`, group: compositeWSGroup("gpt-5.4"),
|
||||
compositeResolver: compositeWSResolver(service.PlatformOpenAI, "responses", "gpt-5.4"), firstFrameCloseExpected: true,
|
||||
})
|
||||
}
|
||||
|
||||
func TestOpenAIResponsesWebSocket_CompositeModelSwitchRequiresReconnect(t *testing.T) {
|
||||
for _, mode := range []string{service.OpenAIWSIngressModePassthrough, service.OpenAIWSIngressModeDedicated} {
|
||||
for _, model := range []string{"gpt-5.4", "grok-4.3", "unknown-alias"} {
|
||||
t.Run(mode+"/"+model, func(t *testing.T) {
|
||||
runOpenAIResponsesWebSocketUsageLogCase(t, openAIResponsesWSUsageLogCase{
|
||||
firstPayload: `{"type":"response.create","model":"public-alias"}`,
|
||||
secondPayload: `{"type":"response.create","model":"` + model + `"}`, group: compositeWSGroup(), ingressMode: mode,
|
||||
compositeResolver: compositeWSResolver(service.PlatformOpenAI, "responses", "gpt-5.4"),
|
||||
secondTurnCloseExpected: true, closeReason: "model switch requires reconnect",
|
||||
})
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenAIResponsesWebSocket_CompositeChannelBilling(t *testing.T) {
|
||||
for _, source := range []string{service.BillingModelSourceRequested, service.BillingModelSourceChannelMapped, service.BillingModelSourceUpstream} {
|
||||
t.Run(source, func(t *testing.T) {
|
||||
got := runOpenAIResponsesWebSocketUsageLogCase(t, openAIResponsesWSUsageLogCase{
|
||||
firstPayload: `{"type":"response.create","model":"gpt-5.6-sol"}`,
|
||||
secondPayload: `{"type":"response.create","model":"gpt-5.6-sol"}`,
|
||||
group: compositeWSGroup("gpt-5.6-sol"),
|
||||
compositeResolver: service.NewCompositeRouteResolver(&compositeWSRouteRepo{routes: []service.CompositeModelRoute{{PublicModel: "gpt-5.6-sol", MatchType: service.CompositeRouteMatchExact, TargetPlatform: service.PlatformOpenAI, Endpoint: service.CompositeRouteEndpointResponses, UpstreamModel: "route-target"}}}),
|
||||
channelMapping: map[string]string{"route-target": "gpt-5.4"}, billingModelSource: source,
|
||||
accountModelMapping: map[string]any{"gpt-5.4": "gpt-5.4"},
|
||||
})
|
||||
for i, log := range got.logs {
|
||||
require.Equal(t, "gpt-5.4", gjson.GetBytes(got.upstreamPayloads[i], "model").String())
|
||||
require.Equal(t, "gpt-5.6-sol", log.RequestedModel)
|
||||
require.Equal(t, "gpt-5.6-sol", log.Model)
|
||||
if source == service.BillingModelSourceRequested {
|
||||
require.InDelta(t, 40e-6, log.TotalCost, 1e-12)
|
||||
} else {
|
||||
require.InDelta(t, 20e-6, log.TotalCost, 1e-12)
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenAIResponsesWebSocket_CompositeDetectorFallback(t *testing.T) {
|
||||
got := runOpenAIResponsesWebSocketUsageLogCase(t, openAIResponsesWSUsageLogCase{
|
||||
firstPayload: `{"type":"response.create","model":"gpt-5.4"}`, group: compositeWSGroup(),
|
||||
compositeResolver: service.NewCompositeRouteResolver(&compositeWSRouteRepo{}),
|
||||
})
|
||||
require.Equal(t, "gpt-5.4", gjson.GetBytes(got.upstreamFirstPayload, "model").String())
|
||||
}
|
||||
|
||||
func TestOpenAIResponsesWebSocket_CompositeSessionModelSwitchRequiresReconnect(t *testing.T) {
|
||||
runOpenAIResponsesWebSocketUsageLogCase(t, openAIResponsesWSUsageLogCase{
|
||||
firstPayload: `{"type":"response.create","model":"public-alias"}`,
|
||||
midPayload: `{"type":"session.update","session":{"model":"grok-4.3"}}`,
|
||||
secondPayload: `{"type":"response.create"}`, group: compositeWSGroup(),
|
||||
compositeResolver: compositeWSResolver(service.PlatformOpenAI, "responses", "gpt-5.4"),
|
||||
secondTurnCloseExpected: true, closeReason: "model switch requires reconnect",
|
||||
})
|
||||
}
|
||||
@@ -132,10 +132,12 @@ func ProvideOpenAIGatewayHandler(
|
||||
grokQuotaService *service.GrokQuotaService,
|
||||
cfg *config.Config,
|
||||
coordinator *securityaudit.Coordinator,
|
||||
compositeResolver *service.CompositeRouteResolver,
|
||||
) *OpenAIGatewayHandler {
|
||||
gatewayService.SetPluginManager(pluginManager)
|
||||
h := NewOpenAIGatewayHandler(gatewayService, concurrencyService, billingCacheService, apiKeyService,
|
||||
usageRecordWorkerPool, errorPassthroughService, contentModerationService, opsService, cfg)
|
||||
h.compositeResolver = compositeResolver
|
||||
h.securityAuditCoordinator = coordinator
|
||||
h.grokMediaEligibilityProber = grokQuotaService
|
||||
return h
|
||||
|
||||
@@ -1776,6 +1776,11 @@ func (s *defaultOpenAIAccountScheduler) isAccountRequestCompatibleReason(ctx con
|
||||
if account == nil {
|
||||
return false, "account_nil"
|
||||
}
|
||||
if source, ok := CompositeRouteSourceFromContext(ctx); ok && source == CompositeRouteSourceAccount {
|
||||
if publicModel, modelOK := RequestedPublicModelFromContext(ctx); modelOK && !explicitModelMappingClaims(*account, publicModel) {
|
||||
return false, "account_model_not_owned"
|
||||
}
|
||||
}
|
||||
if req.RequirePrivacySet && !account.IsPrivacySet() {
|
||||
return false, "privacy_not_set"
|
||||
}
|
||||
|
||||
@@ -44,6 +44,10 @@ func (r schedulerTestOpenAIAccountRepo) ListSchedulableByGroupIDAndPlatform(ctx
|
||||
return result, nil
|
||||
}
|
||||
|
||||
func (r schedulerTestOpenAIAccountRepo) ListSchedulableByGroupID(ctx context.Context, groupID int64) ([]Account, error) {
|
||||
return append([]Account(nil), r.accounts...), nil
|
||||
}
|
||||
|
||||
func (r schedulerTestOpenAIAccountRepo) ListSchedulableByPlatform(ctx context.Context, platform string) ([]Account, error) {
|
||||
var result []Account
|
||||
for _, acc := range r.accounts {
|
||||
@@ -1257,6 +1261,116 @@ func TestOpenAIGatewayService_SelectAccountWithScheduler_StickyWeightedSessionIn
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenAIGatewayService_SelectAccountWithScheduler_AccountModelRouteOnlyUsesClaimingAccount(t *testing.T) {
|
||||
resetOpenAIAdvancedSchedulerSettingCacheForTest()
|
||||
|
||||
ctx := context.Background()
|
||||
groupID := int64(1010711)
|
||||
owner := Account{
|
||||
ID: 371011, Platform: PlatformOpenAI, Type: AccountTypeAPIKey,
|
||||
Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 10,
|
||||
GroupIDs: []int64{groupID},
|
||||
Credentials: map[string]any{"model_mapping": map[string]any{"team-alias": "gpt-5.1"}},
|
||||
}
|
||||
nonOwner := Account{
|
||||
ID: 371012, Platform: PlatformOpenAI, Type: AccountTypeAPIKey,
|
||||
Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 0,
|
||||
GroupIDs: []int64{groupID},
|
||||
}
|
||||
repo := schedulerTestOpenAIAccountRepo{accounts: []Account{owner, nonOwner}}
|
||||
ownership, err := (&GatewayService{accountRepo: repo}).resolveCompositeModelOwnership(ctx, groupID, "team-alias")
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, CompositeModelOwnership{TargetPlatform: PlatformOpenAI, Matched: true}, ownership)
|
||||
ctx = WithCompositeRouteDecision(ctx, CompositeRouteDecision{
|
||||
Matched: true, Source: CompositeRouteSourceAccount, GroupID: groupID,
|
||||
PublicModel: "team-alias", TargetPlatform: ownership.TargetPlatform, UpstreamModel: "team-alias",
|
||||
})
|
||||
|
||||
for _, tt := range []struct {
|
||||
name string
|
||||
sessionHash string
|
||||
sessionBindings map[string]int64
|
||||
}{
|
||||
{name: "initial filter excludes non owner"},
|
||||
{name: "non owner sticky hit is rejected", sessionHash: "non-owner", sessionBindings: map[string]int64{"openai:non-owner": nonOwner.ID}},
|
||||
} {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
cfg := &config.Config{}
|
||||
cfg.Gateway.Scheduling.LoadBatchEnabled = false
|
||||
svc := &OpenAIGatewayService{
|
||||
accountRepo: repo,
|
||||
cache: &schedulerTestGatewayCache{sessionBindings: tt.sessionBindings},
|
||||
cfg: cfg,
|
||||
rateLimitService: newOpenAIAdvancedSchedulerRateLimitService("true"),
|
||||
concurrencyService: NewConcurrencyService(schedulerTestConcurrencyCache{}),
|
||||
}
|
||||
|
||||
selection, _, selectErr := svc.SelectAccountWithScheduler(
|
||||
ctx, &groupID, "", tt.sessionHash, "team-alias", nil, OpenAIUpstreamTransportAny, false,
|
||||
)
|
||||
require.NoError(t, selectErr)
|
||||
require.NotNil(t, selection)
|
||||
require.NotNil(t, selection.Account)
|
||||
require.Equal(t, owner.ID, selection.Account.ID)
|
||||
if selection.ReleaseFunc != nil {
|
||||
selection.ReleaseFunc()
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenAIGatewayService_SelectAccountWithScheduler_AccountModelRouteDBRecheckFailsOverToOwner(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
groupID := int64(1010712)
|
||||
mapping := map[string]any{"model_mapping": map[string]any{"team-alias": "gpt-5.1"}}
|
||||
stalePrimary := &Account{
|
||||
ID: 371021, Platform: PlatformOpenAI, Type: AccountTypeAPIKey,
|
||||
Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 0,
|
||||
GroupIDs: []int64{groupID}, Credentials: mapping,
|
||||
}
|
||||
staleBackup := &Account{
|
||||
ID: 371022, Platform: PlatformOpenAI, Type: AccountTypeAPIKey,
|
||||
Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 10,
|
||||
GroupIDs: []int64{groupID}, Credentials: mapping,
|
||||
}
|
||||
dbPrimary := *stalePrimary
|
||||
dbPrimary.Credentials = nil
|
||||
dbBackup := *staleBackup
|
||||
repo := schedulerTestOpenAIAccountRepo{accounts: []Account{dbPrimary, dbBackup}}
|
||||
ctx = WithCompositeRouteDecision(ctx, CompositeRouteDecision{
|
||||
Matched: true, Source: CompositeRouteSourceAccount, GroupID: groupID,
|
||||
PublicModel: "team-alias", TargetPlatform: PlatformOpenAI, UpstreamModel: "team-alias",
|
||||
})
|
||||
acquiredIDs, releasedIDs := []int64{}, []int64{}
|
||||
svc := &OpenAIGatewayService{
|
||||
accountRepo: repo,
|
||||
cfg: &config.Config{RunMode: config.RunModeStandard},
|
||||
rateLimitService: newOpenAIAdvancedSchedulerRateLimitService("true"),
|
||||
schedulerSnapshot: &SchedulerSnapshotService{cache: &openAISnapshotCacheStub{
|
||||
snapshotAccounts: []*Account{stalePrimary, staleBackup},
|
||||
accountsByID: map[int64]*Account{stalePrimary.ID: stalePrimary, staleBackup.ID: staleBackup},
|
||||
}},
|
||||
concurrencyService: NewConcurrencyService(schedulerTestConcurrencyCache{
|
||||
acquiredIDs: &acquiredIDs, releasedIDs: &releasedIDs,
|
||||
}),
|
||||
}
|
||||
|
||||
selection, _, err := svc.SelectAccountWithScheduler(
|
||||
ctx, &groupID, "", "", "team-alias", nil, OpenAIUpstreamTransportAny, false,
|
||||
)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, selection)
|
||||
require.NotNil(t, selection.Account)
|
||||
require.Equal(t, dbBackup.ID, selection.Account.ID)
|
||||
if containsInt64(acquiredIDs, dbPrimary.ID) {
|
||||
require.Contains(t, releasedIDs, dbPrimary.ID)
|
||||
}
|
||||
require.Contains(t, acquiredIDs, dbBackup.ID)
|
||||
if selection.ReleaseFunc != nil {
|
||||
selection.ReleaseFunc()
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenAIGatewayService_SelectAccountWithScheduler_StickyWeightedPreviousRequiresMovableContext(t *testing.T) {
|
||||
resetOpenAIAdvancedSchedulerSettingCacheForTest()
|
||||
|
||||
@@ -3812,3 +3926,126 @@ func TestOpenAIGatewayService_SelectAccountWithScheduler_SubscriptionPriorityWai
|
||||
require.Equal(t, int64(38011), selection.WaitPlan.AccountID)
|
||||
require.Equal(t, openAIAccountScheduleLayerLoadBalance, decision.Layer)
|
||||
}
|
||||
|
||||
func TestOpenAIGatewayService_SelectAccountWithScheduler_LegacyAccountModelRouteExcludesNonOwner(t *testing.T) {
|
||||
resetOpenAIAdvancedSchedulerSettingCacheForTest()
|
||||
|
||||
ctx := context.Background()
|
||||
groupID := int64(1010713)
|
||||
owner := Account{
|
||||
ID: 371031, Platform: PlatformOpenAI, Type: AccountTypeAPIKey,
|
||||
Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 10,
|
||||
GroupIDs: []int64{groupID},
|
||||
Credentials: map[string]any{"model_mapping": map[string]any{"team-alias": "gpt-5.1"}},
|
||||
// WS transport must be resolvable for the websocket sub-case so the
|
||||
// ownership filter, not transport admission, is the behavior under test.
|
||||
Extra: map[string]any{"responses_websockets_v2_enabled": true},
|
||||
}
|
||||
nonOwner := Account{
|
||||
ID: 371032, Platform: PlatformOpenAI, Type: AccountTypeAPIKey,
|
||||
Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 0,
|
||||
GroupIDs: []int64{groupID},
|
||||
Extra: map[string]any{"responses_websockets_v2_enabled": true},
|
||||
}
|
||||
repo := schedulerTestOpenAIAccountRepo{accounts: []Account{owner, nonOwner}}
|
||||
ownership, err := (&GatewayService{accountRepo: repo}).resolveCompositeModelOwnership(ctx, groupID, "team-alias")
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, CompositeModelOwnership{TargetPlatform: PlatformOpenAI, Matched: true}, ownership)
|
||||
ctx = WithCompositeRouteDecision(ctx, CompositeRouteDecision{
|
||||
Matched: true, Source: CompositeRouteSourceAccount, GroupID: groupID,
|
||||
PublicModel: "team-alias", TargetPlatform: ownership.TargetPlatform, UpstreamModel: "team-alias",
|
||||
})
|
||||
|
||||
for _, tt := range []struct {
|
||||
name string
|
||||
requiredTransport OpenAIUpstreamTransport
|
||||
sessionHash string
|
||||
sessionBindings map[string]int64
|
||||
}{
|
||||
{name: "initial filter excludes non owner", requiredTransport: OpenAIUpstreamTransportAny},
|
||||
{name: "initial filter excludes non owner for websocket transport", requiredTransport: OpenAIUpstreamTransportResponsesWebsocket},
|
||||
{name: "non owner sticky hit is rejected", requiredTransport: OpenAIUpstreamTransportAny, sessionHash: "non-owner", sessionBindings: map[string]int64{"openai:non-owner": nonOwner.ID}},
|
||||
} {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
cfg := &config.Config{}
|
||||
cfg.Gateway.OpenAIWS.Enabled = true
|
||||
cfg.Gateway.OpenAIWS.APIKeyEnabled = true
|
||||
cfg.Gateway.OpenAIWS.ResponsesWebsockets = true
|
||||
svc := &OpenAIGatewayService{
|
||||
accountRepo: repo,
|
||||
cache: &schedulerTestGatewayCache{sessionBindings: tt.sessionBindings},
|
||||
cfg: cfg,
|
||||
rateLimitService: newOpenAIAdvancedSchedulerRateLimitService("false"),
|
||||
concurrencyService: NewConcurrencyService(schedulerTestConcurrencyCache{}),
|
||||
}
|
||||
require.Nil(t, svc.getOpenAIAccountScheduler(ctx), "test precondition: advanced scheduler disabled")
|
||||
|
||||
selection, _, selectErr := svc.SelectAccountWithScheduler(
|
||||
ctx, &groupID, "", tt.sessionHash, "team-alias", nil, tt.requiredTransport, false,
|
||||
)
|
||||
require.NoError(t, selectErr)
|
||||
require.NotNil(t, selection)
|
||||
require.NotNil(t, selection.Account)
|
||||
require.Equal(t, owner.ID, selection.Account.ID)
|
||||
if selection.ReleaseFunc != nil {
|
||||
selection.ReleaseFunc()
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenAIGatewayService_SelectAccountWithScheduler_LegacyAccountModelRouteDBRecheckFailsOverToOwner(t *testing.T) {
|
||||
resetOpenAIAdvancedSchedulerSettingCacheForTest()
|
||||
|
||||
ctx := context.Background()
|
||||
groupID := int64(1010714)
|
||||
mapping := map[string]any{"model_mapping": map[string]any{"team-alias": "gpt-5.1"}}
|
||||
stalePrimary := &Account{
|
||||
ID: 371041, Platform: PlatformOpenAI, Type: AccountTypeAPIKey,
|
||||
Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 0,
|
||||
GroupIDs: []int64{groupID}, Credentials: mapping,
|
||||
}
|
||||
staleBackup := &Account{
|
||||
ID: 371042, Platform: PlatformOpenAI, Type: AccountTypeAPIKey,
|
||||
Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 10,
|
||||
GroupIDs: []int64{groupID}, Credentials: mapping,
|
||||
}
|
||||
dbPrimary := *stalePrimary
|
||||
dbPrimary.Credentials = nil
|
||||
dbBackup := *staleBackup
|
||||
repo := schedulerTestOpenAIAccountRepo{accounts: []Account{dbPrimary, dbBackup}}
|
||||
ctx = WithCompositeRouteDecision(ctx, CompositeRouteDecision{
|
||||
Matched: true, Source: CompositeRouteSourceAccount, GroupID: groupID,
|
||||
PublicModel: "team-alias", TargetPlatform: PlatformOpenAI, UpstreamModel: "team-alias",
|
||||
})
|
||||
acquiredIDs, releasedIDs := []int64{}, []int64{}
|
||||
cfg := &config.Config{RunMode: config.RunModeStandard}
|
||||
svc := &OpenAIGatewayService{
|
||||
accountRepo: repo,
|
||||
cfg: cfg,
|
||||
rateLimitService: newOpenAIAdvancedSchedulerRateLimitService("false"),
|
||||
schedulerSnapshot: &SchedulerSnapshotService{cache: &openAISnapshotCacheStub{
|
||||
snapshotAccounts: []*Account{stalePrimary, staleBackup},
|
||||
accountsByID: map[int64]*Account{stalePrimary.ID: stalePrimary, staleBackup.ID: staleBackup},
|
||||
}},
|
||||
concurrencyService: NewConcurrencyService(schedulerTestConcurrencyCache{
|
||||
acquiredIDs: &acquiredIDs, releasedIDs: &releasedIDs,
|
||||
}),
|
||||
}
|
||||
require.Nil(t, svc.getOpenAIAccountScheduler(ctx), "test precondition: advanced scheduler disabled")
|
||||
|
||||
selection, _, err := svc.SelectAccountWithScheduler(
|
||||
ctx, &groupID, "", "", "team-alias", nil, OpenAIUpstreamTransportAny, false,
|
||||
)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, selection)
|
||||
require.NotNil(t, selection.Account)
|
||||
require.Equal(t, dbBackup.ID, selection.Account.ID)
|
||||
if containsInt64(acquiredIDs, dbPrimary.ID) {
|
||||
require.Contains(t, releasedIDs, dbPrimary.ID)
|
||||
}
|
||||
require.Contains(t, acquiredIDs, dbBackup.ID)
|
||||
if selection.ReleaseFunc != nil {
|
||||
selection.ReleaseFunc()
|
||||
}
|
||||
}
|
||||
|
||||
@@ -392,6 +392,16 @@ func openAICompatibleAccountEligibilityFailureReasonBeforeProfit(ctx context.Con
|
||||
if account == nil {
|
||||
return "account_nil"
|
||||
}
|
||||
// account_model composite routes publish an alias that only accounts with an
|
||||
// explicit model mapping own. Both scheduler modes (advanced and legacy) must
|
||||
// enforce ownership before any priority/sticky/transport consideration; a
|
||||
// non-owner would forward the raw public alias and get model_not_found or
|
||||
// hit an unrelated same-name upstream model.
|
||||
if source, ok := CompositeRouteSourceFromContext(ctx); ok && source == CompositeRouteSourceAccount {
|
||||
if publicModel, modelOK := RequestedPublicModelFromContext(ctx); modelOK && !explicitModelMappingClaims(*account, publicModel) {
|
||||
return "account_model_not_owned"
|
||||
}
|
||||
}
|
||||
if account.Platform != platform || !account.IsOpenAICompatible() {
|
||||
return "platform_mismatch"
|
||||
}
|
||||
|
||||
@@ -585,7 +585,7 @@ func (s *OpenAIGatewayService) ProxyResponsesWebSocketFromClient(
|
||||
bridgeAccountFailoverInputExists := false
|
||||
for turn := 1; ; turn++ {
|
||||
if turn > 1 && hooks != nil && hooks.BeforeRequest != nil {
|
||||
if err := hooks.BeforeRequest(turn, currentBridgePayload.payloadRaw, currentBridgePayload.originalModel); err != nil {
|
||||
if err := hooks.BeforeRequest(turn, currentBridgePayload.rawForHash, currentBridgePayload.originalModel); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
@@ -1272,6 +1272,8 @@ func (s *OpenAIGatewayService) ProxyResponsesWebSocketFromClient(
|
||||
}
|
||||
|
||||
currentPayload := firstPayload.payloadRaw
|
||||
// Admission hooks must see client model candidates before upstream mapping.
|
||||
currentClientPayload := firstPayload.rawForHash
|
||||
currentOriginalModel := firstPayload.originalModel
|
||||
currentImageBillingModel := firstPayload.imageBillingModel
|
||||
currentImageSizeTier := firstPayload.imageSizeTier
|
||||
@@ -1461,7 +1463,7 @@ func (s *OpenAIGatewayService) ProxyResponsesWebSocketFromClient(
|
||||
}
|
||||
for {
|
||||
if turn > 1 && !skipBeforeTurn && hooks != nil && hooks.BeforeRequest != nil {
|
||||
if err := hooks.BeforeRequest(turn, currentPayload, currentOriginalModel); err != nil {
|
||||
if err := hooks.BeforeRequest(turn, currentClientPayload, currentOriginalModel); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
@@ -1970,6 +1972,7 @@ func (s *OpenAIGatewayService) ProxyResponsesWebSocketFromClient(
|
||||
}
|
||||
}
|
||||
currentPayload = nextPayload.payloadRaw
|
||||
currentClientPayload = nextPayload.rawForHash
|
||||
currentOriginalModel = nextPayload.originalModel
|
||||
currentImageBillingModel = nextPayload.imageBillingModel
|
||||
currentImageSizeTier = nextPayload.imageSizeTier
|
||||
|
||||
@@ -63,6 +63,13 @@ payloads, and apply the configured `upstream_model` before dispatch.
|
||||
Codex model manifest requests reuse the existing OpenAI account selection and
|
||||
failover path within the Composite group.
|
||||
|
||||
Responses WebSocket resolves the first `response.create` model in the
|
||||
`responses` route domain (including `any` routes). Only OpenAI and Grok targets
|
||||
are supported. The public model remains the admission and response identity;
|
||||
the route's upstream model is applied before channel/account mapping. Subsequent
|
||||
turns can repeat or omit the public model. Changing it requires reconnecting so
|
||||
the gateway can resolve the route and select an account again.
|
||||
|
||||
## Built-In Detection
|
||||
|
||||
Composite routing detects common public model IDs and provider-prefixed IDs:
|
||||
|
||||
Reference in New Issue
Block a user