diff --git a/backend/internal/service/openai_gateway_passthrough.go b/backend/internal/service/openai_gateway_passthrough.go index 4c9de4b00e..82f63cd9fa 100644 --- a/backend/internal/service/openai_gateway_passthrough.go +++ b/backend/internal/service/openai_gateway_passthrough.go @@ -971,6 +971,7 @@ func (s *OpenAIGatewayService) handleStreamingResponsePassthrough( scanBuf := getSSEScannerBuf64K() scanner.Buffer(scanBuf[:0], maxLineSize) defer putSSEScannerBuf64K(scanBuf) + documentScanner := newOpenAISSEJSONDocumentScanner(scanner) needModelReplace := strings.TrimSpace(originalModel) != "" && strings.TrimSpace(mappedModel) != "" && strings.TrimSpace(originalModel) != strings.TrimSpace(mappedModel) resultWithUsage := func() *openaiStreamingResultPassthrough { @@ -983,8 +984,8 @@ func (s *OpenAIGatewayService) handleStreamingResponsePassthrough( } } - for scanner.Scan() { - line := scanner.Text() + for documentScanner.Scan() { + line := documentScanner.Text() lineStartsClientOutput := false forceFlushFailedEvent := false if data, ok := extractOpenAISSEDataLine(line); ok { @@ -1104,7 +1105,7 @@ func (s *OpenAIGatewayService) handleStreamingResponsePassthrough( } } } - if err := scanner.Err(); err != nil { + if err := documentScanner.Err(); err != nil { if sawTerminalEvent && !sawFailedEvent { return resultWithUsage(), nil } diff --git a/backend/internal/service/openai_gateway_response_handling.go b/backend/internal/service/openai_gateway_response_handling.go index 1d70245eb0..ea65814287 100644 --- a/backend/internal/service/openai_gateway_response_handling.go +++ b/backend/internal/service/openai_gateway_response_handling.go @@ -79,6 +79,7 @@ func (s *OpenAIGatewayService) handleStreamingResponse(ctx context.Context, resp } scanBuf := getSSEScannerBuf64K() scanner.Buffer(scanBuf[:0], maxLineSize) + documentScanner := newOpenAISSEJSONDocumentScanner(scanner) streamInterval := time.Duration(0) if s.cfg != nil && s.cfg.Gateway.StreamDataIntervalTimeout > 0 { @@ -393,13 +394,13 @@ func (s *OpenAIGatewayService) handleStreamingResponse(ctx context.Context, resp // 无超时/无 keepalive 的常见路径走同步扫描,减少 goroutine 与 channel 开销。 if streamInterval <= 0 && keepaliveInterval <= 0 { defer putSSEScannerBuf64K(scanBuf) - for scanner.Scan() { - processSSELine(scanner.Text(), true) + for documentScanner.Scan() { + processSSELine(documentScanner.Text(), true) if streamEarlyErr != nil { return resultWithUsage(), streamEarlyErr } } - if result, err, done := handleScanErr(scanner.Err()); done { + if result, err, done := handleScanErr(documentScanner.Err()); done { return result, err } return finalizeStream() @@ -425,13 +426,13 @@ func (s *OpenAIGatewayService) handleStreamingResponse(ctx context.Context, resp go func(scanBuf *sseScannerBuf64K) { defer putSSEScannerBuf64K(scanBuf) defer close(events) - for scanner.Scan() { + for documentScanner.Scan() { atomic.StoreInt64(&lastReadAt, time.Now().UnixNano()) - if !sendEvent(scanEvent{line: scanner.Text()}) { + if !sendEvent(scanEvent{line: documentScanner.Text()}) { return } } - if err := scanner.Err(); err != nil { + if err := documentScanner.Err(); err != nil { _ = sendEvent(scanEvent{err: err}) } }(scanBuf) diff --git a/backend/internal/service/openai_sse_concatenated_json_test.go b/backend/internal/service/openai_sse_concatenated_json_test.go new file mode 100644 index 0000000000..9b06ef9a39 --- /dev/null +++ b/backend/internal/service/openai_sse_concatenated_json_test.go @@ -0,0 +1,271 @@ +package service + +import ( + "bufio" + "context" + "encoding/json" + "io" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" + + "github.com/Wei-Shaw/sub2api/internal/config" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" +) + +func TestOpenAIStreamingRepairsConcatenatedJSONDocumentsInSingleDataLine(t *testing.T) { + testOpenAIStreamingRepairsConcatenatedJSONDocuments(t, false, 0) +} + +func TestOpenAIStreamingAsyncScannerRepairsConcatenatedJSONDocumentsInSingleDataLine(t *testing.T) { + testOpenAIStreamingRepairsConcatenatedJSONDocuments(t, false, 30) +} + +func TestOpenAIStreamingPassthroughRepairsConcatenatedJSONDocumentsInSingleDataLine(t *testing.T) { + testOpenAIStreamingRepairsConcatenatedJSONDocuments(t, true, 0) +} + +func TestOpenAIWSv2StreamingRepairsConcatenatedJSONDocumentsInSingleMessage(t *testing.T) { + gin.SetMode(gin.TestMode) + + largeInProgress, outputItemAdded, completed := openAIConcatenatedJSONTestEvents(t) + captureConn := &openAIWSCaptureConn{events: [][]byte{ + []byte(largeInProgress + outputItemAdded), + []byte(completed), + }} + + cfg := &config.Config{} + cfg.Security.URLAllowlist.Enabled = false + cfg.Gateway.OpenAIWS.Enabled = true + cfg.Gateway.OpenAIWS.APIKeyEnabled = true + cfg.Gateway.OpenAIWS.ResponsesWebsocketsV2 = true + cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 1 + cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 1 + cfg.Gateway.OpenAIWS.QueueLimitPerConn = 8 + cfg.Gateway.OpenAIWS.DialTimeoutSeconds = 3 + cfg.Gateway.OpenAIWS.ReadTimeoutSeconds = 5 + cfg.Gateway.OpenAIWS.WriteTimeoutSeconds = 3 + + pool := newOpenAIWSConnPool(cfg) + pool.setClientDialerForTest(&openAIWSCaptureDialer{conn: captureConn}) + svc := &OpenAIGatewayService{ + cfg: cfg, + cache: &stubGatewayCache{}, + httpUpstream: &httpUpstreamRecorder{}, + openaiWSResolver: NewOpenAIWSProtocolResolver(cfg), + openaiWSPool: pool, + toolCorrector: NewCodexToolCorrector(), + } + account := &Account{ + ID: 2, + Name: "ws-test", + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + Status: StatusActive, + Schedulable: true, + Concurrency: 1, + Credentials: map[string]any{"api_key": "sk-test"}, + Extra: map[string]any{"responses_websockets_v2_enabled": true}, + } + + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil) + groupID := int64(1) + c.Set("api_key", &APIKey{GroupID: &groupID}) + + result, err := svc.Forward(context.Background(), c, account, []byte(`{"model":"gpt-5.6-sol","stream":true,"input":"hello"}`)) + require.NoError(t, err) + require.NotNil(t, result) + require.Equal(t, 7, result.Usage.InputTokens) + require.Equal(t, 9, result.Usage.OutputTokens) + require.Nil(t, result.FirstTokenMs) + assertOpenAISSEFrames(t, recorder.Body.String(), []string{ + "response.in_progress", + "response.output_item.added", + "response.completed", + }) +} + +func TestSplitOpenAIConcatenatedJSONDocumentsRejectsPayloadOverRepairLimit(t *testing.T) { + first := `{"type":"response.in_progress","padding":"` + strings.Repeat("x", 16*1024*1024) + `"}` + second := `{"type":"response.completed"}` + payload := first + second + + documents, repaired := splitOpenAIConcatenatedJSONDocuments([]byte(payload)) + require.False(t, repaired) + require.Nil(t, documents) + + line := "data: " + payload + scanner := bufio.NewScanner(strings.NewReader(line)) + scanner.Buffer(make([]byte, 1024), len(line)+1) + documentScanner := newOpenAISSEJSONDocumentScanner(scanner) + require.True(t, documentScanner.Scan()) + require.Equal(t, line, documentScanner.Text()) + require.False(t, documentScanner.Scan()) + require.NoError(t, documentScanner.Err()) +} + +func TestOpenAIWSv2StreamingBreaksConnectionWhenTerminalHasTrailingDocument(t *testing.T) { + gin.SetMode(gin.TestMode) + completed := `{"type":"response.completed","response":{"id":"resp_terminal_tail","usage":{"input_tokens":2,"output_tokens":1}}}` + tail := `{"type":"error","error":{"type":"upstream_error","message":"tail"}}` + captureConn := &openAIWSCaptureConn{events: [][]byte{[]byte(completed + tail)}} + + cfg := &config.Config{} + cfg.Security.URLAllowlist.Enabled = false + cfg.Gateway.OpenAIWS.Enabled = true + cfg.Gateway.OpenAIWS.APIKeyEnabled = true + cfg.Gateway.OpenAIWS.ResponsesWebsocketsV2 = true + cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 1 + cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 1 + cfg.Gateway.OpenAIWS.QueueLimitPerConn = 8 + cfg.Gateway.OpenAIWS.DialTimeoutSeconds = 3 + cfg.Gateway.OpenAIWS.ReadTimeoutSeconds = 5 + cfg.Gateway.OpenAIWS.WriteTimeoutSeconds = 3 + + pool := newOpenAIWSConnPool(cfg) + pool.setClientDialerForTest(&openAIWSCaptureDialer{conn: captureConn}) + svc := &OpenAIGatewayService{ + cfg: cfg, + cache: &stubGatewayCache{}, + httpUpstream: &httpUpstreamRecorder{}, + openaiWSResolver: NewOpenAIWSProtocolResolver(cfg), + openaiWSPool: pool, + toolCorrector: NewCodexToolCorrector(), + } + account := &Account{ + ID: 3, + Name: "ws-terminal-tail", + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + Status: StatusActive, + Schedulable: true, + Concurrency: 1, + Credentials: map[string]any{"api_key": "sk-test"}, + Extra: map[string]any{"responses_websockets_v2_enabled": true}, + } + + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil) + groupID := int64(1) + c.Set("api_key", &APIKey{GroupID: &groupID}) + + result, err := svc.Forward(context.Background(), c, account, []byte(`{"model":"gpt-5.6-sol","stream":true,"input":"hello"}`)) + require.NoError(t, err) + require.NotNil(t, result) + require.True(t, captureConn.closed, "a WS message with data after a terminal event must not return to the pool") + assertOpenAISSEFrames(t, recorder.Body.String(), []string{"response.completed"}) +} + +func testOpenAIStreamingRepairsConcatenatedJSONDocuments(t *testing.T, passthrough bool, streamDataIntervalTimeout int) { + t.Helper() + gin.SetMode(gin.TestMode) + + largeInProgress, outputItemAdded, completed := openAIConcatenatedJSONTestEvents(t) + + upstreamBody := strings.Join([]string{ + "event: response.in_progress", + "data: " + largeInProgress + outputItemAdded, + "", + "event: response.completed", + "data: " + completed, + "", + }, "\n") + resp := &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"text/event-stream"}}, + Body: io.NopCloser(strings.NewReader(upstreamBody)), + } + + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil) + svc := &OpenAIGatewayService{ + cfg: &config.Config{Gateway: config.GatewayConfig{ + MaxLineSize: defaultMaxLineSize, + StreamDataIntervalTimeout: streamDataIntervalTimeout, + }}, + toolCorrector: NewCodexToolCorrector(), + } + account := &Account{ID: 1, Name: "test", Platform: PlatformOpenAI} + + var usage *OpenAIUsage + var err error + if passthrough { + result, forwardErr := svc.handleStreamingResponsePassthrough(c.Request.Context(), resp, c, account, time.Now(), "gpt-5.6-sol", "gpt-5.6-sol") + err = forwardErr + if result != nil { + usage = result.usage + } + } else { + result, forwardErr := svc.handleStreamingResponse(c.Request.Context(), resp, c, account, time.Now(), "gpt-5.6-sol", "gpt-5.6-sol") + err = forwardErr + if result != nil { + usage = result.usage + } + } + require.NoError(t, err) + require.NotNil(t, usage) + require.Equal(t, 7, usage.InputTokens) + require.Equal(t, 9, usage.OutputTokens) + + assertOpenAISSEFrames(t, recorder.Body.String(), []string{ + "response.in_progress", + "response.output_item.added", + "response.completed", + }) +} + +func assertOpenAISSEFrames(t *testing.T, body string, expectedTypes []string) { + t.Helper() + var parser openAICompatSSEFrameParser + var eventTypes []string + for _, line := range strings.Split(body, "\n") { + frame, ok := parser.AddLine(strings.TrimSuffix(line, "\r")) + if !ok { + continue + } + require.True(t, json.Valid([]byte(frame.Data)), "each downstream SSE frame must contain exactly one JSON document") + var event struct { + Type string `json:"type"` + } + require.NoError(t, json.Unmarshal([]byte(frame.Data), &event)) + if frame.EventType != "" { + require.Equal(t, event.Type, frame.EventType) + } + eventTypes = append(eventTypes, event.Type) + } + if frame, ok := parser.Finish(); ok { + require.True(t, json.Valid([]byte(frame.Data))) + var event struct { + Type string `json:"type"` + } + require.NoError(t, json.Unmarshal([]byte(frame.Data), &event)) + eventTypes = append(eventTypes, event.Type) + } + require.Equal(t, expectedTypes, eventTypes) +} + +func openAIConcatenatedJSONTestEvents(t *testing.T) (string, string, string) { + t.Helper() + const javascriptErrorPosition = 68106 + prefix := `{"type":"response.in_progress","response":{"id":"resp_large","status":"in_progress","instructions":"` + suffix := `"},"sequence_number":1}` + require.Less(t, len(prefix)+len(suffix), javascriptErrorPosition) + largeInProgress := prefix + strings.Repeat("x", javascriptErrorPosition-len(prefix)-len(suffix)) + suffix + outputItemAdded := `{"type":"response.output_item.added","output_index":0,"item":{"id":"msg_1","type":"message","role":"assistant","status":"in_progress","content":[]},"sequence_number":2}` + completed := `{"type":"response.completed","response":{"id":"resp_large","status":"completed","output":[],"usage":{"input_tokens":7,"output_tokens":9}},"sequence_number":3}` + require.Len(t, largeInProgress, javascriptErrorPosition) + require.True(t, json.Valid([]byte(largeInProgress))) + var decoded any + err := json.Unmarshal([]byte(largeInProgress+outputItemAdded), &decoded) + var syntaxErr *json.SyntaxError + require.ErrorAs(t, err, &syntaxErr) + require.Equal(t, int64(javascriptErrorPosition+1), syntaxErr.Offset) + return largeInProgress, outputItemAdded, completed +} diff --git a/backend/internal/service/openai_sse_json_documents.go b/backend/internal/service/openai_sse_json_documents.go new file mode 100644 index 0000000000..3f6fa28fa8 --- /dev/null +++ b/backend/internal/service/openai_sse_json_documents.go @@ -0,0 +1,115 @@ +package service + +import ( + "bufio" + "bytes" + "encoding/json" + "io" + "strings" +) + +const ( + maxOpenAIConcatenatedJSONDocuments = 16 + maxOpenAIConcatenatedJSONBytes = 16 * 1024 * 1024 +) + +// splitOpenAIConcatenatedJSONDocuments recognizes the narrow corruption shape +// produced when multiple complete Responses events arrive in one transport +// message. Other malformed payloads are left untouched for normal error paths. +func splitOpenAIConcatenatedJSONDocuments(payload []byte) ([][]byte, bool) { + payload = bytes.TrimSpace(payload) + if len(payload) == 0 || len(payload) > maxOpenAIConcatenatedJSONBytes || json.Valid(payload) { + return nil, false + } + + decoder := json.NewDecoder(bytes.NewReader(payload)) + documents := make([][]byte, 0, 2) + for { + var raw json.RawMessage + err := decoder.Decode(&raw) + if err != nil { + if err == io.EOF && len(documents) > 1 { + return documents, true + } + return nil, false + } + raw = bytes.TrimSpace(raw) + var envelope struct { + Type string `json:"type"` + } + if err := json.Unmarshal(raw, &envelope); err != nil { + return nil, false + } + eventType := strings.TrimSpace(envelope.Type) + if eventType == "" || strings.ContainsAny(eventType, "\r\n") { + return nil, false + } + if len(documents) == maxOpenAIConcatenatedJSONDocuments { + return nil, false + } + documents = append(documents, raw) + } +} + +type openAISSEJSONDocumentScanner struct { + scanner *bufio.Scanner + pending []string + current string +} + +func newOpenAISSEJSONDocumentScanner(scanner *bufio.Scanner) *openAISSEJSONDocumentScanner { + return &openAISSEJSONDocumentScanner{scanner: scanner} +} + +func (s *openAISSEJSONDocumentScanner) Scan() bool { + if len(s.pending) > 0 { + s.current = s.pending[0] + s.pending = s.pending[1:] + return true + } + if s.scanner == nil || !s.scanner.Scan() { + return false + } + + line := s.scanner.Text() + data, ok := extractOpenAISSEDataLine(line) + if !ok { + s.current = line + return true + } + if len(data) > maxOpenAIConcatenatedJSONBytes { + s.current = line + return true + } + documents, repaired := splitOpenAIConcatenatedJSONDocuments([]byte(data)) + if !repaired { + s.current = line + return true + } + + expanded := make([]string, 0, len(documents)*3) + for i, document := range documents { + if i > 0 { + var envelope struct { + Type string `json:"type"` + } + _ = json.Unmarshal(document, &envelope) + expanded = append(expanded, "event: "+strings.TrimSpace(envelope.Type)) + } + expanded = append(expanded, "data: "+string(document), "") + } + s.current = expanded[0] + s.pending = expanded[1:] + return true +} + +func (s *openAISSEJSONDocumentScanner) Text() string { + return s.current +} + +func (s *openAISSEJSONDocumentScanner) Err() error { + if s.scanner == nil { + return nil + } + return s.scanner.Err() +} diff --git a/backend/internal/service/openai_ws_forwarder_v2.go b/backend/internal/service/openai_ws_forwarder_v2.go index 90151be93a..9054e60502 100644 --- a/backend/internal/service/openai_ws_forwarder_v2.go +++ b/backend/internal/service/openai_ws_forwarder_v2.go @@ -430,9 +430,30 @@ func (s *OpenAIGatewayService) forwardOpenAIWSV2( } readTimeout := s.openAIWSReadTimeout() + var pendingJSONDocuments [][]byte for { - message, readErr := lease.ReadMessageWithContextTimeout(ctx, readTimeout) + var message []byte + var readErr error + if len(pendingJSONDocuments) > 0 { + message = pendingJSONDocuments[0] + pendingJSONDocuments = pendingJSONDocuments[1:] + } else { + message, readErr = lease.ReadMessageWithContextTimeout(ctx, readTimeout) + if readErr == nil { + if documents, repaired := splitOpenAIConcatenatedJSONDocuments(message); repaired { + logOpenAIWSModeInfo( + "concatenated_json_repaired account_id=%d conn_id=%s documents=%d bytes=%d", + account.ID, + truncateOpenAIWSLogValue(connID, openAIWSIDValueMaxLen), + len(documents), + len(message), + ) + message = documents[0] + pendingJSONDocuments = append(pendingJSONDocuments, documents[1:]...) + } + } + } if readErr != nil { lease.MarkBroken() closeStatus, closeReason := summarizeOpenAIWSReadCloseError(readErr) @@ -629,7 +650,10 @@ func (s *OpenAIGatewayService) forwardOpenAIWSV2( } if isTerminalEvent { - cleanExit = true + // A terminal event must be the final JSON document in its WS message. + // Ignore any tail for the completed client turn, but never reuse the + // ambiguous upstream connection for another request. + cleanExit = len(pendingJSONDocuments) == 0 break } }