mirror of
https://github.com/Wei-Shaw/sub2api.git
synced 2026-10-14 03:52:26 +08:00
fix(openai): isolate passthrough websocket sessions
This commit is contained in:
@@ -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}
|
||||
|
||||
Reference in New Issue
Block a user