Skip to content

Commit a62927b

Browse files
DeliciousBudingDeliciousBuding
committed
feat(proxy): align codex header and continuity defaults
1 parent f0e6c17 commit a62927b

9 files changed

Lines changed: 304 additions & 59 deletions

File tree

‎main.go‎

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -149,10 +149,8 @@ func main() {
149149
r.Use(security.RequestSizeLimiter(security.MaxRequestBodySize))
150150

151151
// handler 不再接收 cfg.APIKeys
152-
// 从环境变量读取设备指纹配置（后续可从数据库配置）
153-
deviceCfg := &proxy.DeviceProfileConfig{
154-
StabilizeDeviceProfile: os.Getenv("STABILIZE_DEVICE_PROFILE") == "true",
155-
}
152+
// 从环境变量读取 Codex 画像与 Beta 配置。
153+
deviceCfg := proxy.DeviceProfileConfigFromEnv(os.Getenv)
156154
handler := proxy.NewHandler(store, db, cfg, deviceCfg)
157155

158156
// 注册 WebSocket 执行函数（避免 proxy ↔ wsrelay 循环依赖）

‎proxy/device_profile.go‎

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,7 @@ type DeviceProfileConfig struct {
4040
OS string
4141
Arch string
4242
StabilizeDeviceProfile bool
43+
BetaFeatures string
4344
}
4445

4546
// CLIVersion 表示 Codex CLI 版本
@@ -95,6 +96,25 @@ func IsDeviceProfileStabilizationEnabled(cfg *DeviceProfileConfig) bool {
9596
return cfg.StabilizeDeviceProfile
9697
}
9798

99+
// DeviceProfileConfigFromEnv loads Codex header defaults from environment variables.
100+
func DeviceProfileConfigFromEnv(lookup func(string) string) *DeviceProfileConfig {
101+
if lookup == nil {
102+
lookup = func(string) string { return "" }
103+
}
104+
trimmed := func(key string) string {
105+
return strings.TrimSpace(lookup(key))
106+
}
107+
return &DeviceProfileConfig{
108+
UserAgent: trimmed("CODEX_USER_AGENT"),
109+
PackageVersion: trimmed("CODEX_PACKAGE_VERSION"),
110+
RuntimeVersion: trimmed("CODEX_RUNTIME_VERSION"),
111+
OS: trimmed("CODEX_OS"),
112+
Arch: trimmed("CODEX_ARCH"),
113+
StabilizeDeviceProfile: strings.EqualFold(trimmed("STABILIZE_DEVICE_PROFILE"), "true"),
114+
BetaFeatures: trimmed("CODEX_BETA_FEATURES"),
115+
}
116+
}
117+
98118
func defaultDeviceProfile(cfg *DeviceProfileConfig) deviceProfile {
99119
// 如果 cfg 为 nil，使用默认空配置兜底
100120
if cfg == nil {

‎proxy/device_profile_env_test.go‎

Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,51 @@
1+
package proxy
2+
3+
import "testing"
4+
5+
func TestDeviceProfileConfigFromEnv(t *testing.T) {
6+
cfg := DeviceProfileConfigFromEnv(func(key string) string {
7+
switch key {
8+
case "STABILIZE_DEVICE_PROFILE":
9+
return "true"
10+
case "CODEX_USER_AGENT":
11+
return "codex_cli_rs/0.120.0 (Mac OS 15.5.0; arm64) Apple_Terminal/464"
12+
case "CODEX_PACKAGE_VERSION":
13+
return "0.120.0"
14+
case "CODEX_RUNTIME_VERSION":
15+
return "0.120.0"
16+
case "CODEX_OS":
17+
return "MacOS"
18+
case "CODEX_ARCH":
19+
return "arm64"
20+
case "CODEX_BETA_FEATURES":
21+
return "multi_agent"
22+
default:
23+
return ""
24+
}
25+
})
26+
27+
if cfg == nil {
28+
t.Fatal("expected config")
29+
}
30+
if !cfg.StabilizeDeviceProfile {
31+
t.Fatal("expected device profile stabilization to be enabled")
32+
}
33+
if cfg.UserAgent != "codex_cli_rs/0.120.0 (Mac OS 15.5.0; arm64) Apple_Terminal/464" {
34+
t.Fatalf("UserAgent = %q", cfg.UserAgent)
35+
}
36+
if cfg.PackageVersion != "0.120.0" {
37+
t.Fatalf("PackageVersion = %q", cfg.PackageVersion)
38+
}
39+
if cfg.RuntimeVersion != "0.120.0" {
40+
t.Fatalf("RuntimeVersion = %q", cfg.RuntimeVersion)
41+
}
42+
if cfg.OS != "MacOS" {
43+
t.Fatalf("OS = %q", cfg.OS)
44+
}
45+
if cfg.Arch != "arm64" {
46+
t.Fatalf("Arch = %q", cfg.Arch)
47+
}
48+
if cfg.BetaFeatures != "multi_agent" {
49+
t.Fatalf("BetaFeatures = %q", cfg.BetaFeatures)
50+
}
51+
}

‎proxy/executor.go‎

Lines changed: 75 additions & 40 deletions
Original file line numberDiff line numberDiff line change
@@ -132,7 +132,7 @@ const (
132132
)
133133

134134
// WebsocketExecuteFunc WebSocket 执行函数（由 wsrelay 包在 main.go 中注册，避免循环依赖）
135-
var WebsocketExecuteFunc func(ctx context.Context, account *auth.Account, requestBody []byte, sessionID string, proxyOverride string) (*http.Response, error)
135+
var WebsocketExecuteFunc func(ctx context.Context, account *auth.Account, requestBody []byte, sessionID string, proxyOverride string, apiKey string, deviceCfg *DeviceProfileConfig, headers http.Header) (*http.Response, error)
136136

137137
// ExecuteRequest 向 Codex 上游发送请求
138138
// sessionID 可选，用于 prompt cache 会话绑定
@@ -141,7 +141,7 @@ var WebsocketExecuteFunc func(ctx context.Context, account *auth.Account, reques
141141
func ExecuteRequest(ctx context.Context, account *auth.Account, requestBody []byte, sessionID string, proxyOverride string, apiKey string, deviceCfg *DeviceProfileConfig, headers http.Header, useWebsocket ...bool) (*http.Response, error) {
142142
// 检查是否使用 WebSocket
143143
if len(useWebsocket) > 0 && useWebsocket[0] && WebsocketExecuteFunc != nil {
144-
return WebsocketExecuteFunc(ctx, account, requestBody, sessionID, proxyOverride)
144+
return WebsocketExecuteFunc(ctx, account, requestBody, sessionID, proxyOverride, apiKey, deviceCfg, headers)
145145
}
146146

147147
if ctx == nil {
@@ -150,7 +150,6 @@ func ExecuteRequest(ctx context.Context, account *auth.Account, requestBody []by
150150

151151
account.Mu().RLock()
152152
accessToken := account.AccessToken
153-
accountID := account.AccountID
154153
proxyURL := account.ProxyURL
155154
account.Mu().RUnlock()
156155

@@ -193,67 +192,103 @@ func ExecuteRequest(ctx context.Context, account *auth.Account, requestBody []by
193192
}
194193

195194
// ==================== 请求头（伪装 Codex CLI） ====================
196-
// 应用设备指纹稳定化
195+
applyCodexRequestHeaders(req, account, accessToken, cacheKey, apiKey, deviceCfg, headers)
196+
197+
// 获取连接池 HTTP 客户端（账号级隔离，复用 TCP/TLS 连接）
198+
client := getPooledClient(account, proxyURL)
199+
200+
resp, err := client.Do(req)
201+
if err != nil {
202+
if shouldRecyclePooledClient(err) {
203+
recyclePooledClient(account, proxyURL)
204+
}
205+
return nil, ErrUpstream(0, "请求上游失败", err)
206+
}
207+
208+
return resp, nil
209+
}
210+
211+
func codexVersionFromProfile(profile deviceProfile, fallback string) string {
212+
if profile.HasVersion {
213+
return fmt.Sprintf("%d.%d.%d", profile.Version.major, profile.Version.minor, profile.Version.patch)
214+
}
215+
return strings.TrimSpace(fallback)
216+
}
217+
218+
func applyCodexRequestHeaders(req *http.Request, account *auth.Account, accessToken, cacheKey, apiKey string, deviceCfg *DeviceProfileConfig, downstreamHeaders http.Header) {
219+
if req == nil {
220+
return
221+
}
222+
223+
accountID := ""
224+
if account != nil {
225+
account.Mu().RLock()
226+
accountID = account.AccountID
227+
account.Mu().RUnlock()
228+
}
229+
230+
var profile deviceProfile
231+
version := ""
197232
if IsDeviceProfileStabilizationEnabled(deviceCfg) {
198-
profile := ResolveDeviceProfile(account, apiKey, headers, deviceCfg)
233+
profile = ResolveDeviceProfile(account, apiKey, downstreamHeaders, deviceCfg)
199234
ApplyDeviceProfileHeaders(req, profile)
200-
// 稳定化时也需要设置 Version 头，保持行为一致
201-
if profile.HasVersion {
202-
req.Header.Set("Version", fmt.Sprintf("%d.%d.%d", profile.Version.major, profile.Version.minor, profile.Version.patch))
203-
}
235+
version = codexVersionFromProfile(profile, strings.TrimSpace(deviceCfg.PackageVersion))
204236
} else {
205-
// 每个账号使用确定性的 ClientProfile（UA + Version），模拟真实用户多样性
206-
profile := ProfileForAccount(account.ID())
207-
req.Header.Set("User-Agent", profile.UserAgent)
208-
req.Header.Set("Version", profile.Version)
237+
clientProfile := ProfileForAccount(account.ID())
238+
req.Header.Set("User-Agent", clientProfile.UserAgent)
239+
version = clientProfile.Version
209240
}
210241

211242
req.Header.Set("Authorization", "Bearer "+accessToken)
212243
req.Header.Set("Content-Type", "application/json")
213244
req.Header.Set("Accept", "text/event-stream")
214-
req.Header.Set("Originator", Originator)
215245
req.Header.Set("Connection", "Keep-Alive")
246+
if version != "" {
247+
req.Header.Set("Version", version)
248+
}
249+
if originator := strings.TrimSpace(downstreamHeaders.Get("Originator")); originator != "" {
250+
req.Header.Set("Originator", originator)
251+
} else {
252+
req.Header.Set("Originator", Originator)
253+
}
216254
if accountID != "" {
217255
req.Header.Set("Chatgpt-Account-Id", accountID)
218256
}
219-
220-
// Session/Conversation 头（用于 prompt cache 绑定）
221-
// 参考 CLIProxyAPI: req.Header.Set("Conversation_id", cache.ID)
222-
// 参考 sub2api: headers.Set("session_id", sessionResolution.SessionID)
223257
if cacheKey != "" {
224258
req.Header.Set("Session_id", cacheKey)
225-
req.Header.Set("Conversation_id", cacheKey)
259+
req.Header.Del("Conversation_id")
226260
}
227-
228-
// 获取连接池 HTTP 客户端（账号级隔离，复用 TCP/TLS 连接）
229-
client := getPooledClient(account, proxyURL)
230-
231-
resp, err := client.Do(req)
232-
if err != nil {
233-
if shouldRecyclePooledClient(err) {
234-
recyclePooledClient(account, proxyURL)
235-
}
236-
return nil, ErrUpstream(0, "请求上游失败", err)
237-
}
238-
239-
return resp, nil
240261
}
241262

242263
// ResolveSessionID 从下游请求提取或生成 session ID
243-
// 优先级（参考 sub2api）：
244-
// 1. Header: session_id
245-
// 2. Header: conversation_id
246-
// 3. Body: prompt_cache_key
247-
// 4. 基于 Bearer API Key 的确定性 UUID（参考 CLIProxyAPI）
248-
func ResolveSessionID(authHeader string, body []byte) string {
249-
// 此函数由 handler 调用，将 gin.Context 的 header 传进来
250-
264+
// 优先级：
265+
// 1. Header: Session_id
266+
// 2. Header: Conversation_id
267+
// 3. Header: Idempotency-Key
268+
// 4. Body: prompt_cache_key
269+
// 5. 基于 Bearer API Key 的确定性 UUID
270+
func ResolveSessionID(headers http.Header, body []byte) string {
271+
if headers != nil {
272+
if v := strings.TrimSpace(headers.Get("Session_id")); v != "" {
273+
return v
274+
}
275+
if v := strings.TrimSpace(headers.Get("Conversation_id")); v != "" {
276+
return v
277+
}
278+
if v := strings.TrimSpace(headers.Get("Idempotency-Key")); v != "" {
279+
return v
280+
}
281+
}
251282
// 优先从 body 的 prompt_cache_key 提取
252283
if v := strings.TrimSpace(gjson.GetBytes(body, "prompt_cache_key").String()); v != "" {
253284
return v
254285
}
255286

256287
// 基于下游用户的 API Key 生成确定性 cache key（参考 CLIProxyAPI codex_executor.go:621）
288+
authHeader := ""
289+
if headers != nil {
290+
authHeader = headers.Get("Authorization")
291+
}
257292
apiKey := strings.TrimPrefix(authHeader, "Bearer ")
258293
apiKey = strings.TrimSpace(apiKey)
259294
if apiKey != "" {

‎proxy/executor_test.go‎

Lines changed: 73 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,8 +3,11 @@ package proxy
33
import (
44
"context"
55
"errors"
6+
"net/http"
67
"strings"
78
"testing"
9+
10+
"github.com/codex2api/auth"
811
)
912

1013
func TestReadSSEStream_MergesMultilineData(t *testing.T) {
@@ -140,3 +143,73 @@ func TestShouldTransparentRetryStream(t *testing.T) {
140143
t.Fatal("expected retry to stop when downstream context is canceled")
141144
}
142145
}
146+
147+
func TestApplyCodexRequestHeadersUsesSessionIDWithoutConversationID(t *testing.T) {
148+
req, err := http.NewRequest(http.MethodPost, "https://example.com/v1/responses", nil)
149+
if err != nil {
150+
t.Fatalf("http.NewRequest() error = %v", err)
151+
}
152+
153+
acc := &auth.Account{
154+
DBID: 42,
155+
AccountID: "acct-42",
156+
}
157+
cfg := &DeviceProfileConfig{
158+
UserAgent: "codex_cli_rs/0.120.0 (Mac OS 15.5.0; arm64) Apple_Terminal/464",
159+
PackageVersion: "0.120.0",
160+
RuntimeVersion: "0.120.0",
161+
OS: "MacOS",
162+
Arch: "arm64",
163+
StabilizeDeviceProfile: true,
164+
}
165+
downstreamHeaders := http.Header{
166+
"Originator": []string{"custom-originator"},
167+
}
168+
169+
applyCodexRequestHeaders(req, acc, "token-123", "cache-key-1", "api-key-1", cfg, downstreamHeaders)
170+
171+
if got := req.Header.Get("Authorization"); got != "Bearer token-123" {
172+
t.Fatalf("Authorization = %q", got)
173+
}
174+
if got := req.Header.Get("Session_id"); got != "cache-key-1" {
175+
t.Fatalf("Session_id = %q", got)
176+
}
177+
if got := req.Header.Get("Conversation_id"); got != "" {
178+
t.Fatalf("Conversation_id = %q, want empty", got)
179+
}
180+
if got := req.Header.Get("User-Agent"); got != cfg.UserAgent {
181+
t.Fatalf("User-Agent = %q", got)
182+
}
183+
if got := req.Header.Get("Version"); got != "0.120.0" {
184+
t.Fatalf("Version = %q", got)
185+
}
186+
if got := req.Header.Get("Originator"); got != "custom-originator" {
187+
t.Fatalf("Originator = %q", got)
188+
}
189+
if got := req.Header.Get("Chatgpt-Account-Id"); got != "acct-42" {
190+
t.Fatalf("Chatgpt-Account-Id = %q", got)
191+
}
192+
}
193+
194+
func TestResolveSessionIDPrefersContinuityHeaders(t *testing.T) {
195+
headers := http.Header{
196+
"Session_id": []string{"session-from-header"},
197+
"Conversation_id": []string{"conversation-from-header"},
198+
"Authorization": []string{"Bearer sk-test-123"},
199+
}
200+
201+
if got := ResolveSessionID(headers, []byte(`{"prompt_cache_key":"body-key"}`)); got != "session-from-header" {
202+
t.Fatalf("ResolveSessionID() = %q, want %q", got, "session-from-header")
203+
}
204+
205+
headers.Del("Session_id")
206+
if got := ResolveSessionID(headers, []byte(`{"prompt_cache_key":"body-key"}`)); got != "conversation-from-header" {
207+
t.Fatalf("ResolveSessionID() = %q, want %q", got, "conversation-from-header")
208+
}
209+
210+
headers.Del("Conversation_id")
211+
headers.Set("Idempotency-Key", "idempotency-key-1")
212+
if got := ResolveSessionID(headers, []byte(`{"prompt_cache_key":"body-key"}`)); got != "idempotency-key-1" {
213+
t.Fatalf("ResolveSessionID() = %q, want %q", got, "idempotency-key-1")
214+
}
215+
}

‎proxy/handler.go‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -426,7 +426,7 @@ func (h *Handler) Responses(c *gin.Context) {
426426

427427
rawBody = normalizeServiceTierField(rawBody)
428428
isStream := gjson.GetBytes(rawBody, "stream").Bool()
429-
sessionID := ResolveSessionID(c.GetHeader("Authorization"), rawBody)
429+
sessionID := ResolveSessionID(c.Request.Header, rawBody)
430430
reasoningEffort := extractReasoningEffort(rawBody)
431431
serviceTier := extractServiceTier(rawBody)
432432
if serviceTier != "" {
@@ -798,7 +798,7 @@ func (h *Handler) ChatCompletions(c *gin.Context) {
798798
return
799799
}
800800

801-
sessionID := ResolveSessionID(c.GetHeader("Authorization"), codexBody)
801+
sessionID := ResolveSessionID(c.Request.Header, codexBody)
802802

803803
// 3. 带重试的上游请求
804804
maxRetries := h.getMaxRetries()

‎proxy/handler_anthropic.go‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -109,7 +109,7 @@ func (h *Handler) Messages(c *gin.Context) {
109109

110110
// 提取 reasoning effort（从翻译后的 codex body 中）
111111
reasoningEffort := extractReasoningEffort(codexBody)
112-
sessionID := ResolveSessionID(c.GetHeader("Authorization"), codexBody)
112+
sessionID := ResolveSessionID(c.Request.Header, codexBody)
113113

114114
// 3. 带重试的上游请求
115115
maxRetries := h.getMaxRetries()
@@ -424,4 +424,3 @@ func (h *Handler) Messages(c *gin.Context) {
424424
sendAnthropicError(c, lastStatusCode, errType, fmt.Sprintf("Upstream returned status %d", lastStatusCode))
425425
}
426426
}
427-

0 commit comments

Comments
 (0)