fix(openai): 修复 Agent Identity 审查问题

This commit is contained in:
cat
2026-07-15 18:30:16 +08:00
parent 12e6e70775
commit 2147da682b
6 changed files with 129 additions and 18 deletions
@@ -137,6 +137,7 @@ type codexModelsManifestRequest struct {
proxyURL string
accountID int64
credentialAccountID int64
credentialAccount *Account
accountConcurrency int
useAPIKeyUpstream bool
}
@@ -310,6 +311,7 @@ func (s *OpenAIGatewayService) FetchCodexModelsManifest(ctx context.Context, acc
proxyURL: proxyURL,
accountID: account.ID,
credentialAccountID: credAccount.ID,
credentialAccount: credAccount,
accountConcurrency: account.Concurrency,
useAPIKeyUpstream: useAPIKeyUpstream,
}
@@ -438,6 +440,7 @@ func (s *OpenAIGatewayService) fetchCodexModelsManifestUpstream(ctx context.Cont
}
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
body, _ := io.ReadAll(io.LimitReader(resp.Body, 2048))
body = s.redactAgentIdentitySensitiveBody(reqCtx, request.credentialAccount, body)
message := strings.TrimSpace(string(body))
if message == "" {
message = resp.Status
@@ -3,6 +3,7 @@ package service
import (
"context"
"errors"
"fmt"
"io"
"net"
"net/http"
@@ -297,6 +298,39 @@ func TestFetchCodexModelsManifestAgentIdentityRecoversInvalidTaskOnce(t *testing
require.Equal(t, "task-models-new", decodeAgentAssertionTask(t, assertions[1]))
}
func TestFetchCodexModelsManifestAgentIdentityRedactsUpstreamErrors(t *testing.T) {
key, privateKey := newTestAgentIdentityKey(t)
account := &Account{
ID: 5,
Platform: PlatformOpenAI,
Type: AccountTypeOAuth,
Credentials: map[string]any{
"auth_mode": OpenAIAuthModeAgentIdentity,
"agent_runtime_id": key.runtimeID,
"agent_private_key": privateKey,
"task_id": key.taskID,
"chatgpt_account_id": "acc-agent-redaction",
},
}
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
w.WriteHeader(http.StatusBadRequest)
_, _ = fmt.Fprintf(w, `{"error":"%s %s %s AgentAssertion leaked"}`, key.runtimeID, key.taskID, privateKey)
}))
defer server.Close()
original := chatgptCodexModelsURL
chatgptCodexModelsURL = server.URL
t.Cleanup(func() { chatgptCodexModelsURL = original })
s := &OpenAIGatewayService{}
_, err := s.FetchCodexModelsManifest(context.Background(), account, "0.137.0", "")
require.Error(t, err)
require.NotContains(t, err.Error(), key.runtimeID)
require.NotContains(t, err.Error(), key.taskID)
require.NotContains(t, err.Error(), privateKey)
require.NotContains(t, err.Error(), "AgentAssertion leaked")
require.Contains(t, err.Error(), "[redacted]")
}
func TestFetchCodexModelsManifestDefaultClientVersion(t *testing.T) {
var gotClientVersion string
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
@@ -158,7 +158,7 @@ func (s *OpenAIQuotaService) QueryUsage(ctx context.Context, accountID int64) (*
var payload OpenAIQuotaUsage
for recovered := false; ; {
quotaHeaders, headerErr := s.buildCodexQuotaHeaders(callCtx, accountID, accessToken, chatGPTAccountID, fedRAMP)
quotaHeaders, expectedTaskID, 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)
}
@@ -173,7 +173,7 @@ func (s *OpenAIQuotaService) QueryUsage(ctx context.Context, accountID int64) (*
if !resp.IsSuccessState() {
if agentIdentity && !recovered && isAgentIdentityTaskInvalidHTTPResponse(resp.StatusCode, []byte(resp.String())) {
recovered = true
if err := s.recoverAgentIdentityTask(ctx, accountID); err != nil {
if err := s.recoverAgentIdentityTask(ctx, accountID, expectedTaskID); err != nil {
return nil, infraerrors.Newf(http.StatusBadGateway, "OPENAI_QUOTA_AUTH_FAILED", "agent identity task recovery failed: %v", err)
}
continue
@@ -207,7 +207,7 @@ func (s *OpenAIQuotaService) QueryUsage(ctx context.Context, accountID int64) (*
}
func (s *OpenAIQuotaService) queryResetCreditDetails(ctx context.Context, client *req.Client, accessToken, chatGPTAccountID string, fedRAMP bool, accountID int64) *openAIRateLimitResetCreditDetails {
quotaHeaders, headerErr := s.buildCodexQuotaHeaders(ctx, accountID, accessToken, chatGPTAccountID, fedRAMP)
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
@@ -281,7 +281,7 @@ func (s *OpenAIQuotaService) ResetCredit(ctx context.Context, accountID int64) (
var payload OpenAIQuotaResetResult
for recovered := false; ; {
headers, headerErr := s.buildCodexQuotaHeaders(callCtx, accountID, accessToken, chatGPTAccountID, fedRAMP)
headers, expectedTaskID, 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)
}
@@ -298,7 +298,7 @@ func (s *OpenAIQuotaService) ResetCredit(ctx context.Context, accountID int64) (
if !resp.IsSuccessState() {
if agentIdentity && !recovered && isAgentIdentityTaskInvalidHTTPResponse(resp.StatusCode, []byte(resp.String())) {
recovered = true
if err := s.recoverAgentIdentityTask(ctx, accountID); err != nil {
if err := s.recoverAgentIdentityTask(ctx, accountID, expectedTaskID); err != nil {
return nil, infraerrors.Newf(http.StatusBadGateway, "OPENAI_QUOTA_AUTH_FAILED", "agent identity task recovery failed: %v", err)
}
continue
@@ -394,7 +394,7 @@ func (s *OpenAIQuotaService) prepareUpstreamCall(ctx context.Context, accountID
return accessToken, chatGPTAccountID, proxyURL, fedRAMP, nil
}
func (s *OpenAIQuotaService) recoverAgentIdentityTask(ctx context.Context, accountID int64) error {
func (s *OpenAIQuotaService) recoverAgentIdentityTask(ctx context.Context, accountID int64, expectedTaskID string) error {
if s == nil || s.accountRepo == nil {
return fmt.Errorf("account repository is unavailable")
}
@@ -411,7 +411,7 @@ func (s *OpenAIQuotaService) recoverAgentIdentityTask(ctx context.Context, accou
if !account.IsOpenAIAgentIdentity() {
return nil
}
return ensureAgentIdentityTaskForAccount(ctx, s.accountRepo, s.agentIdentityWS, &s.agentIdentityTaskMu, account, account.GetCredential("task_id"))
return ensureAgentIdentityTaskForAccount(ctx, s.accountRepo, s.agentIdentityWS, &s.agentIdentityTaskMu, account, expectedTaskID)
}
func (s *OpenAIQuotaService) isAgentIdentityAccount(ctx context.Context, accountID int64) bool {
@@ -431,41 +431,41 @@ func (s *OpenAIQuotaService) isAgentIdentityAccount(ctx context.Context, account
return account.IsOpenAIAgentIdentity()
}
func (s *OpenAIQuotaService) buildCodexQuotaHeaders(ctx context.Context, accountID int64, accessToken, chatGPTAccountID string, fedRAMP bool) (map[string]string, error) {
func (s *OpenAIQuotaService) buildCodexQuotaHeaders(ctx context.Context, accountID int64, accessToken, chatGPTAccountID string, fedRAMP bool) (map[string]string, string, error) {
headers := buildCodexCommonHeaders(accessToken, chatGPTAccountID, fedRAMP)
if s == nil || s.accountRepo == nil {
return headers, 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 nil, "", fmt.Errorf("agent identity account credentials are unavailable")
}
return headers, nil
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")
return nil, "", fmt.Errorf("agent identity shadow credentials are unavailable")
}
}
if !account.IsOpenAIAgentIdentity() {
return headers, nil
return headers, "", nil
}
if err := ensureAgentIdentityTaskForAccount(ctx, s.accountRepo, s.agentIdentityWS, &s.agentIdentityTaskMu, account, ""); err != nil {
return nil, err
return nil, "", err
}
key, err := agentIdentityKeyFromAccount(account)
if err != nil {
return nil, err
return nil, "", err
}
assertion, err := buildAgentAssertion(key, time.Now())
if err != nil {
return nil, err
return nil, "", err
}
headers["authorization"] = assertion
return headers, nil
return headers, key.taskID, nil
}
func (s *OpenAIQuotaService) redactQuotaErrorBody(ctx context.Context, accountID int64, body string) string {
@@ -237,6 +237,61 @@ func TestResetCreditAgentIdentityUsesAssertionAndRecoversInvalidTaskOnce(t *test
require.Equal(t, []int64{account.ID}, invalidator.accountIDs)
}
func TestResetCreditAgentIdentityReusesConcurrentlyRecoveredTask(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: 202,
Platform: PlatformOpenAI,
Type: AccountTypeOAuth,
Credentials: map[string]any{
"auth_mode": OpenAIAuthModeAgentIdentity,
"agent_runtime_id": "runtime-reset-concurrent",
"agent_private_key": base64.StdEncoding.EncodeToString(der),
"task_id": "task-reset-old",
"chatgpt_account_id": "account-reset-concurrent",
},
}
repo := &stubQuotaAccountRepo{accounts: map[int64]*Account{account.ID: account}}
resetCalls := 0
registerCalls := 0
var assertions []string
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-reset-unexpected"}`))
return
}
resetCalls++
assertions = append(assertions, r.Header.Get("authorization"))
if resetCalls == 1 {
credentials := shallowCopyMap(account.Credentials)
credentials["task_id"] = "task-reset-concurrent"
account.Credentials = credentials
w.WriteHeader(http.StatusUnauthorized)
_, _ = w.Write([]byte(`{"error":{"code":"invalid_task_id"}}`))
return
}
_, _ = w.Write([]byte(`{"code":"ok","windows_reset":1}`))
}))
defer srv.Close()
oldBase := openAIAgentIdentityAuthAPIBaseURL
openAIAgentIdentityAuthAPIBaseURL = srv.URL
t.Cleanup(func() { openAIAgentIdentityAuthAPIBaseURL = oldBase })
svc := NewOpenAIQuotaService(repo, nil, nil, newQuotaRedirectingFactory(srv))
result, err := svc.ResetCredit(context.Background(), account.ID)
require.NoError(t, err)
require.Equal(t, "ok", result.Code)
require.Equal(t, 2, resetCalls)
require.Zero(t, registerCalls)
require.Equal(t, "task-reset-old", decodeAgentAssertionTask(t, assertions[0]))
require.Equal(t, "task-reset-concurrent", decodeAgentAssertionTask(t, assertions[1]))
}
// ── Part B: prepareUpstreamCall 影子 resolve ──────────────────────────────
// TestPrepareUpstreamCallShadowResolve 验证影子账号(200QueryUsage 时:
@@ -5515,7 +5515,10 @@ const isAgentIdentityImportContent = (content: string) => {
if (Array.isArray(value)) return value.length > 0 && value.every(isAgentIdentityValue)
if (!value || typeof value !== 'object') return false
const record = value as Record<string, unknown>
return record.auth_mode === 'agentIdentity' && !!record.agent_identity && typeof record.agent_identity === 'object'
const authMode = record.auth_mode ?? record.authMode
const agentIdentity = record.agent_identity ?? record.agentIdentity
return (typeof authMode === 'string' && authMode.toLowerCase() === 'agentidentity')
|| (!!agentIdentity && typeof agentIdentity === 'object')
}
try {
@@ -171,6 +171,22 @@ describe('CreateAccountModal OpenAI long-context billing', () => {
expect(flow.props('titleOverride')).toBe('Agent Identity')
})
it.each([
['camelCase', { authMode: 'agentIdentity', agentIdentity: { agentRuntimeId: 'runtime' } }],
['nested identity without auth_mode', { agent_identity: { agent_runtime_id: 'runtime' } }],
])('accepts backend-compatible %s Agent Identity imports', async (_name, content) => {
const wrapper = mountModal()
await selectButtonByText(wrapper, 'OpenAI')
await wrapper.get('[data-testid="openai-account-type-agent-identity"]').trigger('click')
await wrapper.get('form#create-account-form input[type="text"]').setValue('Agent Identity')
await wrapper.get('form#create-account-form').trigger('submit.prevent')
wrapper.getComponent(OAuthAuthorizationFlowStub).vm.$emit('import-codex-session', JSON.stringify(content))
await flushPromises()
expect(importCodexSessionMock).toHaveBeenCalledTimes(1)
})
it('shows Codex PAT as a separate OpenAI account type', async () => {
const wrapper = mountModal()
await selectButtonByText(wrapper, 'OpenAI')