Skip to content

Commit ed10067

Browse files
clcc2019clcc2019
committed
feat(config): add request compression support for Codex headers
- Introduced `request-compression` field in `CodexHeaderDefaults` to enable zstd compression for requests. - Updated tests to validate the new configuration option and its expected behavior. - Refactored payload handling in Codex and Antigravity executors to utilize lazy loading for original payloads, improving performance and memory usage.
1 parent fbfebe4 commit ed10067

21 files changed

Lines changed: 1435 additions & 227 deletions

‎internal/config/codex_websocket_header_defaults_test.go‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@ func TestLoadConfigOptional_CodexHeaderDefaults(t *testing.T) {
1313
codex-header-defaults:
1414
user-agent: " my-codex-client/1.0 "
1515
beta-features: " feature-a,feature-b "
16+
request-compression: true
1617
`)
1718
if err := os.WriteFile(configPath, configYAML, 0o600); err != nil {
1819
t.Fatalf("failed to write config: %v", err)
@@ -29,4 +30,7 @@ codex-header-defaults:
2930
if got := cfg.CodexHeaderDefaults.BetaFeatures; got != "feature-a,feature-b" {
3031
t.Fatalf("BetaFeatures = %q, want %q", got, "feature-a,feature-b")
3132
}
33+
if cfg.CodexHeaderDefaults.RequestCompression == nil || !*cfg.CodexHeaderDefaults.RequestCompression {
34+
t.Fatal("RequestCompression = nil/false, want true")
35+
}
3236
}

‎internal/config/config.go‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -159,9 +159,12 @@ type ClaudeHeaderDefaults struct {
159159
// requests when neither the incoming request nor the selected credential supplies them.
160160
// UserAgent applies to both auth-file and codex-api-key HTTP/websocket requests.
161161
// BetaFeatures only applies to websocket requests for auth-file/OAuth Codex sessions.
162+
// RequestCompression enables zstd request-body compression for official ChatGPT Codex
163+
// backend HTTP requests when using OAuth/auth-file sessions.
162164
type CodexHeaderDefaults struct {
163-
UserAgent string `yaml:"user-agent" json:"user-agent"`
164-
BetaFeatures string `yaml:"beta-features" json:"beta-features"`
165+
UserAgent string `yaml:"user-agent" json:"user-agent"`
166+
BetaFeatures string `yaml:"beta-features" json:"beta-features"`
167+
RequestCompression *bool `yaml:"request-compression,omitempty" json:"request-compression,omitempty"`
165168
}
166169

167170
// TLSConfig holds HTTPS server settings.

‎internal/runtime/executor/antigravity_executor.go‎

Lines changed: 64 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -92,6 +92,7 @@ var (
9292
antigravityCreditsFailureByAuth sync.Map
9393
antigravityPreferCreditsByModel sync.Map
9494
antigravityShortCooldownByAuth sync.Map
95+
antigravityRefreshFlights sync.Map
9596
antigravityQuotaExhaustedKeywords = []string{
9697
"quota_exhausted",
9798
"quota exhausted",
@@ -112,6 +113,12 @@ var (
112113
}
113114
)
114115

116+
type antigravityRefreshFlight struct {
117+
wg sync.WaitGroup
118+
updated *cliproxyauth.Auth
119+
err error
120+
}
121+
115122
// AntigravityExecutor proxies requests to the antigravity upstream.
116123
type AntigravityExecutor struct {
117124
cfg *config.Config
@@ -706,7 +713,6 @@ func (e *AntigravityExecutor) Execute(ctx context.Context, auth *cliproxyauth.Au
706713
if updatedAuth != nil {
707714
auth = updatedAuth
708715
}
709-
originalTranslated := sdktranslator.TranslateRequest(from, to, baseModel, originalPayload, false)
710716
translated := sdktranslator.TranslateRequest(from, to, baseModel, req.Payload, false)
711717

712718
translated, err = thinking.ApplyThinking(translated, req.Model, from.String(), to.String(), e.Identifier())
@@ -715,7 +721,13 @@ func (e *AntigravityExecutor) Execute(ctx context.Context, auth *cliproxyauth.Au
715721
}
716722

717723
requestedModel := helps.PayloadRequestedModel(opts, req.Model)
718-
translated = helps.ApplyPayloadConfigWithRoot(e.cfg, baseModel, "antigravity", "request", translated, originalTranslated, requestedModel)
724+
var originalTranslated []byte
725+
translated = helps.ApplyPayloadConfigWithRootLazy(e.cfg, baseModel, "antigravity", "request", translated, func() []byte {
726+
if len(originalTranslated) == 0 {
727+
originalTranslated = sdktranslator.TranslateRequest(from, to, baseModel, originalPayload, false)
728+
}
729+
return originalTranslated
730+
}, requestedModel)
719731

720732
baseURLs := antigravityBaseURLFallbackOrder(auth)
721733
httpClient := newAntigravityHTTPClient(ctx, e.cfg, auth, 0)
@@ -919,7 +931,6 @@ func (e *AntigravityExecutor) executeClaudeNonStream(ctx context.Context, auth *
919931
if updatedAuth != nil {
920932
auth = updatedAuth
921933
}
922-
originalTranslated := sdktranslator.TranslateRequest(from, to, baseModel, originalPayload, true)
923934
translated := sdktranslator.TranslateRequest(from, to, baseModel, req.Payload, true)
924935

925936
translated, err = thinking.ApplyThinking(translated, req.Model, from.String(), to.String(), e.Identifier())
@@ -928,7 +939,13 @@ func (e *AntigravityExecutor) executeClaudeNonStream(ctx context.Context, auth *
928939
}
929940

930941
requestedModel := helps.PayloadRequestedModel(opts, req.Model)
931-
translated = helps.ApplyPayloadConfigWithRoot(e.cfg, baseModel, "antigravity", "request", translated, originalTranslated, requestedModel)
942+
var originalTranslated []byte
943+
translated = helps.ApplyPayloadConfigWithRootLazy(e.cfg, baseModel, "antigravity", "request", translated, func() []byte {
944+
if len(originalTranslated) == 0 {
945+
originalTranslated = sdktranslator.TranslateRequest(from, to, baseModel, originalPayload, true)
946+
}
947+
return originalTranslated
948+
}, requestedModel)
932949

933950
baseURLs := antigravityBaseURLFallbackOrder(auth)
934951
httpClient := newAntigravityHTTPClient(ctx, e.cfg, auth, 0)
@@ -1383,7 +1400,6 @@ func (e *AntigravityExecutor) ExecuteStream(ctx context.Context, auth *cliproxya
13831400
if updatedAuth != nil {
13841401
auth = updatedAuth
13851402
}
1386-
originalTranslated := sdktranslator.TranslateRequest(from, to, baseModel, originalPayload, true)
13871403
translated := sdktranslator.TranslateRequest(from, to, baseModel, req.Payload, true)
13881404

13891405
translated, err = thinking.ApplyThinking(translated, req.Model, from.String(), to.String(), e.Identifier())
@@ -1392,7 +1408,13 @@ func (e *AntigravityExecutor) ExecuteStream(ctx context.Context, auth *cliproxya
13921408
}
13931409

13941410
requestedModel := helps.PayloadRequestedModel(opts, req.Model)
1395-
translated = helps.ApplyPayloadConfigWithRoot(e.cfg, baseModel, "antigravity", "request", translated, originalTranslated, requestedModel)
1411+
var originalTranslated []byte
1412+
translated = helps.ApplyPayloadConfigWithRootLazy(e.cfg, baseModel, "antigravity", "request", translated, func() []byte {
1413+
if len(originalTranslated) == 0 {
1414+
originalTranslated = sdktranslator.TranslateRequest(from, to, baseModel, originalPayload, true)
1415+
}
1416+
return originalTranslated
1417+
}, requestedModel)
13961418

13971419
baseURLs := antigravityBaseURLFallbackOrder(auth)
13981420
httpClient := newAntigravityHTTPClient(ctx, e.cfg, auth, 0)
@@ -1792,13 +1814,48 @@ func (e *AntigravityExecutor) ensureAccessToken(ctx context.Context, auth *clipr
17921814
refreshCtx = context.WithValue(refreshCtx, "cliproxy.roundtripper", rt)
17931815
}
17941816
}
1795-
updated, errRefresh := e.refreshToken(refreshCtx, auth.Clone())
1817+
updated, errRefresh := e.refreshTokenShared(refreshCtx, auth)
17961818
if errRefresh != nil {
17971819
return "", nil, errRefresh
17981820
}
17991821
return metaStringValue(updated.Metadata, "access_token"), updated, nil
18001822
}
18011823

1824+
func (e *AntigravityExecutor) refreshTokenShared(ctx context.Context, auth *cliproxyauth.Auth) (*cliproxyauth.Auth, error) {
1825+
if auth == nil {
1826+
return nil, statusErr{code: http.StatusUnauthorized, msg: "missing auth"}
1827+
}
1828+
authID := strings.TrimSpace(auth.ID)
1829+
if authID == "" {
1830+
return e.refreshToken(ctx, auth.Clone())
1831+
}
1832+
1833+
flight := &antigravityRefreshFlight{}
1834+
actual, loaded := antigravityRefreshFlights.LoadOrStore(authID, flight)
1835+
current, _ := actual.(*antigravityRefreshFlight)
1836+
if current == nil {
1837+
antigravityRefreshFlights.Delete(authID)
1838+
return e.refreshToken(ctx, auth.Clone())
1839+
}
1840+
if loaded {
1841+
current.wg.Wait()
1842+
if current.updated == nil {
1843+
return nil, current.err
1844+
}
1845+
return current.updated.Clone(), current.err
1846+
}
1847+
current.wg.Add(1)
1848+
1849+
defer antigravityRefreshFlights.Delete(authID)
1850+
defer current.wg.Done()
1851+
1852+
current.updated, current.err = e.refreshToken(ctx, auth.Clone())
1853+
if current.updated == nil {
1854+
return nil, current.err
1855+
}
1856+
return current.updated.Clone(), current.err
1857+
}
1858+
18021859
func (e *AntigravityExecutor) refreshToken(ctx context.Context, auth *cliproxyauth.Auth) (*cliproxyauth.Auth, error) {
18031860
if auth == nil {
18041861
return nil, statusErr{code: http.StatusUnauthorized, msg: "missing auth"}

0 commit comments

Comments
 (0)