mirror of
https://github.com/Wei-Shaw/sub2api.git
synced 2026-11-09 00:57:17 +00:00
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:
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 ──────────────────────────────
|
||||
|
||||
Reference in new issue
Block a user