fix(openai): isolate passthrough websocket sessions

This commit is contained in:
wucm667
2026-08-30 16:49:01 +08:00
parent b5827cfd54
commit d5a0124637
2 changed files with 58 additions and 0 deletions
@@ -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)
@@ -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}