Skip to content

feat: forward all Codex SSE response headers for codex sub #6098

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
seefs001 wants to merge 1 commit into
base: main
Choose a base branch
from
Open

feat: forward all Codex SSE response headers for codex sub #6098

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
21 changes: 13 additions & 8 deletions relay/helper/stream_scanner.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,14 +18,17 @@ import (
"github.com/QuantumNous/new-api/setting/operation_setting"

"github.com/bytedance/gopkg/util/gopool"
"github.com/samber/lo"

"github.com/gin-gonic/gin"
)

const (
InitialScannerBufferSize = 64 << 10 // 64KB (64*1024)
DefaultMaxScannerBufferSize = 128 << 20 // 64MB (64*1024*1024) default SSE buffer size
DefaultPingInterval = 10 * time.Second
InitialScannerBufferSize = 64 << 10 // 64KB (64*1024)
DefaultMaxScannerBufferSize = 128 << 20 // 64MB (64*1024*1024) default SSE buffer size

@coderabbitai coderabbitai Bot Jul 11, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win

Fix incorrect comment: 128 << 20 is 128MB, not 64MB.

The value 128 << 20 equals 134,217,728 bytes (128MB), but the comment states "64MB (6410241024)". This could mislead developers tuning buffer size.

📝 Proposed fix 
-	DefaultMaxScannerBufferSize  = 128 << 20 // 64MB (64*1024*1024) default SSE buffer size
+	DefaultMaxScannerBufferSize  = 128 << 20 // 128MB (128*1024*1024) default SSE buffer size
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
DefaultMaxScannerBufferSize = 128 << 20 // 64MB (64*1024*1024) default SSE buffer size
DefaultMaxScannerBufferSize = 128 << 20 // 128MB (128*1024*1024) default SSE buffer size
🤖 Prompt for AI Agents 
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@relay/helper/stream_scanner.go` at line 28, Correct the comment for
DefaultMaxScannerBufferSize to state that 128 << 20 is 128MB (128*1024*1024),
matching the configured value.

DefaultPingInterval = 10 * time.Second
codexReasoningIncludedHeader = "X-Reasoning-Included"
codexTurnStateHeader = "X-Codex-Turn-State"
// streamWriteTimeout bounds a single blocked write to a slow client so the
// unconditional wg.Wait() in cleanup can always finish. Without it, a slow
// but connected client (full TCP buffer, no server WriteTimeout) could hang
Expand All @@ -46,13 +49,15 @@ func NewStreamScanner(reader io.Reader) *bufio.Scanner {
return scanner
}

func copyCodexSSEHeaders(c *gin.Context, resp *http.Response) {
func copyCodexSSEHeaders(c *gin.Context, resp *http.Response, copyAll bool) {
if c == nil || c.Writer == nil || resp == nil {
return
}
// codex
for _, name := range []string{"X-Reasoning-Included", "X-Codex-Turn-State"} {
values := resp.Header.Values(name)
headers := resp.Header
if !copyAll {
headers = lo.PickByKeys(resp.Header, []string{codexReasoningIncludedHeader, codexTurnStateHeader})
}
for name, values := range headers {
if !service.ShouldCopyUpstreamHeader(c, name, values) {
continue
}
Expand Down Expand Up @@ -141,7 +146,7 @@ func StreamScannerHandler(c *gin.Context, resp *http.Response, info *relaycommon
defer cleanup()

scanner.Split(bufio.ScanLines)
copyCodexSSEHeaders(c, resp)
copyCodexSSEHeaders(c, resp, info.ChannelMeta.ChannelType == constant.ChannelTypeCodex)
SetEventStreamHeaders(c)

ctx = context.WithValue(ctx, "stop_chan", stopChan)
Expand Down
40 changes: 40 additions & 0 deletions relay/helper/stream_scanner_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -99,6 +99,46 @@ func TestStreamScannerHandler_EmptyBody(t *testing.T) {
assert.False(t, called.Load(), "handler should not be called for empty body")
}

func TestStreamScannerHandler_CodexCopiesEligibleUpstreamHeaders(t *testing.T) {
t.Parallel()

c, resp, info := setupStreamTest(t, strings.NewReader(""))
info.ChannelMeta.ChannelType = constant.ChannelTypeCodex
resp.Header = http.Header{
"X-Reasoning-Included": {"true"},
"X-Codex-Turn-State": {"turn-state"},
"X-Upstream-Trace": {"trace-id"},
"Set-Cookie": {"first=1", "second=2"},
"Content-Length": {"42"},
}

StreamScannerHandler(c, resp, info, func(data string, sr *StreamResult) {})

assert.Equal(t, "true", c.Writer.Header().Get("X-Reasoning-Included"))
assert.Equal(t, "turn-state", c.Writer.Header().Get("X-Codex-Turn-State"))
assert.Equal(t, "trace-id", c.Writer.Header().Get("X-Upstream-Trace"))
assert.Equal(t, []string{"first=1", "second=2"}, c.Writer.Header().Values("Set-Cookie"))
assert.Empty(t, c.Writer.Header().Get("Content-Length"))
}

func TestStreamScannerHandler_NonCodexPreservesLegacyHeaderBehavior(t *testing.T) {
t.Parallel()

c, resp, info := setupStreamTest(t, strings.NewReader(""))
info.ChannelMeta.ChannelType = constant.ChannelTypeOpenAI
resp.Header = http.Header{
"X-Reasoning-Included": {"true"},
"X-Codex-Turn-State": {"turn-state"},
"X-Upstream-Trace": {"trace-id"},
}

StreamScannerHandler(c, resp, info, func(data string, sr *StreamResult) {})

assert.Equal(t, "true", c.Writer.Header().Get("X-Reasoning-Included"))
assert.Equal(t, "turn-state", c.Writer.Header().Get("X-Codex-Turn-State"))
assert.Empty(t, c.Writer.Header().Get("X-Upstream-Trace"))
}

func TestStreamScannerHandler_1000Chunks(t *testing.T) {
t.Parallel()

Expand Down