diff --git a/backend/internal/service/openai_gateway_grok.go b/backend/internal/service/openai_gateway_grok.go index fdfd0bf32..af8240a76 100644 --- a/backend/internal/service/openai_gateway_grok.go +++ b/backend/internal/service/openai_gateway_grok.go @@ -179,11 +179,12 @@ func (s *OpenAIGatewayService) forwardGrokResponses( var firstTokenMs *int responseID := "" if reqStream { + maxLineSize := defaultMaxLineSize + if s.cfg != nil && s.cfg.Gateway.MaxLineSize > 0 { + maxLineSize = s.cfg.Gateway.MaxLineSize + } + resp.Body = newGrokResponsesBillingPingFilterBody(resp.Body, account, maxLineSize) if hasGrokResponsesClientToolMapping(clientToolMapping) { - maxLineSize := defaultMaxLineSize - if s.cfg != nil && s.cfg.Gateway.MaxLineSize > 0 { - maxLineSize = s.cfg.Gateway.MaxLineSize - } resp.Body = newGrokResponsesClientToolStreamBody(resp.Body, clientToolMapping, maxLineSize) } streamResult, err := s.handleStreamingResponse(ctx, resp, c, account, startTime, originalModel, upstreamModel) diff --git a/backend/internal/service/openai_gateway_grok_sse_filter.go b/backend/internal/service/openai_gateway_grok_sse_filter.go new file mode 100644 index 000000000..78159144d --- /dev/null +++ b/backend/internal/service/openai_gateway_grok_sse_filter.go @@ -0,0 +1,248 @@ +package service + +import ( + "bufio" + "bytes" + "encoding/json" + "fmt" + "io" + "strings" + "sync" +) + +// OpenAI Responses SSE event types are a closed enum for strict clients +// (grok CLI, Codex CLI): an unknown `event: ping` frame aborts the whole +// turn. Vendor gateways behind Grok subscriptions inject such frames for +// billing/keepalive, so ping frames are rewritten into an SSE comment that +// every parser ignores while the connection still looks alive downstream. +var grokResponsesPingComment = []byte(": ping\n\n") + +// A vendor ping frame is one event line plus one small data line. Cap what is +// buffered while deciding, so an upstream streaming a frame that never ends +// cannot grow gateway memory; frames over the cap are passed through as-is. +const ( + grokResponsesPingFrameMaxLines = 16 + grokResponsesPingFrameMaxBytes = 16 * 1024 +) + +type grokResponsesBillingPingFilterBody struct { + *io.PipeReader + source io.Closer + closeOnce sync.Once + closeErr error +} + +func (b *grokResponsesBillingPingFilterBody) Close() error { + readerErr := b.PipeReader.Close() + sourceErr := b.closeSource() + if readerErr != nil { + return readerErr + } + return sourceErr +} + +func (b *grokResponsesBillingPingFilterBody) closeSource() error { + b.closeOnce.Do(func() { b.closeErr = b.source.Close() }) + return b.closeErr +} + +func newGrokResponsesBillingPingFilterBody(source io.ReadCloser, account *Account, maxLineSize int) io.ReadCloser { + if account == nil || account.Platform != PlatformGrok { + return source + } + reader, writer := io.Pipe() + body := &grokResponsesBillingPingFilterBody{PipeReader: reader, source: source} + go filterGrokResponsesBillingPings(source, writer, body.closeSource, maxLineSize) + return body +} + +func filterGrokResponsesBillingPings( + source io.Reader, + destination *io.PipeWriter, + closeSource func() error, + maxLineSize int, +) { + defer func() { _ = closeSource() }() + if maxLineSize <= 0 { + maxLineSize = defaultMaxLineSize + } + + scanner := bufio.NewScanner(source) + scanBuf := getSSEScannerBuf64K() + defer putSSEScannerBuf64K(scanBuf) + initialBufferSize := len(scanBuf) + if maxLineSize < initialBufferSize { + initialBufferSize = maxLineSize + } + scanner.Buffer(scanBuf[:0:initialBufferSize], maxLineSize) + scanner.Split(scanSSELinesPreservingEndings) + + // Only frames opened by an `event: ping` line are buffered (pingFrame); + // every other frame streams through line by line without copying. + pingFrame := make([][]byte, 0, 3) + pingFrameBytes := 0 + inPassthroughFrame := false + + replayPingFrame := func() error { + for _, line := range pingFrame { + if _, err := destination.Write(line); err != nil { + return err + } + } + pingFrame = pingFrame[:0] + pingFrameBytes = 0 + return nil + } + // endPingFrame decides a complete buffered candidate: vendor ping frames + // become an SSE comment, everything else is replayed verbatim. blankLine + // is nil when the stream ends inside the frame. + endPingFrame := func(blankLine []byte) error { + if isGrokResponsesPingEventFrame(pingFrame) { + pingFrame = pingFrame[:0] + pingFrameBytes = 0 + _, err := destination.Write(grokResponsesPingComment) + return err + } + if err := replayPingFrame(); err != nil { + return err + } + if blankLine == nil { + return nil + } + _, err := destination.Write(blankLine) + return err + } + abort := func(err error) { _ = destination.CloseWithError(err) } + + for scanner.Scan() { + line := scanner.Bytes() + isBlank := len(trimSSELineEnding(line)) == 0 + + if inPassthroughFrame { + if _, err := destination.Write(line); err != nil { + abort(err) + return + } + if isBlank { + inPassthroughFrame = false + } + continue + } + + if len(pingFrame) > 0 { + if isBlank { + if err := endPingFrame(line); err != nil { + abort(err) + return + } + continue + } + if canExtendGrokResponsesPingFrame(line) && + len(pingFrame) < grokResponsesPingFrameMaxLines && + pingFrameBytes+len(line) <= grokResponsesPingFrameMaxBytes { + pingFrame = append(pingFrame, append([]byte(nil), line...)) + pingFrameBytes += len(line) + continue + } + // Not a filterable ping frame after all (unexpected field line, + // or past the buffering caps): replay it and stream the rest of + // the frame through unchanged. + if err := replayPingFrame(); err != nil { + abort(err) + return + } + if _, err := destination.Write(line); err != nil { + abort(err) + return + } + inPassthroughFrame = true + continue + } + + // Frame start: only `event: ping` opens a buffered candidate. + if !isBlank { + if value, ok := extractOpenAISSEEventLine(string(trimSSELineEnding(line))); ok && value == "ping" { + pingFrame = append(pingFrame, append([]byte(nil), line...)) + pingFrameBytes = len(line) + continue + } + } + if _, err := destination.Write(line); err != nil { + abort(err) + return + } + inPassthroughFrame = !isBlank + } + if len(pingFrame) > 0 { + if err := endPingFrame(nil); err != nil { + abort(err) + return + } + } + if err := scanner.Err(); err != nil { + abort(fmt.Errorf("filter Grok Responses billing ping: %w", err)) + return + } + _ = destination.Close() +} + +func scanSSELinesPreservingEndings(data []byte, atEOF bool) (advance int, token []byte, err error) { + for index, value := range data { + switch value { + case '\n': + return index + 1, data[:index+1], nil + case '\r': + if index+1 == len(data) && !atEOF { + return 0, nil, nil + } + if index+1 < len(data) && data[index+1] == '\n' { + return index + 2, data[:index+2], nil + } + return index + 1, data[:index+1], nil + } + } + if atEOF && len(data) > 0 { + return len(data), data, nil + } + return 0, nil, nil +} + +func trimSSELineEnding(line []byte) []byte { + return bytes.TrimSuffix(bytes.TrimSuffix(line, []byte("\n")), []byte("\r")) +} + +// canExtendGrokResponsesPingFrame reports whether a line may still belong to a +// vendor ping frame: only data lines and SSE comments. Any other field (a +// second event line, id, retry, ...) means the frame is not a plain ping. +func canExtendGrokResponsesPingFrame(rawLine []byte) bool { + line := trimSSELineEnding(rawLine) + if len(line) > 0 && line[0] == ':' { + return true + } + _, ok := extractOpenAISSEDataLine(string(line)) + return ok +} + +// isGrokResponsesPingEventFrame decides a buffered candidate whose first line +// is already `event: ping`. The only candidates replayed verbatim are frames +// whose data payload declares a different event type than the SSE event line; +// every other shape (billing cost, keepalive, no data, malformed JSON) would +// break strict Responses clients and is rewritten into a comment. +func isGrokResponsesPingEventFrame(rawLines [][]byte) bool { + dataParts := make([]string, 0, 1) + for _, rawLine := range rawLines[1:] { + if value, ok := extractOpenAISSEDataLine(string(trimSSELineEnding(rawLine))); ok { + dataParts = append(dataParts, value) + } + } + if len(dataParts) == 0 { + return true + } + var payload struct { + Type *string `json:"type"` + } + if err := json.Unmarshal([]byte(strings.Join(dataParts, "\n")), &payload); err != nil || payload.Type == nil { + return true + } + return *payload.Type == "ping" +} diff --git a/backend/internal/service/openai_gateway_grok_sse_filter_test.go b/backend/internal/service/openai_gateway_grok_sse_filter_test.go new file mode 100644 index 000000000..1dce65e95 --- /dev/null +++ b/backend/internal/service/openai_gateway_grok_sse_filter_test.go @@ -0,0 +1,262 @@ +package service + +import ( + "context" + "errors" + "io" + "net/http" + "net/http/httptest" + "strings" + "sync/atomic" + "testing" + "time" + + "github.com/Wei-Shaw/sub2api/internal/config" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" +) + +func filterGrokPingTestInput(t *testing.T, input string) string { + t.Helper() + body := newGrokResponsesBillingPingFilterBody( + io.NopCloser(strings.NewReader(input)), + &Account{Platform: PlatformGrok}, + defaultMaxLineSize, + ) + output, err := io.ReadAll(body) + require.NoError(t, err) + require.NoError(t, body.Close()) + return string(output) +} + +func TestGrokResponsesBillingPingFilter(t *testing.T) { + input := strings.Join([]string{ + ": upstream keepalive", + "", + "event: response.output_text.delta", + `data: {"type":"response.output_text.delta","delta":"hello"}`, + "", + "event: future.vendor_event", + `data: {"type":"future.vendor_event","value":1}`, + "", + "event: ping", + `data: {"type":"ping","x-opencode-type":"inference-cost","cost":2.75,"input-tokens":42}`, + "", + "event: ping", + `data: {"type":"ping","cost":"0"}`, + "", + "event: response.completed", + `data: {"type":"response.completed","response":{"id":"resp_1","usage":{"input_tokens":3,"output_tokens":5}}}`, + "", + }, "\n") + + result := filterGrokPingTestInput(t, input) + require.NotContains(t, result, "event: ping") + require.NotContains(t, result, `"x-opencode-type":"inference-cost"`) + require.NotContains(t, result, `{"type":"ping","cost":"0"}`) + require.Equal(t, 2, strings.Count(result, ": ping\n\n")) + require.Contains(t, result, ": upstream keepalive\n\n") + require.Contains(t, result, "event: response.output_text.delta") + require.Contains(t, result, `{"type":"response.output_text.delta","delta":"hello"}`) + require.Contains(t, result, "event: future.vendor_event") + require.Contains(t, result, `{"type":"future.vendor_event","value":1}`) + require.Contains(t, result, "event: response.completed") + require.Contains(t, result, `"usage":{"input_tokens":3,"output_tokens":5}`) +} + +// Every `event: ping` frame is outside the Responses closed event enum and +// breaks strict clients regardless of its payload shape, so all variants are +// rewritten into an SSE comment (issue #5105). +func TestGrokResponsesBillingPingFilterConvertsPingVariants(t *testing.T) { + frames := []string{ + "event: ping\ndata: {\"type\":\"ping\",\"x-opencode-type\":\"inference-cost\",\"cost\":\"0.06029240\"}\n\n", + "event: ping\ndata: {\"type\":\"ping\",\"cost\":\"0\"}\n\n", + "event: ping\ndata: {\"type\":\"ping\",\"cost\":\"0.06029240\"}\n\n", + "event: ping\ndata: {\"type\":\"ping\",\"kind\":\"keepalive\"}\n\n", + "event: ping\ndata: {\"type\":\"ping\"}\n\n", + "event: ping\ndata: {\"type\":\"ping\",\"cost\":2}\n\n", + "event: ping\ndata: {\"type\":\"ping\",\"cost\":0.0001}\n\n", + "event: ping\ndata: {\"type\":\"ping\",\"cost\":\" 0 \"}\n\n", + "event: ping\ndata: {\"type\":\"ping\",\"x-opencode-type\":\"keepalive\",\"cost\":0}\n\n", + "event: ping\ndata: {\"type\":\"ping\",\"x-opencode-type\":null,\"cost\":0}\n\n", + "event: ping\ndata: {\"cost\":\"0\"}\n\n", + "event: ping\ndata: {not-json}\n\n", + "event: ping\n\n", + "event: ping\n: vendor note\ndata: {\"type\":\"ping\"}\n\n", + } + result := filterGrokPingTestInput(t, strings.Join(frames, "")) + require.Equal(t, strings.Repeat(": ping\n\n", len(frames)), result) +} + +func TestGrokResponsesBillingPingFilterPreservesNonPingFrames(t *testing.T) { + input := strings.Join([]string{ + "event: ping", + `data: {"type":" ping ","cost":0}`, + "", + "event: ping", + `data: {"type":"response.completed"}`, + "", + "event: custom", + `data: {"type":"ping","x-opencode-type":"inference-cost"}`, + "", + `data: {"type":"ping","cost":"0"}`, + "", + ": keepalive comment", + "", + "retry: 1000", + "", + }, "\n") + + require.Equal(t, input, filterGrokPingTestInput(t, input)) +} + +// A ping candidate that turns out to carry an unexpected SSE field is not a +// vendor billing/keepalive frame; it must be replayed byte for byte. +func TestGrokResponsesBillingPingFilterPassesThroughPingFrameWithUnknownField(t *testing.T) { + input := "event: ping\nid: 7\ndata: {\"type\":\"ping\",\"cost\":\"0\"}\n\n" + require.Equal(t, input, filterGrokPingTestInput(t, input)) +} + +// Buffering caps: a ping candidate that grows past the line or byte limit is +// streamed through unchanged instead of accumulating unbounded memory. +func TestGrokResponsesBillingPingFilterPassesThroughOversizedPingFrame(t *testing.T) { + lines := []string{"event: ping"} + for i := 0; i < grokResponsesPingFrameMaxLines; i++ { + lines = append(lines, ": filler comment") + } + lines = append(lines, `data: {"type":"ping","cost":"0"}`, "") + byLines := strings.Join(lines, "\n") + require.Equal(t, byLines, filterGrokPingTestInput(t, byLines)) + + byBytes := "event: ping\ndata: {\"type\":\"ping\",\"pad\":\"" + + strings.Repeat("x", grokResponsesPingFrameMaxBytes) + "\"}\n\n" + require.Equal(t, byBytes, filterGrokPingTestInput(t, byBytes)) +} + +func TestGrokResponsesBillingPingFilterConvertsMalformedPingFrames(t *testing.T) { + input := "event: ping\r\ndata: {not-json}\r\n\r\n" + + "event: ping\r\ndata: {\"type\":\"ping\",\"cost\":\"0\"} trailing\r\n\r\n" + + "event: future.response.event\r\ndata: {\"type\":\"future.response.event\"}" + want := ": ping\n\n" + ": ping\n\n" + + "event: future.response.event\r\ndata: {\"type\":\"future.response.event\"}" + require.Equal(t, want, filterGrokPingTestInput(t, input)) +} + +func TestGrokResponsesBillingPingFilterHandlesBareCRFrames(t *testing.T) { + input := "event: ping\rdata: {\"type\":\"ping\",\"cost\":\"0\"}\r\r" + + "event: future.event\rdata: {\"type\":\"future.event\"}\r\r" + want := ": ping\n\n" + "event: future.event\rdata: {\"type\":\"future.event\"}\r\r" + require.Equal(t, want, filterGrokPingTestInput(t, input)) +} + +func TestGrokResponsesBillingPingFilterConvertsPartialPingFrameAtEOF(t *testing.T) { + input := "event: ping\ndata: {\"type\":\"ping\",\"cost\":\"0\"}" + require.Equal(t, ": ping\n\n", filterGrokPingTestInput(t, input)) +} + +func TestGrokResponsesBillingPingFilterDoesNotFilterNonGrokAccounts(t *testing.T) { + input := "event: ping\ndata: {\"type\":\"ping\",\"cost\":\"0\"}\n\n" + source := io.NopCloser(strings.NewReader(input)) + body := newGrokResponsesBillingPingFilterBody(source, &Account{Platform: PlatformOpenAI}, defaultMaxLineSize) + + output, err := io.ReadAll(body) + require.NoError(t, err) + require.NoError(t, body.Close()) + require.Equal(t, input, string(output)) +} + +func TestGrokResponsesBillingPingFilterPreservesUsageAndTerminalEvent(t *testing.T) { + gin.SetMode(gin.TestMode) + input := strings.Join([]string{ + "event: ping", + `data: {"type":"ping","x-opencode-type":"inference-cost","cost":"0"}`, + "", + "event: response.completed", + `data: {"type":"response.completed","response":{"id":"resp_1","usage":{"input_tokens":3,"output_tokens":5}}}`, + "", + }, "\n") + account := &Account{ID: 1, Platform: PlatformGrok} + resp := &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{}, + Body: newGrokResponsesBillingPingFilterBody( + io.NopCloser(strings.NewReader(input)), account, defaultMaxLineSize, + ), + } + 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}}, + toolCorrector: NewCodexToolCorrector(), + } + + result, err := svc.handleStreamingResponse(context.Background(), resp, c, account, time.Now(), "grok-4.5", "grok-4.5") + require.NoError(t, err) + require.Equal(t, 3, result.usage.InputTokens) + require.Equal(t, 5, result.usage.OutputTokens) + require.Equal(t, "resp_1", result.responseID) + require.Contains(t, recorder.Body.String(), "response.completed") + require.NotContains(t, recorder.Body.String(), "inference-cost") + require.NotContains(t, recorder.Body.String(), "event: ping") +} + +type grokPingFilterTestReadCloser struct { + reader io.ReadCloser + closeCount atomic.Int32 +} + +func (r *grokPingFilterTestReadCloser) Read(p []byte) (int, error) { return r.reader.Read(p) } +func (r *grokPingFilterTestReadCloser) Close() error { + r.closeCount.Add(1) + return r.reader.Close() +} + +func TestGrokResponsesBillingPingFilterCloseCancelsSourceOnce(t *testing.T) { + upstreamReader, upstreamWriter := io.Pipe() + source := &grokPingFilterTestReadCloser{reader: upstreamReader} + body := newGrokResponsesBillingPingFilterBody(source, &Account{Platform: PlatformGrok}, defaultMaxLineSize) + + require.NoError(t, body.Close()) + require.Eventually(t, func() bool { return source.closeCount.Load() == 1 }, time.Second, time.Millisecond) + _, err := upstreamWriter.Write([]byte("blocked")) + require.Error(t, err) + require.NoError(t, upstreamWriter.Close()) +} + +func TestGrokResponsesBillingPingFilterFlushesCompletedFrames(t *testing.T) { + upstreamReader, upstreamWriter := io.Pipe() + body := newGrokResponsesBillingPingFilterBody(upstreamReader, &Account{Platform: PlatformGrok}, defaultMaxLineSize) + t.Cleanup(func() { require.NoError(t, body.Close()) }) + + go func() { + _, _ = io.WriteString(upstreamWriter, "event: future.event\ndata: {\"type\":\"future.event\"}\n\n") + }() + + result := make(chan error, 1) + go func() { + buffer := make([]byte, 64) + n, err := body.Read(buffer) + if err == nil && !strings.Contains(string(buffer[:n]), "future.event") { + err = errors.New("completed frame was not forwarded") + } + result <- err + }() + select { + case err := <-result: + require.NoError(t, err) + case <-time.After(time.Second): + t.Fatal("completed frame was buffered until upstream EOF") + } +} + +func TestGrokResponsesBillingPingFilterReportsOversizedLine(t *testing.T) { + body := newGrokResponsesBillingPingFilterBody( + io.NopCloser(strings.NewReader("data: 123456789\n\n")), + &Account{Platform: PlatformGrok}, + 8, + ) + _, err := io.ReadAll(body) + require.ErrorContains(t, err, "filter Grok Responses billing ping") + require.NoError(t, body.Close()) +}