Skip to content

Commit 2175f64

Browse files
fix(gateway): extend compact keepalive to non-passthrough OAuth path
The initial implementation only covered the passthrough branch, but production traffic showed compact requests also flow through the standard OAuth HTTP forward path. Add keepalive to both Forward and forwardOpenAIPassthrough, with structured logging for diagnostics.
1 parent d5fdd1c commit 2175f64

2 files changed

Lines changed: 191 additions & 12 deletions

File tree

backend/internal/service/openai_gateway_service.go

Lines changed: 74 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -449,9 +449,21 @@ func NewOpenAIGatewayService(
449449
openAITokenProvider.SetAccountRuntimeBlocker(svc)
450450
}
451451
svc.logOpenAIWSModeBootstrap()
452+
svc.logOpenAICompactNonstreamKeepaliveBootstrap()
452453
return svc
453454
}
454455

456+
func (s *OpenAIGatewayService) logOpenAICompactNonstreamKeepaliveBootstrap() {
457+
interval := s.compactNonstreamKeepaliveInterval()
458+
if interval <= 0 {
459+
return
460+
}
461+
logger.L().With(
462+
zap.String("component", "service.openai_gateway"),
463+
zap.Int("interval_seconds", int(interval.Seconds())),
464+
).Info("OpenAI compact non-stream keepalive enabled")
465+
}
466+
455467
// ResolveChannelMapping 解析渠道级模型映射(代理到 ChannelService)
456468
func (s *OpenAIGatewayService) ResolveChannelMapping(ctx context.Context, groupID int64, model string) ChannelMappingResult {
457469
if s.channelService == nil {
@@ -3264,25 +3276,42 @@ func (s *OpenAIGatewayService) Forward(ctx context.Context, c *gin.Context, acco
32643276
}
32653277

32663278
// Send request
3279+
stopCompactKeepalive := func() {}
3280+
if !reqStream {
3281+
stopCompactKeepalive = s.startCompactNonstreamKeepalive(ctx, c)
3282+
}
3283+
32673284
upstreamStart := time.Now()
32683285
resp, err := s.httpUpstream.Do(upstreamReq, proxyURL, account.ID, account.Concurrency)
32693286
SetOpsLatencyMs(c, OpsUpstreamLatencyMsKey, time.Since(upstreamStart).Milliseconds())
32703287
if err != nil {
3288+
stopCompactKeepalive()
32713289
// Transport-level failure (proxy/DNS/TCP/TLS — no HTTP response). Convert to
32723290
// a failover so the handler switches to a healthy account, and temporarily
32733291
// unschedule the account on durable faults (e.g. rejected proxy credentials).
3274-
return nil, s.handleOpenAIUpstreamTransportError(ctx, c, account, err, false)
3292+
transportErr := s.handleOpenAIUpstreamTransportError(ctx, c, account, err, false)
3293+
if openAICompactKeepaliveCommitted(c) {
3294+
logOpenAICompactKeepaliveCommitted(ctx, c, account, nil)
3295+
writeOpenAICommittedTransportError(c)
3296+
return nil, fmt.Errorf("upstream request failed after compact keepalive: %s", sanitizeUpstreamErrorMessage(err.Error()))
3297+
}
3298+
return nil, transportErr
32753299
}
32763300

32773301
// Handle error response
32783302
if resp.StatusCode >= 400 {
3303+
stopCompactKeepalive()
32793304
respBody := s.readUpstreamErrorBody(resp)
32803305
_ = resp.Body.Close()
32813306
resp.Body = io.NopCloser(bytes.NewReader(respBody))
32823307

32833308
upstreamMsg := strings.TrimSpace(extractUpstreamErrorMessage(respBody))
32843309
upstreamMsg = sanitizeUpstreamErrorMessage(upstreamMsg)
32853310
upstreamCode := extractUpstreamErrorCode(respBody)
3311+
if openAICompactKeepaliveCommitted(c) {
3312+
logOpenAICompactKeepaliveCommitted(ctx, c, account, resp)
3313+
return s.handleErrorResponse(ctx, resp, c, account, body, billingModel)
3314+
}
32863315
if !httpInvalidEncryptedContentRetryTried && resp.StatusCode == http.StatusBadRequest && upstreamCode == "invalid_encrypted_content" {
32873316
decoded, decodeErr := ensureReqBody()
32883317
if decodeErr != nil {
@@ -3345,6 +3374,7 @@ func (s *OpenAIGatewayService) Forward(ctx context.Context, c *gin.Context, acco
33453374
imageCount := 0
33463375
var imageOutputSizes []string
33473376
if reqStream {
3377+
stopCompactKeepalive()
33483378
streamResult, err := s.handleStreamingResponse(ctx, resp, c, account, startTime, originalModel, upstreamModel)
33493379
if err != nil {
33503380
return nil, err
@@ -3355,7 +3385,7 @@ func (s *OpenAIGatewayService) Forward(ctx context.Context, c *gin.Context, acco
33553385
imageCount = streamResult.imageCount
33563386
imageOutputSizes = streamResult.imageOutputSizes
33573387
} else {
3358-
nonStreamResult, err := s.handleNonStreamingResponse(ctx, resp, c, account, originalModel, upstreamModel)
3388+
nonStreamResult, err := s.handleNonStreamingResponse(ctx, resp, c, account, originalModel, upstreamModel, stopCompactKeepalive)
33593389
if err != nil {
33603390
return nil, err
33613391
}
@@ -3976,6 +4006,16 @@ func (s *OpenAIGatewayService) startCompactNonstreamKeepalive(ctx context.Contex
39764006
if ctx == nil {
39774007
ctx = context.Background()
39784008
}
4009+
path := ""
4010+
if c.Request != nil && c.Request.URL != nil {
4011+
path = strings.TrimSpace(c.Request.URL.Path)
4012+
}
4013+
log := logger.FromContext(ctx).With(
4014+
zap.String("component", "service.openai_gateway"),
4015+
zap.String("request_path", path),
4016+
zap.Int("interval_seconds", int(interval.Seconds())),
4017+
)
4018+
log.Info("OpenAI compact non-stream keepalive started")
39794019

39804020
headers := c.Writer.Header()
39814021
headers.Set("Content-Type", "application/json")
@@ -3991,15 +4031,23 @@ func (s *OpenAIGatewayService) startCompactNonstreamKeepalive(ctx context.Contex
39914031
defer wg.Done()
39924032
ticker := time.NewTicker(interval)
39934033
defer ticker.Stop()
4034+
flushedLogged := false
39944035
for {
39954036
select {
39964037
case <-stopCh:
39974038
return
39984039
case <-ctx.Done():
39994040
return
40004041
case <-ticker.C:
4001-
_, _ = c.Writer.Write([]byte("\n"))
4042+
if _, err := c.Writer.Write([]byte("\n")); err != nil {
4043+
log.Warn("OpenAI compact non-stream keepalive write failed", zap.Error(err))
4044+
return
4045+
}
40024046
flusher.Flush()
4047+
if !flushedLogged {
4048+
log.Info("OpenAI compact non-stream keepalive flushed")
4049+
flushedLogged = true
4050+
}
40034051
}
40044052
}
40054053
}()
@@ -5795,17 +5843,22 @@ func openAIUsageFromGJSON(value gjson.Result) (OpenAIUsage, bool) {
57955843
}, true
57965844
}
57975845

5798-
func (s *OpenAIGatewayService) handleNonStreamingResponse(ctx context.Context, resp *http.Response, c *gin.Context, account *Account, originalModel, mappedModel string) (*openaiNonStreamingResult, error) {
5799-
body, err := ReadUpstreamResponseBody(resp.Body, s.cfg, c, openAITooLargeError)
5846+
func (s *OpenAIGatewayService) handleNonStreamingResponse(ctx context.Context, resp *http.Response, c *gin.Context, account *Account, originalModel, mappedModel string, stopBeforeWrite ...func()) (*openaiNonStreamingResult, error) {
5847+
stop := compactStopFunc(stopBeforeWrite...)
5848+
body, err := ReadUpstreamResponseBody(resp.Body, s.cfg, c, func(c *gin.Context) {
5849+
stop()
5850+
openAITooLargeError(c)
5851+
})
58005852
if err != nil {
5853+
stop()
58015854
return nil, err
58025855
}
58035856

58045857
// Detect SSE responses for ALL account types via Content-Type header.
58055858
// Some OpenAI-compatible upstreams (including other sub2api instances)
58065859
// may return SSE even when stream=false was requested.
58075860
if isEventStreamResponse(resp.Header) {
5808-
return s.handleSSEToJSON(resp, c, body, originalModel, mappedModel)
5861+
return s.handleSSEToJSON(resp, c, body, originalModel, mappedModel, stop)
58095862
}
58105863
bodyLooksLikeSSE := bytes.Contains(body, []byte("data:")) || bytes.Contains(body, []byte("event:"))
58115864

@@ -5815,14 +5868,15 @@ func (s *OpenAIGatewayService) handleNonStreamingResponse(ctx context.Context, r
58155868
// positives on JSON responses that coincidentally contain "data:" or
58165869
// "event:" in their text content.
58175870
if account.Type == AccountTypeOAuth && bodyLooksLikeSSE {
5818-
return s.handleSSEToJSON(resp, c, body, originalModel, mappedModel)
5871+
return s.handleSSEToJSON(resp, c, body, originalModel, mappedModel, stop)
58195872
}
58205873

58215874
usageValue, usageOK := extractOpenAIUsageFromJSONBytes(body)
58225875
if !usageOK {
58235876
if bodyLooksLikeSSE {
5824-
return s.handleSSEToJSON(resp, c, body, originalModel, mappedModel)
5877+
return s.handleSSEToJSON(resp, c, body, originalModel, mappedModel, stop)
58255878
}
5879+
stop()
58265880
return nil, fmt.Errorf("parse response: invalid json response")
58275881
}
58285882
usage := &usageValue
@@ -5841,6 +5895,7 @@ func (s *OpenAIGatewayService) handleNonStreamingResponse(ctx context.Context, r
58415895
}
58425896
}
58435897

5898+
stop()
58445899
c.Data(resp.StatusCode, contentType, body)
58455900

58465901
return &openaiNonStreamingResult{
@@ -5852,12 +5907,20 @@ func (s *OpenAIGatewayService) handleNonStreamingResponse(ctx context.Context, r
58525907
}, nil
58535908
}
58545909

5910+
func compactStopFunc(stops ...func()) func() {
5911+
if len(stops) == 0 || stops[0] == nil {
5912+
return func() {}
5913+
}
5914+
return stops[0]
5915+
}
5916+
58555917
func isEventStreamResponse(header http.Header) bool {
58565918
contentType := strings.ToLower(header.Get("Content-Type"))
58575919
return strings.Contains(contentType, "text/event-stream")
58585920
}
58595921

5860-
func (s *OpenAIGatewayService) handleSSEToJSON(resp *http.Response, c *gin.Context, body []byte, originalModel, mappedModel string) (*openaiNonStreamingResult, error) {
5922+
func (s *OpenAIGatewayService) handleSSEToJSON(resp *http.Response, c *gin.Context, body []byte, originalModel, mappedModel string, stopBeforeWrite ...func()) (*openaiNonStreamingResult, error) {
5923+
stop := compactStopFunc(stopBeforeWrite...)
58615924
bodyText := string(body)
58625925
finalResponse, ok := extractCodexFinalResponse(bodyText)
58635926

@@ -5889,6 +5952,7 @@ func (s *OpenAIGatewayService) handleSSEToJSON(resp *http.Response, c *gin.Conte
58895952
if msg == "" {
58905953
msg = "Upstream compact response failed"
58915954
}
5955+
stop()
58925956
return nil, s.writeOpenAINonStreamingProtocolError(resp, c, msg)
58935957
}
58945958
usage = s.parseSSEUsageFromBody(bodyText)
@@ -5907,6 +5971,7 @@ func (s *OpenAIGatewayService) handleSSEToJSON(resp *http.Response, c *gin.Conte
59075971
contentType = "text/event-stream"
59085972
}
59095973
}
5974+
stop()
59105975
c.Data(resp.StatusCode, contentType, body)
59115976

59125977
return &openaiNonStreamingResult{

backend/internal/service/openai_oauth_passthrough_test.go

Lines changed: 117 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -64,12 +64,12 @@ func (u *httpUpstreamRecorder) Do(req *http.Request, proxyURL string, accountID
6464
req.Body = io.NopCloser(bytes.NewReader(b))
6565
}
6666
u.requests = append(u.requests, req)
67-
if u.err != nil {
68-
return nil, u.err
69-
}
7067
if u.delay > 0 {
7168
time.Sleep(u.delay)
7269
}
70+
if u.err != nil {
71+
return nil, u.err
72+
}
7373
if len(u.responses) > 0 {
7474
resp := u.responses[0]
7575
u.responses = u.responses[1:]
@@ -514,6 +514,55 @@ func TestOpenAIGatewayService_OAuthPassthrough_CompactNonstreamKeepaliveWritesLe
514514
require.Contains(t, string(body), `"id":"cmp_keepalive"`)
515515
}
516516

517+
func TestOpenAIGatewayService_OAuthCompactNonPassthroughKeepaliveWritesLeadingNewline(t *testing.T) {
518+
gin.SetMode(gin.TestMode)
519+
520+
rec := httptest.NewRecorder()
521+
c, _ := gin.CreateTestContext(rec)
522+
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses/compact", bytes.NewReader(nil))
523+
c.Request.Header.Set("User-Agent", "Codex Desktop/0.135.0-alpha.1")
524+
c.Request.Header.Set("Content-Type", "application/json")
525+
526+
originalBody := []byte(`{"model":"gpt-5.5","stream":false,"instructions":"local-test-instructions","input":[{"type":"message","role":"user","content":"compact me"}]}`)
527+
resp := &http.Response{
528+
StatusCode: http.StatusOK,
529+
Header: http.Header{"Content-Type": []string{"application/json"}, "x-request-id": []string{"rid-compact-native-keepalive"}},
530+
Body: io.NopCloser(strings.NewReader(`{"id":"cmp_native_keepalive","usage":{"input_tokens":11,"output_tokens":22}}`)),
531+
}
532+
upstream := &httpUpstreamRecorder{resp: resp, delay: 1200 * time.Millisecond}
533+
534+
svc := &OpenAIGatewayService{
535+
cfg: &config.Config{Gateway: config.GatewayConfig{
536+
OpenAICompactNonstreamKeepaliveInterval: 1,
537+
}},
538+
httpUpstream: upstream,
539+
}
540+
account := &Account{
541+
ID: 124,
542+
Name: "acc-native",
543+
Platform: PlatformOpenAI,
544+
Type: AccountTypeOAuth,
545+
Concurrency: 1,
546+
Credentials: map[string]any{
547+
"access_token": "oauth-token",
548+
"chatgpt_account_id": "chatgpt-acc",
549+
},
550+
Status: StatusActive,
551+
Schedulable: true,
552+
RateMultiplier: f64p(1),
553+
}
554+
555+
result, err := svc.Forward(context.Background(), c, account, originalBody)
556+
require.NoError(t, err)
557+
require.NotNil(t, result)
558+
body := rec.Body.Bytes()
559+
require.True(t, bytes.HasPrefix(body, []byte("\n")), "native compact keepalive should write a leading blank line")
560+
require.True(t, json.Valid(bytes.TrimSpace(body)))
561+
require.Contains(t, string(body), `"id":"cmp_native_keepalive"`)
562+
require.NotNil(t, upstream.lastReq)
563+
require.Equal(t, chatgptCodexURL+"/compact", upstream.lastReq.URL.String())
564+
}
565+
517566
func TestOpenAIGatewayService_OAuthPassthrough_CompactNonstreamKeepaliveDisabledHasNoLeadingWhitespace(t *testing.T) {
518567
gin.SetMode(gin.TestMode)
519568

@@ -640,6 +689,71 @@ func TestOpenAIGatewayService_OAuthPassthrough_CompactNonstreamKeepaliveCommitte
640689
}
641690
}
642691

692+
func TestOpenAIGatewayService_OAuthCompactNonstreamKeepaliveCommittedTransportErrorDoesNotFailover(t *testing.T) {
693+
gin.SetMode(gin.TestMode)
694+
695+
testCases := []struct {
696+
name string
697+
body []byte
698+
account *Account
699+
}{
700+
{
701+
name: "passthrough",
702+
body: []byte(`{"model":"gpt-5.1-codex","stream":true,"instructions":"local-test-instructions","input":[{"type":"text","text":"compact me"}]}`),
703+
account: newOpenAICompactPassthroughTestAccount(),
704+
},
705+
{
706+
name: "non_passthrough",
707+
body: []byte(`{"model":"gpt-5.5","stream":false,"instructions":"local-test-instructions","input":[{"type":"message","role":"user","content":"compact me"}]}`),
708+
account: &Account{
709+
ID: 124,
710+
Name: "acc-native",
711+
Platform: PlatformOpenAI,
712+
Type: AccountTypeOAuth,
713+
Concurrency: 1,
714+
Credentials: map[string]any{
715+
"access_token": "oauth-token",
716+
"chatgpt_account_id": "chatgpt-acc",
717+
},
718+
Status: StatusActive,
719+
Schedulable: true,
720+
RateMultiplier: f64p(1),
721+
},
722+
},
723+
}
724+
725+
for _, tc := range testCases {
726+
t.Run(tc.name, func(t *testing.T) {
727+
rec := httptest.NewRecorder()
728+
c, _ := gin.CreateTestContext(rec)
729+
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses/compact", bytes.NewReader(nil))
730+
c.Request.Header.Set("User-Agent", "codex_cli_rs/0.1.0")
731+
c.Request.Header.Set("Content-Type", "application/json")
732+
733+
upstream := &httpUpstreamRecorder{
734+
err: errors.New("dial tcp: i/o timeout"),
735+
delay: 1200 * time.Millisecond,
736+
}
737+
svc := &OpenAIGatewayService{
738+
cfg: &config.Config{Gateway: config.GatewayConfig{
739+
ForceCodexCLI: false,
740+
OpenAICompactNonstreamKeepaliveInterval: 1,
741+
}},
742+
httpUpstream: upstream,
743+
}
744+
745+
result, err := svc.Forward(context.Background(), c, tc.account, tc.body)
746+
require.Error(t, err)
747+
require.Nil(t, result)
748+
var failoverErr *UpstreamFailoverError
749+
require.False(t, errors.As(err, &failoverErr), "committed compact keepalive must suppress transport failover")
750+
body := rec.Body.Bytes()
751+
require.True(t, bytes.HasPrefix(body, []byte("\n")))
752+
require.Contains(t, string(body), "Upstream request failed")
753+
})
754+
}
755+
}
756+
643757
func newOpenAICompactPassthroughTestAccount() *Account {
644758
return &Account{
645759
ID: 123,

0 commit comments

Comments
 (0)