diff --git a/backend/internal/service/openai_ws_forwarder_ingress.go b/backend/internal/service/openai_ws_forwarder_ingress.go index b6893d6c8d..f24dc25697 100644 --- a/backend/internal/service/openai_ws_forwarder_ingress.go +++ b/backend/internal/service/openai_ws_forwarder_ingress.go @@ -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 { diff --git a/backend/internal/service/openai_ws_http_bridge_isolation_test.go b/backend/internal/service/openai_ws_http_bridge_isolation_test.go new file mode 100644 index 0000000000..53711dccd3 --- /dev/null +++ b/backend/internal/service/openai_ws_http_bridge_isolation_test.go @@ -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) +} diff --git a/backend/internal/service/openai_ws_session_preemption.go b/backend/internal/service/openai_ws_session_preemption.go index 3047de6998..95fa49e7f7 100644 --- a/backend/internal/service/openai_ws_session_preemption.go +++ b/backend/internal/service/openai_ws_session_preemption.go @@ -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 := "" diff --git a/backend/internal/service/openai_ws_session_preemption_test.go b/backend/internal/service/openai_ws_session_preemption_test.go index 0e0c6f5622..d00adf63b4 100644 --- a/backend/internal/service/openai_ws_session_preemption_test.go +++ b/backend/internal/service/openai_ws_session_preemption_test.go @@ -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)