diff --git a/backend/internal/service/openai_ws_session_preemption.go b/backend/internal/service/openai_ws_session_preemption.go index 54da7f7a1d..3047de6998 100644 --- a/backend/internal/service/openai_ws_session_preemption.go +++ b/backend/internal/service/openai_ws_session_preemption.go @@ -56,6 +56,10 @@ func (s *OpenAIGatewayService) BeginOpenAIWSIngressSessionPreemption( if armed, _ := ctx.Value(openAIWSSessionPreemptContextKey{}).(bool); armed { return ctx, func() {}, true } + if s != nil && s.cfg != nil && s.cfg.Gateway.OpenAIWS.ModeRouterV2Enabled && + account != nil && account.ResolveOpenAIResponsesWebSocketV2Mode(s.cfg.Gateway.OpenAIWS.IngressModeDefault) == OpenAIWSIngressModePassthrough { + return ctx, func() {}, false + } preemptSessionHash := "" preemptGroupID := getOpenAIGroupIDFromContext(c) diff --git a/backend/internal/service/openai_ws_session_preemption_test.go b/backend/internal/service/openai_ws_session_preemption_test.go index 9df9465944..0e0c6f5622 100644 --- a/backend/internal/service/openai_ws_session_preemption_test.go +++ b/backend/internal/service/openai_ws_session_preemption_test.go @@ -11,6 +11,7 @@ import ( "testing" "time" + "github.com/Wei-Shaw/sub2api/internal/config" "github.com/gin-gonic/gin" "github.com/stretchr/testify/require" ) @@ -151,6 +152,59 @@ func TestOpenAIWSIngressSessionPreemptionSurvivesNestedForwardCleanup(t *testing require.True(t, IsOpenAIWSSessionPreemptedError(context.Cause(firstCtx))) } +func TestOpenAIWSIngressSessionPreemptionRespectsResolvedMode(t *testing.T) { + gin.SetMode(gin.TestMode) + groupID := int64(7) + newContext := func() *gin.Context { + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + c.Request = httptest.NewRequest(http.MethodGet, "/v1/responses", nil) + c.Set("api_key", &APIKey{ID: 11, GroupID: &groupID}) + return c + } + newAccount := func(mode string) *Account { + return &Account{ + ID: 1, + Platform: PlatformOpenAI, + Type: AccountTypeOAuth, + Extra: map[string]any{ + "openai_oauth_responses_websockets_v2_mode": mode, + }, + } + } + firstMessage := []byte(`{"type":"response.create","prompt_cache_key":"session-1","input":"hello"}`) + cfg := &config.Config{} + cfg.Gateway.OpenAIWS.ModeRouterV2Enabled = true + cfg.Gateway.OpenAIWS.IngressModeDefault = OpenAIWSIngressModeCtxPool + svc := &OpenAIGatewayService{cfg: cfg} + + passthrough := newAccount(OpenAIWSIngressModePassthrough) + firstCtx, firstCleanup, armed := svc.BeginOpenAIWSIngressSessionPreemption( + context.Background(), newContext(), passthrough, firstMessage, + ) + require.False(t, armed) + defer firstCleanup() + secondCtx, secondCleanup, armed := svc.BeginOpenAIWSIngressSessionPreemption( + context.Background(), newContext(), passthrough, firstMessage, + ) + require.False(t, armed) + defer secondCleanup() + require.NoError(t, firstCtx.Err(), "concurrent passthrough request must remain isolated") + require.NoError(t, secondCtx.Err()) + + ctxPool := newAccount(OpenAIWSIngressModeCtxPool) + sharedCtx, sharedCleanup, armed := svc.BeginOpenAIWSIngressSessionPreemption( + context.Background(), newContext(), ctxPool, firstMessage, + ) + require.True(t, armed) + defer sharedCleanup() + _, replacementCleanup, armed := svc.BeginOpenAIWSIngressSessionPreemption( + context.Background(), newContext(), ctxPool, firstMessage, + ) + require.True(t, armed) + defer replacementCleanup() + require.True(t, IsOpenAIWSSessionPreemptedError(context.Cause(sharedCtx)), "ctx_pool must retain shared-session preemption") +} + func TestOpenAIWSSessionPreemptRemoteClaimAndStaleReleaseAreAtomic(t *testing.T) { cache := &openAIWSSessionPreemptCacheStub{} svc := &OpenAIGatewayService{cache: cache}