mirror of
https://github.com/Wei-Shaw/sub2api.git
synced 2026-10-14 12:02:27 +08:00
Merge pull request #6704 from aofee/codex/fix-http-bridge-session-isolation
fix(openai-ws): 修复 HTTP bridge 并发连接抢占与状态串用
This commit is contained in:
Wesley Liddick
@@ -521,6 +521,7 @@ func (s *OpenAIGatewayService) ProxyResponsesWebSocketFromClient(
|
||||
return err
|
||||
}
|
||||
|
||||
useHTTPBridge := forceHTTPBridge || s.shouldBridgeOpenAIWSHTTP(account, firstPayload.payloadBytes, firstPayload.previousResponseID)
|
||||
turnState := strings.TrimSpace(c.GetHeader(openAIWSTurnStateHeader))
|
||||
stateStore := s.getOpenAIWSStateStore()
|
||||
groupID := getOpenAIGroupIDFromContext(c)
|
||||
@@ -530,20 +531,25 @@ func (s *OpenAIGatewayService) ProxyResponsesWebSocketFromClient(
|
||||
storeDisabled := false
|
||||
refreshIngressRouteState := func(payload openAIWSClientPayload) {
|
||||
sessionHash = s.GenerateSessionHash(c, payload.rawForHash)
|
||||
preferredConnID = ""
|
||||
storeDisabled = s.isOpenAIWSStoreDisabledInRequestRaw(payload.payloadRaw, account)
|
||||
if useHTTPBridge {
|
||||
// Sticky account affinity may be shared, but an HTTP bridge must not
|
||||
// inherit another connection's native WS turn state or socket binding.
|
||||
return
|
||||
}
|
||||
if turnState == "" && stateStore != nil && sessionHash != "" {
|
||||
if savedTurnState, ok := stateStore.GetSessionTurnState(groupID, sessionHash); ok {
|
||||
turnState = savedTurnState
|
||||
}
|
||||
}
|
||||
|
||||
preferredConnID = ""
|
||||
if stateStore != nil && payload.previousResponseID != "" {
|
||||
if connID, ok := stateStore.GetResponseConn(payload.previousResponseID); ok {
|
||||
preferredConnID = connID
|
||||
}
|
||||
}
|
||||
|
||||
storeDisabled = s.isOpenAIWSStoreDisabledInRequestRaw(payload.payloadRaw, account)
|
||||
if stateStore != nil && storeDisabled && payload.previousResponseID == "" && sessionHash != "" {
|
||||
if connID, ok := stateStore.GetSessionConn(groupID, sessionHash); ok {
|
||||
preferredConnID = connID
|
||||
@@ -552,7 +558,7 @@ func (s *OpenAIGatewayService) ProxyResponsesWebSocketFromClient(
|
||||
}
|
||||
refreshIngressRouteState(firstPayload)
|
||||
|
||||
if forceHTTPBridge || s.shouldBridgeOpenAIWSHTTP(account, firstPayload.payloadBytes, firstPayload.previousResponseID) {
|
||||
if useHTTPBridge {
|
||||
logOpenAIWSModeInfo(
|
||||
"ingress_ws_http_bridge_start account_id=%d account_type=%s payload_bytes=%d threshold_bytes=%d has_session_hash=%v store_disabled=%v",
|
||||
account.ID,
|
||||
@@ -725,10 +731,9 @@ func (s *OpenAIGatewayService) ProxyResponsesWebSocketFromClient(
|
||||
bridgeAccountFailoverInputExists = true
|
||||
}
|
||||
if bridgeTurnState := strings.TrimSpace(result.ResponseHeaders.Get(openAIWSTurnStateHeader)); bridgeTurnState != "" {
|
||||
// Follow-up turns on this bridge retain their own upstream state;
|
||||
// publishing it by session hash would leak it to independent bridges.
|
||||
turnState = bridgeTurnState
|
||||
if stateStore != nil && sessionHash != "" {
|
||||
stateStore.BindSessionTurnState(groupID, sessionHash, bridgeTurnState, s.openAIWSSessionStickyTTL())
|
||||
}
|
||||
}
|
||||
responseID := strings.TrimSpace(result.RequestID)
|
||||
if responseID != "" && stateStore != nil {
|
||||
|
||||
@@ -0,0 +1,249 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/tlsfingerprint"
|
||||
coderws "github.com/coder/websocket"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
"github.com/tidwall/gjson"
|
||||
)
|
||||
|
||||
type httpBridgeIsolationBody struct {
|
||||
ctx context.Context
|
||||
prefix *bytes.Reader
|
||||
suffix *bytes.Reader
|
||||
release <-chan struct{}
|
||||
closed chan struct{}
|
||||
once sync.Once
|
||||
}
|
||||
|
||||
func (b *httpBridgeIsolationBody) Read(p []byte) (int, error) {
|
||||
if b.prefix.Len() > 0 {
|
||||
return b.prefix.Read(p)
|
||||
}
|
||||
select {
|
||||
case <-b.release:
|
||||
return b.suffix.Read(p)
|
||||
case <-b.closed:
|
||||
return 0, io.EOF
|
||||
case <-b.ctx.Done():
|
||||
return 0, b.ctx.Err()
|
||||
}
|
||||
}
|
||||
|
||||
func (b *httpBridgeIsolationBody) Close() error {
|
||||
b.once.Do(func() { close(b.closed) })
|
||||
return nil
|
||||
}
|
||||
|
||||
type httpBridgeIsolationRequest struct {
|
||||
input []string
|
||||
state string
|
||||
}
|
||||
|
||||
type httpBridgeIsolationUpstream struct {
|
||||
ctx context.Context
|
||||
mu sync.Mutex
|
||||
requests []httpBridgeIsolationRequest
|
||||
firstRelease chan struct{}
|
||||
}
|
||||
|
||||
func (u *httpBridgeIsolationUpstream) Do(req *http.Request, _ string, _ int64, _ int) (*http.Response, error) {
|
||||
body, err := io.ReadAll(req.Body)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
_ = req.Body.Close()
|
||||
input := gjson.GetBytes(body, "input")
|
||||
var texts []string
|
||||
if input.Type == gjson.String {
|
||||
texts = append(texts, input.String())
|
||||
} else {
|
||||
for _, item := range input.Array() {
|
||||
if item.Type == gjson.String {
|
||||
texts = append(texts, item.String())
|
||||
continue
|
||||
}
|
||||
if item.Get("type").String() == "function_call" {
|
||||
texts = append(texts, item.Get("call_id").String())
|
||||
continue
|
||||
}
|
||||
if item.Get("type").String() == "function_call_output" {
|
||||
texts = append(texts, item.Get("output").String())
|
||||
continue
|
||||
}
|
||||
content := item.Get("content")
|
||||
if content.Type == gjson.String {
|
||||
texts = append(texts, content.String())
|
||||
} else {
|
||||
for _, part := range content.Array() {
|
||||
if text := part.Get("text"); text.Exists() {
|
||||
texts = append(texts, text.String())
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
if len(texts) == 0 {
|
||||
return nil, fmt.Errorf("missing bridge input")
|
||||
}
|
||||
u.mu.Lock()
|
||||
u.requests = append(u.requests, httpBridgeIsolationRequest{input: texts, state: req.Header.Get(openAIWSTurnStateHeader)})
|
||||
u.mu.Unlock()
|
||||
|
||||
id := texts[0]
|
||||
turn := 1
|
||||
if input.IsArray() {
|
||||
turn = 2
|
||||
}
|
||||
responseID := fmt.Sprintf("resp_%s_%d", id, turn)
|
||||
prefix := fmt.Sprintf("data: {\"type\":\"response.created\",\"response\":{\"id\":%q,\"status\":\"in_progress\"}}\n\n", responseID)
|
||||
if turn == 1 {
|
||||
prefix += "data: {\"type\":\"response.output_item.added\",\"output_index\":0,\"item\":{\"id\":\"rs_test\",\"type\":\"reasoning\",\"encrypted_content\":\"opaque\"}}\n\n"
|
||||
}
|
||||
output := "[]"
|
||||
if turn == 1 {
|
||||
output = fmt.Sprintf(`[{"type":"function_call","id":%q,"call_id":%q,"name":"lookup","arguments":"{}","status":"completed"}]`, "fc_"+id, "call_"+id)
|
||||
}
|
||||
suffix := fmt.Sprintf("data: {\"type\":\"response.completed\",\"response\":{\"id\":%q,\"status\":\"completed\",\"output\":%s,\"usage\":{\"input_tokens\":1,\"output_tokens\":1}}}\n\n", responseID, output)
|
||||
headers := http.Header{"Content-Type": []string{"text/event-stream"}}
|
||||
headers.Set(openAIWSTurnStateHeader, "turn-"+id)
|
||||
responseBody := io.NopCloser(strings.NewReader(prefix + suffix))
|
||||
if turn == 1 {
|
||||
responseBody = &httpBridgeIsolationBody{ctx: u.ctx, prefix: bytes.NewReader([]byte(prefix)), suffix: bytes.NewReader([]byte(suffix)), release: u.firstRelease, closed: make(chan struct{})}
|
||||
}
|
||||
return &http.Response{StatusCode: http.StatusOK, Header: headers, Body: responseBody}, nil
|
||||
}
|
||||
|
||||
func (u *httpBridgeIsolationUpstream) DoWithTLS(req *http.Request, proxyURL string, accountID int64, concurrency int, _ *tlsfingerprint.Profile) (*http.Response, error) {
|
||||
return u.Do(req, proxyURL, accountID, concurrency)
|
||||
}
|
||||
|
||||
func TestOpenAIWSHTTPBridgeSessionIsolationAcrossSameSessionHash(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
||||
defer cancel()
|
||||
cfg := &config.Config{}
|
||||
cfg.Gateway.OpenAIWS.Enabled = true
|
||||
cfg.Gateway.OpenAIWS.OAuthEnabled = true
|
||||
cfg.Gateway.OpenAIWS.ResponsesWebsocketsV2 = true
|
||||
cfg.Gateway.OpenAIWS.ModeRouterV2Enabled = true
|
||||
cfg.Gateway.OpenAIWS.IngressModeDefault = OpenAIWSIngressModeHTTPBridge
|
||||
cfg.Gateway.OpenAIWS.ReadTimeoutSeconds = 5
|
||||
cfg.Gateway.OpenAIWS.WriteTimeoutSeconds = 5
|
||||
upstream := &httpBridgeIsolationUpstream{ctx: ctx, firstRelease: make(chan struct{})}
|
||||
stateStore := NewOpenAIWSStateStore(nil)
|
||||
svc := &OpenAIGatewayService{cfg: cfg, httpUpstream: upstream, openaiWSResolver: NewOpenAIWSProtocolResolver(cfg), openaiWSStateStore: stateStore}
|
||||
account := &Account{ID: 1, Platform: PlatformOpenAI, Type: AccountTypeOAuth,
|
||||
Credentials: map[string]any{"access_token": "test-token"},
|
||||
Extra: map[string]any{"openai_oauth_responses_websockets_v2_mode": OpenAIWSIngressModeHTTPBridge},
|
||||
Concurrency: 2, Status: StatusActive, Schedulable: true}
|
||||
groupID := int64(7)
|
||||
newContext := func(r *http.Request) *gin.Context {
|
||||
c, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
c.Request = r.Clone(ctx)
|
||||
c.Request.Header.Set("session-id", "shared-session")
|
||||
c.Set("api_key", &APIKey{ID: 11, GroupID: &groupID})
|
||||
return c
|
||||
}
|
||||
seedHash := svc.GenerateSessionHash(newContext(httptest.NewRequest(http.MethodGet, "/v1/responses", nil)), nil)
|
||||
stateStore.BindSessionTurnState(groupID, seedHash, "sentinel-turn-state", time.Hour)
|
||||
stateStore.BindSessionConn(groupID, seedHash, "sentinel-native-conn", time.Hour)
|
||||
|
||||
serverResults := make(chan error, 2)
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
conn, err := coderws.Accept(w, r, nil)
|
||||
if err != nil {
|
||||
serverResults <- err
|
||||
return
|
||||
}
|
||||
defer func() { _ = conn.CloseNow() }()
|
||||
_, firstMessage, err := conn.Read(ctx)
|
||||
if err != nil {
|
||||
serverResults <- err
|
||||
return
|
||||
}
|
||||
c := newContext(r)
|
||||
// Match the handler-owned registration and nested forwarding call.
|
||||
preemptCtx, cleanup, _ := svc.BeginOpenAIWSIngressSessionPreemption(ctx, c, account, firstMessage)
|
||||
defer cleanup()
|
||||
hooks := &OpenAIWSIngressHooks{ClientLifecycleContext: ctx}
|
||||
serverResults <- svc.ProxyResponsesWebSocketFromClient(preemptCtx, c, conn, account, "test-token", firstMessage, hooks)
|
||||
}))
|
||||
var releaseOnce sync.Once
|
||||
release := func() { releaseOnce.Do(func() { close(upstream.firstRelease) }) }
|
||||
defer func() {
|
||||
release()
|
||||
cancel()
|
||||
server.Close()
|
||||
}()
|
||||
|
||||
readEvent := func(conn *coderws.Conn, wantType string) []byte {
|
||||
t.Helper()
|
||||
_, event, err := conn.Read(ctx)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, wantType, gjson.GetBytes(event, "type").String(), "event: %s", event)
|
||||
return event
|
||||
}
|
||||
clients := make([]*coderws.Conn, 0, 2)
|
||||
for _, id := range []string{"alpha", "beta"} {
|
||||
conn, _, err := coderws.Dial(ctx, "ws"+strings.TrimPrefix(server.URL, "http"), nil)
|
||||
require.NoError(t, err)
|
||||
defer func() { _ = conn.CloseNow() }()
|
||||
clients = append(clients, conn)
|
||||
payload := fmt.Sprintf(`{"type":"response.create","model":"gpt-5","prompt_cache_key":%q,"client_metadata":{"thread_id":%q},"input":%q}`, id+"-cache", id+"-thread", id)
|
||||
require.NoError(t, conn.Write(ctx, coderws.MessageText, []byte(payload)))
|
||||
readEvent(conn, "response.created")
|
||||
readEvent(conn, "response.output_item.added")
|
||||
}
|
||||
// Both HTTP streams are now in flight. Neither may replace the other.
|
||||
release()
|
||||
for i, id := range []string{"alpha", "beta"} {
|
||||
event := readEvent(clients[i], "response.completed")
|
||||
require.Equal(t, "resp_"+id+"_1", gjson.GetBytes(event, "response.id").String())
|
||||
}
|
||||
for i, id := range []string{"alpha", "beta"} {
|
||||
payload := fmt.Sprintf(`{"type":"response.create","model":"gpt-5","previous_response_id":%q,"input":[{"type":"function_call_output","call_id":%q,"output":%q}]}`, "resp_"+id+"_1", "call_"+id, id+"-result")
|
||||
require.NoError(t, clients[i].Write(ctx, coderws.MessageText, []byte(payload)))
|
||||
readEvent(clients[i], "response.created")
|
||||
event := readEvent(clients[i], "response.completed")
|
||||
require.Equal(t, "resp_"+id+"_2", gjson.GetBytes(event, "response.id").String())
|
||||
require.NoError(t, clients[i].Close(coderws.StatusNormalClosure, "done"))
|
||||
}
|
||||
for range clients {
|
||||
select {
|
||||
case err := <-serverResults:
|
||||
require.NoError(t, err)
|
||||
case <-ctx.Done():
|
||||
t.Fatal(ctx.Err())
|
||||
}
|
||||
}
|
||||
|
||||
upstream.mu.Lock()
|
||||
requests := append([]httpBridgeIsolationRequest(nil), upstream.requests...)
|
||||
upstream.mu.Unlock()
|
||||
require.ElementsMatch(t, []httpBridgeIsolationRequest{
|
||||
{input: []string{"alpha"}},
|
||||
{input: []string{"beta"}},
|
||||
{input: []string{"alpha", "call_alpha", "alpha-result"}, state: "turn-alpha"},
|
||||
{input: []string{"beta", "call_beta", "beta-result"}, state: "turn-beta"},
|
||||
}, requests)
|
||||
gotState, ok := stateStore.GetSessionTurnState(groupID, seedHash)
|
||||
require.True(t, ok)
|
||||
require.Equal(t, "sentinel-turn-state", gotState)
|
||||
gotConn, ok := stateStore.GetSessionConn(groupID, seedHash)
|
||||
require.True(t, ok)
|
||||
require.Equal(t, "sentinel-native-conn", gotConn)
|
||||
}
|
||||
@@ -56,9 +56,14 @@ 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
|
||||
if s != nil && s.cfg != nil && s.cfg.Gateway.OpenAIWS.ModeRouterV2Enabled && account != nil {
|
||||
switch account.ResolveOpenAIResponsesWebSocketV2Mode(s.cfg.Gateway.OpenAIWS.IngressModeDefault) {
|
||||
case OpenAIWSIngressModePassthrough, OpenAIWSIngressModeHTTPBridge:
|
||||
// These modes own their upstream transport and replay state per
|
||||
// inbound connection. Sharing a sticky/cache identity does not mean
|
||||
// a new connection should cancel an independent in-flight request.
|
||||
return ctx, func() {}, false
|
||||
}
|
||||
}
|
||||
|
||||
preemptSessionHash := ""
|
||||
|
||||
@@ -224,6 +224,59 @@ func TestOpenAIWSSessionPreemptRemoteClaimAndStaleReleaseAreAtomic(t *testing.T)
|
||||
require.Equal(t, "owner-b", current, "stale cleanup must preserve the replacement owner")
|
||||
}
|
||||
|
||||
func TestOpenAIWSHTTPBridgeSessionPreemptionEligibility(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
tests := []struct {
|
||||
name string
|
||||
routerEnabled bool
|
||||
defaultMode string
|
||||
accountMode string
|
||||
wantArmed bool
|
||||
}{
|
||||
{name: "explicit HTTP bridge", routerEnabled: true, defaultMode: OpenAIWSIngressModeCtxPool, accountMode: OpenAIWSIngressModeHTTPBridge},
|
||||
{name: "default HTTP bridge", routerEnabled: true, defaultMode: OpenAIWSIngressModeHTTPBridge},
|
||||
{name: "ctx pool overrides bridge default", routerEnabled: true, defaultMode: OpenAIWSIngressModeHTTPBridge, accountMode: OpenAIWSIngressModeCtxPool, wantArmed: true},
|
||||
{name: "disabled router retains legacy preemption", defaultMode: OpenAIWSIngressModeHTTPBridge, accountMode: OpenAIWSIngressModeHTTPBridge, wantArmed: true},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
cfg := &config.Config{}
|
||||
cfg.Gateway.OpenAIWS.ModeRouterV2Enabled = tt.routerEnabled
|
||||
cfg.Gateway.OpenAIWS.IngressModeDefault = tt.defaultMode
|
||||
cache := &openAIWSSessionPreemptCacheStub{}
|
||||
svc := &OpenAIGatewayService{cfg: cfg, cache: cache}
|
||||
groupID := int64(7)
|
||||
newContext := func() *gin.Context {
|
||||
c, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
c.Request = httptest.NewRequest(http.MethodGet, "/v1/responses", nil)
|
||||
c.Request.Header.Set("session-id", "shared-session")
|
||||
c.Set("api_key", &APIKey{ID: 11, GroupID: &groupID})
|
||||
return c
|
||||
}
|
||||
account := &Account{ID: 1, Platform: PlatformOpenAI, Type: AccountTypeOAuth,
|
||||
Extra: map[string]any{"openai_oauth_responses_websockets_v2_mode": tt.accountMode}}
|
||||
firstMessage := []byte(`{"type":"response.create","prompt_cache_key":"cache-a","input":"first request"}`)
|
||||
secondMessage := []byte(`{"type":"response.create","prompt_cache_key":"cache-b","input":"second request"}`)
|
||||
first, cleanupFirst, firstArmed := svc.BeginOpenAIWSIngressSessionPreemption(context.Background(), newContext(), account, firstMessage)
|
||||
defer cleanupFirst()
|
||||
second, cleanupSecond, secondArmed := svc.BeginOpenAIWSIngressSessionPreemption(context.Background(), newContext(), account, secondMessage)
|
||||
defer cleanupSecond()
|
||||
require.Equal(t, tt.wantArmed, firstArmed)
|
||||
require.Equal(t, tt.wantArmed, secondArmed)
|
||||
require.NoError(t, second.Err())
|
||||
if tt.wantArmed {
|
||||
require.True(t, IsOpenAIWSSessionPreemptedError(context.Cause(first)))
|
||||
} else {
|
||||
require.NoError(t, first.Err(), "independent HTTP bridges must not cancel each other")
|
||||
cache.mu.Lock()
|
||||
ownerCount := len(cache.owners)
|
||||
cache.mu.Unlock()
|
||||
require.Zero(t, ownerCount, "HTTP bridges must not claim a distributed preemption owner")
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestNewOpenAIWSSessionPreemptKeyRequiresFullIsolationScope(t *testing.T) {
|
||||
_, ok := newOpenAIWSSessionPreemptKey(0, 11, "sess")
|
||||
require.False(t, ok)
|
||||
|
||||
Reference in New Issue
Block a user