mirror of
https://github.com/Wei-Shaw/sub2api.git
synced 2026-09-24 16:05:44 +08:00
Merge pull request #4269 from catoncat/agent/sub2api-agent-identity
feat(openai): support Codex Agent Identity authentication
This commit is contained in:
@@ -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",
|
||||
// 云服务凭据
|
||||
|
||||
@@ -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 ©
|
||||
}
|
||||
|
||||
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 是上游网络
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user