Merge pull request #4359 from catoncat/agent/fix-agent-identity-import-expiry

fix(openai): make Agent Identity usable end to end
This commit is contained in:
Wesley Liddick authored and GitHub committed 2026-07-15 20:55:05 +08:00
commit 34015a1791
12 files changed
+424 -62

No files matched your search

@@ -1,6 +1,7 @@
package admin
import (
"context"
"crypto/ed25519"
"crypto/rand"
"crypto/x509"
@@ -45,3 +46,38 @@ func TestNormalizeCodexImportEntryAcceptsAgentIdentityAuthJSON(t *testing.T) {
require.NotContains(t, item.Credentials, "refresh_token")
require.NotEmpty(t, item.WarningTexts)
}
func TestImportCodexSessionsCreatesAgentIdentityWithoutOAuthExpiry(t *testing.T) {
_, privateKey, err := ed25519.GenerateKey(rand.Reader)
require.NoError(t, err)
der, err := x509.MarshalPKCS8PrivateKey(privateKey)
require.NoError(t, err)
svc := newCodexImportMemoryAdminService(nil)
handler := NewAccountHandler(svc, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil)
result, err := handler.importCodexSessions(context.Background(), CodexSessionImportRequest{
SkipDefaultGroupBind: boolPtr(true),
}, []codexImportEntry{{
Index: 1,
Value: map[string]any{
"auth_mode": "agentIdentity",
"agent_identity": map[string]any{
"agent_runtime_id": "runtime-import",
"agent_private_key": base64.StdEncoding.EncodeToString(der),
"task_id": "task-import",
"account_id": "account-import",
"chatgpt_user_id": "user-import",
},
},
}})
require.NoError(t, err)
require.Equal(t, 1, result.Created)
require.Zero(t, result.Failed)
require.Len(t, svc.createdAccounts, 1)
created := svc.createdAccounts[0]
require.Nil(t, created.ExpiresAt)
require.Nil(t, created.AutoPauseOnExpired)
require.Equal(t, service.OpenAIAuthModeAgentIdentity, created.Credentials["auth_mode"])
require.NotContains(t, created.Credentials, "access_token")
require.NotContains(t, created.Credentials, "refresh_token")
}
@@ -774,6 +774,12 @@ func resolveCodexImportExpiry(req CodexSessionImportRequest, item *codexImportAc
if item == nil {
return nil, nil, nil, nil, errors.New("导入项为空")
}
// Agent Identity has no OAuth access-token lifetime. Its runtime/task
// lifecycle is handled by the upstream task recovery path, so it must not
// be rejected or auto-paused by the OAuth import expiry policy.
if item.IsAgentIdentity {
return nil, nil, nil, nil, nil
}
var requestExpiresAt *time.Time
if req.ExpiresAt != nil && *req.ExpiresAt > 0 {
@@ -42,8 +42,10 @@ type CodexModelsManifest struct {
}
type codexModelsManifestUpstreamError struct {
err error
retryable bool
err error
retryable bool
statusCode int
body []byte
}
func (e *codexModelsManifestUpstreamError) Error() string { return e.err.Error() }
@@ -135,6 +137,7 @@ type codexModelsManifestRequest struct {
proxyURL string
accountID int64
credentialAccountID int64
credentialAccount *Account
accountConcurrency int
useAPIKeyUpstream bool
}
@@ -308,15 +311,42 @@ func (s *OpenAIGatewayService) FetchCodexModelsManifest(ctx context.Context, acc
proxyURL: proxyURL,
accountID: account.ID,
credentialAccountID: credAccount.ID,
credentialAccount: credAccount,
accountConcurrency: account.Concurrency,
useAPIKeyUpstream: useAPIKeyUpstream,
}
if useAPIKeyUpstream {
return s.fetchCachedAPIKeyCodexModelsManifest(ctx, request, ifNoneMatch)
}
manifest, fetchErr := s.fetchCodexModelsManifestUpstream(ctx, request, ifNoneMatch)
if !credAccount.IsOpenAIAgentIdentity() || !isAgentIdentityTaskInvalidCodexModelsError(fetchErr) {
return manifest, fetchErr
}
expectedTaskID := strings.TrimSpace(credAccount.GetCredential("task_id"))
if recoverErr := s.recoverAgentIdentityTask(ctx, credAccount, expectedTaskID); recoverErr != nil {
return nil, infraerrors.Newf(http.StatusBadGateway, "OPENAI_CODEX_MODELS_AUTH_FAILED", "agent identity task recovery failed: %v", recoverErr)
}
authHeaders, authErr := s.buildOpenAIAuthenticationHeaders(ctx, credAccount, "")
if authErr != nil {
return nil, infraerrors.Newf(http.StatusBadGateway, "OPENAI_CODEX_MODELS_AUTH_FAILED", "build Codex models authentication after task recovery: %v", authErr)
}
request.headers.Del("Authorization")
request.headers.Del("ChatGPT-Account-ID")
for key, values := range authHeaders {
for _, value := range values {
request.headers.Add(key, value)
}
}
setOpenAIChatGPTAccountHeaders(request.headers, credAccount)
return s.fetchCodexModelsManifestUpstream(ctx, request, ifNoneMatch)
}
func isAgentIdentityTaskInvalidCodexModelsError(err error) bool {
var upstreamErr *codexModelsManifestUpstreamError
return errors.As(err, &upstreamErr) &&
isAgentIdentityTaskInvalidHTTPResponse(upstreamErr.statusCode, upstreamErr.body)
}
func (s *OpenAIGatewayService) fetchCachedAPIKeyCodexModelsManifest(ctx context.Context, request codexModelsManifestRequest, ifNoneMatch string) (*CodexModelsManifest, error) {
if err := ctx.Err(); err != nil {
return nil, err
@@ -410,12 +440,15 @@ 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
}
return nil, &codexModelsManifestUpstreamError{
err: infraerrors.Newf(http.StatusBadGateway, "OPENAI_CODEX_MODELS_UPSTREAM_FAILED", "codex models manifest upstream error %d: %s", resp.StatusCode, message),
err: infraerrors.Newf(http.StatusBadGateway, "OPENAI_CODEX_MODELS_UPSTREAM_FAILED", "codex models manifest upstream error %d: %s", resp.StatusCode, message),
statusCode: resp.StatusCode,
body: body,
retryable: resp.StatusCode == http.StatusTooManyRequests ||
(resp.StatusCode >= http.StatusInternalServerError && resp.StatusCode < 600),
}
@@ -3,6 +3,7 @@ package service
import (
"context"
"errors"
"fmt"
"io"
"net"
"net/http"
@@ -17,6 +18,7 @@ import (
"github.com/Wei-Shaw/sub2api/internal/config"
infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
"github.com/Wei-Shaw/sub2api/internal/pkg/tlsfingerprint"
"github.com/stretchr/testify/require"
"golang.org/x/net/http2"
)
@@ -242,6 +244,93 @@ func TestFetchCodexModelsManifestAgentIdentityUsesAssertionWithoutOAuthToken(t *
}
}
func TestFetchCodexModelsManifestAgentIdentityRecoversInvalidTaskOnce(t *testing.T) {
key, privateKey := newTestAgentIdentityKey(t)
account := &Account{
ID: 4,
Platform: PlatformOpenAI,
Type: AccountTypeOAuth,
Credentials: map[string]any{
"auth_mode": OpenAIAuthModeAgentIdentity,
"agent_runtime_id": key.runtimeID,
"agent_private_key": privateKey,
"task_id": "task-models-old",
"chatgpt_account_id": "acc-agent-recovery",
},
}
repo := &stubQuotaAccountRepo{accounts: map[int64]*Account{account.ID: account}}
modelsCalls := 0
registerCalls := 0
var assertions []string
server := 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-models-new"}`))
return
}
modelsCalls++
assertions = append(assertions, r.Header.Get("Authorization"))
if modelsCalls == 1 {
w.WriteHeader(http.StatusUnauthorized)
_, _ = w.Write([]byte(`{"error":{"code":"invalid_task_id"}}`))
return
}
_, _ = w.Write([]byte(`{"models":[]}`))
}))
defer server.Close()
originalModelsURL := chatgptCodexModelsURL
chatgptCodexModelsURL = server.URL
t.Cleanup(func() { chatgptCodexModelsURL = originalModelsURL })
originalAuthBase := openAIAgentIdentityAuthAPIBaseURL
openAIAgentIdentityAuthAPIBaseURL = server.URL
t.Cleanup(func() { openAIAgentIdentityAuthAPIBaseURL = originalAuthBase })
s := &OpenAIGatewayService{accountRepo: repo}
manifest, err := s.FetchCodexModelsManifest(context.Background(), account, "0.137.0", "")
require.NoError(t, err)
require.Equal(t, `{"models":[]}`, string(manifest.Body))
require.Equal(t, 2, modelsCalls)
require.Equal(t, 1, registerCalls)
require.Len(t, assertions, 2)
require.Equal(t, "task-models-old", decodeAgentAssertionTask(t, assertions[0]))
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) {
@@ -22,8 +22,6 @@ 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"
@@ -160,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)
}
@@ -175,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
@@ -209,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
@@ -260,9 +258,6 @@ 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)
@@ -282,26 +277,38 @@ 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
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)
for recovered := false; ; {
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)
}
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, expectedTaskID); 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",
@@ -387,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")
}
@@ -404,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 {
@@ -424,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 {
@@ -176,21 +176,120 @@ func TestResetCreditShadowRejected(t *testing.T) {
"shadow ResetCredit 应映射为 409 Conflict 而非 500")
}
func TestResetCreditAgentIdentityRejectedBeforeUpstream(t *testing.T) {
func TestResetCreditAgentIdentityUsesAssertionAndRecoversInvalidTaskOnce(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: 201,
Platform: PlatformOpenAI,
Type: AccountTypeOAuth,
Credentials: map[string]any{
"auth_mode": OpenAIAuthModeAgentIdentity,
"auth_mode": OpenAIAuthModeAgentIdentity,
"agent_runtime_id": "runtime-reset-recovery",
"agent_private_key": base64.StdEncoding.EncodeToString(der),
"task_id": "task-reset-old",
"chatgpt_account_id": "account-reset-recovery",
},
}
repo := &stubQuotaAccountRepo{accounts: map[int64]*Account{account.ID: account}}
svc := &OpenAIQuotaService{accountRepo: repo}
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-new"}`))
return
}
resetCalls++
assertions = append(assertions, r.Header.Get("authorization"))
require.Equal(t, "account-reset-recovery", r.Header.Get("chatgpt-account-id"))
if resetCalls == 1 {
w.WriteHeader(http.StatusUnauthorized)
_, _ = w.Write([]byte(`{"error":{"code":"invalid_task_id"}}`))
return
}
_, _ = w.Write([]byte(`{"code":"ok","windows_reset":2}`))
}))
defer srv.Close()
oldBase := openAIAgentIdentityAuthAPIBaseURL
openAIAgentIdentityAuthAPIBaseURL = srv.URL
t.Cleanup(func() { openAIAgentIdentityAuthAPIBaseURL = oldBase })
_, err := svc.ResetCredit(context.Background(), account.ID)
require.ErrorIs(t, err, ErrAgentIdentityResetNotSupported)
require.Equal(t, http.StatusConflict, infraerrors.Code(err))
invalidator := &agentIdentityWSInvalidationRecorder{}
svc := NewOpenAIQuotaService(repo, nil, nil, newQuotaRedirectingFactory(srv))
svc.agentIdentityWS = invalidator
result, err := svc.ResetCredit(context.Background(), account.ID)
require.NoError(t, err)
require.Equal(t, "ok", result.Code)
require.Equal(t, 2, result.WindowsReset)
require.Equal(t, 2, resetCalls)
require.Equal(t, 1, registerCalls)
require.Len(t, assertions, 2)
require.True(t, strings.HasPrefix(assertions[0], "AgentAssertion "))
require.True(t, strings.HasPrefix(assertions[1], "AgentAssertion "))
require.NotEqual(t, assertions[0], assertions[1])
require.Equal(t, "task-reset-new", account.GetCredential("task_id"))
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 ──────────────────────────────