diff --git a/backend/cmd/server/wire_gen.go b/backend/cmd/server/wire_gen.go index 37390f92cf..7adddf9ad4 100644 --- a/backend/cmd/server/wire_gen.go +++ b/backend/cmd/server/wire_gen.go @@ -294,7 +294,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) diff --git a/backend/internal/handler/openai_gateway_handler.go b/backend/internal/handler/openai_gateway_handler.go index bf4563f2f2..bfb4475931 100644 --- a/backend/internal/handler/openai_gateway_handler.go +++ b/backend/internal/handler/openai_gateway_handler.go @@ -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) diff --git a/backend/internal/handler/openai_gateway_handler_test.go b/backend/internal/handler/openai_gateway_handler_test.go index d13ac01589..cf9f1d9e7c 100644 --- a/backend/internal/handler/openai_gateway_handler_test.go +++ b/backend/internal/handler/openai_gateway_handler_test.go @@ -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) diff --git a/backend/internal/handler/openai_gateway_ws_composite_test.go b/backend/internal/handler/openai_gateway_ws_composite_test.go new file mode 100644 index 0000000000..7c2d79bbaa --- /dev/null +++ b/backend/internal/handler/openai_gateway_ws_composite_test.go @@ -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", + }) +} diff --git a/backend/internal/handler/wire.go b/backend/internal/handler/wire.go index d35245ecbf..e51b105728 100644 --- a/backend/internal/handler/wire.go +++ b/backend/internal/handler/wire.go @@ -130,10 +130,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 diff --git a/backend/internal/service/openai_ws_forwarder_ingress.go b/backend/internal/service/openai_ws_forwarder_ingress.go index ca6135352b..98e6327c3c 100644 --- a/backend/internal/service/openai_ws_forwarder_ingress.go +++ b/backend/internal/service/openai_ws_forwarder_ingress.go @@ -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 @@ -1460,7 +1462,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 } } @@ -1941,6 +1943,7 @@ func (s *OpenAIGatewayService) ProxyResponsesWebSocketFromClient( } } currentPayload = nextPayload.payloadRaw + currentClientPayload = nextPayload.rawForHash currentOriginalModel = nextPayload.originalModel currentImageBillingModel = nextPayload.imageBillingModel currentImageSizeTier = nextPayload.imageSizeTier diff --git a/docs/COMPOSITE_GROUPS.md b/docs/COMPOSITE_GROUPS.md index bf32c902bc..99456f6efc 100644 --- a/docs/COMPOSITE_GROUPS.md +++ b/docs/COMPOSITE_GROUPS.md @@ -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: