Skip to content

Commit 93cb989

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 7a4eabe commit 93cb989

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
@@ -448,9 +448,21 @@ func NewOpenAIGatewayService(
448448
openAITokenProvider.SetAccountRuntimeBlocker(svc)
449449
}
450450
svc.logOpenAIWSModeBootstrap()
451+
svc.logOpenAICompactNonstreamKeepaliveBootstrap()
451452
return svc
452453
}
453454

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

31853197
// Send request
3198+
stopCompactKeepalive := func() {}
3199+
if !reqStream {
3200+
stopCompactKeepalive = s.startCompactNonstreamKeepalive(ctx, c)
3201+
}
3202+
31863203
upstreamStart := time.Now()
31873204
resp, err := s.httpUpstream.Do(upstreamReq, proxyURL, account.ID, account.Concurrency)
31883205
SetOpsLatencyMs(c, OpsUpstreamLatencyMsKey, time.Since(upstreamStart).Milliseconds())
31893206
if err != nil {
3207+
stopCompactKeepalive()
31903208
// Transport-level failure (proxy/DNS/TCP/TLS — no HTTP response). Convert to
31913209
// a failover so the handler switches to a healthy account, and temporarily
31923210
// unschedule the account on durable faults (e.g. rejected proxy credentials).
3193-
return nil, s.handleOpenAIUpstreamTransportError(ctx, c, account, err, false)
3211+
transportErr := s.handleOpenAIUpstreamTransportError(ctx, c, account, err, false)
3212+
if openAICompactKeepaliveCommitted(c) {
3213+
logOpenAICompactKeepaliveCommitted(ctx, c, account, nil)
3214+
writeOpenAICommittedTransportError(c)
3215+
return nil, fmt.Errorf("upstream request failed after compact keepalive: %s", sanitizeUpstreamErrorMessage(err.Error()))
3216+
}
3217+
return nil, transportErr
31943218
}
31953219

31963220
// Handle error response
31973221
if resp.StatusCode >= 400 {
3222+
stopCompactKeepalive()
31983223
respBody := s.readUpstreamErrorBody(resp)
31993224
_ = resp.Body.Close()
32003225
resp.Body = io.NopCloser(bytes.NewReader(respBody))
32013226

32023227
upstreamMsg := strings.TrimSpace(extractUpstreamErrorMessage(respBody))
32033228
upstreamMsg = sanitizeUpstreamErrorMessage(upstreamMsg)
32043229
upstreamCode := extractUpstreamErrorCode(respBody)
3230+
if openAICompactKeepaliveCommitted(c) {
3231+
logOpenAICompactKeepaliveCommitted(ctx, c, account, resp)
3232+
return s.handleErrorResponse(ctx, resp, c, account, body, billingModel)
3233+
}
32053234
if !httpInvalidEncryptedContentRetryTried && resp.StatusCode == http.StatusBadRequest && upstreamCode == "invalid_encrypted_content" {
32063235
decoded, decodeErr := ensureReqBody()
32073236
if decodeErr != nil {
@@ -3264,6 +3293,7 @@ func (s *OpenAIGatewayService) Forward(ctx context.Context, c *gin.Context, acco
32643293
imageCount := 0
32653294
var imageOutputSizes []string
32663295
if reqStream {
3296+
stopCompactKeepalive()
32673297
streamResult, err := s.handleStreamingResponse(ctx, resp, c, account, startTime, originalModel, upstreamModel)
32683298
if err != nil {
32693299
return nil, err
@@ -3274,7 +3304,7 @@ func (s *OpenAIGatewayService) Forward(ctx context.Context, c *gin.Context, acco
32743304
imageCount = streamResult.imageCount
32753305
imageOutputSizes = streamResult.imageOutputSizes
32763306
} else {
3277-
nonStreamResult, err := s.handleNonStreamingResponse(ctx, resp, c, account, originalModel, upstreamModel)
3307+
nonStreamResult, err := s.handleNonStreamingResponse(ctx, resp, c, account, originalModel, upstreamModel, stopCompactKeepalive)
32783308
if err != nil {
32793309
return nil, err
32803310
}
@@ -3885,6 +3915,16 @@ func (s *OpenAIGatewayService) startCompactNonstreamKeepalive(ctx context.Contex
38853915
if ctx == nil {
38863916
ctx = context.Background()
38873917
}
3918+
path := ""
3919+
if c.Request != nil && c.Request.URL != nil {
3920+
path = strings.TrimSpace(c.Request.URL.Path)
3921+
}
3922+
log := logger.FromContext(ctx).With(
3923+
zap.String("component", "service.openai_gateway"),
3924+
zap.String("request_path", path),
3925+
zap.Int("interval_seconds", int(interval.Seconds())),
3926+
)
3927+
log.Info("OpenAI compact non-stream keepalive started")
38883928

38893929
headers := c.Writer.Header()
38903930
headers.Set("Content-Type", "application/json")
@@ -3900,15 +3940,23 @@ func (s *OpenAIGatewayService) startCompactNonstreamKeepalive(ctx context.Contex
39003940
defer wg.Done()
39013941
ticker := time.NewTicker(interval)
39023942
defer ticker.Stop()
3943+
flushedLogged := false
39033944
for {
39043945
select {
39053946
case <-stopCh:
39063947
return
39073948
case <-ctx.Done():
39083949
return
39093950
case <-ticker.C:
3910-
_, _ = c.Writer.Write([]byte("\n"))
3951+
if _, err := c.Writer.Write([]byte("\n")); err != nil {
3952+
log.Warn("OpenAI compact non-stream keepalive write failed", zap.Error(err))
3953+
return
3954+
}
39113955
flusher.Flush()
3956+
if !flushedLogged {
3957+
log.Info("OpenAI compact non-stream keepalive flushed")
3958+
flushedLogged = true
3959+
}
39123960
}
39133961
}
39143962
}()
@@ -5699,17 +5747,22 @@ func openAIUsageFromGJSON(value gjson.Result) (OpenAIUsage, bool) {
56995747
}, true
57005748
}
57015749

5702-
func (s *OpenAIGatewayService) handleNonStreamingResponse(ctx context.Context, resp *http.Response, c *gin.Context, account *Account, originalModel, mappedModel string) (*openaiNonStreamingResult, error) {
5703-
body, err := ReadUpstreamResponseBody(resp.Body, s.cfg, c, openAITooLargeError)
5750+
func (s *OpenAIGatewayService) handleNonStreamingResponse(ctx context.Context, resp *http.Response, c *gin.Context, account *Account, originalModel, mappedModel string, stopBeforeWrite ...func()) (*openaiNonStreamingResult, error) {
5751+
stop := compactStopFunc(stopBeforeWrite...)
5752+
body, err := ReadUpstreamResponseBody(resp.Body, s.cfg, c, func(c *gin.Context) {
5753+
stop()
5754+
openAITooLargeError(c)
5755+
})
57045756
if err != nil {
5757+
stop()
57055758
return nil, err
57065759
}
57075760

57085761
// Detect SSE responses for ALL account types via Content-Type header.
57095762
// Some OpenAI-compatible upstreams (including other sub2api instances)
57105763
// may return SSE even when stream=false was requested.
57115764
if isEventStreamResponse(resp.Header) {
5712-
return s.handleSSEToJSON(resp, c, body, originalModel, mappedModel)
5765+
return s.handleSSEToJSON(resp, c, body, originalModel, mappedModel, stop)
57135766
}
57145767
bodyLooksLikeSSE := bytes.Contains(body, []byte("data:")) || bytes.Contains(body, []byte("event:"))
57155768

@@ -5719,14 +5772,15 @@ func (s *OpenAIGatewayService) handleNonStreamingResponse(ctx context.Context, r
57195772
// positives on JSON responses that coincidentally contain "data:" or
57205773
// "event:" in their text content.
57215774
if account.Type == AccountTypeOAuth && bodyLooksLikeSSE {
5722-
return s.handleSSEToJSON(resp, c, body, originalModel, mappedModel)
5775+
return s.handleSSEToJSON(resp, c, body, originalModel, mappedModel, stop)
57235776
}
57245777

57255778
usageValue, usageOK := extractOpenAIUsageFromJSONBytes(body)
57265779
if !usageOK {
57275780
if bodyLooksLikeSSE {
5728-
return s.handleSSEToJSON(resp, c, body, originalModel, mappedModel)
5781+
return s.handleSSEToJSON(resp, c, body, originalModel, mappedModel, stop)
57295782
}
5783+
stop()
57305784
return nil, fmt.Errorf("parse response: invalid json response")
57315785
}
57325786
usage := &usageValue
@@ -5745,6 +5799,7 @@ func (s *OpenAIGatewayService) handleNonStreamingResponse(ctx context.Context, r
57455799
}
57465800
}
57475801

5802+
stop()
57485803
c.Data(resp.StatusCode, contentType, body)
57495804

57505805
return &openaiNonStreamingResult{
@@ -5756,12 +5811,20 @@ func (s *OpenAIGatewayService) handleNonStreamingResponse(ctx context.Context, r
57565811
}, nil
57575812
}
57585813

5814+
func compactStopFunc(stops ...func()) func() {
5815+
if len(stops) == 0 || stops[0] == nil {
5816+
return func() {}
5817+
}
5818+
return stops[0]
5819+
}
5820+
57595821
func isEventStreamResponse(header http.Header) bool {
57605822
contentType := strings.ToLower(header.Get("Content-Type"))
57615823
return strings.Contains(contentType, "text/event-stream")
57625824
}
57635825

5764-
func (s *OpenAIGatewayService) handleSSEToJSON(resp *http.Response, c *gin.Context, body []byte, originalModel, mappedModel string) (*openaiNonStreamingResult, error) {
5826+
func (s *OpenAIGatewayService) handleSSEToJSON(resp *http.Response, c *gin.Context, body []byte, originalModel, mappedModel string, stopBeforeWrite ...func()) (*openaiNonStreamingResult, error) {
5827+
stop := compactStopFunc(stopBeforeWrite...)
57655828
bodyText := string(body)
57665829
finalResponse, ok := extractCodexFinalResponse(bodyText)
57675830

@@ -5793,6 +5856,7 @@ func (s *OpenAIGatewayService) handleSSEToJSON(resp *http.Response, c *gin.Conte
57935856
if msg == "" {
57945857
msg = "Upstream compact response failed"
57955858
}
5859+
stop()
57965860
return nil, s.writeOpenAINonStreamingProtocolError(resp, c, msg)
57975861
}
57985862
usage = s.parseSSEUsageFromBody(bodyText)
@@ -5811,6 +5875,7 @@ func (s *OpenAIGatewayService) handleSSEToJSON(resp *http.Response, c *gin.Conte
58115875
contentType = "text/event-stream"
58125876
}
58135877
}
5878+
stop()
58145879
c.Data(resp.StatusCode, contentType, body)
58155880

58165881
return &openaiNonStreamingResult{

backend/internal/service/openai_oauth_passthrough_test.go

Lines changed: 117 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -49,12 +49,12 @@ func (u *httpUpstreamRecorder) Do(req *http.Request, proxyURL string, accountID
4949
req.Body = io.NopCloser(bytes.NewReader(b))
5050
}
5151
u.requests = append(u.requests, req)
52-
if u.err != nil {
53-
return nil, u.err
54-
}
5552
if u.delay > 0 {
5653
time.Sleep(u.delay)
5754
}
55+
if u.err != nil {
56+
return nil, u.err
57+
}
5858
if len(u.responses) > 0 {
5959
resp := u.responses[0]
6060
u.responses = u.responses[1:]
@@ -499,6 +499,55 @@ func TestOpenAIGatewayService_OAuthPassthrough_CompactNonstreamKeepaliveWritesLe
499499
require.Contains(t, string(body), `"id":"cmp_keepalive"`)
500500
}
501501

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

@@ -625,6 +674,71 @@ func TestOpenAIGatewayService_OAuthPassthrough_CompactNonstreamKeepaliveCommitte
625674
}
626675
}
627676

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

0 commit comments

Comments
 (0)