From 1dab126944374f49ffd2014e9087c7276a49e9eb Mon Sep 17 00:00:00 2001 From: cat Date: Tue, 14 Jul 2026 16:46:27 +0800 Subject: [PATCH 1/8] =?UTF-8?q?feat(openai):=20=E6=94=AF=E6=8C=81=20Agent?= =?UTF-8?q?=20Identity=20=E8=AE=A4=E8=AF=81?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- ...ccount_codex_agent_identity_import_test.go | 47 ++ .../handler/admin/account_codex_import.go | 118 ++++- .../handler/dto/credentials_redact_test.go | 4 + .../service/account_credentials_redact.go | 2 +- .../internal/service/account_test_service.go | 89 +++- .../internal/service/account_usage_service.go | 22 +- .../internal/service/openai_agent_identity.go | 431 ++++++++++++++++++ .../openai_agent_identity_compat_test.go | 247 ++++++++++ .../service/openai_agent_identity_test.go | 177 +++++++ .../service/openai_codex_models_service.go | 12 +- .../service/openai_gateway_count_tokens.go | 10 +- .../service/openai_gateway_forward.go | 30 +- .../service/openai_gateway_passthrough.go | 12 +- .../service/openai_gateway_service.go | 4 + .../service/openai_gateway_upstream_errors.go | 1 + backend/internal/service/openai_images.go | 10 +- .../internal/service/openai_quota_service.go | 76 ++- .../service/openai_quota_spark_window_test.go | 43 ++ .../service/openai_ws_forwarder_ingress.go | 15 +- .../service/openai_ws_forwarder_payload.go | 4 +- .../service/openai_ws_forwarder_v2.go | 18 +- backend/internal/service/openai_ws_pool.go | 65 ++- .../openai_ws_v2_passthrough_adapter.go | 26 +- 23 files changed, 1399 insertions(+), 64 deletions(-) create mode 100644 backend/internal/handler/admin/account_codex_agent_identity_import_test.go create mode 100644 backend/internal/service/openai_agent_identity.go create mode 100644 backend/internal/service/openai_agent_identity_compat_test.go create mode 100644 backend/internal/service/openai_agent_identity_test.go diff --git a/backend/internal/handler/admin/account_codex_agent_identity_import_test.go b/backend/internal/handler/admin/account_codex_agent_identity_import_test.go new file mode 100644 index 0000000000..db55182b6b --- /dev/null +++ b/backend/internal/handler/admin/account_codex_agent_identity_import_test.go @@ -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) +} diff --git a/backend/internal/handler/admin/account_codex_import.go b/backend/internal/handler/admin/account_codex_import.go index 01a5fbfa1c..8abd269f23 100644 --- a/backend/internal/handler/admin/account_codex_import.go +++ b/backend/internal/handler/admin/account_codex_import.go @@ -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 { @@ -492,6 +497,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"}, @@ -573,6 +613,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") } @@ -808,6 +851,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 { @@ -838,6 +884,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 { @@ -887,6 +941,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) { @@ -1039,6 +1097,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 != "" { diff --git a/backend/internal/handler/dto/credentials_redact_test.go b/backend/internal/handler/dto/credentials_redact_test.go index 431078fafd..9bf4f78039 100644 --- a/backend/internal/handler/dto/credentials_redact_test.go +++ b/backend/internal/handler/dto/credentials_redact_test.go @@ -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 { diff --git a/backend/internal/service/account_credentials_redact.go b/backend/internal/service/account_credentials_redact.go index 76c2d1de5b..8eb70513ba 100644 --- a/backend/internal/service/account_credentials_redact.go +++ b/backend/internal/service/account_credentials_redact.go @@ -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", // 云服务凭据 diff --git a/backend/internal/service/account_test_service.go b/backend/internal/service/account_test_service.go index b9dae84057..f83a99a7e2 100644 --- a/backend/internal/service/account_test_service.go +++ b/backend/internal/service/account_test_service.go @@ -15,6 +15,7 @@ import ( "net/http/httptest" "regexp" "strings" + "sync" "time" "github.com/Wei-Shaw/sub2api/internal/config" @@ -72,6 +73,7 @@ type AccountTestService struct { httpUpstream HTTPUpstream cfg *config.Config tlsFPProfileService *TLSFingerprintProfileService + agentIdentityTaskMu sync.Mutex } // NewAccountTestService creates a new AccountTestService @@ -544,9 +546,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") } @@ -597,7 +601,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, nil, &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 { @@ -804,16 +820,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" @@ -852,7 +878,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, nil, &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) @@ -863,7 +901,7 @@ func (s *AccountTestService) testOpenAICompactConnection(c *gin.Context, account if isOAuth { req.Host = "chatgpt.com" - setOpenAIChatGPTAccountHeaders(req.Header, account) + setOpenAIChatGPTAccountHeaders(req.Header, credentialAccount) } // 账号级请求头覆写：测试请求与真实转发保持一致的最终头 @@ -1677,8 +1715,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") } @@ -1710,17 +1759,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, nil, &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) diff --git a/backend/internal/service/account_usage_service.go b/backend/internal/service/account_usage_service.go index 281122d4f0..4dc09c4109 100644 --- a/backend/internal/service/account_usage_service.go +++ b/backend/internal/service/account_usage_service.go @@ -291,6 +291,7 @@ type AccountUsageService struct { cache *UsageCache identityCache IdentityCache tlsFPProfileService *TLSFingerprintProfileService + agentIdentityTaskMu sync.Mutex } // NewAccountUsageService 创建AccountUsageService实例 @@ -682,8 +683,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 @@ -701,7 +705,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, nil, &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") diff --git a/backend/internal/service/openai_agent_identity.go b/backend/internal/service/openai_agent_identity.go new file mode 100644 index 0000000000..c8818cd6f3 --- /dev/null +++ b/backend/internal/service/openai_agent_identity.go @@ -0,0 +1,431 @@ +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 + +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, pool *openAIWSConnPool, 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") + } + taskMu.Lock() + defer taskMu.Unlock() + 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 pool != nil { + pool.ClearAccount(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.openaiWSPool, &s.agentIdentityTaskMu, account, expectedTaskID) +} + +func isAgentIdentityTaskInvalidHTTPResponse(statusCode int, body []byte) bool { + if statusCode != http.StatusUnauthorized { + return false + } + lower := strings.ToLower(string(body)) + for _, marker := range []string{ + "invalid task", + "task_id", + "task id", + "task_not_found", + "task_expired", + "unknown task", + } { + if strings.Contains(lower, marker) { + return true + } + } + return false +} + +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.openaiWSPool, &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, pool *openAIWSConnPool, 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, pool, 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.openaiWSPool, &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 (s *OpenAIGatewayService) redactAgentIdentitySensitiveBody(ctx context.Context, account *Account, body []byte) []byte { + if !s.isAgentIdentityAccount(ctx, account) || len(body) == 0 { + return body + } + credAccount := account + if account != nil && account.IsShadow() { + if resolved, err := resolveCredentialAccount(ctx, s.accountRepo, account); err == nil && resolved != nil { + credAccount = resolved + } + } + 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]") + } + } + return []byte(redacted) +} diff --git a/backend/internal/service/openai_agent_identity_compat_test.go b/backend/internal/service/openai_agent_identity_compat_test.go new file mode 100644 index 0000000000..5c0c996c42 --- /dev/null +++ b/backend/internal/service/openai_agent_identity_compat_test.go @@ -0,0 +1,247 @@ +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 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, "cache-agent"), req.Header.Get("session_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+`"}`)) + require.NotContains(t, string(redacted), key.runtimeID) + require.NotContains(t, string(redacted), key.taskID) + require.NotContains(t, string(redacted), oauthValue) + require.Contains(t, string(redacted), "[redacted]") +} + +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))}, + }} + 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) +} + +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 +} + +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 +} diff --git a/backend/internal/service/openai_agent_identity_test.go b/backend/internal/service/openai_agent_identity_test.go new file mode 100644 index 0000000000..73ac0ddf08 --- /dev/null +++ b/backend/internal/service/openai_agent_identity_test.go @@ -0,0 +1,177 @@ +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) + require.True(t, ed25519.Verify(key.privateKey.Public().(ed25519.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(mustJSON(t, redacted)), privateKey) +} + +type agentIdentityCredentialsRepo struct { + AccountRepository + credentials map[string]any + mu sync.Mutex +} + +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 mustJSON(t *testing.T, value any) []byte { + t.Helper() + encoded, err := json.Marshal(value) + require.NoError(t, err) + return encoded +} diff --git a/backend/internal/service/openai_codex_models_service.go b/backend/internal/service/openai_codex_models_service.go index 8a919fa2b0..0eb2e1a071 100644 --- a/backend/internal/service/openai_codex_models_service.go +++ b/backend/internal/service/openai_codex_models_service.go @@ -42,7 +42,7 @@ func (s *OpenAIGatewayService) FetchCodexModelsManifest(ctx context.Context, acc return nil, infraerrors.Newf(http.StatusInternalServerError, "OPENAI_CODEX_MODELS_CREDENTIALS_FAILED", "resolve credential account: %v", err) } accessToken := credAccount.GetOpenAIAccessToken() - if accessToken == "" { + if accessToken == "" && !credAccount.IsOpenAIAgentIdentity() { return nil, infraerrors.New(http.StatusBadGateway, "OPENAI_CODEX_MODELS_TOKEN_MISSING", "account has no Codex backend access token") } @@ -58,7 +58,15 @@ func (s *OpenAIGatewayService) FetchCodexModelsManifest(ctx context.Context, acc if err != nil { return nil, infraerrors.Newf(http.StatusInternalServerError, "OPENAI_CODEX_MODELS_REQUEST_FAILED", "create codex models request: %v", err) } - req.Header.Set("Authorization", "Bearer "+accessToken) + authHeaders, err := s.buildOpenAIAuthenticationHeaders(ctx, credAccount, accessToken) + if err != nil { + return nil, infraerrors.Newf(http.StatusBadGateway, "OPENAI_CODEX_MODELS_AUTH_FAILED", "build Codex models authentication: %v", err) + } + for key, values := range authHeaders { + for _, value := range values { + req.Header.Add(key, value) + } + } req.Header.Set("Accept", "application/json") req.Header.Set("Originator", "codex_cli_rs") req.Header.Set("Version", clientVersion) diff --git a/backend/internal/service/openai_gateway_count_tokens.go b/backend/internal/service/openai_gateway_count_tokens.go index 7518a6073a..46b5a3843a 100644 --- a/backend/internal/service/openai_gateway_count_tokens.go +++ b/backend/internal/service/openai_gateway_count_tokens.go @@ -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") diff --git a/backend/internal/service/openai_gateway_forward.go b/backend/internal/service/openai_gateway_forward.go index 53cc06989e..e1c9aea15f 100644 --- a/backend/internal/service/openai_gateway_forward.go +++ b/backend/internal/service/openai_gateway_forward.go @@ -487,6 +487,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 { @@ -571,6 +572,7 @@ func (s *OpenAIGatewayService) Forward(ctx context.Context, c *gin.Context, acco startTime, attempt, wsLastFailureReason, + &agentTaskRecoveryTried, ) if wsErr == nil { break @@ -578,6 +580,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 != "" { @@ -684,6 +690,7 @@ func (s *OpenAIGatewayService) Forward(ctx context.Context, c *gin.Context, acco } httpInvalidEncryptedContentRetryTried := false + agentTaskRecoveryTried := false for { // Build upstream request upstreamCtx, releaseUpstreamCtx := detachUpstreamContext(ctx) @@ -719,6 +726,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 { @@ -869,8 +886,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 { diff --git a/backend/internal/service/openai_gateway_passthrough.go b/backend/internal/service/openai_gateway_passthrough.go index 238d359b9e..2ea914b537 100644 --- a/backend/internal/service/openai_gateway_passthrough.go +++ b/backend/internal/service/openai_gateway_passthrough.go @@ -338,7 +338,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 +442,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 +486,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 是上游网络安全策略拦截，不冷却账号， diff --git a/backend/internal/service/openai_gateway_service.go b/backend/internal/service/openai_gateway_service.go index 0d03d83dcf..a17a6f6b05 100644 --- a/backend/internal/service/openai_gateway_service.go +++ b/backend/internal/service/openai_gateway_service.go @@ -356,6 +356,7 @@ type OpenAIGatewayService struct { openaiWSStateStoreOnce sync.Once openaiSchedulerOnce sync.Once openaiWSPassthroughDialerOnce sync.Once + agentIdentityTaskMu sync.Mutex openaiWSPool *openAIWSConnPool openaiWSStateStore OpenAIWSStateStore openaiScheduler OpenAIAccountScheduler @@ -1052,6 +1053,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) diff --git a/backend/internal/service/openai_gateway_upstream_errors.go b/backend/internal/service/openai_gateway_upstream_errors.go index 2e661e7417..bc9bf91fef 100644 --- a/backend/internal/service/openai_gateway_upstream_errors.go +++ b/backend/internal/service/openai_gateway_upstream_errors.go @@ -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 diff --git a/backend/internal/service/openai_images.go b/backend/internal/service/openai_images.go index fc7d37a173..b411cb3104 100644 --- a/backend/internal/service/openai_images.go +++ b/backend/internal/service/openai_images.go @@ -752,7 +752,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 diff --git a/backend/internal/service/openai_quota_service.go b/backend/internal/service/openai_quota_service.go index 337f8c1e8f..b8b766fcff 100644 --- a/backend/internal/service/openai_quota_service.go +++ b/backend/internal/service/openai_quota_service.go @@ -10,6 +10,7 @@ import ( "log/slog" "net/http" "strings" + "sync" "time" infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors" @@ -118,6 +119,7 @@ type OpenAIQuotaService struct { proxyRepo ProxyRepository tokenProvider *OpenAITokenProvider privacyClientFactory PrivacyClientFactory + agentIdentityTaskMu sync.Mutex } // NewOpenAIQuotaService constructs a quota service. token provider is required — @@ -154,10 +156,14 @@ func (s *OpenAIQuotaService) QueryUsage(ctx context.Context, accountID int64) (* callCtx, cancel := context.WithTimeout(ctx, openaiQuotaUpstreamTimeout) defer cancel() + 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) + } var payload OpenAIQuotaUsage resp, err := client.R(). SetContext(callCtx). - SetHeaders(buildCodexCommonHeaders(accessToken, chatGPTAccountID, fedRAMP)). + SetHeaders(quotaHeaders). SetSuccessResult(&payload). Get(chatGPTUsageURL) if err != nil { @@ -178,9 +184,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) []OpenAIRateLimitResetCreditDetail { + 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) @@ -239,7 +250,10 @@ func (s *OpenAIQuotaService) ResetCredit(ctx context.Context, accountID int64) ( callCtx, cancel := context.WithTimeout(ctx, openaiQuotaUpstreamTimeout) defer cancel() - headers := buildCodexCommonHeaders(accessToken, chatGPTAccountID, fedRAMP) + 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" var payload OpenAIQuotaResetResult @@ -271,7 +285,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") } @@ -309,12 +323,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() @@ -337,6 +356,43 @@ func (s *OpenAIQuotaService) prepareUpstreamCall(ctx context.Context, accountID return accessToken, chatGPTAccountID, proxyURL, fedRAMP, nil } +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, nil, &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 +} + // 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 { diff --git a/backend/internal/service/openai_quota_spark_window_test.go b/backend/internal/service/openai_quota_spark_window_test.go index c56600d3f5..d2b18af05b 100644 --- a/backend/internal/service/openai_quota_spark_window_test.go +++ b/backend/internal/service/openai_quota_spark_window_test.go @@ -2,6 +2,10 @@ package service import ( "context" + "crypto/ed25519" + "crypto/rand" + "crypto/x509" + "encoding/base64" "encoding/json" "errors" "fmt" @@ -214,6 +218,45 @@ 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 TestParseOpenAIRateLimitResetCreditDetails_CompatibleContainers(t *testing.T) { tests := []struct { name string diff --git a/backend/internal/service/openai_ws_forwarder_ingress.go b/backend/internal/service/openai_ws_forwarder_ingress.go index 45e004b6d0..9a6453892a 100644 --- a/backend/internal/service/openai_ws_forwarder_ingress.go +++ b/backend/internal/service/openai_ws_forwarder_ingress.go @@ -539,6 +539,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() @@ -602,7 +605,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 @@ -611,6 +616,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) && dialErr != nil && dialErr.StatusCode == http.StatusUnauthorized && !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( diff --git a/backend/internal/service/openai_ws_forwarder_payload.go b/backend/internal/service/openai_ws_forwarder_payload.go index a4d47218e7..0aa4adc313 100644 --- a/backend/internal/service/openai_ws_forwarder_payload.go +++ b/backend/internal/service/openai_ws_forwarder_payload.go @@ -67,7 +67,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 { diff --git a/backend/internal/service/openai_ws_forwarder_v2.go b/backend/internal/service/openai_ws_forwarder_v2.go index f0f71648dd..65a6add7a5 100644 --- a/backend/internal/service/openai_ws_forwarder_v2.go +++ b/backend/internal/service/openai_ws_forwarder_v2.go @@ -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) && agentDialErr != nil && agentDialErr.StatusCode == http.StatusUnauthorized && 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", diff --git a/backend/internal/service/openai_ws_pool.go b/backend/internal/service/openai_ws_pool.go index 5950e02841..be81d611f9 100644 --- a/backend/internal/service/openai_ws_pool.go +++ b/backend/internal/service/openai_ws_pool.go @@ -60,9 +60,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: 强制本次获取新连接（避免复用导致连接内续链状态互相污染）。 @@ -517,6 +521,7 @@ type openAIWSAccountPool struct { conns map[string]*openAIWSConn pinnedConns map[string]int creating int + generation uint64 lastCleanupAt time.Time lastAcquire *openAIWSAcquireRequest prewarmActive bool @@ -1283,6 +1288,7 @@ func (p *openAIWSConnPool) ensureTargetIdleAsync(accountID int64) { } var req openAIWSAcquireRequest + generation := uint64(0) need := 0 ap, ok := p.getAccountPool(accountID) if !ok || ap == nil { @@ -1317,6 +1323,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) @@ -1324,7 +1331,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 { @@ -1367,7 +1374,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() @@ -1398,6 +1409,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.mu.Unlock() conn.close() @@ -1410,6 +1426,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 @@ -1485,7 +1530,15 @@ 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 { return nil, &openAIWSDialError{ StatusCode: status, diff --git a/backend/internal/service/openai_ws_v2_passthrough_adapter.go b/backend/internal/service/openai_ws_v2_passthrough_adapter.go index d822d181cd..01b66c376b 100644 --- a/backend/internal/service/openai_ws_v2_passthrough_adapter.go +++ b/backend/internal/service/openai_ws_v2_passthrough_adapter.go @@ -359,10 +359,28 @@ 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 + } + if s.isAgentIdentityAccount(ctx, account) && statusCode == http.StatusUnauthorized && !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, From 57cf2df09fe743db260b8dfa4a446498959dcedb Mon Sep 17 00:00:00 2001 From: cat Date: Tue, 14 Jul 2026 16:54:53 +0800 Subject: [PATCH 2/8] =?UTF-8?q?fix(openai):=20=E6=94=B6=E7=B4=A7=20Agent?= =?UTF-8?q?=20Identity=20=E6=81=A2=E5=A4=8D=E4=B8=8E=E8=84=B1=E6=95=8F?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../internal/service/account_test_service.go | 4 ++ .../internal/service/openai_agent_identity.go | 43 +++++++++++++-- .../openai_agent_identity_compat_test.go | 49 ++++++++++++++++- .../service/openai_agent_identity_test.go | 38 +++++++++++++ .../service/openai_gateway_passthrough.go | 55 +++++++++++++------ .../internal/service/openai_quota_service.go | 15 ++++- backend/internal/service/openai_ws_client.go | 30 +++++++++- .../service/openai_ws_forwarder_ingress.go | 2 +- .../service/openai_ws_forwarder_v2.go | 2 +- backend/internal/service/openai_ws_pool.go | 7 +++ .../openai_ws_v2_passthrough_adapter.go | 9 ++- 11 files changed, 224 insertions(+), 30 deletions(-) diff --git a/backend/internal/service/account_test_service.go b/backend/internal/service/account_test_service.go index f83a99a7e2..5a0bc40f40 100644 --- a/backend/internal/service/account_test_service.go +++ b/backend/internal/service/account_test_service.go @@ -655,6 +655,7 @@ 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 resp.StatusCode == http.StatusTooManyRequests { s.reconcileOpenAI429State(ctx, account, resp.Header, body) } @@ -924,6 +925,7 @@ 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 s.accountRepo != nil { updates := buildOpenAICompactProbeExtraUpdates(resp, body, nil, time.Now()) @@ -1800,6 +1802,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) @@ -1811,6 +1814,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 { diff --git a/backend/internal/service/openai_agent_identity.go b/backend/internal/service/openai_agent_identity.go index c8818cd6f3..8383a10e9c 100644 --- a/backend/internal/service/openai_agent_identity.go +++ b/backend/internal/service/openai_agent_identity.go @@ -29,6 +29,8 @@ const ( var openAIAgentIdentityAuthAPIBaseURL = agentIdentityAuthAPIBaseURL +var agentIdentityTaskLocks sync.Map // map[int64]*sync.Mutex + type agentIdentityKey struct { runtimeID string privateKey ed25519.PrivateKey @@ -251,8 +253,14 @@ func ensureAgentIdentityTaskForAccount(ctx context.Context, repo AccountReposito if taskMu == nil { return errors.New("agent identity task lock is unavailable") } - taskMu.Lock() - defer taskMu.Unlock() + sharedTaskMu := taskMu + if credAccount.ID > 0 { + candidate := &sync.Mutex{} + actual, _ := agentIdentityTaskLocks.LoadOrStore(credAccount.ID, candidate) + sharedTaskMu = actual.(*sync.Mutex) + } + sharedTaskMu.Lock() + defer sharedTaskMu.Unlock() currentTaskID = strings.TrimSpace(credAccount.GetCredential("task_id")) if currentTaskID != "" && (expectedTaskID == "" || currentTaskID != expectedTaskID) { return nil @@ -302,6 +310,10 @@ func isAgentIdentityTaskInvalidHTTPResponse(statusCode int, body []byte) bool { return false } +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") @@ -401,16 +413,19 @@ func (s *OpenAIGatewayService) isAgentIdentityAccount(ctx context.Context, accou // 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 (s *OpenAIGatewayService) redactAgentIdentitySensitiveBody(ctx context.Context, account *Account, body []byte) []byte { - if !s.isAgentIdentityAccount(ctx, account) || len(body) == 0 { +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, s.accountRepo, account); err == nil && resolved != nil { + 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", @@ -427,5 +442,23 @@ func (s *OpenAIGatewayService) redactAgentIdentitySensitiveBody(ctx context.Cont redacted = strings.ReplaceAll(redacted, value, "[redacted]") } } + for { + start := strings.Index(redacted, "AgentAssertion ") + if start < 0 { + break + } + end := start + len("AgentAssertion ") + for end < len(redacted) && !strings.ContainsRune(" \t\r\n\"',}", rune(redacted[end])) { + end++ + } + redacted = redacted[:start] + "AgentAssertion [redacted]" + redacted[end:] + } 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) +} diff --git a/backend/internal/service/openai_agent_identity_compat_test.go b/backend/internal/service/openai_agent_identity_compat_test.go index 5c0c996c42..c034af2bfa 100644 --- a/backend/internal/service/openai_agent_identity_compat_test.go +++ b/backend/internal/service/openai_agent_identity_compat_test.go @@ -104,13 +104,45 @@ func TestOpenAIAgentIdentityErrorRedactionDoesNotLeakCredentialValues(t *testing } svc := &OpenAIGatewayService{} oauthValue := account.GetCredential("access_token") - redacted := svc.redactAgentIdentitySensitiveBody(context.Background(), account, []byte(`{"message":"runtime-test task-test `+oauthValue+`"}`)) + 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 TestOpenAIWSConnPoolHeadersFactoryRunsAtDialAndStalePrewarmIsDiscarded(t *testing.T) { cfg := &config.Config{} cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 1 @@ -218,6 +250,21 @@ func TestOpenAIAgentIdentityTaskInvalidRetriesExactlyOnce(t *testing.T) { 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{"application/json"}}, Body: io.NopCloser(strings.NewReader(successBody))}, + } + 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 decodeAgentAssertionTask(t *testing.T, header string) string { diff --git a/backend/internal/service/openai_agent_identity_test.go b/backend/internal/service/openai_agent_identity_test.go index 73ac0ddf08..6e94f53abe 100644 --- a/backend/internal/service/openai_agent_identity_test.go +++ b/backend/internal/service/openai_agent_identity_test.go @@ -156,6 +156,44 @@ func TestEnsureAgentIdentityTaskPersistsAndRedactsCredentials(t *testing.T) { require.NotContains(t, string(mustJSON(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{} + 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) + for range 2 { + go func() { + <-start + errors <- ensureAgentIdentityTaskForAccount(context.Background(), repo, nil, &sync.Mutex{}, account, "") + }() + } + 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", account.GetCredential("task_id")) +} + type agentIdentityCredentialsRepo struct { AccountRepository credentials map[string]any diff --git a/backend/internal/service/openai_gateway_passthrough.go b/backend/internal/service/openai_gateway_passthrough.go index 2ea914b537..f280f37841 100644 --- a/backend/internal/service/openai_gateway_passthrough.go +++ b/backend/internal/service/openai_gateway_passthrough.go @@ -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) diff --git a/backend/internal/service/openai_quota_service.go b/backend/internal/service/openai_quota_service.go index b8b766fcff..cebab0a1d7 100644 --- a/backend/internal/service/openai_quota_service.go +++ b/backend/internal/service/openai_quota_service.go @@ -171,7 +171,7 @@ func (s *OpenAIQuotaService) QueryUsage(ctx context.Context, accountID int64) (* } if !resp.IsSuccessState() { status := resp.StatusCode - body := truncate(resp.String(), 240) + 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) } @@ -268,7 +268,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) } @@ -393,6 +393,17 @@ func (s *OpenAIQuotaService) buildCodexQuotaHeaders(ctx context.Context, account 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 { diff --git a/backend/internal/service/openai_ws_client.go b/backend/internal/service/openai_ws_client.go index 80b7553083..d30c4a1cb3 100644 --- a/backend/internal/service/openai_ws_client.go +++ b/backend/internal/service/openai_ws_client.go @@ -4,6 +4,7 @@ import ( "context" "errors" "fmt" + "io" "net/http" "net/url" "strings" @@ -61,6 +62,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 @@ -97,7 +120,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)。 diff --git a/backend/internal/service/openai_ws_forwarder_ingress.go b/backend/internal/service/openai_ws_forwarder_ingress.go index 9a6453892a..169919b8b9 100644 --- a/backend/internal/service/openai_ws_forwarder_ingress.go +++ b/backend/internal/service/openai_ws_forwarder_ingress.go @@ -617,7 +617,7 @@ func (s *OpenAIGatewayService) ProxyResponsesWebSocketFromClient( lease, acquireErr := pool.Acquire(acquireCtx, req) acquireCancel() var dialErr *openAIWSDialError - if acquireErr != nil && s.isAgentIdentityAccount(ctx, account) && errors.As(acquireErr, &dialErr) && dialErr != nil && dialErr.StatusCode == http.StatusUnauthorized && !agentTaskRecoveryTried { + 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) diff --git a/backend/internal/service/openai_ws_forwarder_v2.go b/backend/internal/service/openai_ws_forwarder_v2.go index 65a6add7a5..90151be93a 100644 --- a/backend/internal/service/openai_ws_forwarder_v2.go +++ b/backend/internal/service/openai_ws_forwarder_v2.go @@ -192,7 +192,7 @@ func (s *OpenAIGatewayService) forwardOpenAIWSV2( }) if err != nil { var agentDialErr *openAIWSDialError - if s.isAgentIdentityAccount(ctx, account) && errors.As(err, &agentDialErr) && agentDialErr != nil && agentDialErr.StatusCode == http.StatusUnauthorized && agentTaskRecoveryTried != nil && !*agentTaskRecoveryTried { + 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) diff --git a/backend/internal/service/openai_ws_pool.go b/backend/internal/service/openai_ws_pool.go index be81d611f9..3a02da6fe5 100644 --- a/backend/internal/service/openai_ws_pool.go +++ b/backend/internal/service/openai_ws_pool.go @@ -39,6 +39,7 @@ var ( type openAIWSDialError struct { StatusCode int ResponseHeaders http.Header + ResponseBody []byte Err error } @@ -1540,9 +1541,15 @@ func (p *openAIWSConnPool) dialConn(ctx context.Context, req openAIWSAcquireRequ } 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, } } diff --git a/backend/internal/service/openai_ws_v2_passthrough_adapter.go b/backend/internal/service/openai_ws_v2_passthrough_adapter.go index 01b66c376b..bb097e540b 100644 --- a/backend/internal/service/openai_ws_v2_passthrough_adapter.go +++ b/backend/internal/service/openai_ws_v2_passthrough_adapter.go @@ -374,7 +374,8 @@ func (s *OpenAIGatewayService) proxyResponsesWebSocketV2Passthrough( if err == nil { break } - if s.isAgentIdentityAccount(ctx, account) && statusCode == http.StatusUnauthorized && !agentTaskRecoveryTried { + var dialErr *openAIWSDialError + if s.isAgentIdentityAccount(ctx, account) && errors.As(err, &dialErr) && 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) @@ -696,9 +697,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, } } From 3fe7b4da78715389a60c07bcf1ba02da044f50cc Mon Sep 17 00:00:00 2001 From: cat Date: Tue, 14 Jul 2026 17:01:53 +0800 Subject: [PATCH 3/8] =?UTF-8?q?fix(openai):=20=E9=97=AD=E5=90=88=E7=9B=B4?= =?UTF-8?q?=E9=80=9A=20WS=20=E6=81=A2=E5=A4=8D=E4=B8=8E=E9=94=99=E8=AF=AF?= =?UTF-8?q?=E8=84=B1=E6=95=8F?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../internal/service/openai_agent_identity.go | 22 +++++++++++++++++++ .../service/openai_agent_identity_test.go | 20 +++++++++++++---- .../service/openai_gateway_upstream_errors.go | 1 + backend/internal/service/openai_images.go | 1 + .../openai_ws_v2_passthrough_adapter.go | 9 ++++++-- 5 files changed, 47 insertions(+), 6 deletions(-) diff --git a/backend/internal/service/openai_agent_identity.go b/backend/internal/service/openai_agent_identity.go index 8383a10e9c..a5ea9b7612 100644 --- a/backend/internal/service/openai_agent_identity.go +++ b/backend/internal/service/openai_agent_identity.go @@ -261,6 +261,25 @@ func ensureAgentIdentityTaskForAccount(ctx context.Context, repo AccountReposito } 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 @@ -277,6 +296,9 @@ func ensureAgentIdentityTaskForAccount(ctx context.Context, repo AccountReposito if err := persistAccountCredentials(ctx, repo, credAccount, credentials); err != nil { return err } + if !account.IsShadow() && account != credAccount { + account.Credentials = shallowCopyMap(credAccount.Credentials) + } if pool != nil { pool.ClearAccount(credAccount.ID) } diff --git a/backend/internal/service/openai_agent_identity_test.go b/backend/internal/service/openai_agent_identity_test.go index 6e94f53abe..12f5a7766a 100644 --- a/backend/internal/service/openai_agent_identity_test.go +++ b/backend/internal/service/openai_agent_identity_test.go @@ -163,7 +163,7 @@ func TestEnsureAgentIdentityTaskSharesLockAcrossServicesForSameAccount(t *testin "agent_runtime_id": key.runtimeID, "agent_private_key": privateKey, }} - repo := &agentIdentityCredentialsRepo{} + repo := &agentIdentityCredentialsRepo{account: account} registerCalls := 0 var registerMu sync.Mutex server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { @@ -179,10 +179,11 @@ func TestEnsureAgentIdentityTaskSharesLockAcrossServicesForSameAccount(t *testin start := make(chan struct{}) errors := make(chan error, 2) - for range 2 { + requests := []*Account{cloneAgentIdentityTestAccount(account), cloneAgentIdentityTestAccount(account)} + for _, request := range requests { go func() { <-start - errors <- ensureAgentIdentityTaskForAccount(context.Background(), repo, nil, &sync.Mutex{}, account, "") + errors <- ensureAgentIdentityTaskForAccount(context.Background(), repo, nil, &sync.Mutex{}, request, "") }() } close(start) @@ -191,15 +192,26 @@ func TestEnsureAgentIdentityTaskSharesLockAcrossServicesForSameAccount(t *testin registerMu.Lock() defer registerMu.Unlock() require.Equal(t, 1, registerCalls) - require.Equal(t, "task-shared", account.GetCredential("task_id")) + 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() diff --git a/backend/internal/service/openai_gateway_upstream_errors.go b/backend/internal/service/openai_gateway_upstream_errors.go index bc9bf91fef..d8823fe4ed 100644 --- a/backend/internal/service/openai_gateway_upstream_errors.go +++ b/backend/internal/service/openai_gateway_upstream_errors.go @@ -471,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 是上游网络 diff --git a/backend/internal/service/openai_images.go b/backend/internal/service/openai_images.go index b411cb3104..081b5e741b 100644 --- a/backend/internal/service/openai_images.go +++ b/backend/internal/service/openai_images.go @@ -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) diff --git a/backend/internal/service/openai_ws_v2_passthrough_adapter.go b/backend/internal/service/openai_ws_v2_passthrough_adapter.go index bb097e540b..75df7191ac 100644 --- a/backend/internal/service/openai_ws_v2_passthrough_adapter.go +++ b/backend/internal/service/openai_ws_v2_passthrough_adapter.go @@ -374,8 +374,13 @@ func (s *OpenAIGatewayService) proxyResponsesWebSocketV2Passthrough( if err == nil { break } - var dialErr *openAIWSDialError - if s.isAgentIdentityAccount(ctx, account) && errors.As(err, &dialErr) && isAgentIdentityTaskInvalidWSDialError(dialErr) && !agentTaskRecoveryTried { + 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) From 10aa88aab95f19a48b66142ae2803408fd37a2be Mon Sep 17 00:00:00 2001 From: cat Date: Tue, 14 Jul 2026 17:19:15 +0800 Subject: [PATCH 4/8] =?UTF-8?q?fix(openai):=20=E8=A1=A5=E9=BD=90=E8=BA=AB?= =?UTF-8?q?=E4=BB=BD=E5=A4=B1=E6=95=88=E6=81=A2=E5=A4=8D=E8=B7=AF=E5=BE=84?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../internal/service/account_test_service.go | 31 +++- .../internal/service/openai_agent_identity.go | 36 ++++- .../openai_agent_identity_compat_test.go | 147 +++++++++++++++++- .../openai_gateway_chat_completions.go | 7 + .../service/openai_gateway_messages.go | 7 + .../service/openai_images_responses.go | 8 + .../internal/service/openai_quota_service.go | 132 +++++++++++----- .../service/openai_quota_spark_window_test.go | 58 +++++++ 8 files changed, 378 insertions(+), 48 deletions(-) diff --git a/backend/internal/service/account_test_service.go b/backend/internal/service/account_test_service.go index 5a0bc40f40..68a00a5191 100644 --- a/backend/internal/service/account_test_service.go +++ b/backend/internal/service/account_test_service.go @@ -590,8 +590,11 @@ func (s *AccountTestService) testOpenAIAccountConnection(c *gin.Context, account payload := createOpenAITestPayload(testModelID, 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 { @@ -656,6 +659,14 @@ 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, nil, &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) } @@ -718,7 +729,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 { @@ -869,7 +882,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 { @@ -926,6 +941,14 @@ func (s *AccountTestService) testOpenAICompactConnection(c *gin.Context, account 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, nil, &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()) diff --git a/backend/internal/service/openai_agent_identity.go b/backend/internal/service/openai_agent_identity.go index a5ea9b7612..0492a2b82e 100644 --- a/backend/internal/service/openai_agent_identity.go +++ b/backend/internal/service/openai_agent_identity.go @@ -317,13 +317,26 @@ func isAgentIdentityTaskInvalidHTTPResponse(statusCode int, body []byte) bool { return false } lower := strings.ToLower(string(body)) + compact := strings.NewReplacer(" ", "", "\t", "", "\r", "", "\n", "").Replace(lower) for _, marker := range []string{ - "invalid task", - "task_id", - "task id", - "task_not_found", - "task_expired", - "unknown task", + `"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 @@ -332,6 +345,17 @@ func isAgentIdentityTaskInvalidHTTPResponse(statusCode int, body []byte) bool { 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) } diff --git a/backend/internal/service/openai_agent_identity_compat_test.go b/backend/internal/service/openai_agent_identity_compat_test.go index c034af2bfa..bbfa82481d 100644 --- a/backend/internal/service/openai_agent_identity_compat_test.go +++ b/backend/internal/service/openai_agent_identity_compat_test.go @@ -55,6 +55,52 @@ func TestAccountTestServiceOpenAICompactAgentIdentityUsesFreshAssertion(t *testi 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"}`))}, + }} + svc := &AccountTestService{accountRepo: repo, httpUpstream: upstream} + 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) +} + func TestOpenAIAgentIdentityPassthroughKeepsSessionAndPromptCacheHeaders(t *testing.T) { gin.SetMode(gin.TestMode) key, privateKey := newTestAgentIdentityKey(t) @@ -225,6 +271,7 @@ func TestOpenAIAgentIdentityTaskInvalidRetriesExactlyOnce(t *testing.T) { {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) @@ -256,7 +303,7 @@ func TestOpenAIAgentIdentityTaskInvalidRetriesExactlyOnce(t *testing.T) { 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{"application/json"}}, Body: io.NopCloser(strings.NewReader(successBody))}, + {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) @@ -267,6 +314,80 @@ func TestOpenAIAgentIdentityTaskInvalidRetriesExactlyOnce(t *testing.T) { 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 ") @@ -284,6 +405,30 @@ type agentIdentityForwardRepo struct { account *Account } +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 } diff --git a/backend/internal/service/openai_gateway_chat_completions.go b/backend/internal/service/openai_gateway_chat_completions.go index 15c4e47635..5d80292c47 100644 --- a/backend/internal/service/openai_gateway_chat_completions.go +++ b/backend/internal/service/openai_gateway_chat_completions.go @@ -267,6 +267,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) { diff --git a/backend/internal/service/openai_gateway_messages.go b/backend/internal/service/openai_gateway_messages.go index 219b5e4be4..4d8b14cc28 100644 --- a/backend/internal/service/openai_gateway_messages.go +++ b/backend/internal/service/openai_gateway_messages.go @@ -316,6 +316,13 @@ 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.ID, xai.ParseQuotaHeaders(resp.Header, resp.StatusCode)) s.handleGrokAccountUpstreamError(ctx, account, resp.StatusCode, resp.Header, respBody) diff --git a/backend/internal/service/openai_images_responses.go b/backend/internal/service/openai_images_responses.go index 226cbf8e88..04347d6f90 100644 --- a/backend/internal/service/openai_images_responses.go +++ b/backend/internal/service/openai_images_responses.go @@ -1560,6 +1560,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) diff --git a/backend/internal/service/openai_quota_service.go b/backend/internal/service/openai_quota_service.go index cebab0a1d7..c736768b72 100644 --- a/backend/internal/service/openai_quota_service.go +++ b/backend/internal/service/openai_quota_service.go @@ -155,25 +155,36 @@ func (s *OpenAIQuotaService) QueryUsage(ctx context.Context, accountID int64) (* callCtx, cancel := context.WithTimeout(ctx, openaiQuotaUpstreamTimeout) defer cancel() + agentIdentity := s.isAgentIdentityAccount(ctx, accountID) - 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) - } var payload OpenAIQuotaUsage - 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() { - 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) + 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() @@ -249,28 +260,38 @@ func (s *OpenAIQuotaService) ResetCredit(ctx context.Context, accountID int64) ( callCtx, cancel := context.WithTimeout(ctx, openaiQuotaUpstreamTimeout) defer cancel() - - 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" + agentIdentity := s.isAgentIdentityAccount(ctx, accountID) var payload OpenAIQuotaResetResult - resp, err := client.R(). - SetContext(callCtx). - SetHeaders(headers). - SetBody(map[string]string{"redeem_request_id": redeemRequestID}). - SetSuccessResult(&payload). - Post(chatGPTRateLimitResetURL) - if err != nil { - return nil, infraerrors.Newf(http.StatusBadGateway, "OPENAI_QUOTA_RESET_REQUEST_FAILED", "upstream request failed: %v", err) - } - if !resp.IsSuccessState() { - status := resp.StatusCode - 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) + for recovered := false; ; { + 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). + SetBody(map[string]string{"redeem_request_id": redeemRequestID}). + SetSuccessResult(&payload). + Post(chatGPTRateLimitResetURL) + if err != nil { + return nil, infraerrors.Newf(http.StatusBadGateway, "OPENAI_QUOTA_RESET_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(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) + } + break } slog.Info("openai_quota_reset_success", @@ -356,6 +377,43 @@ 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, nil, &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 { diff --git a/backend/internal/service/openai_quota_spark_window_test.go b/backend/internal/service/openai_quota_spark_window_test.go index d2b18af05b..99bfe15ca1 100644 --- a/backend/internal/service/openai_quota_spark_window_test.go +++ b/backend/internal/service/openai_quota_spark_window_test.go @@ -38,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 @@ -257,6 +266,55 @@ func TestQueryUsageAgentIdentityUsesAssertionWithoutOAuthToken(t *testing.T) { 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 + } + 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 }) + + svc := NewOpenAIQuotaService(repo, nil, nil, newQuotaRedirectingFactory(srv)) + 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")) +} + func TestParseOpenAIRateLimitResetCreditDetails_CompatibleContainers(t *testing.T) { tests := []struct { name string From bd1399f81e33292df5720b8c3bcd073efb7ed161 Mon Sep 17 00:00:00 2001 From: cat Date: Tue, 14 Jul 2026 17:41:17 +0800 Subject: [PATCH 5/8] =?UTF-8?q?fix(openai):=20=E5=85=81=E8=AE=B8=20Agent?= =?UTF-8?q?=20Identity=20=E6=97=A0=E4=BB=A4=E7=89=8C=E6=8B=A8=E5=8F=B7=20W?= =?UTF-8?q?S?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../openai_agent_identity_compat_test.go | 26 +++++++++++++++++++ .../service/openai_ws_forwarder_ingress.go | 4 +-- .../service/openai_ws_forwarder_payload.go | 10 +++++++ .../openai_ws_v2_passthrough_adapter.go | 4 +-- 4 files changed, 40 insertions(+), 4 deletions(-) diff --git a/backend/internal/service/openai_agent_identity_compat_test.go b/backend/internal/service/openai_agent_identity_compat_test.go index bbfa82481d..ff58276310 100644 --- a/backend/internal/service/openai_agent_identity_compat_test.go +++ b/backend/internal/service/openai_agent_identity_compat_test.go @@ -189,6 +189,32 @@ func TestOpenAIWSAgentIdentityRecoveryRequiresTaskInvalidBody(t *testing.T) { })) } +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 diff --git a/backend/internal/service/openai_ws_forwarder_ingress.go b/backend/internal/service/openai_ws_forwarder_ingress.go index 169919b8b9..f7a54fb2b5 100644 --- a/backend/internal/service/openai_ws_forwarder_ingress.go +++ b/backend/internal/service/openai_ws_forwarder_ingress.go @@ -39,8 +39,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 diff --git a/backend/internal/service/openai_ws_forwarder_payload.go b/backend/internal/service/openai_ws_forwarder_payload.go index 0aa4adc313..67fdddc9f7 100644 --- a/backend/internal/service/openai_ws_forwarder_payload.go +++ b/backend/internal/service/openai_ws_forwarder_payload.go @@ -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") diff --git a/backend/internal/service/openai_ws_v2_passthrough_adapter.go b/backend/internal/service/openai_ws_v2_passthrough_adapter.go index 75df7191ac..72fbbd0f35 100644 --- a/backend/internal/service/openai_ws_v2_passthrough_adapter.go +++ b/backend/internal/service/openai_ws_v2_passthrough_adapter.go @@ -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()) From f479a5d10cb4d61379a1174b809948b5e911065d Mon Sep 17 00:00:00 2001 From: cat Date: Tue, 14 Jul 2026 18:09:59 +0800 Subject: [PATCH 6/8] =?UTF-8?q?fix(openai):=20=E6=94=B6=E7=B4=A7=20Agent?= =?UTF-8?q?=20Identity=20task=20=E7=94=9F=E5=91=BD=E5=91=A8=E6=9C=9F?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- backend/cmd/server/wire_gen.go | 22 +++---- .../internal/service/account_test_service.go | 11 ++-- .../internal/service/account_usage_service.go | 3 +- .../internal/service/openai_agent_identity.go | 20 +++--- .../openai_agent_identity_compat_test.go | 12 +++- .../service/openai_gateway_service.go | 6 ++ .../internal/service/openai_quota_service.go | 21 +++---- .../service/openai_quota_spark_window_test.go | 20 ++++++ backend/internal/service/wire.go | 63 ++++++++++++++++++- 9 files changed, 137 insertions(+), 41 deletions(-) diff --git a/backend/cmd/server/wire_gen.go b/backend/cmd/server/wire_gen.go index e148c5c363..088046de96 100644 --- a/backend/cmd/server/wire_gen.go +++ b/backend/cmd/server/wire_gen.go @@ -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) @@ -190,10 +182,10 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) { claudeUsageFetcher := repository.NewClaudeUsageFetcher(httpUpstream) antigravityQuotaFetcher := service.NewAntigravityQuotaFetcher(proxyRepository) grokQuotaFetcher := service.NewGrokQuotaFetcher() - 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, openAIQuotaService, usageCache, identityCache, tlsFingerprintProfileService) - accountTestService := service.NewAccountTestService(accountRepository, geminiTokenProvider, claudeTokenProvider, grokTokenProvider, antigravityGatewayService, httpUpstream, configConfig, tlsFingerprintProfileService) + accountUsageService := service.ProvideAccountUsageService(accountRepository, usageLogRepository, claudeUsageFetcher, geminiQuotaService, antigravityQuotaFetcher, grokQuotaFetcher, 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.NewAccountHandler(adminService, oAuthService, openAIOAuthService, geminiOAuthService, antigravityOAuthService, rateLimitService, accountUsageService, accountTestService, concurrencyService, crsSyncService, sessionLimitCache, rpmCache, compositeTokenCacheInvalidator) adminAnnouncementHandler := admin.NewAnnouncementHandler(announcementService) @@ -267,6 +259,13 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) { handlerPaymentHandler := handler.NewPaymentHandler(paymentService, paymentConfigService, channelService) 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) diff --git a/backend/internal/service/account_test_service.go b/backend/internal/service/account_test_service.go index 68a00a5191..bb57a93cd7 100644 --- a/backend/internal/service/account_test_service.go +++ b/backend/internal/service/account_test_service.go @@ -74,6 +74,7 @@ type AccountTestService struct { cfg *config.Config tlsFPProfileService *TLSFingerprintProfileService agentIdentityTaskMu sync.Mutex + agentIdentityWS agentIdentityWSConnectionInvalidator } // NewAccountTestService creates a new AccountTestService @@ -605,7 +606,7 @@ func (s *AccountTestService) testOpenAIAccountConnection(c *gin.Context, account // Set common headers req.Header.Set("Content-Type", "application/json") if credentialAccount.IsOpenAIAgentIdentity() { - authHeaders, authErr := buildAgentIdentityAuthenticationHeaders(ctx, s.accountRepo, nil, &s.agentIdentityTaskMu, credentialAccount) + authHeaders, authErr := buildAgentIdentityAuthenticationHeaders(ctx, s.accountRepo, s.agentIdentityWS, &s.agentIdentityTaskMu, credentialAccount) if authErr != nil { return s.sendErrorAndEnd(c, "Failed to build Agent Identity authentication") } @@ -661,7 +662,7 @@ func (s *AccountTestService) testOpenAIAccountConnection(c *gin.Context, account 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, nil, &s.agentIdentityTaskMu, credentialAccount, expectedTaskID); err != nil { + 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)) @@ -895,7 +896,7 @@ func (s *AccountTestService) testOpenAICompactConnection(c *gin.Context, account req.Header.Set("Content-Type", "application/json") req.Header.Set("Accept", "application/json") if credentialAccount.IsOpenAIAgentIdentity() { - authHeaders, authErr := buildAgentIdentityAuthenticationHeaders(ctx, s.accountRepo, nil, &s.agentIdentityTaskMu, credentialAccount) + authHeaders, authErr := buildAgentIdentityAuthenticationHeaders(ctx, s.accountRepo, s.agentIdentityWS, &s.agentIdentityTaskMu, credentialAccount) if authErr != nil { return s.sendErrorAndEnd(c, "Failed to build Agent Identity authentication") } @@ -943,7 +944,7 @@ func (s *AccountTestService) testOpenAICompactConnection(c *gin.Context, account 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, nil, &s.agentIdentityTaskMu, credentialAccount, expectedTaskID); err != nil { + 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)) @@ -1785,7 +1786,7 @@ func (s *AccountTestService) testOpenAIImageOAuth(c *gin.Context, ctx context.Co req = req.WithContext(WithHTTPUpstreamProfile(req.Context(), HTTPUpstreamProfileOpenAI)) req.Host = "chatgpt.com" if credentialAccount.IsOpenAIAgentIdentity() { - authHeaders, authErr := buildAgentIdentityAuthenticationHeaders(ctx, s.accountRepo, nil, &s.agentIdentityTaskMu, credentialAccount) + authHeaders, authErr := buildAgentIdentityAuthenticationHeaders(ctx, s.accountRepo, s.agentIdentityWS, &s.agentIdentityTaskMu, credentialAccount) if authErr != nil { return s.sendErrorAndEnd(c, "Failed to build Agent Identity authentication") } diff --git a/backend/internal/service/account_usage_service.go b/backend/internal/service/account_usage_service.go index 4dc09c4109..622323c617 100644 --- a/backend/internal/service/account_usage_service.go +++ b/backend/internal/service/account_usage_service.go @@ -292,6 +292,7 @@ type AccountUsageService struct { identityCache IdentityCache tlsFPProfileService *TLSFingerprintProfileService agentIdentityTaskMu sync.Mutex + agentIdentityWS agentIdentityWSConnectionInvalidator } // NewAccountUsageService 创建AccountUsageService实例 @@ -706,7 +707,7 @@ func (s *AccountUsageService) probeOpenAICodexSnapshot(ctx context.Context, acco req.Host = "chatgpt.com" req.Header.Set("Content-Type", "application/json") if account.IsOpenAIAgentIdentity() { - authHeaders, authErr := buildAgentIdentityAuthenticationHeaders(ctx, s.accountRepo, nil, &s.agentIdentityTaskMu, account) + authHeaders, authErr := buildAgentIdentityAuthenticationHeaders(ctx, s.accountRepo, s.agentIdentityWS, &s.agentIdentityTaskMu, account) if authErr != nil { return nil, fmt.Errorf("build Agent Identity authentication: %w", authErr) } diff --git a/backend/internal/service/openai_agent_identity.go b/backend/internal/service/openai_agent_identity.go index 0492a2b82e..4bf447ad64 100644 --- a/backend/internal/service/openai_agent_identity.go +++ b/backend/internal/service/openai_agent_identity.go @@ -31,6 +31,10 @@ 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 @@ -231,7 +235,7 @@ func registerAgentIdentityTask(ctx context.Context, account *Account) (string, e return decryptAgentTaskID(key, encrypted) } -func ensureAgentIdentityTaskForAccount(ctx context.Context, repo AccountRepository, pool *openAIWSConnPool, taskMu *sync.Mutex, account *Account, expectedTaskID string) error { +func ensureAgentIdentityTaskForAccount(ctx context.Context, repo AccountRepository, wsInvalidator agentIdentityWSConnectionInvalidator, taskMu *sync.Mutex, account *Account, expectedTaskID string) error { if account == nil || !account.IsOpenAIAgentIdentity() { return nil } @@ -299,8 +303,8 @@ func ensureAgentIdentityTaskForAccount(ctx context.Context, repo AccountReposito if !account.IsShadow() && account != credAccount { account.Credentials = shallowCopyMap(credAccount.Credentials) } - if pool != nil { - pool.ClearAccount(credAccount.ID) + if wsInvalidator != nil { + wsInvalidator.InvalidateAgentIdentityWSConnections(credAccount.ID) } return nil } @@ -309,7 +313,7 @@ func (s *OpenAIGatewayService) ensureAgentIdentityTask(ctx context.Context, acco if s == nil { return errors.New("openai gateway service is nil") } - return ensureAgentIdentityTaskForAccount(ctx, s.accountRepo, s.openaiWSPool, &s.agentIdentityTaskMu, account, expectedTaskID) + return ensureAgentIdentityTaskForAccount(ctx, s.accountRepo, s, &s.agentIdentityTaskMu, account, expectedTaskID) } func isAgentIdentityTaskInvalidHTTPResponse(statusCode int, body []byte) bool { @@ -374,7 +378,7 @@ func (s *OpenAIGatewayService) buildOpenAIAuthenticationHeaders(ctx context.Cont } headers := make(http.Header) if credAccount != nil && credAccount.IsOpenAIAgentIdentity() { - agentHeaders, err := buildAgentIdentityAuthenticationHeaders(ctx, s.accountRepo, s.openaiWSPool, &s.agentIdentityTaskMu, credAccount) + agentHeaders, err := buildAgentIdentityAuthenticationHeaders(ctx, s.accountRepo, s, &s.agentIdentityTaskMu, credAccount) if err != nil { return nil, err } @@ -384,11 +388,11 @@ func (s *OpenAIGatewayService) buildOpenAIAuthenticationHeaders(ctx context.Cont return headers, nil } -func buildAgentIdentityAuthenticationHeaders(ctx context.Context, repo AccountRepository, pool *openAIWSConnPool, taskMu *sync.Mutex, account *Account) (http.Header, error) { +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, pool, taskMu, account, ""); err != nil { + if err := ensureAgentIdentityTaskForAccount(ctx, repo, wsInvalidator, taskMu, account, ""); err != nil { return nil, err } key, err := agentIdentityKeyFromAccount(account) @@ -423,7 +427,7 @@ func (s *OpenAIGatewayService) refreshOpenAIAgentIdentityHeaders(ctx context.Con if refreshed == nil { refreshed = make(http.Header) } - authHeaders, err := buildAgentIdentityAuthenticationHeaders(ctx, s.accountRepo, s.openaiWSPool, &s.agentIdentityTaskMu, credAccount) + authHeaders, err := buildAgentIdentityAuthenticationHeaders(ctx, s.accountRepo, s, &s.agentIdentityTaskMu, credAccount) if err != nil { return nil, err } diff --git a/backend/internal/service/openai_agent_identity_compat_test.go b/backend/internal/service/openai_agent_identity_compat_test.go index ff58276310..eb90dfb56e 100644 --- a/backend/internal/service/openai_agent_identity_compat_test.go +++ b/backend/internal/service/openai_agent_identity_compat_test.go @@ -89,7 +89,8 @@ func TestAccountTestServiceOpenAICompactAgentIdentityRecoversInvalidTaskOnce(t * {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"}`))}, }} - svc := &AccountTestService{accountRepo: repo, httpUpstream: upstream} + 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)) @@ -99,6 +100,7 @@ func TestAccountTestServiceOpenAICompactAgentIdentityRecoversInvalidTaskOnce(t * 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) { @@ -431,6 +433,14 @@ type agentIdentityForwardRepo struct { 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 diff --git a/backend/internal/service/openai_gateway_service.go b/backend/internal/service/openai_gateway_service.go index a17a6f6b05..24f4786e09 100644 --- a/backend/internal/service/openai_gateway_service.go +++ b/backend/internal/service/openai_gateway_service.go @@ -554,6 +554,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 diff --git a/backend/internal/service/openai_quota_service.go b/backend/internal/service/openai_quota_service.go index c736768b72..735a580570 100644 --- a/backend/internal/service/openai_quota_service.go +++ b/backend/internal/service/openai_quota_service.go @@ -24,6 +24,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" @@ -120,6 +122,7 @@ type OpenAIQuotaService struct { tokenProvider *OpenAITokenProvider privacyClientFactory PrivacyClientFactory agentIdentityTaskMu sync.Mutex + agentIdentityWS agentIdentityWSConnectionInvalidator } // NewOpenAIQuotaService constructs a quota service. token provider is required — @@ -241,6 +244,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) @@ -260,10 +266,8 @@ func (s *OpenAIQuotaService) ResetCredit(ctx context.Context, accountID int64) ( callCtx, cancel := context.WithTimeout(ctx, openaiQuotaUpstreamTimeout) defer cancel() - agentIdentity := s.isAgentIdentityAccount(ctx, accountID) - var payload OpenAIQuotaResetResult - for recovered := false; ; { + for { 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) @@ -279,13 +283,6 @@ func (s *OpenAIQuotaService) ResetCredit(ctx context.Context, accountID int64) ( return nil, infraerrors.Newf(http.StatusBadGateway, "OPENAI_QUOTA_RESET_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(callCtx, accountID, resp.String()), 240) slog.Warn("openai_quota_reset_failed", "account_id", accountID, "status", status, "body", body) @@ -394,7 +391,7 @@ func (s *OpenAIQuotaService) recoverAgentIdentityTask(ctx context.Context, accou if !account.IsOpenAIAgentIdentity() { return nil } - return ensureAgentIdentityTaskForAccount(ctx, s.accountRepo, nil, &s.agentIdentityTaskMu, account, account.GetCredential("task_id")) + return ensureAgentIdentityTaskForAccount(ctx, s.accountRepo, s.agentIdentityWS, &s.agentIdentityTaskMu, account, account.GetCredential("task_id")) } func (s *OpenAIQuotaService) isAgentIdentityAccount(ctx context.Context, accountID int64) bool { @@ -436,7 +433,7 @@ func (s *OpenAIQuotaService) buildCodexQuotaHeaders(ctx context.Context, account if !account.IsOpenAIAgentIdentity() { return headers, nil } - if err := ensureAgentIdentityTaskForAccount(ctx, s.accountRepo, nil, &s.agentIdentityTaskMu, account, ""); err != nil { + if err := ensureAgentIdentityTaskForAccount(ctx, s.accountRepo, s.agentIdentityWS, &s.agentIdentityTaskMu, account, ""); err != nil { return nil, err } key, err := agentIdentityKeyFromAccount(account) diff --git a/backend/internal/service/openai_quota_spark_window_test.go b/backend/internal/service/openai_quota_spark_window_test.go index 99bfe15ca1..4e5fbc4a84 100644 --- a/backend/internal/service/openai_quota_spark_window_test.go +++ b/backend/internal/service/openai_quota_spark_window_test.go @@ -176,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 时: @@ -306,13 +323,16 @@ func TestQueryUsageAgentIdentityRecoversInvalidTaskOnce(t *testing.T) { 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) { diff --git a/backend/internal/service/wire.go b/backend/internal/service/wire.go index 7258ff05a3..8c20af661b 100644 --- a/backend/internal/service/wire.go +++ b/backend/internal/service/wire.go @@ -131,8 +131,65 @@ 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, + openAIQuotaService *OpenAIQuotaService, + cache *UsageCache, + identityCache IdentityCache, + tlsFPProfileService *TLSFingerprintProfileService, + openAIGatewayService *OpenAIGatewayService, +) *AccountUsageService { + service := NewAccountUsageService( + accountRepo, + usageLogRepo, + usageFetcher, + geminiQuotaService, + antigravityQuotaFetcher, + grokQuotaFetcher, + 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( @@ -601,8 +658,8 @@ var ProviderSet = wire.NewSet( ProvideClaudeTokenProvider, NewAntigravityGatewayService, ProvideRateLimitService, - NewAccountUsageService, - NewAccountTestService, + ProvideAccountUsageService, + ProvideAccountTestService, ProvideSettingService, NewDataManagementService, ProvideBackupService, From ec7c1b6f726584f777493c9e4ce6210c3feff2ad Mon Sep 17 00:00:00 2001 From: cat Date: Tue, 14 Jul 2026 19:24:56 +0800 Subject: [PATCH 7/8] =?UTF-8?q?fix(openai):=20=E4=BF=AE=E5=A4=8D=20Agent?= =?UTF-8?q?=20Identity=20CI=20=E9=97=AE=E9=A2=98?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../handler/admin/account_codex_import.go | 4 +- .../internal/service/openai_agent_identity.go | 6 ++- .../service/openai_agent_identity_test.go | 8 ++-- .../internal/service/openai_quota_service.go | 41 +++++++++---------- 4 files changed, 31 insertions(+), 28 deletions(-) diff --git a/backend/internal/handler/admin/account_codex_import.go b/backend/internal/handler/admin/account_codex_import.go index a6a07af0c1..57b97f8476 100644 --- a/backend/internal/handler/admin/account_codex_import.go +++ b/backend/internal/handler/admin/account_codex_import.go @@ -515,10 +515,10 @@ func normalizeCodexImportEntry(entry codexImportEntry) (*codexImportAccount, err 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 缺少必要字段") + return nil, errors.New("agent identity 缺少必要字段") } if err := service.ValidateOpenAIAgentIdentityPrivateKey(item.AgentPrivateKey); err != nil { - return nil, errors.New("Agent Identity private key 格式无效") + return nil, errors.New("agent identity private key 格式无效") } item.Credentials["auth_mode"] = service.OpenAIAuthModeAgentIdentity item.Credentials["agent_runtime_id"] = item.AgentRuntimeID diff --git a/backend/internal/service/openai_agent_identity.go b/backend/internal/service/openai_agent_identity.go index f375eac85c..aa96227f3d 100644 --- a/backend/internal/service/openai_agent_identity.go +++ b/backend/internal/service/openai_agent_identity.go @@ -261,7 +261,11 @@ func ensureAgentIdentityTaskForAccount(ctx context.Context, repo AccountReposito if credAccount.ID > 0 { candidate := &sync.Mutex{} actual, _ := agentIdentityTaskLocks.LoadOrStore(credAccount.ID, candidate) - sharedTaskMu = actual.(*sync.Mutex) + loadedTaskMu, ok := actual.(*sync.Mutex) + if !ok { + return errors.New("agent identity task lock has invalid type") + } + sharedTaskMu = loadedTaskMu } sharedTaskMu.Lock() defer sharedTaskMu.Unlock() diff --git a/backend/internal/service/openai_agent_identity_test.go b/backend/internal/service/openai_agent_identity_test.go index 12f5a7766a..a906f044f0 100644 --- a/backend/internal/service/openai_agent_identity_test.go +++ b/backend/internal/service/openai_agent_identity_test.go @@ -56,7 +56,9 @@ func TestBuildAgentAssertionMatchesCodexEnvelopeAndSignature(t *testing.T) { require.Equal(t, "2026-07-14T00:09:10Z", envelope.Timestamp) signature, err := base64.StdEncoding.DecodeString(envelope.Signature) require.NoError(t, err) - require.True(t, ed25519.Verify(key.privateKey.Public().(ed25519.PublicKey), []byte("runtime-test:task-test:2026-07-14T00:09:10Z"), signature)) + 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) { @@ -153,7 +155,7 @@ func TestEnsureAgentIdentityTaskPersistsAndRedactsCredentials(t *testing.T) { redacted[key] = value } } - require.NotContains(t, string(mustJSON(t, redacted)), privateKey) + require.NotContains(t, string(mustAgentIdentityJSON(t, redacted)), privateKey) } func TestEnsureAgentIdentityTaskSharesLockAcrossServicesForSameAccount(t *testing.T) { @@ -219,7 +221,7 @@ func (r *agentIdentityCredentialsRepo) UpdateCredentials(_ context.Context, _ in return nil } -func mustJSON(t *testing.T, value any) []byte { +func mustAgentIdentityJSON(t *testing.T, value any) []byte { t.Helper() encoded, err := json.Marshal(value) require.NoError(t, err) diff --git a/backend/internal/service/openai_quota_service.go b/backend/internal/service/openai_quota_service.go index 3669154b64..63ce173639 100644 --- a/backend/internal/service/openai_quota_service.go +++ b/backend/internal/service/openai_quota_service.go @@ -281,28 +281,25 @@ func (s *OpenAIQuotaService) ResetCredit(ctx context.Context, accountID int64) ( callCtx, cancel := context.WithTimeout(ctx, openaiQuotaUpstreamTimeout) defer cancel() var payload OpenAIQuotaResetResult - for { - 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). - SetBody(map[string]string{"redeem_request_id": redeemRequestID}). - SetSuccessResult(&payload). - Post(chatGPTRateLimitResetURL) - if err != nil { - return nil, infraerrors.Newf(http.StatusBadGateway, "OPENAI_QUOTA_RESET_REQUEST_FAILED", "upstream request failed: %v", err) - } - if !resp.IsSuccessState() { - status := resp.StatusCode - 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) - } - break + 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). + SetBody(map[string]string{"redeem_request_id": redeemRequestID}). + SetSuccessResult(&payload). + Post(chatGPTRateLimitResetURL) + if err != nil { + return nil, infraerrors.Newf(http.StatusBadGateway, "OPENAI_QUOTA_RESET_REQUEST_FAILED", "upstream request failed: %v", err) + } + if !resp.IsSuccessState() { + status := resp.StatusCode + 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) } slog.Info("openai_quota_reset_success", From 64850811224ea421caceb2f20fa36b3047372f37 Mon Sep 17 00:00:00 2001 From: cat Date: Tue, 14 Jul 2026 19:25:13 +0800 Subject: [PATCH 8/8] =?UTF-8?q?feat(frontend):=20=E6=A0=87=E6=98=8E=20Open?= =?UTF-8?q?AI=20=E8=AE=A4=E8=AF=81=E6=A8=A1=E5=BC=8F?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../account/OAuthAuthorizationFlow.vue | 2 +- .../components/common/PlatformTypeBadge.vue | 9 ++++++++ .../__tests__/PlatformTypeBadge.grok.spec.ts | 21 +++++++++++++++++++ .../src/i18n/locales/en/admin/accounts.ts | 12 +++++------ .../src/i18n/locales/zh/admin/accounts.ts | 12 +++++------ frontend/src/views/admin/AccountsView.vue | 7 +++++++ 6 files changed, 50 insertions(+), 13 deletions(-) diff --git a/frontend/src/components/account/OAuthAuthorizationFlow.vue b/frontend/src/components/account/OAuthAuthorizationFlow.vue index a23ec798d2..395b4ae923 100644 --- a/frontend/src/components/account/OAuthAuthorizationFlow.vue +++ b/frontend/src/components/account/OAuthAuthorizationFlow.vue @@ -276,7 +276,7 @@ - +
{ 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' diff --git a/frontend/src/components/common/__tests__/PlatformTypeBadge.grok.spec.ts b/frontend/src/components/common/__tests__/PlatformTypeBadge.grok.spec.ts index 753fb1a5f2..0bad3d9551 100644 --- a/frontend/src/components/common/__tests__/PlatformTypeBadge.grok.spec.ts +++ b/frontend/src/components/common/__tests__/PlatformTypeBadge.grok.spec.ts @@ -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') + }) +}) diff --git a/frontend/src/i18n/locales/en/admin/accounts.ts b/frontend/src/i18n/locales/en/admin/accounts.ts index e64a51ea20..64a86dad17 100644 --- a/frontend/src/i18n/locales/en/admin/accounts.ts +++ b/frontend/src/i18n/locales/en/admin/accounts.ts @@ -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}', diff --git a/frontend/src/i18n/locales/zh/admin/accounts.ts b/frontend/src/i18n/locales/zh/admin/accounts.ts index 49d4b8564f..65fb089acd 100644 --- a/frontend/src/i18n/locales/zh/admin/accounts.ts +++ b/frontend/src/i18n/locales/zh/admin/accounts.ts @@ -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}', diff --git a/frontend/src/views/admin/AccountsView.vue b/frontend/src/views/admin/AccountsView.vue index e8bac39af4..140b7c934f 100644 --- a/frontend/src/views/admin/AccountsView.vue +++ b/frontend/src/views/admin/AccountsView.vue @@ -234,6 +234,7 @@
@@ -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