Merge pull request #4269 from catoncat/agent/sub2api-agent-identity

feat(openai): support Codex Agent Identity authentication
This commit is contained in:
Wesley Liddick
2026-07-15 09:47:00 +08:00
committed by GitHub
36 changed files with 2294 additions and 139 deletions
+11 -11
View File
@@ -96,9 +96,6 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) {
usageLogRepository := repository.NewUsageLogRepository(client, db)
usageService := service.NewUsageService(usageLogRepository, userRepository, client, apiKeyAuthCacheInvalidator)
opsRepository := repository.NewOpsRepository(db)
batchImageRepository := repository.NewBatchImageRepository(db)
batchImageQueue := repository.NewBatchImageQueue(redisClient, configConfig)
batchImageDownloadLimiter := repository.NewBatchImageDownloadLimiter(redisClient, configConfig)
usageBillingRepository := repository.NewUsageBillingRepository(client, db)
gatewayCache := repository.NewGatewayCache(redisClient)
schedulerOutboxRepository := repository.NewSchedulerOutboxRepository(db)
@@ -137,11 +134,6 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) {
channelRepository := repository.NewChannelRepository(db)
channelService := service.NewChannelService(channelRepository, groupRepository, apiKeyAuthCacheInvalidator, pricingService)
modelPricingResolver := service.NewModelPricingResolver(channelService, billingService)
batchImageModelPricingResolver := service.ProvideBatchImageModelPricingResolver(modelPricingResolver)
batchImagePublicService := service.NewBatchImagePublicService(batchImageRepository, accountRepository, groupRepository, userGroupRateRepository, batchImageQueue, batchImageModelPricingResolver, usageBillingRepository, apiKeyAuthCacheInvalidator, configConfig)
batchImageDownloadService := service.NewBatchImageDownloadService(batchImageRepository, accountRepository, batchImageDownloadLimiter, configConfig)
batchImageCleanupService := service.ProvideBatchImageCleanupService(batchImageRepository, accountRepository, configConfig)
batchImageWorkerRuntime := service.ProvideBatchImageWorkerRuntime(batchImageRepository, accountRepository, batchImageQueue, usageBillingRepository, usageLogRepository, batchImageModelPricingResolver, apiKeyAuthCacheInvalidator, configConfig)
notificationEmailService := service.NewNotificationEmailService(settingRepository, emailService)
balanceNotifyService := service.ProvideBalanceNotifyService(emailService, settingRepository, accountRepository, notificationEmailService)
gatewayService := service.NewGatewayService(accountRepository, groupRepository, usageLogRepository, usageBillingRepository, userRepository, userSubscriptionRepository, userGroupRateRepository, gatewayCache, configConfig, schedulerSnapshotService, concurrencyService, billingService, rateLimitService, billingCacheService, identityService, httpUpstream, deferredService, claudeTokenProvider, sessionLimitCache, rpmCache, digestSessionStore, settingService, tlsFingerprintProfileService, channelService, modelPricingResolver, balanceNotifyService, serviceUserPlatformQuotaRepository)
@@ -191,10 +183,10 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) {
antigravityQuotaFetcher := service.NewAntigravityQuotaFetcher(proxyRepository)
grokQuotaFetcher := service.NewGrokQuotaFetcher()
grokQuotaService := service.ProvideGrokQuotaService(accountRepository, proxyRepository, grokTokenProvider, httpUpstream, usageLogRepository)
openAIQuotaService := service.ProvideOpenAIQuotaService(accountRepository, proxyRepository, openAITokenProvider, privacyClientFactory)
openAIQuotaService := service.ProvideOpenAIQuotaService(accountRepository, proxyRepository, openAITokenProvider, privacyClientFactory, openAIGatewayService)
usageCache := service.NewUsageCache()
accountUsageService := service.NewAccountUsageService(accountRepository, usageLogRepository, claudeUsageFetcher, geminiQuotaService, antigravityQuotaFetcher, grokQuotaFetcher, grokQuotaService, openAIQuotaService, usageCache, identityCache, tlsFingerprintProfileService)
accountTestService := service.NewAccountTestService(accountRepository, geminiTokenProvider, claudeTokenProvider, grokTokenProvider, antigravityGatewayService, httpUpstream, configConfig, tlsFingerprintProfileService)
accountUsageService := service.ProvideAccountUsageService(accountRepository, usageLogRepository, claudeUsageFetcher, geminiQuotaService, antigravityQuotaFetcher, grokQuotaFetcher, grokQuotaService, openAIQuotaService, usageCache, identityCache, tlsFingerprintProfileService, openAIGatewayService)
accountTestService := service.ProvideAccountTestService(accountRepository, geminiTokenProvider, claudeTokenProvider, grokTokenProvider, antigravityGatewayService, httpUpstream, configConfig, tlsFingerprintProfileService, openAIGatewayService)
crsSyncService := service.NewCRSSyncService(accountRepository, proxyRepository, oAuthService, openAIOAuthService, geminiOAuthService, configConfig)
accountHandler := admin.ProvideAccountHandler(adminService, oAuthService, openAIOAuthService, geminiOAuthService, antigravityOAuthService, rateLimitService, accountUsageService, accountTestService, concurrencyService, crsSyncService, sessionLimitCache, rpmCache, compositeTokenCacheInvalidator, grokQuotaService)
adminAnnouncementHandler := admin.NewAnnouncementHandler(announcementService)
@@ -267,6 +259,13 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) {
handlerPaymentHandler := handler.NewPaymentHandler(paymentService, paymentConfigService)
paymentWebhookHandler := handler.NewPaymentWebhookHandler(paymentService, registry)
availableChannelHandler := handler.NewAvailableChannelHandler(channelService, apiKeyService, settingService)
batchImageRepository := repository.NewBatchImageRepository(db)
batchImageQueue := repository.NewBatchImageQueue(redisClient, configConfig)
batchImageModelPricingResolver := service.ProvideBatchImageModelPricingResolver(modelPricingResolver)
batchImagePublicService := service.NewBatchImagePublicService(batchImageRepository, accountRepository, groupRepository, userGroupRateRepository, batchImageQueue, batchImageModelPricingResolver, usageBillingRepository, apiKeyAuthCacheInvalidator, configConfig)
batchImageDownloadLimiter := repository.NewBatchImageDownloadLimiter(redisClient, configConfig)
batchImageDownloadService := service.NewBatchImageDownloadService(batchImageRepository, accountRepository, batchImageDownloadLimiter, configConfig)
batchImageCleanupService := service.ProvideBatchImageCleanupService(batchImageRepository, accountRepository, configConfig)
batchImageHandler := handler.NewBatchImageHandler(batchImagePublicService, batchImageDownloadService, batchImageCleanupService)
idempotencyCoordinator := service.ProvideIdempotencyCoordinator(idempotencyRepository, configConfig)
idempotencyCleanupService := service.ProvideIdempotencyCleanupService(idempotencyRepository, configConfig)
@@ -285,6 +284,7 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) {
accountExpiryService := service.ProvideAccountExpiryService(accountRepository)
proxyExpiryService := service.ProvideProxyExpiryService(proxyRepository)
subscriptionExpiryService := service.ProvideSubscriptionExpiryService(userSubscriptionRepository, settingRepository, notificationEmailService, leaderLockCache, db)
batchImageWorkerRuntime := service.ProvideBatchImageWorkerRuntime(batchImageRepository, accountRepository, batchImageQueue, usageBillingRepository, usageLogRepository, batchImageModelPricingResolver, apiKeyAuthCacheInvalidator, configConfig)
scheduledTestRunnerService := service.ProvideScheduledTestRunnerService(scheduledTestPlanRepository, scheduledTestService, accountTestService, rateLimitService, configConfig)
paymentOrderExpiryService := service.ProvidePaymentOrderExpiryService(paymentService, leaderLockCache, db)
channelMonitorRunner := service.ProvideChannelMonitorRunner(channelMonitorService, settingService)
@@ -0,0 +1,47 @@
package admin
import (
"crypto/ed25519"
"crypto/rand"
"crypto/x509"
"encoding/base64"
"testing"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/stretchr/testify/require"
)
func TestNormalizeCodexImportEntryAcceptsAgentIdentityAuthJSON(t *testing.T) {
_, privateKey, err := ed25519.GenerateKey(rand.Reader)
require.NoError(t, err)
der, err := x509.MarshalPKCS8PrivateKey(privateKey)
require.NoError(t, err)
privateKeyBase64 := base64.StdEncoding.EncodeToString(der)
item, err := normalizeCodexImportEntry(codexImportEntry{
Index: 1,
Value: map[string]any{
"auth_mode": "agentIdentity",
"agent_identity": map[string]any{
"agent_runtime_id": "runtime-import",
"agent_private_key": privateKeyBase64,
"account_id": "account-import",
"chatgpt_user_id": "user-import",
"email": "agent@example.invalid",
"plan_type": "pro",
"chatgpt_account_is_fedramp": false,
},
},
})
require.NoError(t, err)
require.NotNil(t, item)
require.True(t, item.IsAgentIdentity)
require.Equal(t, service.OpenAIAuthModeAgentIdentity, item.Credentials["auth_mode"])
require.Equal(t, "runtime-import", item.Credentials["agent_runtime_id"])
require.Equal(t, privateKeyBase64, item.Credentials["agent_private_key"])
require.Equal(t, "account-import", item.Credentials["chatgpt_account_id"])
require.Equal(t, "user-import", item.Credentials["chatgpt_user_id"])
require.NotContains(t, item.Credentials, "access_token")
require.NotContains(t, item.Credentials, "refresh_token")
require.NotEmpty(t, item.WarningTexts)
}
@@ -72,20 +72,25 @@ type codexImportEntry struct {
}
type codexImportAccount struct {
Name string
AccessToken string
RefreshToken string
IDToken string
Email string
AccountID string
UserID string
PlanType string
Organization string
Credentials map[string]any
Extra map[string]any
TokenExpiresAt *time.Time
IdentityKeys []string
WarningTexts []string
Name string
AccessToken string
RefreshToken string
IDToken string
Email string
AccountID string
UserID string
PlanType string
Organization string
AgentRuntimeID string
AgentPrivateKey string
AgentTaskID string
AgentFedRAMP bool
IsAgentIdentity bool
Credentials map[string]any
Extra map[string]any
TokenExpiresAt *time.Time
IdentityKeys []string
WarningTexts []string
}
type codexJWTClaims struct {
@@ -496,6 +501,41 @@ func normalizeCodexImportEntry(entry codexImportEntry) (*codexImportAccount, err
case string:
item.AccessToken = strings.TrimSpace(raw)
case map[string]any:
if agentIdentity, ok := firstCodexMap(raw, []string{"agent_identity"}, []string{"agentIdentity"}); ok || strings.EqualFold(firstCodexString(raw, []string{"auth_mode"}, []string{"authMode"}), service.OpenAIAuthModeAgentIdentity) {
if !ok {
agentIdentity = raw
}
item.IsAgentIdentity = true
item.AgentRuntimeID = firstCodexString(agentIdentity, []string{"agent_runtime_id"}, []string{"agentRuntimeId"})
item.AgentPrivateKey = firstCodexString(agentIdentity, []string{"agent_private_key"}, []string{"agentPrivateKey"})
item.AgentTaskID = firstCodexString(agentIdentity, []string{"task_id"}, []string{"taskId"})
item.AccountID = firstCodexString(agentIdentity, []string{"account_id"}, []string{"accountId"})
item.UserID = firstCodexString(agentIdentity, []string{"chatgpt_user_id"}, []string{"chatgptUserId"})
item.Email = firstCodexString(agentIdentity, []string{"email"})
item.PlanType = firstCodexString(agentIdentity, []string{"plan_type"}, []string{"planType"})
item.AgentFedRAMP = firstCodexBool(agentIdentity, []string{"chatgpt_account_is_fedramp"}, []string{"chatgptAccountIsFedramp"})
if item.AgentRuntimeID == "" || item.AgentPrivateKey == "" || item.AccountID == "" || item.UserID == "" {
return nil, errors.New("agent identity 缺少必要字段")
}
if err := service.ValidateOpenAIAgentIdentityPrivateKey(item.AgentPrivateKey); err != nil {
return nil, errors.New("agent identity private key 格式无效")
}
item.Credentials["auth_mode"] = service.OpenAIAuthModeAgentIdentity
item.Credentials["agent_runtime_id"] = item.AgentRuntimeID
item.Credentials["agent_private_key"] = item.AgentPrivateKey
item.Credentials["chatgpt_account_id"] = item.AccountID
item.Credentials["chatgpt_user_id"] = item.UserID
item.Credentials["chatgpt_account_is_fedramp"] = item.AgentFedRAMP
setCodexCredentialIfNotEmpty(item.Credentials, "task_id", item.AgentTaskID)
setCodexCredentialIfNotEmpty(item.Credentials, "email", item.Email)
setCodexCredentialIfNotEmpty(item.Credentials, "plan_type", item.PlanType)
if item.AgentTaskID == "" {
item.WarningTexts = append(item.WarningTexts, "未包含 task_id,首次请求会使用现有 runtime 注册新 task")
}
item.IdentityKeys = buildCodexAgentIdentityKeys(item.AccountID, item.UserID, item.Email, item.AgentRuntimeID)
item.Name = buildCodexImportAccountName(item, entry.Index)
return item, nil
}
item.AccessToken = firstCodexString(raw,
[]string{"tokens", "access_token"},
[]string{"tokens", "accessToken"},
@@ -577,6 +617,9 @@ func normalizeCodexImportEntry(entry codexImportEntry) (*codexImportAccount, err
return nil, fmt.Errorf("第 %d 条格式不支持", entry.Index)
}
if item.IsAgentIdentity {
return item, nil
}
if item.AccessToken == "" {
return nil, errors.New("缺少 accessToken/access_token")
}
@@ -812,6 +855,9 @@ func sanitizeCodexImportCredentialExtras(input map[string]any) map[string]any {
"openai_auth_mode": {},
"token_type": {},
"chatgpt_account_is_fedramp": {},
"agent_runtime_id": {},
"agent_private_key": {},
"task_id": {},
}
out := make(map[string]any, len(input))
for key, value := range input {
@@ -842,6 +888,14 @@ func buildCodexImportIdentityKeys(accountID, userID, email, accessToken, refresh
return buildCodexStoredIdentityKeys(accountID, userID, email, accessToken)
}
func buildCodexAgentIdentityKeys(accountID, userID, email, runtimeID string) []string {
keys := buildCodexStoredIdentityKeys(accountID, userID, email, "")
if runtimeID = strings.TrimSpace(runtimeID); runtimeID != "" {
keys = append([]string{"agent:" + runtimeID}, keys...)
}
return keys
}
// buildCodexStoredIdentityKeys 生成存量账号索引键,保留 user/account 维度,
// 让 accessToken-only 账号后续升级为完整 OAuth 时仍能命中并更新原账号。
func buildCodexStoredIdentityKeys(accountID, userID, email, accessToken string) []string {
@@ -891,6 +945,10 @@ func (i *codexAccountIndex) Add(account service.Account) {
for _, key := range keys {
i.accountsByKey[key] = upsertCodexAccount(i.accountsByKey[key], account)
}
if runtimeID := codexCredentialString(account.Credentials, "agent_runtime_id"); runtimeID != "" {
key := "agent:" + runtimeID
i.accountsByKey[key] = upsertCodexAccount(i.accountsByKey[key], account)
}
}
func (i *codexAccountIndex) remove(accountID int64) {
@@ -1043,6 +1101,38 @@ func firstCodexString(obj map[string]any, paths ...[]string) string {
return ""
}
func firstCodexMap(obj map[string]any, paths ...[]string) (map[string]any, bool) {
for _, path := range paths {
value, ok := codexPathValue(obj, path)
if !ok || value == nil {
continue
}
if mapped, ok := value.(map[string]any); ok {
return mapped, true
}
}
return nil, false
}
func firstCodexBool(obj map[string]any, paths ...[]string) bool {
for _, path := range paths {
value, ok := codexPathValue(obj, path)
if !ok {
continue
}
switch value := value.(type) {
case bool:
return value
case string:
parsed, err := strconv.ParseBool(strings.TrimSpace(value))
if err == nil {
return parsed
}
}
}
return false
}
func copyCodexExtraString(obj map[string]any, extra map[string]any, key string, path []string) {
value := firstCodexString(obj, path)
if value != "" {
@@ -20,6 +20,7 @@ func TestRedactCredentials_StripsSensitiveKeysAndReportsStatus(t *testing.T) {
"aws_secret_access_key": "aws-secret",
"service_account_json": map[string]any{"private_key": "..."},
"private_key": "raw-key",
"agent_private_key": "agent-key-secret",
// 非敏感
"base_url": "https://api.example.com",
"model_mapping": map[string]any{"foo": "bar"},
@@ -35,6 +36,7 @@ func TestRedactCredentials_StripsSensitiveKeysAndReportsStatus(t *testing.T) {
require.NotContains(t, out, "aws_secret_access_key")
require.NotContains(t, out, "service_account_json")
require.NotContains(t, out, "private_key")
require.NotContains(t, out, "agent_private_key")
require.Equal(t, "https://api.example.com", out["base_url"])
require.Equal(t, map[string]any{"foo": "bar"}, out["model_mapping"])
@@ -47,6 +49,7 @@ func TestRedactCredentials_StripsSensitiveKeysAndReportsStatus(t *testing.T) {
require.True(t, status["has_aws_secret_access_key"])
require.True(t, status["has_service_account_json"])
require.True(t, status["has_private_key"])
require.True(t, status["has_agent_private_key"])
// 状态 map 不应携带非敏感键的 has_*
require.NotContains(t, status, "has_base_url")
@@ -84,6 +87,7 @@ func TestRedactCredentials_AllKnownSensitiveKeys(t *testing.T) {
"api_key", "session_key", "cookie",
"aws_secret_access_key", "aws_session_token",
"service_account_json", "service_account", "private_key",
"agent_private_key",
}
in := make(map[string]any, len(keys))
for _, k := range keys {
@@ -4,7 +4,7 @@ package service
// dto 层做响应脱敏、service 层做更新合并都引用此清单——新增凭证类型时务必同步。
var SensitiveCredentialKeys = []string{
// OAuth
"access_token", "refresh_token", "id_token",
"access_token", "refresh_token", "id_token", "agent_private_key",
// API Key 类
"api_key", "session_key", "cookie",
// 云服务凭据
+107 -18
View File
@@ -15,6 +15,7 @@ import (
"net/http/httptest"
"regexp"
"strings"
"sync"
"time"
"github.com/Wei-Shaw/sub2api/internal/config"
@@ -72,6 +73,8 @@ type AccountTestService struct {
httpUpstream HTTPUpstream
cfg *config.Config
tlsFPProfileService *TLSFingerprintProfileService
agentIdentityTaskMu sync.Mutex
agentIdentityWS agentIdentityWSConnectionInvalidator
}
// NewAccountTestService creates a new AccountTestService
@@ -544,9 +547,11 @@ func (s *AccountTestService) testOpenAIAccountConnection(c *gin.Context, account
if credentialAccount.IsOAuth() {
isOAuth = true
// OAuth - use Bearer token with ChatGPT internal API
authToken = credentialAccount.GetOpenAIAccessToken()
if authToken == "" {
// Agent Identity signs each request and does not retain the OAuth token.
if !credentialAccount.IsOpenAIAgentIdentity() {
authToken = credentialAccount.GetOpenAIAccessToken()
}
if authToken == "" && !credentialAccount.IsOpenAIAgentIdentity() {
return s.sendErrorAndEnd(c, "No access token available")
}
@@ -591,8 +596,11 @@ func (s *AccountTestService) testOpenAIAccountConnection(c *gin.Context, account
payload := createOpenAITestPayload(upstreamTestModelID, isOAuth)
payloadBytes, _ := json.Marshal(payload)
// Send test_start event
s.sendEvent(c, TestEvent{Type: "test_start", Model: testModelID})
// Send test_start event once. A task-invalid Agent Identity response may
// restart this probe after registering a replacement task.
if !agentIdentityTaskRecoveryWasTried(ctx) {
s.sendEvent(c, TestEvent{Type: "test_start", Model: testModelID})
}
req, err := http.NewRequestWithContext(ctx, "POST", apiURL, bytes.NewReader(payloadBytes))
if err != nil {
@@ -602,7 +610,19 @@ func (s *AccountTestService) testOpenAIAccountConnection(c *gin.Context, account
// Set common headers
req.Header.Set("Content-Type", "application/json")
req.Header.Set("Authorization", "Bearer "+authToken)
if credentialAccount.IsOpenAIAgentIdentity() {
authHeaders, authErr := buildAgentIdentityAuthenticationHeaders(ctx, s.accountRepo, s.agentIdentityWS, &s.agentIdentityTaskMu, credentialAccount)
if authErr != nil {
return s.sendErrorAndEnd(c, "Failed to build Agent Identity authentication")
}
for key, values := range authHeaders {
for _, value := range values {
req.Header.Add(key, value)
}
}
} else {
req.Header.Set("Authorization", "Bearer "+authToken)
}
// Set OAuth-specific headers for ChatGPT internal API
if isOAuth {
@@ -644,6 +664,15 @@ func (s *AccountTestService) testOpenAIAccountConnection(c *gin.Context, account
if resp.StatusCode != http.StatusOK {
body, _ := io.ReadAll(resp.Body)
body = redactAgentIdentitySensitiveBodyForAccount(ctx, s.accountRepo, credentialAccount, body)
if !agentIdentityTaskRecoveryWasTried(ctx) && credentialAccount.IsOpenAIAgentIdentity() && isAgentIdentityTaskInvalidHTTPResponse(resp.StatusCode, body) {
expectedTaskID := credentialAccount.GetCredential("task_id")
if err := ensureAgentIdentityTaskForAccount(ctx, s.accountRepo, s.agentIdentityWS, &s.agentIdentityTaskMu, credentialAccount, expectedTaskID); err != nil {
return s.sendErrorAndEnd(c, fmt.Sprintf("Agent Identity task recovery failed: %s", err.Error()))
}
c.Request = c.Request.WithContext(markAgentIdentityTaskRecoveryTried(ctx))
return s.testOpenAIAccountConnection(c, account, modelID, prompt, mode)
}
if resp.StatusCode == http.StatusTooManyRequests {
s.reconcileOpenAI429State(ctx, account, resp.Header, body)
}
@@ -715,7 +744,9 @@ func (s *AccountTestService) testGrokAccountConnection(c *gin.Context, account *
return s.sendErrorAndEnd(c, "Failed to create Grok test payload")
}
s.sendEvent(c, TestEvent{Type: "test_start", Model: testModelID})
if !agentIdentityTaskRecoveryWasTried(ctx) {
s.sendEvent(c, TestEvent{Type: "test_start", Model: testModelID})
}
req, err := http.NewRequestWithContext(ctx, http.MethodPost, apiURL, bytes.NewReader(payloadBytes))
if err != nil {
@@ -833,16 +864,26 @@ func (s *AccountTestService) testOpenAIChatCompletionsConnection(
// resulting capability state on the account.
func (s *AccountTestService) testOpenAICompactConnection(c *gin.Context, account *Account, testModelID string) error {
ctx := c.Request.Context()
credentialAccount := account
if account.IsShadow() {
resolved, err := resolveCredentialAccount(ctx, s.accountRepo, account)
if err != nil {
return s.sendErrorAndEnd(c, "Failed to resolve account credentials")
}
credentialAccount = resolved
}
authToken := ""
apiURL := ""
isOAuth := false
switch {
case account.IsOAuth():
case credentialAccount.IsOAuth():
isOAuth = true
authToken = account.GetOpenAIAccessToken()
if authToken == "" {
if !credentialAccount.IsOpenAIAgentIdentity() {
authToken = credentialAccount.GetOpenAIAccessToken()
}
if authToken == "" && !credentialAccount.IsOpenAIAgentIdentity() {
return s.sendErrorAndEnd(c, "No access token available")
}
apiURL = chatgptCodexAPIURL + "/compact"
@@ -871,7 +912,9 @@ func (s *AccountTestService) testOpenAICompactConnection(c *gin.Context, account
c.Writer.Flush()
payloadBytes, _ := json.Marshal(createOpenAICompactProbePayload(testModelID))
s.sendEvent(c, TestEvent{Type: "test_start", Model: testModelID})
if !agentIdentityTaskRecoveryWasTried(ctx) {
s.sendEvent(c, TestEvent{Type: "test_start", Model: testModelID})
}
req, err := http.NewRequestWithContext(ctx, "POST", apiURL, bytes.NewReader(payloadBytes))
if err != nil {
@@ -881,7 +924,19 @@ func (s *AccountTestService) testOpenAICompactConnection(c *gin.Context, account
req.Header.Set("Content-Type", "application/json")
req.Header.Set("Accept", "application/json")
req.Header.Set("Authorization", "Bearer "+authToken)
if credentialAccount.IsOpenAIAgentIdentity() {
authHeaders, authErr := buildAgentIdentityAuthenticationHeaders(ctx, s.accountRepo, s.agentIdentityWS, &s.agentIdentityTaskMu, credentialAccount)
if authErr != nil {
return s.sendErrorAndEnd(c, "Failed to build Agent Identity authentication")
}
for key, values := range authHeaders {
for _, value := range values {
req.Header.Add(key, value)
}
}
} else {
req.Header.Set("Authorization", "Bearer "+authToken)
}
req.Header.Set("OpenAI-Beta", "responses=experimental")
req.Header.Set("Originator", "codex_cli_rs")
req.Header.Set("User-Agent", codexCLIUserAgent)
@@ -892,7 +947,7 @@ func (s *AccountTestService) testOpenAICompactConnection(c *gin.Context, account
if isOAuth {
req.Host = "chatgpt.com"
setOpenAIChatGPTAccountHeaders(req.Header, account)
setOpenAIChatGPTAccountHeaders(req.Header, credentialAccount)
}
// 账号级请求头覆写:测试请求与真实转发保持一致的最终头
@@ -915,6 +970,15 @@ func (s *AccountTestService) testOpenAICompactConnection(c *gin.Context, account
defer func() { _ = resp.Body.Close() }()
body, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20))
body = redactAgentIdentitySensitiveBodyForAccount(ctx, s.accountRepo, credentialAccount, body)
if !agentIdentityTaskRecoveryWasTried(ctx) && credentialAccount.IsOpenAIAgentIdentity() && isAgentIdentityTaskInvalidHTTPResponse(resp.StatusCode, body) {
expectedTaskID := credentialAccount.GetCredential("task_id")
if err := ensureAgentIdentityTaskForAccount(ctx, s.accountRepo, s.agentIdentityWS, &s.agentIdentityTaskMu, credentialAccount, expectedTaskID); err != nil {
return s.sendErrorAndEnd(c, fmt.Sprintf("Agent Identity task recovery failed: %s", err.Error()))
}
c.Request = c.Request.WithContext(markAgentIdentityTaskRecoveryTried(ctx))
return s.testOpenAICompactConnection(c, account, testModelID)
}
if s.accountRepo != nil {
updates := buildOpenAICompactProbeExtraUpdates(resp, body, nil, time.Now())
@@ -1706,8 +1770,19 @@ func (s *AccountTestService) testOpenAIImageAPIKey(c *gin.Context, ctx context.C
// testOpenAIImageOAuth tests OpenAI image generation using an OAuth account via Codex /responses API.
func (s *AccountTestService) testOpenAIImageOAuth(c *gin.Context, ctx context.Context, account *Account, modelID, prompt string) error {
authToken := account.GetOpenAIAccessToken()
if authToken == "" {
credentialAccount := account
if account.IsShadow() {
resolved, err := resolveCredentialAccount(ctx, s.accountRepo, account)
if err != nil {
return s.sendErrorAndEnd(c, "Failed to resolve account credentials")
}
credentialAccount = resolved
}
authToken := ""
if !credentialAccount.IsOpenAIAgentIdentity() {
authToken = credentialAccount.GetOpenAIAccessToken()
}
if authToken == "" && !credentialAccount.IsOpenAIAgentIdentity() {
return s.sendErrorAndEnd(c, "No access token available")
}
@@ -1739,17 +1814,29 @@ func (s *AccountTestService) testOpenAIImageOAuth(c *gin.Context, ctx context.Co
}
req = req.WithContext(WithHTTPUpstreamProfile(req.Context(), HTTPUpstreamProfileOpenAI))
req.Host = "chatgpt.com"
req.Header.Set("Authorization", "Bearer "+authToken)
if credentialAccount.IsOpenAIAgentIdentity() {
authHeaders, authErr := buildAgentIdentityAuthenticationHeaders(ctx, s.accountRepo, s.agentIdentityWS, &s.agentIdentityTaskMu, credentialAccount)
if authErr != nil {
return s.sendErrorAndEnd(c, "Failed to build Agent Identity authentication")
}
for key, values := range authHeaders {
for _, value := range values {
req.Header.Add(key, value)
}
}
} else {
req.Header.Set("Authorization", "Bearer "+authToken)
}
req.Header.Set("Content-Type", "application/json")
req.Header.Set("Accept", "text/event-stream")
req.Header.Set("OpenAI-Beta", "responses=experimental")
req.Header.Set("originator", "codex_cli_rs")
if customUA := strings.TrimSpace(account.GetOpenAIUserAgent()); customUA != "" {
if customUA := strings.TrimSpace(credentialAccount.GetOpenAIUserAgent()); customUA != "" {
req.Header.Set("User-Agent", customUA)
} else {
req.Header.Set("User-Agent", codexCLIUserAgent)
}
setOpenAIChatGPTAccountHeaders(req.Header, account)
setOpenAIChatGPTAccountHeaders(req.Header, credentialAccount)
// 与真实转发一致:originator 与最终 User-Agent 首段配套(原 opencode 与 Codex UA 错配会 404,issue #3901)。
enforceCodexIdentityHeaders(req.Header)
@@ -1768,6 +1855,7 @@ func (s *AccountTestService) testOpenAIImageOAuth(c *gin.Context, ctx context.Co
}()
if resp.StatusCode >= 400 {
body, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20))
body = redactAgentIdentitySensitiveBodyForAccount(ctx, s.accountRepo, credentialAccount, body)
message := strings.TrimSpace(extractUpstreamErrorMessage(body))
if message == "" {
message = fmt.Sprintf("Responses API returned %d", resp.StatusCode)
@@ -1779,6 +1867,7 @@ func (s *AccountTestService) testOpenAIImageOAuth(c *gin.Context, ctx context.Co
if err != nil {
return s.sendErrorAndEnd(c, fmt.Sprintf("Failed to read image response: %s", err.Error()))
}
body = redactAgentIdentitySensitiveBodyForAccount(ctx, s.accountRepo, credentialAccount, body)
results, _, _, _, _, err := collectOpenAIImagesFromResponsesBody(body)
if err != nil {
@@ -299,6 +299,8 @@ type AccountUsageService struct {
cache *UsageCache
identityCache IdentityCache
tlsFPProfileService *TLSFingerprintProfileService
agentIdentityTaskMu sync.Mutex
agentIdentityWS agentIdentityWSConnectionInvalidator
}
// NewAccountUsageService 创建AccountUsageService实例
@@ -692,8 +694,11 @@ func (s *AccountUsageService) probeOpenAICodexSnapshot(ctx context.Context, acco
if account == nil || !account.IsOAuth() {
return nil, nil
}
accessToken := account.GetOpenAIAccessToken()
if accessToken == "" {
accessToken := ""
if !account.IsOpenAIAgentIdentity() {
accessToken = account.GetOpenAIAccessToken()
}
if accessToken == "" && !account.IsOpenAIAgentIdentity() {
return nil, fmt.Errorf("no access token available")
}
modelID := openaipkg.DefaultTestModel
@@ -711,7 +716,19 @@ func (s *AccountUsageService) probeOpenAICodexSnapshot(ctx context.Context, acco
}
req.Host = "chatgpt.com"
req.Header.Set("Content-Type", "application/json")
req.Header.Set("Authorization", "Bearer "+accessToken)
if account.IsOpenAIAgentIdentity() {
authHeaders, authErr := buildAgentIdentityAuthenticationHeaders(ctx, s.accountRepo, s.agentIdentityWS, &s.agentIdentityTaskMu, account)
if authErr != nil {
return nil, fmt.Errorf("build Agent Identity authentication: %w", authErr)
}
for key, values := range authHeaders {
for _, value := range values {
req.Header.Add(key, value)
}
}
} else {
req.Header.Set("Authorization", "Bearer "+accessToken)
}
req.Header.Set("Accept", "text/event-stream")
req.Header.Set("OpenAI-Beta", "responses=experimental")
req.Header.Set("Originator", "codex_cli_rs")
@@ -0,0 +1,522 @@
package service
import (
"context"
"crypto"
"crypto/ed25519"
"crypto/sha512"
"crypto/x509"
"encoding/base64"
"encoding/json"
"errors"
"fmt"
"io"
"net/http"
"strings"
"sync"
"time"
"github.com/Wei-Shaw/sub2api/internal/pkg/httpclient"
"golang.org/x/crypto/curve25519"
"golang.org/x/crypto/nacl/box"
)
const (
OpenAIAuthModeAgentIdentity = "agentIdentity"
agentIdentityAuthAPIBaseURL = "https://auth.openai.com/api/accounts"
agentIdentityTaskRegistrationTimeout = 30 * time.Second
)
var openAIAgentIdentityAuthAPIBaseURL = agentIdentityAuthAPIBaseURL
var agentIdentityTaskLocks sync.Map // map[int64]*sync.Mutex
type agentIdentityWSConnectionInvalidator interface {
InvalidateAgentIdentityWSConnections(accountID int64)
}
type agentIdentityKey struct {
runtimeID string
privateKey ed25519.PrivateKey
taskID string
}
type agentIdentityTaskRegistrationResponse struct {
TaskID string `json:"task_id"`
TaskIDCamel string `json:"taskId"`
EncryptedTaskID string `json:"encrypted_task_id"`
EncryptedTaskIDCamel string `json:"encryptedTaskId"`
}
type agentIdentityTaskRecoveredError struct{}
func (e *agentIdentityTaskRecoveredError) Error() string {
return "agent identity task recovered"
}
func (a *Account) IsOpenAIAgentIdentity() bool {
if a == nil || !a.IsOpenAIOAuth() {
return false
}
return strings.EqualFold(strings.TrimSpace(a.GetCredential(openAIAuthModeCredentialKey)), OpenAIAuthModeAgentIdentity)
}
func agentIdentityPrivateKey(account *Account) (ed25519.PrivateKey, error) {
if account == nil {
return nil, errors.New("agent identity account is nil")
}
raw := strings.TrimSpace(account.GetCredential("agent_private_key"))
if raw == "" {
return nil, errors.New("agent identity private key is missing")
}
der, err := base64.StdEncoding.DecodeString(raw)
if err != nil {
return nil, errors.New("agent identity private key is not valid base64")
}
key, err := x509.ParsePKCS8PrivateKey(der)
if err != nil {
return nil, errors.New("agent identity private key is not valid PKCS#8")
}
privateKey, ok := key.(ed25519.PrivateKey)
if !ok || len(privateKey) != ed25519.PrivateKeySize {
return nil, errors.New("agent identity private key is not Ed25519")
}
return privateKey, nil
}
// ValidateOpenAIAgentIdentityPrivateKey validates the stored PKCS#8 Ed25519
// form without returning or logging the key material.
func ValidateOpenAIAgentIdentityPrivateKey(encoded string) error {
account := &Account{Credentials: map[string]any{"agent_private_key": encoded}}
_, err := agentIdentityPrivateKey(account)
return err
}
func agentIdentityKeyFromAccount(account *Account) (agentIdentityKey, error) {
privateKey, err := agentIdentityPrivateKey(account)
if err != nil {
return agentIdentityKey{}, err
}
runtimeID := strings.TrimSpace(account.GetCredential("agent_runtime_id"))
if runtimeID == "" {
return agentIdentityKey{}, errors.New("agent identity runtime id is missing")
}
return agentIdentityKey{
runtimeID: runtimeID,
privateKey: privateKey,
taskID: strings.TrimSpace(account.GetCredential("task_id")),
}, nil
}
func buildAgentAssertion(key agentIdentityKey, now time.Time) (string, error) {
if key.runtimeID == "" || key.taskID == "" {
return "", errors.New("agent identity runtime or task id is missing")
}
timestamp := now.UTC().Format(time.RFC3339)
payload := []byte(key.runtimeID + ":" + key.taskID + ":" + timestamp)
signature, err := key.privateKey.Sign(nil, payload, crypto.Hash(0))
if err != nil {
return "", errors.New("failed to sign agent assertion")
}
envelope := map[string]string{
"agent_runtime_id": key.runtimeID,
"task_id": key.taskID,
"timestamp": timestamp,
"signature": base64.StdEncoding.EncodeToString(signature),
}
encoded, err := json.Marshal(envelope)
if err != nil {
return "", errors.New("failed to serialize agent assertion")
}
return "AgentAssertion " + base64.RawURLEncoding.EncodeToString(encoded), nil
}
func signAgentTaskRegistration(key agentIdentityKey, timestamp time.Time) (string, string, error) {
if key.runtimeID == "" {
return "", "", errors.New("agent identity runtime id is missing")
}
formatted := timestamp.UTC().Format(time.RFC3339)
signature, err := key.privateKey.Sign(nil, []byte(key.runtimeID+":"+formatted), crypto.Hash(0))
if err != nil {
return "", "", errors.New("failed to sign agent task registration")
}
return formatted, base64.StdEncoding.EncodeToString(signature), nil
}
func decryptAgentTaskID(key agentIdentityKey, encoded string) (string, error) {
ciphertext, err := base64.StdEncoding.DecodeString(strings.TrimSpace(encoded))
if err != nil {
return "", errors.New("encrypted agent task id is not valid base64")
}
seed := key.privateKey.Seed()
digest := sha512.Sum512(seed)
var curvePrivate [32]byte
copy(curvePrivate[:], digest[:32])
curvePrivate[0] &= 248
curvePrivate[31] &= 127
curvePrivate[31] |= 64
curvePublicBytes, err := curve25519.X25519(curvePrivate[:], curve25519.Basepoint)
if err != nil {
return "", errors.New("failed to derive agent identity decryption key")
}
var curvePublic [32]byte
copy(curvePublic[:], curvePublicBytes)
plaintext, ok := box.OpenAnonymous(nil, ciphertext, &curvePublic, &curvePrivate)
if !ok {
return "", errors.New("failed to decrypt encrypted agent task id")
}
taskID := strings.TrimSpace(string(plaintext))
if taskID == "" {
return "", errors.New("decrypted agent task id is empty")
}
return taskID, nil
}
func registerAgentIdentityTask(ctx context.Context, account *Account) (string, error) {
key, err := agentIdentityKeyFromAccount(account)
if err != nil {
return "", err
}
timestamp, signature, err := signAgentTaskRegistration(key, time.Now())
if err != nil {
return "", err
}
proxyURL := ""
if account.ProxyID != nil && account.Proxy != nil {
proxyURL = account.Proxy.URL()
}
client, err := httpclient.GetClient(httpclient.Options{
ProxyURL: proxyURL,
Timeout: agentIdentityTaskRegistrationTimeout,
ResponseHeaderTimeout: 15 * time.Second,
})
if err != nil {
return "", errors.New("invalid proxy configuration for agent task registration")
}
body, err := json.Marshal(map[string]string{
"timestamp": timestamp,
"signature": signature,
})
if err != nil {
return "", errors.New("failed to serialize agent task registration")
}
url := strings.TrimRight(strings.TrimSpace(openAIAgentIdentityAuthAPIBaseURL), "/") + "/v1/agent/" + key.runtimeID + "/task/register"
req, err := http.NewRequestWithContext(ctx, http.MethodPost, url, strings.NewReader(string(body)))
if err != nil {
return "", errors.New("failed to build agent task registration request")
}
req.Header.Set("Content-Type", "application/json")
req.Header.Set("Accept", "application/json")
resp, err := client.Do(req)
if err != nil {
return "", errors.New("agent task registration request failed")
}
defer func() { _ = resp.Body.Close() }()
if resp.StatusCode < http.StatusOK || resp.StatusCode >= http.StatusMultipleChoices {
return "", fmt.Errorf("agent task registration returned status %d", resp.StatusCode)
}
var result agentIdentityTaskRegistrationResponse
if err := json.NewDecoder(io.LimitReader(resp.Body, 64*1024)).Decode(&result); err != nil {
return "", errors.New("agent task registration response is invalid")
}
if taskID := strings.TrimSpace(result.TaskID); taskID != "" {
return taskID, nil
}
if taskID := strings.TrimSpace(result.TaskIDCamel); taskID != "" {
return taskID, nil
}
encrypted := strings.TrimSpace(result.EncryptedTaskID)
if encrypted == "" {
encrypted = strings.TrimSpace(result.EncryptedTaskIDCamel)
}
if encrypted == "" {
return "", errors.New("agent task registration response omitted task id")
}
return decryptAgentTaskID(key, encrypted)
}
func ensureAgentIdentityTaskForAccount(ctx context.Context, repo AccountRepository, wsInvalidator agentIdentityWSConnectionInvalidator, taskMu *sync.Mutex, account *Account, expectedTaskID string) error {
if account == nil || !account.IsOpenAIAgentIdentity() {
return nil
}
credAccount := account
if account.IsShadow() {
resolved, err := resolveCredentialAccount(ctx, repo, account)
if err != nil {
return err
}
credAccount = resolved
}
if credAccount == nil || !credAccount.IsOpenAIAgentIdentity() {
return errors.New("agent identity credentials are unavailable")
}
currentTaskID := strings.TrimSpace(credAccount.GetCredential("task_id"))
if currentTaskID != "" && (expectedTaskID == "" || currentTaskID != expectedTaskID) {
return nil
}
if taskMu == nil {
return errors.New("agent identity task lock is unavailable")
}
sharedTaskMu := taskMu
if credAccount.ID > 0 {
candidate := &sync.Mutex{}
actual, _ := agentIdentityTaskLocks.LoadOrStore(credAccount.ID, candidate)
loadedTaskMu, ok := actual.(*sync.Mutex)
if !ok {
return errors.New("agent identity task lock has invalid type")
}
sharedTaskMu = loadedTaskMu
}
sharedTaskMu.Lock()
defer sharedTaskMu.Unlock()
// Re-read inside the shared lock. Different request paths often receive
// independent repository snapshots; checking only the caller's snapshot
// would allow sequential duplicate registrations after the first writer
// has already persisted a new task.
if repo != nil && credAccount.ID > 0 {
if refreshed, refreshErr := repo.GetByID(ctx, credAccount.ID); refreshErr == nil && refreshed != nil {
if refreshed.IsShadow() {
if resolved, resolveErr := resolveCredentialAccount(ctx, repo, refreshed); resolveErr == nil && resolved != nil {
refreshed = resolved
}
}
if refreshed.IsOpenAIAgentIdentity() {
credAccount = refreshed
if !account.IsShadow() {
account.Credentials = shallowCopyMap(credAccount.Credentials)
}
}
}
}
currentTaskID = strings.TrimSpace(credAccount.GetCredential("task_id"))
if currentTaskID != "" && (expectedTaskID == "" || currentTaskID != expectedTaskID) {
return nil
}
newTaskID, err := registerAgentIdentityTask(ctx, credAccount)
if err != nil {
return err
}
credentials := make(map[string]any, len(credAccount.Credentials)+1)
for key, value := range credAccount.Credentials {
credentials[key] = value
}
credentials["task_id"] = newTaskID
if err := persistAccountCredentials(ctx, repo, credAccount, credentials); err != nil {
return err
}
if !account.IsShadow() && account != credAccount {
account.Credentials = shallowCopyMap(credAccount.Credentials)
}
if wsInvalidator != nil {
wsInvalidator.InvalidateAgentIdentityWSConnections(credAccount.ID)
}
return nil
}
func (s *OpenAIGatewayService) ensureAgentIdentityTask(ctx context.Context, account *Account, expectedTaskID string) error {
if s == nil {
return errors.New("openai gateway service is nil")
}
return ensureAgentIdentityTaskForAccount(ctx, s.accountRepo, s, &s.agentIdentityTaskMu, account, expectedTaskID)
}
func isAgentIdentityTaskInvalidHTTPResponse(statusCode int, body []byte) bool {
if statusCode != http.StatusUnauthorized {
return false
}
lower := strings.ToLower(string(body))
compact := strings.NewReplacer(" ", "", "\t", "", "\r", "", "\n", "").Replace(lower)
for _, marker := range []string{
`"code":"invalid_task_id"`,
`"code":"task_not_found"`,
`"code":"task_expired"`,
`"error":"invalid_task_id"`,
} {
if strings.Contains(compact, marker) {
return true
}
}
for _, marker := range []string{
"invalid task_id",
"invalid task id",
"task_id is invalid",
"task id is invalid",
"task not found",
"task expired",
"unknown task_id",
"unknown task id",
} {
if strings.Contains(lower, marker) {
return true
}
}
return false
}
type agentIdentityTaskRecoveryContextKey struct{}
func markAgentIdentityTaskRecoveryTried(ctx context.Context) context.Context {
return context.WithValue(ctx, agentIdentityTaskRecoveryContextKey{}, true)
}
func agentIdentityTaskRecoveryWasTried(ctx context.Context) bool {
tried, _ := ctx.Value(agentIdentityTaskRecoveryContextKey{}).(bool)
return tried
}
func isAgentIdentityTaskInvalidWSDialError(err *openAIWSDialError) bool {
return err != nil && isAgentIdentityTaskInvalidHTTPResponse(err.StatusCode, err.ResponseBody)
}
func (s *OpenAIGatewayService) buildOpenAIAuthenticationHeaders(ctx context.Context, account *Account, token string) (http.Header, error) {
if account == nil {
return nil, errors.New("account is nil")
}
credAccount := account
if account.IsShadow() {
resolved, err := resolveCredentialAccount(ctx, s.accountRepo, account)
if err != nil {
return nil, err
}
credAccount = resolved
}
headers := make(http.Header)
if credAccount != nil && credAccount.IsOpenAIAgentIdentity() {
agentHeaders, err := buildAgentIdentityAuthenticationHeaders(ctx, s.accountRepo, s, &s.agentIdentityTaskMu, credAccount)
if err != nil {
return nil, err
}
return agentHeaders, nil
}
headers.Set("Authorization", "Bearer "+token)
return headers, nil
}
func buildAgentIdentityAuthenticationHeaders(ctx context.Context, repo AccountRepository, wsInvalidator agentIdentityWSConnectionInvalidator, taskMu *sync.Mutex, account *Account) (http.Header, error) {
if account == nil || !account.IsOpenAIAgentIdentity() {
return nil, errors.New("agent identity account is required")
}
if err := ensureAgentIdentityTaskForAccount(ctx, repo, wsInvalidator, taskMu, account, ""); err != nil {
return nil, err
}
key, err := agentIdentityKeyFromAccount(account)
if err != nil {
return nil, err
}
assertion, err := buildAgentAssertion(key, time.Now())
if err != nil {
return nil, err
}
headers := make(http.Header)
headers.Set("Authorization", assertion)
return headers, nil
}
func (s *OpenAIGatewayService) refreshOpenAIAgentIdentityHeaders(ctx context.Context, account *Account, headers http.Header) (http.Header, error) {
if account == nil {
return cloneHeader(headers), nil
}
credAccount := account
if account.IsShadow() {
resolved, err := resolveCredentialAccount(ctx, s.accountRepo, account)
if err != nil {
return nil, err
}
credAccount = resolved
}
if !credAccount.IsOpenAIAgentIdentity() {
return cloneHeader(headers), nil
}
refreshed := cloneHeader(headers)
if refreshed == nil {
refreshed = make(http.Header)
}
authHeaders, err := buildAgentIdentityAuthenticationHeaders(ctx, s.accountRepo, s, &s.agentIdentityTaskMu, credAccount)
if err != nil {
return nil, err
}
refreshed.Set("Authorization", authHeaders.Get("Authorization"))
return refreshed, nil
}
func (s *OpenAIGatewayService) recoverAgentIdentityTask(ctx context.Context, account *Account, expectedTaskID string) error {
if account != nil && account.IsShadow() {
if resolved, err := resolveCredentialAccount(ctx, s.accountRepo, account); err == nil && resolved != nil && strings.TrimSpace(expectedTaskID) == "" {
expectedTaskID = strings.TrimSpace(resolved.GetCredential("task_id"))
}
}
return s.ensureAgentIdentityTask(ctx, account, expectedTaskID)
}
func (s *OpenAIGatewayService) isAgentIdentityAccount(ctx context.Context, account *Account) bool {
if account == nil {
return false
}
credAccount := account
if account.IsShadow() {
resolved, err := resolveCredentialAccount(ctx, s.accountRepo, account)
if err != nil {
return false
}
credAccount = resolved
}
return credAccount != nil && credAccount.IsOpenAIAgentIdentity()
}
// redactAgentIdentitySensitiveBody removes credential values before an
// upstream error can reach logs, ops events, or returned error text. Agent
// Identity responses should not echo these values, but keeping this boundary
// defensive prevents accidental disclosure if an upstream error does.
func redactAgentIdentitySensitiveBodyForAccount(ctx context.Context, repo AccountRepository, account *Account, body []byte) []byte {
if account == nil || len(body) == 0 {
return body
}
credAccount := account
if account != nil && account.IsShadow() {
if resolved, err := resolveCredentialAccount(ctx, repo, account); err == nil && resolved != nil {
credAccount = resolved
}
}
if credAccount == nil || !credAccount.IsOpenAIAgentIdentity() {
return body
}
redacted := string(body)
for _, key := range []string{
"agent_private_key",
"agent_runtime_id",
"task_id",
"access_token",
"refresh_token",
"id_token",
"api_key",
"session_key",
"cookie",
} {
if value := strings.TrimSpace(credAccount.GetCredential(key)); value != "" {
redacted = strings.ReplaceAll(redacted, value, "[redacted]")
}
}
const assertionPrefix = "AgentAssertion "
for offset := 0; offset < len(redacted); {
relativeStart := strings.Index(redacted[offset:], assertionPrefix)
if relativeStart < 0 {
break
}
start := offset + relativeStart
valueStart := start + len(assertionPrefix)
end := valueStart
for end < len(redacted) && !strings.ContainsRune(" \t\r\n\"',}", rune(redacted[end])) {
end++
}
redacted = redacted[:valueStart] + "[redacted]" + redacted[end:]
offset = valueStart + len("[redacted]")
}
return []byte(redacted)
}
func (s *OpenAIGatewayService) redactAgentIdentitySensitiveBody(ctx context.Context, account *Account, body []byte) []byte {
if !s.isAgentIdentityAccount(ctx, account) {
return body
}
return redactAgentIdentitySensitiveBodyForAccount(ctx, s.accountRepo, account, body)
}
@@ -0,0 +1,500 @@
package service
import (
"bytes"
"context"
"encoding/base64"
"encoding/json"
"io"
"net/http"
"net/http/httptest"
"strings"
"testing"
"github.com/Wei-Shaw/sub2api/internal/config"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/require"
)
func TestAccountTestServiceOpenAICompactAgentIdentityUsesFreshAssertion(t *testing.T) {
gin.SetMode(gin.TestMode)
key, privateKey := newTestAgentIdentityKey(t)
account := Account{
ID: 21,
Name: "agent-identity",
Platform: PlatformOpenAI,
Type: AccountTypeOAuth,
Status: StatusActive,
Schedulable: true,
Concurrency: 1,
Credentials: map[string]any{
"auth_mode": OpenAIAuthModeAgentIdentity,
"agent_runtime_id": key.runtimeID,
"agent_private_key": privateKey,
"task_id": key.taskID,
"chatgpt_account_id": "account-agent-test",
"chatgpt_account_is_fedramp": true,
},
}
repo := &snapshotUpdateAccountRepo{stubOpenAIAccountRepo: stubOpenAIAccountRepo{accounts: []Account{account}}}
upstream := &httpUpstreamRecorder{resp: &http.Response{
StatusCode: http.StatusOK,
Header: http.Header{"Content-Type": []string{"application/json"}},
Body: io.NopCloser(strings.NewReader(`{"id":"compact-agent","status":"completed"}`)),
}}
svc := &AccountTestService{accountRepo: repo, httpUpstream: upstream}
rec := httptest.NewRecorder()
c, _ := gin.CreateTestContext(rec)
c.Request = httptest.NewRequest(http.MethodPost, "/api/v1/admin/accounts/21/test", bytes.NewReader(nil))
require.NoError(t, svc.TestAccountConnection(c, account.ID, "gpt-5.4", "", AccountTestModeCompact))
require.Equal(t, "AgentAssertion", strings.SplitN(upstream.lastReq.Header.Get("Authorization"), " ", 2)[0])
require.Equal(t, "account-agent-test", upstream.lastReq.Header.Get("chatgpt-account-id"))
require.Equal(t, "true", upstream.lastReq.Header.Get("x-openai-fedramp"))
require.NotContains(t, upstream.lastReq.Header.Get("Authorization"), privateKey)
}
func TestAccountTestServiceOpenAICompactAgentIdentityRecoversInvalidTaskOnce(t *testing.T) {
gin.SetMode(gin.TestMode)
key, privateKey := newTestAgentIdentityKey(t)
account := &Account{
ID: 22,
Name: "agent-identity-recovery",
Platform: PlatformOpenAI,
Type: AccountTypeOAuth,
Status: StatusActive,
Schedulable: true,
Concurrency: 1,
Credentials: map[string]any{
"auth_mode": OpenAIAuthModeAgentIdentity,
"agent_runtime_id": key.runtimeID,
"agent_private_key": privateKey,
"task_id": "task-compact-old",
"chatgpt_account_id": "account-agent-compact-recovery",
},
}
repo := &accountTestAgentIdentityRepo{account: account}
registerCalls := 0
registerServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
registerCalls++
_, _ = io.WriteString(w, `{"task_id":"task-compact-new"}`)
}))
defer registerServer.Close()
oldBase := openAIAgentIdentityAuthAPIBaseURL
openAIAgentIdentityAuthAPIBaseURL = registerServer.URL
t.Cleanup(func() { openAIAgentIdentityAuthAPIBaseURL = oldBase })
upstream := &httpUpstreamRecorder{responses: []*http.Response{
{StatusCode: http.StatusUnauthorized, Header: http.Header{"Content-Type": []string{"application/json"}}, Body: io.NopCloser(strings.NewReader(`{"error":{"code":"invalid_task_id"}}`))},
{StatusCode: http.StatusOK, Header: http.Header{"Content-Type": []string{"application/json"}}, Body: io.NopCloser(strings.NewReader(`{"id":"compact-agent","status":"completed"}`))},
}}
invalidator := &agentIdentityWSInvalidationRecorder{}
svc := &AccountTestService{accountRepo: repo, httpUpstream: upstream, agentIdentityWS: invalidator}
rec := httptest.NewRecorder()
c, _ := gin.CreateTestContext(rec)
c.Request = httptest.NewRequest(http.MethodPost, "/api/v1/admin/accounts/22/test", bytes.NewReader(nil))
require.NoError(t, svc.TestAccountConnection(c, account.ID, "gpt-5.4", "", AccountTestModeCompact))
require.Equal(t, 1, registerCalls)
require.Len(t, upstream.requests, 2)
require.Equal(t, "task-compact-new", account.GetCredential("task_id"))
require.Equal(t, 0, repo.setErrorCalls)
require.Equal(t, []int64{account.ID}, invalidator.accountIDs)
}
func TestOpenAIAgentIdentityPassthroughKeepsSessionAndPromptCacheHeaders(t *testing.T) {
gin.SetMode(gin.TestMode)
key, privateKey := newTestAgentIdentityKey(t)
account := &Account{
ID: 24,
Platform: PlatformOpenAI,
Type: AccountTypeOAuth,
Credentials: map[string]any{
"auth_mode": OpenAIAuthModeAgentIdentity,
"agent_runtime_id": key.runtimeID,
"agent_private_key": privateKey,
"task_id": key.taskID,
"chatgpt_account_id": "account-agent-passthrough",
},
}
rec := httptest.NewRecorder()
c, _ := gin.CreateTestContext(rec)
body := []byte(`{"model":"gpt-5.4","instructions":"Reply OK","input":[],"stream":true,"prompt_cache_key":"cache-agent"}`)
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", bytes.NewReader(body))
c.Request.Header.Set("session_id", "client-session")
c.Request.Header.Set("conversation_id", "client-conversation")
c.Request.Header.Set("Authorization", "Bearer inbound-must-not-forward")
svc := &OpenAIGatewayService{}
req, err := svc.buildUpstreamRequestOpenAIPassthrough(context.Background(), c, account, body, "")
require.NoError(t, err)
require.Equal(t, "AgentAssertion", strings.SplitN(req.Header.Get("Authorization"), " ", 2)[0])
require.Equal(t, "account-agent-passthrough", req.Header.Get("chatgpt-account-id"))
require.NotEqual(t, "client-session", req.Header.Get("session_id"))
require.NotEqual(t, "client-conversation", req.Header.Get("conversation_id"))
require.Equal(t, isolateOpenAISessionID(0, "client-session"), req.Header.Get("session_id"))
require.Equal(t, isolateOpenAISessionID(0, "client-conversation"), req.Header.Get("conversation_id"))
requestBody, err := io.ReadAll(req.Body)
require.NoError(t, err)
require.Contains(t, string(requestBody), `"prompt_cache_key":"cache-agent"`)
// Authentication mode must not affect session isolation or prompt-cache
// behavior. Compare the same request with the existing OAuth path instead
// of pinning this test to an implementation-specific hash.
oauthAccount := &Account{
ID: 26,
Platform: PlatformOpenAI,
Type: AccountTypeOAuth,
Credentials: map[string]any{
"chatgpt_account_id": "account-oauth-passthrough",
},
}
oauthRecorder := httptest.NewRecorder()
oauthContext, _ := gin.CreateTestContext(oauthRecorder)
oauthContext.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", bytes.NewReader(body))
oauthContext.Request.Header.Set("session_id", "client-session")
oauthContext.Request.Header.Set("conversation_id", "client-conversation")
oauthReq, err := svc.buildUpstreamRequestOpenAIPassthrough(context.Background(), oauthContext, oauthAccount, body, "oauth-token")
require.NoError(t, err)
require.Equal(t, oauthReq.Header.Get("session_id"), req.Header.Get("session_id"))
require.Equal(t, oauthReq.Header.Get("conversation_id"), req.Header.Get("conversation_id"))
}
func TestOpenAIAgentIdentityErrorRedactionDoesNotLeakCredentialValues(t *testing.T) {
key, privateKey := newTestAgentIdentityKey(t)
account := &Account{
ID: 25,
Platform: PlatformOpenAI,
Type: AccountTypeOAuth,
Credentials: map[string]any{
"auth_mode": OpenAIAuthModeAgentIdentity,
"agent_runtime_id": key.runtimeID,
"agent_private_key": privateKey,
"task_id": key.taskID,
"access_token": key.runtimeID + "-oauth-value",
},
}
svc := &OpenAIGatewayService{}
oauthValue := account.GetCredential("access_token")
redacted := svc.redactAgentIdentitySensitiveBody(context.Background(), account, []byte(`{"message":"runtime-test task-test `+oauthValue+` AgentAssertion abc123"}`))
require.NotContains(t, string(redacted), key.runtimeID)
require.NotContains(t, string(redacted), key.taskID)
require.NotContains(t, string(redacted), oauthValue)
require.NotContains(t, string(redacted), "AgentAssertion abc123")
require.Contains(t, string(redacted), "[redacted]")
}
func TestOpenAIAuthenticationHeadersPreserveOAuthPATAndAPIKeyBearerModes(t *testing.T) {
svc := &OpenAIGatewayService{}
tests := []struct {
name string
account *Account
token string
}{
{name: "oauth", account: &Account{Platform: PlatformOpenAI, Type: AccountTypeOAuth}, token: "oauth-runtime-token"},
{name: "personal access token", account: &Account{Platform: PlatformOpenAI, Type: AccountTypeOAuth, Credentials: map[string]any{"auth_mode": OpenAIAuthModePersonalAccessToken}}, token: "pat-runtime-token"},
{name: "api key", account: &Account{Platform: PlatformOpenAI, Type: AccountTypeAPIKey}, token: "api-key-runtime-token"},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
headers, err := svc.buildOpenAIAuthenticationHeaders(context.Background(), tt.account, tt.token)
require.NoError(t, err)
require.Equal(t, "Bearer "+tt.token, headers.Get("Authorization"))
})
}
}
func TestOpenAIWSAgentIdentityRecoveryRequiresTaskInvalidBody(t *testing.T) {
require.False(t, isAgentIdentityTaskInvalidWSDialError(&openAIWSDialError{
StatusCode: http.StatusUnauthorized,
ResponseBody: []byte(`{"error":{"code":"invalid_signature"}}`),
}))
require.True(t, isAgentIdentityTaskInvalidWSDialError(&openAIWSDialError{
StatusCode: http.StatusUnauthorized,
ResponseBody: []byte(`{"error":{"code":"invalid_task_id"}}`),
}))
}
func TestValidateOpenAIWSBearerTokenAllowsAgentIdentityWithoutStoredToken(t *testing.T) {
t.Run("Given Agent Identity When a WS path receives no bearer token Then dial-time assertion auth is allowed", func(t *testing.T) {
account := &Account{
Platform: PlatformOpenAI,
Type: AccountTypeOAuth,
Credentials: map[string]any{
"auth_mode": OpenAIAuthModeAgentIdentity,
},
}
require.NoError(t, validateOpenAIWSBearerToken(account, ""))
})
t.Run("Given bearer credentials When a WS path receives no token Then the request is rejected", func(t *testing.T) {
accounts := []*Account{
{Platform: PlatformOpenAI, Type: AccountTypeOAuth},
{Platform: PlatformOpenAI, Type: AccountTypeOAuth, Credentials: map[string]any{"auth_mode": OpenAIAuthModePersonalAccessToken}},
{Platform: PlatformOpenAI, Type: AccountTypeAPIKey},
}
for _, account := range accounts {
require.EqualError(t, validateOpenAIWSBearerToken(account, ""), "token is empty")
}
})
}
func TestOpenAIWSConnPoolHeadersFactoryRunsAtDialAndStalePrewarmIsDiscarded(t *testing.T) {
cfg := &config.Config{}
cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 1
pool := newOpenAIWSConnPool(cfg)
defer pool.Close()
pool.setClientDialerForTest(&openAIWSFakeDialer{})
accountID := int64(22)
ap := pool.getOrCreateAccountPool(accountID)
factoryCalls := 0
latestHeader := ""
req := openAIWSAcquireRequest{
Account: &Account{ID: accountID, Platform: PlatformOpenAI, Type: AccountTypeOAuth},
WSURL: "wss://example.com/v1/responses",
HeadersFactory: func(_ context.Context, headers http.Header) (http.Header, error) {
factoryCalls++
latestHeader = "AgentAssertion dial-" + string(rune('0'+factoryCalls))
if headers == nil {
headers = make(http.Header)
}
headers.Set("Authorization", latestHeader)
return headers, nil
},
}
ap.mu.Lock()
ap.lastAcquire = &req
generation := ap.generation
ap.mu.Unlock()
pool.prewarmConns(accountID, req, 1, generation)
require.Equal(t, 1, factoryCalls, "prewarm must generate authorization inside the actual dial")
require.Equal(t, "AgentAssertion dial-1", latestHeader)
pool.ClearAccount(accountID)
ap.mu.Lock()
require.Empty(t, ap.conns, "credential recovery must remove pooled connections")
require.Nil(t, ap.lastAcquire, "credential recovery must discard delayed acquire state")
require.Equal(t, generation+1, ap.generation)
ap.mu.Unlock()
// A prewarm captured before ClearAccount must not be admitted after recovery.
pool.prewarmConns(accountID, req, 1, generation)
ap.mu.Lock()
require.Empty(t, ap.conns)
ap.mu.Unlock()
}
func TestOpenAIAgentIdentityTaskInvalidRetriesExactlyOnce(t *testing.T) {
gin.SetMode(gin.TestMode)
key, privateKey := newTestAgentIdentityKey(t)
account := &Account{
ID: 23,
Name: "agent-identity",
Platform: PlatformOpenAI,
Type: AccountTypeOAuth,
Status: StatusActive,
Schedulable: true,
Concurrency: 1,
Credentials: map[string]any{
"auth_mode": OpenAIAuthModeAgentIdentity,
"agent_runtime_id": key.runtimeID,
"agent_private_key": privateKey,
"task_id": "task-old",
"chatgpt_account_id": "account-agent-retry",
},
}
repo := &agentIdentityForwardRepo{account: account}
registerCalls := 0
registerServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
registerCalls++
_, _ = io.WriteString(w, `{"task_id":"task-new"}`)
}))
defer registerServer.Close()
oldBase := openAIAgentIdentityAuthAPIBaseURL
openAIAgentIdentityAuthAPIBaseURL = registerServer.URL
t.Cleanup(func() { openAIAgentIdentityAuthAPIBaseURL = oldBase })
successBody := `{"id":"resp-agent-retry","object":"response","model":"gpt-5.4","output":[],"usage":{"input_tokens":1,"output_tokens":1,"total_tokens":2}}`
upstream := &httpUpstreamRecorder{responses: []*http.Response{
{StatusCode: http.StatusUnauthorized, Header: http.Header{"Content-Type": []string{"application/json"}}, Body: io.NopCloser(strings.NewReader(`{"error":{"code":"invalid_task_id"}}`))},
{StatusCode: http.StatusOK, Header: http.Header{"Content-Type": []string{"application/json"}}, Body: io.NopCloser(strings.NewReader(successBody))},
}}
require.True(t, isAgentIdentityTaskInvalidHTTPResponse(http.StatusUnauthorized, []byte(`{"error":{"code":"invalid_task_id"}}`)))
svc := &OpenAIGatewayService{cfg: &config.Config{}, accountRepo: repo, httpUpstream: upstream}
rec := httptest.NewRecorder()
c, _ := gin.CreateTestContext(rec)
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", strings.NewReader(`{"model":"gpt-5.4","instructions":"Reply OK","input":[],"stream":false}`))
_, err := svc.Forward(context.Background(), c, account, []byte(`{"model":"gpt-5.4","instructions":"Reply OK","input":[],"stream":false}`))
require.NoError(t, err)
require.Equal(t, 1, registerCalls)
require.Len(t, upstream.requests, 2)
require.NotEqual(t, upstream.requests[0].Header.Get("Authorization"), upstream.requests[1].Header.Get("Authorization"))
require.Equal(t, "task-new", decodeAgentAssertionTask(t, upstream.requests[1].Header.Get("Authorization")))
// Two consecutive invalid responses still produce only one retry for this
// request; the recovery path must not loop indefinitely.
upstream.responses = []*http.Response{
{StatusCode: http.StatusUnauthorized, Header: http.Header{"Content-Type": []string{"application/json"}}, Body: io.NopCloser(strings.NewReader(`{"error":{"code":"invalid_task_id"}}`))},
{StatusCode: http.StatusUnauthorized, Header: http.Header{"Content-Type": []string{"application/json"}}, Body: io.NopCloser(strings.NewReader(`{"error":{"code":"invalid_task_id"}}`))},
}
rec2 := httptest.NewRecorder()
c2, _ := gin.CreateTestContext(rec2)
c2.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", strings.NewReader(`{"model":"gpt-5.4","instructions":"Reply OK","input":[],"stream":false}`))
_, err = svc.Forward(context.Background(), c2, account, []byte(`{"model":"gpt-5.4","instructions":"Reply OK","input":[],"stream":false}`))
require.Error(t, err)
require.Equal(t, 2, registerCalls)
require.Len(t, upstream.requests, 4)
// Passthrough uses the same one-shot task recovery contract.
account.Extra = map[string]any{"openai_passthrough": true}
account.Credentials["task_id"] = "task-old-passthrough"
upstream.responses = []*http.Response{
{StatusCode: http.StatusUnauthorized, Header: http.Header{"Content-Type": []string{"application/json"}}, Body: io.NopCloser(strings.NewReader(`{"error":{"code":"invalid_task_id"}}`))},
{StatusCode: http.StatusOK, Header: http.Header{"Content-Type": []string{"text/event-stream"}}, Body: io.NopCloser(strings.NewReader("data: {\"type\":\"response.completed\",\"response\":{\"usage\":{\"input_tokens\":1,\"output_tokens\":1,\"total_tokens\":2}}}\n\ndata: [DONE]\n\n"))},
}
rec3 := httptest.NewRecorder()
c3, _ := gin.CreateTestContext(rec3)
c3.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", strings.NewReader(`{"model":"gpt-5.4","instructions":"Reply OK","input":[],"stream":false}`))
_, err = svc.Forward(context.Background(), c3, account, []byte(`{"model":"gpt-5.4","instructions":"Reply OK","input":[],"stream":false}`))
require.NoError(t, err)
require.Equal(t, 3, registerCalls)
require.Len(t, upstream.requests, 6)
}
func TestOpenAIAgentIdentityCompatRoutesRecoverInvalidTaskOnce(t *testing.T) {
gin.SetMode(gin.TestMode)
tests := []struct {
name string
path string
body []byte
call func(*OpenAIGatewayService, context.Context, *gin.Context, *Account, []byte) (*OpenAIForwardResult, error)
}{
{
name: "chat completions",
path: "/v1/chat/completions",
body: []byte(`{"model":"gpt-5.4","stream":false,"messages":[{"role":"user","content":"hi"}]}`),
call: func(s *OpenAIGatewayService, ctx context.Context, c *gin.Context, account *Account, body []byte) (*OpenAIForwardResult, error) {
return s.ForwardAsChatCompletions(ctx, c, account, body, "", "gpt-5.4")
},
},
{
name: "anthropic messages",
path: "/v1/messages",
body: []byte(`{"model":"gpt-5.4","stream":false,"max_tokens":32,"messages":[{"role":"user","content":"hi"}]}`),
call: func(s *OpenAIGatewayService, ctx context.Context, c *gin.Context, account *Account, body []byte) (*OpenAIForwardResult, error) {
return s.ForwardAsAnthropic(ctx, c, account, body, "", "gpt-5.4")
},
},
}
for index, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
key, privateKey := newTestAgentIdentityKey(t)
account := &Account{
ID: int64(40 + index),
Name: "agent-identity-compat",
Platform: PlatformOpenAI,
Type: AccountTypeOAuth,
Status: StatusActive,
Schedulable: true,
Concurrency: 1,
Credentials: map[string]any{
"auth_mode": OpenAIAuthModeAgentIdentity,
"agent_runtime_id": key.runtimeID,
"agent_private_key": privateKey,
"task_id": "task-compat-old",
"chatgpt_account_id": "account-compat-recovery",
},
}
repo := &agentIdentityForwardRepo{account: account}
registerCalls := 0
registerServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
registerCalls++
_, _ = io.WriteString(w, `{"task_id":"task-compat-new"}`)
}))
defer registerServer.Close()
oldBase := openAIAgentIdentityAuthAPIBaseURL
openAIAgentIdentityAuthAPIBaseURL = registerServer.URL
t.Cleanup(func() { openAIAgentIdentityAuthAPIBaseURL = oldBase })
upstream := &httpUpstreamRecorder{responses: []*http.Response{
{StatusCode: http.StatusUnauthorized, Header: http.Header{"Content-Type": []string{"application/json"}}, Body: io.NopCloser(strings.NewReader(`{"error":{"code":"invalid_task_id"}}`))},
{StatusCode: http.StatusUnauthorized, Header: http.Header{"Content-Type": []string{"application/json"}}, Body: io.NopCloser(strings.NewReader(`{"error":{"code":"invalid_task_id"}}`))},
}}
svc := &OpenAIGatewayService{cfg: &config.Config{}, accountRepo: repo, httpUpstream: upstream}
rec := httptest.NewRecorder()
c, _ := gin.CreateTestContext(rec)
c.Request = httptest.NewRequest(http.MethodPost, tt.path, bytes.NewReader(tt.body))
_, err := tt.call(svc, context.Background(), c, account, tt.body)
require.Error(t, err)
require.Equal(t, 1, registerCalls)
require.Len(t, upstream.requests, 2)
require.Equal(t, "task-compat-new", account.GetCredential("task_id"))
})
}
}
func decodeAgentAssertionTask(t *testing.T, header string) string {
t.Helper()
encoded := strings.TrimPrefix(header, "AgentAssertion ")
decoded, err := base64.RawURLEncoding.DecodeString(encoded)
require.NoError(t, err)
var envelope struct {
TaskID string `json:"task_id"`
}
require.NoError(t, json.Unmarshal(decoded, &envelope))
return envelope.TaskID
}
type agentIdentityForwardRepo struct {
AccountRepository
account *Account
}
type agentIdentityWSInvalidationRecorder struct {
accountIDs []int64
}
func (r *agentIdentityWSInvalidationRecorder) InvalidateAgentIdentityWSConnections(accountID int64) {
r.accountIDs = append(r.accountIDs, accountID)
}
type accountTestAgentIdentityRepo struct {
AccountRepository
account *Account
setErrorCalls int
}
func (r *accountTestAgentIdentityRepo) GetByID(_ context.Context, _ int64) (*Account, error) {
return r.account, nil
}
func (r *accountTestAgentIdentityRepo) UpdateCredentials(_ context.Context, _ int64, credentials map[string]any) error {
r.account.Credentials = credentials
return nil
}
func (r *accountTestAgentIdentityRepo) UpdateExtra(_ context.Context, _ int64, _ map[string]any) error {
return nil
}
func (r *accountTestAgentIdentityRepo) SetError(_ context.Context, _ int64, _ string) error {
r.setErrorCalls++
return nil
}
func (r *agentIdentityForwardRepo) GetByID(_ context.Context, _ int64) (*Account, error) {
return r.account, nil
}
func (r *agentIdentityForwardRepo) UpdateCredentials(_ context.Context, _ int64, credentials map[string]any) error {
r.account.Credentials = credentials
return nil
}
@@ -0,0 +1,229 @@
package service
import (
"context"
"crypto/ed25519"
"crypto/rand"
"crypto/sha512"
"crypto/x509"
"encoding/base64"
"encoding/json"
"fmt"
"net/http"
"net/http/httptest"
"strings"
"sync"
"testing"
"time"
"github.com/stretchr/testify/require"
"golang.org/x/crypto/curve25519"
"golang.org/x/crypto/nacl/box"
)
func newTestAgentIdentityKey(t *testing.T) (agentIdentityKey, string) {
t.Helper()
_, privateKey, err := ed25519.GenerateKey(rand.Reader)
require.NoError(t, err)
der, err := x509.MarshalPKCS8PrivateKey(privateKey)
require.NoError(t, err)
return agentIdentityKey{
runtimeID: "runtime-test",
privateKey: privateKey,
taskID: "task-test",
}, base64.StdEncoding.EncodeToString(der)
}
func TestBuildAgentAssertionMatchesCodexEnvelopeAndSignature(t *testing.T) {
key, _ := newTestAgentIdentityKey(t)
now := time.Date(2026, 7, 14, 8, 9, 10, 0, time.FixedZone("UTC+8", 8*60*60))
assertion, err := buildAgentAssertion(key, now)
require.NoError(t, err)
require.True(t, strings.HasPrefix(assertion, "AgentAssertion "))
encoded := strings.TrimPrefix(assertion, "AgentAssertion ")
decoded, err := base64.RawURLEncoding.DecodeString(encoded)
require.NoError(t, err)
var envelope struct {
AgentRuntimeID string `json:"agent_runtime_id"`
TaskID string `json:"task_id"`
Timestamp string `json:"timestamp"`
Signature string `json:"signature"`
}
require.NoError(t, json.Unmarshal(decoded, &envelope))
require.Equal(t, "runtime-test", envelope.AgentRuntimeID)
require.Equal(t, "task-test", envelope.TaskID)
require.Equal(t, "2026-07-14T00:09:10Z", envelope.Timestamp)
signature, err := base64.StdEncoding.DecodeString(envelope.Signature)
require.NoError(t, err)
publicKey, ok := key.privateKey.Public().(ed25519.PublicKey)
require.True(t, ok)
require.True(t, ed25519.Verify(publicKey, []byte("runtime-test:task-test:2026-07-14T00:09:10Z"), signature))
}
func TestDecryptAgentTaskIDSupportsCodexSealedBoxResponse(t *testing.T) {
key, _ := newTestAgentIdentityKey(t)
digest := sha512.Sum512(key.privateKey.Seed())
var curvePrivate [32]byte
copy(curvePrivate[:], digest[:32])
curvePrivate[0] &= 248
curvePrivate[31] &= 127
curvePrivate[31] |= 64
curvePublicBytes, err := curve25519.X25519(curvePrivate[:], curve25519.Basepoint)
require.NoError(t, err)
var curvePublic [32]byte
copy(curvePublic[:], curvePublicBytes)
ciphertext, err := box.SealAnonymous(nil, []byte("task-sealed"), &curvePublic, rand.Reader)
require.NoError(t, err)
got, err := decryptAgentTaskID(key, base64.StdEncoding.EncodeToString(ciphertext))
require.NoError(t, err)
require.Equal(t, "task-sealed", got)
}
func TestRegisterAgentIdentityTaskAcceptsPlaintextAndEncryptedResponses(t *testing.T) {
key, privateKey := newTestAgentIdentityKey(t)
requestCount := 0
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
require.Equal(t, http.MethodPost, r.Method)
require.Equal(t, "/v1/agent/runtime-test/task/register", r.URL.Path)
var request map[string]string
require.NoError(t, json.NewDecoder(r.Body).Decode(&request))
require.NotEmpty(t, request["timestamp"])
require.NotEmpty(t, request["signature"])
requestCount++
if requestCount == 2 {
digest := sha512.Sum512(key.privateKey.Seed())
var curvePrivate [32]byte
copy(curvePrivate[:], digest[:32])
curvePrivate[0] &= 248
curvePrivate[31] &= 127
curvePrivate[31] |= 64
curvePublicBytes, curveErr := curve25519.X25519(curvePrivate[:], curve25519.Basepoint)
require.NoError(t, curveErr)
var curvePublic [32]byte
copy(curvePublic[:], curvePublicBytes)
ciphertext, sealErr := box.SealAnonymous(nil, []byte("task-encrypted"), &curvePublic, rand.Reader)
require.NoError(t, sealErr)
_, _ = fmt.Fprintf(w, `{"encrypted_task_id":%q}`, base64.StdEncoding.EncodeToString(ciphertext))
return
}
_, _ = w.Write([]byte(`{"task_id":"task-plain"}`))
}))
defer server.Close()
oldBase := openAIAgentIdentityAuthAPIBaseURL
openAIAgentIdentityAuthAPIBaseURL = server.URL
t.Cleanup(func() { openAIAgentIdentityAuthAPIBaseURL = oldBase })
account := &Account{ID: 1, Type: AccountTypeOAuth, Platform: PlatformOpenAI, Credentials: map[string]any{
"auth_mode": OpenAIAuthModeAgentIdentity,
"agent_runtime_id": key.runtimeID,
"agent_private_key": privateKey,
}}
taskID, err := registerAgentIdentityTask(context.Background(), account)
require.NoError(t, err)
require.Equal(t, "task-plain", taskID)
taskID, err = registerAgentIdentityTask(context.Background(), account)
require.NoError(t, err)
require.Equal(t, "task-encrypted", taskID)
}
func TestEnsureAgentIdentityTaskPersistsAndRedactsCredentials(t *testing.T) {
key, privateKey := newTestAgentIdentityKey(t)
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
_, _ = w.Write([]byte(`{"task_id":"task-persisted"}`))
}))
defer server.Close()
oldBase := openAIAgentIdentityAuthAPIBaseURL
openAIAgentIdentityAuthAPIBaseURL = server.URL
t.Cleanup(func() { openAIAgentIdentityAuthAPIBaseURL = oldBase })
repo := &agentIdentityCredentialsRepo{}
account := &Account{ID: 7, Type: AccountTypeOAuth, Platform: PlatformOpenAI, Credentials: map[string]any{
"auth_mode": OpenAIAuthModeAgentIdentity,
"agent_runtime_id": key.runtimeID,
"agent_private_key": privateKey,
"chatgpt_account_id": "account-test",
}}
service := &OpenAIGatewayService{accountRepo: repo}
require.NoError(t, service.ensureAgentIdentityTask(context.Background(), account, ""))
require.Equal(t, "task-persisted", account.GetCredential("task_id"))
require.Equal(t, "task-persisted", repo.credentials["task_id"])
require.True(t, IsSensitiveCredentialKey("agent_private_key"))
redacted := make(map[string]any)
for key, value := range account.Credentials {
if !IsSensitiveCredentialKey(key) {
redacted[key] = value
}
}
require.NotContains(t, string(mustAgentIdentityJSON(t, redacted)), privateKey)
}
func TestEnsureAgentIdentityTaskSharesLockAcrossServicesForSameAccount(t *testing.T) {
key, privateKey := newTestAgentIdentityKey(t)
account := &Account{ID: 9001, Type: AccountTypeOAuth, Platform: PlatformOpenAI, Credentials: map[string]any{
"auth_mode": OpenAIAuthModeAgentIdentity,
"agent_runtime_id": key.runtimeID,
"agent_private_key": privateKey,
}}
repo := &agentIdentityCredentialsRepo{account: account}
registerCalls := 0
var registerMu sync.Mutex
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
registerMu.Lock()
registerCalls++
registerMu.Unlock()
_, _ = w.Write([]byte(`{"task_id":"task-shared"}`))
}))
defer server.Close()
oldBase := openAIAgentIdentityAuthAPIBaseURL
openAIAgentIdentityAuthAPIBaseURL = server.URL
t.Cleanup(func() { openAIAgentIdentityAuthAPIBaseURL = oldBase })
start := make(chan struct{})
errors := make(chan error, 2)
requests := []*Account{cloneAgentIdentityTestAccount(account), cloneAgentIdentityTestAccount(account)}
for _, request := range requests {
go func() {
<-start
errors <- ensureAgentIdentityTaskForAccount(context.Background(), repo, nil, &sync.Mutex{}, request, "")
}()
}
close(start)
require.NoError(t, <-errors)
require.NoError(t, <-errors)
registerMu.Lock()
defer registerMu.Unlock()
require.Equal(t, 1, registerCalls)
require.Equal(t, "task-shared", repo.account.GetCredential("task_id"))
}
func cloneAgentIdentityTestAccount(account *Account) *Account {
copy := *account
copy.Credentials = shallowCopyMap(account.Credentials)
return &copy
}
type agentIdentityCredentialsRepo struct {
AccountRepository
credentials map[string]any
account *Account
mu sync.Mutex
}
func (r *agentIdentityCredentialsRepo) GetByID(_ context.Context, _ int64) (*Account, error) {
return r.account, nil
}
func (r *agentIdentityCredentialsRepo) UpdateCredentials(_ context.Context, _ int64, credentials map[string]any) error {
r.mu.Lock()
defer r.mu.Unlock()
r.credentials = credentials
return nil
}
func mustAgentIdentityJSON(t *testing.T, value any) []byte {
t.Helper()
encoded, err := json.Marshal(value)
require.NoError(t, err)
return encoded
}
@@ -241,7 +241,7 @@ func (s *OpenAIGatewayService) FetchCodexModelsManifest(ctx context.Context, acc
switch {
case credAccount.IsOpenAIOAuth():
authToken = strings.TrimSpace(credAccount.GetOpenAIAccessToken())
if authToken == "" {
if authToken == "" && !credAccount.IsOpenAIAgentIdentity() {
return nil, infraerrors.New(http.StatusBadGateway, "OPENAI_CODEX_MODELS_TOKEN_MISSING", "account has no Codex backend access token")
}
case credAccount.IsOpenAIApiKey():
@@ -277,16 +277,25 @@ func (s *OpenAIGatewayService) FetchCodexModelsManifest(ctx context.Context, acc
}
headers := make(http.Header)
headers.Set("Authorization", "Bearer "+authToken)
if useAPIKeyUpstream {
headers.Set("Authorization", "Bearer "+authToken)
credAccount.ApplyHeaderOverrides(headers)
} else {
authHeaders, authErr := s.buildOpenAIAuthenticationHeaders(ctx, credAccount, authToken)
if authErr != nil {
return nil, infraerrors.Newf(http.StatusBadGateway, "OPENAI_CODEX_MODELS_AUTH_FAILED", "build Codex models authentication: %v", authErr)
}
for key, values := range authHeaders {
for _, value := range values {
headers.Add(key, value)
}
}
setOpenAIChatGPTAccountHeaders(headers, credAccount)
}
headers.Set("Accept", "application/json")
headers.Set("Originator", "codex_cli_rs")
headers.Set("Version", clientVersion)
headers.Set("User-Agent", codexCLIUserAgent)
if useAPIKeyUpstream {
credAccount.ApplyHeaderOverrides(headers)
} else {
setOpenAIChatGPTAccountHeaders(headers, credAccount)
}
proxyURL := ""
if account.ProxyID != nil && account.Proxy != nil {
@@ -199,6 +199,49 @@ func TestFetchCodexModelsManifestPassthrough(t *testing.T) {
}
}
func TestFetchCodexModelsManifestAgentIdentityUsesAssertionWithoutOAuthToken(t *testing.T) {
key, privateKey := newTestAgentIdentityKey(t)
account := &Account{
ID: 3,
Platform: PlatformOpenAI,
Type: AccountTypeOAuth,
Credentials: map[string]any{
"auth_mode": OpenAIAuthModeAgentIdentity,
"agent_runtime_id": key.runtimeID,
"agent_private_key": privateKey,
"task_id": key.taskID,
"chatgpt_account_id": "acc-agent",
},
}
var gotAuth, gotAccountID string
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
gotAuth = r.Header.Get("Authorization")
gotAccountID = r.Header.Get("chatgpt-account-id")
_, _ = w.Write([]byte(`{"models":[]}`))
}))
defer server.Close()
original := chatgptCodexModelsURL
chatgptCodexModelsURL = server.URL
defer func() { chatgptCodexModelsURL = original }()
s := &OpenAIGatewayService{}
manifest, err := s.FetchCodexModelsManifest(context.Background(), account, "0.137.0", "")
if err != nil {
t.Fatalf("FetchCodexModelsManifest returned error: %v", err)
}
if string(manifest.Body) != `{"models":[]}` {
t.Fatalf("unexpected manifest body: %q", manifest.Body)
}
if !strings.HasPrefix(gotAuth, "AgentAssertion ") {
t.Fatalf("authorization scheme: got %q", strings.SplitN(gotAuth, " ", 2)[0])
}
if gotAccountID != "acc-agent" {
t.Fatalf("chatgpt-account-id header: got %q", gotAccountID)
}
}
func TestFetchCodexModelsManifestDefaultClientVersion(t *testing.T) {
var gotClientVersion string
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
@@ -277,6 +277,13 @@ func (s *OpenAIGatewayService) ForwardAsChatCompletions(
// 8. Handle error response with failover
if resp.StatusCode >= 400 {
respBody, upstreamMsg := s.readOpenAIUpstreamError(resp)
if !agentIdentityTaskRecoveryWasTried(ctx) && s.isAgentIdentityAccount(ctx, account) && isAgentIdentityTaskInvalidHTTPResponse(resp.StatusCode, respBody) {
expectedTaskID := account.GetCredential("task_id")
if err := s.recoverAgentIdentityTask(ctx, account, expectedTaskID); err != nil {
return nil, fmt.Errorf("agent identity task recovery failed: %w", err)
}
return s.ForwardAsChatCompletions(markAgentIdentityTaskRecoveryTried(ctx), c, account, body, promptCacheKey, defaultMappedModel)
}
if account.Type == AccountTypeAPIKey &&
openai_compat.ResolveResponsesSupport(account.Extra) == openai_compat.ResponsesSupportUnknown &&
!isResponsesEndpointSupportedByStatus(resp.StatusCode) {
@@ -215,7 +215,15 @@ func (s *OpenAIGatewayService) buildInputTokensUpstreamRequest(
return nil, err
}
req = req.WithContext(WithHTTPUpstreamProfile(req.Context(), HTTPUpstreamProfileOpenAI))
req.Header.Set("authorization", "Bearer "+token)
authHeaders, err := s.buildOpenAIAuthenticationHeaders(ctx, account, token)
if err != nil {
return nil, err
}
for key, values := range authHeaders {
for _, value := range values {
req.Header.Add(key, value)
}
}
req.Header.Set("content-type", "application/json")
req.Header.Set("accept", "application/json")
@@ -499,6 +499,7 @@ func (s *OpenAIGatewayService) Forward(ctx context.Context, c *gin.Context, acco
var wsResult *OpenAIForwardResult
var wsErr error
wsLastFailureReason := ""
agentTaskRecoveryTried := false
wsPrevResponseRecoveryTried := false
wsInvalidEncryptedContentRecoveryTried := false
recoverPrevResponseNotFound := func(attempt int) bool {
@@ -583,6 +584,7 @@ func (s *OpenAIGatewayService) Forward(ctx context.Context, c *gin.Context, acco
startTime,
attempt,
wsLastFailureReason,
&agentTaskRecoveryTried,
)
if wsErr == nil {
break
@@ -590,6 +592,10 @@ func (s *OpenAIGatewayService) Forward(ctx context.Context, c *gin.Context, acco
if c != nil && c.Writer != nil && c.Writer.Written() {
break
}
var taskRecoveredErr *agentIdentityTaskRecoveredError
if errors.As(wsErr, &taskRecoveredErr) {
continue
}
reason, retryable := classifyOpenAIWSReconnectReason(wsErr)
if reason != "" {
@@ -696,6 +702,7 @@ func (s *OpenAIGatewayService) Forward(ctx context.Context, c *gin.Context, acco
}
httpInvalidEncryptedContentRetryTried := false
agentTaskRecoveryTried := false
for {
// Build upstream request
upstreamCtx, releaseUpstreamCtx := detachUpstreamContext(ctx)
@@ -731,6 +738,16 @@ func (s *OpenAIGatewayService) Forward(ctx context.Context, c *gin.Context, acco
upstreamMsg := strings.TrimSpace(extractUpstreamErrorMessage(respBody))
upstreamMsg = sanitizeUpstreamErrorMessage(upstreamMsg)
upstreamCode := extractUpstreamErrorCode(respBody)
if !agentTaskRecoveryTried && s.isAgentIdentityAccount(ctx, account) && isAgentIdentityTaskInvalidHTTPResponse(resp.StatusCode, respBody) {
agentTaskRecoveryTried = true
expectedTaskID := account.GetCredential("task_id")
if err := s.recoverAgentIdentityTask(ctx, account, expectedTaskID); err != nil {
return nil, fmt.Errorf("agent identity task recovery failed: %w", err)
}
continue
}
respBody = s.redactAgentIdentitySensitiveBody(ctx, account, respBody)
resp.Body = io.NopCloser(bytes.NewReader(respBody))
if !httpInvalidEncryptedContentRetryTried && resp.StatusCode == http.StatusBadRequest && upstreamCode == "invalid_encrypted_content" {
decoded, decodeErr := ensureReqBody()
if decodeErr != nil {
@@ -881,8 +898,17 @@ func (s *OpenAIGatewayService) buildUpstreamRequest(ctx context.Context, c *gin.
}
req = req.WithContext(WithHTTPUpstreamProfile(req.Context(), HTTPUpstreamProfileOpenAI))
// Set authentication header
req.Header.Set("authorization", "Bearer "+token)
// Build authentication for this request. Agent Identity signs a fresh
// assertion here; OAuth/PAT/API-key keep their existing Bearer behavior.
authHeaders, err := s.buildOpenAIAuthenticationHeaders(ctx, account, token)
if err != nil {
return nil, fmt.Errorf("build openai authentication headers: %w", err)
}
for key, values := range authHeaders {
for _, value := range values {
req.Header.Add(key, value)
}
}
// Set headers specific to OAuth accounts (ChatGPT internal API)
if account.Type == AccountTypeOAuth {
@@ -323,6 +323,17 @@ func (s *OpenAIGatewayService) ForwardAsAnthropic(
// 8. Handle error response with failover
if resp.StatusCode >= 400 {
respBody, upstreamMsg := s.readOpenAIUpstreamError(resp)
if !agentIdentityTaskRecoveryWasTried(ctx) && s.isAgentIdentityAccount(ctx, account) && isAgentIdentityTaskInvalidHTTPResponse(resp.StatusCode, respBody) {
expectedTaskID := account.GetCredential("task_id")
if err := s.recoverAgentIdentityTask(ctx, account, expectedTaskID); err != nil {
return nil, fmt.Errorf("agent identity task recovery failed: %w", err)
}
return s.ForwardAsAnthropic(markAgentIdentityTaskRecoveryTried(ctx), c, account, body, promptCacheKey, defaultMappedModel)
}
if account.Platform == PlatformGrok {
s.updateGrokUsageSnapshot(ctx, account, xai.ParseQuotaHeaders(resp.Header, resp.StatusCode))
s.handleGrokAccountUpstreamError(ctx, account, resp.StatusCode, resp.Header, respBody)
}
if previousResponseID != "" && (isOpenAICompatPreviousResponseNotFound(resp.StatusCode, upstreamMsg, respBody) || isOpenAICompatPreviousResponseUnsupported(resp.StatusCode, upstreamMsg, respBody)) {
if isOpenAICompatPreviousResponseUnsupported(resp.StatusCode, upstreamMsg, respBody) {
s.disableOpenAICompatSessionContinuation(ctx, c, account, promptCacheKey)
@@ -11,6 +11,7 @@ import (
"encoding/json"
"errors"
"fmt"
"io"
"net/http"
"sort"
"strings"
@@ -161,13 +162,6 @@ func (s *OpenAIGatewayService) forwardOpenAIPassthrough(
return nil, err
}
upstreamCtx, releaseUpstreamCtx := detachUpstreamContext(ctx)
upstreamReq, err := s.buildUpstreamRequestOpenAIPassthrough(upstreamCtx, c, account, body, token)
releaseUpstreamCtx()
if err != nil {
return nil, err
}
proxyURL := ""
if account.ProxyID != nil && account.Proxy != nil {
proxyURL = account.Proxy.URL()
@@ -177,18 +171,42 @@ func (s *OpenAIGatewayService) forwardOpenAIPassthrough(
c.Set("openai_passthrough", true)
}
upstreamStart := time.Now()
resp, err := s.httpUpstream.Do(upstreamReq, proxyURL, account.ID, account.Concurrency)
SetOpsLatencyMs(c, OpsUpstreamLatencyMsKey, time.Since(upstreamStart).Milliseconds())
if err != nil {
// Transport-level failure (proxy/DNS/TCP/TLS — no HTTP response). Convert to
// a failover so the handler switches to a healthy account, and temporarily
// unschedule the account on durable faults (e.g. rejected proxy credentials).
return nil, s.handleOpenAIUpstreamTransportError(ctx, c, account, err, true)
}
defer func() { _ = resp.Body.Close() }()
agentTaskRecoveryTried := false
var resp *http.Response
for {
upstreamCtx, releaseUpstreamCtx := detachUpstreamContext(ctx)
upstreamReq, buildErr := s.buildUpstreamRequestOpenAIPassthrough(upstreamCtx, c, account, body, token)
releaseUpstreamCtx()
if buildErr != nil {
return nil, buildErr
}
upstreamStart := time.Now()
resp, err = s.httpUpstream.Do(upstreamReq, proxyURL, account.ID, account.Concurrency)
SetOpsLatencyMs(c, OpsUpstreamLatencyMsKey, time.Since(upstreamStart).Milliseconds())
if err != nil {
// Transport-level failure (proxy/DNS/TCP/TLS — no HTTP response). Convert to
// a failover so the handler switches to a healthy account.
return nil, s.handleOpenAIUpstreamTransportError(ctx, c, account, err, true)
}
if resp.StatusCode < 400 {
break
}
// Peek only to identify an invalid task. Restore the body so the existing
// passthrough error handling sees the same response after recovery fails.
probeBody := s.readUpstreamErrorBody(resp)
_ = resp.Body.Close()
resp.Body = io.NopCloser(bytes.NewReader(probeBody))
if !agentTaskRecoveryTried && s.isAgentIdentityAccount(ctx, account) && isAgentIdentityTaskInvalidHTTPResponse(resp.StatusCode, probeBody) {
agentTaskRecoveryTried = true
expectedTaskID := account.GetCredential("task_id")
if recoveryErr := s.recoverAgentIdentityTask(ctx, account, expectedTaskID); recoveryErr != nil {
return nil, fmt.Errorf("agent identity task recovery failed: %w", recoveryErr)
}
continue
}
if resp.StatusCode >= 400 {
// 透传模式默认保持原样代理;但 429/529 属于网关必须兜底的
// 上游容量类错误,应先触发多账号 failover 以维持基础 SLA。
if shouldFailoverOpenAIPassthroughResponse(resp.StatusCode) {
@@ -196,6 +214,7 @@ func (s *OpenAIGatewayService) forwardOpenAIPassthrough(
}
return nil, s.handleErrorResponsePassthrough(ctx, resp, c, account, body)
}
defer func() { _ = resp.Body.Close() }()
serviceTier := extractOpenAIServiceTierFromBody(body)
@@ -338,7 +357,15 @@ func (s *OpenAIGatewayService) buildUpstreamRequestOpenAIPassthrough(
req.Header.Del("authorization")
req.Header.Del("x-api-key")
req.Header.Del("x-goog-api-key")
req.Header.Set("authorization", "Bearer "+token)
authHeaders, err := s.buildOpenAIAuthenticationHeaders(ctx, account, token)
if err != nil {
return nil, fmt.Errorf("build openai authentication headers: %w", err)
}
for key, values := range authHeaders {
for _, value := range values {
req.Header.Add(key, value)
}
}
// OAuth 透传到 ChatGPT internal API 时补齐必要头。
if account.Type == AccountTypeOAuth {
@@ -434,6 +461,7 @@ func (s *OpenAIGatewayService) handleFailoverErrorResponsePassthrough(
requestBody []byte,
) error {
body := s.readUpstreamErrorBody(resp)
body = s.redactAgentIdentitySensitiveBody(ctx, account, body)
upstreamMsg := strings.TrimSpace(extractUpstreamErrorMessage(body))
upstreamMsg = sanitizeUpstreamErrorMessage(upstreamMsg)
@@ -477,6 +505,7 @@ func (s *OpenAIGatewayService) handleErrorResponsePassthrough(
) error {
MarkResponseCommitted(c)
body := s.readUpstreamErrorBody(resp)
body = s.redactAgentIdentitySensitiveBody(ctx, account, body)
// cyber_policy:透传账号本就把原始 body 回给客户端(下方 c.Data),此处仅打标记,
// 供 handler 事后写风控/邮件。cyber 是上游网络安全策略拦截,不冷却账号,
@@ -393,6 +393,7 @@ type OpenAIGatewayService struct {
openaiWSStateStoreOnce sync.Once
openaiSchedulerOnce sync.Once
openaiWSPassthroughDialerOnce sync.Once
agentIdentityTaskMu sync.Mutex
openaiWSPool *openAIWSConnPool
openaiWSStateStore OpenAIWSStateStore
openaiScheduler OpenAIAccountScheduler
@@ -595,6 +596,12 @@ func (s *OpenAIGatewayService) CloseOpenAIWSPool() {
}
}
func (s *OpenAIGatewayService) InvalidateAgentIdentityWSConnections(accountID int64) {
if pool := s.getOpenAIWSConnPool(); pool != nil {
pool.ClearAccount(accountID)
}
}
func (s *OpenAIGatewayService) logOpenAIWSModeBootstrap() {
if s == nil || s.cfg == nil {
return
@@ -1094,6 +1101,9 @@ func (s *OpenAIGatewayService) GetAccessToken(ctx context.Context, account *Acco
}
switch account.Type {
case AccountTypeOAuth:
if account.IsOpenAIAgentIdentity() {
return "", OpenAIAuthModeAgentIdentity, nil
}
if account.Platform == PlatformGrok {
if s.grokTokenProvider != nil {
accessToken, err := s.grokTokenProvider.GetAccessToken(ctx, account)
@@ -279,6 +279,7 @@ func (s *OpenAIGatewayService) handleErrorResponse(
requestedModel ...string,
) (*OpenAIForwardResult, error) {
body := s.readUpstreamErrorBody(resp)
body = s.redactAgentIdentitySensitiveBody(ctx, account, body)
// cyber_policy 硬阻断:透传上游原始错误体给客户端(不重包成通用 502),不冷却账号。
// 当前请求恒透传(需求1);标记供 handler 事后写风控/邮件。400 cyber 不可 failover
@@ -470,6 +471,7 @@ func (s *OpenAIGatewayService) handleCompatErrorResponse(
requestedModel ...string,
) (*OpenAIForwardResult, error) {
body := s.readUpstreamErrorBody(resp)
body = s.redactAgentIdentitySensitiveBody(context.Background(), account, body)
// cyber_policy:兼容路径(Chat Completions / Anthropic)以各自格式回写错误,
// 不原样透传 responses 格式的 cyber body(否则对下游格式不合法)。cyber 是上游网络
+10 -1
View File
@@ -632,6 +632,7 @@ func (s *OpenAIGatewayService) forwardOpenAIImagesAPIKey(
if resp.StatusCode >= 400 {
respBody := s.readUpstreamErrorBody(resp)
_ = resp.Body.Close()
respBody = s.redactAgentIdentitySensitiveBody(upstreamCtx, account, respBody)
resp.Body = io.NopCloser(bytes.NewReader(respBody))
upstreamMsg := strings.TrimSpace(extractUpstreamErrorMessage(respBody))
upstreamMsg = sanitizeUpstreamErrorMessage(upstreamMsg)
@@ -752,7 +753,15 @@ func (s *OpenAIGatewayService) buildOpenAIImagesRequest(
return nil, err
}
req = req.WithContext(WithHTTPUpstreamProfile(req.Context(), HTTPUpstreamProfileOpenAI))
req.Header.Set("Authorization", "Bearer "+token)
authHeaders, err := s.buildOpenAIAuthenticationHeaders(ctx, account, token)
if err != nil {
return nil, fmt.Errorf("build openai authentication headers: %w", err)
}
for key, values := range authHeaders {
for _, value := range values {
req.Header.Add(key, value)
}
}
for key, values := range c.Request.Header {
if !openaiPassthroughAllowedHeaders[strings.ToLower(key)] {
continue
@@ -1571,6 +1571,14 @@ func (s *OpenAIGatewayService) forwardOpenAIImagesOAuth(
if resp.StatusCode >= 400 {
respBody := s.readUpstreamErrorBody(resp)
_ = resp.Body.Close()
respBody = s.redactAgentIdentitySensitiveBody(upstreamCtx, account, respBody)
if !agentIdentityTaskRecoveryWasTried(ctx) && s.isAgentIdentityAccount(ctx, account) && isAgentIdentityTaskInvalidHTTPResponse(resp.StatusCode, respBody) {
expectedTaskID := account.GetCredential("task_id")
if err := s.recoverAgentIdentityTask(ctx, account, expectedTaskID); err != nil {
return nil, fmt.Errorf("agent identity task recovery failed: %w", err)
}
return s.forwardOpenAIImagesOAuth(markAgentIdentityTaskRecoveryTried(ctx), c, account, parsed, channelMappedModel)
}
resp.Body = io.NopCloser(bytes.NewReader(respBody))
upstreamMsg := strings.TrimSpace(extractUpstreamErrorMessage(respBody))
upstreamMsg = sanitizeUpstreamErrorMessage(upstreamMsg)
+145 -26
View File
@@ -8,6 +8,7 @@ import (
"log/slog"
"net/http"
"strings"
"sync"
"time"
infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
@@ -21,6 +22,8 @@ import (
// errors.Is still matches it by identity since ResetCredit returns this var.
var ErrSparkShadowResetNotSupported = infraerrors.New(http.StatusConflict, "SPARK_SHADOW_RESET_NOT_SUPPORTED", "spark shadow account does not support credit reset; reset the parent account")
var ErrAgentIdentityResetNotSupported = infraerrors.New(http.StatusConflict, "AGENT_IDENTITY_RESET_NOT_SUPPORTED", "agent identity does not support rate-limit reset credit consumption")
// Endpoints used by the OpenAI/ChatGPT/Codex quota query and reset feature.
const (
chatGPTUsageURL = "https://chatgpt.com/backend-api/wham/usage"
@@ -116,6 +119,8 @@ type OpenAIQuotaService struct {
proxyRepo ProxyRepository
tokenProvider *OpenAITokenProvider
privacyClientFactory PrivacyClientFactory
agentIdentityTaskMu sync.Mutex
agentIdentityWS agentIdentityWSConnectionInvalidator
}
// NewOpenAIQuotaService constructs a quota service. token provider is required —
@@ -151,21 +156,36 @@ func (s *OpenAIQuotaService) QueryUsage(ctx context.Context, accountID int64) (*
callCtx, cancel := context.WithTimeout(ctx, openaiQuotaUpstreamTimeout)
defer cancel()
agentIdentity := s.isAgentIdentityAccount(ctx, accountID)
var payload OpenAIQuotaUsage
resp, err := client.R().
SetContext(callCtx).
SetHeaders(buildCodexCommonHeaders(accessToken, chatGPTAccountID, fedRAMP)).
SetSuccessResult(&payload).
Get(chatGPTUsageURL)
if err != nil {
return nil, infraerrors.Newf(http.StatusBadGateway, "OPENAI_QUOTA_REQUEST_FAILED", "upstream request failed: %v", err)
}
if !resp.IsSuccessState() {
status := resp.StatusCode
body := truncate(resp.String(), 240)
slog.Warn("openai_quota_query_failed", "account_id", accountID, "status", status, "body", body)
return nil, infraerrors.Newf(mapUpstreamStatus(status), "OPENAI_QUOTA_UPSTREAM_ERROR", "upstream returned %d: %s", status, body)
for recovered := false; ; {
quotaHeaders, headerErr := s.buildCodexQuotaHeaders(callCtx, accountID, accessToken, chatGPTAccountID, fedRAMP)
if headerErr != nil {
return nil, infraerrors.Newf(http.StatusBadGateway, "OPENAI_QUOTA_AUTH_FAILED", "failed to build upstream authentication: %v", headerErr)
}
resp, err := client.R().
SetContext(callCtx).
SetHeaders(quotaHeaders).
SetSuccessResult(&payload).
Get(chatGPTUsageURL)
if err != nil {
return nil, infraerrors.Newf(http.StatusBadGateway, "OPENAI_QUOTA_REQUEST_FAILED", "upstream request failed: %v", err)
}
if !resp.IsSuccessState() {
if agentIdentity && !recovered && isAgentIdentityTaskInvalidHTTPResponse(resp.StatusCode, []byte(resp.String())) {
recovered = true
if err := s.recoverAgentIdentityTask(ctx, accountID); err != nil {
return nil, infraerrors.Newf(http.StatusBadGateway, "OPENAI_QUOTA_AUTH_FAILED", "agent identity task recovery failed: %v", err)
}
continue
}
status := resp.StatusCode
body := truncate(s.redactQuotaErrorBody(ctx, accountID, resp.String()), 240)
slog.Warn("openai_quota_query_failed", "account_id", accountID, "status", status, "body", body)
return nil, infraerrors.Newf(mapUpstreamStatus(status), "OPENAI_QUOTA_UPSTREAM_ERROR", "upstream returned %d: %s", status, body)
}
break
}
payload.FetchedAt = time.Now().Unix()
@@ -189,9 +209,14 @@ func (s *OpenAIQuotaService) QueryUsage(ctx context.Context, accountID int64) (*
}
func (s *OpenAIQuotaService) queryResetCreditDetails(ctx context.Context, client *req.Client, accessToken, chatGPTAccountID string, fedRAMP bool, accountID int64) *openAIRateLimitResetCreditDetails {
quotaHeaders, headerErr := s.buildCodexQuotaHeaders(ctx, accountID, accessToken, chatGPTAccountID, fedRAMP)
if headerErr != nil {
slog.Warn("openai_quota_reset_credit_details_auth_failed", "account_id", accountID, "error", headerErr)
return nil
}
resp, err := client.R().
SetContext(ctx).
SetHeaders(buildCodexCommonHeaders(accessToken, chatGPTAccountID, fedRAMP)).
SetHeaders(quotaHeaders).
Get(chatGPTRateLimitCreditsURL)
if err != nil {
slog.Warn("openai_quota_reset_credit_details_failed", "account_id", accountID, "error", err)
@@ -233,6 +258,9 @@ func (s *OpenAIQuotaService) ResetCredit(ctx context.Context, accountID int64) (
if acc.IsShadow() {
return nil, ErrSparkShadowResetNotSupported
}
if acc.IsOpenAIAgentIdentity() {
return nil, ErrAgentIdentityResetNotSupported
}
}
accessToken, chatGPTAccountID, proxyURL, fedRAMP, err := s.prepareUpstreamCall(ctx, accountID)
@@ -252,11 +280,12 @@ func (s *OpenAIQuotaService) ResetCredit(ctx context.Context, accountID int64) (
callCtx, cancel := context.WithTimeout(ctx, openaiQuotaUpstreamTimeout)
defer cancel()
headers := buildCodexCommonHeaders(accessToken, chatGPTAccountID, fedRAMP)
headers["content-type"] = "application/json"
var payload OpenAIQuotaResetResult
headers, headerErr := s.buildCodexQuotaHeaders(callCtx, accountID, accessToken, chatGPTAccountID, fedRAMP)
if headerErr != nil {
return nil, infraerrors.Newf(http.StatusBadGateway, "OPENAI_QUOTA_AUTH_FAILED", "failed to build upstream authentication: %v", headerErr)
}
headers["content-type"] = "application/json"
resp, err := client.R().
SetContext(callCtx).
SetHeaders(headers).
@@ -268,7 +297,7 @@ func (s *OpenAIQuotaService) ResetCredit(ctx context.Context, accountID int64) (
}
if !resp.IsSuccessState() {
status := resp.StatusCode
body := truncate(resp.String(), 240)
body := truncate(s.redactQuotaErrorBody(callCtx, accountID, resp.String()), 240)
slog.Warn("openai_quota_reset_failed", "account_id", accountID, "status", status, "body", body)
return nil, infraerrors.Newf(mapUpstreamStatus(status), "OPENAI_QUOTA_RESET_UPSTREAM_ERROR", "upstream returned %d: %s", status, body)
}
@@ -285,7 +314,7 @@ func (s *OpenAIQuotaService) ResetCredit(ctx context.Context, accountID int64) (
// token via the shared TokenProvider, and resolves the chatgpt-account-id and
// proxy URL. Centralized so QueryUsage / ResetCredit share validation.
func (s *OpenAIQuotaService) prepareUpstreamCall(ctx context.Context, accountID int64) (accessToken, chatGPTAccountID, proxyURL string, fedRAMP bool, err error) {
if s == nil || s.accountRepo == nil || s.tokenProvider == nil || s.privacyClientFactory == nil {
if s == nil || s.accountRepo == nil || s.privacyClientFactory == nil {
return "", "", "", false, infraerrors.New(http.StatusInternalServerError, "OPENAI_QUOTA_NOT_CONFIGURED", "openai quota service is not configured")
}
@@ -323,12 +352,17 @@ func (s *OpenAIQuotaService) prepareUpstreamCall(ctx context.Context, accountID
return "", "", "", false, infraerrors.New(http.StatusBadRequest, "OPENAI_QUOTA_MISSING_ACCOUNT_ID", "chatgpt_account_id is missing; please re-authorize this account")
}
accessToken, err = s.tokenProvider.GetAccessToken(ctx, account)
if err != nil {
return "", "", "", false, infraerrors.Newf(http.StatusBadGateway, "OPENAI_QUOTA_TOKEN_UNAVAILABLE", "failed to acquire access token: %v", err)
}
if strings.TrimSpace(accessToken) == "" {
return "", "", "", false, infraerrors.New(http.StatusBadGateway, "OPENAI_QUOTA_TOKEN_UNAVAILABLE", "access token is empty")
if !account.IsOpenAIAgentIdentity() {
if s.tokenProvider == nil {
return "", "", "", false, infraerrors.New(http.StatusInternalServerError, "OPENAI_QUOTA_NOT_CONFIGURED", "openai quota token provider is not configured")
}
accessToken, err = s.tokenProvider.GetAccessToken(ctx, account)
if err != nil {
return "", "", "", false, infraerrors.Newf(http.StatusBadGateway, "OPENAI_QUOTA_TOKEN_UNAVAILABLE", "failed to acquire access token: %v", err)
}
if strings.TrimSpace(accessToken) == "" {
return "", "", "", false, infraerrors.New(http.StatusBadGateway, "OPENAI_QUOTA_TOKEN_UNAVAILABLE", "access token is empty")
}
}
fedRAMP = account.IsChatGPTAccountFedRAMP()
@@ -351,6 +385,91 @@ func (s *OpenAIQuotaService) prepareUpstreamCall(ctx context.Context, accountID
return accessToken, chatGPTAccountID, proxyURL, fedRAMP, nil
}
func (s *OpenAIQuotaService) recoverAgentIdentityTask(ctx context.Context, accountID int64) error {
if s == nil || s.accountRepo == nil {
return fmt.Errorf("account repository is unavailable")
}
account, err := s.accountRepo.GetByID(ctx, accountID)
if err != nil || account == nil {
return fmt.Errorf("account is unavailable")
}
if account.IsShadow() {
account, err = resolveCredentialAccount(ctx, s.accountRepo, account)
if err != nil || account == nil {
return fmt.Errorf("credential account is unavailable")
}
}
if !account.IsOpenAIAgentIdentity() {
return nil
}
return ensureAgentIdentityTaskForAccount(ctx, s.accountRepo, s.agentIdentityWS, &s.agentIdentityTaskMu, account, account.GetCredential("task_id"))
}
func (s *OpenAIQuotaService) isAgentIdentityAccount(ctx context.Context, accountID int64) bool {
if s == nil || s.accountRepo == nil {
return false
}
account, err := s.accountRepo.GetByID(ctx, accountID)
if err != nil || account == nil {
return false
}
if account.IsShadow() {
account, err = resolveCredentialAccount(ctx, s.accountRepo, account)
if err != nil || account == nil {
return false
}
}
return account.IsOpenAIAgentIdentity()
}
func (s *OpenAIQuotaService) buildCodexQuotaHeaders(ctx context.Context, accountID int64, accessToken, chatGPTAccountID string, fedRAMP bool) (map[string]string, error) {
headers := buildCodexCommonHeaders(accessToken, chatGPTAccountID, fedRAMP)
if s == nil || s.accountRepo == nil {
return headers, nil
}
account, err := s.accountRepo.GetByID(ctx, accountID)
if err != nil || account == nil {
if strings.TrimSpace(accessToken) == "" {
return nil, fmt.Errorf("agent identity account credentials are unavailable")
}
return headers, nil
}
if account.IsShadow() {
if resolved, resolveErr := resolveCredentialAccount(ctx, s.accountRepo, account); resolveErr == nil && resolved != nil {
account = resolved
} else if strings.TrimSpace(accessToken) == "" {
return nil, fmt.Errorf("agent identity shadow credentials are unavailable")
}
}
if !account.IsOpenAIAgentIdentity() {
return headers, nil
}
if err := ensureAgentIdentityTaskForAccount(ctx, s.accountRepo, s.agentIdentityWS, &s.agentIdentityTaskMu, account, ""); err != nil {
return nil, err
}
key, err := agentIdentityKeyFromAccount(account)
if err != nil {
return nil, err
}
assertion, err := buildAgentAssertion(key, time.Now())
if err != nil {
return nil, err
}
headers["authorization"] = assertion
return headers, nil
}
func (s *OpenAIQuotaService) redactQuotaErrorBody(ctx context.Context, accountID int64, body string) string {
if s == nil || s.accountRepo == nil {
return body
}
account, err := s.accountRepo.GetByID(ctx, accountID)
if err != nil || account == nil {
return body
}
return string(redactAgentIdentitySensitiveBodyForAccount(ctx, s.accountRepo, account, []byte(body)))
}
// buildCodexCommonHeaders sets the request headers expected by the chatgpt.com
// backend so calls succeed past Cloudflare/WASM checks.
func buildCodexCommonHeaders(accessToken, chatGPTAccountID string, fedRAMP bool) map[string]string {
@@ -2,6 +2,10 @@ package service
import (
"context"
"crypto/ed25519"
"crypto/rand"
"crypto/x509"
"encoding/base64"
"encoding/json"
"errors"
"fmt"
@@ -34,6 +38,15 @@ func (r *stubQuotaAccountRepo) GetByID(_ context.Context, id int64) (*Account, e
return acc, nil
}
func (r *stubQuotaAccountRepo) UpdateCredentials(_ context.Context, id int64, credentials map[string]any) error {
acc, ok := r.accounts[id]
if !ok {
return fmt.Errorf("account %d not found", id)
}
acc.Credentials = credentials
return nil
}
// stubQuotaTokenCache 实现 OpenAITokenCache,返回预设静态 token。
type stubQuotaTokenCache struct {
tokens map[string]string
@@ -163,6 +176,23 @@ func TestResetCreditShadowRejected(t *testing.T) {
"shadow ResetCredit 应映射为 409 Conflict 而非 500")
}
func TestResetCreditAgentIdentityRejectedBeforeUpstream(t *testing.T) {
account := &Account{
ID: 201,
Platform: PlatformOpenAI,
Type: AccountTypeOAuth,
Credentials: map[string]any{
"auth_mode": OpenAIAuthModeAgentIdentity,
},
}
repo := &stubQuotaAccountRepo{accounts: map[int64]*Account{account.ID: account}}
svc := &OpenAIQuotaService{accountRepo: repo}
_, err := svc.ResetCredit(context.Background(), account.ID)
require.ErrorIs(t, err, ErrAgentIdentityResetNotSupported)
require.Equal(t, http.StatusConflict, infraerrors.Code(err))
}
// ── Part B: prepareUpstreamCall 影子 resolve ──────────────────────────────
// TestPrepareUpstreamCallShadowResolve 验证影子账号(200)QueryUsage 时:
@@ -214,6 +244,101 @@ func TestPrepareUpstreamCallShadowResolve(t *testing.T) {
"prepareUpstreamCall should use parent's chatgpt_account_id after shadow resolve")
}
func TestQueryUsageAgentIdentityUsesAssertionWithoutOAuthToken(t *testing.T) {
_, privateKey, err := ed25519.GenerateKey(rand.Reader)
require.NoError(t, err)
der, err := x509.MarshalPKCS8PrivateKey(privateKey)
require.NoError(t, err)
account := &Account{
ID: 300,
Platform: PlatformOpenAI,
Type: AccountTypeOAuth,
Credentials: map[string]any{
"auth_mode": OpenAIAuthModeAgentIdentity,
"agent_runtime_id": "runtime-quota",
"agent_private_key": base64.StdEncoding.EncodeToString(der),
"task_id": "task-quota",
"chatgpt_account_id": "account-quota",
"chatgpt_account_is_fedramp": true,
},
}
repo := &stubQuotaAccountRepo{accounts: map[int64]*Account{account.ID: account}}
var authorization string
var accountHeader string
var fedrampHeader string
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
authorization = r.Header.Get("authorization")
accountHeader = r.Header.Get("chatgpt-account-id")
fedrampHeader = r.Header.Get("x-openai-fedramp")
w.Header().Set("content-type", "application/json")
_, _ = w.Write([]byte(`{"plan_type":"pro","rate_limit":{"allowed":true}}`))
}))
defer srv.Close()
svc := NewOpenAIQuotaService(repo, nil, nil, newQuotaRedirectingFactory(srv))
usage, err := svc.QueryUsage(context.Background(), account.ID)
require.NoError(t, err)
require.NotNil(t, usage)
require.True(t, strings.HasPrefix(authorization, "AgentAssertion "))
require.Equal(t, "account-quota", accountHeader)
require.Equal(t, "true", fedrampHeader)
}
func TestQueryUsageAgentIdentityRecoversInvalidTaskOnce(t *testing.T) {
_, privateKey, err := ed25519.GenerateKey(rand.Reader)
require.NoError(t, err)
der, err := x509.MarshalPKCS8PrivateKey(privateKey)
require.NoError(t, err)
account := &Account{
ID: 301,
Platform: PlatformOpenAI,
Type: AccountTypeOAuth,
Credentials: map[string]any{
"auth_mode": OpenAIAuthModeAgentIdentity,
"agent_runtime_id": "runtime-quota-recovery",
"agent_private_key": base64.StdEncoding.EncodeToString(der),
"task_id": "task-quota-old",
"chatgpt_account_id": "account-quota-recovery",
},
}
repo := &stubQuotaAccountRepo{accounts: map[int64]*Account{account.ID: account}}
usageCalls := 0
registerCalls := 0
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("content-type", "application/json")
if strings.Contains(r.URL.Path, "/task/register") {
registerCalls++
_, _ = w.Write([]byte(`{"task_id":"task-quota-new"}`))
return
}
if strings.Contains(r.URL.Path, "rate-limit-reset-credits") {
_, _ = w.Write([]byte(`{}`))
return
}
usageCalls++
if usageCalls == 1 {
w.WriteHeader(http.StatusUnauthorized)
_, _ = w.Write([]byte(`{"error":{"code":"invalid_task_id"}}`))
return
}
_, _ = w.Write([]byte(`{"plan_type":"pro","rate_limit":{"allowed":true}}`))
}))
defer srv.Close()
oldBase := openAIAgentIdentityAuthAPIBaseURL
openAIAgentIdentityAuthAPIBaseURL = srv.URL
t.Cleanup(func() { openAIAgentIdentityAuthAPIBaseURL = oldBase })
invalidator := &agentIdentityWSInvalidationRecorder{}
svc := NewOpenAIQuotaService(repo, nil, nil, newQuotaRedirectingFactory(srv))
svc.agentIdentityWS = invalidator
usage, err := svc.QueryUsage(context.Background(), account.ID)
require.NoError(t, err)
require.NotNil(t, usage)
require.Equal(t, 2, usageCalls)
require.Equal(t, 1, registerCalls)
require.Equal(t, "task-quota-new", account.GetCredential("task_id"))
require.Equal(t, []int64{account.ID}, invalidator.accountIDs)
}
func TestParseOpenAIRateLimitResetCreditDetails_CompatibleContainers(t *testing.T) {
tests := []struct {
name string
+29 -1
View File
@@ -4,6 +4,7 @@ import (
"context"
"errors"
"fmt"
"io"
"net/http"
"net/url"
"strings"
@@ -68,6 +69,28 @@ type coderOpenAIWSClientDialer struct {
proxyMisses atomic.Int64
}
// openAIWSHandshakeError keeps a bounded, non-logged HTTP error body so the
// Agent Identity recovery path can distinguish an invalid task from other
// 401 handshake failures.
type openAIWSHandshakeError struct {
Body []byte
Err error
}
func (e *openAIWSHandshakeError) Error() string {
if e == nil || e.Err == nil {
return "openai ws handshake failed"
}
return e.Err.Error()
}
func (e *openAIWSHandshakeError) Unwrap() error {
if e == nil {
return nil
}
return e.Err
}
type openAIWSProxyClientEntry struct {
client *http.Client
lastUsedUnixNano int64
@@ -104,7 +127,12 @@ func (d *coderOpenAIWSClientDialer) Dial(
status = resp.StatusCode
respHeaders = cloneHeader(resp.Header)
}
return nil, status, respHeaders, err
var body []byte
if resp != nil && resp.Body != nil {
body, _ = io.ReadAll(io.LimitReader(resp.Body, 8<<10))
_ = resp.Body.Close()
}
return nil, status, respHeaders, &openAIWSHandshakeError{Body: body, Err: err}
}
// coder/websocket 默认单消息读取上限为 32KB,Codex WS 事件(如 rate_limits/大 delta)
// 可能超过该阈值,需显式提高上限,避免本地 read_fail(message too big)。
@@ -46,8 +46,8 @@ func (s *OpenAIGatewayService) ProxyResponsesWebSocketFromClient(
if account == nil {
return errors.New("account is nil")
}
if strings.TrimSpace(token) == "" {
return errors.New("token is empty")
if err := validateOpenAIWSBearerToken(account, token); err != nil {
return err
}
// 预取一次 OpenAI Fast Policy settings,绑定到 ctx,让该 WS session
@@ -577,6 +577,9 @@ func (s *OpenAIGatewayService) ProxyResponsesWebSocketFromClient(
Account: account,
WSURL: wsURL,
Headers: wsHeaders,
HeadersFactory: func(factoryCtx context.Context, headers http.Header) (http.Header, error) {
return s.refreshOpenAIAgentIdentityHeaders(factoryCtx, account, headers)
},
ProxyURL: func() string {
if account.ProxyID != nil && account.Proxy != nil {
return account.Proxy.URL()
@@ -640,7 +643,9 @@ func (s *OpenAIGatewayService) ProxyResponsesWebSocketFromClient(
acquireTimeout = 30 * time.Second
}
acquireTurnLease := func(turn int, preferred string, forcePreferredConn bool) (*openAIWSConnLease, error) {
agentTaskRecoveryTried := false
var acquireTurnLease func(int, string, bool) (*openAIWSConnLease, error)
acquireTurnLease = func(turn int, preferred string, forcePreferredConn bool) (*openAIWSConnLease, error) {
req := cloneOpenAIWSAcquireRequest(baseAcquireReq)
req.PreferredConnID = strings.TrimSpace(preferred)
req.ForcePreferredConn = forcePreferredConn
@@ -649,6 +654,14 @@ func (s *OpenAIGatewayService) ProxyResponsesWebSocketFromClient(
acquireCtx, acquireCancel := context.WithTimeout(ctx, acquireTimeout)
lease, acquireErr := pool.Acquire(acquireCtx, req)
acquireCancel()
var dialErr *openAIWSDialError
if acquireErr != nil && s.isAgentIdentityAccount(ctx, account) && errors.As(acquireErr, &dialErr) && isAgentIdentityTaskInvalidWSDialError(dialErr) && !agentTaskRecoveryTried {
agentTaskRecoveryTried = true
if recoveryErr := s.recoverAgentIdentityTask(ctx, account, account.GetCredential("task_id")); recoveryErr != nil {
return nil, fmt.Errorf("agent identity task recovery failed: %w", recoveryErr)
}
return acquireTurnLease(turn, preferred, forcePreferredConn)
}
if acquireErr != nil {
dialStatus, dialClass, dialCloseStatus, dialCloseReason, dialRespServer, dialRespVia, dialRespCFRay, dialRespReqID := summarizeOpenAIWSDialError(acquireErr)
logOpenAIWSModeInfo(
@@ -15,6 +15,16 @@ import (
"github.com/tidwall/sjson"
)
func validateOpenAIWSBearerToken(account *Account, token string) error {
if account == nil {
return errors.New("account is nil")
}
if strings.TrimSpace(token) == "" && !account.IsOpenAIAgentIdentity() {
return errors.New("token is empty")
}
return nil
}
func (s *OpenAIGatewayService) buildOpenAIResponsesWSURL(account *Account) (string, error) {
if account == nil {
return "", errors.New("account is nil")
@@ -67,7 +77,9 @@ func (s *OpenAIGatewayService) buildOpenAIWSHeaders(
promptCacheKey string,
) (http.Header, openAIWSSessionHeaderResolution, error) {
headers := make(http.Header)
headers.Set("authorization", "Bearer "+token)
if account == nil || !account.IsOpenAIAgentIdentity() {
headers.Set("authorization", "Bearer "+token)
}
sessionResolution := resolveOpenAIWSSessionHeaders(c, promptCacheKey)
if c != nil && c.Request != nil {
@@ -30,6 +30,7 @@ func (s *OpenAIGatewayService) forwardOpenAIWSV2(
startTime time.Time,
attempt int,
lastFailureReason string,
agentTaskRecoveryTried *bool,
) (*OpenAIForwardResult, error) {
if s == nil || account == nil {
return nil, wrapOpenAIWSFallback("invalid_state", errors.New("service or account is nil"))
@@ -174,9 +175,12 @@ func (s *OpenAIGatewayService) forwardOpenAIWSV2(
defer acquireCancel()
lease, err := s.getOpenAIWSConnPool().Acquire(acquireCtx, openAIWSAcquireRequest{
Account: account,
WSURL: wsURL,
Headers: wsHeaders,
Account: account,
WSURL: wsURL,
Headers: wsHeaders,
HeadersFactory: func(factoryCtx context.Context, headers http.Header) (http.Header, error) {
return s.refreshOpenAIAgentIdentityHeaders(factoryCtx, account, headers)
},
PreferredConnID: preferredConnID,
ForceNewConn: forceNewConn,
ProxyURL: func() string {
@@ -187,6 +191,14 @@ func (s *OpenAIGatewayService) forwardOpenAIWSV2(
}(),
})
if err != nil {
var agentDialErr *openAIWSDialError
if s.isAgentIdentityAccount(ctx, account) && errors.As(err, &agentDialErr) && isAgentIdentityTaskInvalidWSDialError(agentDialErr) && agentTaskRecoveryTried != nil && !*agentTaskRecoveryTried {
*agentTaskRecoveryTried = true
if recoveryErr := s.recoverAgentIdentityTask(ctx, account, account.GetCredential("task_id")); recoveryErr != nil {
return nil, fmt.Errorf("agent identity task recovery failed: %w", recoveryErr)
}
return nil, &agentIdentityTaskRecoveredError{}
}
dialStatus, dialClass, dialCloseStatus, dialCloseReason, dialRespServer, dialRespVia, dialRespCFRay, dialRespReqID := summarizeOpenAIWSDialError(err)
logOpenAIWSModeInfo(
"acquire_fail account_id=%d account_type=%s transport=%s reason=%s dial_status=%d dial_class=%s dial_close_status=%s dial_close_reason=%s dial_resp_server=%s dial_resp_via=%s dial_resp_cf_ray=%s dial_resp_x_request_id=%s cause=%s preferred_conn_id=%s force_new_conn=%v ws_host=%s ws_path=%s proxy_enabled=%v",
+66 -6
View File
@@ -39,6 +39,7 @@ var (
type openAIWSDialError struct {
StatusCode int
ResponseHeaders http.Header
ResponseBody []byte
Err error
}
@@ -60,9 +61,13 @@ func (e *openAIWSDialError) Unwrap() error {
}
type openAIWSAcquireRequest struct {
Account *Account
WSURL string
Headers http.Header
Account *Account
WSURL string
Headers http.Header
// HeadersFactory is evaluated inside dialConn. It exists so credentials
// whose authorization is per-dial (Agent Identity) are never cached in
// lastAcquire or delayed prewarm state.
HeadersFactory func(context.Context, http.Header) (http.Header, error)
ProxyURL string
PreferredConnID string
// ForceNewConn: 强制本次获取新连接(避免复用导致连接内续链状态互相污染)。
@@ -544,6 +549,7 @@ type openAIWSAccountPool struct {
pinnedConns map[string]int
changedCh chan struct{}
creating int
generation uint64
lastCleanupAt time.Time
lastAcquire *openAIWSAcquireRequest
prewarmActive bool
@@ -1391,6 +1397,7 @@ func (p *openAIWSConnPool) ensureTargetIdleAsync(accountID int64) {
}
var req openAIWSAcquireRequest
generation := uint64(0)
need := 0
ap, ok := p.getAccountPool(accountID)
if !ok || ap == nil {
@@ -1425,6 +1432,7 @@ func (p *openAIWSConnPool) ensureTargetIdleAsync(accountID int64) {
return
}
req = cloneOpenAIWSAcquireRequest(*ap.lastAcquire)
generation = ap.generation
ap.prewarmActive = true
if cooldown := p.prewarmCooldown(); cooldown > 0 {
ap.prewarmUntil = now.Add(cooldown)
@@ -1432,7 +1440,7 @@ func (p *openAIWSConnPool) ensureTargetIdleAsync(accountID int64) {
ap.creating += need
p.metrics.scaleUpTotal.Add(int64(need))
go p.prewarmConns(accountID, req, need)
go p.prewarmConns(accountID, req, need, generation)
}
func (p *openAIWSConnPool) targetConnCountLocked(ap *openAIWSAccountPool, maxConns int) int {
@@ -1475,7 +1483,11 @@ func (p *openAIWSConnPool) targetConnCountLocked(ap *openAIWSAccountPool, maxCon
return target
}
func (p *openAIWSConnPool) prewarmConns(accountID int64, req openAIWSAcquireRequest, total int) {
func (p *openAIWSConnPool) prewarmConns(accountID int64, req openAIWSAcquireRequest, total int, generations ...uint64) {
generation := uint64(0)
if len(generations) > 0 {
generation = generations[0]
}
defer func() {
if ap, ok := p.getAccountPool(accountID); ok && ap != nil {
ap.mu.Lock()
@@ -1507,6 +1519,11 @@ func (p *openAIWSConnPool) prewarmConns(accountID int64, req openAIWSAcquireRequ
ap.mu.Unlock()
continue
}
if ap.generation != generation || ap.lastAcquire == nil {
ap.mu.Unlock()
conn.close()
continue
}
if len(ap.conns) >= p.effectiveMaxConnsByAccount(req.Account) {
ap.signalChangedLocked()
ap.mu.Unlock()
@@ -1521,6 +1538,35 @@ func (p *openAIWSConnPool) prewarmConns(accountID int64, req openAIWSAcquireRequ
}
}
// ClearAccount closes all pooled connections and discards delayed prewarm
// state for one account. The generation guard prevents an in-flight prewarm
// started before credential recovery from re-entering the pool afterwards.
func (p *openAIWSConnPool) ClearAccount(accountID int64) {
if p == nil || accountID <= 0 {
return
}
ap, ok := p.getAccountPool(accountID)
if !ok || ap == nil {
return
}
ap.mu.Lock()
ap.generation++
conns := make([]*openAIWSConn, 0, len(ap.conns))
for id, conn := range ap.conns {
delete(ap.conns, id)
delete(ap.pinnedConns, id)
if conn != nil {
conns = append(conns, conn)
}
}
ap.lastAcquire = nil
ap.prewarmUntil = time.Time{}
ap.prewarmFails = 0
ap.prewarmFailAt = time.Time{}
ap.mu.Unlock()
closeOpenAIWSConns(conns)
}
func (p *openAIWSConnPool) evictConn(accountID int64, connID string) {
if p == nil || accountID <= 0 || stringsTrim(connID) == "" {
return
@@ -1599,11 +1645,25 @@ func (p *openAIWSConnPool) dialConn(ctx context.Context, req openAIWSAcquireRequ
if p == nil || p.clientDialer == nil {
return nil, errors.New("openai ws client dialer is nil")
}
conn, status, handshakeHeaders, err := p.clientDialer.Dial(ctx, req.WSURL, req.Headers, req.ProxyURL)
headers := cloneHeader(req.Headers)
var err error
if req.HeadersFactory != nil {
headers, err = req.HeadersFactory(ctx, headers)
if err != nil {
return nil, err
}
}
conn, status, handshakeHeaders, err := p.clientDialer.Dial(ctx, req.WSURL, headers, req.ProxyURL)
if err != nil {
var handshakeErr *openAIWSHandshakeError
var responseBody []byte
if errors.As(err, &handshakeErr) && handshakeErr != nil {
responseBody = append([]byte(nil), handshakeErr.Body...)
}
return nil, &openAIWSDialError{
StatusCode: status,
ResponseHeaders: cloneHeader(handshakeHeaders),
ResponseBody: responseBody,
Err: err,
}
}
@@ -243,8 +243,8 @@ func (s *OpenAIGatewayService) proxyResponsesWebSocketV2Passthrough(
if account == nil {
return errors.New("account is nil")
}
if strings.TrimSpace(token) == "" {
return errors.New("token is empty")
if err := validateOpenAIWSBearerToken(account, token); err != nil {
return err
}
requestModel := strings.TrimSpace(gjson.GetBytes(firstClientMessage, "model").String())
requestPreviousResponseID := strings.TrimSpace(gjson.GetBytes(firstClientMessage, "previous_response_id").String())
@@ -359,10 +359,34 @@ func (s *OpenAIGatewayService) proxyResponsesWebSocketV2Passthrough(
return errors.New("openai ws passthrough dialer is nil")
}
dialCtx, cancelDial := context.WithTimeout(ctx, s.openAIWSDialTimeout())
defer cancelDial()
upstreamConn, statusCode, handshakeHeaders, err := dialer.Dial(dialCtx, wsURL, headers, proxyURL)
if err != nil {
agentTaskRecoveryTried := false
var upstreamConn openAIWSClientConn
statusCode := 0
var handshakeHeaders http.Header
for {
headers, err = s.refreshOpenAIAgentIdentityHeaders(ctx, account, headers)
if err != nil {
return fmt.Errorf("refresh ws authentication headers: %w", err)
}
dialCtx, cancelDial := context.WithTimeout(ctx, s.openAIWSDialTimeout())
upstreamConn, statusCode, handshakeHeaders, err = dialer.Dial(dialCtx, wsURL, headers, proxyURL)
cancelDial()
if err == nil {
break
}
var handshakeErr *openAIWSHandshakeError
responseBody := []byte(nil)
if errors.As(err, &handshakeErr) && handshakeErr != nil {
responseBody = handshakeErr.Body
}
dialErr := &openAIWSDialError{StatusCode: statusCode, ResponseBody: responseBody, Err: err}
if s.isAgentIdentityAccount(ctx, account) && isAgentIdentityTaskInvalidWSDialError(dialErr) && !agentTaskRecoveryTried {
agentTaskRecoveryTried = true
if recoveryErr := s.recoverAgentIdentityTask(ctx, account, account.GetCredential("task_id")); recoveryErr != nil {
return fmt.Errorf("agent identity task recovery failed: %w", recoveryErr)
}
continue
}
logOpenAIWSV2Passthrough(
"relay_dial_failed account_id=%d status_code=%d err=%s",
account.ID,
@@ -678,9 +702,15 @@ func (s *OpenAIGatewayService) mapOpenAIWSPassthroughDialError(
wrappedErr := err
var dialErr *openAIWSDialError
if !errors.As(err, &dialErr) {
var handshakeErr *openAIWSHandshakeError
var responseBody []byte
if errors.As(err, &handshakeErr) && handshakeErr != nil {
responseBody = append([]byte(nil), handshakeErr.Body...)
}
wrappedErr = &openAIWSDialError{
StatusCode: statusCode,
ResponseHeaders: cloneHeader(handshakeHeaders),
ResponseBody: responseBody,
Err: err,
}
}
+62 -3
View File
@@ -131,8 +131,67 @@ func ProvideOpenAIQuotaService(
proxyRepo ProxyRepository,
tokenProvider *OpenAITokenProvider,
privacyClientFactory PrivacyClientFactory,
openAIGatewayService *OpenAIGatewayService,
) *OpenAIQuotaService {
return NewOpenAIQuotaService(accountRepo, proxyRepo, tokenProvider, privacyClientFactory)
service := NewOpenAIQuotaService(accountRepo, proxyRepo, tokenProvider, privacyClientFactory)
service.agentIdentityWS = openAIGatewayService
return service
}
func ProvideAccountUsageService(
accountRepo AccountRepository,
usageLogRepo UsageLogRepository,
usageFetcher ClaudeUsageFetcher,
geminiQuotaService *GeminiQuotaService,
antigravityQuotaFetcher *AntigravityQuotaFetcher,
grokQuotaFetcher *GrokQuotaFetcher,
grokQuotaService *GrokQuotaService,
openAIQuotaService *OpenAIQuotaService,
cache *UsageCache,
identityCache IdentityCache,
tlsFPProfileService *TLSFingerprintProfileService,
openAIGatewayService *OpenAIGatewayService,
) *AccountUsageService {
service := NewAccountUsageService(
accountRepo,
usageLogRepo,
usageFetcher,
geminiQuotaService,
antigravityQuotaFetcher,
grokQuotaFetcher,
grokQuotaService,
openAIQuotaService,
cache,
identityCache,
tlsFPProfileService,
)
service.agentIdentityWS = openAIGatewayService
return service
}
func ProvideAccountTestService(
accountRepo AccountRepository,
geminiTokenProvider *GeminiTokenProvider,
claudeTokenProvider *ClaudeTokenProvider,
grokTokenProvider *GrokTokenProvider,
antigravityGatewayService *AntigravityGatewayService,
httpUpstream HTTPUpstream,
cfg *config.Config,
tlsFPProfileService *TLSFingerprintProfileService,
openAIGatewayService *OpenAIGatewayService,
) *AccountTestService {
service := NewAccountTestService(
accountRepo,
geminiTokenProvider,
claudeTokenProvider,
grokTokenProvider,
antigravityGatewayService,
httpUpstream,
cfg,
tlsFPProfileService,
)
service.agentIdentityWS = openAIGatewayService
return service
}
func ProvideGrokQuotaService(
@@ -602,8 +661,8 @@ var ProviderSet = wire.NewSet(
ProvideClaudeTokenProvider,
NewAntigravityGatewayService,
ProvideRateLimitService,
NewAccountUsageService,
NewAccountTestService,
ProvideAccountUsageService,
ProvideAccountTestService,
ProvideSettingService,
NewDataManagementService,
ProvideBackupService,
@@ -276,7 +276,7 @@
</div>
</div>
<!-- Codex OAuth/session JSON batch import -->
<!-- Codex auth.json / session credential batch import -->
<div v-if="inputMethod === 'codex_session'" class="space-y-4">
<div
class="rounded-lg border border-blue-300 bg-white/80 p-4 dark:border-blue-600 dark:bg-gray-800/80"
@@ -77,6 +77,7 @@ const { t } = useI18n()
interface Props {
platform: AccountPlatform
type: AccountType
authMode?: string
planType?: string
privacyMode?: string
subscriptionExpiresAt?: string
@@ -92,7 +93,15 @@ const platformLabel = computed(() => {
return 'Gemini'
})
const normalizedAuthMode = computed(() =>
(props.authMode || '').trim().toLowerCase().replace(/[\s_-]+/g, '')
)
const typeLabel = computed(() => {
if (props.platform === 'openai' && props.type === 'oauth') {
if (normalizedAuthMode.value === 'agentidentity') return 'Agent Identity'
if (normalizedAuthMode.value === 'personalaccesstoken') return 'PAT'
}
switch (props.type) {
case 'oauth':
return 'OAuth'
@@ -62,3 +62,24 @@ describe('PlatformTypeBadge Grok plans', () => {
expect(wrapper.findAll('path')).toHaveLength(2)
})
})
describe('PlatformTypeBadge OpenAI authentication modes', () => {
it('distinguishes Agent Identity, PAT, and OAuth accounts', async () => {
const wrapper = mount(PlatformTypeBadge, {
props: {
platform: 'openai',
type: 'oauth',
authMode: 'agentIdentity',
},
})
expect(wrapper.text()).toContain('Agent Identity')
await wrapper.setProps({ authMode: 'personalAccessToken' })
expect(wrapper.text()).toContain('PAT')
expect(wrapper.text()).not.toContain('Agent Identity')
await wrapper.setProps({ authMode: undefined })
expect(wrapper.text()).toContain('OAuth')
})
})
@@ -822,13 +822,13 @@ export default {
refreshTokenAuth: 'Manual RT Input',
refreshTokenDesc: 'Enter your existing OpenAI Refresh Token(s). Supports batch input (one per line). The system will automatically validate and create accounts.',
refreshTokenPlaceholder: 'Paste your OpenAI Refresh Token...\nSupports multiple, one per line',
codexSessionAuth: 'Codex JSON / AT Batch Input',
codexSessionDesc: 'Paste Codex JSON or an accessToken. Accounts use the step 1 settings.',
codexSessionInputLabel: 'Codex JSON or accessToken',
codexSessionPlaceholder: 'Multiple lines supported, one token or JSON per line',
codexSessionHint: 'sessionToken will not be saved as refresh_token. Without refresh_token, the account expires with the accessToken expiry; import is rejected if the expiry cannot be parsed and step 1 has no expiration.',
codexSessionAuth: 'Codex auth.json / AT Import',
codexSessionDesc: 'Paste a Codex auth.json (OAuth or Agent Identity) or an accessToken. Accounts use the step 1 settings.',
codexSessionInputLabel: 'Codex auth.json or accessToken',
codexSessionPlaceholder: 'Multiple lines supported, one token or auth.json object per line',
codexSessionHint: 'Agent Identity keeps no OAuth token and signs each upstream request dynamically. Session/access-token imports retain their existing expiration behavior.',
codexSessionImportAndCreate: 'Import & Create Account',
codexSessionEmpty: 'Please enter Codex JSON or accessToken',
codexSessionEmpty: 'Please enter a Codex auth.json or accessToken',
codexSessionImportFailed: 'Failed to import Codex account',
codexSessionImportSuccess: 'Import completed: created {created}, updated {updated}, skipped {skipped}',
codexSessionImportPartial: 'Partial success: created {created}, updated {updated}, skipped {skipped}, failed {failed}',
@@ -909,13 +909,13 @@ export default {
refreshTokenAuth: '手动输入 RT',
refreshTokenDesc: '输入您已有的 OpenAI Refresh Token,支持批量输入(每行一个),系统将自动验证并创建账号。',
refreshTokenPlaceholder: '粘贴您的 OpenAI Refresh Token...\n支持多个,每行一个',
codexSessionAuth: 'Codex JSON / AT 批量输入',
codexSessionDesc: '粘贴 Codex JSON 或 accessToken,按第一步配置创建账号。',
codexSessionInputLabel: 'Codex JSON 或 accessToken',
codexSessionPlaceholder: '支持多行,每行一个 token 或 JSON',
codexSessionHint: 'sessionToken 不会作为 refresh_token 保存;未包含 refresh_token 时会按 accessToken 过期时间设置账号过期,无法解析且第一步未设置过期时间时会拒绝导入。',
codexSessionAuth: 'Codex auth.json / AT 导入',
codexSessionDesc: '粘贴 Codex auth.json(OAuth 或 Agent Identity)或 accessToken,按第一步配置创建账号。',
codexSessionInputLabel: 'Codex auth.json 或 accessToken',
codexSessionPlaceholder: '支持多行,每行一个 token 或 auth.json 对象',
codexSessionHint: 'Agent Identity 不保存 OAuth token,并为每次上游请求动态签名;session/accessToken 导入继续沿用原有过期规则。',
codexSessionImportAndCreate: '导入并创建账号',
codexSessionEmpty: '请输入 Codex JSON 或 accessToken',
codexSessionEmpty: '请输入 Codex auth.json 或 accessToken',
codexSessionImportFailed: 'Codex 账号导入失败',
codexSessionImportSuccess: '导入完成:新增 {created},更新 {updated},跳过 {skipped}',
codexSessionImportPartial: '部分成功:新增 {created},更新 {updated},跳过 {skipped},失败 {failed}',
@@ -234,6 +234,7 @@
<div class="flex min-w-0 flex-col gap-1">
<div class="flex flex-wrap items-center gap-1">
<PlatformTypeBadge :platform="row.platform" :type="row.type"
:auth-mode="getOpenAIAuthMode(row)"
:plan-type="getAccountPlanType(row)"
:privacy-mode="row.extra?.privacy_mode || row.parent_privacy_mode"
:subscription-expires-at="row.credentials?.subscription_expires_at || row.parent_subscription_expires_at" />
@@ -1167,6 +1168,12 @@ function getAccountPlanType(row: any): string | undefined {
return row.credentials?.plan_type || row.parent_plan_type || undefined
}
function getOpenAIAuthMode(row: any): string | undefined {
if (!row || row.platform !== 'openai' || row.type !== 'oauth') return undefined
const authMode = row.credentials?.auth_mode
return typeof authMode === 'string' && authMode.trim() ? authMode : undefined
}
// Antigravity 订阅等级辅助函数
function getAntigravityTierFromRow(row: any): string | null {
if (row.platform !== 'antigravity') return null